Coverage for src/cosmic_toolbox/TransformedGaussianMixture.py: 97%
74 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-10 10:31 +0000
1"""
2Transformed Gaussian Mixture Model.
4Provides a Gaussian Mixture Model that operates in a transformed parameter space
5to handle bounded parameters more effectively.
6"""
8import numpy as np
9from scipy.stats import norm
10from sklearn.mixture import GaussianMixture
13def scale_fwd(x, param_bounds, param_bounds_trans):
14 """
15 Scale values from original bounds to transformed bounds.
17 :param x: Values to scale.
18 :type x: numpy.ndarray
19 :param param_bounds: Original parameter bounds [min, max].
20 :type param_bounds: array-like
21 :param param_bounds_trans: Transformed parameter bounds [min, max].
22 :type param_bounds_trans: array-like
23 :return: Scaled values.
24 :rtype: numpy.ndarray
25 """
26 xs = x - param_bounds[0]
27 xs = (
28 xs
29 / (param_bounds[1] - param_bounds[0])
30 * (param_bounds_trans[1] - param_bounds_trans[0])
31 )
32 xs = xs + param_bounds_trans[0]
33 return xs
36def scale_inv(x, param_bounds, param_bounds_trans):
37 """
38 Scale values from transformed bounds back to original bounds.
40 :param x: Values to scale.
41 :type x: numpy.ndarray
42 :param param_bounds: Original parameter bounds [min, max].
43 :type param_bounds: array-like
44 :param param_bounds_trans: Transformed parameter bounds [min, max].
45 :type param_bounds_trans: array-like
46 :return: Scaled values.
47 :rtype: numpy.ndarray
48 """
49 xs = x - param_bounds_trans[0]
50 xs = (
51 xs
52 * (param_bounds[1] - param_bounds[0])
53 / (param_bounds_trans[1] - param_bounds_trans[0])
54 )
55 xs = xs + param_bounds[0]
56 return xs
59def trans_fwd(x, param_bounds, param_bounds_trans):
60 """
61 Transform values forward using normal PPF (probit transform).
63 :param x: Values to transform.
64 :type x: numpy.ndarray
65 :param param_bounds: Original parameter bounds [min, max].
66 :type param_bounds: array-like
67 :param param_bounds_trans: Transformed parameter bounds [min, max].
68 :type param_bounds_trans: array-like
69 :return: Transformed values.
70 :rtype: numpy.ndarray
71 """
72 xs = scale_fwd(x, param_bounds, param_bounds_trans)
73 ppfx = norm.ppf(xs)
74 if np.any(~np.isfinite(ppfx)):
75 import ipdb
77 ipdb.set_trace()
78 return ppfx
81def trans_inv(x, param_bounds, param_bounds_trans):
82 """
83 Transform values back using normal CDF (inverse probit transform).
85 :param x: Values to transform.
86 :type x: numpy.ndarray
87 :param param_bounds: Original parameter bounds [min, max].
88 :type param_bounds: array-like
89 :param param_bounds_trans: Transformed parameter bounds [min, max].
90 :type param_bounds_trans: array-like
91 :return: Inverse-transformed values.
92 :rtype: numpy.ndarray
93 """
94 xi = norm.cdf(x)
95 xi = scale_inv(xi, param_bounds, param_bounds_trans)
96 return xi
99class TransformedGaussianMixture:
100 """
101 Gaussian Mixture Model with parameter transformation.
103 This class wraps sklearn's GaussianMixture to work with bounded parameters
104 by transforming them to an unbounded space using a probit (normal CDF/PPF)
105 transformation.
107 :param param_bounds: Parameter bounds for each dimension, shape (n_dims, 2).
108 If None, bounds are inferred from the data during fit.
109 :type param_bounds: numpy.ndarray or None
110 :param args: Positional arguments passed to GaussianMixture.
111 :param kwargs: Keyword arguments passed to GaussianMixture.
112 """
114 def __init__(self, param_bounds=None, *args, **kwargs):
115 self.eps = 1e-8
116 self.bounds_trans = [self.eps, 1 - self.eps]
117 self.param_bounds = param_bounds
119 self.gm = GaussianMixture(*args, **kwargs)
121 def set_bounds(self, X):
122 """
123 Set parameter bounds from data if not already set.
125 :param X: Training data, shape (n_samples, n_features).
126 :type X: numpy.ndarray
127 """
128 if self.param_bounds is None:
129 self.param_bounds = np.array(
130 [
131 np.min(X, axis=0) - 10 * self.eps,
132 np.max(X, axis=0) + 10 * self.eps,
133 ]
134 ).T
136 def fit(self, X, y=None):
137 """
138 Fit the Gaussian Mixture Model.
140 :param X: Training data, shape (n_samples, n_features).
141 :type X: numpy.ndarray
142 :param y: Ignored (for sklearn compatibility).
143 :return: self
144 """
145 self.set_bounds(X)
146 X_trans = self._transform_params_forward(X)
147 return self.gm.fit(X_trans, y)
149 def fit_predict(self, X, y=None):
150 """
151 Fit and predict cluster labels.
153 :param X: Training data, shape (n_samples, n_features).
154 :type X: numpy.ndarray
155 :param y: Ignored (for sklearn compatibility).
156 :return: Component labels for each sample.
157 :rtype: numpy.ndarray
158 """
159 self.set_bounds(X)
160 X_trans = self._transform_params_forward(X)
161 return self.gm.fit_predict(X_trans, y)
163 def predict_proba(self, X):
164 """
165 Predict posterior probability of each component given the data.
167 :param X: Data points, shape (n_samples, n_features).
168 :type X: numpy.ndarray
169 :return: Posterior probabilities, shape (n_samples, n_components).
170 :rtype: numpy.ndarray
171 """
172 X_trans = self._transform_params_forward(X)
173 return self.gm.predict_proba(X_trans)
175 def sample(self, n_samples=1):
176 """
177 Generate random samples from the fitted Gaussian mixture.
179 :param n_samples: Number of samples to generate.
180 :type n_samples: int
181 :return: Tuple of (samples, component_labels).
182 :rtype: tuple(numpy.ndarray, numpy.ndarray)
183 """
184 X_trans, y = self.gm.sample(n_samples)
185 X = self._transform_params_inverse(X_trans)
186 return X, y
188 def score(self, X, y=None):
189 """
190 Compute the per-sample average log-likelihood.
192 :param X: Data points, shape (n_samples, n_features).
193 :type X: numpy.ndarray
194 :param y: Ignored (for sklearn compatibility).
195 :return: Log-likelihood of the data.
196 :rtype: float
197 """
198 X_trans = self._transform_params_forward(X)
199 return self.gm.score(X_trans)
201 def score_samples(self, X):
202 """
203 Compute the log-likelihood of each sample.
205 :param X: Data points, shape (n_samples, n_features).
206 :type X: numpy.ndarray
207 :return: Log-likelihood for each sample.
208 :rtype: numpy.ndarray
209 """
210 X_trans = self._transform_params_forward(X)
211 return self.gm.score_samples(X_trans)
213 def set_params(self, **params):
214 """
215 Set parameters of the underlying GaussianMixture.
217 :param params: Parameters to set.
218 :return: self
219 """
220 return self.gm.set_params(**params)
222 def get_params(self, deep=True):
223 """
224 Get parameters of the underlying GaussianMixture.
226 :param deep: If True, return parameters for sub-objects.
227 :type deep: bool
228 :return: Parameter names mapped to their values.
229 :rtype: dict
230 """
231 return self.gm.get_params(deep)
233 def bic(self, X):
234 """
235 Compute the Bayesian Information Criterion.
237 :param X: Data points, shape (n_samples, n_features).
238 :type X: numpy.ndarray
239 :return: BIC score.
240 :rtype: float
241 """
242 X_trans = self._transform_params_forward(X)
243 return self.gm.bic(X_trans)
245 def aic(self, X):
246 """
247 Compute the Akaike Information Criterion.
249 :param X: Data points, shape (n_samples, n_features).
250 :type X: numpy.ndarray
251 :return: AIC score.
252 :rtype: float
253 """
254 X_trans = self._transform_params_forward(X)
255 return self.gm.aic(X_trans)
257 def _transform_params_forward(self, X):
258 """
259 Transform parameters from original to unbounded space.
261 :param X: Data in original space.
262 :type X: numpy.ndarray
263 :return: Data in transformed space.
264 :rtype: numpy.ndarray
265 """
266 X_trans = X.copy()
267 for i in range(X.shape[1]):
268 X_trans[:, i] = trans_fwd(
269 x=np.array(X[:, i]),
270 param_bounds=self.param_bounds[i],
271 param_bounds_trans=self.bounds_trans,
272 )
274 return X_trans
276 def _transform_params_inverse(self, X_trans):
277 """
278 Transform parameters from unbounded back to original space.
280 :param X_trans: Data in transformed space.
281 :type X_trans: numpy.ndarray
282 :return: Data in original space.
283 :rtype: numpy.ndarray
284 """
285 X = X_trans.copy()
286 for i in range(X.shape[1]):
287 X[:, i] = trans_inv(
288 x=np.array(X_trans[:, i]),
289 param_bounds=self.param_bounds[i],
290 param_bounds_trans=self.bounds_trans,
291 )
293 return X