Coverage for src/lsqfitgp/copula/_distr.py: 94%

1# lsqfitgp/copula/_distr.py 

2# 

3# Copyright (c) 2023, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

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/>. 

19 

20"""Define Distr and distribution.""" 

21 

22import abc 3 ctx1cba

23import collections 3 ctx1cba

24import functools 3 ctx1cba

25import inspect 3 ctx1cba

26import math 3 ctx1cba

27import numbers 3 ctx1cba

28import types 3 ctx1cba

29 

30import jax 3 ctx1cba

31import numpy 3 ctx1cba

32from jax import numpy as jnp 3 ctx1cba

33 

34from lsqfitgp import _gvarext, _signature 3 ctx1cba

35from lsqfitgp.copula import _base 3 ctx1cba

36 

37######### The following 5 functions are adapted from numpy.lib.mixins ######### 

38 

39 

40def _disables_array_ufunc(obj): 3 ctx1cba

41 """Check if __array_ufunc__ is set to None.""" 

42 return getattr(obj, '__array_ufunc__', NotImplemented) is None 2 ctx1ba

43 

44 

45def _binary_method(ufunc, name): 3 ctx1cba

46 """Implement a forward binary method with a ufunc, e.g., __add__.""" 

47 

48 def func(self, other): 3 ctx1cba

49 if _disables_array_ufunc(other): 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true1 ctx1a

50 return NotImplemented 

51 return ufunc(self, other) 1 ctx1a

52 

53 func.__name__ = f'__{name}__' 3 ctx1cba

54 return func 3 ctx1cba

55 

56 

57def _reflected_binary_method(ufunc, name): 3 ctx1cba

58 """Implement a reflected binary method with a ufunc, e.g., __radd__.""" 

59 

60 def func(self, other): 3 ctx1cba

61 if _disables_array_ufunc(other): 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true2 ctx1ba

62 return NotImplemented 

63 return ufunc(other, self) 2 ctx1ba

64 

65 func.__name__ = f'__r{name}__' 3 ctx1cba

66 return func 3 ctx1cba

67 

68 

69def _numeric_methods(ufunc, name): 3 ctx1cba

70 """Implement forward and reflected binary methods with a ufunc.""" 

71 return (_binary_method(ufunc, name), _reflected_binary_method(ufunc, name)) 3 ctx1cba

72 

73 

74def _unary_method(ufunc, name): 3 ctx1cba

75 """Implement a unary special method with a ufunc.""" 

76 

77 def func(self): 3 ctx1cba

78 return ufunc(self) 1 ctx1a

79 

80 func.__name__ = f'__{name}__' 3 ctx1cba

81 return func 3 ctx1cba

82 

83 

84############################################################################### 

85 

86 

87class Distr(_base.DistrBase): 3 ctx1cba

