Coverage for src/lsqfitgp/copula/_base.py: 95%
65 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-04 11:43 +0000
1# lsqfitgp/copula/_base.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"""Define DistrBase."""
22import abc 1cba
23import collections 1cba
24import functools 1cba
26import gvar 1cba
27import numpy 1cba
30class DistrBase(metaclass=abc.ABCMeta): 1cba
31 r"""
33 Abstract base class to represent (trees of) probability distributions.
35 Attributes
36 ----------
37 in_shape : tuple of int
38 The core shape of the input array to `partial_invfcn`.
39 shape : (tree of) tuple of int
40 The core shape of the output array of `partial_invfcn`.
41 dtype : (tree of) dtype
42 The dtype of the output array of `partial_invfcn`.
43 distrshape : (tree of) tuple of int
44 The sub-core shape of the output array of `partial_invfcn` that
45 represents the atomic shape of the distribution.
47 Methods
48 -------
49 partial_invfcn :
50 Transformation function from a (multivariate) Normal variable to the
51 target random variable.
52 add_distribution :
53 Register the distribution for usage with `gvar.BufferDict`.
54 gvars :
55 Return an array of gvars with the appropriate shape for usage with
56 `gvar.BufferDict`.
58 See Also
59 --------
60 Distr, Copula
62 """
64 def __init_subclass__(cls, **kw): 1cba
65 super().__init_subclass__(**kw) 1cba
66 cls._named = {} 1cba
68 in_shape = NotImplemented 1cba
69 shape = NotImplemented 1cba
70 dtype = NotImplemented 1cba
71 distrshape = NotImplemented 1cba
73 def partial_invfcn(self, x): 1cba
74 """
75 Map independent Normal variables to the desired distribution.
77 This function is a generalized ufunc. It is jax traceable and
78 differentiable one time. It supports arrays of gvars as input. The
79 attributes `in_shape` and `shape` give the core shapes.
81 Parameters
82 ----------
83 x : ``(..., *in_shape)`` array
84 An array of values representing draws of i.i.d. Normal variates.
86 Returns
87 -------
88 y : (tree of) ``(..., *shape)`` array
89 An array of values representing draws of the desired distribution.
90 """
91 return self._partial_invfcn(x) 1cba
93 @abc.abstractmethod 1cba
94 def _partial_invfcn(self, x): 1cba
95 pass
97 def _is_same_family(self, invfcn): 1cba
98 return getattr(invfcn, '__self__', None).__class__ is self.__class__ 1ba
100 def add_distribution(self, name): 1cba
101 """
102 Register the distribution for usage with `gvar.BufferDict`.
104 Parameters
105 ----------
106 name : str
107 The name to use for the distribution. It must be globally unique,
108 and it should not contain parentheses. To redefine a distribution
109 with the same name, use `gvar.BufferDict.del_distribution` first.
110 However, it is allowed to reuse the name if the distribution family,
111 shape and parameters are identical to those used for the existing
112 definition.
114 See Also
115 --------
116 gvar.BufferDict.add_distribution, gvar.BufferDict.del_distribution
117 """
118 if gvar.BufferDict.has_distribution(name): 1cba
119 invfcn = gvar.BufferDict.invfcn[name] 1ba
120 if not self._is_same_family(invfcn): 1ba
121 msg = f'distribution {name} already defined' 1a
122 raise ValueError(msg) 1a
123 existing = self._named[name] 1ba
124 if existing != self._staticdescr: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true1ba
125 msg = (
126 'Attempt to overwrite existing'
127 f' {self.__class__.__name__} distribution with name {name}'
128 )
129 raise ValueError(msg)
130 # cls._named is not updated by
131 # gvar.BufferDict.del_distribution, but it is not a problem
133 else:
134 gvar.BufferDict.add_distribution(name, self.partial_invfcn) 1cba
135 self._named[name] = self._staticdescr 1cba
137 def gvars(self): 1cba
138 """
139 Return an array of gvars intended as value in a `gvar.BufferDict`.
141 Returns
142 -------
143 gvars : array of gvars
144 An array of i.i.d. standard Normal primary gvars with shape
145 `in_shape`.
146 """
147 return gvar.gvar(numpy.zeros(self.in_shape), numpy.ones(self.in_shape)) 1cba
149 @abc.abstractmethod 1cba
150 def __repr__(self, path='', cache=None): 1cba
151 """
152 Produce a representation where no object appears more than once.
154 Later appearances are replaced by a user-friendly identifier.
155 """
156 if cache is None: 1cba
157 cache = {} 1a
158 if self in cache: 1cba
159 return cache[self] 1a
160 cache[self] = f'<{path}>' 1cba
161 return cache 1cba
163 class _Path(collections.namedtuple('Path', ['path'])): 1cba
164 pass 1cba
166 @abc.abstractmethod 1cba
167 def _compute_staticdescr(self, path, cache): 1cba
168 """Compute static description of self, can be compared."""
169 if self in cache: 1cba
170 return cache[self] 1a
171 cache[self] = self._Path(path) 1cba
172 return None 1cba
174 @functools.cached_property 1cba
175 def _staticdescr(self): 1cba
176 return self._compute_staticdescr([], {}) 1cba
178 @abc.abstractmethod 1cba
179 def _compute_in_size(self, cache): 1cba
180 """Compute input size to partial_invfcn, without double counting."""
181 if self in cache: 1ba
182 return 0 1a
183 cache.add(self) 1ba
184 return None 1ba
186 @abc.abstractmethod 1cba
187 def _partial_invfcn_internal(self, x, i, cache): 1cba
188 assert x.ndim == 1 1cba
189 if self in cache: 1cba
190 return cache[self], i 1a
191 return None 1cba