Coverage for src/lsqfitgp/copula/_copula.py: 89%
134 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-04 11:43 +0000
1# lsqfitgp/copula/_copula.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/>.
20"""Defines Copula."""
22import functools 1bca
23import pprint 1bca
25import gvar 1bca
26import jax 1bca
27from jax import numpy as jnp 1bca
28from jax import tree, tree_util, vmap 1bca
30from lsqfitgp import _array, _gvarext 1bca
31from lsqfitgp.copula import _base 1bca
34class Copula(_base.DistrBase): 1bca
35 """
36 Represents a tree of probability distributions.
38 By "tree" it is intended an arbitrarily nested structure of containers
39 (e.g., `dict` or `list`) where each leaf node is either a `Distr` object or
40 another `Copula`.
42 The function of a `Copula` is to keep into account the relationships
43 between all the `Distr` objects when defining the function `partial_invfcn`
44 that maps Normal variates to the desired random variable. The same `Distr`
45 object appearing in multiple places will not be sampled more than once.
47 This class inherits the functionality defined in `DistrBase` without
48 additions. The attributes `shape`, `dtype`, `distrshape`, that represents
49 properties of the output of `partial_invfcn`, are trees matching the
50 structure of the tree of variables. The same holds for the output of
51 `partial_invfcn`. `in_shape` instead is an ordinary array shape, indicating
52 that the input Normal variables are a raveled single array.
54 Parameters
55 ----------
56 variables : tree of Distr or Copula
57 The tree of distributions to wrap. The containers are copied, so
58 successive modifications will not be reflected on the `Copula`.
60 See Also
61 --------
62 DistrBase, Distr, jax.tree_util
64 Examples
65 --------
66 Define a model into a dictionary one piece at a time, then wrap it as a
67 `Copula`:
69 >>> m = {}
70 >>> m['a'] = lgp.copula.halfnorm(1.5)
71 >>> m['b'] = lgp.copula.halfcauchy(2)
72 >>> m['c'] = [
73 ... lgp.copula.uniform(m['a'], m['b']),
74 ... lgp.copula.uniform(m['b'], m['a']),
75 ... ]
76 >>> cop = lgp.copula.Copula(m)
77 >>> cop
78 Copula({'a': halfnorm(1.5),
79 'b': halfcauchy(2),
80 'c': [uniform(<a>, <b>), uniform(<b>, <a>)]})
82 Notice how, when showing the object on the REPL, multiple appearances of
83 the same variables are replaced by identifiers derived from the dictionary
84 keys.
86 The model may then be extended to create a variant. This does not affect
87 `cop`:
89 >>> m['d'] = lgp.copula.invgamma(m['c'][0], m['c'][1])
90 >>> cop2 = lgp.copula.Copula(m)
91 >>> cop2
92 Copula({'a': halfnorm(1.5),
93 'b': halfcauchy(2),
94 'c': [uniform(<a>, <b>), uniform(<b>, <a>)],
95 'd': invgamma(<c.0>, <c.1>)})
96 >>> cop
97 Copula({'a': halfnorm(1.5),
98 'b': halfcauchy(2),
99 'c': [uniform(<a>, <b>), uniform(<b>, <a>)]})
100 """
102 @staticmethod 1bca
103 def _tree_path_str(path): 1bca
104 """Format a jax pytree key path as a compact, readable string."""
106 def parsekey(key): 1a
107 if hasattr(key, 'key'): 107 ↛ 109line 107 didn't jump to line 109 because the condition on line 107 was always true1a
108 return key.key 1a
109 elif hasattr(key, 'idx'):
110 return key.idx
111 else:
112 return key
114 def keystr(key): 1a
115 key = parsekey(key) 1a
116 return str(key).replace('.', r'\.') 1a
118 return '.'.join(map(keystr, path)) 1a
120 @classmethod 1bca
121 def _jaxext_dict_sorting(cls, pytree): 1bca
122 """
123 Replace dicts in pytree with a custom dict subclass.
125 The subclass is such that the insertion order is maintained, see
126 https://github.com/google/jax/issues/4085.
127 """
129 def is_dict(obj): 1a
130 return obj.__class__ is dict 1a
132 def patch_dict(obj): 1a
133 if is_dict(obj): 1a
134 return tree.map(patch_dict, cls._Dict(obj)) 1a
135 else:
136 return obj 1a
138 return tree.map(patch_dict, pytree, is_leaf=is_dict) 1a
140 @tree_util.register_pytree_with_keys_class 1bca
141 class _Dict(dict): 1bca
142 def tree_flatten_with_keys(self): 1bca
143 treedef = dict.fromkeys(self) 1a
144 keys_values = [(tree_util.DictKey(k), v) for k, v in self.items()] 1a
145 return keys_values, treedef 1a
147 @classmethod 1bca
148 def tree_unflatten(cls, treedef, values): 1bca
149 return cls(zip(treedef, values, strict=True)) 1a
151 def __init__(self, variables): 1bca
152 variables = self._jaxext_dict_sorting(variables) 1a
154 def check_type(path, obj): 1a
155 if not isinstance(obj, _base.DistrBase): 155 ↛ 156line 155 didn't jump to line 156 because the condition on line 155 was never true1a
156 msg = (
157 f'only Distr or Copula objects can be '
158 f'contained in a Copula, found {obj!r} at '
159 f'<{self._tree_path_str(path)}>'
160 )
161 raise TypeError(msg)
162 return obj 1a
164 self._variables = tree.map_with_path(check_type, variables) 1a
165 cache = set() 1a
166 self.in_shape = (self._compute_in_size(cache),) 1a
167 self._ancestor_count = len(cache) - 1 1a
168 self.shape = self._map_getattr('shape') 1a
169 self.distrshape = self._map_getattr('distrshape') 1a
170 self.dtype = self._map_getattr('dtype') 1a
172 def _compute_in_size(self, cache): 1bca
173 if (out := super()._compute_in_size(cache)) is not None: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true1a
174 return out
176 def accumulate(in_size, obj): 1a
177 return in_size + obj._compute_in_size(cache) 1a
179 return tree.reduce(accumulate, self._variables, 0) 1a
181 def _map_getattr(self, attr): 1bca
182 def get_attr(obj): 1a
183 if isinstance(obj, __class__): # ty: ignore[unresolved-reference] 1a
184 return obj._map_getattr(attr) 1a
185 else:
186 return getattr(obj, attr) 1a
188 return tree.map(get_attr, self._variables) 1a
190 def _partial_invfcn_internal(self, x, i, cache): 1bca
191 if (out := super()._partial_invfcn_internal(x, i, cache)) is not None: 191 ↛ 192line 191 didn't jump to line 192 because the condition on line 191 was never true1a
192 return out
194 distributions, treedef = tree.flatten(self._variables) 1a
195 outputs = [] 1a
196 for distr in distributions: 1a
197 out, i = distr._partial_invfcn_internal(x, i, cache) 1a
198 outputs.append(out) 1a
199 out = tree.unflatten(treedef, outputs) 1a
201 cache[self] = out 1a
202 return out, i 1a
204 @functools.cached_property 1bca
205 def _partial_invfcn(self): 1bca
207 # non vectorized version, check core shapes and call recursive impl
208 # @jit
209 def partial_invfcn_0(x): 1a
210 assert x.shape == self.in_shape 1a
211 cache = {} 1a
212 y, i = self._partial_invfcn_internal(x, 0, cache) 1a
213 assert i == x.size 1a
214 assert len(cache) == 1 + self._ancestor_count 1a
215 return y 1a
217 partial_invfcn_0_deriv = jax.jacfwd(partial_invfcn_0) 1a
219 # add 1-axis vectorization
220 partial_invfcn_1 = vmap(partial_invfcn_0) 1a
221 partial_invfcn_1_deriv = vmap(partial_invfcn_0_deriv) 1a
223 # add gvar support
224 def partial_invfcn_2(x): 1a
226 if x.dtype == object: 1a
227 # unpack the gvars
228 in_mean = gvar.mean(x) 1a
229 in_jac, indices = _gvarext.jacobian(x) 1a
231 # apply function
232 out_mean = partial_invfcn_1(in_mean) 1a
233 jac = partial_invfcn_1_deriv(in_mean) 1a
235 # concatenate derivatives and repack as gvars
236 def contract_and_pack(out_mean, jac): 1a
237 # indices:
238 # b = broadcast
239 # i = input
240 # ... = output
241 # g = gvar indices
242 out_jac = jnp.einsum('b...i,big->b...g', jac, in_jac) 1a
243 return _gvarext.from_jacobian(out_mean, out_jac, indices) 1a
245 return tree.map(contract_and_pack, out_mean, jac) 1a
247 else:
248 return partial_invfcn_1(x) 1a
250 # add full vectorization
251 def partial_invfcn_3(x): 1a
252 x = _array.asarray(x) 1a
253 assert x.shape[-1:] == self.in_shape 1a
254 head = x.shape[:-1] 1a
255 x = x.reshape((-1, *self.in_shape)) 1a
256 y = partial_invfcn_2(x) 1a
258 def reshape_y(y, shape): 1a
259 assert y.shape[1:] == shape 1a
260 y = y.reshape(head + shape) 1a
261 if y.dtype == object and not y.ndim: 1a
262 y = y.item() 1a
263 return y 1a
265 return tree.map(reshape_y, y, self.shape) 1a
267 return partial_invfcn_3 1a
269 def __repr__(self, path='', cache=None): 1bca
271 if isinstance(cache := super().__repr__(path, cache), str): 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true1a
272 return cache
274 def subrepr(k, obj): 1a
275 if isinstance(obj, _base.DistrBase): 275 ↛ 279line 275 didn't jump to line 279 because the condition on line 275 was always true1a
276 k = self._tree_path_str(k) 1a
277 return obj.__repr__(f'{path}.{k}'.lstrip('.'), cache) 1a
278 else:
279 return repr(obj)
281 class NoQuotesRepr: 1a
282 def __init__(self, s): 1a
283 self.s = s 1a
285 def __repr__(self): 1a
286 return self.s 1a
288 out = tree.map_with_path(subrepr, self._variables) 1a
289 out = tree.map(NoQuotesRepr, out) 1a
290 out = pprint.pformat(out, sort_dicts=False) 1a
291 return f'{self.__class__.__name__}({out})' 1a
293 def _compute_staticdescr(self, path, cache): 1bca
294 def compute(key, x): 1a
295 return x._compute_staticdescr([*path, key], cache) 1a
297 return tree.map_with_path(compute, self._variables) 1a