88 r""" 

89 

90 Abstract base class to represent probability distributions. 

91 

92 A `Distr` object represents a probability distribution of a variable in 

93 :math:`\mathbb R^n`, and provides a transformation function from a 

94 (multivariate) Normal variable to the target random variable. 

95 

96 The main functionality is defined in `DistrBase`. The additional attributes 

97 and methods `params`, `signature`, and `invfcn` are not intended for common 

98 usage. 

99 

100 Parameters 

101 ---------- 

102 *params : tuple of scalar, array or Distr 

103 The parameters of the distribution. If the parameters have leading axes 

104 other than those required, the distribution is repeated i.i.d. 

105 over those axes. If a parameter is an instance of `Distr` itself, it 

106 is a random parameter and its distribution is accounted for. 

107 shape : int or tuple of int 

108 The shape of the array of i.i.d. variables to be represented, scalar by 

109 default. If the variable is multivariate, this shape adds as leading 

110 axes in the array. This shape broadcasts with the non-core shapes of the 

111 parameters. 

112 name : str, optional 

113 If specified, the distribution is defined for usage with 

114 `gvar.BufferDict` using `gvar.BufferDict.add_distribution`, and for 

115 convenience the constructor returns an array of gvars with the 

116 appropriate shape instead of the `Distr` object. See `add_distribution`. 

117 

118 Returns 

119 ------- 

120 distr : Distr 

121 If `name` is None (default), an object representing the distribution. 

122 gvars : array of gvars 

123 If `name` is specified, an array of primary gvars that can be set as 

124 value in a `gvar.BufferDict` under a key that uses the just defined 

125 name. 

126 

127 Attributes 

128 ---------- 

129 params : tuple 

130 The parameters as passed to the constructor. 

131 signature : Signature 

132 An object representing the signature of `invfcn`. This is a class 

133 attribute. 

134 

135 Methods 

136 ------- 

137 invfcn : classmethod 

138 Transformation function from a (multivariate) Normal variable to the 

139 target random variable. 

140 

141 Examples 

142 -------- 

143 Use directly with `gvar.BufferDict` by setting `name`: 

144 

145 >>> copula = gvar.BufferDict({ 

146 ... 'A(x)': lgp.copula.beta(1, 1, name='A'), 

147 ... 'B(y)': lgp.copula.beta(3, 5, name='B'), 

148 ... }) 

149 >>> copula['x'] 

150 0.50(40) 

151 >>> copula['y'] 

152 0.36(18) 

153 

154 Corresponding "unrolled" usage: 

155 

156 >>> A = lgp.copula.beta(1, 1) 

157 >>> B = lgp.copula.beta(3, 5) 

158 >>> A.add_distribution('A') 

159 >>> B.add_distribution('B') 

160 >>> copula = gvar.BufferDict({ 

161 ... 'A(x)': A.gvars(), 

162 ... 'B(y)': B.gvars(), 

163 ... }) 

164 

165 Notice that, although the name used for `add_distribution` must be globally 

166 unique, for convenience it is permitted to redefine the same distribution 

167 family with the same parameters, even from another `Distr` instance. 

168 

169 To generate automatically sensible names and avoid repeating them twice, use 

170 `makedict`: 

171 

172 >>> lgp.copula.makedict({ 

173 ... 'x': lgp.copula.beta(1, 1), 

174 ... 'y': lgp.copula.beta(3, 5), 

175 ... }) 

176 BufferDict({'__copula_beta{1, 1}(x)': 0.0(1.0), '__copula_beta{3, 5}(y)': 0.0(1.0)}) 

177 

178 Define a distribution with a random parameter: 

179 

180 >>> X = lgp.copula.halfnorm(np.sqrt(lgp.copula.invgamma(1, 1))) 

181 >>> X 

182 halfnorm(sqrt(invgamma(1, 1))) 

183 

184 Now `X` represents the model 

185 

186 .. math:: 

187 \sigma^2 &\sim \mathrm{InvGamma}(1, 1), \\ 

188 X \mid \sigma &\sim \mathrm{HalfNorm}(\sigma). 

189 

190 In general it is possible to transform a `Distr` with `numpy` ufuncs and 

191 continuous arithmetic operations. 

192 

193 Repeated usage of `Distr` instances for random parameters will share 

194 those parameters in the distributions. The following code: 

195 

196 >>> sigma2 = lgp.copula.invgamma(1, 1) 

197 >>> X = lgp.copula.halfnorm(np.sqrt(sigma2)) 

198 >>> Y = lgp.copula.halfcauchy(np.sqrt(sigma2)) 

199 

200 Corresponds to the model 

201 

202 .. math:: 

203 \sigma^2 &\sim \mathrm{InvGamma}(1, 1), \\ 

204 X \mid \sigma &\sim \mathrm{HalfNorm}(\sigma), \\ 

205 Y \mid \sigma &\sim \mathrm{HalfCauchy}(\sigma), 

206 

207 with the same parameter :math:`\sigma^2` shared between the two 

208 distributions. However, if the distributions are now put into a 

209 `gvar.BufferDict`, with 

210 

211 >>> sigma2.add_distribution('distr_sigma2') 

212 >>> X.add_distribution('distr_X') 

213 >>> Y.add_distribution('distr_Y') 

214 >>> bd = gvar.BufferDict({ 

215 ... 'distr_sigma2(sigma2)': sigma2.gvars(), 

216 ... 'distr_X(X)': X.gvars(), 

217 ... 'distr_Y(Y)': Y.gvars(), 

218 ... }) 

219 

220 then this relationship breaks down; the model represented by the dictionary 

221 `bd` is 

222 

223 .. math:: 

224 \sigma^2 &\sim \mathrm{InvGamma}(1, 1), \\ 

225 X \mid \sigma_X &\sim \mathrm{HalfNorm}(\sigma_X), \quad 

226 & \sigma_X^2 &\sim \mathrm{InvGamma}(1, 1), \\ 

227 Y \mid \sigma_Y &\sim \mathrm{HalfCauchy}(\sigma_Y), \quad 

228 & \sigma_Y^2 &\sim \mathrm{InvGamma}(1, 1), 

229 

230 with separate, independent parameters :math:`\sigma,\sigma_X,\sigma_Y`, 

231 because each dictionary entry is evaluated separately. Indeed, trying to do 

232 this with `makedict` will raise an error: 

233 

234 >>> bd = lgp.copula.makedict({'sigma2': sigma2, 'X': X, 'Y': Y}) 

235 ValueError: cross-key occurrences of object(s): 

236 invgamma with id 6201535248: <sigma2>, <X.0.0>, <Y.0.0> 

237 

238 To use all the distributions at once while preserving the relationships, 

239 put them into a container of choice and wrap it as a `Copula` object: 

240 

241 >>> sigmaXY = lgp.copula.Copula({'sigma2': sigma2, 'X': X, 'Y': Y}) 

242 

243 The `Copula` provides a `partial_invfcn` function to map Normal variables 

244 to a structure, with the same layout as the input one, of desired variates. 

245 The whole `Copula` can be used in `gvar.BufferDict`: 

246 

247 >>> bd = lgp.copula.makedict({'sigmaXY': sigmaXY}) 

248 >>> bd 

249 BufferDict({"__copula_{'sigma2': invgamma{1, 1}, 'X': halfnorm{sqrt{_Path{path=[{DictKey{key='sigma2'},}]}}}, 'Y': halfcauchy{sqrt{_Path{path=[{DictKey{key='sigma2'},}]}}}}(sigmaXY)": array([0.0(1.0), 0.0(1.0), 0.0(1.0)], dtype=object)}) 

250 >>> bd['sigmaXY'] 

251 {'sigma2': 1.4(1.7), 'X': 0.81(89), 'Y': 1.2(1.7)} 

252 >>> gvar.corr(bd['sigmaXY']['X'], bd['sigmaXY']['Y']) 

253 0.21950577757757836 

254 

255 Although the actual dictionary value is a flat array, getting the unwrapped 

256 key reproduces the original structure. 

257 

258 To apply arbitrary transformations, use manually `invfcn`: 

259 

260 >>> @functools.partial(lgp.gvar_gufunc, signature='(n)->(n)') 

261 >>> @functools.partial(jnp.vectorize, signature='(n)->(n)') 

262 >>> def model_invfcn(normal_params): 

263 ... sigma2 = lgp.copula.invgamma.invfcn(normal_params[0], 1, 1) 

264 ... sigma = jnp.sqrt(sigma2) 

265 ... X = lgp.copula.halfnorm.invfcn(normal_params[1], sigma) 

266 ... Y = lgp.copula.halfcauchy.invfcn(normal_params[2], sigma) 

267 ... return jnp.stack([sigma, X, Y]) 

268 

269 The `jax.numpy.vectorize` decorator makes `model_invfcn` support 

270 broadcasting on additional input axes, while `gvar_gufunc` makes it accept 

271 gvars as input. 

272 

273 See Also 

274 -------- 

275 DistrBase, Copula, gvar.BufferDict.uniform 

276 

277 Notes 

278 ----- 

279 Concrete subclasses must define `invfcn`, and define the class attribute 

280 `signature` to the numpy signature string of `invfcn`, unless `invfcn` is an 

281 ufunc and its number of parameters can be inferred. `invfcn` must be 

282 vectorized. 

283 

284 """ 

