Coverage for src/lsqfitgp/_linalg/

1# lsqfitgp/_linalg/_seqalg.py

5# This file is part of lsqfitgp.

7# lsqfitgp is free software: you can redistribute it and/or modify

8# it under the terms of the GNU General Public License as published by

9# the Free Software Foundation, either version 3 of the License, or

10# (at your option) any later version.

11#

12# lsqfitgp is distributed in the hope that it will be useful,

13# but WITHOUT ANY WARRANTY; without even the implied warranty of

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

15# GNU General Public License for more details.

16#

17# You should have received a copy of the GNU General Public License

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>.

20import abc 3 ctx1bca

22from jax import lax 3 ctx1bca

23from jax import numpy as jnp 3 ctx1bca

25from lsqfitgp._linalg import _pytree 3 ctx1bca

28class SequentialOperation(_pytree.AutoPyTree, metaclass=abc.ABCMeta): 3 ctx1bca

29 """See jax.lax.fori_loop for semantics."""

31 @abc.abstractmethod 3 ctx1bca

32 def __init__(self, *args): # pragma: no cover 3 ctx1bca

33 pass

35 @property 3 ctx1bca

36 @abc.abstractmethod 3 ctx1bca

37 def inputs(self): # pragma: no cover 3 ctx1bca

38 """Tuple of indices of other ops to be used as input."""

39 pass

41 @abc.abstractmethod 3 ctx1bca

42 def init(self, n, *inputs): # pragma: no cover 3 ctx1bca

43 """Initialize before the cycle starts, with the requested inputs."""

44 pass

46 @abc.abstractmethod 3 ctx1bca

47 def iter_out(self, i): # pragma: no cover 3 ctx1bca

48 """

49 Output passed to other ops who request it through `inputs`.

51 Guaranteed to be called after `init`.

52 """

53 pass

55 @abc.abstractmethod 3 ctx1bca

56 def iter(self, i, *inputs): # pragma: no cover 3 ctx1bca

57 """Update for iteration."""

58 pass

60 @abc.abstractmethod 3 ctx1bca

61 def finalize(self): # pragma: no cover 3 ctx1bca

62 """Return final product."""

63 pass

66def sequential_algorithm(n, ops): 3 ctx1bca

67 """

68 Define and execute a sequential algorithm on matrices.

70 Parameters

71 ----------

72 n : int

73 Number of steps of the algorithm.

74 ops : list of SequentialOperation

75 Instantiated `SequentialOperation`s that represent the algorithm.

77 Return

78 ------

79 results : tuple

80 The sequence of final outputs of each operation in `ops`.

81 """

82 for i, op in enumerate(ops): 2 ctx1ba

83 inputs = op.inputs 2 ctx1ba

84 if any(j >= i for j in inputs): 2 ctx1ba

85 msg = f'{i}-th operation {op.__class__.__name__} requested inputs {inputs!r} with forward references' 1 ctx1a

86 raise ValueError(msg) 1 ctx1a

87 args = (ops[j].iter_out(0) for j in inputs) 2 ctx1ba

88 op.init(n, *args) 2 ctx1ba

90 def body_fun(i, ops): 2 ctx1ba

91 for op in ops: 2 ctx1ba

92 args = (ops[j].iter_out(i) for j in op.inputs) 2 ctx1ba

93 op.iter(i, *args) 2 ctx1ba

94 return ops 2 ctx1ba

96 ops = lax.fori_loop(1, n, body_fun, ops) 2 ctx1ba

97 return tuple(op.finalize() for op in ops) 2 ctx1ba

100class Producer(SequentialOperation): 3 ctx1bca

101 """Produces something at each iteration but no final output."""

102

103 def finalize(self): 3 ctx1bca

104 pass 2 ctx1ba

105

106

107class Consumer(SequentialOperation): 3 ctx1bca

108 """Produces a final output but no iteration output."""

109

110 iter_out = NotImplemented 3 ctx1bca

111

112

113class SingleInput(SequentialOperation): 3 ctx1bca

114 def __init__(self, input): # noqa: A002, `inputs` are the ops consumed 3 ctx1bca

115 self.inputs = (input,) 1 ctx1a

116

117 inputs = NotImplemented 3 ctx1bca

118

119

120class Stack(Consumer, SingleInput): 3 ctx1bca

121 """`input` = an operation producing arrays."""

122

123 def init(self, n, a0): 3 ctx1bca

124 out = jnp.zeros((n, *a0.shape), a0.dtype) 1 ctx1a

125 self.out = out.at[0, ...].set(a0) 1 ctx1a

126

127 def iter(self, i, ai): 3 ctx1bca

128 self.out = self.out.at[i, ...].set(ai) 1 ctx1a

129

130 def finalize(self): 3 ctx1bca

131 """Return the stacked arrays."""

132 return self.out 1 ctx1a

133

134

135class MatMulIterByFull(Consumer, SingleInput): 3 ctx1bca

136 def __init__(self, input, b): # noqa: A002 3 ctx1bca

137 """

138 Initialize the matrix multiplication.

139

140 `input` = an operation producing pieces of left operand (a)

141 b = right operand.

142 """

143 self.inputs = (input,) 2 ctx1ba

