Coverage for /opt/conda/lib/python3.14/site-packages/medil/models.py: 89%
247 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 02:29 +0000
1"""MeDIL causal model classes for linear Gaussian and deep generative settings."""
3import copy
4import os
5import pickle
6import random
7import warnings
8from datetime import datetime
9from pathlib import Path
11import numpy as np
12import numpy.typing as npt
13from numpy.random import default_rng
14from scipy.optimize import minimize
15from sklearn.model_selection import train_test_split
17from ._ecc_algorithms import _find_heuristic_1pc
18from ._independence_testing import _estimate_UDG
20try:
21 import torch
22 import torch.nn.functional as F
23 from torch.utils.data import DataLoader, TensorDataset
24 from tqdm import tqdm
25 from ._vae import VariationalAutoencoder
26 _TORCH_AVAILABLE = True
27except ImportError:
28 _TORCH_AVAILABLE = False
31class _MedilCausalModel(object):
32 """Base class using principle of polymorphism to establish common
33 interface for derived parametric estimators.
34 """
36 def __init__(
37 self,
38 biadj: npt.NDArray = np.array([]),
39 udg: npt.NDArray = np.array([]),
40 one_pure_child: bool = True,
41 rng=default_rng(0),
42 ) -> None:
43 self.biadj = biadj
44 self.udg = udg
45 self.one_pure_child = one_pure_child
46 self.rng = rng
48 def fit(self, dataset: npt.NDArray) -> "_MedilCausalModel":
49 raise NotImplementedError
51 def sample(self, sample_size: int) -> npt.NDArray:
52 raise NotImplementedError
55class _Parameters(object):
56 "Different parameterizations of MeDIL causal Models."
58 def __init__(self, parameterization: str) -> None:
59 self.parameterization = parameterization
61 if parameterization == "Gaussian":
62 self.error_means = np.array([])
63 self.error_variances = np.array([])
64 self.biadj_weights = np.array([])
65 elif parameterization == "VAE":
66 self.weights = np.array([])
67 self.vae = None
69 def __str__(self) -> str:
70 return "\n".join(
71 f"parameters.{attr}: {val}" for attr, val in vars(self).items()
72 )
75class GaussianMCM(_MedilCausalModel):
76 """Linear Gaussian MeDIL causal model.
78 Learns a bipartite latent→measurement causal structure and estimates
79 linear Gaussian parameters (edge weights, error means, error variances)
80 by constraint-based structure learning and least-squares optimization
81 of the covariance matrix.
83 Parameters
84 ----------
85 biadj : ndarray of shape (num_latent, num_meas), optional
86 Boolean biadjacency matrix. If empty (default), estimated from data
87 during :meth:`fit`.
88 udg : ndarray of shape (num_meas, num_meas), optional
89 Boolean undirected dependence graph over observed variables. If empty
90 (default), estimated from data during :meth:`fit`.
91 rng : numpy.random.Generator, optional
92 Random number generator used during :meth:`sample`.
94 Attributes
95 ----------
96 biadj : ndarray of shape (num_latent, num_meas)
97 Boolean biadjacency matrix (set after :meth:`fit` or at init).
98 parameters : object
99 Learned parameters with attributes ``biadj_weights``
100 (shape ``(num_latent, num_meas)``), ``error_means``
101 (shape ``(num_meas,)``), and ``error_variances``
102 (shape ``(num_meas,)``).
103 """
105 def __init__(self, **kwargs):
106 super().__init__(**kwargs)
107 self.parameters = _Parameters("Gaussian")
109 def fit(self, dataset: npt.NDArray) -> "GaussianMCM":
110 """Fit a GaussianMCM to a dataset.
112 Estimates the biadjacency matrix via constraint-based structure
113 learning (if not pre-specified), then estimates edge weights and
114 error variances by least-squares optimization of the covariance.
116 Parameters
117 ----------
118 dataset : ndarray of shape (n_samples, n_features)
119 Observed data matrix. Rows are observations, columns are variables.
121 Returns
122 -------
123 self : GaussianMCM
124 """
125 self.dataset = dataset
126 if self.biadj.size == 0:
127 self._compute_biadj()
129 self.parameters.error_means = self.dataset.mean(0)
131 cov = np.cov(self.dataset, rowvar=False)
133 num_weights = self.biadj.sum()
134 num_err_vars = self.biadj.shape[1]
136 def _objective(weights_and_err_vars):
137 weights = weights_and_err_vars[:num_weights]
138 err_vars = weights_and_err_vars[num_weights:]
140 biadj_weights = np.zeros_like(self.biadj, float)
141 biadj_weights[self.biadj] = weights
143 return (
144 (cov - biadj_weights.T @ biadj_weights - np.diagflat(err_vars)) ** 2
145 ).sum()
147 result = minimize(_objective, np.ones(num_weights + num_err_vars))
148 if not result.success:
149 warnings.warn(f"Optimization failed: {result.message}")
151 self.parameters.error_variances = result.x[num_weights:]
153 self.parameters.biadj_weights = np.zeros_like(self.biadj, float)
154 self.parameters.biadj_weights[self.biadj] = result.x[:num_weights]
156 return self
158 def _compute_biadj(self):
159 """Constraint-based structure learning."""
160 if self.udg.size == 0:
161 self._estimate_udg()
162 self.biadj = _find_heuristic_1pc(self.udg)
164 def _estimate_udg(self):
165 """Constraint-based structure learning."""
166 samp_size = len(self.dataset)
167 cov = np.cov(self.dataset, rowvar=False)
168 corr = np.corrcoef(self.dataset, rowvar=False)
169 inner_numerator = 1 - cov * corr # should never be <= 0?
170 inner_numerator = inner_numerator.clip(min=0.00001)
171 inner_numerator[np.tril_indices_from(inner_numerator)] = 1
172 udg_triu = np.log(inner_numerator) < (-np.log(samp_size) / samp_size)
173 udg = udg_triu + udg_triu.T
174 self.udg = udg
176 def sample(self, sample_size: int, include_latent: bool = False) -> npt.NDArray:
177 """Sample observations from a GaussianMCM.
179 Requires ``biadj`` and ``parameters`` to be set, either by calling
180 :meth:`fit` or by constructing the model via :func:`medil.sample.mcm`.
182 Parameters
183 ----------
184 sample_size : int
185 Number of observations to draw.
186 include_latent : bool, optional
187 If True, also return the sampled latent variables.
189 Returns
190 -------
191 samples : ndarray of shape (sample_size, num_meas)
192 latent_samples : ndarray of shape (sample_size, num_latent), only if include_latent=True
193 """
194 num_latent = len(self.biadj)
195 latent_sample = self.rng.multivariate_normal(
196 np.zeros(num_latent), np.eye(num_latent), sample_size
197 )
198 error_sample = self.rng.multivariate_normal(
199 self.parameters.error_means,
200 np.diagflat(self.parameters.error_variances),
201 sample_size,
202 )
203 sample = latent_sample @ self.parameters.biadj_weights + error_sample
205 return (sample, latent_sample) if include_latent else sample
208class NeuroCausalFactorAnalysis(_MedilCausalModel):
209 """Nonlinear MeDIL causal model represented by a masked variational autoencoder.
211 Learns a nonlinear MeDIL causal model in two phases: pairwise independence
212 tests identify the causal factor structure (``biadj``), then a masked VAE
213 whose decoder connectivity encodes that structure is trained to learn
214 nonlinear generative mechanisms :cite:`markham2023neuro`.
216 Requires PyTorch: ``pip install medil[ncfa]``.
218 For continuous data, input should be standardized (zero mean, unit variance
219 per feature) before calling :meth:`fit`. For categorical data, pass raw
220 class indices (integers 0 to K-1) and set ``hyperparams["num_classes"] = K``.
222 Parameters
223 ----------
224 seed : int, optional
225 Random seed for reproducibility. Default 0.
226 log_path : str, optional
227 Directory for training artifacts (model weights, loss history).
228 Created if it does not exist. No artifacts written if empty (default).
229 verbose : bool, optional
230 Print timestamped training log entries to stdout. Default False.
231 biadj : ndarray of shape (num_latent, num_meas), optional
232 Boolean biadjacency matrix. If empty (default), estimated from data
233 during :meth:`fit` using xi correlation.
234 **kwargs
235 Additional keyword arguments passed to the base class (``udg``, ``rng``).
237 Attributes
238 ----------
239 biadj : ndarray of shape (num_latent, num_meas)
240 Boolean biadjacency matrix (set after :meth:`fit` or at init).
241 parameters : object
242 Learned VAE, accessible as ``parameters.vae``.
243 loss : dict or None
244 Train/validation ELBO and reconstruction losses after :meth:`fit`,
245 keyed by ``"elbo_train"``, ``"elbo_valid"``, ``"recon_train"``,
246 ``"recon_valid"``.
247 hyperparams : dict
248 Training hyperparameters. Modify via ``model.hyperparams.update({...})``
249 before calling :meth:`fit`. Keys:
251 - ``"method"`` : independence test for structure learning
252 (``"xicor"``, ``"dcov_fast"``, or ``"g-test"`` for integer data;
253 default ``"xicor"``)
254 - ``"alpha"`` : significance level for independence tests (default 0.05)
255 - ``"num_epochs"`` : maximum training epochs (default 200)
256 - ``"lr"`` : AdamW learning rate (default 1e-3)
257 - ``"beta"`` : KL weight in the ELBO (default 1.0)
258 - ``"latent_width"`` : hidden units per latent variable in the decoder (default 2)
259 - ``"meas_width"`` : hidden units per measurement variable in the decoder (default 2)
260 - ``"num_hidden_layers"`` : number of decoder hidden layers (default 1)
261 - ``"encoder_hidden_dim"`` : hidden dimension of the encoder MLP (default 64)
262 - ``"batch_size"`` : mini-batch size (default 128)
263 - ``"early_stopping"`` : stop when validation ELBO stagnates (default True)
264 - ``"patience"`` : early stopping patience in epochs (default 20)
265 - ``"min_delta"`` : minimum ELBO improvement to reset patience (default 1e-4)
266 - ``"num_classes"`` : number of categories per measurement (1 = continuous
267 with MSE reconstruction, K ≥ 2 = categorical with cross-entropy loss;
268 a single K is applied uniformly to all measurements — if some variables
269 have fewer than K categories the model trains correctly but
270 :meth:`sample` may return out-of-range class indices; default 1)
271 """
273 def __init__(
274 self,
275 seed: int = 0,
276 log_path: str = "",
277 verbose: bool = False,
278 **kwargs,
279 ):
280 if not _TORCH_AVAILABLE:
281 raise ImportError(
282 "NeuroCausalFactorAnalysis requires PyTorch. "
283 "Install it with: pip install medil[ncfa]"
284 )
285 super().__init__(**kwargs)
287 if log_path:
288 Path(log_path).mkdir(exist_ok=True)
290 self.log_path = log_path
291 self.verbose = verbose
292 self.seed = seed
294 self.hyperparams = {
295 "method": "xicor",
296 "alpha": 0.05,
297 "batch_size": 128,
298 "num_epochs": 200,
299 "lr": 1e-3,
300 "beta": 1.0,
301 "latent_width": 2,
302 "meas_width": 2,
303 "num_hidden_layers": 1,
304 "encoder_hidden_dim": 64,
305 "shuffle": True,
306 "early_stopping": True,
307 "patience": 20,
308 "min_delta": 1e-4,
309 "num_classes": 1,
310 }
312 self.parameters = _Parameters("VAE")
313 self.loss = None
314 self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
316 def _log(self, entry: str) -> None:
317 if not (self.log_path or self.verbose):
318 return
320 time_stamped_entry = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} {entry}"
322 if self.log_path:
323 with open(os.path.join(self.log_path, "training.log"), "a") as log_file:
324 log_file.write(time_stamped_entry + "\n")
326 if self.verbose:
327 print(time_stamped_entry)
329 def _set_deterministic_seed(self):
330 os.environ["PYTHONHASHSEED"] = str(self.seed)
331 random.seed(self.seed)
332 np.random.seed(self.seed)
334 torch.manual_seed(self.seed)
335 if torch.cuda.is_available():
336 torch.cuda.manual_seed(self.seed)
337 torch.cuda.manual_seed_all(self.seed)
339 torch.backends.cudnn.deterministic = True
340 torch.backends.cudnn.benchmark = False
341 torch.use_deterministic_algorithms(True)
343 def fit(self, dataset: npt.NDArray, split_idcs=None) -> "NeuroCausalFactorAnalysis":
344 """Fit a NeuroCausalFactorAnalysis model to a dataset using a masked VAE.
346 Parameters
347 ----------
348 dataset : ndarray of shape (n_samples, n_features)
349 Observed data matrix. Should be standardized (zero mean, unit
350 variance per feature) for best results.
351 split_idcs : tuple of index arrays, optional
352 (train_indices, valid_indices). If None, a 70/30 train/valid split
353 is created automatically using self.seed.
355 Returns
356 -------
357 self : NeuroCausalFactorAnalysis
358 """
359 self._set_deterministic_seed()
360 self.dataset = dataset
362 if self.biadj.size == 0:
363 self._compute_biadj()
365 if split_idcs is None:
366 train_split, valid_split = train_test_split(
367 dataset, train_size=0.7, random_state=self.seed
368 )
369 else:
370 train_split = dataset[split_idcs[0]]
371 valid_split = dataset[split_idcs[1]]
373 train_loader = self._data_loader(train_split)
374 valid_loader = self._data_loader(valid_split)
376 model_recon, loss_recon, error_recon = self._train_vae(
377 train_loader, valid_loader
378 )
380 if self.log_path:
381 torch.save(
382 model_recon.state_dict(), os.path.join(self.log_path, "model_recon.pt")
383 )
384 with open(os.path.join(self.log_path, "loss_recon.pkl"), "wb") as handle:
385 pickle.dump(loss_recon, handle, protocol=pickle.HIGHEST_PROTOCOL)
386 with open(os.path.join(self.log_path, "error_recon.pkl"), "wb") as handle:
387 pickle.dump(error_recon, handle, protocol=pickle.HIGHEST_PROTOCOL)
389 self.parameters.vae = model_recon
390 self.loss = {
391 "elbo_train": loss_recon[0],
392 "elbo_valid": loss_recon[1],
393 "recon_train": error_recon[0],
394 "recon_valid": error_recon[1],
395 }
397 return self
399 def _compute_biadj(self):
400 if self.udg.size == 0:
401 self._estimate_udg()
402 self.biadj = _find_heuristic_1pc(self.udg)
404 def _estimate_udg(self):
405 self.udg, _ = _estimate_UDG(
406 self.dataset,
407 method=self.hyperparams["method"],
408 significance_level=self.hyperparams["alpha"],
409 )
411 def _data_loader(self, sample):
412 sample_x = sample.astype(np.float32)
413 dataset = TensorDataset(torch.tensor(sample_x))
414 return DataLoader(
415 dataset,
416 batch_size=self.hyperparams["batch_size"],
417 shuffle=self.hyperparams["shuffle"],
418 num_workers=0,
419 )
421 def _train_vae(self, train_loader, valid_loader):
422 num_meas = self.dataset.shape[1]
423 biadj = torch.tensor(self.biadj.T, dtype=torch.float32)
424 num_classes = self.hyperparams["num_classes"]
426 model = VariationalAutoencoder(
427 num_latent=biadj.shape[1],
428 num_meas=num_meas,
429 num_hidden_layers=self.hyperparams["num_hidden_layers"],
430 latent_width=self.hyperparams["latent_width"],
431 meas_width=self.hyperparams["meas_width"],
432 biadj=biadj,
433 encoder_hidden_dim=self.hyperparams["encoder_hidden_dim"],
434 num_classes=num_classes,
435 ).to(self.device)
437 optimizer = torch.optim.AdamW(model.parameters(), lr=self.hyperparams["lr"])
439 num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
440 self._log(f"Number of parameters: {num_params}")
442 train_elbo, train_error = [], []
443 valid_elbo, valid_error = [], []
445 best_valid = float("inf")
446 best_state = copy.deepcopy(model.state_dict())
447 epochs_without_improvement = 0
449 pbar = tqdm(
450 range(self.hyperparams["num_epochs"]), desc="Training NCFA", unit="epoch"
451 )
453 for epoch in pbar:
454 model.train()
456 for (x_batch,) in train_loader:
457 x_batch = x_batch.to(self.device)
459 x_recon, mu, logvar = model(x_batch)
460 loss = self._vae_loss(
461 x_batch, x_recon, mu, logvar, beta=self.hyperparams["beta"],
462 num_classes=num_classes,
463 )
465 optimizer.zero_grad()
466 loss.backward()
467 optimizer.step()
469 train_lb, train_er = self._eval_loss(model, train_loader, num_classes)
470 train_elbo.append(train_lb)
471 train_error.append(train_er)
473 valid_lb, valid_er = self._eval_loss(model, valid_loader, num_classes)
474 valid_elbo.append(valid_lb)
475 valid_error.append(valid_er)
477 pbar.set_postfix({"train": train_lb, "valid": valid_lb})
479 improved = valid_lb < (best_valid - self.hyperparams["min_delta"])
480 if improved:
481 best_valid = valid_lb
482 best_state = copy.deepcopy(model.state_dict())
483 epochs_without_improvement = 0
484 else:
485 epochs_without_improvement += 1
487 if (
488 self.hyperparams["early_stopping"]
489 and epochs_without_improvement >= self.hyperparams["patience"]
490 ):
491 self._log(f"Early stopping at epoch {epoch}")
492 break
494 model.load_state_dict(best_state)
496 return (
497 model,
498 [np.array(train_elbo), np.array(valid_elbo)],
499 [np.array(train_error), np.array(valid_error)],
500 )
502 def _eval_loss(self, model, loader, num_classes=1):
503 model.eval()
504 total_loss = 0.0
505 total_recon = 0.0
506 n = 0
508 with torch.no_grad():
509 for (x_batch,) in loader:
510 x_batch = x_batch.to(self.device)
511 x_recon, mu, logvar = model(x_batch)
512 loss = self._vae_loss(
513 x_batch, x_recon, mu, logvar, beta=self.hyperparams["beta"],
514 num_classes=num_classes,
515 )
516 recon = self._recon_error(x_batch, x_recon, num_classes=num_classes)
518 bs = x_batch.shape[0]
519 total_loss += loss.item()
520 total_recon += recon.item()
521 n += bs
523 return total_loss / n, total_recon / n
525 @staticmethod
526 def _vae_loss(x, x_recon, mu, logvar, beta=1.0, num_classes=1):
527 if num_classes >= 2:
528 recon_loss = F.cross_entropy(
529 x_recon.view(-1, num_classes), x.long().view(-1), reduction="sum"
530 )
531 else:
532 recon_loss = F.mse_loss(x_recon, x, reduction="sum")
533 kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
534 return recon_loss + beta * kl_div
536 @staticmethod
537 def _recon_error(x, x_recon, num_classes=1):
538 if num_classes >= 2:
539 return F.cross_entropy(
540 x_recon.view(-1, num_classes), x.long().view(-1), reduction="sum"
541 )
542 return torch.linalg.norm(x - x_recon, ord=2)
544 def sample(self, sample_size: int, include_latent: bool = False) -> npt.NDArray:
545 """Sample observations from a fitted NeuroCausalFactorAnalysis model.
547 Parameters
548 ----------
549 sample_size : int
550 Number of samples to generate.
551 include_latent : bool, optional
552 If True, also return the latent codes z drawn from the prior.
553 Note: z has shape (sample_size, num_latent * latent_width).
555 Returns
556 -------
557 sample : ndarray of shape (sample_size, num_meas)
558 latent_sample : ndarray of shape (sample_size, latent_dim), only if include_latent=True
559 """
560 if self.parameters.vae is None:
561 raise ValueError("Model must be fitted before sampling.")
562 vae = self.parameters.vae
563 latent_dim = vae.decoder.latent_dim
564 num_classes = self.hyperparams["num_classes"]
565 z = torch.randn(sample_size, latent_dim, device=self.device)
566 with torch.no_grad():
567 vae.eval()
568 x_recon = vae.decoder(z)
569 if num_classes >= 2:
570 num_meas = vae.decoder.num_meas
571 probs = torch.softmax(x_recon.view(sample_size, num_meas, num_classes), dim=-1)
572 x_recon = torch.multinomial(probs.view(-1, num_classes), 1).view(sample_size, num_meas).float()
573 out = x_recon.cpu().numpy()
574 return (out, z.cpu().numpy()) if include_latent else out
576 def _set_full_decoder_mask(self, num_meas=None):
577 if num_meas is None:
578 if not hasattr(self, "dataset"):
579 raise ValueError("Provide num_meas or set dataset first.")
580 num_meas = self.dataset.shape[1]
582 num_latent = num_meas
583 self.biadj = np.ones((num_latent, num_meas), dtype=float)