285 

286 @classmethod 3 ctx1cba

287 @abc.abstractmethod 3 ctx1cba

288 def invfcn(cls, x, *params): 3 ctx1cba

289 r""" 

290 

291 Transform a Normal variable to the desired distribution. 

292 

293 Maps a (multivariate) Normal variable to a variable with the desired 

294 marginal distribution. In symbols: :math:`y = F^{-1}(\Phi(x))`. This 

295 function is a generalized ufunc, jax traceable, vmappable one time, and 

296 differentiable one time. The signature is accessible through the 

297 class attribute `signature`. 

298 

299 Parameters 

300 ---------- 

301 x : array_like 

302 The input Normal variable. 

303 *params : array_like 

304 The parameters of the distribution. 

305 

306 Returns 

307 ------- 

308 y : array_like 

309 The output variable with the desired marginal distribution. 

310 

311 """ # noqa: DOC202 

312 pass 

313 

314 def _get_x_core_shape(self, *preprocessed_params): 3 ctx1cba

315 sig = self.signature.eval(None, *preprocessed_params) 3 ctx1cba

316 return sig.core_in_shapes[0] 3 ctx1cba

317 

318 def _eval_shapes(self, shape): 3 ctx1cba

319 

320 # check number of parameters 

321 if self.signature.nin != 1 + len(self.params): 3 ctx1cba