144 b = jnp.asarray(b) 2 ctx1ba

145 assert b.ndim in (1, 2) 2 ctx1ba

146 vec = b.ndim < 2 2 ctx1ba

147 if vec: 2 ctx1ba

148 b = b[:, None] 2 ctx1ba

149 self.vec = vec 2 ctx1ba

150 self.b = b 2 ctx1ba

151

152 @abc.abstractmethod 3 ctx1bca

153 def init(self, n, a0): # pragma: no cover 3 ctx1bca

154 self.ab = ...

155

156 @abc.abstractmethod 3 ctx1bca

157 def iter(self, i, ai): # pragma: no cover 3 ctx1bca

158 self.ab = ...

159

160 def finalize(self): 3 ctx1bca

161 ab = self.ab 1 ctx1a

162 if self.vec: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true1 ctx1a

163 ab = jnp.squeeze(ab, -1)

164 return ab 1 ctx1a

165

166

167class MatMulRowByFull(Producer, MatMulIterByFull): 3 ctx1bca

168 def init(self, n, a0): # noqa: ARG002, interface method 3 ctx1bca

169 assert a0.ndim == 1 2 ctx1ba

170 assert self.b.shape[0] == len(a0) 2 ctx1ba

171 self.abi = a0 @ self.b 2 ctx1ba

172

173 def iter_out(self, i): # noqa: ARG002, interface method 3 ctx1bca

174 abi = self.abi 2 ctx1ba

175 if self.vec: 2 ctx1ba

176 abi = jnp.squeeze(abi, -1) 2 ctx1ba

177 return abi 2 ctx1ba

178

179 def iter(self, i, ai): # noqa: ARG002, interface method 3 ctx1bca

180 """`i`-th row of `input @ b`."""

181 self.abi = ai @ self.b 2 ctx1ba

182

183

184class SolveTriLowerColByFull(MatMulIterByFull): 3 ctx1bca

185 # x[0] /= a[0, 0]

186 # for i in range(1, len(x)):

187 # x[i:] -= x[i - 1] * a[i:, i - 1]

188 # x[i] /= a[i, i]

189

190 def init(self, n, a0): 3 ctx1bca

191 b = self.b 1 ctx1a

192 del self.b 1 ctx1a

193 assert a0.shape == (n,) 1 ctx1a

194 assert b.shape[0] == n 1 ctx1a

195 self.prevai = a0.at[0].set(0) 1 ctx1a

196 self.ab = b.at[0, :].divide(a0[0]) 1 ctx1a

197

198 def iter(self, i, ai): 3 ctx1bca

199 ab = self.ab 1 ctx1a

200 ab = ab - ab[i - 1, :] * self.prevai[:, None] 1 ctx1a

201 self.ab = ab.at[i, :].divide(ai[i]) 1 ctx1a

202 self.prevai = ai.at[i].set(0) 1 ctx1a

203

204

205class Rows(Producer): 3 ctx1bca

206 def __init__(self, x): 3 ctx1bca

207 self.x = x 1 ctx1a

208

209 inputs = () 3 ctx1bca

210

211 def init(self, n): 3 ctx1bca

212 pass 1 ctx1a

213

214 def iter_out(self, i): 3 ctx1bca

215 return self.x[i, ...] 1 ctx1a

216

217 def iter(self, i): 3 ctx1bca

218 pass 1 ctx1a

219

220

221class MatMulColByRow(Consumer): 3 ctx1bca

222 def __init__(self, inputa, inputb): 3 ctx1bca

223 self.inputs = (inputa, inputb) 2 ctx1ba

224

225 inputs = None 3 ctx1bca

226

227 def init(self, n, a0, b0): # noqa: ARG002, interface method 3 ctx1bca

228 assert a0.ndim == 1 2 ctx1ba

229 assert b0.ndim <= 1 2 ctx1ba

230 self.vec = b0.ndim > 0 2 ctx1ba

231 if self.vec: 2 ctx1ba

232 self.ab = a0[:, None] * b0[None, :] 1 ctx1a

233 else:

234 self.ab = a0 * b0 2 ctx1ba

235

236 def iter(self, i, ai, bi): # noqa: ARG002, interface method 3 ctx1bca

237 if self.vec: 2 ctx1ba

238 self.ab = self.ab + ai[:, None] * bi[None, :] 1 ctx1a

239 else:

240 self.ab = self.ab + ai * bi 2 ctx1ba

241

242 def finalize(self): 3 ctx1bca

243 return self.ab 2 ctx1ba

244

245

246class SumLogDiag(Consumer, SingleInput): 3 ctx1bca

247 """`input` = operation producing the rows/columns of a square matrix."""

248

249 def init(self, n, m0): 3 ctx1bca

250 assert m0.shape == (n,) 1 ctx1a

251 self.sld = jnp.log(m0[0]) 1 ctx1a

252

253 def iter(self, i, mi): 3 ctx1bca

254 self.sld = self.sld + jnp.log(mi[i]) 1 ctx1a

255

256 def finalize(self): 3 ctx1bca

257 """`sum(log(diag(m)))`."""

258 return self.sld 1 ctx1a

Coverage for src/lsqfitgp/_linalg/_seqalg.py: 98%

112 statements