class AMICA:
"""
Adaptive Mixture ICA (AMICA) implementation.
This class implements the AMICA algorithm for blind source separation using
adaptive mixtures of independent component analyzers.
The algorithm provides two progress reporting modes:
1. A modern tqdm progress bar (default) showing overall progress and key metrics
2. Detailed per-line progress output in the style of the original Fortran implementation
(enabled with verbose=True or use_tqdm=False)
"""
def __init__(
self,
params_file: Optional[str] = None,
use_tqdm: bool = True,
verbose: bool = False,
**kwargs,
):
"""
Initialize AMICA with parameters.
Parameters
----------
params_file : str, optional
Path to a JSON or Fortran-format (``input.param``) parameter
file with default values (issue #304).
use_tqdm : bool, default=True
Whether to use tqdm progress bar (False will use per-line printing)
verbose : bool, default=False
Whether to enable verbose output (will use per-line printing regardless of use_tqdm)
**kwargs : dict
Override default parameters with these values.
``do_opt_block`` (False), ``blk_min`` (4096), ``blk_max`` (32768),
``blk_step`` (4096) carry AMICATorchNG's names, defaults and
semantics (issue #232): with ``do_opt_block`` on, ``fit`` times each
candidate block size on the real data and keeps the fastest instead
of using ``block_size`` as given. The choice is timing-based and so
machine-dependent, which is why it is off by default -- a run
compared against the reference binary must leave it off and pin
``block_size``. Unlike Fortran, a candidate that cannot be
allocated is skipped rather than aborting the run. See
:mod:`pamica.blocktune`.
``n_restarts`` (1) and ``restart_seeds`` (None) carry
AMICATorchNG's names, defaults and semantics (issue #198): run the
fit from ``n_restarts`` different seeds and keep the one with the
highest ``ll[-1]``. ``1`` is the parity-preserving default and
bypasses the restart machinery entirely. ``n_restarts > 1`` requires
a base ``seed`` (or an explicit ``restart_seeds`` list of exactly
that length) so the winner is reproducible, and costs
``n_restarts`` times as long -- restarts run serially. This is a
pamica extension: Fortran has no search over seeds, and it is
unrelated to ``maxrestarts``/``restartiter``, its recovery path
after an early non-finite likelihood. See :mod:`pamica.restarts`.
``pcakeep`` (None) and ``pcadb`` (None) carry AMICATorchNG's
names, defaults, validation and semantics (issue #323): explicit
PCA reduction capped by the detected numerical rank. ``pcakeep``
must be an integer >= 1 and ``pcadb`` a finite number > 0, or the
constructor raises ``ValueError``; when both are set ``pcakeep``
takes precedence and ``pcadb`` is ignored, as in the reference
(which parses ``pcadb`` but never uses it). Both are ignored, with
one WARNING, when ``do_sphere`` is False, as in the reference. See
:mod:`pamica.rank`.
``newt_start`` (20), ``rejstart`` (2), ``writestep`` (100) and
``histstep`` (10) name iterations counted from 1, as the
reference's do (issue #335): Newton takes its first step on
iteration ``newt_start``, the first rejection follows iteration
``rejstart``, and a checkpoint lands on every ``writestep``-th
iteration. ``newt_start`` must be an integer >= 0 (0 and 1 fit
identically), ``rejstart`` an integer >= 1 when ``do_reject`` is
on, and ``histstep`` an integer >= 1 when ``do_history`` is on,
or the constructor raises ``ValueError`` (for ``newt_start`` and
``rejstart``, with the same messages as the PyTorch and MLX
backends). See :mod:`pamica.schedule`.
``restartiter`` (10) and ``maxrestarts`` (3) configure the
restart-on-NaN recovery (Fortran amica15.f90:1022-1050): a
non-finite likelihood within the first ``restartiter`` iterations,
counted from 1, redraws the mixing matrix and continues, up to
``maxrestarts`` times; ``restartiter=0`` disables it, as in the
reference. Both must be integers >= 0. The restart count, and what
the fit does after a restart, differ from the reference; see
``docs/guides/amica-differences.md``. The PyTorch and MLX backends
have no restart-on-NaN path: they stop on a non-finite likelihood.
``doscaling`` (True) and ``scalestep`` (1) carry AMICATorchNG's
names, defaults, validation and semantics (issue #333): each
component's mixing vector (a row of the stored ``A``, issue #334)
is rescaled to unit norm, with the matching ``mu``/``beta``
rescale, on iterations ``scalestep``, ``2*scalestep``, ...
counted from 1 (see :meth:`_rescale_components`). ``scalestep`` is
validated only when ``doscaling`` is on (an integer >= 1, or the
constructor raises ``ValueError``); with ``doscaling`` off it is
inert, never read.
``outdir`` (None) is where the fit writes its ``out.txt`` log,
its ``writestep`` checkpoints, its ``do_history`` snapshots and its
final results, in the Fortran ``amicaout`` layout. The default
writes no files at all, like the PyTorch and MLX backends; set it
(here, in a params file, or with the CLI's ``--outdir``) to get the
reference's on-disk output.
``min_nd``, ``maxdecs`` and ``share_iter`` -- the canonical
spelling of ``min_grad_norm``, ``max_decs`` and ``share_int``
(this backend's own names), which matches ``AMICATorchNG``'s
constructor and the spelling a params file already resolves
either way -- are also accepted here directly (issue #346) and
translated to this backend's own attribute. Passing both
spellings of the same setting at once raises ``TypeError``
rather than silently picking one.
``files``, ``data_dim`` and ``field_dim`` -- data-location
metadata (used by :meth:`fit` when it is called with no data),
normally read from ``params_file`` -- are also accepted here
directly, with the same meaning (issue #346). Setting both
``params_file`` and a keyword argument to a *different* value
for the same one of these three raises ``TypeError``; the same
value from both is not a conflict.
Raises
------
TypeError
If ``**kwargs`` holds a name this constructor does not recognize
(issue #346), or both ``params_file`` and a keyword argument set
``files``/``data_dim``/``field_dim`` to different values. An
unrecognized name gets one of: the correct name, for
``n_models``/``n_mix`` (this backend's own spelling of the same
setting); an explanation, for ``n_channels`` (inferred from the
data passed to :meth:`fit`, here and in the ``AMICA`` wrapper,
which builds the raw PyTorch and MLX constructors with it); a
name implemented on the PyTorch
backend but not this one (e.g. ``keep_best``, ``device``,
``dtype``), pointing to ``AMICA(backend='torch')``; or a generic
"unexpected keyword argument" message with a ``difflib``-based
"did you mean" suggestion when one exists. Every offending name
is checked and named in ONE error, whatever mix of these applies.
Checked before any other keyword is applied, so a typo (e.g.
``max_iters=50``) can no longer construct silently.
"""
# Reject an unrecognized keyword before anything else runs (issue
# #346): a typo or a cross-backend-only option used to construct
# silently and have no effect.
_reject_unknown_kwargs(kwargs)
# Store progress bar settings
self.use_tqdm = use_tqdm
self.verbose = verbose
# Load default parameters (issue #304: both the bundled default and a
# user-supplied params_file are read through the same translator, so
# a literal Fortran input.param works here too, not just JSON). Kept
# unstripped (raw_params) so the data-location keys below can be read
# off the same parse instead of a second file read.
raw_params = _read_numpy_keyed_params(params_file)
params = {k: v for k, v in raw_params.items() if k not in _DATA_LOCATION_KEYS}
# A params-file setting this backend does not consume is silently
# dropped nowhere else in this constructor, so name it here rather
# than letting it vanish (issue #304). Checked before the **kwargs
# override below: kwargs are explicit Python-level arguments, not
# file settings, so they are not subject to this warning.
unconsumed = sorted(set(raw_params) - _CONSUMED_KEYS)
if unconsumed:
logger.warning(
"AMICA (NumPy): %d parameter-file setting(s) have no "
"NumPy-backend equivalent and are not applied: %s",
len(unconsumed),
unconsumed,
)
# Override with any provided parameters. A canonical-spelling kwarg
# (min_nd/maxdecs/share_iter) is translated to this backend's own
# name first, exactly as a params file's setting already is (issue
# #346), so it takes effect instead of landing under a key nothing
# reads.
params.update(_translate_canonical_kwargs(kwargs))
# Store parameters
self.num_models = params.get("num_models", 1)
if self.num_models < 1:
raise ValueError(f"num_models must be >= 1, got {self.num_models}")
self.num_mix = params.get("num_mix", 3)
self.max_iter = params.get("max_iter", 100)
if self.max_iter < 1:
raise ValueError(f"max_iter must be >= 1, got {self.max_iter}")
self.do_newton = params.get("do_newton", False)
self.newt_start = params.get("newt_start", 20)
# Validated whether or not do_newton is on: newt_start also gates the
# rho-rate ceiling ratchet on the natural-gradient path
# (schedule.past_newton_start, amica15.f90:1067).
schedule.validate_iteration_setting("newt_start", self.newt_start, 0)
self.newt_ramp = params.get("newt_ramp", 10)
self.newtrate = params.get("newtrate", 0.5)
self.do_reject = params.get("do_reject", False)
self.rejsig = params.get("rejsig", 3.0)
self.rejstart = params.get("rejstart", 2)
self.rejint = params.get("rejint", 3)
self.maxrej = params.get("maxrej", 1)
if self.do_reject:
# Validate up front, matching AMICATorchNG, rather than letting a bad
# value produce a nonsensical run deep in the EM loop: rejint<1 would
# be a ZeroDivisionError in the reject schedule; rejsig<=0 breaks the
# reject-below-the-mean semantics (at 0 the threshold is the mean, so
# ~half the samples drop every pass, and negative values invert it);
# maxrej<0 is a sanity guard (it would just make rejection inert via
# the maxrej>0 schedule gate); rejstart counts from 1, so
# rejstart<=0 would silently disable the reference's unconditional
# ``iter == rejstart`` pass.
if self.rejint < 1:
raise ValueError(f"rejint must be >= 1, got {self.rejint}")
if self.rejsig <= 0:
raise ValueError(f"rejsig must be > 0, got {self.rejsig}")
if self.maxrej < 0:
raise ValueError(f"maxrej must be >= 0, got {self.maxrej}")
schedule.validate_iteration_setting("rejstart", self.rejstart, 1)
self.num_comps = params.get("num_comps", -1)
self.lrate = params.get("lrate", 0.1)
self.lrate0 = self.lrate
self.minlrate = params.get("minlrate", 1e-12)
self.lratefact = params.get("lratefact", 0.5)
self.rho0 = params.get("rho0", 1.5)
self.minrho = params.get("minrho", 1.0)
self.maxrho = params.get("maxrho", 2.0)
# Working rho rate, its ceiling and the constructor value, as the other
# backends keep them (the reference's rholrate/rholrate0,
# amica15.f90:1063-1068, 1806/1813): a decrease scales the working
# rate, a maxdecs ratchet the ceiling, and every A update resets the
# working rate to the ceiling.
self.rholrate = params.get("rholrate", 0.05)
self.rholrate_cap = self.rholrate
self.rholrate0 = self.rholrate
self.rholratefact = params.get("rholratefact", 0.1)
self.invsigmax = params.get("invsigmax", 1000.0)
self.invsigmin = params.get("invsigmin", 1e-4)
self.do_history = params.get("do_history", False)
self.histstep = params.get("histstep", 10)
if self.do_history:
# A history snapshot fires on ``mod(iter, histstep) == 0``; 0 was a
# bare ZeroDivisionError on the first iteration of the fit.
schedule.validate_iteration_setting("histstep", self.histstep, 1)
# Block-size search (issue #232). OFF by default, unlike Fortran, whose
# header default is .true.: the choice is timing-based and therefore
# machine-dependent, so a parity run has to be able to pin block_size.
# The sweep bounds are re-derived rather than copied from Fortran's
# 128-1024, which sits far below where any pamica backend peaks; the
# stepping stays Fortran's arithmetic range so a literal input.param
# means the same thing on both sides. See pamica/blocktune.py.
self.do_opt_block = params.get("do_opt_block", False)
self.block_size = params.get("block_size", 8192)
self.blk_min = params.get("blk_min", blocktune.DEFAULT_BLK_MIN)
self.blk_max = params.get("blk_max", blocktune.DEFAULT_BLK_MAX)
self.blk_step = params.get("blk_step", blocktune.DEFAULT_BLK_STEP)
if self.do_opt_block:
# Validated only when the search is on, matching do_reject and
# share_comps below: inert otherwise, and a Fortran input.param
# carrying these alongside do_opt_block=0 must stay loadable.
blocktune.validate_block_tune_params(
self.blk_min, self.blk_max, self.blk_step
)
self.share_comps = params.get("share_comps", False)
self.comp_thresh = params.get("comp_thresh", 0.99)
self.share_start = params.get("share_start", 100)
self.share_int = params.get("share_int", 100)
# The reference's A-freeze reads share_start/share_int whether or not
# share_comps is on (issue #345), so both are validated always, with
# AMICATorchNG's messages (pamica.schedule): share_start >= 1, and a
# share_int below 7, which would hold A for every iteration of every
# cycle, is rejected under both of this backend's spellings.
schedule.validate_share_start(self.share_start)
schedule.validate_share_iter(self.share_int, "share_int/share_iter")
if self.share_comps:
# Same validation (and the same reason) as AMICATorchNG: comp_thresh
# is a cosine cutoff, so it is only meaningful in (0, 1]; at 0 every
# pair of columns merges.
if not 0.0 < self.comp_thresh <= 1.0:
raise ValueError(
f"comp_thresh must be in (0, 1], got {self.comp_thresh}"
)
self.doscaling = params.get("doscaling", True)
self.scalestep = params.get("scalestep", 1)
if self.doscaling:
# Same check and message as AMICATorchNG: a zero cadence divides by
# zero mid-fit; the reference never reads scalestep.
schedule.validate_iteration_setting("scalestep", self.scalestep, 1)
self.do_sphere = params.get("do_sphere", True)
self.do_mean = params.get("do_mean", True)
self.do_approx_sphere = params.get("do_approx_sphere", True)
self.pcakeep = params.get("pcakeep")
self.pcadb = params.get("pcadb")
# Explicit PCA reduction, validated by the policy shared with the
# PyTorch and MLX backends (pamica/rank.py, issue #323): a bad value
# fails here rather than as a silently wrongly sized or degenerate fit;
# then one log line for any part of it a fit will ignore.
validate_pca_reduction(self.pcakeep, self.pcadb)
log_ignored_pca_request(self.pcakeep, self.pcadb, self.do_sphere)
# Numerical-rank floors (issue #223); see pamica/rank.py and ADR 0004.
self.mineig = params.get("mineig", MINEIG)
self.mineig_rel = params.get("mineig_rel", MINEIG_REL)
self.writestep = params.get("writestep", 100)
self.max_decs = params.get("max_decs", 5)
# Consecutive small-increase iterations tolerated before stopping
# (Fortran maxincs, amica17.f90:1087).
self.maxincs = params.get("maxincs", 5)
# Restart-on-NaN (Fortran amica15.f90:1022-1052): if the LL goes
# non-finite within the first restartiter iterations, counted from 1 as
# the reference counts them (``iter .le. restartiter``), reinitialize and
# start over, up to maxrestarts times; a later NaN stops the fit (Fortran
# exits too). restartiter=0 therefore disables the recovery, as in the
# reference. The restart count differs from the reference's; see
# docs/guides/amica-differences.md.
self.restartiter = params.get("restartiter", 10)
self.maxrestarts = params.get("maxrestarts", 3)
schedule.validate_iteration_setting("restartiter", self.restartiter, 0)
schedule.validate_iteration_setting("maxrestarts", self.maxrestarts, 0)
self.numrestarts = 0
# Set by fit(): whether the fit ended usable, and the reason it stopped.
# converged=False signals a terminal non-finite LL or non-finite fitted
# parameters (diverged or degenerate; no results written), which
# callers/CLI must surface.
self.converged = False
self.stop_reason = None
self.min_dll = params.get("min_dll", 1e-9)
# min_grad_norm is Fortran-faithful but not reachable on small
# recordings: the reference binary's own gradient norm plateaus two
# orders above it on the bundled sample, so min_dll is what ends a fit
# there. Kept rather than retuned; see "Which convergence criterion
# actually stops a fit" in docs/guides/validation.md (issue #218).
# AMICATorchNG spells this same threshold min_nd.
self.min_grad_norm = params.get("min_grad_norm", 1e-7)
self.use_min_dll = params.get("use_min_dll", True)
self.use_grad_norm = params.get("use_grad_norm", True)
# This backend implements only the generalized-Gaussian source
# density (pdftype 0): _compute_log_pdf takes no pdftype and always
# runs the GG update, so a non-zero pdftype would silently configure
# nothing (issue #304; the narrow-exception rule in
# .rules/backend_parity.md). pdftype 1 (extended-Infomax adaptive
# switcher), 2 (Gaussian), 3 (logistic) and 4 (sub-Gaussian cosh+) are
# available on AMICA/AMICATorchNG and AMICAMLXNG instead of silently
# doing nothing here.
self.pdftype = params.get("pdftype", 0)
if self.pdftype != 0:
raise NotImplementedError(
"AMICA_NumPy implements only the generalized-Gaussian source "
f"density (pdftype=0); got pdftype={self.pdftype}. pdftype "
"1 (extended-Infomax adaptive switcher), 2 (Gaussian), 3 "
"(logistic) and 4 (sub-Gaussian cosh+) are available on "
"AMICA (AMICATorchNG) and AMICAMLXNG."
)
# Where results are written. None (the default) writes nothing: no
# out.txt log, no writestep checkpoints, no history and no final
# results, the same as the PyTorch and MLX backends, which never write
# unless asked. An explicit outdir (keyword, params file or the CLI's
# --outdir) writes all of them there, as the reference binary does.
outdir = params.get("outdir")
self.outdir: Optional[Path] = None if outdir is None else Path(outdir)
# Data-source config (used by fit() when called without explicit
# data): files/data_dim/field_dim are data-specific, not
# hyperparameters, so load_default_params()/the _DATA_LOCATION_KEYS
# filter above strips them out of `params`, and **kwargs is never
# routed through `params`/`raw_params` for these three -- they are
# read directly off `raw_params` (a params file, issue #304: no
# second file read) and off `kwargs` itself here instead. Accepting
# them as keyword arguments directly, with the same meaning as in a
# params file, is PR #347 review item 1: before this, a kwarg of one
# of these three names passed the unknown-kwarg check (they are in
# _CONSUMED_KEYS/_ACCEPTED_KWARGS) but was silently inert, so
# fit()'s no-data path still said "no 'files' configured" even
# though the caller had just set one. Setting both a params file and
# a keyword argument to different values is ambiguous, so this
# raises rather than picking one, following the same precedent as
# the min_nd/maxdecs/share_iter alias conflict above
# (_translate_canonical_kwargs) -- unlike that guard, this one is
# "different values", not "either given", since a params file and an
# explicit kwarg agreeing is not a conflict.
file_files = raw_params.get("files") if params_file is not None else None
file_data_dim = raw_params.get("data_dim") if params_file is not None else None
file_field_dim = (
raw_params.get("field_dim") if params_file is not None else None
)
kwarg_files = kwargs.get("files")
kwarg_data_dim = kwargs.get("data_dim")
kwarg_field_dim = kwargs.get("field_dim")
for name, from_file, from_kwarg in (
("files", file_files, kwarg_files),
("data_dim", file_data_dim, kwarg_data_dim),
("field_dim", file_field_dim, kwarg_field_dim),
):
if (
from_file is not None
and from_kwarg is not None
and from_file != from_kwarg
):
raise TypeError(
f"AMICA_NumPy got {name!r} from both params_file and a "
f"keyword argument, with different values ({from_file!r} "
f"vs {from_kwarg!r}); keep only one."
)
self._config_files = kwarg_files if kwarg_files is not None else file_files
self._config_data_dim = (
kwarg_data_dim if kwarg_data_dim is not None else file_data_dim
)
self._config_field_dim = (
kwarg_field_dim if kwarg_field_dim is not None else file_field_dim
)
# Initialize random state
self.seed = params.get("seed")
self.rng = np.random.RandomState(self.seed)
# Best-of-N restarts (issue #198), a pamica extension: run the fit from
# several seeds and keep the highest-likelihood one. Distinct from
# ``maxrestarts``/``numrestarts`` above, which is Fortran's *recovery*
# path (redraw A after an early non-finite LL, amica15.f90:1022-1052)
# and never compares two completed fits. Resolved here so a bad
# configuration fails before any data is touched, and derived from the
# constructor seed so a second fit() repeats the same seeds even though
# fit() leaves self.seed on the winning restart.
n_restarts = params.get("n_restarts", restarts.DEFAULT_N_RESTARTS)
restart_seeds = params.get("restart_seeds")
self._restart_seeds = restarts.resolve_seeds(
n_restarts, restart_seeds, self.seed
)
self.n_restarts = int(n_restarts)
self.restart_seeds = None if restart_seeds is None else list(restart_seeds)
# Per-restart records, set by fit(): index-aligned lists of the seed each
# restart ran from, the log-likelihood it returned (NaN for a degenerate
# restart) and why it stopped. A degenerate restart is excluded from
# selection but kept here -- it is a fact about that seed.
self.restart_seeds_: List[Optional[int]] = []
self.restart_lls_: List[float] = []
self.restart_stop_reasons_: List[Optional[str]] = []
# Pristine learning-rate ceilings and block size, captured before any
# fit can ratchet them. Unlike the other two backends, this one does not
# reset the rates at fit start (``lrate0`` and ``newtrate`` are ratcheted
# in place by _check_convergence and never restored), so a restart has to
# put them back itself or restart k+1 would inherit restart k's annealed
# schedule. See _reset_for_restart.
self._pristine_state = {
"lrate": self.lrate,
"lrate0": self.lrate0,
"newtrate": self.newtrate,
"rholrate": self.rholrate,
"rholrate_cap": self.rholrate_cap,
"block_size": self.block_size,
}
# Initialize model parameters
# Mixing matrix, one component per row: (num_comps, data_dim) (#334)
self.A: Optional[np.ndarray] = None
self.W: Optional[np.ndarray] = None # Unmixing matrix
self.c: Optional[np.ndarray] = None # Bias terms
self.mu: Optional[np.ndarray] = None # Means of mixture components
self.alpha: Optional[np.ndarray] = None # Mixture weights
self.beta: Optional[np.ndarray] = None # Scale parameters
self.rho: Optional[np.ndarray] = None # Shape parameters
self.gm: Optional[np.ndarray] = None # Model weights
# Initialize data parameters
self.data_dim: Optional[int] = None
self.num_samples: Optional[int] = None
self.mean: Optional[np.ndarray] = None
self.sphere: Optional[np.ndarray] = None
# Cached pinv(sphere), the sensor-space back-map (see _pinv_sphere).
# Invalidated wherever self.sphere is (re)assigned.
self._sphere_pinv: Optional[np.ndarray] = None
self.sldet = 0.0
self.comp_list: Optional[np.ndarray] = None
self.comp_used: Optional[np.ndarray] = None
# Outlier rejection (do_reject), mirroring AMICATorchNG's good_idx: an
# index array of the currently-kept samples that only ever shrinks, plus
# its count (num_good_samples), which normalizes gm (the model weights)
# and self.ll (see _get_updates_and_likelihood, matching Fortran's
# LL(iter) = LLtmp2 / dble(numgoodsum*nw), amica15.f90:1770). Both
# None/full until fit() sets them up. numrej counts rejection passes
# (the maxrej budget). _last_ll_samples
# holds the pre-update per-sample LL that _reject_outliers thresholds
# (captured in the E-step, applied after the parameter update, matching
# the torch ordering).
self.good_idx: Optional[np.ndarray] = None
self.num_good_samples: Optional[int] = None
self.numrej = 0
self._last_ll_samples: Optional[np.ndarray] = None
# LLt stash (issue #157): the per-sample/per-model log-likelihood the
# E-step already computes, kept for the write path instead of being
# recomputed there by a second full-dataset forward pass. This is
# Fortran's design -- ``modloglik(num_models,N)``/``loglik(N)`` are
# allocated once (amica15.f90:2617-2620), filled by every E-step
# (amica15.f90:1406-1411) and simply dumped by ``write_output``
# (amica15.f90:2338-2343). Zero-filled so a ``do_reject`` sample keeps
# the zero that ``load_rej`` reads as the rejection sentinel. Memory is
# ``(num_models + 1) * n_samples * 8`` bytes, which grows with the data
# but stays far below ``self.data`` itself (``data_dim x n_samples``).
self._llt_logv: Optional[np.ndarray] = None
self._llt_ll: Optional[np.ndarray] = None
# Initialize optimization state
self.iter = 0
self.ll = [] # Log likelihood history
# ``self.ll[-1]`` is this backend's equivalent of AMICATorchNG's
# ``final_ll_``: the LL of the returned model. Exactly so when the fit
# ended on a convergence stop, which exits before that iteration's
# update as the reference does (issue #339); after max_iter the last
# iteration's update follows its LL, so it is the LL one update before
# the returned parameters, also as in the reference. Under ``share_comps``,
# if a merge fires on the LAST fit iteration, ``self.A``/``comp_list``
# are already post-merge but ``self.ll[-1]`` still reports the
# pre-merge log-likelihood -- the merge runs after that iteration's LL
# is stored, so its effect on the LL only shows up in the next
# iteration's E-step, which never runs. This matches the reference
# ordering (Fortran identify_shared_comps runs after the iteration's
# LL accumulation, amica15.f90:1856-1858), so it is documented behavior,
# not a bug (issue #269).
self.nd = [] # Gradient norm history
# Initialize Newton optimization parameters
self.sigma2: Optional[np.ndarray] = None
self.lambda_: Optional[np.ndarray] = None
self.kappa: Optional[np.ndarray] = None
# Setup logging
self._setup_logging()
def _setup_logging(self):
"""Setup logging configuration."""
# Create main logger
self.logger = logging.getLogger("AMICA")
self.logger.setLevel(logging.INFO)
# Remove any existing handlers
for handler in self.logger.handlers[:]:
self.logger.removeHandler(handler)
# Add console handler for stdout
console_handler = logging.StreamHandler()
console_formatter = logging.Formatter("%(message)s")
console_handler.setFormatter(console_formatter)
self.logger.addHandler(console_handler)
# Add file handler for out.txt, only when an output directory was given.
self.file_path: Optional[Path] = None
if self.outdir is not None:
self.outdir.mkdir(parents=True, exist_ok=True)
self.file_path = self.outdir / "out.txt"
file_handler = logging.FileHandler(self.file_path, mode="w")
file_formatter = logging.Formatter("%(message)s")
file_handler.setFormatter(file_formatter)
self.logger.addHandler(file_handler)
# Prevent propagation to avoid duplicate logging
self.logger.propagate = False
def _append_to_log_file(self, line: str) -> None:
"""Append ``line`` to ``outdir/out.txt``; a no-op without an outdir."""
if self.file_path is not None:
with open(self.file_path, "a") as f:
f.write(line + "\n")
@classmethod
def from_params_file(cls, params_file: str, **kwargs) -> "AMICA":
"""
Construct an AMICA model from a parameter file.
Equivalent to ``AMICA(params_file=params_file, **kwargs)``. Accepts
both pamica's own JSON schema and the literal Fortran ``input.param``
text format, auto-detected from the file's content (issue #304; see
:func:`pamica.fortran_params.read_params_file`) -- matches the
PyTorch wrapper's :meth:`~pamica.amica.AMICA.from_params_file`. If the
parameter file defines ``files``/``data_dim``/``field_dim``, a
subsequent call to :meth:`fit` with no arguments will load the data
described there (see :meth:`fit`).
Parameters
----------
params_file : str
Path to a JSON or Fortran-format parameter file.
**kwargs
Additional overrides passed through to the constructor.
Returns
-------
model : AMICA
An unfitted model configured from the parameter file.
"""
return cls(params_file=params_file, **kwargs)
@property
def data_dim_in(self) -> int:
"""Input channel count, i.e. the width of the sphere.
Differs from ``data_dim`` only when rank reduction shrank the model to
the detected numerical rank (issue #223). Derived rather than stored, so
it cannot drift from the sphere it describes. Mirrors
``AMICATorchNG.n_channels_in``.
"""
if self.sphere is None:
return self.data_dim if self.data_dim is not None else 0
return int(self.sphere.shape[1])
def _pinv_sphere(self) -> np.ndarray:
"""Cached ``pinv(sphere)``: the back-map from sphered to input-channel space.
This is the Fortran ``Spinv`` (amica15.f90:568-578), which the reference
also builds as a pseudo-inverse, ``Spinv(nx, numeigs)``, under rank/PCA
reduction. A pseudo-inverse rather than an inverse because reduction
leaves the sphere non-square (issue #223) and a square sphere fitted on
rank-deficient data is singular; for a full-rank square sphere the two
agree to ~1e-15. Built on first use and invalidated wherever
``self.sphere`` is (re)assigned (:meth:`_preprocess_data`), so it can
never describe a sphere other than the current one. Mirrors
``AMICATorchNG._pinv_sphere``.
"""
assert self.sphere is not None
if self._sphere_pinv is None:
if not np.isfinite(self.sphere).all():
# Only a degenerate fit (non-finite input data) gets here. Say
# so, rather than letting LAPACK report a confusing
# "ill-conditioned / repeated singular values" SVD failure.
raise RuntimeError(
"The sphere holds non-finite values, so it has no "
"pseudo-inverse: the fit is degenerate. Check the input "
"data for NaN/inf."
)
self._sphere_pinv = np.linalg.pinv(self.sphere)
return self._sphere_pinv
def _check_usable(self, action: str) -> None:
"""Refuse to serve output from a degenerate fit (issue #306; port of
``AMICATorchNG._check_usable`` for the legacy NumPy backend).
Callers first check their own unfitted marker(s) and raise the
existing ``"Model has not been fitted yet"`` ``RuntimeError``
(unchanged); this assumes a fit has actually run. NumPy's single
``converged`` flag already folds together every degenerate cause
(non-finite LL, non-finite exit parameters, or a restart error, see
:meth:`_fit_once`) that the PyTorch/MLX backends track via
``stop_reason``/``_DEGENERATE_STOP_REASONS``, so the equivalent gate
here is simply ``not self.converged``; the defense-in-depth isfinite
sweep reuses :meth:`_nonfinite_params`, the same predicate ``fit()``'s
own outcome bookkeeping and every checkpoint write already trust.
"""
if not self.converged:
raise RuntimeError(
f"Refusing to {action}: fit ended degenerate (stop_reason="
f"{self.stop_reason!r}), so the model holds non-finite "
f"parameters and would produce NaN output. Lower lrate, "
f"disable Newton, or check data conditioning, then refit."
)
nonfinite = self._nonfinite_params()
if nonfinite:
raise RuntimeError(
f"Refusing to {action}: parameters {nonfinite} hold "
f"non-finite values (stop_reason={self.stop_reason!r})."
)
def _check_input_shape(self, data: np.ndarray) -> None:
"""Validate a data array against the fitted input channel count,
mirroring :meth:`fit`'s own ``data`` validation (issue #306): a raw
matmul/broadcast error deep inside a method is less useful than this
named ``ValueError`` at the entry point."""
if data.ndim != 2:
raise ValueError(
f"data must be a 2D array (n_channels, n_samples), got shape "
f"{data.shape}"
)
if data.shape[0] != self.data_dim_in:
raise ValueError(
f"data has {data.shape[0]} channels, model expects {self.data_dim_in}"
)
def get_sensor_mixing_matrix(self, model_idx: int = 0) -> np.ndarray:
"""Mixing matrix mapped back to input-channel space.
``pinv(sphere) @ A``, of shape ``(data_dim_in, data_dim)`` -- the Fortran
``Spinv`` mapping (amica15.f90:568-578), and the only way to recover
sensor maps when rank reduction has made the sphere non-square
(issue #223). Mirrors ``AMICATorchNG.get_sensor_mixing_matrix``.
``A`` here is the true mixing matrix, the model's component rows of the
stored ``A`` transposed (issue #334 layout): the stored ``W =
inv(A[comp_list[:, h], :])`` is itself the transpose of the unmixing
(:meth:`get_weights`, issue #24 convention). So the result inverts the
spatial filter the sources come from, ``get_weights() @ sphere @ result
== I``, and each column is one component's sensor map.
"""
if self.sphere is None or self.A is None or self.comp_list is None:
raise RuntimeError("Model has not been fitted yet; call fit() first.")
self._check_usable("get the sensor mixing matrix")
A = self.A[self.comp_list[:, model_idx], :].T
return self._pinv_sphere() @ A
def get_weights(self) -> np.ndarray:
"""
Return the learned unmixing matrix for the first model.
Returns
-------
W : ndarray of shape (n_components, n_components)
Unmixing matrix for model 0.
"""
if self.W is None:
raise RuntimeError("Model has not been fitted yet; call fit() first.")
self._check_usable("get the weights (unmixing matrix)")
# Internal W = inv(A) is stored transposed relative to the true unmixing
# (the E-step forms activations as (X-c)^T @ W), so return W^T (issue #24).
# This is the raw unmixing matrix; it does not account for the per-model
# data-space center c (issue #27) -- use transform() for c-corrected
# sources. Harmless for model 0 single-model fits where c == 0.
return self.W[:, :, 0].T
def fit(self, data: Optional[np.ndarray] = None) -> "AMICA":
"""
Fit the AMICA model to the data.
Parameters
----------
data : ndarray of shape (n_channels, n_samples), optional
The input data to fit the model to. If omitted, the data is
loaded from the ``files``/``data_dim``/``field_dim`` parameters
supplied via ``params_file`` (see :meth:`from_params_file`).
Returns
-------
self : AMICA
The fitted model.
Notes
-----
Iteration order (issue #339), as in ``AMICATorchNG.fit`` and
``AMICAMLXNG.fit``: each iteration runs the E-step, then the
restart-on-NaN window, the likelihood-decrease response and the
stopping checks, and only then, unless a check fired, the update with
the rates just set, followed by the share merge, the checkpoints and
rejection. ``self.ll`` and ``self.nd`` record every iteration whose
likelihood was finite, and nothing else. A convergence stop takes no
update on its stopping iteration, so the returned parameters are the
ones whose likelihood is ``self.ll[-1]``; a fit that runs to
``max_iter`` updates on its last iteration, as the reference does. A
non-finite likelihood outside the restart window, a non-finite update
direction, and non-finite parameters after an update each stop the fit
with ``converged=False`` (see ``_DEGENERATE_STOP_REASONS``).
Under ``share_comps``, if a merge fires on the LAST iteration, the
returned ``A``/``comp_list`` are already post-merge but ``self.ll[-1]``
still reports the pre-merge log-likelihood -- the merge's effect on
the LL only shows up in the next E-step, which never runs. This
matches the reference ordering (issue #269); see the ``self.ll``
attribute's comment for detail.
With ``n_restarts > 1`` (issue #198) the fit runs once per seed in
``restart_seeds`` (serially) and the model is left holding the
highest-``ll[-1]`` non-degenerate restart -- the state a single fit from
that seed would have left. ``restart_seeds_``/``restart_lls_``/
``restart_stop_reasons_`` record every restart (NaN likelihood for a
degenerate one) and one INFO line names the winner. The default
``n_restarts=1`` bypasses the restart machinery entirely: no reseeding,
no state copy, bit-identical to a pre-#198 fit.
Only the winner is written to ``outdir`` at the end of the fit, but the
periodic ``writestep`` checkpoints of *every* restart pass through the
same files while it runs, so a losing restart's intermediate output can
appear on disk mid-fit; the final write replaces it.
"""
if data is None:
if not self._config_files:
raise ValueError(
"No data provided and no 'files' configured (via "
"params_file or as a keyword argument); either pass "
"data explicitly or set 'files'/'data_dim'/'field_dim'."
)
if self._config_data_dim is None or self._config_field_dim is None:
raise ValueError(
"No data provided and 'data_dim'/'field_dim' are not "
"both configured (via params_file or as a keyword "
"argument); either pass data explicitly or set "
"'files'/'data_dim'/'field_dim'."
)
if len(self._config_files) != len(self._config_field_dim):
raise ValueError(
f"'files' has {len(self._config_files)} entries but "
f"'field_dim' has {len(self._config_field_dim)}; "
"load_multiple_files() requires one field_dim per file "
"(a length mismatch would silently truncate to the "
"shorter list via zip())."
)
from .data import load_multiple_files
data = load_multiple_files(
self._config_files, self._config_data_dim, self._config_field_dim
)
if data.ndim != 2:
raise ValueError(
f"data must be a 2D array (n_channels, n_samples), got shape {data.shape}"
)
if data.size == 0:
raise ValueError("data must not be empty")
# __init__ already rejects a non-positive max_iter (line ~189), but
# max_iter is a plain public attribute a caller can reassign after
# construction (unlike torch/mlx, where it is a fit() parameter
# re-validated every call) -- re-check at fit entry so that path is
# not a silent bypass (PR #318 review item 5).
if self.max_iter < 1:
raise ValueError(f"max_iter must be >= 1, got {self.max_iter}")
# Log initial message
self.logger.info("Starting AMICA fitting...")
# Initialize dimensions
self.data_dim = data.shape[0]
self.num_samples = data.shape[1]
if self.num_comps == -1:
self.num_comps = self.data_dim * self.num_models
# Preprocess data
self._preprocess_data(data)
seeds = self._restart_seeds
if len(seeds) == 1:
# Single-restart path: nothing is reseeded or reset unless the
# caller passed an explicit one-element restart_seeds, so this is
# byte-for-byte the pre-#198 fit.
if self.restart_seeds is not None:
self._reset_for_restart(seeds[0])
self._fit_once()
self.restart_seeds_ = list(seeds)
self.restart_lls_ = [self._returned_ll()]
self.restart_stop_reasons_ = [self.stop_reason]
else:
self._fit_restarts(seeds)
# Always persist the final converged result. _write_results is otherwise
# only called on writestep boundaries during the loop, so a run whose
# last iteration is not a writestep multiple (or that stops early) would
# never save the final state. Guarded by the same finiteness predicate as
# every checkpoint, so a run that diverged to a non-finite LL (issue #39)
# or ended holding non-finite parameters (issue #240) cannot overwrite
# the last good on-disk result with NaNs. Under best-of-N restarts this
# writes the WINNER, whose state is live by the time it runs. Nothing
# is written without an outdir.
if self.converged and self.outdir is not None:
self._write_results()
return self
def _returned_ll(self) -> float:
"""Log-likelihood of the model this fit returns -- ``AMICATorchNG``'s
``final_ll_`` on this backend (see the ``self.ll`` comment in
``__init__``). NaN when no iteration recorded one."""
return float(self.ll[-1]) if len(self.ll) > 0 else float("nan")
def _fit_restarts(self, seeds: List[Optional[int]]) -> None:
"""Run one full fit per restart seed and keep the winner (issue #198).
Each restart is a complete :meth:`_fit_once` preceded by
:meth:`_reset_for_restart`, so nothing leaks from one restart into the
next; the winner's state is copied with :meth:`_capture_restart_state`
and reapplied at the end unless it is already live (the last restart).
"""
lls: List[float] = []
degenerate: List[bool] = []
stop_reasons: List[Optional[str]] = []
states: Dict[int, Dict[str, object]] = {}
for index, seed in enumerate(seeds):
self._reset_for_restart(seed)
crashed = False
try:
self._fit_once()
except np.linalg.LinAlgError as exc:
# A truly singular A makes get_unmixing_matrices raise
# numpy.linalg.LinAlgError instead of producing the non-finite
# likelihood the loop guards catch. Letting it propagate would
# throw away the restarts that already succeeded, so record it
# as a degenerate restart and continue. Deliberately narrow:
# LinAlgError specifically, NOT ValueError at large (its base
# class) and NOT RuntimeError (this backend's numeric fit path
# has no RuntimeError-raising failure; torch/MLX catch it in
# their loops because their linalg fails that way), so a
# caller or programming mistake still propagates.
# ``converged`` is this backend's degeneracy verdict and
# _fit_once never got to set it, so set it here.
crashed = True
self.converged = False
self.stop_reason = restarts.ERROR_STOP_REASON
self.logger.warning(
restarts.error_message(index, len(seeds), seed, exc)
)
# A crashed restart records NaN, not the last likelihood it happened
# to reach before raising: self.ll stays the true trajectory, but the
# RECORD has to say "this restart produced no result", the same value
# the other two backends put there.
ll = float("nan") if crashed else self._returned_ll()
# converged is this backend's degeneracy verdict: False means a
# non-finite likelihood or non-finite fitted parameters (issue #240).
is_degenerate = not self.converged
lls.append(ll)
degenerate.append(is_degenerate)
stop_reasons.append(self.stop_reason)
self.logger.info(
restarts.progress_message(
index, len(seeds), seed, ll, self.stop_reason, is_degenerate
)
)
# Keep only the best state seen so far: one copy at a time.
if restarts.select_best(lls, degenerate) == index:
states = {index: self._capture_restart_state()}
winner = restarts.select_best(lls, degenerate)
if winner is None:
# WARNING, matching the other two backends (see the message helper):
# the terminal "did not converge" ERROR is _fit_once's to emit.
self.logger.warning(
restarts.all_degenerate_message(len(seeds), stop_reasons)
)
else:
self.logger.info(
restarts.winner_message(winner, len(seeds), seeds[winner], lls[winner])
)
if winner != len(seeds) - 1:
self._apply_restart_state(states[winner])
self.restart_seeds_ = list(seeds)
self.restart_lls_ = lls
self.restart_stop_reasons_ = stop_reasons
def _capture_restart_state(self) -> Dict[str, object]:
"""Independent copy of every attribute the fit path writes.
The list is :data:`_RESTART_STATE_ATTRS`; ``test_restart_policy.py``
cross-checks it against the attributes the fit-path methods actually
assign, so a field added later cannot be silently dropped.
"""
return {
name: restarts.copy_state_value(getattr(self, name))
for name in self._RESTART_STATE_ATTRS
}
def _apply_restart_state(self, state: Dict[str, object]) -> None:
"""Restore the state captured by :meth:`_capture_restart_state`."""
for name, value in state.items():
setattr(self, name, value)
def _reset_for_restart(self, seed: Optional[int]) -> None:
"""Return the per-fit state to what a freshly constructed model holds,
and reseed the RNG (issue #198).
The other two backends need no such method: their
``_initialize_parameters`` redraws unconditionally and their ``fit``
resets every counter. This one initializes a parameter only ``if <param>
is None`` (so it can honor externally supplied starting values) and its
learning-rate ceilings are ratcheted in place by ``_check_convergence``
and never restored, so a second fit on the same instance would otherwise
continue from the first fit's parameters and annealed schedule. Nulling
the parameters is what makes restart *k* equal to a fresh model fit from
``seed`` -- which is the property the acceptance test pins.
"""
self.seed = seed
self.rng = np.random.RandomState(seed)
# Redrawn by _initialize_parameters once they are None.
self.A = None
self.W = None
self.mu = None
self.alpha = None
self.beta = None
self.rho = None
self.gm = None
self.c = None
self.comp_list = None
self.comp_used = None
self.sigma2 = None
self.lambda_ = None
self.kappa = None
self._llt_logv = None
self._llt_ll = None
# Per-fit bookkeeping.
self.good_idx = None
self.num_good_samples = None
self.numrej = 0
self.numrestarts = 0
self._last_ll_samples = None
self.iter = 0
self.ll = []
self.nd = []
self.converged = False
self.stop_reason = None
# Learning-rate ceilings and block size, back to their pristine values.
for name, value in self._pristine_state.items():
setattr(self, name, value)
def _fit_once(self) -> None:
"""One fit: initialize, optionally tune the block size, run the EM loop,
and record the outcome in ``converged``/``stop_reason``. Called once per
restart by :meth:`fit`, on data that is already preprocessed."""
# Initialize parameters
self._initialize_parameters()
# Block-size search (issue #232): after preprocessing and parameter
# initialization, before the first EM iteration, so it times the real
# data with the parameters the fit starts from. A no-op when off, and
# its probes leave no state behind, so a fit with the search off is
# byte-for-byte what it was before this existed.
if self.do_opt_block:
self._tune_block_size()
# Main optimization loop
self._optimize()
# Record the outcome: a degenerate stop (a non-finite likelihood that
# restart-on-NaN could not recover, a non-finite update direction, or
# non-finite parameters after an update) means the fit diverged, which
# callers/CLI must be able to detect rather than silently trusting
# model.A/W. self.ll holds only finite values, so the stop reason, not
# its last entry, is what says so.
degenerate_stop = self._is_degenerate_stop(self.stop_reason)
self.converged = len(self.ll) > 0 and not degenerate_stop
if degenerate_stop:
self.logger.error(
"AMICA did not converge: %s (after %d restart(s)); results were "
"not written.",
self.stop_reason,
self.numrestarts,
)
elif not self.converged:
self.logger.error(
"AMICA did not converge: no iteration completed; results were "
"not written."
)
else:
# A finite likelihood is not on its own proof of a usable fit: a
# parameter can go non-finite in a way the LL does not see (a
# collapsed mixture component leaves NaN in mu/beta while the LL of
# the remaining components stays finite, issue #240). Returning that
# as a success is the silent failure the project rules single out, so
# check the fitted parameters themselves, mirroring the PyTorch
# wrapper's degenerate-fit contract (issue #50): converged=False,
# stop_reason names what went non-finite, and nothing is written.
degenerate = self._nonfinite_params()
if degenerate:
self.converged = False
self.stop_reason = "Non-finite parameters at exit: " + ", ".join(
degenerate
)
self.logger.error(
"AMICA did not converge: %s; results were not written.",
self.stop_reason,
)
# Everything a fit writes, and therefore everything a restart snapshot must
# copy for the winning restart to be indistinguishable from a single fit
# from that seed (issue #198). Together with the invariants below this must
# account for every ``self.x =`` in the fit path, which
# ``test_restart_policy.py`` enforces by parsing this module -- so a field
# added to a fit-path method fails the suite until it is classified here.
_RESTART_STATE_ATTRS = (
# Fitted parameters and the Newton buffers ...
"A", "W", "c", "mu", "alpha", "beta", "rho", "gm", "comp_list",
"comp_used", "sigma2", "lambda_", "kappa",
# ... the per-fit trajectory, counters and outcome ...
"iter", "ll", "nd", "converged", "stop_reason", "numrej", "numrestarts",
"good_idx", "num_good_samples", "_last_ll_samples", "_llt_logv", "_llt_ll",
# ... the ratcheted learning-rate ceilings (this backend anneals lrate0
# and newtrate in place), the tuned block size, and the seed/RNG the
# winning restart ran from.
"lrate", "lrate0", "newtrate", "rholrate", "rholrate_cap",
"block_size", "seed", "rng",
) # fmt: skip
# Written by the fit path but identical across the restarts of one fit()
# call, because they are functions of the data alone: the preprocessed data
# itself, the preprocessing outputs, and the model sizing derived from them.
_RESTART_INVARIANT_ATTRS = (
"data", "data_dim", "num_samples", "num_comps",
"mean", "sphere", "sldet", "_sphere_pinv",
) # fmt: skip
# The stop reasons that mark a fit degenerate, in this backend's prose
# vocabulary: a non-finite likelihood (the PyTorch/MLX "nan_ll" and
# "singular_ll"), a non-finite update direction ("nan_direction"),
# non-finite parameters after an update ("nan_params", followed by ": " and
# the offenders' names, as the exit check's reason is), and a restart that
# raised. fit() reports converged=False for each (_is_degenerate_stop).
_NONFINITE_LL_REASON = "Non-finite likelihood (NaN/-inf) encountered"
_NONFINITE_DIRECTION_REASON = "Non-finite update direction (NaN/inf) encountered"
_NONFINITE_PARAMS_REASON = "Non-finite parameters after an update"
_DEGENERATE_STOP_REASONS = (
_NONFINITE_LL_REASON,
_NONFINITE_DIRECTION_REASON,
_NONFINITE_PARAMS_REASON,
restarts.ERROR_STOP_REASON,
)
# What the post-update guard checks, in the order AMICAMLXNG's guard names
# them (it adds its cached log-determinant of W).
_POST_UPDATE_CHECKED = ("A", "mu", "alpha", "beta", "rho", "gm", "c", "W")
# Every parameter a caller can read back off disk or off the object. `A` and
# `W` are the decomposition; `c`/`mu`/`alpha`/`beta`/`rho`/`gm` are the model
# a downstream `loadmodout` reads. `sphere`/`mean` are preprocessing outputs
# and are checked too because write_amicaout persists them.
_FITTED_PARAMS = (
"A",
"W",
"c",
"mu",
"alpha",
"beta",
"rho",
"gm",
"sphere",
"mean",
)
@classmethod
def _is_degenerate_stop(cls, reason: Optional[str]) -> bool:
"""Whether ``reason`` is one of :attr:`_DEGENERATE_STOP_REASONS`,
with or without the ``": names"`` suffix the parameter reason carries."""
return reason is not None and any(
reason == known or reason.startswith(f"{known}: ")
for known in cls._DEGENERATE_STOP_REASONS
)
def _nonfinite_params(self, names: Optional[Tuple[str, ...]] = None) -> List[str]:
"""Names of fitted parameters currently holding a non-finite value.
The single definition of "this state is safe to persist or report as a
success", used by :meth:`fit`'s outcome and by every write site. An empty
list is the finiteness predicate; the names themselves go into
``stop_reason`` and the skipped-write log, so a failure says *what* went
non-finite rather than only *that* something did.
Parameters not yet allocated (before ``_initialize_parameters``) are
skipped rather than treated as bad. ``names`` narrows the sweep (the
post-update guard in :meth:`_optimize` passes
:attr:`_POST_UPDATE_CHECKED`); the default is :attr:`_FITTED_PARAMS`.
"""
return [
name
for name in (self._FITTED_PARAMS if names is None else names)
if getattr(self, name, None) is not None
and not np.all(np.isfinite(np.asarray(getattr(self, name))))
]
def _write_checkpoint(self, kind: str) -> bool:
"""Write results/history for a mid-fit checkpoint, unless degenerate.
A checkpoint written from a non-finite state persists NaN parameters that
``loadmodout`` reads back without complaint, so the run's only on-disk
artifact would be corrupt while the fit continued (issue #240). Skip the
write instead -- loudly, never silently, and without touching whatever
valid output an earlier checkpoint already left on disk.
Returns True if the write happened.
"""
degenerate = self._nonfinite_params()
if degenerate:
self.logger.error(
"Skipping the %s checkpoint at iter %d: non-finite %s. The "
"previously written output on disk is left untouched.",
kind,
self.iter + 1,
", ".join(degenerate),
)
return False
if kind == "history":
self._write_history()
else:
self._write_results()
return True
def _preprocess_data(self, data: np.ndarray):
"""Preprocess the data by removing mean and sphering."""
assert self.data_dim is not None
# Remove mean if requested
if self.do_mean:
self.mean = np.mean(data, axis=1, keepdims=True)
data = data - self.mean
else:
self.mean = np.zeros((self.data_dim, 1))
# Compute sphering matrix if requested
if self.do_sphere:
# Population covariance (divide by N, bias=True), matching Fortran's
# DSYRK scatter/N -- not np.cov's default /(N-1). The two differ by a
# scalar sqrt(N/(N-1)); /(N-1) leaves a ~5e-6 sphere mismatch vs the
# reference (issue #24).
cov = np.cov(data, bias=True)
# Eigenvalue decomposition
evals, evecs = linalg.eigh(cov)
# Sort in descending order
idx = np.argsort(evals)[::-1]
evals = evals[idx]
evecs = evecs[:, idx]
# Numerical rank + explicit PCA reduction, decided by the policy
# shared with the PyTorch and MLX backends (pamica/rank.py, issue
# #223) so the three cannot disagree about how many dimensions are
# real. Fortran: numeigs = min(pcakeep, count(eigs > mineig)).
n_comp = numerical_rank(
evals,
mineig=self.mineig,
mineig_rel=self.mineig_rel,
pcakeep=self.pcakeep,
pcadb=self.pcadb,
)
V = evecs[:, :n_comp]
inv_sqrt = np.diag(1.0 / np.sqrt(evals[:n_comp]))
# Create sphering matrix
if n_comp < self.data_dim:
# Rank-reduced: the sphere is (n_comp, data_dim), so the sphered
# data come out at the kept rank instead of staying
# rank-deficient at full width (Fortran nw = numeigs,
# amica15.f90:563).
w_pca = inv_sqrt @ V.T
if self.do_approx_sphere:
# Fortran symmetrizes the reduced whitening by the orthogonal
# polar factor of the leading n_comp block of V^T
# (amica15.f90:501-508).
U_b, _, Vt_b = np.linalg.svd(evecs.T[:n_comp, :n_comp])
self.sphere = (Vt_b.T @ U_b.T) @ w_pca
else:
self.sphere = w_pca
elif self.do_approx_sphere:
# Symmetric ZCA sphere V diag(1/sqrt(eval)) V^T (Fortran
# do_approx_sphere=True, amica17.f90:480-481) -- the parity form.
# The old diag(1/sqrt)@V^T (PCA whitening) is a different,
# non-symmetric transform that breaks activation parity.
self.sphere = V @ inv_sqrt @ V.T
else:
# Non-symmetric PCA whitening D^-1/2 V^T (amica17.f90:495).
self.sphere = inv_sqrt @ V.T
# Rank reduction shrank the sphered space: size the model to the
# kept rank before _initialize_parameters allocates against
# data_dim (Fortran nw = numeigs, amica15.f90:563). No-op, and so
# bit-exact, for full-rank data.
n_kept = self.sphere.shape[0]
if n_kept != self.data_dim:
self.logger.info(
"Data covariance has numerical rank %d of %d; fitting %d "
"sources and mapping back via the sphere pseudo-inverse.",
n_kept,
self.data_dim,
n_kept,
)
if self.num_comps == self.data_dim * self.num_models:
# num_comps was derived from the channel count, so it follows
# the rank; an explicitly configured num_comps is left alone.
self.num_comps = n_kept * self.num_models
self.data_dim = n_kept
# Apply sphering
data = np.dot(self.sphere, data)
# Sphering log-determinant term of the data log-likelihood
# (Fortran sldet, amica17.f90:474): sum over kept eigenvalues of
# -0.5*log(eval). Required so the reported LL matches Fortran; its
# omission was why the NumPy LL sat ~ +1.5 instead of ~ -3.4.
self.sldet = float(-0.5 * np.sum(np.log(evals[:n_comp])))
else:
self.sphere = np.eye(self.data_dim)
self.sldet = 0.0
# self.sphere was just (re)built: any cached back-map now describes a
# stale sphere. Covers every assignment above regardless of branch.
self._sphere_pinv = None
self.data = data
def _initialize_parameters(self):
"""Initialize all model parameters."""
assert self.data_dim is not None
# Initialize component assignments
self.comp_list = np.zeros((self.data_dim, self.num_models), dtype=int)
self.comp_used = np.ones(self.num_comps, dtype=bool)
for h in range(self.num_models):
self.comp_list[:, h] = np.arange(h * self.data_dim, (h + 1) * self.data_dim)
# Initialize mixing/unmixing matrices. A holds one component per row
# (issue #334, pamica.component_layout): model h's block is rows
# h*data_dim..(h+1)*data_dim-1 under the default comp_list above. A
# drawn A is normalized to unit-norm components as the reference does
# (issue #341, pamica.initialization), here and in the restart after a
# non-finite likelihood, which redraws through this method. An A
# supplied before fit() (or left by a previous fit on this instance) is
# used as is, like the reference's loaded A, once it is checked: a
# wrong shape, a non-finite entry or a singular model block raises
# ValueError here instead of failing deep in the fit.
if self.A is None:
drawn = initial_mixing(
self.rng,
self.data_dim,
self.num_models,
fix_init=bool(getattr(self, "fix_init", False)),
)
self.A = np.zeros((self.num_comps, self.data_dim))
self.A[: drawn.shape[0], :] = drawn
else:
validate_supplied_mixing(
self.A,
(self.num_comps, self.data_dim),
self.comp_list,
owner="AMICA_NumPy",
)
# Outlier-rejection state (do_reject): start with every sample good
# (good_idx = all indices), mirroring AMICATorchNG. num_good_samples
# drives the gm/LL normalization and equals num_samples until the first
# rejection shrinks good_idx.
assert self.num_samples is not None
if self.do_reject:
self.good_idx = np.arange(self.num_samples)
self.num_good_samples = self.num_samples
# LLt stash (issue #157), Fortran's modloglik/loglik. Allocated here so
# a restart (_reinitialize_for_restart calls back into this method)
# cannot leave the pre-restart basin's per-sample values behind; the
# next E-step refills every good column, and rejected columns stay at
# the zero sentinel either way.
self._llt_logv = np.zeros((self.num_samples, self.num_models))
self._llt_ll = np.zeros(self.num_samples)
# Initialize mixture parameters
if self.mu is None:
self.mu = np.zeros((self.num_mix, self.num_comps))
for k in range(self.num_comps):
self.mu[:, k] = np.linspace(-1, 1, self.num_mix)
if not hasattr(self, "fix_init") or not self.fix_init:
self.mu[:, k] += 0.05 * (1 - 2 * self.rng.rand(self.num_mix))
if self.alpha is None:
self.alpha = np.ones((self.num_mix, self.num_comps)) / self.num_mix
if self.beta is None:
self.beta = np.ones((self.num_mix, self.num_comps))
if not hasattr(self, "fix_init") or not self.fix_init:
self.beta += 0.1 * (0.5 - self.rng.rand(self.num_mix, self.num_comps))
if self.rho is None:
self.rho = self.rho0 * np.ones((self.num_mix, self.num_comps))
if self.gm is None:
self.gm = np.ones(self.num_models) / self.num_models
# Initialize bias terms
if self.c is None:
self.c = np.zeros((self.data_dim, self.num_models))
# Initialize Newton optimization parameters
if self.do_newton:
self.sigma2 = np.ones((self.data_dim, self.num_models))
self.lambda_ = np.zeros((self.data_dim, self.num_models))
self.kappa = np.zeros((self.data_dim, self.num_models))
# Get initial unmixing matrices
self._update_unmixing_matrices()
def _reinitialize_for_restart(self):
"""Redraw the mixing matrix after a non-finite likelihood.
Matches Fortran's restart path (amica15.f90:1026-1046): it re-draws
*only* the mixing matrix ``A`` (from the already-advanced RNG, a new
random basin), normalizes each redrawn component to unit norm as the
reference does (:1039-1040, issue #341), and recomputes
``comp_list``/``W``; the last-successful
mixture parameters (``mu``/``alpha``/``beta``/``rho``/``gm``/``c``) are
kept, not cold-reset. The learning rate and the LL/gradient-norm history
are reset so the restarted run is judged from scratch; preprocessing
(mean/sphere) and the RNG are preserved.
"""
# Only A is nulled; _initialize_parameters redraws it (and unconditionally
# rebuilds comp_list and W) while leaving the still-finite mixture params.
# Preserve the outlier-rejection state across a restart: Fortran's
# startover redraws A but never reverts already-applied rejections
# (amica17.f90:1121-1148), so snapshot good_idx/num_good_samples around
# the reinit (which would otherwise reset them to the full set) and
# restore them; numrej is an instance attribute and already survives.
saved_good_idx = self.good_idx
saved_num_good = self.num_good_samples
self.A = None
self._initialize_parameters()
if self.do_reject:
self.good_idx = saved_good_idx
self.num_good_samples = saved_num_good
self.lrate = self.lrate0
# rholrate_cap is the (maxdecs-ratcheted) rho-rate ceiling; restore it,
# and the working rate with it, to the pristine rholrate0 so a re-fit
# starts fresh (issue #193).
self.rholrate_cap = self.rholrate0
self.rholrate = self.rholrate0
self.ll = []
self.nd = []
def _update_unmixing_matrices(self):
"""Update unmixing matrices from mixing matrix."""
self.W = get_unmixing_matrices(self.A, self.comp_list)
def _compute_log_pdf(
self, y: np.ndarray, rho: float
) -> Tuple[np.ndarray, np.ndarray]:
"""
Compute log PDF value and its derivative for given activation.
Parameters
----------
y : ndarray
Activation values
rho : float
Shape parameter
Returns
-------
log_pdf : ndarray
Log PDF values
dpdf : ndarray
PDF derivatives (not in log space)
"""
if rho == 1.0:
# Laplace distribution
log_pdf = -np.abs(y) - LOG2
pdf = np.exp(log_pdf)
dpdf = -np.sign(y) * pdf
elif rho == 2.0:
# Gaussian distribution, with the reference's single-precision
# normalizer log(dble(1.772453851)) (amica15.f90:1313, issue #344)
log_pdf = -y * y - LOG_SQRT_PI
pdf = np.exp(log_pdf)
dpdf = -2 * y * pdf
else:
# Generalized Gaussian distribution
log_pdf = -np.power(np.abs(y), rho) - LOG2 - gammaln(1.0 + 1.0 / rho)
pdf = np.exp(log_pdf)
dpdf = -rho * np.power(np.abs(y), rho - 1) * np.sign(y) * pdf
return log_pdf, dpdf
def _compute_score(self, y: np.ndarray, rho: float) -> np.ndarray:
"""Generalized-Gaussian score ``fp = d|y|^rho/dy`` (Fortran ``fp``,
amica17.f90:1455-1467): ``sign(y)`` for Laplace, ``2y`` for Gaussian,
``rho*sign(y)*|y|^(rho-1)`` otherwise. Used by the Newton curvature
statistics; distinct from the density derivative ``dpdf``."""
if rho == 1.0:
return np.sign(y)
if rho == 2.0:
return 2.0 * y
return rho * np.sign(y) * np.power(np.abs(y), rho - 1.0)
def _tune_block_size(self) -> None:
"""Set ``self.block_size`` to the fastest timed candidate (issue #232).
The probe is one ``_get_updates_and_likelihood`` pass -- the same
E-step-plus-sufficient-statistics work every EM iteration does, so it
times what the fit will actually spend its time on. This replaces the
old ``determine_block_size`` helper, which timed a bare ``X.T @ X``
(not the shape of any work AMICA does) over Fortran's 128-1024 range
(far below where any pamica backend peaks) and had no fallback at all
when an allocation failed. See :mod:`pamica.blocktune`.
The pass reads model state and consumes no RNG; the one thing it does
write, ``_last_ll_samples`` under ``do_reject``, is restored here, so
the fit that follows is bit-identical to one started directly at the
chosen block size.
"""
assert self.data is not None and self.data_dim is not None
saved_block_size = self.block_size
saved_ll_samples = self._last_ll_samples
n_samples = (
int(self.good_idx.size)
if self.do_reject and self.good_idx is not None
else int(self.data.shape[1])
)
def probe(size: int) -> float:
self.block_size = size
try:
start = time.perf_counter()
self._get_updates_and_likelihood()
return time.perf_counter() - start
finally:
# Never leave the model holding a candidate -- least of all one
# that just failed to allocate -- or a probe's throwaway state.
self.block_size = saved_block_size
self._last_ll_samples = saved_ll_samples
self.block_size = blocktune.search(
probe=probe,
fallback=saved_block_size,
blk_min=self.blk_min,
blk_max=self.blk_max,
blk_step=self.blk_step,
n_samples=n_samples,
n_channels=self.data_dim,
n_mix=self.num_mix,
n_models=self.num_models,
itemsize=self.data.dtype.itemsize,
available_bytes=blocktune.host_memory_bytes(),
log=self.logger,
)
def _get_updates_and_likelihood(self) -> Dict:
"""
Compute parameter updates and data likelihood.
Returns
-------
updates : dict
Dictionary containing parameter updates and likelihood
"""
assert self.data_dim is not None
# Initialize update accumulators
updates = {
"dgm": np.zeros(self.num_models),
"dalpha_n": np.zeros((self.num_mix, self.num_comps)),
"dmu_n": np.zeros((self.num_mix, self.num_comps)),
"dmu_d": np.zeros((self.num_mix, self.num_comps)),
"dbeta_n": np.zeros((self.num_mix, self.num_comps)),
"dbeta_d": np.zeros((self.num_mix, self.num_comps)),
"drho_n": np.zeros((self.num_mix, self.num_comps)),
"dWtmp": np.zeros((self.data_dim, self.data_dim, self.num_models)),
"dc_numer": np.zeros((self.data_dim, self.num_models)),
"ll": 0.0,
}
if self.do_newton:
updates.update(
{
"dsigma2": np.zeros((self.data_dim, self.num_models)),
"dlambda": np.zeros((self.data_dim, self.num_models)),
"dkappa": np.zeros((self.data_dim, self.num_models)),
}
)
# Restrict the E-step to the currently-good samples under do_reject
# (mirrors AMICATorchNG's ``X_use = X_t[:, good_idx]``); the default path
# uses the full array with no copy, so it stays bit-identical.
good_idx = self.good_idx if self.do_reject else None
assert not self.do_reject or good_idx is not None
data_use = self.data[:, good_idx] if good_idx is not None else self.data
assert self._llt_logv is not None and self._llt_ll is not None
# Process data in blocks
for start in range(0, data_use.shape[1], self.block_size):
end = min(start + self.block_size, data_use.shape[1])
X = data_use[:, start:end]
# Get block updates
block_updates = self._get_block_updates(X)
# Accumulate updates. block_updates also carries the per-sample
# "logV"/"ll_samples", which are stashed below rather than summed
# here -- this loop only walks the accumulator keys, so they are
# skipped.
for key in updates:
updates[key] += block_updates[key]
# LLt stash (issue #157), Fortran's per-block modloglik/loglik write
# (amica15.f90:1406-1411). Under do_reject the block came from
# data[:, good_idx], so block row r is dataset sample good_idx[r].
rows = good_idx[start:end] if good_idx is not None else slice(start, end)
self._llt_logv[rows] = block_updates["logV"]
self._llt_ll[rows] = block_updates["ll_samples"]
# Per-sample log-likelihood of the good set, in good_idx order, so
# _reject_outliers' keep-mask maps back onto good_idx. Read straight out
# of the stash just written (the same values the blocks produced) rather
# than concatenated a second time.
if good_idx is not None:
self._last_ll_samples = self._llt_ll[good_idx]
# Normalize the accumulated total LL by (good-sample count x working
# dimensionality), matching Fortran's LL(iter) = LLtmp2 / dble(numgoodsum*nw)
# (amica15.f90:1770) at the point the full-data sum becomes available.
# numgoodsum is self.num_good_samples: Fortran initializes numgoodsum to
# all_blks (all samples) and only shrinks it under do_reject (amica15.f90:252,
# 2252), and self.num_good_samples mirrors that exactly (set to num_samples
# in _initialize_parameters, only shrunk by _reject_outliers), so this one
# expression is correct whether or not do_reject is set -- no branch needed.
# nw is the per-model working dimensionality (Fortran's numeigs, amica15.f90:563);
# self.data_dim is its NumPy analogue -- comp_list is (data_dim, num_models),
# matching Fortran's comp_list(nw, num_models) (amica15.f90:617).
assert self.num_good_samples is not None and self.data_dim is not None
updates["ll"] /= self.num_good_samples * self.data_dim
return updates
def _forward_block(self, X: np.ndarray):
"""Shared E-step forward pass for one data block (Fortran
amica17.f90:1280-1372): activations ``b``, the raw (unnormalized)
mixture log-probabilities ``z``, their per-source max ``z0max``, the
per-model log-likelihood ``logV`` (with the log|det W| + sldet
Jacobian), and the per-sample total log-likelihood (``Vmax``/
``ll_samples``, log-sum-exp over models).
Called only from ``_get_block_updates``, whose per-block ``logV``/
``ll_samples`` serve both the M-step (after normalization into
responsibilities) and the ``LLt`` stash (:meth:`_llt_arrays`). Since
issue #157 there is no second caller: the write path reads the stash
instead of re-running this pass, so it cannot silently drift from the
training E-step -- this codebase has had one-sided forward-pass bugs
before (issue #24's fp-vs-dpdf sign bug).
Returns
-------
b : ndarray of shape (batch, data_dim, num_models)
z : ndarray of shape (batch, data_dim, num_mix, num_models)
Raw (unnormalized) mixture log-probabilities.
z0max : ndarray of shape (batch, data_dim, 1, num_models)
logV : ndarray of shape (batch, num_models)
Vmax : ndarray of shape (batch, 1)
ll_samples : ndarray of shape (batch,)
"""
assert (
self.data_dim is not None
and self.c is not None
and self.W is not None
and self.comp_list is not None
and self.beta is not None
and self.mu is not None
and self.rho is not None
and self.alpha is not None
and self.gm is not None
)
batch_size = X.shape[1]
# Compute activations for each model. c is the per-model data-space
# center: b = W(x - c). Fortran subtracts wc in the E-step
# (amica17.f90:1280-1292), where wc = W@c is precomputed in
# get_unmixing_matrices (amica17.f90:2178). For n_models=1, c == 0, so
# this is bit-identical to X.T @ W.
b = np.zeros((batch_size, self.data_dim, self.num_models))
for h in range(self.num_models):
b[:, :, h] = np.dot((X - self.c[:, h][:, None]).T, self.W[:, :, h])
# Compute mixture probabilities and responsibilities
z = np.zeros((batch_size, self.data_dim, self.num_mix, self.num_models))
for h in range(self.num_models):
for i in range(self.data_dim):
k = self.comp_list[i, h]
for j in range(self.num_mix):
y = self.beta[j, k] * (b[:, i, h] - self.mu[j, k])
# Compute log PDF and its derivative
log_pdf, dpdf = self._compute_log_pdf(y, self.rho[j, k])
# Compute log probability directly in log space
z[:, i, j, h] = (
np.log(self.alpha[j, k]) + np.log(self.beta[j, k]) + log_pdf
)
# Per-source log-density = logsumexp over mixtures of the pre-norm logits
# z0, then the per-model log-likelihood logV adds the log|det W| + sldet
# Jacobian (Fortran amica17.f90:1341-1350). This is the correct
# pre-normalization log-likelihood; the earlier post-normalization
# np.sum(np.exp(z_normalized)) did not recover a real log-density and
# omitted the Jacobian (so LL was positive and ~4.9 off per channel).
z0max = np.max(z, axis=2, keepdims=True)
ll_src = z0max[:, :, 0, :] + np.log(
np.sum(np.exp(z - z0max), axis=2)
) # (batch, data_dim, num_models)
logV = np.zeros((batch_size, self.num_models))
for h in range(self.num_models):
# A near-singular W (a transient the natural gradient can pass
# through) makes slogdet emit divide/overflow/invalid FP warnings.
# Suppress the numpy console noise, but DON'T rely on that as the
# guard: the explicit isfinite check below is the real diagnostic --
# it fires for a -inf logdet (singular W) AND a NaN logdet (genuinely
# broken W), so silencing `invalid` does not hide a NaN W. The
# fit-loop LL check (_check_convergence) also stops on a -inf LL.
with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
_, logdet_W = np.linalg.slogdet(self.W[:, :, h])
if not np.isfinite(logdet_W):
self.logger.warning(
"Non-finite logdet(W) for model %d at iter %d (logdet=%s); "
"W is singular or corrupt.",
h,
getattr(self, "iter", -1),
logdet_W,
)
logV[:, h] = (
np.log(self.gm[h])
+ logdet_W
+ self.sldet
+ np.sum(ll_src[:, :, h], axis=1)
)
# Block log-likelihood = sum_t logsumexp_h logV (Fortran :1372).
Vmax = np.max(logV, axis=1, keepdims=True)
ll_samples = Vmax[:, 0] + np.log(np.sum(np.exp(logV - Vmax), axis=1))
return b, z, z0max, logV, Vmax, ll_samples
def _get_block_updates(self, X: np.ndarray) -> Dict:
"""
Compute parameter updates for a data block.
Parameters
----------
X : ndarray
Data block to process
Returns
-------
updates : dict
Parameter updates for this block
"""
assert (
self.data_dim is not None
and self.c is not None
and self.W is not None
and self.comp_list is not None
and self.beta is not None
and self.mu is not None
and self.rho is not None
and self.alpha is not None
and self.gm is not None
)
batch_size = X.shape[1]
tiny = np.finfo(np.float64).tiny
updates = {
"dgm": np.zeros(self.num_models),
"dalpha_n": np.zeros((self.num_mix, self.num_comps)),
"dmu_n": np.zeros((self.num_mix, self.num_comps)),
"dmu_d": np.zeros((self.num_mix, self.num_comps)),
"dbeta_n": np.zeros((self.num_mix, self.num_comps)),
"dbeta_d": np.zeros((self.num_mix, self.num_comps)),
"drho_n": np.zeros((self.num_mix, self.num_comps)),
"dWtmp": np.zeros((self.data_dim, self.data_dim, self.num_models)),
"dc_numer": np.zeros((self.data_dim, self.num_models)),
"ll": 0.0,
}
if self.do_newton:
updates.update(
{
"dsigma2": np.zeros((self.data_dim, self.num_models)),
"dlambda": np.zeros((self.data_dim, self.num_models)),
"dkappa": np.zeros((self.data_dim, self.num_models)),
}
)
b, z, z0max, logV, Vmax, ll_samples = self._forward_block(X)
updates["ll"] = np.sum(ll_samples)
# Per-sample E-step outputs, for outlier rejection (issue #123) and the
# LLt stash (issue #157). Not accumulators: _get_updates_and_likelihood
# scatters them into the full-dataset buffers instead of summing them
# (its accumulate loop walks only the keys initialized above).
updates["logV"] = logV
updates["ll_samples"] = ll_samples
# Model responsibilities v = softmax(logV); mixture responsibilities z.
v = np.exp(logV - Vmax)
v /= np.sum(v, axis=1, keepdims=True)
z = np.exp(z - z0max)
z /= np.sum(z, axis=2, keepdims=True)
for h in range(self.num_models):
# Model weights
updates["dgm"][h] = np.sum(v[:, h])
for i in range(self.data_dim):
k = self.comp_list[i, h]
# Newton second moment sigma2 = E_v[b_i^2] is a per-source
# quantity (no mixture index), accumulated once per (i, h)
# (Fortran amica17.f90:1419). It must NOT sit inside the
# mixture loop below, or it is inflated num_mix-fold.
if self.do_newton:
updates["dsigma2"][i, h] += np.sum(v[:, h] * b[:, i, h] ** 2)
# Mixture exact-EM sufficient statistics (Fortran
# amica17.f90:1524-1578). These use the score
# fp = rho*sign(y)*|y|^(rho-1) (Fortran fp), NOT the density
# derivative dpdf, and produce numerator/denominator pairs for a
# fixed-point (not first-order gradient) update. Assumes rho <= 2
# (the maxrho default); the rho > 2 denominator branches are
# unreachable and not implemented.
for j in range(self.num_mix):
y = self.beta[j, k] * (b[:, i, h] - self.mu[j, k])
fp = self._compute_score(y, self.rho[j, k])
u = v[:, h] * z[:, i, j, h] # model x mixture responsibility
ufp = u * fp
updates["dalpha_n"][j, k] += np.sum(u)
updates["dmu_n"][j, k] += np.sum(ufp) # sum(ufp) (:1532)
updates["dmu_d"][j, k] += self.beta[j, k] * np.sum(
ufp / y
) # sbeta*sum(ufp/y) (:1537)
updates["dbeta_n"][j, k] += np.sum(u) # sum(u) (:1550)
updates["dbeta_d"][j, k] += np.sum(ufp * y) # sum(ufp*y) (:1556)
# drho_numer = rho*sum(u*|y|^rho*ln|y|) (:1560-1578). Leading
# rho from ln(|y|^rho)=rho*ln|y| (issue #24 Bug 1); no
# per-component rho!=1&rho!=2 mask (Bug 2), only the
# per-sample underflow guard (:1570).
ay = np.abs(y)
ayrho = np.power(ay, self.rho[j, k])
logab = self.rho[j, k] * np.log(np.maximum(ay, tiny))
# Fortran zeros the term when |y|^rho < epsdble
# (amica15.f90:1558 / amica15_header.f90:73), not at denormal
# underflow; use its value to match, not np.finfo.tiny.
logab = np.where(ayrho < EPSDBLE, 0.0, logab)
updates["drho_n"][j, k] += np.sum(u * ayrho * logab)
if self.do_newton:
# Newton curvature terms use the score fp (Fortran
# :1500-1512): kappa carries sbeta^2, lambda folds in the
# mu^2 curvature term so lambda=dlambda/dgm matches Fortran.
dkap = np.sum(u * fp**2) * self.beta[j, k] ** 2
updates["dkappa"][i, h] += dkap
updates["dlambda"][i, h] += (
np.sum(u * (fp * y - 1) ** 2) + dkap * self.mu[j, k] ** 2
)
# Natural-gradient accumulator: g_i = sum_j sbeta*u*fp, then the
# source-space sum dWtmp = g^T b (Fortran :1493/:1592). Uses the
# score fp, not dpdf.
g = np.zeros((batch_size, self.data_dim))
for i in range(self.data_dim):
k = self.comp_list[i, h]
for j in range(self.num_mix):
y = self.beta[j, k] * (b[:, i, h] - self.mu[j, k])
fp = self._compute_score(y, self.rho[j, k])
g[:, i] += self.beta[j, k] * (v[:, h] * z[:, i, j, h]) * fp
updates["dWtmp"][:, :, h] += np.dot(g.T, b[:, :, h])
# Data-space bias numerator dc_numer[i,h] = sum_t v_h(t)*x(i,t)
# (Fortran :1423-1429); denominator is dgm[h] = sum_t v_h(t). Replaces
# the old gradient-style bias sum(g), which was accumulated but never
# applied (c was frozen at 0); the Fortran update is the data-space
# responsibility-weighted mean (issue #27).
updates["dc_numer"][:, h] += np.dot(X, v[:, h])
return updates
def _llt_arrays(self) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
"""The stashed per-model/per-sample log-likelihood, in ``LLt`` layout.
Returns what the last E-step computed (issue #157), reshaped for
:func:`pamica.numpy_impl.load.write_amicaout`; it performs no forward
pass of its own, so a ``writestep`` checkpoint costs no more than it did
before ``LLt`` output existed. This is Fortran's design: ``modloglik``/
``loglik`` are allocated once (amica15.f90:2617-2620), filled by every
E-step (amica15.f90:1406-1411) and dumped verbatim by ``write_output``
(amica15.f90:2338-2343).
Ordering, adopted from the reference on purpose. Fortran's main loop is
``get_updates_and_likelihood`` -> ``update_params`` -> ``write_output``
(amica15.f90:996, 1122, 1124-1127), so the ``LLt`` written at any
checkpoint is the E-step of the parameters as they stood BEFORE that
iteration's M-step: one M-step older than the ``W``/``A`` written beside
it. The same holds for the final write after a fit that ran to
``max_iter``; after a convergence stop the loop exits before the
update (amica15.f90:1111, issue #339), so the final ``LLt`` is the
E-step of the written parameters themselves. pamica's loop has the
same steps in the same order, so stashing reproduces that exactly, and
the on-disk ``LLt`` is comparable with the binary's. The invariant this pins down, on both sides, is
``Lt.sum() / (num_good_samples * data_dim) == self.ll[-1]``.
That equality holds bit for bit except on a write whose iteration also
fired a rejection: ``self.ll[-1]`` was normalized over the good set as
it stood *before* ``_reject_outliers`` ran, and the dropped samples'
stash entries have since been zeroed, so the two sides no longer count
the same samples and a small residual remains. Reference-faithful
rather than a defect -- Fortran computes ``LL(iter) =
LLtmp2/dble(numgoodsum*nw)`` (amica15.f90:1770) before ``reject_data``
(amica15.f90:1138) shrinks ``numgoodsum`` (amica15.f90:2252) -- and any
later iteration re-normalizes over the shrunk set and restores it.
Under ``do_reject`` a rejected sample's entries are zero -- they are
never scored again after ``_reject_outliers`` zeroes them (mirroring
amica15.f90:2231-2234), and ``load_rej`` reads exactly that zero as the
rejection sentinel (``sum(modloglik(:,i)) == 0.0``, amica15.f90:907).
Returns
-------
Lht : ndarray of shape (num_models, n_samples), or None
Per-model log-likelihood, Fortran's ``modloglik``. Zero for
rejected samples under ``do_reject``. ``None`` before any E-step
has run, so the caller omits the file rather than writing an
all-zero array that ``load_rej`` would read as "everything
rejected".
Lt : ndarray of shape (n_samples,), or None
Total log-likelihood, Fortran's ``loglik``; ``None`` on the same
condition as ``Lht``.
"""
if self._llt_logv is None or self._llt_ll is None or not self.ll:
return None, None
return self._llt_logv.T, self._llt_ll
def _update_direction(self, updates: Dict) -> _UpdateStep:
"""The natural-gradient or Newton step for ``A`` from this iteration's
sufficient statistics, and its norm, without changing ``A`` or any
density parameter.
The reference computes both in ``accum_updates_and_likelihood``
(amica15.f90:1666-1761), with ``LL(iter)``, before the likelihood-
decrease response and the stopping checks read ``ndtmpsum`` and before
``update_params`` applies the step. It finalizes the Newton curvature
(``sigma2``/``lambda_``/``kappa``) as it goes but records nothing:
:meth:`_optimize` appends the iteration's likelihood and the returned
norm to ``self.ll``/``self.nd`` once the likelihood is known to be
finite, as the PyTorch and MLX loops append ``ll_history``, and hands
the step to :meth:`_update_parameters` once the checks have run.
"""
assert (
self.data_dim is not None
and self.gm is not None
and self.comp_list is not None
and self.A is not None
)
used = (
self.comp_used
if self.comp_used is not None
else np.ones(self.num_comps, dtype=bool)
)
newton_active = schedule.newton_active(
self.do_newton, self.iter, self.newt_start
)
if newton_active:
# Finalize Newton curvature statistics (Fortran amica17.f90:1762-1776).
# The dsigma2/dkappa/dlambda accumulators already carry the sbeta^2
# and baralpha-weighted mu^2 factors, so finalization is a plain
# division by the model mass dgm = sum_t v_h.
#
# dgm is (num_models,) and the accumulators are (data_dim,
# num_models), so the model mass broadcasts along the LAST axis
# (issue #267). The old ``[:, None]`` made it (num_models, 1), which
# only happens to broadcast when num_models == 1: every multi-model
# Newton fit raised "operands could not be broadcast together with
# shapes (data_dim, num_models) (num_models, 1)". Same as the torch
# backend's ``dgm.unsqueeze(0)`` (``AMICATorchNG._finalize_newton_stats``).
dgm = updates["dgm"][None, :]
self.sigma2 = updates["dsigma2"] / dgm
self.lambda_ = updates["dlambda"] / dgm
self.kappa = updates["dkappa"] / dgm
# Per-model direction: Newton H if the model is positive definite,
# otherwise natural gradient. Matching Fortran (amica17.f90:1814-1837),
# if any off-diagonal pair fails sk1*sk2 > 1 the whole model falls
# back to the natural gradient and the ramp targets lrate0, not newtrate.
directions = []
no_newt = False
for h in range(self.num_models):
dA = -updates["dWtmp"][:, :, h] / updates["dgm"][h]
dA[np.diag_indices_from(dA)] += 1
if newton_active:
assert (
self.lambda_ is not None
and self.sigma2 is not None
and self.kappa is not None
)
H = np.zeros_like(dA)
posdef = True
for i in range(self.data_dim):
for j in range(self.data_dim):
if i == j:
H[i, i] = dA[i, i] / self.lambda_[i, h]
else:
sk1 = self.sigma2[i, h] * self.kappa[j, h]
sk2 = self.sigma2[j, h] * self.kappa[i, h]
if sk1 * sk2 > 1.0:
H[i, j] = (sk1 * dA[i, j] - dA[j, i]) / (
sk1 * sk2 - 1.0
)
else:
posdef = False
if posdef:
directions.append(H)
else:
no_newt = True
directions.append(dA)
else:
directions.append(dA)
# Weight-gradient norm (Fortran ndtmpsum, amica15.f90:1749-1761). Computed
# HERE, before the A step and before rescaling, because Fortran builds dAk
# inside accum_updates_and_likelihood (:1731-1743) strictly before
# update_params applies it (:1789). Using the post-update, post-rescale A
# would measure a different quantity. Computed every iteration, including
# a frozen one, because Fortran computes it in the accumulation pass that
# runs unconditionally -- the grad-norm stop needs the true gradient
# magnitude, not just the magnitude on iterations where A moves.
#
# dAk is the gm-weighted average of the per-model directions mapped
# through A (Fortran dAk/zeta, amica15.f90:1749-1761): each contributing
# model adds gm[h]/zeta of its step to the shared component's row, where
# zeta = sum of gm over the models that reference the component. Weighted
# by the PRE-update model weights (gm is not reassigned until
# _update_parameters), because Fortran builds dAk before update_params
# reassigns gm (issue #219).
#
# The weights are normalized per model (gm[h]/zeta) rather than summed and
# then divided (Fortran's literal sum-then-divide, and AMICATorchNG's):
# mathematically the same average, but a component with a single
# contributor then has weight exactly gm[h]/gm[h] == 1.0, so the step is
# bit-identical to the pre-#242 per-model update instead of drifting by
# the ULP that a multiply-then-divide round trip can introduce. Every
# component has exactly one contributor unless share_comps merged one,
# so this keeps the default multi-model trajectory byte-for-byte.
#
# ndtmpsum is then the RMS of the used rows of dAk:
# ||dAk[used, :]|| / sqrt(nw * n_used), with NO lrate factor. Fortran
# measures the gradient direction before the step, not the applied update
# lrate*dAk, and does not divide by lrate either (amica15.f90:1760-1761) --
# there is no missing factor here.
zeta = np.zeros(self.num_comps)
for h in range(self.num_models):
zeta[self.comp_list[:, h]] += self.gm[h]
dAk = np.zeros_like(self.A)
for h in range(self.num_models):
# comp_list[:, h] holds distinct indices within a model
# (identify_shared_components never merges two components that
# appear in the same model), so buffered `+=` on fancy indices cannot
# drop a contribution here. Row i of the model's step is source i's
# component, comp_list[i, h] (issue #334).
idx = self.comp_list[:, h]
weight = self.gm[h] / np.maximum(zeta[idx], np.finfo(np.float64).tiny)
dAk[idx, :] += weight[:, None] * np.dot(directions[h].T, self.A[idx, :])
nd_value = float(
np.sqrt(np.sum(dAk[used, :] ** 2) / (self.data_dim * int(used.sum())))
)
return _UpdateStep(dAk, newton_active, no_newt, nd_value)
def _update_parameters(self, updates: Dict, step: Optional[_UpdateStep] = None):
"""
Update model parameters using computed updates.
Parameters
----------
updates : dict
Dictionary containing parameter updates
step : _UpdateStep, optional
The mixing-matrix step :meth:`_update_direction` computed from the
same ``updates``: :meth:`_optimize` passes the one its stopping
checks already read, so nothing is computed twice. A direct call
may omit it, and it is then computed here first; neither path
records anything in ``self.ll``/``self.nd``, which only
:meth:`_optimize` appends to.
"""
assert (
self.data_dim is not None
and self.num_samples is not None
and self.c is not None
and self.mu is not None
and self.alpha is not None
and self.beta is not None
and self.rho is not None
and self.comp_list is not None
and self.A is not None
)
if step is None:
step = self._update_direction(updates)
# The step was built with the pre-update gm (_update_direction), so gm
# can be reassigned now (Fortran reassigns it in update_params,
# amica15.f90:1788, after accum_updates_and_likelihood; issue #219).
# Update model weights, normalizing by the number of samples the E-step
# actually summed over: the good set under do_reject, else all samples.
if self.do_reject:
assert self.num_good_samples is not None
self.gm = updates["dgm"] / self.num_good_samples
else:
self.gm = updates["dgm"] / self.num_samples
# Per-model data-space bias c (Fortran's `update_c` flag, amica17.f90:1423-
# 1429 numerator / :1899-1901 division): c[i,h] = sum_t v_h*x / sum_t v_h,
# the responsibility-weighted data mean for model h (the E-step centers
# each model at its own mean, b = W(x - c)). Skipped for a single model to
# keep the issue #24 parity bit-exact: with v==1 the update would add a
# ~1e-13 float-sum residual of the (mean-removed) data. dgm[h] = sum_t v_h
# is the denominator (Fortran `dc_denom`). A dead model (dgm[h]==0) gives
# 0/0; keep its prior c so a NaN cannot poison the next iteration's
# cross-model softmax for every model (mirrors the mu/beta/rho guards).
if self.num_models > 1:
dgm = updates["dgm"]
live = dgm > 0.0
new_c = (
updates["dc_numer"]
/ np.maximum(dgm, np.finfo(np.float64).tiny)[None, :]
)
self.c = np.where(live[None, :], new_c, self.c)
if not np.all(live):
self.logger.warning(
"Zero-responsibility model(s) at iter %d; kept prior bias c "
"(dead-model guard).",
self.iter,
)
# A component merged away by share_comps is no longer referenced by
# comp_list, so no sufficient statistic accumulates into its column and
# its mixture divisions would be 0/0 = NaN. Update only the used columns
# and leave the rest frozen at their last finite value, as AMICATorchNG
# does (Fortran instead carries the NaN harmlessly behind its comp_used
# mask; keeping them finite means a fit cannot report success while
# holding NaN parameters, issue #240).
#
# The dead columns are INDEXED OUT rather than computed and masked: a
# 0/0 that is never evaluated raises no RuntimeWarning, so nothing has to
# be suppressed with np.errstate and a genuine 0/0 in a LIVE column (a
# component whose responsibility mass collapses to exactly zero) still
# warns, as it does with sharing off. ``cols`` is a full slice whenever
# every column is live -- always so with the default comp_list -- which
# keeps the ordinary path bit-identical and copy-free. (comp_used is None
# until fit() sets it up, and the M-step is exercised directly in tests
# before that happens, so fall back to "all used".)
used = (
self.comp_used
if self.comp_used is not None
else np.ones(self.num_comps, dtype=bool)
)
cols = slice(None) if used.all() else np.flatnonzero(used)
# Update mixture weights
dalpha_n = updates["dalpha_n"][:, cols]
self.alpha[:, cols] = dalpha_n / np.sum(dalpha_n, axis=0)
# Fortran has no alpha canary, but a collapsed live component gives the
# same 0/0 the mu/beta canary below reports; surface it here too so the
# origin is not lost to a later unattributable nan-LL stop.
if not np.all(np.isfinite(self.alpha)):
self.logger.warning(
"Non-finite alpha at iter %d (component responsibility mass "
"collapsed).",
self.iter,
)
# The A branch, where the reference has it: after gm/alpha/c and before
# mu/sbeta/rho (amica15.f90:1803-1816).
#
# A is stored as Fortran's A^T, one component per row (issue #334; true
# unmixing = W^T = inv(block)^T), so the Fortran step on a model's block
# A_fort(:, comp_list(:,h)) @ dir becomes dir^T @ A[comp_list[:, h], :]
# (LEFT-multiply by the TRANSPOSED direction), which _update_direction
# built into step.dAk. Right-multiply by the untransposed dir is
# invisible at the fixed point but sends the fit downhill -- issue #24
# root cause.
#
# ONE application of the averaged dAk (Fortran's single DAXPY,
# amica15.f90:1807/1814), not the per-model loop this used to run: a
# component shared by two models took one step per contributing model,
# the second against an already-stepped A, which is a different
# operation from Fortran's single weighted average (issue #242). A
# merged-away component has no contributor, so its dAk row stays
# exactly zero and it holds the value it was merged away with. On an
# iteration the reference holds A (pamica.schedule.share_freeze),
# everything inside the branch is skipped together: the
# Newton-fallback message, the lrate ramp, the reset of the working rho
# rate to its ceiling and the step itself. The step was built from the
# parameters the E-step saw, so taking it before the mixture updates
# changes nothing they read.
if not self._a_frozen():
if step.newton_active and step.no_newt:
# Fortran prints this whenever a model is not positive definite
# (amica17.f90:1911-1913); surface it rather than falling back
# silently.
self.logger.info(
"Hessian not positive definite at iter %d; using natural gradient.",
self.iter,
)
# The ramp starts from this iteration's lrate, which a likelihood
# decrease has already halved (_optimize runs the response before
# the update, as the reference does, issue #339).
if step.newton_active and not step.no_newt:
self.lrate = min(
self.newtrate, self.lrate + min(1.0 / self.newt_ramp, self.lrate)
)
else:
self.lrate = min(
self.lrate0, self.lrate + min(1.0 / self.newt_ramp, self.lrate)
)
# rho moves at the ceiling (``rholrate = rholrate0``, :1806/:1813),
# so a decrease's scaling of the working rate reaches it only on a
# held iteration.
self.rholrate = self.rholrate_cap
self.A = self.A - self.lrate * step.dAk
# Exact-EM mixture location/scale (Fortran :1978/:1993). These are
# fixed-point updates -- mu += dmu_n/dmu_d, beta *= sqrt(dbeta_n/dbeta_d)
# -- NOT first-order gradient steps, so they carry no lrate.
dmu = updates["dmu_n"][:, cols] / updates["dmu_d"][:, cols]
self.mu[:, cols] = self.mu[:, cols] + dmu
self.beta[:, cols] = self.beta[:, cols] * np.sqrt(
updates["dbeta_n"][:, cols] / updates["dbeta_d"][:, cols]
)
self.beta = np.clip(self.beta, self.invsigmin, self.invsigmax)
# Fortran keeps a live "NaN in sbeta!" canary here (amica17.f90:1996-2000);
# the exact-EM mu/beta divisions are unguarded (matching Fortran), so
# surface a non-finite value immediately instead of letting it propagate
# to a later, unattributable nan-LL stop.
if not np.all(np.isfinite(self.mu)) or not np.all(np.isfinite(self.beta)):
self.logger.warning(
"Non-finite mu/beta at iter %d (mixture component mass likely "
"collapsed).",
self.iter,
)
# GG shape update with the 1/psi(1+1/rho) digamma factor (Fortran
# :2013-2014); the divisor is the per-component responsibility mass
# dalpha_n (floored so a near-empty component cannot poison rho). A NaN
# is reset to rho0 -- but logged first, so the reset does not silently
# erase the failure origin.
# Restricted to the live columns for the same reason as mu/beta/alpha
# (and as AMICATorchNG, which masks rho with ``used`` too): a merged-away
# column has no responsibility mass, so its rho would ratchet up on the
# floored denominator every iteration instead of staying at the value it
# was merged away with.
if not np.all(self.rho == 1.0) and not np.all(self.rho == 2.0):
rho_cols = self.rho[:, cols]
drho = updates["drho_n"][:, cols] / np.maximum(
updates["dalpha_n"][:, cols], 1e-8
)
psi = digamma(1.0 + 1.0 / rho_cols)
new_rho = rho_cols + self.rholrate * (1.0 - (rho_cols / psi) * drho)
nan_mask = np.isnan(new_rho)
if nan_mask.any():
self.logger.warning(
"NaN in rho update at iter %d for %d component(s); resetting "
"to rho0=%g.",
self.iter,
int(nan_mask.sum()),
self.rho0,
)
new_rho = np.where(nan_mask, self.rho0, new_rho)
self.rho[:, cols] = np.clip(new_rho, self.minrho, self.maxrho)
# (c was updated above, before the mixture/A updates, from dc_numer/dgm.)
# Rescale parameters if requested. The reference rescales every
# iteration (it parses ``scalestep`` but never reads it,
# amica15.f90:1843/3686); pamica keeps ``scalestep`` as an extension
# counted from 1 like the reference's other cadences (iterations s, 2s,
# ...; ``schedule.every``), so the default 1 is the reference.
if self.doscaling and schedule.every(self.iter, self.scalestep):
self._rescale_components()
# Update unmixing matrices
self._update_unmixing_matrices()
def _component_sensor_maps(self) -> np.ndarray:
"""Every component's mixing vector in input-channel (sensor) space.
``pinv(sphere) @ A.T``, shape ``(data_dim_in, num_comps)``: column
``comp_list[i, h]`` is column ``i`` of ``get_sensor_mixing_matrix(h)``
(issue #334). These are the vectors the share metric compares.
"""
assert self.A is not None
return self._pinv_sphere() @ self.A.T
def _identify_shared_comps(self) -> None:
"""Merge near-collinear components across models (Fortran
``identify_shared_comps``, amica15.f90:1916).
Runs the shared kernel
:func:`pamica.numpy_impl.utils.identify_shared_components` on the
sensor-space component maps (:meth:`_component_sensor_maps`, issue
#258), the metric ``AMICATorchNG._identify_shared_comps`` computes, so
both backends make the same merge decision from the same fitted state.
A merge folds one component id into another in ``comp_list`` and
``comp_used``; the caller rebuilds ``W``, as the reference's
``get_unmixing_matrices`` follows the scan (amica15.f90:1858,1863).
"""
assert self.comp_list is not None
unique_before = int(np.unique(self.comp_list).size)
self.comp_list, self.comp_used = identify_shared_components(
self._component_sensor_maps(), self.comp_list, self.comp_thresh
)
unique_after = int(np.unique(self.comp_list).size)
# identify_shared_components is stateless, so the merge log (matching
# AMICATorchNG's) is emitted here from the before/after unique-component
# count instead.
if unique_after < unique_before:
self.logger.info(
"Component sharing (iter %d): %d merge(s), %d unique components.",
self.iter,
unique_before - unique_after,
unique_after,
)
def _rescale_components(self) -> None:
"""Rescale every component to a unit-norm mixing vector (Fortran
``doscaling``, amica15.f90:1843-1851), an exact change of scale.
Component ``k`` is row ``k`` of ``A`` (issue #334, ADR 0007), the
reference's column ``A(:,k)``. Dividing that row by its norm scales the
source up by the norm; ``mu[:, k] *= norm`` and ``beta[:, k] /= norm``
rescale its density to match, so the log-likelihood is unchanged. A
zero-norm row is left untouched, as in the reference (``Anrmk > 0``).
Each component is rescaled once, including a shared one, and a
merged-away row is left untouched. Same rule, norms and guard as
``AMICATorchNG._rescale_components``.
"""
assert self.A is not None and self.mu is not None and self.beta is not None
assert self.comp_list is not None
scale = np.ones(self.A.shape[0])
for h in range(self.num_models):
# The per-model block, as the A-update and W gather it; a shared row
# gets the same norm from every block that holds it. Laid out
# column-major, numpy sums each row's squares column by column, the
# order the pre-#334 block (a column fancy-index, which numpy lays
# out column-major) was reduced in, so the default trajectory stays
# bit-identical; a row-major block would sum pairwise instead.
idx = self.comp_list[:, h]
block = np.asfortranarray(self.A[idx, :])
norm = np.sqrt(np.sum(block**2, axis=1)) # (data_dim,)
scale[idx] = np.where(norm > 0, norm, 1.0)
self.A = self.A / scale[:, None]
self.mu = self.mu * scale
self.beta = self.beta / scale
def _a_frozen(self) -> bool:
"""Whether the reference holds the A update (with its lrate ramp and
rho-rate reset) this iteration: once ``iter >= share_start``, every
iteration with ``mod(iter, share_int) <= 5``, counted from 1
(amica15.f90:1803, :func:`pamica.schedule.share_freeze`), exactly as
``AMICATorchNG._a_frozen`` decides it.
The reference applies this whether or not ``share_comps`` is on and for
any number of models (issue #345), so this does too: with the defaults
``share_start = share_int = 100``, every fit of 100 or more iterations
holds A on iterations 100-105, 200-205, and so on. The constructor
rejects ``share_int < 7``, since a shorter cycle would hold A for good.
"""
return schedule.share_freeze(self.iter, self.share_start, self.share_int)
def _optimize(self):
"""Main optimization loop."""
# Log optimization start
self.logger.info("Starting optimization...")
# Initialize optimization variables
numdecs = 0
numincs = 0
start_time = time.time()
convergence_reason = None
final_iter = 0
# Per-fit counters (reset here so a refit on the same instance gets a
# fresh budget). numrej is an instance attribute so the reject schedule
# and tests can read the pass count; it survives a restart (which
# preserves prior rejections) but not a fresh fit().
self.numrestarts = 0
self.numrej = 0
# Determine whether to use tqdm or per-line printing
use_tqdm_progress = self.use_tqdm and not self.verbose
# Create iterator (with or without tqdm)
if use_tqdm_progress:
# Use minimal progress bar with dynamic width and ASCII characters for better compatibility
progress_bar = tqdm(
range(self.max_iter),
desc="AMICA",
unit="it",
ncols=60, # Smaller fixed width
dynamic_ncols=True, # Adapt to terminal width
ascii=True, # Use ASCII characters for better compatibility
miniters=1, # Update on every iteration
leave=True,
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]", # Simpler format
)
iterator = progress_bar
else:
iterator = range(self.max_iter)
try:
for iter in iterator:
self.iter = iter
final_iter = iter
# Every schedule below counts iterations from 1, as the reference
# does (pamica.schedule converts this 0-indexed loop counter), so
# an identical setting fires on the same iterations here, in
# AMICATorchNG, in AMICAMLXNG, and in the binary (issue #335).
# One iteration follows the reference's main loop
# (amica15.f90:949-1142, issue #339): the E-step, which yields
# LL(iter) and the step with its norm; the restart-on-NaN window;
# the likelihood-decrease response and the stopping checks; an
# exit BEFORE any parameter moves if a check fired; otherwise the
# update with the rates the response just set, then the share
# merge, the checkpoints and rejection.
updates = self._get_updates_and_likelihood()
ll = float(updates["ll"])
# A non-finite likelihood is never recorded in self.ll, as it is
# never recorded in the PyTorch and MLX ll_history: the history
# stays the finite trajectory of the parameters the fit
# actually visited, and every backend's has the same length for
# the same event. A singular W makes logdet -> -inf (not NaN), so
# the guard is isfinite, not isnan.
if not np.isfinite(ll):
# Restart-on-NaN (Fortran amica15.f90:1022-1050): an early
# non-finite LL usually means an unlucky init, so redraw A
# and start over, up to maxrestarts times, within the first
# restartiter iterations (Fortran's absolute
# `iter <= restartiter` window over its 1-based counter; the
# iteration counter is not reset on restart here). Checked
# before the update, which the reference skips on a restart
# too (``startover``, :1115).
if (
schedule.within_restart_window(iter, self.restartiter)
and self.numrestarts < self.maxrestarts
):
self.numrestarts += 1
self.logger.warning(
"Non-finite LL at iter %d; reinitializing and "
"starting over (restart %d of %d).",
iter + 1,
self.numrestarts,
self.maxrestarts,
)
self._reinitialize_for_restart()
# The reference's checks still run on its restart
# iteration, and its min_dll comparison with the NaN
# likelihood is false (amica15.f90:1078-1090), so it
# zeroes numincs; its decrease comparison is false too,
# so numdecs is left as it was. Match that state, so
# small gains before the restart cannot count toward a
# min_dll stop after it (issue #339 review).
numincs = 0
continue
# Past the window, or out of restarts: stop, as the
# reference does (amica15.f90:1052-1055).
self.logger.warning(
"Non-finite log-likelihood (%s) at iteration %d; stopping.",
ll,
iter,
)
convergence_reason = self._NONFINITE_LL_REASON
break
step = self._update_direction(updates)
self.ll.append(ll)
self.nd.append(step.nd)
# A non-finite step or norm would pass both gradient-norm checks
# (NaN <= min_nd is False) and then be applied, so stop on it
# here, before any check reads it and before the update, with
# the parameters whose (finite) likelihood was just recorded.
if not (np.isfinite(step.nd) and np.all(np.isfinite(step.dAk))):
self.logger.warning(
"Non-finite update direction (ndtmpsum %s) at iteration "
"%d; stopping before the update.",
step.nd,
iter,
)
convergence_reason = self._NONFINITE_DIRECTION_REASON
break
# Calculate metrics for logging/progress
elapsed_time = time.time() - start_time
seconds_per_iter = (
elapsed_time / (iter + 1) if iter > 0 else elapsed_time
)
total_seconds = seconds_per_iter * self.max_iter
total_hours = total_seconds / 3600
current_seconds = (
elapsed_time / 3600 - int(elapsed_time / 3600)
) * 3600
if len(self.ll) > 1:
ll_diff = self.ll[-1] - self.ll[-2]
# Always log detailed metrics to the file logger
if self.use_grad_norm:
detailed_log = (
f" iter {iter + 1:5d} lrate = {self.lrate:12.10f} "
f"LL = {self.ll[-1]:13.10f} "
f"nd = {self.nd[-1]:11.10f}, "
f"D = {ll_diff:11.5e} {ll_diff:11.5e} "
f"({current_seconds:5.2f} s, {total_hours:4.1f} h)"
)
# Always write detailed logs to the file
self._append_to_log_file(detailed_log)
# Also log to console if verbose or not using tqdm
if self.verbose or not self.use_tqdm:
self.logger.info(detailed_log)
# Check convergence (threads numdecs/numincs back so they
# accumulate across iterations, and applies the likelihood-
# decrease response, including the ceiling ratchets).
converged, reason, numdecs, numincs = self._check_convergence(
numdecs, numincs
)
# Reset the decrease counter when Newton turns on (Fortran
# amica15.f90:1099-1102).
if schedule.newton_switches_on(self.do_newton, iter, self.newt_start):
numdecs = 0
# Stop before this iteration's update, as the reference does
# (amica15.f90:1111, ahead of update_params at :1122): the
# returned parameters are the ones whose likelihood is
# self.ll[-1].
if converged:
convergence_reason = reason
break
# Update parameters, with the rates the response above just set.
self._update_parameters(updates, step)
# Surface a corrupted update (a collapsed mixture component, a
# singular inverse) on the iteration it happens, with the same
# check and message as AMICAMLXNG and AMICATorchNG: otherwise a
# corruption on the last iteration would end as a max_iter fit
# holding non-finite parameters. One difference, for this
# backend's restart-on-NaN, which redraws A (and so W) and keeps
# the mixture parameters, as the reference's does
# (amica15.f90:1026-1046): when only A/W went non-finite and the
# next iteration still falls inside the restart window with a
# restart left, the next E-step's non-finite likelihood restarts
# the fit, as the reference's would, so the loop goes on to it.
# A restart could not repair any other parameter, so those stop.
bad = self._nonfinite_params(self._POST_UPDATE_CHECKED)
if bad:
restartable = (
set(bad) <= {"A", "W"}
and iter + 1 < self.max_iter
and schedule.within_restart_window(iter + 1, self.restartiter)
and self.numrestarts < self.maxrestarts
)
if not restartable:
self.logger.warning(
"Non-finite %s at iter %d (a mixture component likely "
"collapsed); stopping.",
", ".join(bad),
iter,
)
convergence_reason = (
f"{self._NONFINITE_PARAMS_REASON}: {', '.join(bad)}"
)
break
self.logger.warning(
"Non-finite %s at iter %d (a mixture component likely "
"collapsed); the next iteration's restart-on-NaN will "
"reinitialize.",
", ".join(bad),
iter,
)
# Share components if requested (Fortran identify_shared_comps
# schedule, amica15.f90:1856): once per share_int cycle from
# share_start, merging near-collinear components across
# models using the just-updated A, then rebuilding W from the
# merged comp_list (Fortran runs identify_shared_comps before
# get_unmixing_matrices, amica15.f90:1858,1863) -- otherwise the
# next E-step would read a stale W while indexing the densities
# by the merged comp_list. The schedule is pamica.schedule's, the
# one AMICATorchNG and AMICAMLXNG use, so an identical
# (share_start, share_int) fires on the same iterations in every
# backend.
#
# This runs AFTER self.ll recorded this iteration's likelihood
# (above), so a merge on the final iteration lands
# in self.A/comp_list but not in the ll value already stored --
# see the final_ll_ note on self.ll's init (issue #269).
if self.share_comps and schedule.periodic_due(
iter, self.share_start, self.share_int
):
self._identify_shared_comps()
self._update_unmixing_matrices()
# Write intermediate results/history if requested, on Fortran's
# cadence: `mod(iter, writestep) == 0` over its 1-indexed
# iteration counter (amica15.f90:1124/1130), so the first
# checkpoint lands at iteration `writestep`, not before it. This
# loop's `iter` is 0-indexed, and the literal transcription fired
# at iter 0 -- so EVERY fit wrote a checkpoint after its first
# iteration whatever writestep said, including fits far shorter
# than one interval. Both are skipped (loudly) from a non-finite
# state so a checkpoint cannot persist NaN parameters that
# loadmodout reads back without complaint (issue #240).
if (
self.outdir is not None
and self.writestep > 0
and schedule.every(iter, self.writestep)
):
self._write_checkpoint("results")
if (
self.outdir is not None
and self.do_history
and schedule.every(iter, self.histstep)
):
self._write_checkpoint("history")
# Reject outliers if requested, after the update and the
# checkpoints, as the reference does (amica15.f90:1136; the
# schedule, including its max(1, ...) clamp, is shared with
# AMICATorchNG and AMICAMLXNG in pamica.schedule). The statistic
# is this iteration's pre-update per-sample likelihood
# (_last_ll_samples, captured in the E-step).
if schedule.rejection_due(
self.do_reject,
iter,
self.rejstart,
self.rejint,
self.numrej,
self.maxrej,
):
self._reject_outliers()
self.numrej += 1
finally:
# Close the progress bar if using tqdm
if use_tqdm_progress:
progress_bar.close()
# Display final metrics after progress bar is closed
if len(self.ll) > 0 and self.use_grad_norm and len(self.nd) > 0:
final_metrics = (
f"Final LL: {self.ll[-1]:.6e}, Gradient norm: {self.nd[-1]:.6e}"
)
self.logger.info(final_metrics)
# Also log to file if using tqdm (since it wouldn't be logged during iterations)
self._append_to_log_file(final_metrics)
# Record and log the reason the loop stopped (None if it ran to
# max_iter). fit() uses self.converged for the terminal outcome.
self.stop_reason = convergence_reason
if convergence_reason:
self.logger.info(convergence_reason)
self._append_to_log_file(convergence_reason)
# Log final message (only once)
final_message = f"Optimization finished after {final_iter + 1} iterations"
self.logger.info(final_message)
def _check_convergence(
self, numdecs: int, numincs: int
) -> Tuple[bool, Optional[str], int, int]:
"""
Check convergence criteria and ratchet the learning rate.
Mirrors Fortran's per-iteration convergence handling
(amica17.f90:1062-1103). On a likelihood decrease Fortran does NOT stop
at ``maxdecs``; it lowers the learning-rate *ceilings* (``lrate0``; the
rho rate once ``iter > newt_start``; ``newtrate`` under Newton) and
continues, which is what keeps a long run from oscillating and drifting
past its converged solution (issue #41). The rho ceiling is
``self.rholrate_cap`` here (reset to ``rholrate0`` each fit), ratcheted
only at ``maxdecs`` (issue #193); each decrease scales the working
``self.rholrate``, which every A update resets to the ceiling. It runs
before the iteration's update, as the reference's checks do, so the
update already uses the new rates (issue #339). The updated
``numdecs``/``numincs`` counters are
returned so they accumulate across iterations (they previously did not).
Parameters
----------
numdecs : int
Consecutive-decrease counter (Fortran ``numdecs``).
numincs : int
Consecutive small-increase counter (Fortran ``numincs``).
Returns
-------
converged : bool
True if optimization should stop.
reason : str or None
Reason for stopping, or None.
numdecs, numincs : int
The updated counters (must be threaded back into the loop).
"""
if len(self.ll) == 0:
return False, None, numdecs, numincs
# Check for non-finite LL: a singular W makes logdet -> -inf (not NaN),
# so guard on isfinite, not isnan alone, or a degenerate model would run
# to max_iter undetected. _optimize stops on a non-finite likelihood
# before recording it, so this guards a direct caller's history.
if not np.isfinite(self.ll[-1]):
return (True, self._NONFINITE_LL_REASON, numdecs, numincs)
# The remaining checks compare consecutive iterations; skip until there
# are two LL values -- the first iteration, or the first iteration after
# a restart cleared the LL history (guard on len, not self.iter, which
# keeps counting across restarts).
if len(self.ll) < 2:
return False, None, numdecs, numincs
# Gradient norm is computed every iteration, so it is the decrease-stop
# floor unconditionally (Fortran's ndtmpsum); use_grad_norm only gates
# the separate final gradient-norm stop below.
grad_norm = self.nd[-1] if len(self.nd) > 0 else None
# Likelihood decrease (Fortran amica17.f90:1062-1083): reduce the current
# lrate, and once maxdecs decreases have accrued, ratchet the ceilings
# (lrate0; the rho rate; newtrate under Newton) down and continue --
# NOT stop. Only a lrate/gradient floor terminates on a decrease.
if self.ll[-1] < self.ll[-2]:
if self.lrate <= self.minlrate or (
grad_norm is not None and grad_norm <= self.min_grad_norm
):
return (
True,
"Converged: minimum change threshold met",
numdecs,
numincs,
)
self.lrate *= self.lratefact
# The working rho rate, as the reference scales it (:1063); every A
# update resets it to the ceiling before rho moves, so this reaches
# rho only on an iteration on which A is held.
self.rholrate *= self.rholratefact
numdecs += 1
if numdecs >= self.max_decs:
self.lrate0 *= self.lratefact
if schedule.past_newton_start(self.iter, self.newt_start):
# The rho-rate CEILING ratchets only here, at maxdecs
# (amica15.f90:1068), never per LL-decrease: a per-decrease
# decay of the rate rho actually uses, with no reset, collapsed
# the rho rate and froze rho at a stale shape (issue #193).
self.rholrate_cap *= self.rholratefact
if self.do_newton and schedule.past_newton_start(
self.iter, self.newt_start
):
self.newtrate *= self.lratefact
numdecs = 0
# Small likelihood increase (Fortran :1084-1096): stop after maxincs
# consecutive tiny gains; reset on any larger gain.
if self.use_min_dll:
if self.ll[-1] - self.ll[-2] < self.min_dll:
numincs += 1
if numincs > self.maxincs:
return (
True,
"Converged: small likelihood increase",
numdecs,
numincs,
)
else:
numincs = 0
# Gradient-norm floor (Fortran :1097-1103).
if (
self.use_grad_norm
and grad_norm is not None
and grad_norm <= self.min_grad_norm
):
return True, "Converged: small gradient norm", numdecs, numincs
return False, None, numdecs, numincs
def _reject_outliers(self):
"""Permanently drop samples whose (pre-update) per-sample log-likelihood
is a low outlier, mirroring ``AMICATorchNG._reject_outliers`` (Fortran
``reject_data``, amica17.f90:2380-2464): drop any currently-good sample
with ``loglik < mean - rejsig*std`` (population std). Rejection is
one-directional -- ``good_idx`` only shrinks -- and ``num_good_samples``
normalizes ``gm`` and ``self.ll`` thereafter (issue #212). Dropping the
most-negative samples raises the mean of what remains regardless of
normalization (removing below-average points can only raise a mean), so
the reject iteration is an LL increase and does not spuriously trip the
convergence checks.
``self._last_ll_samples`` is this iteration's per-sample LL over the
current good set (captured pre-update in ``_get_updates_and_likelihood``,
in ``good_idx`` order), so the keep-mask indexes straight into
``good_idx``.
"""
if not self.do_reject:
return
assert self.good_idx is not None and self._last_ll_samples is not None
ll_vec = self._last_ll_samples
mean = ll_vec.mean()
# Population std (ddof=0), matching np.std's default and the torch path.
std = np.sqrt(max(float(np.mean(ll_vec**2) - mean**2), 0.0))
keep = ll_vec >= (mean - self.rejsig * std)
if not bool(keep.any()):
# For finite log-likelihoods the max sample is always >= mean >=
# mean - rejsig*std (rejsig>0 is validated at construction), so it is
# always kept; the only way every sample is dropped is a non-finite
# per-sample LL (one NaN poisons mean/std, making every comparison
# False). Report that accurately instead of blaming rejsig (issue
# #127), which a user cannot fix by tuning rejsig. In a normal fit()
# the earlier aggregate non-finite-LL guard (the sum is non-finite
# iff a term is) stops the loop first, so this mainly serves direct
# callers of _reject_outliers and is defense in depth.
n_bad = int(np.count_nonzero(~np.isfinite(ll_vec)))
if n_bad:
raise ValueError(
f"{n_bad} of {ll_vec.size} samples have a non-finite "
"log-likelihood; this indicates numerical instability "
"upstream (singular W / overflow), not a rejsig "
"miscalibration. Check for rank-deficient or "
"average-referenced data, or reduce the learning rate."
)
raise ValueError( # defensive: unreachable for finite LL, rejsig>0
f"Outlier rejection removed all {self.good_idx.size} samples "
f"(rejsig={self.rejsig} too aggressive for this data)."
)
# Zero the LLt stash for the samples being dropped, as Fortran's
# reject_data does (amica15.f90:2231-2234): they are never scored again,
# so otherwise they would keep the log-likelihood from the last
# iteration that still counted them good, and load_rej's
# ``sum(modloglik(:,i)) == 0`` sentinel would not see them as rejected.
if self._llt_logv is not None and self._llt_ll is not None:
dropped = self.good_idx[~keep]
self._llt_logv[dropped] = 0.0
self._llt_ll[dropped] = 0.0
n_before = self.good_idx.size
self.good_idx = self.good_idx[keep]
self.num_good_samples = int(self.good_idx.size)
self.logger.info(
"Rejected %d samples (%d good remaining).",
n_before - self.num_good_samples,
self.num_good_samples,
)
def _write_results(self):
"""Write current results to disk in the Fortran AMICA binary format.
Writes raw little-endian float64 (and int32 ``comp_list``) files with no
extension, in the layout that ``load.loadmodout`` reads (and that
``load_results`` reads back), so pamica output is loadable by the same
reader as the Fortran reference (issue #30).
This shares the ``write_amicaout`` writer with the torch backend, so its
on-disk guarantees apply here too: single-model output is byte-identical
to the Fortran ``amicaout`` files, and since #159 multi-model output is
written in genuine Fortran layout as well (each array's model axis
slowest, column-major within a model), so EEGLAB's ``loadmodout15.m``
reads both correctly. See :func:`pamica.numpy_impl.load.write_amicaout`.
Also writes ``LLt`` (issue #155) from the stash the E-step already
filled (:meth:`_llt_arrays`, issue #157). A ``writestep`` checkpoint
therefore pays no forward pass of its own -- and, exactly as in the
reference, writes the E-step of the parameters as they stood before
this iteration's M-step. See :meth:`_llt_arrays` for that ordering and
its Fortran citation.
"""
# A is written in the reference's layout, as the binary writes it
# (loadmodout derives A from W and S instead); load_results reads it
# back for the viz helpers. The Fortran 'nd' file (per-component
# weight-change history) is a different quantity from pamica's scalar
# self.nd, so it is not emitted (loadmodout treats 'nd' as optional).
from .load import write_amicaout
# Raises rather than asserts: ``assert`` is stripped under ``python -O``.
if self.outdir is None:
raise RuntimeError("_write_results needs an outdir")
if self.A is None:
raise RuntimeError("_write_results needs a fitted A; call fit() first")
Lht, Lt = self._llt_arrays()
write_amicaout(
self.outdir,
gm=self.gm,
W=self.W,
sphere=self.sphere,
mean=self.mean,
c=self.c,
alpha=self.alpha,
mu=self.mu,
sbeta=self.beta, # Fortran's 'sbeta' is pamica's beta (scale)
rho=self.rho,
comp_list=self.comp_list,
ll=np.asarray(self.ll),
# The reference layout, (nw, num_comps) with component k in column
# k: the component-row A transposed (issue #334).
A=self.A.T,
Lht=Lht,
Lt=Lt,
)
def _write_history(self):
"""Write optimization history at current iteration."""
if not self.do_history:
return
assert (
self.A is not None
and self.W is not None
and self.c is not None
and self.mu is not None
and self.alpha is not None
and self.beta is not None
and self.rho is not None
and self.gm is not None
and self.mean is not None
and self.sphere is not None
and self.comp_list is not None
)
assert self.outdir is not None, "_write_history needs an outdir"
hist_dir = self.outdir / "history" / f"{self.iter:06d}"
if not hist_dir.exists():
hist_dir.mkdir(parents=True)
# Save current state
np.save(hist_dir / "A.npy", self.A)
np.save(hist_dir / "W.npy", self.W)
np.save(hist_dir / "c.npy", self.c)
np.save(hist_dir / "mu.npy", self.mu)
np.save(hist_dir / "alpha.npy", self.alpha)
np.save(hist_dir / "beta.npy", self.beta)
np.save(hist_dir / "rho.npy", self.rho)
np.save(hist_dir / "gm.npy", self.gm)
np.save(hist_dir / "mean.npy", self.mean)
np.save(hist_dir / "sphere.npy", self.sphere)
np.save(hist_dir / "comp_list.npy", self.comp_list)
np.save(hist_dir / "ll.npy", self.ll)
if self.use_grad_norm:
np.save(hist_dir / "nd.npy", self.nd)
def transform(self, data: np.ndarray) -> np.ndarray:
"""
Apply the learned unmixing matrices to new data.
Parameters
----------
data : ndarray of shape (n_channels, n_samples)
The data to transform
Returns
-------
S : ndarray of shape (n_components, n_samples, n_models)
The unmixed sources for each model
Raises
------
RuntimeError
If the model is unfitted, or the fit ended degenerate (issue #306).
ValueError
If ``data`` is not a 2D array of the fitted input channel count.
"""
if self.W is None or self.comp_list is None or self.c is None:
raise RuntimeError("Model has not been fitted yet; call fit() first.")
self._check_usable("transform")
self._check_input_shape(data)
if self.mean is not None:
data = data - self.mean
if self.sphere is not None:
data = np.dot(self.sphere, data)
S = np.zeros((self.num_comps, data.shape[1], self.num_models))
for h in range(self.num_models):
idx = self.comp_list[:, h]
# W^T is the true unmixing (issue #24 transpose convention); c is the
# per-model data-space center, so unmix as W(x - c) (issue #27).
S[idx, :, h] = np.dot(self.W[:, :, h].T, data - self.c[:, h][:, None])
return S