322 msg = ( 1 ctx1a

323 f'{self.__class__.__name__} distribution has ' 

324 f'{self.signature.nin - 1} parameters, but {len(self.params)} ' 

325 'parameters were passed to the constructor' 

326 ) 

327 raise TypeError(msg) 1 ctx1a

328 

329 # convert shape to tuple 

330 if isinstance(shape, numbers.Integral): 3 ctx1cba

331 shape = (shape,) 1 ctx1a

332 else: 

333 shape = tuple(shape) 3 ctx1cba

334 

335 # make sure parameters have a shape 

336 array_params = [ 3 ctx1cba

337 p if hasattr(p, 'shape') else jnp.asarray(p) for p in self.params 

338 ] 

339 

340 # parse signature of cls.invfcn 

341 x_core_shape = self._get_x_core_shape(*array_params) 3 ctx1cba

342 x = jax.ShapeDtypeStruct(shape + x_core_shape, 'd') 3 ctx1cba

343 sig = self.signature.eval(x, *array_params) 3 ctx1cba

344 self._in_shape_1 = sig.in_shapes[0] 3 ctx1cba

345 (self.distrshape,) = sig.core_out_shapes 3 ctx1cba

346 (self.shape,) = sig.out_shapes 3 ctx1cba

347 

348 self._compute_in_shape() 3 ctx1cba

349 

350 def _compute_in_shape(self): 3 ctx1cba

351 in_size = math.prod(self._in_shape_1) 3 ctx1cba

352 cache = set() 3 ctx1cba

353 for p in self.params: 3 ctx1cba

354 if isinstance(p, __class__): 3 ctx1cba

355 in_size += p._compute_in_size(cache) 2 ctx1ba

356 if in_size == 1: 3 ctx1cba

357 self.in_shape = () 3 ctx1cba

358 else: 

359 self.in_shape = (in_size,) 1 ctx1a

360 self._ancestor_count = len(cache) 3 ctx1cba

361 

362 def _compute_in_size(self, cache): 3 ctx1cba

363 if (out := super()._compute_in_size(cache)) is not None: 2 ctx1ba

364 return out 1 ctx1a

365 in_size = math.prod(self._in_shape_1) 2 ctx1ba

366 for p in self.params: 2 ctx1ba

367 if isinstance(p, __class__): 2 ctx1ba

368 in_size += p._compute_in_size(cache) 1 ctx1a

369 return in_size 2 ctx1ba

370 

371 def _partial_invfcn_internal(self, x, i, cache): 3 ctx1cba

372 if (out := super()._partial_invfcn_internal(x, i, cache)) is not None: 3 ctx1cba

373 return out 1 ctx1a

374 

375 concrete_params = [] 3 ctx1cba

376 for p in self.params: 3 ctx1cba

377 if isinstance(p, __class__): 3 ctx1cba

378 p, i = p._partial_invfcn_internal(x, i, cache) 2 ctx1ba

379 else: 

380 p = jnp.asarray(p) 3 ctx1cba

381 

382 concrete_params.append(p) 3 ctx1cba

