Coverage for src/proximal_lq/proximal.py: 100%

1"""Fast solver for 0.5 ||mat @ x - vec||^2 s. t. {x >= 0, sum(x) = 1}. 

2 

3This module implements proximal gradient descent for constrained linear least squares 

4optimization on the probability simplex. The algorithm is based on iterative projection 

5using the efficient simplex projection from Duchi et al. (2008). 

6 

7References: 

8---------- 

9Duchi, J., Shalev-Shwartz, S., Singer, Y., & Chandra, T. (2008). 

10"Efficient Projections onto the l1-Ball for Learning in High Dimensions." 

11Proceedings of the 25th International Conference on Machine Learning (ICML). 

12""" 

13 

14from __future__ import annotations 

15 

16from typing import TYPE_CHECKING 

17 

18import numpy as np 

19 

20if TYPE_CHECKING: 

21 from numpy.typing import NDArray 

22 

23 

24def proj_simplex( 

25 vec: NDArray[np.floating], 

26 rad: float = 1.0, 

27) -> NDArray[np.floating]: 

28 """Project a vector onto the probability simplex. 

29 

30 This function computes the Euclidean projection of a given vector onto the probability 

31 simplex. The simplex is defined as the set of non-negative vectors that sum to a 

32 given radius, typically 1. The projection ensures that the resulting vector satisfies 

33 these constraints. 

34 

35 The algorithm is based on Duchi et al. (2008) "Efficient Projections onto the 

36 l1-Ball for Learning in High Dimensions". 

37 

38 Parameters 

39 ---------- 

40 vec : NDArray[np.floating] 

41 Input vector that is to be projected onto the simplex. 

42 rad : float, optional 

43 Radius of the simplex. The projected vector will have components summing 

44 to this value. Default is 1.0. 

45 

46 Returns: 

47 ------- 

48 NDArray[np.floating] 

49 The projected vector that lies on the probability simplex. 

50 

51 Raises: 

52 ------ 

53 ValueError 

54 If the input vector is empty. 

55 

56 Examples: 

57 -------- 

58 >>> import numpy as np 

59 >>> vec = np.array([1.0, 2.0, 3.0]) 

60 >>> result = proj_simplex(vec) 

61 >>> bool(np.isclose(result.sum(), 1.0)) 

62 True 

63 >>> bool(np.all(result >= 0)) 

64 True 

65 

66 """ 

67 if vec.size == 0: 

68 msg = "vec must be non-empty" 

69 raise ValueError(msg) 

70 

71 # Duchi et al. (2008): sort descending, find the largest index rho whose 

72 # sorted value still exceeds the running mean, then shift by that mean and 

73 # clip negatives to zero. 

74 sorted_desc = np.sort(vec)[::-1] 

75 running_mean = (np.cumsum(sorted_desc) - rad) / np.arange(1, len(vec) + 1) 

76 rho = np.max(np.where(sorted_desc > running_mean)[0]) 

77 threshold = running_mean[rho] 

78 result: NDArray[np.floating] = np.maximum(vec - threshold, 0) 

79 return result 

80 

81 

82def _validate_inputs(mat: NDArray[np.floating], vec: NDArray[np.floating]) -> None: 

83 """Validate the shapes of ``mat`` and ``vec`` for :func:`prox_gradient`. 

84 

85 Parameters 

86 ---------- 

87 mat : NDArray[np.floating] 

88 The matrix argument; must be a non-empty 2-D array. 

89 vec : NDArray[np.floating] 

90 The vector argument; must be a non-empty 1-D array whose length matches 

91 ``mat.shape[0]``. 

92 

93 Raises: 

94 ------ 

95 ValueError 

96 If ``mat`` is not 2-D, ``vec`` is not 1-D, either input is empty, or 

97 ``vec.shape[0]`` does not match ``mat.shape[0]``. 

98 

99 """ 

100 if mat.ndim != 2: 

101 msg = f"mat must be a 2-D array (n_samples, n_features), got {mat.ndim}-D with shape {mat.shape}" 

102 raise ValueError(msg) 

103 if vec.ndim != 1: 

104 msg = f"vec must be a 1-D array (n_samples,), got {vec.ndim}-D with shape {vec.shape}" 

105 raise ValueError(msg) 

106 if min(mat.size, vec.size) == 0: 

107 msg = f"mat and vec must be non-empty, got shapes {mat.shape} and {vec.shape}" 

108 raise ValueError(msg) 

109 if vec.shape[0] != mat.shape[0]: 