383 

384 in_size = math.prod(self._in_shape_1) 3 ctx1cba

385 assert i + in_size <= x.size 3 ctx1cba

386 last = x[i : i + in_size].reshape(self._in_shape_1) 3 ctx1cba

387 

388 y = self.invfcn(last, *concrete_params) 3 ctx1cba

389 if y.shape != self.shape or y.dtype != self.dtype: 389 ↛ 390line 389 didn't jump to line 390 because the condition on line 389 was never true3 ctx1cba

390 msg = ( 

391 f'{self.__class__.__name__}.invfcn returned ' 

392 f'array with shape {y.shape} and dtype {y.dtype}, while ' 

393 f'{self.shape} and {self.dtype} were expected' 

394 ) 

395 raise ValueError(msg) 

396 

397 cache[self] = y 3 ctx1cba

398 return y, i + in_size 3 ctx1cba

399 

400 @functools.cached_property 3 ctx1cba

401 def _partial_invfcn(self): 3 ctx1cba

402 

403 # determine signature 

404 shapestr = lambda shape: ','.join(map(str, shape)) 3 ctx1cba

405 signature = f'({shapestr(self.in_shape)})->({shapestr(self.shape)})' 3 ctx1cba

406 

407 # wrap to support gvars 

408 @functools.partial(_gvarext.gvar_gufunc, signature=signature) 3 ctx1cba

409 # @jit 

410 @functools.partial(jnp.vectorize, signature=signature) 3 ctx1cba

411 def _partial_invfcn(x): 3 ctx1cba

412 assert x.shape == self.in_shape 3 ctx1cba

413 if not self.in_shape: 3 ctx1cba

414 x = x[None] 3 ctx1cba

415 cache = {} 3 ctx1cba

416 y, i = self._partial_invfcn_internal(x, 0, cache) 3 ctx1cba

417 assert i == x.size 3 ctx1cba

418 assert len(cache) == 1 + self._ancestor_count 3 ctx1cba

419 return y 3 ctx1cba

420 

421 return _partial_invfcn 3 ctx1cba

422 

423 def __init_subclass__(cls, **kw): 3 ctx1cba

424 super().__init_subclass__(**kw) 3 ctx1cba

425 

426 # check and/or set signature attribute (the gufunc signature of invfcn) 

427 if not hasattr(cls, 'signature'): 3 ctx1cba

428 sig = inspect.signature(cls.invfcn) 3 ctx1cba

429 if not all( 429 ↛ 437line 429 didn't jump to line 437 because the condition on line 429 was never true3 ctx1cba

430 p.kind 

431 in ( 

432 inspect.Parameter.POSITIONAL_ONLY, 

433 inspect.Parameter.POSITIONAL_OR_KEYWORD, 

434 ) 

435 for p in sig.parameters.values() 

436 ): 

437 msg = ( 

438 'can not automatically infer signature of ' 

439 f'{cls.__qualname__}.invfcn' 

440 ) 

441 raise ValueError(msg) 

442 cls.signature = ','.join(['()'] * len(sig.parameters)) + '->()' 3 ctx1cba

443 if not isinstance(cls.signature, _signature.Signature): 443 ↛ 445line 443 didn't jump to line 445 because the condition on line 443 was always true3 ctx1cba

444 cls.signature = _signature.Signature(cls.signature) 3 ctx1cba

445 cls.signature.check_nargs(cls.invfcn) 3 ctx1cba

446 

447 # set dtype to float if not specified 

448 if getattr(cls, 'dtype', NotImplemented) is NotImplemented: 3 ctx1cba

449 cls.dtype = jax.dtypes.canonicalize_dtype(jnp.float64) 3 ctx1cba

450 

451 # set __signature__ to take positional parameters from invfcn 

452 sig = inspect.signature(cls.invfcn) 3 ctx1cba

453 pos_params = list(sig.parameters.values())[1:] 3 ctx1cba

454 sig = inspect.signature(cls.__new__) 3 ctx1cba

455 key_params = [ 3 ctx1cba

456 p 

457 for i, p in enumerate(sig.parameters.values()) 

458 if p.kind 

459 in (inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) 

460 and i > 0 

461 ] 

462 cls.__signature__ = inspect.Signature(pos_params + key_params) 3 ctx1cba

463 

464 def __new__(cls, *params, shape=(), name=None): 3 ctx1cba

465 

466 self = super().__new__(cls) 3 ctx1cba

467 self.params = params 3 ctx1cba

468 self._eval_shapes(shape) 3 ctx1cba

469 

470 if name is None: 3 ctx1cba

471 return self 3 ctx1cba

472 else: 

473 self.add_distribution(name) 1 ctx1a

474 return self.gvars() 1 ctx1a

475 

476 class _Descr(collections.namedtuple('Distr', 'family shape params')): 3 ctx1cba

477 """Static representation of a Distr object.""" 

478 

479 def __repr__(self): 3 ctx1cba

480 args = list(map(repr, self.params)) 3 ctx1cba

481 if len(self.shape) == 1: 3 ctx1cba

482 args += [f'shape={self.shape[0]}'] 1 ctx1a

483 elif self.shape: 3 ctx1cba

484 args += [f'shape={self.shape}'] 1 ctx1a

485 arglist = ', '.join(args) 3 ctx1cba

486 return f'{self.family.__name__}({arglist})' 3 ctx1cba

487 

488 def _compute_staticdescr(self, path, cache): 3 ctx1cba

489 if (obj := super()._compute_staticdescr(path, cache)) is not None: 3 ctx1cba

490 return obj 1 ctx1a

491 

492 params = [] 3 ctx1cba

493 for i, p in enumerate(self.params): 3 ctx1cba

494 if isinstance(p, __class__): 3 ctx1cba

495 p = p._compute_staticdescr([*path, i], cache) 2 ctx1ba

496 else: 

497 p = numpy.asarray(p).tolist() 3 ctx1cba

498 params.append(p) 3 ctx1cba

499 

500 return self._Descr(self.__class__, self.shape, tuple(params)) 3 ctx1cba

501 

502 def _shapestr(self, shape): 3 ctx1cba

503 if shape: 2 ctx1ba

504 return ( 1 ctx1a

505 str(shape) 

506 .replace(',)', ')') 

507 .replace('(', '[') 

508 .replace(')', ']') 

509 .replace(' ', '') 

510 ) 

511 else: 

512 return '' 1 ctx1b

513 

514 def __repr__(self, path='', cache=None): 3 ctx1cba

515 

516 if isinstance(cache := super().__repr__(path, cache), str): 3 ctx1cba

517 return cache 1 ctx1a

518 

519 args = [] 3 ctx1cba

520 for i, p in enumerate(self.params): 3 ctx1cba

521 if isinstance(p, __class__): 3 ctx1cba

522 p = p.__repr__('.'.join((path, str(i))).lstrip('.'), cache) 2 ctx1ba

523 elif hasattr(p, 'shape'): 3 ctx1cba

524 p = f'Array{self._shapestr(p.shape)}' 2 ctx1ba

525 else: 

526 p = repr(p) 3 ctx1cba

527 args.append(p) 3 ctx1cba

528 

529 if len(self.shape) == 1: 3 ctx1cba

530 args += [f'shape={self.shape[0]}'] 1 ctx1a

531 elif self.shape: 3 ctx1cba

532 args += [f'shape={self.shape}'] 1 ctx1a

533 

534 return f'{self.__class__.__name__}({", ".join(args)})' 3 ctx1cba

535 

536 def __array_ufunc__(self, ufunc, method, *inputs, **kw): 3 ctx1cba

537 if method != '__call__' or kw or ufunc.signature: 537 ↛ 538line 537 didn't jump to line 538 because the condition on line 537 was never true2 ctx1ba

538 return NotImplemented 

539 ufunc_class = UFunc.make_subclass(ufunc) 2 ctx1ba

540 return ufunc_class(*inputs) 2 ctx1ba

541 

542 # continuous binary operations 

543 __add__, __radd__ = _numeric_methods(numpy.add, 'add') 3 ctx1cba

544 __sub__, __rsub__ = _numeric_methods(numpy.subtract, 'sub') 3 ctx1cba

545 __mul__, __rmul__ = _numeric_methods(numpy.multiply, 'mul') 3 ctx1cba