110 msg = ( 

111 f"vec length ({vec.shape[0]}) must match mat.shape[0] ({mat.shape[0]}); " 

112 f"vec has one entry per row (sample) of mat" 

113 ) 

114 raise ValueError(msg) 

115 

116 

117def _iterate( 

118 prim_var: NDArray[np.floating], 

119 sym_mat: NDArray[np.floating], 

120 out_prod: NDArray[np.floating], 

121 step: float, 

122 eps_rel: float, 

123 max_iter: int, 

124) -> NDArray[np.floating]: 

125 """Run the projected-gradient iteration loop until convergence. 

126 

127 Parameters 

128 ---------- 

129 prim_var : NDArray[np.floating] 

130 The initial primal variable. 

131 sym_mat : NDArray[np.floating] 

132 The Gram matrix ``mat.T @ mat``. 

133 out_prod : NDArray[np.floating] 

134 The vector ``mat.T @ vec``. 

135 step : float 

136 The gradient step size. 

137 eps_rel : float 

138 The relative error threshold for the stopping criterion. 

139 max_iter : int 

140 The maximum number of iterations. 

141 

142 Returns: 

143 ------- 

144 NDArray[np.floating] 

145 The primal variable after convergence or reaching ``max_iter``. 

146 

147 """ 

148 err_rel = eps_rel + 1 

149 for _ in range(max_iter): 

150 if err_rel <= eps_rel: 

151 break 

152 prim_var_new = proj_simplex(prim_var - step * (sym_mat @ prim_var - out_prod)) 

153 err_rel = float(np.linalg.norm(prim_var - prim_var_new, 2)) 

154 prim_var = prim_var_new.copy() 

155 return prim_var 

156 

157 

158def prox_gradient( 

159 mat: NDArray[np.floating], 

160 vec: NDArray[np.floating], 

161 eps_rel: float = 1e-6, 

162 max_iter: int = 1000, 

163 seed: int | None = None, 

164) -> NDArray[np.floating]: 

165 """Perform proximal gradient descent to solve a constrained optimization problem. 

166 

167 Solves the optimization problem: 

168 minimize 0.5 ||mat @ x - vec||^2 

169 subject to x >= 0, sum(x) = 1 

170 

171 The function uses proximal gradient descent with simplex projection to find 

172 the solution. The step size is determined by the Lipschitz constant of the 

173 gradient. 

174 

175 Parameters 

176 ---------- 

177 mat : NDArray[np.floating] 

178 A matrix of shape (n_samples, n_features) used in the optimization 

179 problem. 

180 vec : NDArray[np.floating] 

181 A vector of shape (n_samples,) used in the optimization problem. 

182 Its length must equal ``mat.shape[0]``. 

183 eps_rel : float, optional 

184 The relative error threshold for stopping criteria. Default is 1e-6. 

185 max_iter : int, optional 

186 The maximum number of iterations for the algorithm. Default is 1000. 

187 seed : int | None, optional 

188 Seed for the random number generator used to initialise the primal 

189 variable. Pass an integer for reproducible results (the problem is 

190 convex, so the optimum is independent of the seed). Default is None, 

191 which draws a fresh, unseeded initialisation. 

192 

193 Returns: 

194 ------- 

195 NDArray[np.floating] 

196 The solution vector of shape (n_features,) obtained after the 

197 optimization process. 

198 

199 Raises: 

200 ------ 

201 ValueError 

202 If ``mat`` is not a 2-D array, ``vec`` is not a 1-D array, either input 

203 is empty, or ``vec.shape[0]`` does not match ``mat.shape[0]``. 

204 

205 Examples: 

206 -------- 

207 >>> import numpy as np 

208 >>> mat = np.array([[1.0, 0.5], [0.5, 1.0]]) 

209 >>> vec = np.ones(2) 

210 >>> result = prox_gradient(mat, vec) 

211 >>> bool(np.isclose(result.sum(), 1.0)) 

212 True 

213 

214 """ 

215 _validate_inputs(mat, vec) 

216 

217 rng = np.random.default_rng(seed) 

218 prim_var: NDArray[np.floating] = np.asarray(rng.standard_normal(size=mat.shape[1])) 

219 sym_mat = mat.T @ mat 

220 lip = np.linalg.norm(sym_mat, 2) # Lipschitz constant of the gradient 

221 step = float(0.5 / lip) if abs(lip) > 1e-15 else 1.0 

222 

223 out_prod = mat.T @ vec 

224 return _iterate(prim_var, sym_mat, out_prod, step, eps_rel, max_iter)