546 # __matmul__, __rmatmul__ = _numeric_methods(numpy.matmul, 'matmul') 

547 __truediv__, __rtruediv__ = _numeric_methods(numpy.divide, 'truediv') 3 ctx1cba

548 __mod__, __rmod__ = _numeric_methods(numpy.remainder, 'mod') 3 ctx1cba

549 __divmod__, __rdivmod__ = _numeric_methods(numpy.divmod, 'divmod') 3 ctx1cba

550 __pow__, __rpow__ = _numeric_methods(numpy.power, 'pow') 3 ctx1cba

551 

552 # continuous unary operations 

553 __neg__ = _unary_method(numpy.negative, 'neg') 3 ctx1cba

554 __pos__ = _unary_method(numpy.positive, 'pos') 3 ctx1cba

555 __abs__ = _unary_method(numpy.absolute, 'abs') 3 ctx1cba

556 

557 

558class UFunc: 3 ctx1cba

559 """Base class of objects representing ufuncs applied to Distr instances.""" 

560 

561 def __new__(cls, *args): 3 ctx1cba

562 return super().__new__(cls, *args) 2 ctx1ba

563 # this __new__ serves to forbid keyword arguments 

564 

565 @classmethod 3 ctx1cba

566 def invfcn(cls, x, *args): # noqa: ARG003, the signature is inspected 3 ctx1cba

567 return cls._ufunc(*args) 2 ctx1ba

568 

569 def _get_x_core_shape(self, *_): 3 ctx1cba

570 return (0,) 2 ctx1ba

571 

572 @classmethod 3 ctx1cba

573 @functools.cache # functools.cache not available in 3.8 3 ctx1cba

574 def make_subclass(cls, ufunc): 3 ctx1cba

575 def exec_body(ns): 2 ctx1ba

576 ns['_ufunc'] = getattr(jnp, ufunc.__name__) 2 ctx1ba

577 ns['signature'] = ','.join(['(0)'] + ufunc.nin * ['()']) + '->()' 2 ctx1ba

578 

579 return types.new_class(ufunc.__name__, (__class__, Distr), exec_body=exec_body) 2 ctx1ba

580 

581 

582def distribution(invfcn, signature=None, dtype=None): 3 ctx1cba

583 r""" 

584 

585 Decorate a transformation function to define a distribution. 

586 

587 Parameters 

588 ---------- 

589 invfcn : function 

590 The transformation function from a (multivariate) standard Normal 

591 variable to the target random variable. The signature must be 

592 ``invfcn(x, *params)``. It must be jax-traceable. It does not need to 

593 be vectorized. 

594 signature : str, optional 

595 The signature of `invfcn`, as a numpy signature string. If not 

596 specified, `invfcn` is assumed to take and output scalars. 

597 dtype : dtype, optional 

598 The dtype of the output of `invfcn`. If not specified, it is assumed to 

599 be floating point. 

600 

601 Returns 

602 ------- 

603 cls : Distr 

604 The new distribution class. 

605 

606 Examples 

607 -------- 

608 >>> @lgp.copula.distribution 

609 ... def uniform(x, a, b): 

610 ... return a + (b - a) * jax.scipy.stats.norm.cdf(x) 

611 

612 >>> @functools.partial(lgp.copula.distribution, signature='(n,m)->(n)') 

613 ... def wishart(x): 

614 ... " this parametrization is terrible, do not use " 

615 ... return x @ x.T 

616 

617 """ 

618 

619 def exec_body(ns): 1 ctx1a

620 if signature is not None: 620 ↛ 622line 620 didn't jump to line 622 because the condition on line 620 was always true1 ctx1a

621 ns['signature'] = signature 1 ctx1a

622 if dtype is not None: 622 ↛ 624line 622 didn't jump to line 624 because the condition on line 622 was always true1 ctx1a

623 ns['dtype'] = dtype 1 ctx1a

624 ns['invfcn'] = staticmethod(jnp.vectorize(invfcn, signature=signature)) 1 ctx1a

625 

626 return types.new_class(invfcn.__name__, (Distr,), exec_body=exec_body) 1 ctx1a