class AMICATorchNG:
"""
Natural-gradient EM AMICA, ported from ``pamica.numpy_impl.core.AMICA``.
Not an ``nn.Module``: there are no learnable ``nn.Parameter``s and no
autograd. Parameters (``A``, ``W``, ``c``, ``mu``, ``alpha``, ``beta``,
``rho``, ``gm``) are plain tensors mutated in place by closed-form
E-step/M-step updates each iteration, mirroring
``pamica.AMICA._get_block_updates``/``_update_parameters``.
Parameters
----------
n_channels : int
Number of input channels (``data_dim`` in the NumPy/Fortran code).
n_models : int, default=1
Number of ICA mixture models.
n_mix : int, default=3
Number of mixture components per source.
block_size : int, default=8192
Number of samples processed per accumulation block. Peak memory
during the E-step scales with this, not with the total sample count.
Larger blocks give bigger tensor ops (less Python/dispatch overhead,
better threading/GPU utilization) at higher memory. Every backend is
dispatch-bound at small blocks, making this the largest throughput knob:
raising it from 512 to 8192 is ~6x on CPU float64 for the bundled sample
(issue #216). 8192 rather than larger because peak block memory scales
with it and 8192 stays near 240 MB even at 256 channels.
Fortran pins no comparable value (header default 128, auto-tuned over
128-1024 via ``do_opt_block``). Per-iteration sufficient statistics are
block-size-independent to ~1e-8 (``test_blocking_invariance``); the
multi-iteration trajectory shifts ~1e-6, inside parity tolerance but
enough that a bit-for-bit Fortran comparison must match ``block_size`` on
both sides (the bundled ``input.param`` uses 512).
do_opt_block : bool, default=False
Time a few candidate block sizes on the real data and device at the
start of ``fit`` and keep the fastest, instead of using ``block_size``
as given (issue #232; Fortran ``do_opt_block``). The measured optimum
moves with host, device and data, so a fixed default necessarily leaves
16-60% on the table depending on backend -- but the choice is
**timing-based and therefore machine-dependent**, so two hosts can pick
different sizes and their trajectories then differ at the same ~1e-6
level any ``block_size`` change produces. Off by default for that
reason: a run compared bit-for-bit against the reference binary must
leave this off and pin ``block_size``. When on, ``block_size`` is the
fallback the search keeps if no candidate can be timed.
Unlike Fortran, whose ``determine_block_size`` aborts the run when a
candidate cannot be allocated, a failing candidate here is skipped, the
upward search stops there, and the fit continues at the largest size
that ran. See :mod:`pamica.blocktune`.
blk_min, blk_max, blk_step : int, defaults 4096, 32768, 4096
Candidate sweep for ``do_opt_block``: ``blk_min``, ``blk_min +
blk_step``, ..., ``<= blk_max``, Fortran's arithmetic stepping, each
clamped to ``n_samples`` and to a conservative memory estimate.
Validated (and only used) when ``do_opt_block`` is on. The defaults are
re-derived rather than copied from Fortran's 128-1024, which sits far
below where any pamica backend peaks; they bracket the measured CPU
optimum and include the 8192 default.
lrate : float, default=0.1
Initial/maximum natural-gradient learning rate (``lrate0`` in NumPy).
minlrate : float, default=1e-12
Hard learning-rate floor: once ``lrate`` anneals to it, ``fit`` stops
(``stop_reason="lrate_floor"``).
lratefact : float, default=0.5
Factor by which ``lrate`` (and the ceiling ``lrate_cap``/``newtrate``)
are annealed when the log-likelihood decreases; see ``fit`` for the
Fortran-style ``numdecs``/``maxdecs`` ratchet. The response runs on
the iteration whose E-step saw the decrease, before that iteration's
update, so the step it takes already uses the halved rate, as in the
reference (amica15.f90:1056-1122, issue #339).
maxdecs : int, default=5
Number of log-likelihood decreases after which the learning-rate
*ceiling* is ratcheted down (Fortran ``maxdecs``). The decreases need
not be consecutive: the count resets only at each ratchet and when
Newton switches on, as in the reference (amica15.f90:1062-1076).
use_min_dll : bool, default=True
Enable the small-likelihood-increase stop (Fortran ``use_min_dll``,
amica15_header.f90:24/74; amica15.f90:1078-1090): once the per-
sample-channel log-likelihood gain ``ll_history[-1] - ll_history[-2]``
falls below ``min_dll`` for more than ``maxincs`` *consecutive*
iterations, ``fit`` stops (``stop_reason="min_dll"``). The counter
resets to 0 on any iteration with a larger gain (including a
likelihood decrease, which is always "less than" a positive
``min_dll``, so it also increments the counter). Checked every
iteration once two log-likelihood values exist (never on the first).
min_dll : float, default=1e-9
Threshold for ``use_min_dll``, on the log-likelihood's own scale
(mean log-likelihood per sample-channel, matching ``ll_history`` --
see ``amica15.f90:1770``, which normalizes ``LL(iter)`` by
``numgoodsum*nw`` before this comparison in the reference).
maxincs : int, default=5
Number of consecutive small-gain iterations tolerated before
``use_min_dll`` stops the fit (Fortran ``maxincs``, not itself
configurable from the Fortran param file -- fixed at its header
default).
use_grad_norm : bool, default=True
Enable the weight-gradient-norm stop (Fortran ``use_grad_norm``,
amica15_header.f90:24/74; amica15.f90:1091-1097): once the RMS
weight-update norm ``ndtmpsum`` (see ``min_nd``) falls to or below
``min_nd``, ``fit`` stops (``stop_reason="grad_norm"``). This is
independent of ``use_min_dll`` and of whether the log-likelihood
just decreased; it is also folded into the likelihood-decrease
branch unconditionally (``stop_reason="grad_norm_floor"``, Fortran
amica15.f90:1058's ``.or. (ndtmpsum .le. min_nd)``, alongside the
existing ``lrate <= minlrate`` check) -- this decrease-branch half is
what fixes the reported CUDA/``do_newton=True`` case where ``lrate``
sits at ``newtrate`` and oscillates instead of annealing, so the old
``lrate_floor``-only check never fired and ``max_iter`` was the only
stop (issue #207). Checked every iteration once two log-likelihood
values exist (never on the first).
CAUTION: with ``use_grad_norm`` at this default (``True``), the
stop that actually surfaces for the fixed CUDA scenario is
``"grad_norm"``, not ``"grad_norm_floor"``. The standalone check
above runs every iteration regardless of LL direction and (in
``fit``'s per-iteration ordering) is evaluated after the
likelihood-decrease branch, with no ``elif``/``leave`` gate between
them; whichever iteration first satisfies ``ndtmpsum <= min_nd``
also satisfies the standalone check that same iteration, so it always
overwrites ``stop_reason`` before a decrease-gated
``"grad_norm_floor"`` could be the value ``fit`` finally reports.
``"grad_norm_floor"`` is therefore only distinctly reachable as the
*final* ``stop_reason`` when ``use_grad_norm=False`` (isolating the
decrease-branch half, as ``test_grad_norm_floor_fires_on_likelihood_decrease``
does); ``test_grad_norm_shadows_grad_norm_floor_under_shipped_defaults``
(same setup, ``use_grad_norm`` left at its default) confirms the
shadowing directly. ``"min_dll"`` can likewise be shadowed by
``"grad_norm"`` if both conditions happen to hold in the same
iteration -- Fortran has this same structure (independent
``leave=.true.`` assignments with no declared precedence among them),
so this is not a fidelity bug, just a reporting nuance worth knowing
before reading ``stop_reason`` as a precise diagnosis.
min_nd : float, default=1e-7
Threshold for ``use_grad_norm`` (and the decrease-branch grad-norm
check). Matches Fortran's ``ndtmpsum`` (amica15.f90:1760-1761): the
RMS, over ``comp_used`` components only, of the per-iteration
weight-update direction ``dAk`` (the natural-gradient/Newton step
before the ``lrate`` scaling and before the reference's A-freeze
(see ``share_iter``) may discard it) -- ``sqrt(sum(dAk**2, axis=1)[comp_used].sum() /
(n_channels * comp_used.sum()))``, one squared norm per component
row, as the reference sums each component's column. The
``comp_used`` mask only differs from all-True when ``share_comps``
has merged components away (issue #60); it is a no-op otherwise. Computed every iteration
regardless of ``use_grad_norm``/``use_min_dll`` (both stops read the
same per-iteration value; Fortran computes ``ndtmpsum`` unconditionally
too, in ``accum_updates_and_likelihood``, before either check runs).
Not reachable on small recordings, in any implementation: the reference
binary's own gradient norm plateaus at 1.0-1.65e-5 on the bundled
32-channel sample, two orders above this threshold, so the stop never
fires there and ``min_dll`` is what ends the fit. The default is kept
Fortran-faithful rather than retuned; see the convergence-criteria
section of ``docs/guides/validation.md`` (issue #218).
newt_ramp : int, default=10
Denominator of the per-iteration learning-rate ramp toward the current
ceiling: ``lrate = min(ceiling, lrate + min(1/newt_ramp, lrate))``
(ceiling is ``lrate_cap`` for natural gradient, ``newtrate`` for
Newton).
do_newton : bool, default=False
Enable the Newton preconditioner for the ``A``/``W`` update from
iteration ``newt_start`` on, ported from the Fortran reference: a
per-source-pair 2x2 solve that preconditions the natural gradient by
an approximate Hessian, which converges faster near the optimum. An
iteration whose Hessian is not positive definite falls back to the
natural gradient (counted in ``n_newton_fallbacks``). Off by default,
as in the reference's compiled default; the bundled reference
``input.param`` turns it on (``do_newton 1``).
newt_start : int, default=20
Iteration at which the Newton step switches on (natural gradient is
used before it, letting the mixture parameters settle first). Counted
from 1, as the reference counts it (``iter .ge. newt_start``): the
Newton step is taken on the ``newt_start``-th iteration, i.e. once
``iteration + 1 >= newt_start`` for the 0-based ``iteration``
attribute (issue #335). ``newt_start=0`` fits exactly as ``1`` does,
Newton from the first iteration. It also gates the ``maxdecs`` ratchet
of the ``rholrate`` ceiling (and, under Newton, of ``newtrate``) to
iterations after ``newt_start`` (Fortran amica15.f90:1067/1070), so it
is validated (an integer >= 0) whether or not ``do_newton`` is on.
newtrate : float, default=0.5
Maximum learning rate the ramp climbs to while Newton is active
(the natural-gradient phase is capped at ``lrate``/``lrate0``).
do_reject : bool, default=False
Enable Fortran-style outlier rejection: after the parameter update,
samples whose total log-likelihood falls below
``mean - rejsig*std`` are permanently excluded from subsequent
sufficient-statistic accumulation and from the sample count used to
normalize ``gm`` and the reported log-likelihood.
rejsig : float, default=3.0
Rejection threshold in standard deviations of the per-sample
log-likelihood.
rejstart, rejint, maxrej : int
First rejection iteration, interval between rejections, and maximum
number of rejection passes (matching ``amica15.f90:1136``).
``rejstart`` counts from 1, as the reference counts it (issue #335),
and must be an integer >= 1 when ``do_reject`` is on: ``rejstart <= 0``
would silently skip the reference's unconditional first pass.
rho0, minrho, maxrho, rholrate : float
Generalized-Gaussian shape-parameter initialization, clamp bounds,
and learning rate. As in the reference, the rate has a working value
(``rholrate``), which each likelihood decrease multiplies by
``rholratefact``, and a ceiling (``rholrate_cap``), which the
``maxdecs`` ratchet multiplies by ``rholratefact`` after
``newt_start``; every update of ``A`` resets the working value to the
ceiling before ``rho`` moves, so the decrease scaling reaches ``rho``
only on an iteration on which ``A`` is held (see ``share_iter``).
keep_best : bool, default=True
Return the highest-log-likelihood iterate instead of the last one
(issue #51). The lrate schedule is non-monotone (it anneals only after
an LL *decrease*), so a late Newton-fallback overshoot can leave the
final iterate below a peak the run already reached. When the final LL
falls more than a small tolerance below that peak, ``fit`` restores the
peak's parameters. A monotone single-model run (issue #24 parity) is a
bit-exact no-op. Automatically inactive under ``do_reject`` (the
good-sample set, and thus the LL normalization, changes across
iterations, making per-iteration LLs incomparable) and under
``share_comps`` (a merge changes the parameter count, so pre- and
post-merge LLs are not comparable and reverting to an earlier snapshot
would silently undo the merge; issue #269).
pdftype : int, default=0
Source-density family (issue #26), matching Fortran ``amica15.f90``'s
``pdtype`` codes: 0 generalized Gaussian (default; rho adapts), 2
Gaussian, 3 logistic, 4 sub-Gaussian cosh+. ``pdftype=1`` enables the
extended-Infomax adaptive switcher, which flips each source between the
super-Gaussian (code 1) and sub-Gaussian (code 4) cosh densities by
kurtosis sign. For every non-GG family the GG shape update is frozen
(Fortran ``dorho=.false.``); the single-component families 1/4 (and the
adaptive mode) require ``n_mix=1``.
kurt_start, num_kurt, kurt_int : int
Adaptive-switch schedule (only used when ``pdftype=1``): first iteration
to re-estimate kurtosis (counted from 1), number of switch passes, and
the iteration interval between them. ``num_kurt=0`` disables switching
(the family stays at its super-Gaussian init).
invsigmin, invsigmax : float
Clamp bounds for the mixture scale parameter ``beta``.
doscaling, scalestep : bool, int
Whether/how often to rescale each component's mixing vector (a row of
the stored ``A``) to unit norm, with the matching
``mu``/``beta`` rescale, an exact change of scale
(see :meth:`_rescale_components`). The rescale runs on iterations
``scalestep``, ``2*scalestep``, ... counted from 1; the default 1
rescales every iteration, as the reference always does (it ignores
``scalestep``). ``scalestep`` is validated only when ``doscaling`` is
on (an integer >= 1, or the constructor raises ``ValueError``); with
``doscaling`` off it is inert, never read.
share_comps : bool, default=False
Enable multi-model component sharing (Fortran ``share_comps`` /
``identify_shared_comps``, amica15.f90:1916): components that are
near-collinear across different models are merged so they share one
mixing vector (one row of ``A``) and one density. Requires
``n_models >= 2`` (a model cannot share with itself); a no-op
otherwise. OFF by default, so single-model (#24) and default
multi-model (#27) results are unchanged. The reference's similarity
metric is never initialized (like ``do_choose_pdfs``, #26), so the
metric has no bit-exact oracle; the merged state it produces does, and
the update from it matches the reference to round-off (issue #334,
see :meth:`_identify_shared_comps`). A merge that
fires on the LAST fit iteration is reflected in the returned model but
trails in ``final_ll_``; see that attribute's comment (issue #269).
share_start, share_iter : int
Sharing schedule: first iteration to attempt merges (counted from 1)
and the interval between attempts (Fortran
``share_start``/``share_iter``). They also set the reference's A-freeze,
which applies to every fit, sharing on or off (issue #345): from
iteration ``share_start`` on, the update of ``A`` (with its lrate ramp)
is held on every iteration whose number, counted from 1, has a
remainder of 0 to 5 modulo ``share_iter`` (amica15.f90:1803), so with
the defaults on iterations 100-105, 200-205, and so on.
Both are validated whether or not ``share_comps`` is on:
``share_start`` must be an integer >= 1, and ``share_iter`` an
integer >= 7, or the freeze would hold A permanently from
``share_start`` on.
comp_thresh : float, default=0.99
Cosine-similarity cutoff (in the de-sphered/sensor-space metric) above
which two components' mixing vectors are identified and merged. The
de-sphering uses
``pinv(sphere)``, so sharing also works on rank-reduced and
rank-deficient fits (issues #253, #221); see
:meth:`_identify_shared_comps`.
do_mean, do_sphere, do_approx_sphere : bool
Preprocessing options, matching ``pamica.AMICA._preprocess_data``.
pcakeep, pcadb : int, float, optional
Explicit PCA dimensionality reduction. ``pcakeep`` keeps that many
principal dimensions (Fortran ``pcakeep``); ``pcadb`` keeps those whose
covariance eigenvalue lies within ``pcadb`` dB of the largest (a pamica
extension: the reference parses ``pcadb`` but never uses it). Both are
capped by the detected numerical rank (``mineig``/``mineig_rel``),
matching Fortran's ``numeigs = min(pcakeep, count(eigs > mineig))``.
``pcakeep`` must be an integer >= 1 and ``pcadb`` a finite number > 0;
anything else raises ``ValueError`` at construction. When both are set,
``pcakeep`` takes precedence and ``pcadb`` is ignored (one INFO log
line), as in the reference. Both are ignored, with one WARNING, when
``do_sphere=False``: reduction happens only while sphering, as in the
reference. ``None`` (the default) for both leaves only automatic rank
detection. See :mod:`pamica.rank`.
mineig : float, default=1e-15
Absolute floor on data-covariance eigenvalues used to detect the
numerical rank (Fortran ``mineig``, amica15.f90:413 and
amica15_header.f90:66). Eigen-directions at or below it are dropped, the
model is sized to the surviving rank, and sensor-space maps come from
:meth:`get_sensor_mixing_matrix`. Full-rank data keep every eigenvalue,
so this is a no-op there and single-model parity is byte-for-byte.
Being absolute, it is unit-dependent: EEG in microvolts gives
eigenvalues of order 1-100 and the default behaves, but MEG in Tesla
gives ~1e-26 and every eigenvalue falls below it, which Fortran would
turn into ``numeigs = 0``. pamica raises instead of fitting an empty
model. Use ``mineig_rel`` (or rescale) for such data.
mineig_rel : float, optional
Scale-free alternative to ``mineig``: when set, the threshold becomes
``mineig_rel * largest_eigenvalue`` and ``mineig`` is ignored. Off by
default so rank detection stays Fortran-exact. It is also the more
accurate detector -- the absolute floor sits amid the numerical-zero
eigenvalues of rank-deficient data and over-retains, while a relative
floor recovers the true rank (issue #223).
seed : int, optional
Seed for parameter initialization. Uses ``numpy.random.RandomState``
internally (not ``torch``'s RNG) with the exact same draw order as
``pamica.AMICA._initialize_parameters``, so the same seed produces
bit-identical starting parameters to the NumPy reference.
n_restarts : int, default=1
Number of independent fits to run from different seeds, keeping the one
with the highest ``final_ll_`` (issue #198). ``1`` (the default) is the
parity-preserving setting: the restart machinery is bypassed entirely
and the fit is bit-identical to a pre-#198 run. With ``n_restarts > 1``
a base ``seed`` (or explicit ``restart_seeds``) is required, so the
winning fit can be reproduced. Restarts run serially, so a fit costs
``n_restarts`` times as long. Fortran has no equivalent; see
``docs/guides/amica-differences.md`` and :mod:`pamica.restarts`.
restart_seeds : sequence of int, optional
Explicit per-restart seeds; must have exactly ``n_restarts`` entries.
When omitted the seeds are ``seed, seed + 1, ..., seed + n_restarts - 1``.
device : str or torch.device, optional
Compute device for the block loop. ``None`` picks MPS, then CUDA,
then CPU, whichever is available first. MPS has no float64, so when
``None`` picks MPS for a float64 run (the default ``dtype``) the
constructor uses the CPU instead and logs a warning; with
``dtype=torch.float32``, ``None`` keeps MPS. An explicit
``device="mps"`` with float64 raises ``ValueError``. Preprocessing
(mean/cov/eigh) is always done in float64 on CPU regardless of device,
since eigh is not reliably supported on MPS.
dtype : torch.dtype, default=torch.float64
Parameter/computation dtype. float64 is the parity default (Fortran
bit-parity) and ~4.5x on CUDA over CPU (issue #63). float32 converges on
full-size data across seeds (issue #75 guarded the one float32-only
divide-by-zero -- a sample rounding an activation to exactly 0 gave
``0/0`` in the mu denominator) and is the required precision on MPS,
which has no float64. float32 is NOT bit-parity with float64 (~7
significant digits), so use float64 for Fortran-parity runs and float32
for speed / Apple-GPU.
"""
def __init__(
self,
n_channels: int,
n_models: int = 1,
n_mix: int = 3,
block_size: int = 8192,
do_opt_block: bool = False,
blk_min: int = blocktune.DEFAULT_BLK_MIN,
blk_max: int = blocktune.DEFAULT_BLK_MAX,
blk_step: int = blocktune.DEFAULT_BLK_STEP,
lrate: float = 0.1,
minlrate: float = 1e-12,
lratefact: float = 0.5,
maxdecs: int = 5,
use_min_dll: bool = True,
min_dll: float = 1e-9,
maxincs: int = 5,
use_grad_norm: bool = True,
min_nd: float = 1e-7,
newt_ramp: int = 10,
do_newton: bool = False,
newt_start: int = 20,
newtrate: float = 0.5,
do_reject: bool = False,
rejsig: float = 3.0,
rejstart: int = 2,
rejint: int = 3,
maxrej: int = 1,
rho0: float = 1.5,
minrho: float = 1.0,
maxrho: float = 2.0,
rholrate: float = 0.05,
rholratefact: float = 0.1,
keep_best: bool = True,
pdftype: int = 0,
kurt_start: int = 3,
num_kurt: int = 5,
kurt_int: int = 1,
invsigmin: float = 1e-4,
invsigmax: float = 1000.0,
doscaling: bool = True,
scalestep: int = 1,
share_comps: bool = False,
share_start: int = 100,
share_iter: int = 100,
comp_thresh: float = 0.99,
do_mean: bool = True,
do_sphere: bool = True,
do_approx_sphere: bool = True,
pcakeep: Optional[int] = None,
pcadb: Optional[float] = None,
mineig: float = MINEIG,
mineig_rel: Optional[float] = MINEIG_REL,
seed: Optional[int] = None,
n_restarts: int = restarts.DEFAULT_N_RESTARTS,
restart_seeds: Optional[Sequence[int]] = None,
device: Optional[Union[str, torch.device]] = None,
dtype: torch.dtype = torch.float64,
):
self.n_channels = n_channels
# The input channel count, kept apart from n_channels, which
# _preprocess shrinks to the kept rank on any rank reduction. fit()
# validates X against this and resets n_channels/n_comps from it, so a
# refit or a later restart starts from the constructor's geometry.
self._n_input_channels = n_channels
self.n_models = n_models
self.n_mix = n_mix
self.n_comps = n_channels * n_models
self.mineig = mineig
self.mineig_rel = mineig_rel
self.block_size = block_size
self.do_opt_block = do_opt_block
self.blk_min = blk_min
self.blk_max = blk_max
self.blk_step = blk_step
if do_opt_block:
# Only validated when the search is on, matching how share_comps /
# do_reject validate their own schedules: these three are inert
# otherwise, and a literal Fortran input.param carrying them
# alongside do_opt_block=0 must stay loadable (issue #232).
blocktune.validate_block_tune_params(blk_min, blk_max, blk_step)
self.lrate0 = lrate
self.lrate = lrate
self.minlrate = minlrate
self.lratefact = lratefact
self.maxdecs = maxdecs
# Convergence stops (issue #207), Fortran-faithful defaults (both
# amica15_header.f90:24/74 flags default True): the small-
# likelihood-increase stop (use_min_dll/min_dll/maxincs) and the
# weight-gradient-norm stop (use_grad_norm/min_nd). See fit() for the
# per-iteration checks and _update_direction for the ndtmpsum
# computation these both read.
if maxincs < 0:
raise ValueError(f"maxincs must be >= 0, got {maxincs}")
self.use_min_dll = use_min_dll
self.min_dll = min_dll
self.maxincs = maxincs
self.use_grad_norm = use_grad_norm
self.min_nd = min_nd
self.newt_ramp = newt_ramp
self.do_newton = do_newton
# 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", newt_start, 0)
self.newt_start = newt_start
self.newtrate = newtrate
self.newtrate0 = newtrate
self.do_reject = do_reject
self.rejsig = rejsig
self.rejstart = rejstart
self.rejint = rejint
self.maxrej = maxrej
if do_reject:
if rejint < 1:
raise ValueError(f"rejint must be >= 1, got {rejint}")
if rejsig <= 0:
raise ValueError(f"rejsig must be > 0, got {rejsig}")
if maxrej < 0:
raise ValueError(f"maxrej must be >= 0, got {maxrej}")
# Counted from 1, so rejstart <= 0 would silently disable the
# reference's unconditional ``iter == rejstart`` pass.
schedule.validate_iteration_setting("rejstart", rejstart, 1)
self.rho0 = rho0
self.minrho = minrho
self.maxrho = maxrho
# The reference keeps a working rho rate (``rholrate``) and its ceiling
# (``rholrate0``): a likelihood decrease scales the working rate, a
# maxdecs ratchet the ceiling, and every A update resets the working
# rate to the ceiling before the rho update (amica15.f90:1063-1068,
# 1806/1813). ``rholrate_cap`` is that ceiling (as ``lrate_cap`` is
# lrate's); ``rholrate0`` keeps the constructor value for a refit.
self.rholrate = rholrate
self.rholrate_cap = rholrate
self.rholrate0 = rholrate
self.rholratefact = rholratefact
# Best-iterate safeguard (issue #51). When True, fit() restores the
# highest-log-likelihood iterate if the run ends more than _KEEP_BEST_TOL
# below it (a late Newton-fallback overshoot). Disabled automatically
# under do_reject, where the good-sample set (and the LL normalization)
# changes across iterations, so per-iteration LLs are not comparable.
self.keep_best = keep_best
# Source-density family selection (Fortran ``pdftype``, amica15.f90). Values
# match Fortran's per-source ``pdtype`` codes: 0 generalized Gaussian (the
# default, GG-mixture with adaptive rho), 2 Gaussian mixture, 3 logistic
# (sech^2) mixture, 4 sub-Gaussian cosh+ (single component). pdftype=1 enables
# the extended-Infomax adaptive switcher (Fortran's do_choose_pdfs trigger),
# which flips each source between the super-Gaussian (code 1) and sub-Gaussian
# (code 4) cosh densities by kurtosis sign on the kurt_start/num_kurt/kurt_int
# schedule. Families 1 and 4 are single-component (no alpha mixture).
if pdftype not in (0, 1, 2, 3, 4):
raise ValueError(f"pdftype must be one of 0,1,2,3,4; got {pdftype}")
self.pdftype = pdftype
# Fortran freezes the GG shape update for every non-GG family (amica15.f90:
# `if (pdftype /= 0) dorho = .false.`, lines 3704-3705).
self.dorho = pdftype == 0
# pdftype==1 is Fortran's adaptive trigger (amica15.f90:612).
self.do_choose_pdfs = pdftype == 1
self.kurt_start = kurt_start
self.num_kurt = num_kurt
self.kurt_int = kurt_int
# Families 1/4 (and the adaptive mode, which uses only codes 1 and 4) are
# single-component densities: Fortran's z0 references only mixture component
# j=1 and omits log(alpha). They are meaningful only with n_mix == 1.
if pdftype in (1, 4) and n_mix != 1:
raise ValueError(
f"pdftype={pdftype} is a single-component density (adaptive mode "
f"uses codes 1 and 4); it requires n_mix=1, got n_mix={n_mix}."
)
# Validate the adaptive-switch schedule up front (mirrors the do_reject
# checks below): kurt_int==0 would otherwise raise a bare ZeroDivisionError
# deep in fit(), and a negative kurt_int silently changes the schedule.
if self.do_choose_pdfs:
if kurt_int < 1:
raise ValueError(f"kurt_int must be >= 1, got {kurt_int}")
if kurt_start < 1:
raise ValueError(f"kurt_start must be >= 1, got {kurt_start}")
if num_kurt < 0:
raise ValueError(f"num_kurt must be >= 0, got {num_kurt}")
self.invsigmin = invsigmin
self.invsigmax = invsigmax
self.doscaling = doscaling
self.scalestep = scalestep
if doscaling:
# A zero cadence divides by zero mid-fit; a fractional one fires on
# no meaningful schedule. The reference never reads scalestep.
schedule.validate_iteration_setting("scalestep", scalestep, 1)
# Component sharing (Fortran share_comps / identify_shared_comps trigger
# amica15.f90:1856, subroutine :1916-1963): periodically merge
# components near-collinear across DIFFERENT models so they share one
# density and one mixing vector (one row of A, issue #334). Multi-model
# only (a model cannot share with itself); OFF by default so
# single-model (#24) and default multi-model (#27) parity stay
# byte-for-byte. The reference's Spinv2 metric is declared but never
# allocated, so its own scan never merges (every similarity is NaN; the
# do_choose_pdfs situation, #26); the merged state it would produce is
# checked against the reference through load_comp_list instead.
self.share_comps = share_comps
self.share_start = share_start
self.share_iter = share_iter
self.comp_thresh = comp_thresh
# Cached sphere pseudo-inverse (issues #223, #253): the sensor-space
# back-map, shared by get_sensor_mixing_matrix and the sharing metric.
self._sphere_pinv = None
# The A-freeze schedule reads share_start/share_iter whether or not
# share_comps is on (issue #345), so both are validated always.
schedule.validate_share_start(share_start)
schedule.validate_share_iter(share_iter)
if share_comps:
if not 0.0 < comp_thresh <= 1.0:
raise ValueError(f"comp_thresh must be in (0, 1], got {comp_thresh}")
self.do_mean = do_mean
self.do_sphere = do_sphere
self.do_approx_sphere = do_approx_sphere
# Explicit PCA reduction, validated by the policy shared with the NumPy
# and MLX backends (pamica/rank.py, issue #323) so a bad value fails
# here rather than as a silently wrongly sized or nan_ll fit; then one
# log line for any part of it a fit will ignore.
validate_pca_reduction(pcakeep, pcadb)
log_ignored_pca_request(pcakeep, pcadb, do_sphere)
self.pcakeep = pcakeep
self.pcadb = pcadb
self.seed = seed
# Best-of-N restarts (issue #198), a pamica extension: Fortran has no
# search over seeds. Resolved here so a bad configuration fails before
# any data is touched, and derived from the CONSTRUCTOR seed so that a
# second fit() on the same instance repeats the same seeds even though
# fit() leaves self.seed on the winning restart.
self._restart_seeds = restarts.resolve_seeds(n_restarts, restart_seeds, 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]] = []
if device is None:
device = setup_device()
# MPS has no float64, so an automatic MPS pick for a float64 run
# (the parity default) moves to the CPU instead of raising below.
# An explicit device="mps" is the caller's choice and still raises.
if device.type == "mps" and dtype == torch.float64:
logger.warning(
"AMICATorchNG: auto-selected MPS cannot represent float64 "
"(the parity default), so running on CPU. Pass "
"dtype=torch.float32 with device='mps' to run on MPS."
)
device = torch.device("cpu")
elif isinstance(device, str):
device = torch.device(device)
self.device = device
self.dtype = dtype
if self.device.type == "mps" and self.dtype == torch.float64:
raise ValueError(
"MPS does not support float64. Use dtype=torch.float32 for "
"device='mps', or device='cpu'/'cuda' for float64 parity runs."
)
self.iteration = 0
self.ll_history: list[float] = []
# Log-likelihood of the *returned* parameters (issue #51). With
# keep_best, ``ll_history`` stays the true per-iteration trajectory
# (which can include a late overshoot), while ``final_ll_`` is the LL of
# the iterate fit() actually kept -- use this, not ``ll_history[-1]``, as
# the model's fitted log-likelihood. Set by fit(). Exact for a fit that
# ended on a convergence stop, which exits before that iteration's
# update (issue #339); after max_iter it is the LL one update before
# the returned parameters, as in the reference (see _fit_once).
#
# Under share_comps, a merge that fires on the LAST fit iteration is
# reflected in the returned A/W/comp_list but NOT in final_ll_: the
# merge runs after that iteration's LL has already been computed and
# recorded (Fortran identify_shared_comps runs after the iteration's
# LL accumulation, amica15.f90:1856-1858 vs the earlier LL accumulation), so
# the merge's effect on the likelihood only shows up in the next
# iteration's E-step -- which never runs. keep_best is disabled under
# share_comps (see fit()), so this is not a keep_best artifact; it
# holds even with keep_best=False. Fortran-faithful, so this is
# documented behavior, not a bug (issue #269).
self.final_ll_: Optional[float] = None
# Mutual Information Reduction (MIR) waypoint trajectory (issue #137),
# populated by fit() when
# mir_step > 0: (iteration, mir_nats, variance) tuples from the
# CURRENT (mid-fit) W/sphere. Like ll_history, this is a true
# trajectory that a keep_best restore does NOT rewrite -- the
# fit-end MIR is mir() on the returned parameters, not
# mir_history_[-1]. Not part of state_dict(): it's a diagnostic,
# not a fitted parameter.
#
# Not index-aligned with ll_history: the entry for iteration i is
# computed AFTER that iteration's _update_parameters, while
# ll_history[i] is the likelihood of the parameters BEFORE it (the
# E-step accumulator that produced the update). The two therefore
# describe states one update apart, so zipping them by index compares
# different parameters (issue #161). An iteration that ends the fit on
# a stop takes no update and records no waypoint.
self.mir_history_: list[tuple[int, float, float]] = []
# Outlier-rejection bookkeeping (set up in fit()).
self.numrej = 0
self.good_idx: Optional[torch.Tensor] = None
# Set by fit(): why fitting stopped ("max_iter", "nan_ll", "lrate_floor",
# "grad_norm_floor", "min_dll", "grad_norm" -- issue #207 added the last
# three) and how many iterations reverted Newton to natural gradient
# (Fortran prints this; here it is exposed for parity debugging, see
# issue #21).
self.stop_reason: Optional[str] = None
self.n_newton_fallbacks = 0
# Weight-gradient-norm (Fortran ndtmpsum), recomputed every iteration by
# _update_direction and read by fit()'s convergence checks (issue #207).
# None before the first _update_direction call.
self._ndtmpsum: Optional[float] = None
# Populated by fit()/_initialize_parameters().
self.A: Optional[torch.Tensor] = None
self.W: Optional[torch.Tensor] = None
self.c: Optional[torch.Tensor] = None
self.mu: Optional[torch.Tensor] = None
self.alpha: Optional[torch.Tensor] = None
self.beta: Optional[torch.Tensor] = None
self.rho: Optional[torch.Tensor] = None
# Per-source density-family codes (n_channels, n_models); set in
# _initialize_parameters and mutated by the adaptive switcher.
self.pdtype: Optional[torch.Tensor] = None
# Number of adaptive-switch passes already performed (Fortran numchpdf).
self.n_kurt_done = 0
self.gm: Optional[torch.Tensor] = None
self.comp_list: Optional[torch.Tensor] = None
self.mean: Optional[torch.Tensor] = None
self.sphere: Optional[torch.Tensor] = None
self.sldet = 0.0
# Full-dataset per-sample/per-model log-likelihood (Fortran's LLt,
# issue #155), STASHED as the training E-step computes it rather than
# recomputed by a separate forward pass at write time (issue #157) --
# which is what Fortran does, keeping ``modloglik(num_models,N)`` and
# ``loglik(N)`` permanently allocated (amica15.f90:2617-2620) so that
# ``write_output`` just dumps them (amica15.f90:2338-2343).
#
# ``_llt_logv``/``_llt_ll`` are the live per-fit buffers (device
# tensors; Fortran's ``modloglik``/``loglik``), zero-filled so a
# ``do_reject`` sample keeps Fortran's zero sentinel. ``fit`` converts
# them into the compact numpy ``_llt_lht``/``_llt_lt`` that
# ``write_amica_output`` consumes, then drops the device buffers.
# Memory is ``(n_models + 1) * n_samples * itemsize`` and so grows with
# the data, but stays far below the sphered dataset already resident
# (``n_channels x n_samples``) since ``n_channels >> n_models + 1`` in
# any real fit: ~0.7 MB for the bundled 30504-sample two-model case.
# Not fitted parameters (absent from state_dict()/_PARAM_TENSORS): a
# model restored via from_state_dict() has none, so write_amica_output
# writes no LLt for it.
self._llt_logv: Optional[torch.Tensor] = None
self._llt_ll: Optional[torch.Tensor] = None
self._llt_lht: Optional[np.ndarray] = None
self._llt_lt: Optional[np.ndarray] = None
# ------------------------------------------------------------------
# Preprocessing
# ------------------------------------------------------------------
def _preprocess(self, X: np.ndarray) -> torch.Tensor:
"""Mean-removal + sphering, matching ``pamica.AMICA._preprocess_data``.
Done in float64 on CPU (eigh is not reliably supported on MPS and
this is a one-time O(n_channels^3) cost, not the per-block hot
path), then cast/moved to ``self.device``/``self.dtype``.
"""
X_cpu = torch.from_numpy(np.ascontiguousarray(X)).to(torch.float64)
data_dim = X_cpu.shape[0]
if self.do_mean:
mean = X_cpu.mean(dim=1, keepdim=True)
X_cpu = X_cpu - mean
else:
mean = torch.zeros(data_dim, 1, dtype=torch.float64)
if self.do_sphere:
# Population covariance (divide by N), matching Fortran's DSYRK
# scatter/N -- NOT torch.cov's default sample covariance (/(N-1)).
# The two differ by a pure scalar sqrt(N/(N-1)); using /(N-1) leaves
# a ~5e-6 sphere mismatch vs the reference (issue #24, check [1] of
# .context/issue-24/root_cause_Aupdate.py).
cov = torch.cov(X_cpu, correction=0)
evals, evecs = torch.linalg.eigh(cov)
order = torch.argsort(evals, descending=True)
evals = evals[order]
evecs = evecs[:, order]
# Numerical-rank detection (Fortran amica15.f90:413). The policy is
# shared with the NumPy and MLX backends so they cannot drift
# (pamica/rank.py); only the eigenvalues cross the boundary, as a
# read-only copy, so the sphere below stays bit-exact.
n_comp = numerical_rank(
evals.cpu().numpy(),
mineig=self.mineig,
mineig_rel=self.mineig_rel,
pcakeep=self.pcakeep,
pcadb=self.pcadb,
)
V = evecs[:, :n_comp]
inv_sqrt = torch.diag(1.0 / torch.sqrt(evals[:n_comp]))
if n_comp < data_dim:
# Rank-reduced sphere: (n_comp, data_dim), so the sphered data
# come out at the kept rank rather than staying rank-deficient
# at data_dim rows (Fortran nw = numeigs, amica15.f90:563).
# Vt rows are eigenvectors in descending-eigenvalue order,
# matching Fortran's reversed Stmp2 (amica15.f90:473-479).
w_pca = inv_sqrt @ V.T
if self.do_approx_sphere:
# Fortran amica15.f90:501-508 symmetrizes the reduced
# whitening by the orthogonal polar factor of the leading
# n_comp x n_comp block of V^T:
# B = (V^T)[:n, :n] = U_b S_b Vt_b
# S = (V_b U_b^T) @ w_pca
B = evecs.T[:n_comp, :n_comp]
U_b, _, Vt_b = torch.linalg.svd(B)
sphere = (Vt_b.T @ U_b.T) @ w_pca
else:
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). This is the
# Fortran default and the parity-validated form; the old
# diag(1/sqrt)@V^T (PCA whitening) is a different, non-symmetric
# transform that breaks activation parity.
sphere = V @ inv_sqrt @ V.T
else:
# Non-symmetric PCA whitening D^-1/2 V^T (Fortran
# do_approx_sphere=False path, amica17.f90:495).
sphere = inv_sqrt @ V.T
X_cpu = sphere @ X_cpu
# Sphering log-determinant term of the data log-likelihood
# (Fortran ``sldet``, amica17.f90:474): sum over the kept
# eigenvalues of -0.5*log(eval). For the PCA-reduced-rank case
# this is a pseudo-determinant, matching Fortran which sums over
# numeigs kept eigenvalues regardless of full rank.
sldet = float(-0.5 * torch.log(evals[:n_comp]).sum().item())
else:
sphere = torch.eye(data_dim, dtype=torch.float64)
sldet = 0.0
self.mean = mean.to(device=self.device, dtype=self.dtype)
self.sphere = sphere.to(device=self.device, dtype=self.dtype)
self.sldet = sldet
# Rank reduction shrank the sphered space, so size the model to the kept
# rank before _initialize_parameters allocates against n_channels
# (Fortran ``nw = numeigs``, amica15.f90:563). No-op, and therefore
# bit-exact, whenever the data are full rank.
n_kept = sphere.shape[0]
if n_kept != self.n_channels:
logger.info(
"Data covariance has numerical rank %d of %d; fitting %d "
"sources and mapping back to %d channels via the sphere "
"pseudo-inverse.",
n_kept,
data_dim,
n_kept,
data_dim,
)
self.n_channels = n_kept
self.n_comps = n_kept * self.n_models
self._sphere_pinv = None # rebuilt on demand for this fit's sphere
return X_cpu.to(device=self.device, dtype=self.dtype)
# ------------------------------------------------------------------
# Initialization
# ------------------------------------------------------------------
def _initialize_parameters(self):
"""Initialize parameters, mirroring ``pamica.AMICA._initialize_parameters``
exactly (same RNG draws, same order) so the same seed gives
bit-identical starting parameters to the NumPy reference.
"""
rng = np.random.RandomState(self.seed)
n, m, ncomp, nmix = self.n_channels, self.n_models, self.n_comps, self.n_mix
# One row per component (issue #334): model h's block is rows
# h*n..(h+1)*n-1. Drawn and normalized to unit-norm components as the
# reference does (issue #341, pamica.initialization).
A_np = initial_mixing(rng, n, m)
comp_list_np = np.zeros((n, m), dtype=np.int64)
for h in range(m):
comp_list_np[:, h] = np.arange(h * n, (h + 1) * n)
mu_np = np.zeros((nmix, ncomp), dtype=np.float64)
for k in range(ncomp):
mu_np[:, k] = np.linspace(-1, 1, nmix)
mu_np[:, k] += 0.05 * (1 - 2 * rng.rand(nmix))
alpha_np = np.ones((nmix, ncomp), dtype=np.float64) / nmix
beta_np = np.ones((nmix, ncomp), dtype=np.float64)
beta_np += 0.1 * (0.5 - rng.rand(nmix, ncomp))
rho_np = self.rho0 * np.ones((nmix, ncomp), dtype=np.float64)
gm_np = np.ones(m, dtype=np.float64) / m
c_np = np.zeros((n, m), dtype=np.float64)
self.A = torch.from_numpy(A_np).to(self.device, self.dtype)
self.comp_list = torch.from_numpy(comp_list_np).to(self.device)
self.mu = torch.from_numpy(mu_np).to(self.device, self.dtype)
self.alpha = torch.from_numpy(alpha_np).to(self.device, self.dtype)
self.beta = torch.from_numpy(beta_np).to(self.device, self.dtype)
self.rho = torch.from_numpy(rho_np).to(self.device, self.dtype)
self.gm = torch.from_numpy(gm_np).to(self.device, self.dtype)
self.c = torch.from_numpy(c_np).to(self.device, self.dtype)
# Per-source density-family codes, Fortran ``pdtype = pdftype`` (amica15.f90:
# 611). In adaptive mode (pdftype==1) every source starts as the
# super-Gaussian code 1 and the switcher may flip it to 4.
self.pdtype = torch.full(
(n, m), self.pdftype, dtype=torch.long, device=self.device
)
self.n_kurt_done = 0
# Reset the mutable optimization state to the pristine constructor
# values (lrate/lrate_cap, newtrate and rholrate/rholrate_cap are
# annealed or ratcheted down during fit; restore them so a re-fit
# starts fresh).
self.lrate = self.lrate0
self.lrate_cap = self.lrate0
self.newtrate = self.newtrate0
self.rholrate = self.rholrate0
self.rholrate_cap = self.rholrate0
self.iteration = 0
self._update_unmixing_matrices()
def _update_unmixing_matrices(self):
"""Recompute W from A via direct (batched) inversion -- never pinv.
Model ``h``'s block is the component rows ``A[comp_list[:, h], :]``.
"""
assert self.A is not None and self.comp_list is not None
A_stack = torch.stack(
[self.A[self.comp_list[:, h], :] for h in range(self.n_models)], dim=0
)
try:
W_stack = torch.linalg.inv(A_stack)
except torch.linalg.LinAlgError:
# Mid-loop invariant raise (PR #318 review): this is called from
# inside _update_parameters, after A/mu/beta/rho/alpha/gm/c have
# already been reassigned to this iterate's values -- so a
# singular A here would otherwise strand the instance holding an
# inconsistent mix of new params and stale self.W, with
# stop_reason still "max_iter", while the exception propagates
# uncaught through the single-restart fit() path (no try/except
# there). Set the degenerate marker before re-raising so every
# downstream state_dict()/write_amica_output() refusal check
# catches it regardless of whether the caller catches this.
# _fit_restarts's except block also sets this -- now
# redundant-but-harmless there, kept so that path does not
# depend on this site doing it correctly.
self.stop_reason = restarts.ERROR_STOP_REASON
raise
self.W = W_stack.permute(1, 2, 0).contiguous()
def _pdtype_h(self, h: int) -> Optional[torch.Tensor]:
"""Per-source density-family codes for model ``h``, shaped for
broadcasting against ``(batch, n_channels, num_mix)`` tensors, or
``None`` on the default ``pdftype=0`` (GG-only) fast path so the E-step
stays bit-identical to the pre-#26 implementation.
"""
if self.pdftype == 0:
return None
assert self.pdtype is not None
return self.pdtype[:, h].view(1, -1, 1)
# ------------------------------------------------------------------
# E-step / M-step sufficient statistics (the hot path)
# ------------------------------------------------------------------
def _forward(self, X: torch.Tensor):
"""Run the E-step forward pass for one data block.
Computes, for every model ``h``, the activations ``b``, scaled
activations ``y``, normalized mixture responsibilities ``z``, the power
``|y|^rho`` (reused by the rho update), and the per-sample per-model
log-likelihood ``logV`` (including the ``log|det W|`` and ``sldet``
Jacobian terms, matching Fortran's ``Ptmp`` seed, amica17.f90:1273).
Shared by ``_get_block_updates`` (which reduces it into sufficient
statistics) and ``_block_sample_ll`` (which only needs ``logV``).
Returns
-------
logV : torch.Tensor of shape (batch, n_models)
b_list, z_list, y_list, azrho_list : lists (one entry per model) of
per-model tensors (``b``: (batch, n_channels); ``z``/``y``/``azrho``:
(batch, n_channels, n_mix)).
"""
assert (
self.comp_list is not None
and self.c is not None
and self.W is not None
and self.mu is not None
and self.beta 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]
num_models = self.n_models
b_list, z_list, y_list, azrho_list = [], [], [], []
logV = torch.empty(batch_size, num_models, dtype=self.dtype, device=self.device)
for h in range(num_models):
idx = self.comp_list[:, h]
# Activation b = W(x - c): c is the per-model data-space center.
# Fortran subtracts wc in the E-step (amica17.f90:1280-1292), where
# wc = W@c is precomputed in get_unmixing_matrices (amica17.f90:2178).
# Subtracting c in data space before W is equivalent and keeps c's
# semantics identical to Fortran's. For n_models=1, c == 0, so this is
# bit-identical to the old X.T @ W.
b = (X - self.c[:, h].unsqueeze(1)).T @ self.W[:, :, h] # (batch, n_ch)
mu_h = self.mu[:, idx].T.unsqueeze(0) # (1, n_channels, num_mix)
beta_h = self.beta[:, idx].T.unsqueeze(0)
rho_h = self.rho[:, idx].T.unsqueeze(0)
alpha_h = self.alpha[:, idx].T.unsqueeze(0)
y = beta_h * (b.unsqueeze(-1) - mu_h) # (batch, n_channels, num_mix)
# Only log_pdf is needed here; the score fp (and drho's |y|^rho) are
# reused in _get_block_updates. az_rho = |y|^rho is threaded through
# so the rho-update does not recompute it (issue #63).
log_pdf, az_rho = _log_pdf_only(y, rho_h, self._pdtype_h(h))
# z0 = log(alpha) + log(beta) + log_pdf. For the single-component
# families (codes 1/4) n_mix==1 so alpha==1 and log(alpha)==0, which
# reproduces Fortran's alpha-free z0 (amica15.f90:1358/1370).
z0 = torch.log(alpha_h) + torch.log(beta_h) + log_pdf
ll_i = torch.logsumexp(
z0, dim=-1
) # (batch, n_channels) -- per-source log-density
z = torch.softmax(z0, dim=-1) # normalized responsibilities
logdet_W = torch.linalg.slogdet(self.W[:, :, h])[1]
logV[:, h] = (
torch.log(self.gm[h]) + logdet_W + self.sldet + ll_i.sum(dim=-1)
)
b_list.append(b)
z_list.append(z)
y_list.append(y)
azrho_list.append(az_rho)
return logV, b_list, z_list, y_list, azrho_list
def _block_sample_ll(self, X: torch.Tensor) -> torch.Tensor:
"""Per-sample total log-likelihood for a data block (the rejection
statistic; Fortran ``P``/``loglik``, amica17.f90:1372)."""
logV, *_ = self._forward(X)
return torch.logsumexp(logV, dim=1) # (batch,)
def _get_block_updates(self, X: torch.Tensor) -> Dict[str, torch.Tensor]:
"""Compute sufficient-statistic accumulators for one data block.
Fortran-faithful exact-EM statistics (amica17.f90:1437-1592), validated
against the reference binary to machine precision (issue #24). Unlike a
first-order gradient M-step, the mixture updates use exact-EM numerator/
denominator pairs and the score ``fp = rho*sign(y)*|y|^(rho-1)`` (``_score``,
Fortran ``fp``) rather than the density derivative ``dpdf``:
* ``dmu_n = sum(u*fp)``, ``dmu_d = sbeta*sum(u*fp/y)`` (mu += dmu_n/dmu_d)
* ``dbeta_n = sum(u)``, ``dbeta_d = sum(u*fp*y)`` (beta *= sqrt(n/d))
* ``drho_n = rho*sum(u*|y|^rho*ln|y|)`` (rho digamma update)
* ``dWtmp = g^T b`` with ``g = sum_j sbeta*u*fp`` (natural gradient)
where ``u = v*z`` (model x mixture responsibility). ``ll`` is the correct
pre-normalization ``logsumexp`` (see module docstring).
Assumes ``rho <= 2`` (the ``maxrho`` default); the ``rho > 2`` denominator
branches of Fortran (:1539/:1551) are unreachable and not implemented.
Returns
-------
updates : dict with ``dgm`` (n_models,), ``dalpha_n``/``dmu_n``/``dmu_d``/
``dbeta_n``/``dbeta_d``/``drho_n`` (n_mix, n_comps), ``dWtmp``
(n_channels, n_channels, n_models), ``dc_numer`` (n_channels,
n_models; the data-space bias numerator ``sum_t v_h*x``, issue #27),
``ll`` (scalar), and -- when ``do_newton`` -- ``dsigma2_numer``,
``dkappa_numer``, ``dlambda_numer`` (see ``_finalize_newton_stats``).
Plus two per-sample (non-summable) entries consumed and removed by
``_accumulate_blocks``: ``logV`` (batch, n_models) and
``ll_samples`` (batch,) -- Fortran's ``modloglik``/``loglik``
columns for this block, stashed for the LLt output (issue #157).
"""
assert (
self.comp_list is not None
and self.beta is not None
and self.rho is not None
)
num_mix, num_models = self.n_mix, self.n_models
dev, dt = self.device, self.dtype
logV, b_list, z_list, y_list, azrho_list = self._forward(X)
# Per-sample total log-likelihood (Fortran ``P``/``loglik``,
# amica15.f90:1402). Kept as a vector rather than folded straight into
# the scalar sum so the LLt stash can reuse it (issue #157); ``block_ll``
# is the same summation as before, bit for bit.
block_ll_samples = torch.logsumexp(logV, dim=1)
block_ll = block_ll_samples.sum()
v = torch.softmax(logV, dim=1) # (batch, num_models)
def zeros(*shape):
return torch.zeros(*shape, dtype=dt, device=dev)
dgm = zeros(num_models)
dalpha_n = zeros(num_mix, self.n_comps)
dmu_n = zeros(num_mix, self.n_comps)
dmu_d = zeros(num_mix, self.n_comps)
dbeta_n = zeros(num_mix, self.n_comps)
dbeta_d = zeros(num_mix, self.n_comps)
drho_n = zeros(num_mix, self.n_comps)
dWtmp = zeros(self.n_channels, self.n_channels, num_models)
dc_numer = zeros(self.n_channels, num_models)
do_newton = self.do_newton
if do_newton:
dsigma2_numer = zeros(self.n_channels, num_models)
dkappa_numer = zeros(num_mix, self.n_channels, num_models)
dlambda_numer = zeros(num_mix, self.n_channels, num_models)
tiny = torch.finfo(dt).tiny
for h in range(num_models):
idx = self.comp_list[:, h]
b, zr, y = b_list[h], z_list[h], y_list[h]
v_h = v[:, h]
beta_h = self.beta[:, idx].T.unsqueeze(0) # sbeta, (1, n_ch, num_mix)
rho_h = self.rho[:, idx].T # (n_ch, num_mix)
# score fp; the family select-case is amica15.f90:1467-1491 (amica17
# is GG-only, so cite the binary's source explicitly here).
fp = _score(y, rho_h.unsqueeze(0), self._pdtype_h(h))
u = v_h.unsqueeze(-1).unsqueeze(-1) * zr # u = v*z (:1439)
ufp = u * fp # (:1485)
dgm[h] = v_h.sum()
dalpha_n.index_add_(1, idx, u.sum(0).T) # sum(u) (:1524)
dmu_n.index_add_(1, idx, ufp.sum(0).T) # sum(ufp) (:1532)
# mu denominator sbeta*sum(ufp/y) (:1537). In float32 a sample sitting
# on a mixture mean can round y to *exactly* 0; the score fp(0)=0 (for
# the supported rho>=1), so the raw ufp/y is 0/0 = NaN -- the sole
# trigger of the full-data float32 divergence (issue #75; NOT a
# summation-precision problem, so compensated accumulation does not
# help). The true term ufp/y = u*rho*|y|^(rho-2) is NOT 0 in the limit:
# a nonzero constant at rho==2, and an integrable singularity that
# diverges as y->0 for rho<2 -- so once y underflows to exactly 0 the
# real contribution is unrepresentable. Substituting 0 (ufp==0 there,
# so 0/1) drops that one sample instead of poisoning all of dmu_d with a
# NaN: a bounded, empirically negligible bias (it fires <=1 sample per
# iteration on the sample EEG, and float32 still matches the float64 LL
# to ~5 sig digits). float64 never rounds y to exactly 0, so the guard
# is a bit-identical no-op there (single-model #24 parity preserved),
# and it needs no float64, so it also stabilizes the MPS/float32 path.
safe_y = torch.where(y == 0, torch.ones_like(y), y)
dmu_d.index_add_(
1, idx, (beta_h.squeeze(0) * (ufp / safe_y).sum(0)).T
) # (:1537)
dbeta_n.index_add_(1, idx, u.sum(0).T) # sum(u) (:1550)
dbeta_d.index_add_(1, idx, (ufp * y).sum(0).T) # sum(ufp*y) (:1556)
# drho_numer = rho * sum(u*|y|^rho*ln|y|) (:1560-1578). The leading
# rho comes from ln(|y|^rho)=rho*ln|y| in the Fortran logab chain
# (issue #24 Bug 1). Guard only the per-sample underflow (:1570) --
# no per-component (rho!=1&rho!=2) mask (Bug 2): |y|^rho*ln|y| is 0 at
# y=0, and clamping the log input makes the product collapse there.
ay = y.abs()
ayrho = azrho_list[h] # |y|^rho reused from _forward (issue #63)
logab = rho_h.unsqueeze(0) * torch.log(ay.clamp_min(tiny)) # rho*ln|y|
logab = torch.where(ayrho < EPSDBLE, torch.zeros_like(logab), logab)
drho_n.index_add_(1, idx, (u * (ayrho * logab)).sum(0).T)
g = (beta_h * ufp).sum(-1) # g_i = sum_j sbeta*ufp (:1493)
dWtmp[:, :, h] = g.T @ b # source-space sum g_t b_t^T (:1592)
# Data-space bias accumulator: dc_numer[i,h] = sum_t v_h(t)*x(i,t)
# (Fortran :1423-1429). The denominator is dgm[h] = sum_t v_h(t).
# NOTE: this replaces the old gradient-style bias g.sum(0), which was
# accumulated but never applied (c was frozen at 0); the Fortran
# update is the data-space responsibility-weighted mean (issue #27).
dc_numer[:, h] = X @ v_h
if do_newton:
# Newton curvature accumulators (Fortran amica17.f90:1419,
# 1500-1514), in terms of the score fp (not dpdf).
dsigma2_numer[:, h] = (v_h.unsqueeze(-1) * b.pow(2)).sum(0) # (:1419)
dkappa_numer[:, :, h] = (
(u * fp.pow(2)).sum(0) * beta_h.squeeze(0).pow(2)
).T # (:1500)
dlambda_numer[:, :, h] = (u * (fp * y - 1.0).pow(2)).sum(0).T # (:1511)
updates = {
"dgm": dgm,
"dalpha_n": dalpha_n,
"dmu_n": dmu_n,
"dmu_d": dmu_d,
"dbeta_n": dbeta_n,
"dbeta_d": dbeta_d,
"drho_n": drho_n,
"dWtmp": dWtmp,
"dc_numer": dc_numer,
"ll": block_ll,
# Per-sample E-step outputs for the LLt stash (issue #157). These
# are NOT summable accumulators -- _accumulate_blocks pops them
# before folding the rest -- and cost nothing extra: both are
# already computed above for ``ll``/``v``.
"logV": logV,
"ll_samples": block_ll_samples,
}
if do_newton:
updates["dsigma2_numer"] = dsigma2_numer
updates["dkappa_numer"] = dkappa_numer
updates["dlambda_numer"] = dlambda_numer
return updates
def _accumulate_blocks(
self, X: torch.Tensor, stash_llt: bool = False
) -> Dict[str, torch.Tensor]:
"""Sum sufficient statistics over all blocks of ``X``.
Peak memory scales with ``block_size`` (each block's intermediates
are freed once accumulated), not with ``X.shape[1]``.
Parameters
----------
X : torch.Tensor
The (sphered) data to accumulate over, already restricted to the
good set under ``do_reject``.
stash_llt : bool, default=False
Scatter each block's per-sample ``logV``/``ll_samples`` into the
``_llt_logv``/``_llt_ll`` buffers as it goes, so the LLt output
never needs a second pass over the data (issue #157). Only the
training loop sets this; the ``_tune_block_size`` probes leave the
buffers untouched, so the tuner still leaves no state behind. The
per-sample values are dropped from the returned dict either way --
they are per-block quantities, not accumulators, and summing them
would be meaningless.
"""
n_samples = X.shape[1]
acc: Optional[Dict[str, torch.Tensor]] = None
for start in range(0, n_samples, self.block_size):
end = min(start + self.block_size, n_samples)
block_acc = self._get_block_updates(X[:, start:end])
logv = block_acc.pop("logV")
ll_samples = block_acc.pop("ll_samples")
if stash_llt:
assert self._llt_logv is not None and self._llt_ll is not None
# Map this block's rows back onto the full-dataset index. Under
# do_reject the caller passed X_t[:, good_idx], so block
# [start:end] of X is good_idx[start:end] of the dataset;
# otherwise the block index is the sample index.
rows = (
self.good_idx[start:end]
if self.good_idx is not None
else slice(start, end)
)
self._llt_logv[rows] = logv
self._llt_ll[rows] = ll_samples
if acc is None:
acc = block_acc
else:
for key in acc:
acc[key] = acc[key] + block_acc[key]
assert acc is not None
return acc
def _available_memory_bytes(self) -> Optional[int]:
"""Memory the current device reports as usable, for the search's cap.
``None`` (no cap) when the device cannot report it; the search then
relies solely on catching the allocation failure.
The three branches do not report the same quantity. CUDA's
``mem_get_info`` gives currently-FREE memory, while MPS's
``recommended_max_memory`` and the host branch give total CAPACITY --
neither accounts for what is already allocated. The cap is therefore an
upper bound on what the device could ever give, not on what is free
right now, which is why it is only ever a first filter:
:data:`~pamica.blocktune.MEMORY_BUDGET_FRACTION` keeps it conservative
and catching the real allocation failure is what actually makes the
search safe.
"""
dev = self.device.type
if dev == "cuda":
try:
return int(torch.cuda.mem_get_info(self.device)[0])
except (RuntimeError, AttributeError) as exc:
logger.debug(
"could not query CUDA memory (%s: %s); block-size search "
"runs without a memory cap",
type(exc).__name__,
exc,
)
return None
if dev == "mps":
try:
# Capacity, not free memory (see the docstring).
return int(torch.mps.recommended_max_memory())
except (RuntimeError, AttributeError) as exc:
logger.debug(
"could not query MPS memory (%s: %s); block-size search "
"runs without a memory cap",
type(exc).__name__,
exc,
)
return None
# Total host RAM, likewise capacity rather than free.
return blocktune.host_memory_bytes()
def _tune_block_size(self, X: torch.Tensor) -> None:
"""Set ``self.block_size`` to the fastest timed candidate (issue #232).
The probe is one ``_accumulate_blocks`` 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. ``_accumulate_blocks`` only reads
model state and consumes no RNG, and ``block_size`` is restored around
every probe, so the fit that follows is bit-identical to one started
directly at the chosen size (``test_post_tune_fit_is_bit_identical``).
"""
saved = self.block_size
def probe(size: int) -> float:
self.block_size = size
try:
start = time.perf_counter()
acc = self._accumulate_blocks(X)
# Force completion before stopping the clock: CUDA/MPS queue
# work asynchronously, so reading a result is what makes the
# elapsed time mean anything. Cheap and correct on CPU too.
float(acc["ll"])
return time.perf_counter() - start
finally:
# Never leave the model holding a candidate -- least of all one
# that just failed to allocate (issue #232).
self.block_size = saved
self.block_size = blocktune.search(
probe=probe,
fallback=saved,
blk_min=self.blk_min,
blk_max=self.blk_max,
blk_step=self.blk_step,
n_samples=int(X.shape[1]),
n_channels=self.n_channels,
n_mix=self.n_mix,
n_models=self.n_models,
itemsize=torch.finfo(self.dtype).bits // 8,
available_bytes=self._available_memory_bytes(),
log=logger,
)
# ------------------------------------------------------------------
# M-step parameter update
# ------------------------------------------------------------------
def _finalize_newton_stats(self, acc: Dict[str, torch.Tensor]):
"""Reduce the Newton block accumulators into ``(sigma2, lambda, kappa)``.
Ports the Fortran finalization (amica17.f90:1762-1776). The Fortran
``baralpha``/``dkappa_denom``/``dlambda_denom`` responsibility masses
all cancel algebraically against the per-mixture ``dalpha`` weighting,
leaving simply (with ``dgm = sum_t v_h`` the raw model mass):
sigma2[i,h] = dsigma2_numer[i,h] / dgm[h]
kappa[i,h] = sum_j dkappa_numer[j,i,h] / dgm[h]
lambda[i,h] = sum_j (dlambda_numer[j,i,h]
+ dkappa_numer[j,i,h] * mu[j,comp(i,h)]^2) / dgm[h]
Returns (sigma2, lambda_, kappa), each (n_channels, n_models).
"""
assert self.mu is not None and self.comp_list is not None
dgm = acc["dgm"].unsqueeze(0) # (1, n_models)
sigma2 = acc["dsigma2_numer"] / dgm
kappa = acc["dkappa_numer"].sum(dim=0) / dgm
# mu at each source's component: mu[j, comp_list[i,h]] -> (n_mix, n_ch, n_models)
mu_at = self.mu[:, self.comp_list]
lambda_ = (acc["dlambda_numer"] + acc["dkappa_numer"] * mu_at.pow(2)).sum(
dim=0
) / dgm
return sigma2, lambda_, kappa
def _newton_direction(self, dA_h, sigma2_h, lambda_h, kappa_h):
"""Per-model Newton direction ``H`` from the natural gradient ``dA_h``.
Vectorized port of the per-source-pair 2x2 solve (amica17.f90:1817-1832;
the NumPy backend's loop in ``AMICA_NumPy._update_parameters``):
H[i,i] = dA_h[i,i] / lambda[i]
sk1 = sigma2[i]*kappa[k]; sk2 = sigma2[k]*kappa[i] (i != k)
H[i,k] = (sk1*dA_h[i,k] - dA_h[k,i]) / (sk1*sk2 - 1) if sk1*sk2 > 1
Returns ``(H, posdef)``. ``posdef`` is False if any off-diagonal pair
fails ``sk1*sk2 > 1`` (the positive-definiteness guard); the caller
then falls back to the natural gradient for this model.
"""
n = self.n_channels
sk1 = sigma2_h.unsqueeze(1) * kappa_h.unsqueeze(0) # [i,k] = sigma2[i]*kappa[k]
sk2 = sigma2_h.unsqueeze(0) * kappa_h.unsqueeze(1) # [i,k] = sigma2[k]*kappa[i]
prod = sk1 * sk2
valid = prod > 1.0
denom = torch.where(valid, prod - 1.0, torch.ones_like(prod))
h_off = (sk1 * dA_h - dA_h.T) / denom
H = torch.where(valid, h_off, torch.zeros_like(h_off))
# Diagonal overrides (uses lambda, not the off-diagonal formula).
diag = torch.diagonal(dA_h) / lambda_h
H = H - torch.diag(torch.diagonal(H)) + torch.diag(diag)
# Positive-definite iff every off-diagonal pair passed the guard.
offdiag = ~torch.eye(n, dtype=torch.bool, device=dA_h.device)
posdef = bool(valid[offdiag].all().item())
return H, posdef
def _update_direction(self, acc: Dict[str, torch.Tensor]) -> _UpdateStep:
"""The natural-gradient or Newton step for ``A`` from this iteration's
sufficient statistics, and its norm, without changing any 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; :meth:`fit` calls this first and
hands the result to :meth:`_update_parameters` once the checks have
run. Sets ``self._ndtmpsum``, the norm those checks read.
Everything here reads the parameters as the E-step saw them: the Newton
curvature folds in the pre-update ``mu`` (as the reference does during
accumulation, amica15.f90:1666-1680), and ``dAk`` weights the models by
the pre-update ``gm`` (the reference does not reassign ``gm`` until
``update_params``, :1788; issue #219).
"""
assert (
self.mu is not None
and self.A is not None
and self.comp_list is not None
and self.gm is not None
)
# Finalize the Newton curvature with the PRE-update mu. Fortran folds the
# mu^2 term into lambda during E-step accumulation, before the M-step
# moves mu (amica17.f90:1762-1774), and the NumPy port bakes it in at
# accumulation time. Nothing has moved mu yet when this runs.
newton_active = schedule.newton_active(
self.do_newton, self.iteration, self.newt_start
)
if newton_active:
sigma2, lambda_, kappa = self._finalize_newton_stats(acc)
# A is stored as Fortran's A^T (one component per row, issue #334; the
# true unmixing is W^T = inv(block)^T), so Fortran's per-model
# A_fort(:, comp_list(:,h)) @ dir becomes, transposed,
# dir^T @ A[comp_list[:, h], :] (LEFT-multiply by the TRANSPOSED
# direction). The direction ``dir`` (natural gradient I - <g b^T>/dgm,
# or its Newton precondition) is built in Fortran's untransposed
# convention. Getting this wrong (right-multiply by the untransposed
# dir) is invisible at the fixed point but sends the free-running fit
# downhill -- issue #24 root cause
# (.context/issue-24/root_cause_Aupdate.py, machine-exact check).
#
# Computed every iteration, including one on which A is held: Fortran
# computes dAk and ndtmpsum in accum_updates_and_likelihood
# (amica15.f90:1749-1761), before and apart from the guarded A step in
# update_params (:1803), and the grad-norm stop needs the true gradient
# magnitude every iteration (issue #207).
eye = torch.eye(self.n_channels, dtype=self.dtype, device=self.device)
directions = []
no_newt = False
for h in range(self.n_models):
dA_h = -acc["dWtmp"][:, :, h] / acc["dgm"][h] + eye # I - <g b^T>/dgm
if newton_active:
H, posdef = self._newton_direction(
dA_h, sigma2[:, h], lambda_[:, h], kappa[:, h]
)
if posdef:
directions.append(H)
else:
no_newt = True
directions.append(dA_h) # fall back to natural gradient
else:
directions.append(dA_h)
# Accumulate each model's natural-gradient/Newton contribution per
# COMPONENT (a row of A) as a gm-WEIGHTED AVERAGE (Fortran dAk/zeta,
# amica15.f90:1749-1761): dAk = sum_h gm[h]*dir_h^T@block_h scattered
# into the rows comp_list names, zeta = sum_h gm[h] per component, then
# dAk /= zeta. For the default disjoint comp_list every component has
# exactly one contributor, so gm cancels (dAk = dir) and single-model
# (gm=[1]) is byte-for-byte unchanged; for a SHARED component (issue
# #60) the step is Fortran's responsibility-weighted average of the
# models' steps for that one mixing vector, NOT a raw sum (a raw sum
# would over-step by the contributor count and destabilize the fit). A
# merged-away component has no contributor, so its zeta is 0 and its
# dAk is 0/tiny = 0: its row takes no step. gm is still the pre-update
# weight here (see the docstring).
dAk = torch.zeros_like(self.A)
zeta = torch.zeros(self.n_comps, dtype=self.dtype, device=self.device)
for h in range(self.n_models):
idx = self.comp_list[:, h]
dAk.index_add_(0, idx, self.gm[h] * (directions[h].T @ self.A[idx, :]))
zeta.index_add_(0, idx, self.gm[h].expand(idx.shape[0]))
dAk = dAk / zeta.clamp_min(torch.finfo(self.dtype).tiny).unsqueeze(1)
# Weight-gradient norm (Fortran ndtmpsum, amica15.f90:1760-1761):
# ``sqrt(sum(dAk**2, mask=comp_used) / (nw*count(comp_used)))``, one
# squared norm per component row, as the reference sums each
# component's column (``nd(iter,:) = sum(dAk*dAk,1)``). Read by fit()'s
# convergence checks (issue #207); the comp_used mask matters only when
# share_comps has merged components away (all-True otherwise, so
# ``comp_used_mask`` covers every component and this is a plain RMS over
# dAk).
comp_used_mask = self.comp_used
n_used = int(comp_used_mask.sum().item())
nd = (dAk**2).sum(dim=1) # (n_comps,)
self._ndtmpsum = float(
torch.sqrt(
nd[comp_used_mask].sum() / (self.n_channels * max(n_used, 1))
).item()
)
return _UpdateStep(dAk, newton_active, no_newt)
def _update_parameters(
self,
acc: Dict[str, torch.Tensor],
n_samples: int,
step: Optional[_UpdateStep] = None,
):
"""Apply the M-step parameter update, matching
``pamica.AMICA._update_parameters`` (natural-gradient and Newton).
``step`` is the mixing-matrix step :meth:`_update_direction` computed
from the same ``acc``: :meth:`fit` passes the one its stopping checks
already read, so nothing is computed twice. A direct call may omit it,
and the step is then computed here first, from the parameters as they
stand.
``n_samples`` is the number of samples that fed the accumulators (the
good-sample count when ``do_reject`` is active), so ``gm`` and the
reported log-likelihood are normalized by the effective sample count.
The mixture parameters use exact-EM fixed-point updates (no ``lrate``);
only the ``A``/``W`` step is scaled by ``lrate`` (Fortran amica17.f90:
1890-2035). The per-model data-space bias ``c`` uses Fortran's exact-EM
``update_c`` (amica17.f90:1423-1429/1899-1901): ``c[i,h] = sum_t v_h*x /
sum_t v_h``, the responsibility-weighted data mean for model ``h``. For
``n_models=1`` on mean-removed data ``v == 1`` so ``c`` collapses to the
(zero) data mean; the update is skipped there so single-model parity stays
bit-exact (issue #24). For ``n_models>1`` the per-model ``v`` is
non-uniform and ``c`` moves each iteration (issue #27).
"""
assert (
self.c is not None
and self.alpha is not None
and self.mu is not None
and self.beta is not None
and self.rho is not None
and self.A is not None
and self.comp_list is not None
)
if step is None:
step = self._update_direction(acc)
# The step was built with the pre-update gm (_update_direction), so gm
# can be reassigned now.
self.gm = acc["dgm"] / n_samples
# Per-model data-space bias (Fortran's `update_c` flag, amica17.f90:1423-
# 1429 numerator / :1899-1901 division). 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, perturbing the
# otherwise-exact single-model trajectory. dgm[h] = sum_t v_h(t) is the
# denominator (Fortran `dc_denom`). A fully-dead model (dgm[h]==0 => v_h==0
# for all t, so dc_numer[:,h]==0 too) gives 0/0; keep its PRIOR c rather
# than write a NaN. A NaN c would poison the NEXT iteration's cross-model
# softmax for EVERY model (unlike log(gm[h])=-inf, which softmax tolerates,
# so a dead model was previously inert) -- this containment mirrors the
# mu/beta/rho non-finite guards below. `dgm>0` is also False for a NaN dgm
# from upstream corruption, so that is contained too.
if self.n_models > 1:
dgm = acc["dgm"]
live = dgm > 0.0
new_c = acc["dc_numer"] / dgm.clamp_min(torch.finfo(self.dtype).tiny)
self.c = torch.where(live.unsqueeze(0), new_c, self.c)
if not bool(live.all()):
logger.warning(
"Zero-responsibility model(s) at iter %d; kept their prior "
"bias c (dead-model guard).",
self.iteration,
)
# Component sharing (#60): a component that was merged away is no longer
# referenced by comp_list, so no sufficient statistic accumulates into
# its density column (dalpha_n/dmu_d/dbeta_d == 0) and the divisions
# below would be 0/0 = NaN. Update only USED components and freeze the
# rest at their last finite value (Fortran carries NaN there harmlessly
# behind its comp_used mask; we keep them finite so save/the degenerate
# guard are not tripped). With the default full comp_list every
# component is used, so ``used`` is all-True and every update below is
# byte-for-byte unchanged.
used = self.comp_used.unsqueeze(0) # (1, n_comps)
self.alpha = torch.where(
used, acc["dalpha_n"] / acc["dalpha_n"].sum(dim=0, keepdim=True), self.alpha
)
# The A branch, where the reference has it: after gm/alpha/c and before
# mu/sbeta/rho (amica15.f90:1803-1816). On an iteration the reference
# holds A (:func:`pamica.schedule.share_freeze`), everything inside the
# branch is skipped together: the Newton-fallback bookkeeping (so a
# discarded Newton direction cannot pollute the fallback counter), the
# lrate ramp, the reset of the working rho rate to its ceiling, and the
# step itself. The step was built by _update_direction 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 "Hessian not positive definite, using natural
# gradient" (amica15.f90:1809-1811). Surface the same signal so an
# all-fallback run (issue #21) is visible without re-instrumenting.
self.n_newton_fallbacks += 1
logger.warning(
"Newton not positive definite at iter %d; using natural gradient.",
self.iteration,
)
# Learning-rate ramp: toward newtrate while Newton is active and
# stable, otherwise toward the lrate ceiling (Fortran
# amica15.f90:1804-1815). It starts from this iteration's lrate, which
# a likelihood decrease has already halved (fit runs the response
# before this 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.lrate_cap, self.lrate + min(1.0 / self.newt_ramp, self.lrate)
)
# The rho update below uses the ceiling (``rholrate = rholrate0``,
# :1806/:1813), so a decrease's per-iteration scaling of the working
# rho rate reaches the rho update 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). No lrate.
# ``used`` masks merged-away components (no-op for the default comp_list).
self.mu = torch.where(used, self.mu + acc["dmu_n"] / acc["dmu_d"], self.mu)
self.beta = torch.where(
used,
torch.clamp(
self.beta * torch.sqrt(acc["dbeta_n"] / acc["dbeta_d"]),
self.invsigmin,
self.invsigmax,
),
self.beta,
)
# The exact-EM mu/beta divisions are unguarded (matching Fortran, whose own
# mu/beta guard is commented out; it keeps a "NaN in sbeta!" canary,
# amica17.f90:1996-2000). fit() checks every parameter right after this
# update and stops, naming the offenders, on the iteration a value goes
# non-finite (stop_reason "nan_params"), so the failure is attributed
# where it happens rather than at a later nan-LL stop.
# 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
# here (e.g. from upstream mu/beta corruption) is reset to rho0 -- but
# logged first, so the reset does not silently erase the failure origin.
# Skipped for every non-GG family: Fortran sets dorho=.false. when
# pdftype/=0 (amica15.f90:3704), freezing rho at rho0.
if (
self.dorho
and not torch.all(self.rho == 1.0)
and not torch.all(self.rho == 2.0)
):
drho = acc["drho_n"] / acc["dalpha_n"].clamp_min(1e-8)
psi = torch.special.digamma(1.0 + 1.0 / self.rho)
new_rho = self.rho + self.rholrate * (1.0 - (self.rho / psi) * drho)
nan_mask = torch.isnan(new_rho)
if nan_mask.any():
logger.warning(
"NaN in rho update at iter %d for %d component(s); resetting "
"to rho0=%g.",
self.iteration,
int(nan_mask.sum()),
self.rho0,
)
new_rho = torch.where(
nan_mask, torch.full_like(new_rho, self.rho0), new_rho
)
self.rho = torch.where(
used, torch.clamp(new_rho, self.minrho, self.maxrho), self.rho
)
# 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.iteration, self.scalestep):
self._rescale_components()
self._update_unmixing_matrices()
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.
Normalizing anything but a component's mixing vector is not a change of
scale and perturbs the fit (issue #333, ADR 0006). A zero-norm row is
left untouched, as in the reference (``Anrmk > 0``), and so is a NaN
norm (``NaN > 0`` is false).
Each component is rescaled exactly once, including one that
``share_comps`` merged into several models. A merged-away row, which no
model reads, is left untouched: the reference's is NaN by then (its
``dAk/zeta`` is ``0/0``), so its own ``Anrmk > 0`` skips it too. The
norms are taken from the per-model blocks ``A[comp_list[:, h], :]``, the
same ``n x n`` gathers the per-block rule of issue #333 used, so an
unshared ``comp_list`` rescales bit for bit as before.
"""
assert (
self.A is not None
and self.mu is not None
and self.beta is not None
and self.comp_list is not None
)
scale = torch.ones(self.n_comps, dtype=self.dtype, device=self.device)
for h in range(self.n_models):
idx = self.comp_list[:, h]
norm = torch.sqrt((self.A[idx, :] ** 2).sum(dim=1)) # (n_channels,)
# A shared row gets the same norm from every block that holds it.
scale[idx] = torch.where(norm > 0, norm, torch.ones_like(norm))
self.A = self.A / scale.unsqueeze(1)
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_iter) <= 5``, counted from 1
(amica15.f90:1803, :func:`pamica.schedule.share_freeze`).
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_iter = 100``, every fit of 100 or more iterations
holds A on iterations 100-105, 200-205, and so on. The reference's own
scan never merges (see :meth:`_identify_shared_comps`), so this
unconditional schedule is the only freeze it ever shows. The
constructor requires ``share_iter >= 7``, since a shorter cycle would
hold A for good (:func:`pamica.schedule.validate_share_iter`).
"""
return schedule.share_freeze(self.iteration, self.share_start, self.share_iter)
def _identify_shared_comps(self) -> None:
"""Merge near-collinear components across models (Fortran
``identify_shared_comps``, amica15.f90:1916).
Two components (model ``h`` source ``i`` and model ``hh`` source ``ii``,
``h < hh``) are identified when the angle between their mixing vectors,
measured in the original (de-sphered) data space, is below the
``comp_thresh`` cutoff::
t0 = |a . b| / (||a|| ||b||), a = Spinv A[ci, :], b = Spinv A[cj, :]
with ``ci = comp_list[i, h]`` and ``cj = comp_list[ii, hh]``. Row ``k`` of
``A`` is component ``k``'s mixing vector (issue #334, ADR 0007), the
reference's column ``A(:, k)``, and ``Spinv = pinv(sphere)`` de-spheres
it back to input-channel (sensor) space, so ``a`` is exactly column
``i`` of ``get_sensor_mixing_matrix(h)``: the similarity compares the
two components' scalp maps. (Before issue #334 the scan compared stored
COLUMNS of a component-column layout, which are not components.) The
reference weights its inner products with ``Spinv2``, which is this
cosine for ``Spinv2 = Spinv^T Spinv``.
The pseudo-inverse -- not a true inverse -- is the faithful back-map:
the reference carries exactly this, ``Spinv(nx, numeigs)``, whenever
rank/PCA reduction is active (amica15.f90:568-578). Invertibility was
never a mathematical requirement of the merge metric, only of the way it
used to be computed (issue #253). Two consequences:
* Full rank, square sphere: ``pinv == inv`` to ~1e-15, orders of
magnitude below any ``comp_thresh`` (~0.99) decision boundary, so
merge decisions on well-conditioned data are unchanged.
* Rank-reduced sphere ``(n_kept, n_channels)`` (issue #223), or a square
sphere fitted on rank-deficient data (Maxwell-filtered MEG,
average-referenced EEG): ``pinv`` maps each component into the
retained sensor subspace instead of failing. For the reduced PCA
sphere ``S = D^-1/2 V_r^T`` this is ``pinv(S) = V_r D^1/2``, i.e. a
de-whitening followed by the orthonormal ``U_r = V_r`` embedding
proposed in issue #221; the embedding leaves the cosine untouched, so
this is the same comparison the full-rank path makes, evaluated in the
subspace the data actually occupy.
On a match, ``cj`` is folded into ``ci`` as the reference folds it:
every ``comp_list`` entry equal to ``cj`` is reassigned to ``ci``, so the
two sources now share component ``ci``'s mixing vector (row ``ci``) and
density, and nothing is copied or averaged at the merge itself. Row
``cj`` is retired: no model reads it, the A-update gives it no step and
its density is frozen (see :meth:`_update_parameters`). The M-step
accumulates every sufficient statistic through ``comp_list`` via
index_add, so a shared component's statistics sum automatically.
Greedy and order-dependent, matching the reference's quadruple loop.
Skips a pair already merged, or one whose two components coexist in some
single model (a model cannot share a component with itself).
No bit-exact oracle for the metric: ``Spinv2`` is *declared* in the
reference headers but never *allocated* anywhere in
``amica15.f90``/``amica17.f90`` (unlike ``Spinv``, allocated at
amica15.f90:569), so the pinned binary's scan reads an unallocated
array through ``DGEMV``.
It does not crash: every similarity comes out NaN, so it never merges,
at any ``comp_thresh`` (measured, epic #324 Phase 8; ``comp_list`` is
unchanged even at ``comp_thresh=0``). The merged STATE does have one: the
reference's ``load_comp_list`` seeds a merged ``comp_list`` directly,
and the update from it matches this backend to float64 round-off
(``pamica/tests/test_component_rows.py``).
"""
if self.n_models < 2:
return
assert self.A is not None and self.comp_list is not None
# _pinv_sphere raises on a non-finite sphere, so the metric below can
# only be garbage if A itself is (guarded per-pair in the scan).
# De-sphered component mixing vectors in sensor space (column k is
# component k), on CPU for the small greedy scan (n_models^2 *
# n_channels^2 pairs; avoids per-element GPU syncs).
atil = self._component_sensor_maps()
norms = np.linalg.norm(atil, axis=0)
cl = self.comp_list.detach().cpu().numpy().copy() # (nw, n_models)
nw, m = cl.shape
tiny = np.finfo(atil.dtype).tiny
merged = 0
for h in range(m):
for hh in range(h + 1, m):
for i in range(nw):
for ii in range(nw):
ci, cj = int(cl[i, h]), int(cl[ii, hh])
if ci == cj:
continue
t0 = abs(atil[:, ci] @ atil[:, cj]) / (
norms[ci] * norms[cj] + tiny
)
# NaN t0 (e.g. a zero-norm component) must NOT merge:
# `NaN < thresh` is False, so guard finiteness explicitly.
if not np.isfinite(t0) or t0 < self.comp_thresh:
continue
# A model cannot share a component with itself: skip if
# any single model already uses both components.
if any(
(cl[:, k] == ci).any() and (cl[:, k] == cj).any()
for k in range(m)
):
continue
cl[cl == cj] = ci # fold cj into ci everywhere
merged += 1
if merged:
self.comp_list = torch.from_numpy(cl).to(self.comp_list.device)
logger.info(
"Component sharing (iter %d): %d merge(s), %d unique components.",
self.iteration,
merged,
int(np.unique(cl).size),
)
def _component_sensor_maps(self) -> np.ndarray:
"""Every component's mixing vector in input-channel (sensor) space.
``pinv(sphere) @ A.T`` as a NumPy array of shape ``(n_channels_in,
n_comps)``: column ``comp_list[i, h]`` is column ``i`` of
:meth:`get_sensor_mixing_matrix` for model ``h`` (issue #334). These
are the vectors the share metric compares. Raises like
:meth:`_pinv_sphere` on a non-finite sphere.
"""
assert self.A is not None
return (self._pinv_sphere() @ self.A.T).detach().cpu().numpy()
def _pinv_sphere(self) -> torch.Tensor:
"""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 per fit in
:meth:`_preprocess` (and on :meth:`_load_params`), so it can never
describe a sphere other than the current one.
"""
assert self.sphere is not None
if self._sphere_pinv is None:
if not torch.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.
# Mid-loop invariant raise (PR #318 review): _pinv_sphere is
# called from _identify_shared_comps, itself called from
# inside _fit_once's iteration body under share_comps -- so
# this can fire with the instance already holding this
# iterate's updated A/mu/etc, mid-fit, propagating uncaught
# through the single-restart fit() path. Set the degenerate
# marker before raising, same reasoning as
# _update_unmixing_matrices above.
self.stop_reason = restarts.ERROR_STOP_REASON
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 = torch.linalg.pinv(self.sphere)
return self._sphere_pinv
@property
def comp_used(self) -> torch.Tensor:
"""Boolean mask (n_comps,) of components still referenced by comp_list.
A component drops out of use when it is folded into another by
:meth:`_identify_shared_comps`; an unused component (its row of ``A``
and its density columns) receives no update and is never read by the
E-step. Derived from ``comp_list`` (not stored).
Raises rather than asserts: ``assert`` is stripped under ``python -O``,
which would turn a pre-fit read into an obscure ``NoneType`` error
instead of the same message every other accessor gives.
"""
if self.comp_list is None:
raise RuntimeError(
"AMICATorchNG.comp_used requires a fitted model; call fit() first."
)
used = torch.zeros(self.n_comps, dtype=torch.bool, device=self.comp_list.device)
used[self.comp_list.reshape(-1)] = True
return used
def _choose_pdfs(self, X: torch.Tensor) -> None:
"""Extended-Infomax adaptive PDF switch (Fortran ``do_choose_pdfs``).
Re-estimates each source's kurtosis from the current model activations
and sets its density family to the super-Gaussian (code 1) or
sub-Gaussian (code 4) cosh density by kurtosis sign. This is the
extended-Infomax rule that ``runamica15.m`` documents for the
``kurt_start``/``num_kurt``/``kurt_int`` schedule (the super/sub-Gaussian
scores ``y +/- tanh(y)`` are exactly the two families 1/4). The
reference binary declares this (``pdftype==1`` sets ``do_choose_pdfs``,
amica15.f90:612) but never runs the switch (``m2sum``/``m4sum`` are
never accumulated), so there is no bit-exact oracle; validated by
real-data log-likelihood (must not decrease vs the fixed GG default).
"""
n_ch, n_models = self.n_channels, self.n_models
m2 = torch.zeros(n_ch, n_models, dtype=self.dtype, device=self.device)
m4 = torch.zeros_like(m2)
nsub = torch.zeros(n_models, dtype=self.dtype, device=self.device)
n_samples = X.shape[1]
for start in range(0, n_samples, self.block_size):
block = X[:, start : start + self.block_size]
logV, b_list, *_ = self._forward(block)
v = torch.softmax(logV, dim=1) # (batch, n_models)
for h in range(n_models):
b = b_list[h] # (batch, n_ch)
vh = v[:, h].unsqueeze(1)
m2[:, h] += (vh * b.pow(2)).sum(0)
m4[:, h] += (vh * b.pow(4)).sum(0)
nsub[h] += v[:, h].sum()
# Kurtosis = E[b^4]/E[b^2]^2 - 3 = nsub * m4 / m2^2 - 3, per (source, model).
tiny = torch.finfo(self.dtype).tiny
kurt = nsub.unsqueeze(0) * m4 / m2.pow(2).clamp_min(tiny) - 3.0
self.pdtype = self._pdtype_from_kurtosis(kurt, nsub)
def _pdtype_from_kurtosis(
self, kurt: torch.Tensor, nsub: torch.Tensor
) -> torch.Tensor:
"""Map per-source excess kurtosis to a density-family code (pure).
Super-Gaussian (positive kurtosis) -> code 1; sub-Gaussian -> code 4.
Only sources with a meaningful signal switch: a dead model
(``nsub[h]==0`` => ``kurt==-3.0``, finite) or a numerically blown-up
source (``kurt`` NaN, and ``NaN>0`` is False) would otherwise be silently
assigned code 4 with no diagnostic, so those keep their prior ``pdtype``
and are logged -- mirroring the dead-model / non-finite guards in
``_update_parameters``. Split out from ``_choose_pdfs`` so the decision
(including the sub-Gaussian branch, which real EEG rarely triggers) is
unit-testable on a constructed ``kurt`` tensor.
"""
assert self.pdtype is not None
ones = torch.ones_like(self.pdtype)
new_pdtype = torch.where(kurt > 0.0, ones, ones * 4)
valid = torch.isfinite(kurt) & (nsub.unsqueeze(0) > 0.0)
result = torch.where(valid, new_pdtype, self.pdtype)
if not bool(valid.all()):
logger.warning(
"Non-finite or zero-mass kurtosis for %d source/model pair(s) at "
"iter %d; kept their prior pdtype (adaptive-switch guard).",
int((~valid).sum()),
self.iteration,
)
return result
def _snapshot_params(self) -> Dict[str, object]:
"""Snapshot the fitted state for the best-iterate safeguard (issue #51).
Clones each ``_PARAM_TENSORS`` tensor (not an alias) so the live in-place
M-step updates do not roll the snapshot forward; the constant
preprocessing tensors (``mean``/``sphere``) are included so a restore is a
total rollback. Also captures the scalar ``n_kurt_done`` (the adaptive-PDF
switch counter that gates ``pdtype``) so a restored model's switch count
stays consistent with its rolled-back ``pdtype`` -- otherwise a switch
applied after the peak iterate would leave the two out of sync in a saved
model (silent-failure review).
Also captures the LLt stash (issue #157). ``fit`` snapshots immediately
after the E-step that produced both the candidate ``best_ll`` and this
iteration's ``_llt_logv``/``_llt_ll``, so a restore rolls the on-disk
LLt back to the E-step of the restored iterate rather than leaving the
last (discarded) iterate's per-sample values behind. Fortran has no
best-iterate safeguard to reconcile here; this is what keeps the
exported LLt the one belonging to the exported parameters. Costs one
``(n_models + 1) * n_samples`` clone per improving iteration, which is
negligible against the E-step's ``O(n_samples * n_channels^2 * n_mix)``.
"""
snap: Dict[str, object] = {
name: getattr(self, name).clone() for name in self._PARAM_TENSORS
}
snap["n_kurt_done"] = self.n_kurt_done
# Present for every in-fit call (fit allocates the buffers before the
# loop and frees them only after the restore); absent only for a direct
# call on an already-returned model, where there is nothing to roll back.
if self._llt_logv is not None and self._llt_ll is not None:
snap["_llt_logv"] = self._llt_logv.clone()
snap["_llt_ll"] = self._llt_ll.clone()
return snap
def _restore_params(self, snapshot: Dict[str, object]) -> None:
"""Restore the state captured by :meth:`_snapshot_params`."""
for name, value in snapshot.items():
setattr(self, name, value)
# ------------------------------------------------------------------
# Best-of-N restarts (issue #198)
# ------------------------------------------------------------------
# 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. Split from the invariants below purely to document *why*
# an attribute is or is not copied; together the two 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 (the state_dict params) ...
"A", "W", "c", "mu", "alpha", "beta", "rho", "gm", "comp_list", "pdtype",
# ... the schedule/counters a fit mutates (the state_dict extras) ...
"iteration", "ll_history", "final_ll_", "stop_reason", "mir_history_",
"n_newton_fallbacks", "n_kurt_done", "numrej", "good_idx", "_ndtmpsum",
"lrate", "lrate_cap", "newtrate", "rholrate", "rholrate_cap",
# ... the LLt stash and its materialized arrays (issue #157) ...
"_llt_logv", "_llt_ll", "_llt_lht", "_llt_lt",
# ... the tuned block size (do_opt_block re-times per restart) and the
# seed the winning restart ran from.
"block_size", "seed",
) # fmt: skip
# Written by the fit path but identical across the restarts of one fit()
# call, because they are functions of the data alone: preprocessing outputs
# and the model sizing derived from the numerical rank. Copying them would
# be harmless but pointless, and would suggest they could differ.
_RESTART_INVARIANT_ATTRS = (
"mean", "sphere", "sldet", "_sphere_pinv", "n_channels", "n_comps",
) # fmt: skip
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def fit(
self,
X: np.ndarray,
max_iter: int = 100,
verbose: bool = True,
mir_step: int = 0,
) -> "AMICATorchNG":
"""Fit the model to data, running ``n_restarts`` fits and keeping the best.
With the default ``n_restarts=1`` this is exactly :meth:`_fit_once` --
the restart machinery draws nothing, copies nothing and changes nothing,
so the trajectory is bit-identical to a pre-issue-#198 fit. With
``n_restarts > 1`` the model is fit once per seed in ``restart_seeds``
(serially), and the highest-``final_ll_`` non-degenerate restart's
complete state is what the returned model holds.
See :meth:`_fit_once` for the semantics of a single fit (parameters,
``keep_best``, LLt, ``share_comps`` ordering); every one of them applies
unchanged to each restart.
Parameters
----------
X, max_iter, verbose, mir_step
As :meth:`_fit_once`.
Returns
-------
self : AMICATorchNG
Holding the winning restart's parameters, ``ll_history``,
``final_ll_``, ``stop_reason``, ``mir_history_``, LLt arrays and
rejection state -- the state a single fit from that seed would have
left, bit for bit.
Notes
-----
Records (index-aligned, one entry per restart, always populated --
including the single-restart path): ``restart_seeds_``,
``restart_lls_`` (NaN where a restart ended degenerate) and
``restart_stop_reasons_``. The winner is named in one INFO log line.
A restart that ends degenerate (any stop in
``_DEGENERATE_STOP_REASONS``, including ``"restart_error"`` for a
restart that raised) is excluded from selection but recorded. If *every* restart is degenerate the model
is left holding the last one, so issue #50's degenerate-fit contract
applies exactly as it does to a single degenerate fit.
"""
seeds = self._restart_seeds
if len(seeds) == 1:
# Single-restart path: no snapshot, no reseeding of anything that is
# not already the constructor's seed (seeds[0] IS self.seed unless
# the caller passed an explicit one-element restart_seeds), so this
# is byte-for-byte the pre-#198 fit.
self.seed = seeds[0]
self._fit_once(X, max_iter=max_iter, verbose=verbose, mir_step=mir_step)
self.restart_seeds_ = list(seeds)
self.restart_lls_ = [
float("nan") if self.final_ll_ is None else float(self.final_ll_)
]
self.restart_stop_reasons_ = [self.stop_reason]
return self
return self._fit_restarts(X, max_iter, verbose, mir_step)
def _fit_restarts(
self, X: np.ndarray, max_iter: int, verbose: bool, mir_step: int
) -> "AMICATorchNG":
"""Run one full fit per restart seed and keep the winner (issue #198).
Each restart is a complete :meth:`_fit_once`, which re-runs
``_initialize_parameters`` and resets every per-fit counter, so no state
leaks from one restart into the next. The winning restart's state is
captured with :meth:`_capture_restart_state` (a copy of every attribute
the fit path writes) and reapplied at the end unless the winner happens
to be the last restart, whose state is already live.
"""
seeds = self._restart_seeds
lls: List[float] = []
degenerate: List[bool] = []
stop_reasons: List[Optional[str]] = []
states: Dict[int, Dict[str, object]] = {}
for index, seed in enumerate(seeds):
self.seed = seed
try:
self._fit_once(X, max_iter=max_iter, verbose=verbose, mir_step=mir_step)
except RuntimeError as exc:
# A truly singular A makes torch.linalg.inv raise
# torch.linalg.LinAlgError (a RuntimeError) instead of producing
# the non-finite likelihood the in-loop guard catches. Letting
# that propagate would throw away the restarts that already
# succeeded -- the precise opposite of what best-of-N is for --
# so the failure is recorded as a degenerate restart and the
# search moves to the next seed. Deliberately narrow: only
# RuntimeError (which covers every torch numerical failure,
# LinAlgError included). A ValueError from the argument checks at
# the top of _fit_once is a caller mistake, not a bad basin, and
# still propagates on the first restart.
self.stop_reason = restarts.ERROR_STOP_REASON
self.final_ll_ = float("nan")
# The crash can land before _fit_once resets the LLt arrays, in
# which case they still describe an EARLIER restart's E-step.
# Drop them: if a later restart wins, its snapshot restores its
# own; if every restart crashes, the model must not carry
# per-sample likelihoods belonging to a different seed's fit.
self._llt_lht = None
self._llt_lt = None
logger.warning(
"%s", restarts.error_message(index, len(seeds), seed, exc)
)
ll = float("nan") if self.final_ll_ is None else float(self.final_ll_)
is_degenerate = self.stop_reason in self._DEGENERATE_STOP_REASONS
lls.append(ll)
degenerate.append(is_degenerate)
stop_reasons.append(self.stop_reason)
logger.info(
"%s",
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, never
# n_restarts of them.
best_so_far = restarts.select_best(lls, degenerate)
if best_so_far == index:
states = {index: self._capture_restart_state()}
winner = restarts.select_best(lls, degenerate)
if winner is None:
logger.warning(
"%s", restarts.all_degenerate_message(len(seeds), stop_reasons)
)
else:
logger.info(
"%s",
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
return self
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 from a restart
snapshot.
"""
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 _fit_once(
self,
X: np.ndarray,
max_iter: int = 100,
verbose: bool = True,
mir_step: int = 0,
) -> "AMICATorchNG":
"""Run one fit (one initialization, one EM loop) -- what :meth:`fit`
calls once per restart.
Parameters
----------
X : np.ndarray of shape (n_channels, n_samples)
Input data.
max_iter : int, default=100
Number of natural-gradient EM iterations.
verbose : bool, default=True
Show a tqdm progress bar.
mir_step : int, default=0
If > 0, compute MIR (issue #137) from the current ``W``/``sphere``
every ``mir_step`` iterations and append it to ``mir_history_`` as
``(iteration, mir_nats, variance)``. ``0`` (default) disables the
waypoints and leaves fit behavior byte-for-byte unchanged.
``mir_history_`` is a true trajectory like ``ll_history``: a
``keep_best`` (issue #51) restore does not rewrite it, so the
fit-end MIR is ``self.mir(X)`` on the returned parameters, not
``mir_history_[-1]``. Not index-aligned with ``ll_history``:
entry ``i`` is computed after iteration ``i``'s parameter update,
while ``ll_history[i]`` is the likelihood of the parameters
before it, so the two are one update apart (issue #161). An
iteration that ends the fit on a stop takes no update, so it has
no waypoint either.
Incompatible with PCA reduction, same as :meth:`mir` itself. This
upfront gate only sees an explicit reduction request: ``pcakeep <
n_channels`` or any ``pcadb`` while sphering (``pcakeep >=
n_channels`` keeps every dimension, and ``do_sphere=False`` never
reduces, so neither is one; issue #323). The sphere for THIS
fit does not exist yet, so automatic ``mineig``/``mineig_rel`` rank
reduction cannot be checked here; that case is instead caught once
the sphere exists, inside the per-waypoint :meth:`mir` call below,
whose ``ValueError`` is already caught and logged rather than
propagated (issue #283).
Returns
-------
self : AMICATorchNG
Notes
-----
Iteration order (issue #339). Each iteration follows the reference's
main loop (amica15.f90:949-1142): the E-step computes the likelihood
of the current parameters, which is appended to ``ll_history``, and
the step for ``A`` with its norm; then the likelihood-decrease
response (``lrate`` halved, the working ``rholrate`` scaled,
``maxdecs`` ratchets) and the stopping checks run; then, unless a check
fired, the parameters are updated with the rates just set. So after
``fit``:
* ``iteration`` is the 0-based index of the last iteration whose
E-step ran, and ``ll_history`` has ``iteration + 1`` entries (one
fewer after a ``nan_ll``/``singular_ll`` stop, whose likelihood is
not recorded).
* On a convergence stop (``stop_reason`` ``"min_dll"``,
``"grad_norm"``, ``"grad_norm_floor"`` or ``"lrate_floor"``) the
stopping iteration takes no update, as in the reference, so the
returned parameters are the ones whose likelihood is
``ll_history[-1]``, and ``final_ll_`` (without a keep-best restore)
is exactly their log-likelihood. ``iteration`` updates were applied.
* A ``"nan_direction"`` stop (a non-finite step or gradient norm,
caught before any check reads it or the update applies it) and a
``"nan_params"`` stop (non-finite parameters right after an update)
both record their iteration's finite likelihood. Both are
degenerate (``_DEGENERATE_STOP_REASONS``), like ``"nan_ll"`` and
``"singular_ll"``, so ``final_ll_`` is NaN and the model is refused
by every output path.
* At ``max_iter`` (``stop_reason == "max_iter"``) the last iteration
does take its update, as in the reference, so ``max_iter`` updates
were applied and ``final_ll_ == ll_history[-1]`` is the likelihood
of the parameters one update before the returned ones.
Under ``share_comps``, if a merge fires on the LAST iteration, the
returned ``A``/``W``/``comp_list`` are already post-merge but
``final_ll_`` 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 ``final_ll_``'s
comment for detail.
LLt semantics (issue #157). The exported ``LLt``
(``_llt_lht``/``_llt_lt``, written by :meth:`write_amica_output`) is
the per-sample log-likelihood **stashed by the E-step that produced**
``final_ll_``, never a separate post-fit forward pass. Equivalently it
satisfies, bit for bit,
``Lt.sum() / (n_good_samples * n_channels) == final_ll_``
where ``n_good_samples`` is the sample count that E-step ran over
(``good_idx.numel()`` under ``do_reject``, ``n_samples`` otherwise) --
the same normalization ``ll_history`` uses. Three consequences worth
stating plainly:
* Without a keep-best restore, ``final_ll_`` is ``ll_history[-1]``.
After a fit that ran to ``max_iter``, that is the LL of the
parameters as they stood *before* the last M-step -- so the exported
``LLt`` is one M-step older than the exported ``W``/``A``. That is
the reference's own behavior (Fortran fills ``modloglik`` during
iteration i's E-step, ``update_params`` then moves the parameters,
and ``write_output`` writes both -- amica15.f90:996, 1122,
1124-1127), adopted deliberately so pamica's on-disk ``LLt`` is
comparable with the binary's. It is verifiable on the committed
reference output: ``sum(Lt)/(N*nw)`` there equals its ``LL[-1]``
exactly, not the LL of the ``W`` written next to it. After a
convergence stop no M-step follows the last E-step (the reference
exits first too, :1111), so ``LLt`` and ``final_ll_`` belong to the
exported ``W``/``A`` themselves.
* With a keep-best restore (issue #51), the restore rolls the stash
back with the parameters, so ``LLt`` is the restored iterate's own
E-step -- the one that measured ``best_ll`` -- and not the discarded
last iterate's. Because ``_snapshot_params`` is taken right after
that E-step, before any check or update, the restored parameters and
the restored ``LLt`` come from the same point in the loop, so there
is no staleness at all in this case.
* The one case where the equality above does NOT hold is a
``do_reject`` fit whose rejection fires on its own last executed
iteration: ``ll`` was normalized over the good set as it stood
*before* that rejection, and ``_reject_outliers`` then zeroed the
dropped samples' stash entries, so numerator and denominator no
longer refer to the same set and a small residual remains (0.011 on
the bundled sample with one 68-sample pass; it scales with how much
that pass drops). This too is reference-faithful, not 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 zeroes the rejected
``modloglik``/``loglik`` (amica15.f90:2231-2234), and the binary shows
the same residual on the same schedule. Any later iteration
re-normalizes over the shrunk good set and the equality returns.
"""
if X.ndim != 2:
raise ValueError(
f"X must be a 2D array (n_channels, n_samples), got shape {X.shape}"
)
if X.shape[0] != self._n_input_channels:
raise ValueError(
f"X has {X.shape[0]} channels, model expects {self._n_input_channels}"
)
if mir_step < 0:
raise ValueError(f"mir_step must be >= 0, got {mir_step}")
if max_iter < 1:
# PR #318 review: max_iter=0 used to run the loop zero times and
# complete "successfully" with stop_reason="max_iter" (not a
# _DEGENERATE_STOP_REASONS marker) and final_ll_=NaN -- an
# untrained model that every state_dict()/write_amica_output()
# degenerate-fit guard then accepted, since none of them check
# "did an E-step ever actually run", only "did stop_reason end
# up degenerate". Reject up front instead.
raise ValueError(f"max_iter must be >= 1, got {max_iter}")
if mir_step > 0 and self._pca_reduction_requested(X.shape[0]):
raise ValueError(
"mir_step > 0 is incompatible with PCA reduction "
"(pcakeep/pcadb): the sphere is rank-deficient, so MIR's "
"log-Jacobian term is undefined. Rejected up front rather "
"than failing mid-fit at the first waypoint."
)
# Size every fit from the input geometry. _preprocess shrinks
# n_channels/n_comps to the kept rank, so without this a refit, or the
# second of n_restarts, would start from the previous fit's rank.
# A no-op for full-rank data, whose sizes never change.
self.n_channels = self._n_input_channels
self.n_comps = self.n_channels * self.n_models
X_t = self._preprocess(X)
n_total = X_t.shape[1]
self._initialize_parameters()
self.ll_history = []
self.mir_history_ = []
self.numrej = 0
self.n_newton_fallbacks = 0
self.stop_reason = "max_iter"
self.good_idx = (
torch.arange(n_total, device=self.device) if self.do_reject else None
)
# Block-size search (issue #232): after preprocessing and parameter
# initialization, before the first EM iteration, so it times the real
# data on the real device with the parameters the fit will start 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(X_t[:, self.good_idx] if self.do_reject else X_t)
# LLt buffers (issue #157), Fortran's permanently-allocated
# modloglik/loglik (amica15.f90:2617-2620). Zero-filled: a do_reject
# sample that is never scored again keeps the zero that Fortran's
# load_rej reads as the rejection sentinel. Re-allocated per fit so a
# refit on a different dataset cannot serve a stale array.
self._llt_logv = torch.zeros(
(n_total, self.n_models), dtype=self.dtype, device=self.device
)
self._llt_ll = torch.zeros(n_total, dtype=self.dtype, device=self.device)
self._llt_lht = None
self._llt_lt = None
numdecs = 0
# Consecutive-small-likelihood-gain counter for the min_dll stop (Fortran
# numincs, amica15.f90:1079-1089; issue #207). Reset here so a refit on
# the same instance gets a fresh count, matching numdecs.
numincs = 0
# MIR waypoint flood guard (PR #318 review): a ValueError from mir()
# (PCA reduction, or metrics.mir's own near-singular-unmixing check)
# is a per-fit-geometry condition, not a per-iteration one -- it does
# not spontaneously resolve, so leaving the schedule running would
# log the identical warning on every remaining waypoint of a long
# fit. A local (not self.<attr>): it only matters within this one
# _fit_once call, never needs to survive a restart snapshot or be
# inspected after fit() returns.
mir_waypoints_disabled = False
# Best-iterate safeguard (issue #51): track the highest-LL iterate so a
# late Newton-fallback overshoot cannot leave the returned model below a
# peak it already reached. Inactive under do_reject (the good set, and so
# the LL normalization, changes across iterations) and under share_comps
# (a merge drops parameters, so pre- and post-merge LLs are not
# comparable AND the snapshot's comp_list would revert the merge -- the
# returned model would silently be unmerged; #60). In both cases fit()
# returns the last iterate, matching Fortran.
track_best = self.keep_best and not self.do_reject and not self.share_comps
best_ll = -math.inf
best_snapshot: Optional[Dict[str, object]] = None
if self.keep_best and (self.do_reject or self.share_comps):
# keep_best defaults on, so a user enabling rejection/sharing would
# otherwise silently lose the safeguard; surface it once.
reason = "do_reject" if self.do_reject else "share_comps"
logger.warning(
"keep_best is inactive under %s: best-iterate selection by LL is "
"not well-defined (%s), so fit() returns the last iterate.",
reason,
"the good-sample set / LL normalization changes across iterations"
if self.do_reject
else "a merge changes the parameter count and reverting to an "
"earlier snapshot would undo the merge",
)
# One iteration follows the reference's main loop (amica15.f90:949-1142,
# issue #339): (1) the E-step, which yields LL(iter), the sufficient
# statistics and the step with its norm ndtmpsum; (2) the likelihood-
# decrease response and the stopping checks, which read those; (3) an
# exit, BEFORE any parameter moves, if a check fired; otherwise (4) the
# update, with the rates the response just set, then the iteration's
# remaining hooks and rejection. So a stop returns the parameters whose
# likelihood is ll_history[-1], and the step right after a decrease is
# already taken at the halved rate.
iterator = tqdm(range(max_iter), desc="AMICA-NG", disable=not verbose)
for it in iterator:
self.iteration = it
X_use = X_t[:, self.good_idx] if self.do_reject else X_t
n_use = X_use.shape[1]
acc = self._accumulate_blocks(X_use, stash_llt=True)
# Log-likelihood of the CURRENT parameters: acc["ll"] is this
# iteration's E-step total. A singular W makes logdet -> -inf (not
# NaN), so guard on isfinite, not isnan alone: a -inf LL would
# otherwise sail past as a mere "decrease" and the run would
# "complete" (stop_reason=max_iter) on a degenerate model. Checking
# here, before the update, stops on the last parameters instead of
# overwriting them with a garbage update first.
ll = (acc["ll"] / (n_use * self.n_channels)).item()
if not math.isfinite(ll):
self.stop_reason = "nan_ll" if math.isnan(ll) else "singular_ll"
logger.warning(
"Non-finite log-likelihood (%s) at iteration %d; stopping.",
ll,
it,
)
break
# The step and its norm, from the same E-step, before any check reads
# ndtmpsum (the reference computes both in
# accum_updates_and_likelihood, amica15.f90:1749-1761).
step = self._update_direction(acc)
# Best-iterate safeguard (issue #51): remember the parameters that
# produced this LL when it is the best seen, so a later overshoot
# does not leave the returned model below this peak. Nothing has
# moved them since the E-step, so the snapshot pairs them with ll.
if track_best and ll > best_ll:
best_ll = ll
best_snapshot = self._snapshot_params()
self.ll_history.append(ll)
# A non-finite step or norm would pass both gradient-norm checks
# below (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 (
self._ndtmpsum is not None
and math.isfinite(self._ndtmpsum)
and bool(torch.isfinite(step.dAk).all())
):
self.stop_reason = "nan_direction"
logger.warning(
"Non-finite update direction (ndtmpsum %s) at iteration %d; "
"stopping before the update.",
self._ndtmpsum,
it,
)
break
# Learning-rate control, ported from Fortran (amica15.f90:1051-1097).
# Natural-gradient/Newton ascent is not monotonic at a fixed rate:
# when the log-likelihood decreases, anneal the working lrate (and the
# working rho rate). If decreases persist for maxdecs iterations,
# ratchet the *ceilings* down (lrate_cap; rholrate_cap and, under
# Newton, newtrate once past newt_start) so the per-iteration ramp can
# no longer re-inflate lrate back to the overshooting value -- without
# this the ramp and a one-shot halving just oscillate and the LL
# drifts down. All of it runs before this iteration's update, as in
# the reference, so the update below already uses the new rates.
#
# The working rho rate is scaled on every decrease, as the
# reference's is (:1063), but every update of A resets it to its
# ceiling before rho moves (_update_parameters, amica15.f90:1806/
# 1813), so the scaling only reaches rho on an iteration on which A
# is held. It is not a monotone decay (the issue #193 collapse):
# nothing but the maxdecs ratchet lowers the ceiling.
# have_prev mirrors Fortran's outer ``if (iter > 1)`` (amica15.f90:1051),
# which wraps the decrease branch AND the two stops below: none of
# the three checks can fire on the first iteration (no LL(iter-1)
# yet).
#
# PRECEDENCE NOTE (PR #213 review, issue #207): the three blocks
# below (decrease branch; min_dll; grad_norm) are independent --
# none is gated on ``leave`` already being True from an earlier
# block this same iteration, matching Fortran's own structure of
# independent ``leave=.true.`` assignments with no declared
# precedence. Whichever block runs LAST and finds its own
# condition true wins (its ``self.stop_reason =`` is what
# ``fit`` ultimately reports), so with this fixed source order
# (decrease branch, then min_dll, then grad_norm) the standalone
# grad_norm block always has final say when its condition holds.
# In particular, under the shipped ``use_grad_norm=True`` default
# this makes the decrease-branch's ``"grad_norm_floor"`` outcome
# unreachable: its condition (``ndtmpsum <= min_nd`` during a
# decrease) is strictly narrower than the standalone block's
# (``ndtmpsum <= min_nd``, any iteration), so whenever
# ``"grad_norm_floor"`` would fire, the standalone block fires
# too, that same iteration, and overwrites it with
# ``"grad_norm"``. See ``use_grad_norm``'s docstring above and
# ``test_grad_norm_shadows_grad_norm_floor_under_shipped_defaults``.
# This is a reporting nuance, not a behavior change -- deliberately
# NOT restructured into an explicit precedence, to keep this
# section a direct, reviewable port of amica15.f90:1051-1098.
have_prev = len(self.ll_history) > 1
leave = False
if have_prev and ll < self.ll_history[-2]:
# ndtmpsum is the SAME per-iteration value use_grad_norm reads
# below (amica15.f90:1058's ``.or. (ndtmpsum .le. min_nd)``,
# issue #207 gap 3): this is what makes lrate stopping robust
# under do_newton, where lrate sits at newtrate/oscillates
# instead of annealing toward minlrate, so the old
# lrate<=minlrate-only check could never fire (the reported bug).
if self.lrate <= self.minlrate:
logger.warning(
"lrate floor (%g) reached at iter %d; stopping.",
self.minlrate,
it,
)
self.stop_reason = "lrate_floor"
leave = True
elif self._ndtmpsum is not None and self._ndtmpsum <= self.min_nd:
logger.warning(
"gradient-norm floor (%g) reached at iter %d on a "
"likelihood decrease; stopping.",
self.min_nd,
it,
)
self.stop_reason = "grad_norm_floor"
leave = True
else:
self.lrate *= self.lratefact
self.rholrate *= self.rholratefact
numdecs += 1
if numdecs >= self.maxdecs:
self.lrate_cap *= self.lratefact
if schedule.past_newton_start(it, self.newt_start):
self.rholrate_cap *= self.rholratefact
if self.do_newton and schedule.past_newton_start(
it, self.newt_start
):
self.newtrate *= self.lratefact
numdecs = 0
# Small-likelihood-increase stop (Fortran amica15.f90:1078-1090,
# use_min_dll/min_dll/maxincs -- issue #207 gap 1). Independent of
# the decrease branch above: it runs every iteration once have_prev,
# including iterations where the LL just decreased (a decrease is
# always "less than" a positive min_dll, so it also increments
# numincs there, matching Fortran exactly). numincs resets to 0 on
# any gain >= min_dll; stops only after MORE than maxincs
# *consecutive* small gains.
if have_prev and self.use_min_dll:
if ll - self.ll_history[-2] < self.min_dll:
numincs += 1
if numincs > self.maxincs:
logger.warning(
"likelihood increasing by less than %g for more than "
"%d iterations; stopping at iter %d.",
self.min_dll,
self.maxincs,
it,
)
self.stop_reason = "min_dll"
leave = True
else:
numincs = 0
# Weight-gradient-norm stop (Fortran amica15.f90:1091-1097,
# use_grad_norm/min_nd -- issue #207 gap 2). Also independent of the
# decrease branch: this is the unconditional every-iteration check
# (as opposed to the decrease-branch's grad_norm_floor above, which
# only applies alongside a likelihood decrease).
if (
have_prev
and self.use_grad_norm
and self._ndtmpsum is not None
and self._ndtmpsum <= self.min_nd
):
logger.warning(
"norm of weight gradient <= %g at iter %d; stopping.",
self.min_nd,
it,
)
self.stop_reason = "grad_norm"
leave = True
if schedule.newton_switches_on(self.do_newton, it, self.newt_start):
numdecs = 0
# Stop before this iteration's update, as the reference does
# (``if (leave) exit``, amica15.f90:1111, ahead of update_params at
# :1122): the returned parameters are the ones whose LL was just
# recorded, so final_ll_ describes them exactly.
if leave:
break
# Whether rejection fires this iteration (Fortran schedule,
# amica15.f90:1136, with rejstart counted from 1 as the reference
# counts it; issue #335). Fortran rejects using the per-sample
# log-likelihood from THIS iteration's E-step, i.e. the PRE-update
# parameters (loglik is stored in get_updates_and_likelihood before
# update_params runs). Capture it here, before _update_parameters,
# to match that ordering.
will_reject = schedule.rejection_due(
self.do_reject, it, self.rejstart, self.rejint, self.numrej, self.maxrej
)
if will_reject:
assert self.good_idx is not None
reject_ll = self._sample_ll(self.good_idx, X_t)
else:
reject_ll = None
self._update_parameters(acc, n_use, step)
# Surface a corrupted update (a collapsed mixture component, a
# singular inverse) on the iteration it happens, with AMICAMLXNG's
# check and message. The likelihood check above only sees a
# corruption through the NEXT iteration's E-step, so one on the last
# iteration would otherwise end as stop_reason="max_iter" with
# non-finite parameters. One sync: the per-parameter flags are
# stacked and reduced together, and the names are only read out once
# that reduction has failed.
post_update = [getattr(self, name) for name in self._POST_UPDATE_CHECKED]
if not bool(
torch.stack(
[torch.isfinite(value).all() for value in post_update]
).all()
):
bad = [
name
for name, value in zip(self._POST_UPDATE_CHECKED, post_update)
if not bool(torch.isfinite(value).all())
]
logger.warning(
"Non-finite %s at iter %d (a mixture component likely "
"collapsed); stopping.",
", ".join(bad),
it,
)
self.stop_reason = "nan_params"
break
# Extended-Infomax adaptive PDF switch (Fortran do_choose_pdfs). Runs
# on the kurt_start/num_kurt/kurt_int schedule using the just-updated
# W; the new per-source families take effect from the next E-step.
# kurt_start counts from 1, like every reference schedule. num_kurt=0
# disables switching (the family stays at its pdftype=1
# super-Gaussian init).
if self.do_choose_pdfs and self.n_kurt_done < self.num_kurt:
if schedule.periodic_due(it, self.kurt_start, self.kurt_int):
self._choose_pdfs(X_use)
self.n_kurt_done += 1
# Component sharing (Fortran identify_shared_comps schedule,
# amica15.f90:1856): once per share_iter cycle from share_start,
# merge near-collinear components across models using the
# just-updated A. Fortran runs identify_shared_comps BEFORE
# get_unmixing_matrices (amica15.f90:1858,1863), so rebuild W from the
# merged comp_list -- otherwise the next E-step would read a stale W
# (pre-merge comp_list) while indexing the densities by the merged
# comp_list. No-op when share_comps is off or n_models == 1.
#
# This runs AFTER ``ll`` (this iteration's LL) was recorded above,
# so a merge on the final iteration lands in the returned
# A/W/comp_list but not in ``ll_history``/``final_ll_`` -- see
# final_ll_'s comment (issue #269).
if self.share_comps and schedule.periodic_due(
it, self.share_start, self.share_iter
):
self._identify_shared_comps()
self._update_unmixing_matrices()
# MIR waypoint (issue #137), following the NumPy backend's
# writestep/histstep idiom (numpy_impl/core.py). Computed from the
# CURRENT W/sphere (just rebuilt above by _update_parameters /
# the share_comps block) against the raw, un-preprocessed X.
#
# A failed waypoint must never kill the fit. `metrics.mir` raises on
# a near-singular unmixing, and a near-singular W mid-fit is a
# transient the natural gradient can pass through (the same
# condition is only a warning on the training path, see
# numpy_impl/core.py's logdet_W check). Letting that propagate would
# let a purely diagnostic flag destroy an otherwise-recoverable
# decomposition. Warn and record NaN instead: the gap stays visible
# in mir_history_ rather than being silently absent, so a plotted
# trajectory shows a hole exactly where the transient was.
#
# ValueError vs LinAlgError get different treatment (PR #318
# review): a ValueError (PCA reduction, or metrics.mir's own
# near-singular-unmixing check) reflects the fit's GEOMETRY --
# the sphere shape or the current unmixing's conditioning as a
# structural fact -- not a one-off numerical hiccup, so it will
# keep firing identically on every remaining scheduled waypoint
# of a long fit. Warn once, then stop scheduling waypoints for
# the rest of THIS fit (mir_history_ simply gets no more
# entries -- every one it would have gotten is the same NaN
# anyway, so nothing is lost). LinAlgError stays per-waypoint:
# it is the genuinely transient case the comment above already
# describes, which the natural gradient can pass through.
if mir_waypoints_disabled:
pass
elif mir_step > 0 and it % mir_step == 0:
try:
mir_nats, mir_var = self.mir(X)
except np.linalg.LinAlgError as exc:
logger.warning(
"MIR waypoint failed at iter %d (%s: %s); recording NaN "
"and continuing. The fit itself is unaffected.",
it,
type(exc).__name__,
exc,
)
mir_nats = mir_var = float("nan")
self.mir_history_.append((it, mir_nats, mir_var))
except ValueError as exc:
logger.warning(
"MIR waypoint failed at iter %d (%s: %s); this "
"condition will not resolve mid-fit, so MIR "
"waypoints are now disabled for the rest of this "
"fit (mir_history_ gets no further entries). The "
"fit itself is unaffected.",
it,
type(exc).__name__,
exc,
)
self.mir_history_.append((it, float("nan"), float("nan")))
mir_waypoints_disabled = True
else:
self.mir_history_.append((it, mir_nats, mir_var))
# Outlier rejection, after the parameter update (Fortran order,
# amica15.f90:1136-1140) but using the pre-update per-sample LL
# captured above.
if will_reject:
assert reject_ll is not None
self._reject_outliers(reject_ll)
iterator.set_postfix({"LL": f"{ll:.4f}", "lrate": f"{self.lrate:.4g}"})
# Log-likelihood of the parameters fit() returns. A degenerate stop
# leaves the model on the diverged parameters, whose LL is NOT the last
# finite ll_history value (the guard breaks before appending), so report
# NaN there rather than a stale healthy-looking number (silent-failure
# review). Otherwise it is the last trajectory value, overwritten with the
# best iterate's LL below if the safeguard restores it.
if self.stop_reason in self._DEGENERATE_STOP_REASONS:
self.final_ll_ = float("nan")
else:
self.final_ll_ = self.ll_history[-1] if self.ll_history else float("nan")
# Restore the best iterate if the run ended materially below it (issue
# #51). Skipped for a degenerate stop -- not because the parameters are
# necessarily non-finite (a singular_ll stop leaves A/W finite but
# singular) but because salvaging a diverged run here would pre-empt issue
# #50's degenerate-fit contract; state_dict() already refuses to persist
# any model whose stop_reason is degenerate. Also skipped when the final
# LL is within _KEEP_BEST_TOL of the best -- a monotone single-model run
# has final == best, so no restore fires and issue #24 parity stays
# bit-exact.
if (
track_best
and best_snapshot is not None
and self.stop_reason not in self._DEGENERATE_STOP_REASONS
and self.ll_history
and best_ll - self.ll_history[-1] > _KEEP_BEST_TOL
):
logger.info(
"Restoring best iterate (LL %.6f) over final LL %.6f "
"(issue #51 best-iterate safeguard).",
best_ll,
self.ll_history[-1],
)
self._restore_params(best_snapshot)
self.final_ll_ = best_ll
# LLt (Fortran's per-sample/per-model log-likelihood, issue #155):
# materialized from the stash the training E-step filled, with NO extra
# forward pass (issue #157). Runs strictly after the keep-best restore
# above, which rolls the stash back alongside the parameters, so these
# arrays are always the E-step of the iterate fit() returns -- i.e. the
# very E-step whose total is ``final_ll_``:
# Lt.sum() / (n_good_samples * n_channels) == final_ll_ (bit-exact)
# where n_good_samples is the count that E-step ran over. This is
# Fortran's own convention (see fit()'s docstring on staleness) and
# holds for the reference binary's own output too. The single exception,
# also reference-faithful, is a do_reject fit that rejects on its own
# last iteration: ll was normalized before that rejection and the stash
# zeroed after it (Fortran: amica15.f90:1770 precedes :1138/:2252).
# Converted to compact numpy here so the device buffers can be freed;
# a refit reallocates them.
if self._llt_logv is not None and self._llt_ll is not None and self.ll_history:
self._llt_lht = self._llt_logv.T.detach().cpu().numpy()
self._llt_lt = self._llt_ll.detach().cpu().numpy()
# Else: no iteration ever completed an E-step whose LL was recorded
# (max_iter=0, or a degenerate first iteration that broke before
# ll_history.append). The buffers hold nothing but zeros, which
# load_rej would misread as "every sample rejected", so leave
# _llt_lht/_llt_lt None and let write_amica_output omit the file with
# its existing warning rather than write a misleading one.
self._llt_logv = None
self._llt_ll = None
return self
def _sample_ll(self, good_idx: torch.Tensor, X_t: torch.Tensor) -> torch.Tensor:
"""Per-sample total log-likelihood over ``good_idx``, block by block, in
``good_idx`` order (so a keep-mask over the result maps back correctly)."""
parts = [
self._block_sample_ll(X_t[:, good_idx[start : start + self.block_size]])
for start in range(0, int(good_idx.numel()), self.block_size)
]
return torch.cat(parts)
def _reject_outliers(self, ll_vec: torch.Tensor):
"""Permanently drop samples whose (pre-update) log-likelihood is a low
outlier.
Fortran ``reject_data`` (amica17.f90:2380-2464): reject any currently-good
sample with ``loglik < mean - rejsig*std`` (population std). The rejection
is one-directional; ``good_idx`` only ever shrinks, and the good-sample
count drives the ``gm``/LL normalization thereafter. ``ll_vec`` is the
per-sample log-likelihood over the current good set, in ``good_idx`` order.
"""
assert self.good_idx is not None
good = self.good_idx
mean = ll_vec.mean()
std = torch.sqrt((ll_vec.pow(2).mean() - mean.pow(2)).clamp_min(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((~torch.isfinite(ll_vec)).sum())
if n_bad:
raise ValueError(
f"{n_bad} of {ll_vec.numel()} 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 {good.numel()} samples "
f"(rejsig={self.rejsig} too aggressive for this data)."
)
# Zero the LLt stash for the samples being dropped, exactly as Fortran's
# reject_data does (amica15.f90:2231-2234): they are never scored again,
# so without this they would keep the log-likelihood from the last
# iteration that still considered 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 = good[~keep]
self._llt_logv[dropped] = 0.0
self._llt_ll[dropped] = 0.0
self.good_idx = good[keep]
self.numrej += 1
n_rejected = int(good.numel() - self.good_idx.numel())
logger.info(
"Rejection %d at iter %d: dropped %d samples (%d good remaining).",
self.numrej,
self.iteration,
n_rejected,
int(self.good_idx.numel()),
)
def _check_model_idx(self, model_idx: int) -> None:
"""Validate a model index against the fitted ``n_models``.
Raises a clear ``ValueError`` (rejecting negatives, which torch's
negative indexing would otherwise turn into a silent wrong-model result)
instead of an opaque tensor ``IndexError``.
"""
if not isinstance(model_idx, (int, np.integer)):
raise TypeError(
f"model_idx must be an int, got {type(model_idx).__name__}."
)
if not (0 <= model_idx < self.n_models):
raise ValueError(
f"model_idx={model_idx} out of range for a {self.n_models}-model "
f"fit (valid: 0..{self.n_models - 1})."
)
def _nonfinite_params(self) -> list:
"""Names of :attr:`_PARAM_TENSORS` currently holding a non-finite
value (matching the legacy NumPy backend's own
``_nonfinite_params``; issue #306 cross-backend review). Parameters
not yet allocated (``None``, e.g. a partially initialized instance)
are skipped rather than treated as bad.
Single-sync fast path: every parameter's ``isfinite().all()`` is
stacked into one small tensor and reduced with exactly one
``.item()``/``bool()`` device sync, instead of one sync per
parameter (measured ~2.2 ms of avoidable host-device round trips
across a 12-tensor sweep before this, PR #329 review -- CUDA/MPS
pay for each ``bool(tensor)``, not just MLX). The per-parameter
breakdown -- one further, small sync -- is only computed once that
reduction is already ``False``.
"""
names = [
name for name in self._PARAM_TENSORS if getattr(self, name) is not None
]
if not names:
return []
flags = torch.stack(
[torch.isfinite(getattr(self, name)).all() for name in names]
)
if bool(flags.all()):
return []
bad = flags.logical_not().tolist()
return [name for name, is_bad in zip(names, bad) if is_bad]
def _check_usable(self, action: str) -> None:
"""Refuse to serve output from a degenerate fit (issue #306).
Callers first check their own unfitted marker(s) and raise the
existing ``requires a fitted model`` ``RuntimeError`` (unchanged);
this assumes a fit has actually run and adds the two layers
:meth:`state_dict`/:meth:`write_amica_output` already use beyond
that: the ``stop_reason`` gate, then a defense-in-depth isfinite
sweep via :meth:`_nonfinite_params`. Mirrors the
:class:`~pamica.AMICA` wrapper's ``_check_usable`` (issue #50) for
callers using this backend directly.
"""
if self.stop_reason in self._DEGENERATE_STOP_REASONS:
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, X: np.ndarray) -> None:
"""Validate a data array against the fitted input channel count,
mirroring :meth:`fit`'s own ``X`` validation (issue #306): a raw
matmul/broadcast error deep inside a method is less useful than this
named ``ValueError`` at the entry point."""
if X.ndim != 2:
raise ValueError(
f"X must be a 2D array (n_channels, n_samples), got shape {X.shape}"
)
if X.shape[0] != self.n_channels_in:
raise ValueError(
f"X has {X.shape[0]} channels, model expects {self.n_channels_in}"
)
def transform(self, X: np.ndarray, model_idx: int = 0) -> np.ndarray:
"""Apply the learned unmixing matrix to (new) data.
The internal ``W = inv(A)`` is stored transposed relative to the true
unmixing (the E-step forms activations as ``(X-c)^T @ W``, see
``_forward``), so the unmixing applied here is ``W^T`` (issue #24
transpose convention) with the per-model data-space center ``c``
subtracted first (issue #27).
"""
if self.sphere is None or self.mean is None or self.W is None or self.c is None:
raise RuntimeError(
"AMICATorchNG.transform() requires a fitted model; call fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("transform")
self._check_input_shape(X)
X_t = torch.from_numpy(np.ascontiguousarray(X)).to(self.device, self.dtype)
X_t = self.sphere @ (X_t - self.mean)
# c is the per-model data-space center: unmix as W(x - c) (issue #27).
S = self.W[:, :, model_idx].T @ (X_t - self.c[:, model_idx : model_idx + 1])
return S.cpu().numpy()
def get_mixing_matrix(self, model_idx: int = 0) -> np.ndarray:
"""True mixing matrix of model ``model_idx``: the reference's
``A(:, comp_list(:, h))``, i.e. that model's component rows of the
stored ``A`` transposed (issue #24 convention; issue #334 layout), so
column ``i`` is source ``i``'s mixing vector in the sphered space."""
if self.A is None or self.comp_list is None:
raise RuntimeError(
"AMICATorchNG.get_mixing_matrix() requires a fitted model; call "
"fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("get the mixing matrix")
return self.A[self.comp_list[:, model_idx], :].T.cpu().numpy()
@property
def n_channels_in(self) -> int:
"""Input channel count, i.e. the width of the sphere.
Differs from ``n_channels`` only when rank reduction shrank the model to
the detected numerical rank (issue #223); equal to it for full-rank data
and before :meth:`fit`. Read off the sphere whenever one exists, so it
cannot drift from the sphere it describes; before the first fit it is
the constructor's channel count, the width :meth:`fit` accepts.
"""
if self.sphere is None:
return self._n_input_channels
return int(self.sphere.shape[1])
def get_sensor_mixing_matrix(self, model_idx: int = 0) -> np.ndarray:
"""Mixing matrix mapped back to input-channel space.
:meth:`get_mixing_matrix` returns ``A`` in the *sphered* space. These are
the corresponding sensor-space maps (EEGLAB/MNE scalp maps),
``pinv(sphere) @ A``, of shape ``(n_channels_in, n_channels)``. This is
the Fortran ``Spinv`` mapping (amica15.f90:568-578), and it is the only
way to recover sensor maps when rank reduction is active, since the
sphere is then non-square (issue #223).
"""
if self.sphere is None:
raise RuntimeError(
"AMICATorchNG.get_sensor_mixing_matrix() requires a fitted "
"model; call fit() first."
)
if self.A is None or self.comp_list is None:
raise RuntimeError(
"AMICATorchNG.get_sensor_mixing_matrix() requires a fitted "
"model; call fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("get the sensor mixing matrix")
A = self.A[self.comp_list[:, model_idx], :].T
return (self._pinv_sphere() @ A).cpu().numpy()
def get_unmixing_matrix(self, model_idx: int = 0) -> np.ndarray:
"""True unmixing matrix ``W_fort`` = (stored W)^T (issue #24 convention)."""
if self.W is None:
raise RuntimeError(
"AMICATorchNG.get_unmixing_matrix() requires a fitted model; call "
"fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("get the unmixing matrix")
return self.W[:, :, model_idx].T.cpu().numpy()
# ------------------------------------------------------------------
# Preprocessing accessors (issue #313). Same names, shapes and float64
# return type as AMICAMLXNG's, so a consumer that composes the transform
# itself (the MNE export, pamica.mne_compat) reads every backend alike.
# ------------------------------------------------------------------
def get_sphere(self) -> np.ndarray:
"""Fitted sphering matrix, shape ``(n_channels, n_channels_in)``.
Square for a full-rank fit and ``(n_kept, n_channels_in)`` after rank
reduction (issue #223). Together with :meth:`get_mean`,
:meth:`get_model_center` and :meth:`get_unmixing_matrix` it composes
:meth:`transform`: ``S = W_fort @ (sphere @ (X - mean) - c)``.
Returned as an independent float64 copy whatever the fit ``dtype``.
"""
if self.sphere is None:
raise RuntimeError(
"AMICATorchNG.get_sphere() requires a fitted model; call fit() first."
)
self._check_usable("get the sphere")
return self.sphere.detach().cpu().numpy().astype(np.float64)
def get_mean(self) -> np.ndarray:
"""Per-channel mean removed before sphering, shape ``(n_channels_in,)``.
All zeros for a ``do_mean=False`` fit. Returned as an independent
float64 copy whatever the fit ``dtype``.
"""
if self.mean is None:
raise RuntimeError(
"AMICATorchNG.get_mean() requires a fitted model; call fit() first."
)
self._check_usable("get the mean")
return self.mean.detach().cpu().numpy().astype(np.float64).ravel()
def get_model_center(self, model_idx: int = 0) -> np.ndarray:
"""Model ``model_idx``'s center ``c`` in sphered space, shape ``(n_channels,)``.
The per-model offset :meth:`transform` subtracts after sphering
(issue #27). Identically zero for a single-model fit, since the ``c`` update
is gated to ``n_models > 1``. Returned as an independent float64 copy
whatever the fit ``dtype``.
"""
if self.c is None:
raise RuntimeError(
"AMICATorchNG.get_model_center() requires a fitted model; call "
"fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("get the model center")
return self.c[:, model_idx].detach().cpu().numpy().astype(np.float64)
def _pca_reduction_requested(self, n_channels: int) -> bool:
"""Whether the explicit ``pcakeep``/``pcadb`` asks to fit fewer than
``n_channels`` dimensions (the shared predicate,
:func:`pamica.rank.pca_reduction_requested`, issue #323).
``n_channels`` is the channel count of the data being fitted, not
``self.n_channels``, which :meth:`_preprocess` shrinks to the kept rank.
``pcakeep >= n_channels`` is not a request (it keeps every dimension);
any ``pcadb`` is, since whether it cuts anything depends on the
eigenvalues; ``pcakeep`` wins when both are set; and with
``do_sphere=False`` nothing is, since no reduction happens then.
Config-only, not geometry: used solely by :meth:`_fit_once`'s upfront
``mir_step`` gate, which runs BEFORE :meth:`_preprocess` builds this
fit's sphere, so the sphere's actual shape (and therefore any
AUTOMATIC ``mineig``/``mineig_rel`` rank reduction) is not yet
knowable. Use :meth:`_pca_reduced` instead wherever a fitted sphere
already exists (issue #283).
"""
return pca_reduction_requested(
self.pcakeep, self.pcadb, n_channels, self.do_sphere
)
def _pca_reduced(self) -> bool:
"""Whether the fitted sphere is rank-reduced (non-square).
Derived from the fitted geometry -- ``sphere.shape[0] !=
sphere.shape[1]`` -- rather than from which parameter caused the
reduction, so this also catches rank reduction from AUTOMATIC
numerical-rank detection (``mineig``/``mineig_rel``), not just
explicit ``pcakeep``/``pcadb`` (issue #283: the old parameter-only
check let an auto-detected reduction slip past :meth:`mir`'s guard,
which then failed with an opaque ``LinAlgError`` instead of the
documented ``ValueError``). ``False`` before :meth:`fit` (``sphere``
is ``None``) and for a full-rank fit, matching the pre-#283 behavior
in both of those cases.
"""
return self.sphere is not None and self.sphere.shape[0] != self.sphere.shape[1]
def mir(
self, X: np.ndarray, *, model_idx: int = 0, nbins: Optional[int] = None
) -> Tuple[float, float]:
"""Mutual Information Reduction (issue #137) of this model's unmixing on ``X``.
Composes the linear part of the raw-data-to-sources transform, ``W_fort @ sphere`` --
i.e. ``get_unmixing_matrix(model_idx) @ sphere`` -- and delegates to
:func:`pamica.metrics.mir`. MIR is shift-invariant, so the data-space
mean/``c`` centering ``transform`` applies is irrelevant here.
Parameters
----------
X : np.ndarray of shape (n_channels, n_samples)
Raw (unpreprocessed) data.
model_idx : int, default=0
Which model's unmixing to use.
nbins : int, optional
Histogram bin count; see :func:`pamica.metrics.mir`.
Returns
-------
mir_nats : float
variance : float
Raises
------
RuntimeError
If the model is unfitted, or the fit ended degenerate
(issue #306).
ValueError
If ``X`` is not a 2D array of the fitted input channel count, or
if the fitted sphere is rank-reduced (non-square): whether from
explicit ``pcakeep``/``pcadb`` or from automatic ``mineig``/
``mineig_rel`` numerical-rank detection, the sphere is
rank-deficient, so MIR's log-Jacobian term is undefined
(issue #283).
"""
if self.A is None or self.W is None or self.sphere is None:
raise RuntimeError(
"AMICATorchNG.mir() requires a fitted model; call fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("compute MIR")
self._check_input_shape(X)
if self._pca_reduced():
raise ValueError(
"mir() is incompatible with PCA reduction: the fitted "
f"sphere is rank-deficient ({self.n_channels} of "
f"{self.n_channels_in} channels kept), whether from explicit "
"pcakeep/pcadb or automatic mineig/mineig_rel numerical-rank "
"detection, so MIR's log-Jacobian term is undefined for the "
"resulting non-square/non-invertible unmixing."
)
unmixing = (self.W[:, :, model_idx].T @ self.sphere).cpu().numpy()
return mir_metric(unmixing, X, nbins)
def pmi(
self, X: np.ndarray, *, model_idx: int = 0, nbins: Optional[int] = None
) -> np.ndarray:
"""Pairwise Mutual Information (issue #137) between this model's sources on ``X``.
Delegates to :func:`pamica.metrics.pairwise_mi` on
``transform(X, model_idx)``.
Parameters
----------
X : np.ndarray of shape (n_channels, n_samples)
Raw (unpreprocessed) data.
model_idx : int, default=0
Which model's sources to use.
nbins : int, optional
Histogram bin count; see :func:`pamica.metrics.pairwise_mi`.
Returns
-------
mi_matrix : np.ndarray of shape (n_sources, n_sources)
Raises
------
RuntimeError
If the model is unfitted, or the fit ended degenerate
(issue #306), both via :meth:`transform`.
ValueError
If ``X`` is not a 2D array of the fitted input channel count
(via :meth:`transform`).
"""
return pairwise_mi(self.transform(X, model_idx=model_idx), nbins)
# ------------------------------------------------------------------
# Multi-model posterior (issue #141)
# ------------------------------------------------------------------
def model_loglik(self, X: np.ndarray) -> np.ndarray:
"""Per-model, per-sample log-likelihood ``Lht`` on (new) data.
For each model ``h`` and sample ``t`` this is the joint log-likelihood
``log(gm[h]) + log|det W_h| + sldet + sum_i log p_h(s_i)`` (Fortran's
``Lht``/``modloglik``), evaluated on arbitrary raw data via the STORED
sphere/mean -- never re-preprocessing, which would overwrite them. The
per-sample posterior over models (model dominance) is
``softmax(Lht, axis=0)``; see :meth:`model_probability`.
This does not replicate a training-time ``do_reject`` mask: it scores
every sample of ``X``. On a ``do_reject`` fit's own training data it
therefore returns real values where the stored ``_llt_lht`` carries
Fortran's sentinel zeros for rejected samples (issue #155), so the two
agree bit-for-bit only when the fit did not use ``do_reject``. Like
:meth:`transform`, it assumes a usable (non-degenerate) fit; the
:class:`~pamica.AMICA` wrapper enforces that via ``_check_usable``.
Parameters
----------
X : np.ndarray of shape (n_channels, n_samples)
Raw (unpreprocessed) data.
Returns
-------
Lht : np.ndarray of shape (n_models, n_samples)
Raises
------
RuntimeError
If the model is unfitted, or the fit ended degenerate
(issue #306).
ValueError
If ``X`` is not a 2D array of the fitted input channel count, or
contains non-finite (NaN/Inf) values.
"""
if self.sphere is None or self.mean is None or self.W is None:
raise RuntimeError(
"AMICATorchNG.model_loglik() requires a fitted model; call fit() first."
)
self._check_usable("compute the model log-likelihood")
self._check_input_shape(X)
return self._model_loglik_unchecked(X)
def _model_loglik_unchecked(self, X: np.ndarray) -> np.ndarray:
"""Core ``Lht`` computation for :meth:`model_loglik`, with no
degenerate-fit guard or shape validation of its own (issue #306
PR #329 review): :meth:`model_loglik` and :meth:`model_probability` each
do their own single guard + shape check, with their own action
wording, then both call this -- so the guard no longer runs twice
(~2x cost) on a :meth:`model_probability` call, which used to run
its own ``_check_usable`` and then :meth:`model_loglik`'s."""
# Narrows Optional[Tensor] for the type checker: both callers already
# checked these are set before calling this private helper.
assert self.sphere is not None
assert self.mean is not None
X = np.ascontiguousarray(X)
if not np.isfinite(X).all():
bad = np.flatnonzero(~np.isfinite(X).all(axis=1))
raise ValueError(
"AMICATorchNG.model_loglik(): input contains non-finite (NaN/Inf) "
f"values in {bad.size} channel(s) {bad.tolist()}; clean bad "
"segments before scoring."
)
X_t = torch.from_numpy(X).to(self.device, self.dtype)
X_t = self.sphere @ (X_t - self.mean)
n_samples = X_t.shape[1]
Lht = np.zeros((self.n_models, n_samples))
for start in range(0, n_samples, self.block_size):
end = min(start + self.block_size, n_samples)
logV, *_ = self._forward(X_t[:, start:end])
Lht[:, start:end] = logV.T.detach().cpu().numpy()
return Lht
def model_probability(self, X: np.ndarray) -> np.ndarray:
"""Per-sample posterior probability of each model (model dominance).
The column-wise ``softmax`` over models of :meth:`model_loglik`, i.e.
``P(model h | x_t)``; each column sums to 1. For a single model this is
all ones.
Parameters
----------
X : np.ndarray of shape (n_channels, n_samples)
Raw (unpreprocessed) data.
Returns
-------
prob : np.ndarray of shape (n_models, n_samples)
Raises
------
RuntimeError
If the model is unfitted, or the fit ended degenerate
(issue #306).
ValueError
If ``X`` is not a 2D array of the fitted input channel count, if
``X`` is non-finite, if every model underflows to ``-inf``
log-likelihood at some sample (the posterior is undefined
there), or if a log-likelihood is NaN (numerical corruption,
distinct from the ``-inf`` underflow case above).
"""
if self.sphere is None or self.mean is None or self.W is None:
raise RuntimeError(
"AMICATorchNG.model_probability() requires a fitted model; "
"call fit() first."
)
self._check_usable("compute the model probability")
self._check_input_shape(X)
Lht = self._model_loglik_unchecked(X)
# NaN and -inf are different failure modes and must not share a
# message: -inf is every model underflowing at a real sample (an
# extreme outlier), while NaN is numerical corruption. isfinite alone
# conflates them (PR #311 review scope extension, issue #306). The
# diagnosis + normalization is shared with AMICAMLXNG (PR #329
# review) rather than duplicated per backend.
return model_probability_from_loglik(
Lht, caller="AMICATorchNG.model_probability()"
)
# ------------------------------------------------------------------
# Fitted-parameter metadata (issue #142)
# ------------------------------------------------------------------
def get_pdftype(self, model_idx: int = 0) -> np.ndarray:
"""Per-source density-family code for model ``model_idx``.
One integer per source component (0-4; see
:data:`pamica.torch_impl.PDFTYPE_NAMES`): 0 generalized Gaussian, 1
super-Gaussian cosh, 2 Gaussian, 3 logistic, 4 sub-Gaussian cosh. All
sources share ``pdftype`` unless the adaptive switcher (``pdftype=1``)
moved them individually (issue #26).
Returns
-------
np.ndarray of int, shape (n_sources,)
"""
if self.pdtype is None:
raise RuntimeError(
"AMICATorchNG.get_pdftype() requires a fitted model; call fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("get the density family")
return self.pdtype[:, model_idx].detach().cpu().numpy()
def get_rho(self, model_idx: int = 0) -> np.ndarray:
"""Generalized-Gaussian shape parameter ``rho`` for model ``model_idx``.
One value per (mixture component, source): ``rho == 2`` is Gaussian-
shaped, ``rho == 1`` Laplacian, ``rho < 1`` heavier-tailed. Only the
generalized-Gaussian family (``pdftype=0``) updates ``rho``; for every
non-zero code (1-4, the fixed and adaptive cosh families) it stays frozen
at ``rho0`` and does not describe the fitted density.
Returns
-------
np.ndarray of float, shape (n_mix, n_sources)
"""
if self.rho is None or self.comp_list is None:
raise RuntimeError(
"AMICATorchNG.get_rho() requires a fitted model; call fit() first."
)
self._check_model_idx(model_idx)
# Folded into the shared guard (issue #306): a degenerate multi-model
# fit can leave one model's rho non-finite without the aggregate LL
# tripping a _DEGENERATE_STOP_REASONS marker, which _check_usable's
# defense-in-depth isfinite sweep over _PARAM_TENSORS (rho included)
# still catches. Refuse rather than return a silent NaN.
self._check_usable("get rho")
idx = self.comp_list[:, model_idx]
return self.rho[:, idx].detach().cpu().numpy()
def shared_components(self) -> list:
"""Components shared across models by ``share_comps`` (issue #60).
``share_comps`` folds near-collinear components of different models onto
one shared component (one row of ``A`` and one density), recorded as a
repeated index in ``comp_list``. Returns one group per shared
component: a list of ``(model_idx, source_idx)`` pairs that all
reference it, whose columns of :meth:`get_sensor_mixing_matrix` are
therefore identical. Empty when no
component is shared across two or more models (always for one model, and
for a default multi-model fit with ``share_comps`` off).
Note that a merge synchronizes only the parameters routed through
``comp_list`` (the mixing vector and ``mu``/``alpha``/``beta``/``rho``);
the per-source density
*family* code ``pdtype`` is a separate tensor and is not synchronized, so
under the adaptive switcher (``pdftype=1``) a shared pair can still report
different :meth:`get_pdftype` codes.
Returns
-------
list of list of tuple(int, int)
"""
if self.comp_list is None:
raise RuntimeError(
"AMICATorchNG.shared_components() requires a fitted model; call "
"fit() first."
)
self._check_usable("get the shared components")
cl = self.comp_list.detach().cpu().numpy() # (n_sources, n_models)
groups = []
for col in np.unique(cl):
src, mdl = np.where(cl == col)
if np.unique(mdl).size >= 2:
groups.append([(int(h), int(i)) for i, h in zip(src, mdl)])
return groups
# ------------------------------------------------------------------
# EEGLAB drop-in output (issue #92)
# ------------------------------------------------------------------
def variance_order(
self, model_idx: int = 0, return_svar: bool = False
) -> Union[np.ndarray, tuple]:
"""EEGLAB back-projected-variance component order (IC1 = highest variance).
Returns the source indices sorted by descending back-projected variance,
the ordering EEGLAB's ``loadmodout15.m`` applies on load (so ``order[0]``
is IC1). The de-sphered sensor-space mixing column ``a_i = pinv(W S)[:, i]``
contributes ``||a_i||^2 * sum_k alpha_ki (mu_ki^2 + r_ki / sbeta_ki^2)``
with ``r_ki = gamma(3/rho_ki)/gamma(1/rho_ki)`` (the source's mixture
variance), matching ``loadmodout15`` exactly. Non-mutating: the stored
parameters keep their fit order; this only reports the display order.
Parameters
----------
model_idx : int, default=0
Which model's components to order.
return_svar : bool, default=False
If True, also return the per-source variance sorted to ``order``.
Returns
-------
order : np.ndarray of int, shape (n_sources,)
Source indices, highest back-projected variance first.
svar : np.ndarray, optional
Present only when ``return_svar``; the sorted variances.
"""
from scipy.special import gamma
if (
self.comp_list is None
or self.alpha is None
or self.mu is None
or self.beta is None
or self.rho is None
or self.W is None
or self.sphere is None
):
raise RuntimeError(
"AMICATorchNG.variance_order() requires a fitted model; call "
"fit() first."
)
self._check_model_idx(model_idx)
self._check_usable("compute the variance order")
cl = self.comp_list[:, model_idx].cpu().numpy()
alpha = self.alpha[:, cl].cpu().numpy()
mu = self.mu[:, cl].cpu().numpy()
sbeta = self.beta[:, cl].cpu().numpy()
rho = self.rho[:, cl].cpu().numpy()
# source mixture variance (sum over the mixture components); unused
# mixtures carry alpha == 0 and drop out, matching loadmodout15.
ratio = gamma(3.0 / rho) / gamma(1.0 / rho)
mix_var = (alpha * (mu**2 + ratio / sbeta**2)).sum(axis=0)
# de-sphered sensor-space mixing: A = pinv(W_fort @ S), columns = maps.
w_fort = self.W[:, :, model_idx].T.cpu().numpy()
sphere = self.sphere.cpu().numpy()
a_sensor = np.linalg.pinv(w_fort @ sphere)
svar = mix_var * (a_sensor**2).sum(axis=0)
order = np.argsort(-svar)
if return_svar:
return order, svar[order]
return order
def write_amica_output(self, outdir) -> None:
"""Write this fitted model as the Fortran/EEGLAB AMICA output directory.
Produces the raw binary files that EEGLAB's ``loadmodout15.m`` (and the
Python port :func:`pamica.numpy_impl.load.loadmodout`) read: ``gm``,
``W``, ``S``, ``mean``, ``c``, ``alpha``, ``mu``, ``sbeta``, ``rho``,
``comp_list``, ``LL``, so a PyTorch NG fit drops directly into an
EEGLAB workflow (issue #92). ``loadmodout15`` performs the
variance-ordering and unit-norm normalization on load, so the on-disk
parameters are written in fit order. Single-model output is
byte-compatible with the Fortran reference.
Also writes ``LLt`` (the per-sample/per-model log-likelihood,
issue #155) for a model that was just ``fit()`` in this process, from the
stash the training E-step filled (issue #157) -- so, exactly as in the
reference, ``LLt`` is the E-step of the returned iterate: one M-step
older than the ``W``/``A`` written beside it after a fit that ran to
``max_iter``, their own after a convergence stop (see
:meth:`_fit_once`'s Notes). A model restored via :meth:`from_state_dict` carries no stash,
so ``LLt`` is omitted for it (a warning is logged) -- the rest of the
output is unaffected.
Raises if the model is unfitted or degenerate (a fit that ended on a
non-finite log-likelihood): a NaN model must not be written silently.
The scikit-learn-style :class:`~pamica.AMICA` wrapper already refuses
this via its own usability gate, but a caller using
:class:`AMICATorchNG` directly has no such gate in front of this
method -- mirrors :meth:`state_dict`'s two-layer guard (stop_reason
refusal, then a defense-in-depth isfinite sweep over the parameter
tensors) so the same protection applies here (PR #311 review).
Parameters
----------
outdir : str or path-like
Destination directory (created if absent).
"""
if self.A is None:
raise RuntimeError(
"write_amica_output requires a fitted model; call fit() first."
)
if self.stop_reason in self._DEGENERATE_STOP_REASONS:
raise RuntimeError(
f"Refusing to write output for a degenerate model (stop_reason="
f"{self.stop_reason!r}): fit() hit a non-finite value at "
f"iteration {self.iteration}. Fix the instability (lower lrate, "
f"disable Newton, or check data conditioning) before writing."
)
# Defense-in-depth, mirroring state_dict(): catch a non-finite
# parameter even if stop_reason bookkeeping ever misses it. Also
# neutralizes a stale LLt stash: a failed final iteration's
# _llt_lht/_llt_lt (from before the nan_params/nan_ll break) can no
# longer reach disk once this guard refuses the write outright.
nonfinite = self._nonfinite_params()
if nonfinite:
raise RuntimeError(
f"Refusing to write output for a model with non-finite "
f"parameters {nonfinite} (stop_reason={self.stop_reason!r})."
)
from ..numpy_impl.load import write_amicaout
def _np(t):
return t.detach().cpu().numpy()
# The exported parameters are the fit()-kept iterate (LL == final_ll_).
# Under the keep_best safeguard (#51) that can be an earlier iterate than
# the last, so end the written LL trajectory at that iterate rather than
# at a later, discarded overshoot -- otherwise LL[-1] would not match the
# model just written. Monotone runs keep the full trajectory unchanged.
ll = np.asarray(self.ll_history, dtype=np.float64)
if (
self.final_ll_ is not None
and np.isfinite(self.final_ll_)
and ll.size
and not np.isclose(ll[-1], self.final_ll_)
):
ll = ll[: int(np.argmax(ll)) + 1]
# LLt (Fortran's per-sample/per-model log-likelihood, issue #155):
# computed once at the end of fit() (after any keep-best restore) and
# stored compactly on self. A model restored via from_state_dict()
# never ran fit() in this process, so it has neither -- warn rather
# than silently omitting the file (silent-failure review).
if self._llt_lht is not None and self._llt_lt is not None:
Lht, Lt = self._llt_lht, self._llt_lt
else:
logger.warning(
"No LLt data available (model was restored via "
"from_state_dict(), not freshly fit()); writing output "
"without the LLt file."
)
Lht = Lt = None
write_amicaout(
outdir,
gm=_np(self.gm),
W=_np(self.W),
sphere=_np(self.sphere),
mean=_np(self.mean),
c=_np(self.c),
alpha=_np(self.alpha),
mu=_np(self.mu),
sbeta=_np(self.beta), # Fortran's 'sbeta' is pamica's beta (scale)
rho=_np(self.rho),
comp_list=_np(self.comp_list),
ll=ll,
# The reference layout, (nw, num_comps) with component k in column
# k: the component-row A transposed (issue #334).
A=_np(self.A).T,
Lht=Lht,
Lt=Lt,
)
# ------------------------------------------------------------------
# Persistence (issue #36)
# ------------------------------------------------------------------
# Full fitted-parameter snapshot. A/W/c/comp_list/mean/sphere are what
# transform()/get_*matrix() read back; mu/alpha/beta/rho/gm are the
# mixture-PDF EM state, included for a complete snapshot (and for parity/
# continued-analysis) even though no public method currently reads them.
# pdtype is the per-source density-family code (issue #26): a non-default
# pdftype model, or the adaptive switcher's chosen 1/4 assignments, would
# otherwise silently revert to GG on reload. comp_list and pdtype are integer
# tensors (dtype preserved on load); the rest follow self.dtype.
_PARAM_TENSORS = (
"A", "W", "c", "mu", "alpha", "beta", "rho", "gm",
"comp_list", "mean", "sphere", "pdtype",
) # fmt: skip
# Integer tensors in _PARAM_TENSORS: keep their dtype on load, only move device.
_INT_PARAM_TENSORS = ("comp_list", "pdtype")
# The state_dict format. 4 (issue #334) stores A with one component per
# row, shape (n_comps, n_channels); see from_state_dict for what older
# versions load as.
_STATE_FORMAT_VERSION = 4
# Stop reasons that mark a fit as degenerate: a non-finite log-likelihood
# ("nan_ll"/"singular_ll"), a non-finite update direction ("nan_direction"),
# non-finite parameters after an update ("nan_params"), or -- only reachable
# under best-of-N restarts, issue #198 -- a fit that raised before it could
# finish. Such a model yields NaN sources or was never fitted to the end, so
# state_dict() refuses to persist it rather than let it round-trip silently
# (silent-failure review, PR #44). The same set as AMICAMLXNG's.
_DEGENERATE_STOP_REASONS = (
"nan_ll",
"singular_ll",
"nan_direction",
"nan_params",
restarts.ERROR_STOP_REASON,
)
# What fit()'s 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")
def state_dict(self) -> dict:
"""Serialize the fitted model to a plain, device-agnostic dict.
The returned dict has three parts: ``config`` (the constructor
arguments needed to rebuild the object), ``params`` (the fitted
tensors, moved to CPU), and ``extra`` (scalar/schedule state, plus the
optional ``good_idx`` index tensor). Every value is a tensor or a plain
Python primitive, so the dict round-trips through
``torch.save``/``torch.load`` with ``weights_only=True`` (no custom
classes or ``torch.dtype`` objects: dtype is stored by name). Rebuild
with :meth:`from_state_dict`.
Raises if the model is unfitted or degenerate (a fit that ended on a
non-finite log-likelihood): a NaN model must not be persisted silently.
"""
if self.A is None:
raise RuntimeError(
"AMICATorchNG.state_dict() requires a fitted model; call fit() first."
)
if self.stop_reason in self._DEGENERATE_STOP_REASONS:
raise RuntimeError(
f"Refusing to serialize a degenerate model (stop_reason="
f"{self.stop_reason!r}): fit() hit a non-finite value at "
f"iteration {self.iteration}. Fix the instability (lower lrate, "
f"disable Newton, or check data conditioning) before saving."
)
# Defense-in-depth: catch a non-finite parameter even if stop_reason
# bookkeeping ever misses it (the codebase has known NaN-suppression
# risks). isfinite on the integer comp_list is trivially all-True.
nonfinite = self._nonfinite_params()
if nonfinite:
raise RuntimeError(
f"Refusing to serialize a model with non-finite parameters "
f"{nonfinite} (stop_reason={self.stop_reason!r})."
)
config = {
"n_channels": self.n_channels,
"n_models": self.n_models,
"n_mix": self.n_mix,
# block_size is the value the fit actually ran at -- which, under
# do_opt_block, is the size the search chose rather than the one the
# constructor was given (issue #232). Persisting the tuned value is
# what makes a reloaded model reproduce the run it came from; the
# sweep bounds ride along so a re-fit can search again if asked.
"block_size": self.block_size,
"do_opt_block": self.do_opt_block,
"blk_min": self.blk_min,
"blk_max": self.blk_max,
"blk_step": self.blk_step,
# lrate/newtrate/rholrate are annealed during fit; persist the
# original constructor values (lrate0/newtrate0/rholrate0) and
# restore the mutated ones from ``extra`` below.
"lrate": self.lrate0,
"minlrate": self.minlrate,
"lratefact": self.lratefact,
"maxdecs": self.maxdecs,
# Convergence stops (issue #207); fixed hyperparameters, not
# annealed during fit, so no mutated counterpart in ``extra``
# (unlike lrate/newtrate/rholrate) is needed.
"use_min_dll": self.use_min_dll,
"min_dll": self.min_dll,
"maxincs": self.maxincs,
"use_grad_norm": self.use_grad_norm,
"min_nd": self.min_nd,
"newt_ramp": self.newt_ramp,
"do_newton": self.do_newton,
"newt_start": self.newt_start,
"newtrate": self.newtrate0,
"do_reject": self.do_reject,
"rejsig": self.rejsig,
"rejstart": self.rejstart,
"rejint": self.rejint,
"maxrej": self.maxrej,
"rho0": self.rho0,
"minrho": self.minrho,
"maxrho": self.maxrho,
"rholrate": self.rholrate0,
"rholratefact": self.rholratefact,
# Best-iterate safeguard flag (issue #51); only affects a re-fit, but
# persisted so a reloaded model reconstructs its exact configuration.
"keep_best": self.keep_best,
# Density-family selection (issue #26): needed so a reloaded model
# rebuilds with the right pdftype/dorho/do_choose_pdfs and switch
# schedule instead of the GG default.
"pdftype": self.pdftype,
"kurt_start": self.kurt_start,
"num_kurt": self.num_kurt,
"kurt_int": self.kurt_int,
"invsigmin": self.invsigmin,
"invsigmax": self.invsigmax,
"doscaling": self.doscaling,
"scalestep": self.scalestep,
# Component sharing (issue #60): persisted so a reloaded multi-model
# run keeps its schedule; the merged comp_list itself is in params.
"share_comps": self.share_comps,
"share_start": self.share_start,
"share_iter": self.share_iter,
"comp_thresh": self.comp_thresh,
"do_mean": self.do_mean,
"do_sphere": self.do_sphere,
"do_approx_sphere": self.do_approx_sphere,
# Plain int/float: the validator accepts numpy scalars (np.int64),
# which torch.load(weights_only=True) refuses to unpickle (issue
# #323). Same cast as AMICAMLXNG.state_dict, whose JSON needs it.
"pcakeep": None if self.pcakeep is None else int(self.pcakeep),
"pcadb": None if self.pcadb is None else float(self.pcadb),
"mineig": self.mineig,
"mineig_rel": self.mineig_rel,
"seed": self.seed,
# Best-of-N restarts (issue #198). Like keep_best, this only affects
# a re-fit, but it is persisted so a reloaded model reconstructs its
# exact configuration; the restart the fit actually kept is in
# ``extra`` below. self.seed is the winner's seed, and
# restart_seeds is the constructor's list, so a reload re-runs the
# same search rather than re-deriving seeds from the winner.
"n_restarts": self.n_restarts,
"restart_seeds": self.restart_seeds,
# Store dtype by name (e.g. "float64") to keep the payload
# weights_only-safe; rebuilt via getattr(torch, ...) on load.
"dtype": str(self.dtype).split(".")[-1],
}
# .clone() forces an independent copy even when self.device is already
# CPU (where .cpu() would alias): fit() mutates A/mu/beta in place each
# iteration, so an aliased snapshot would silently roll forward if
# state_dict() were ever called mid-fit (e.g. best-so-far checkpointing).
params = {
name: getattr(self, name).detach().cpu().clone()
for name in self._PARAM_TENSORS
}
extra = {
"sldet": float(self.sldet),
"iteration": int(self.iteration),
"ll_history": [float(v) for v in self.ll_history],
"final_ll": None if self.final_ll_ is None else float(self.final_ll_),
"stop_reason": self.stop_reason,
"n_newton_fallbacks": int(self.n_newton_fallbacks),
"n_kurt_done": int(self.n_kurt_done),
"numrej": int(self.numrej),
"good_idx": None
if self.good_idx is None
else self.good_idx.detach().cpu().clone(),
"lrate": float(self.lrate),
"lrate_cap": float(self.lrate_cap),
"newtrate": float(self.newtrate),
"rholrate": float(self.rholrate),
"rholrate_cap": float(self.rholrate_cap),
# Per-restart records (issue #198): which seeds ran, what each
# returned, and why each stopped. Persisted so a reloaded best-of-N
# model can still say how its parameters were chosen instead of
# reporting an empty search.
"restart_seeds_": list(self.restart_seeds_),
"restart_lls_": [float(v) for v in self.restart_lls_],
"restart_stop_reasons_": list(self.restart_stop_reasons_),
}
return {
"format_version": self._STATE_FORMAT_VERSION,
"config": config,
"params": params,
"extra": extra,
}
@classmethod
def from_state_dict(
cls, state: dict, device: Optional[Union[str, torch.device]] = None
) -> "AMICATorchNG":
"""Rebuild a fitted :class:`AMICATorchNG` from :meth:`state_dict` output.
``device`` overrides where the restored tensors live (the constructor
picks a default when ``None``); ``dtype`` always comes from the saved
``config``.
A ``format_version`` 3 state (components as columns of ``A``, before
issue #334) loads unchanged in every other respect: its ``A`` is
converted to component rows without loss, unless ``share_comps`` had
merged components, which raises ``ValueError`` asking for a refit
(:func:`pamica.component_layout.rows_from_legacy_columns`).
"""
# format_version 4 (issue #334) stores A with one component per row,
# shape (n_comps, n_channels). A version 3 payload stored it as
# (n_channels, n_comps) with comp_list indexing columns; it is converted
# without loss when unmerged and refused when share_comps had merged
# components (pamica.component_layout, ADR 0007). Earlier bumps were
# PR #52's 1->2 (adaptive PDF) and PR #53's 2->3 (keep_best); versions 1
# and 2 stay unreadable. Additive keys do not bump the version: issue
# #207's five convergence config keys (use_min_dll/min_dll/maxincs/
# use_grad_norm/min_nd) are simply absent from an older payload's
# ``config`` dict, so ``cls(device=device, **config)`` below falls back
# to the constructor's own Fortran-faithful defaults for whichever keys
# are missing -- see
# test_missing_convergence_keys_fall_back_to_fortran_defaults in
# test_ng_convergence.py.
version = state.get("format_version")
if version not in (cls._STATE_FORMAT_VERSION, _COLUMN_LAYOUT_FORMAT_VERSION):
raise ValueError(
f"unsupported AMICATorchNG state format_version: {version!r} "
f"(expected {cls._STATE_FORMAT_VERSION}, or "
f"{_COLUMN_LAYOUT_FORMAT_VERSION} from before the component-row "
"layout)"
)
for section in ("config", "params", "extra"):
if section not in state:
raise ValueError(
f"malformed AMICATorchNG state: missing {section!r} section "
f"(format_version={version}); the payload may be truncated."
)
config = dict(state["config"])
config["dtype"] = getattr(torch, config["dtype"])
# A missing/unexpected key in a malformed or foreign-version payload
# surfaces as a bare TypeError from the constructor call; every other
# validation step in this method already names the payload as the
# culprit with a ValueError, so wrap this one the same way instead of
# letting a mismatched-keyword TypeError propagate unexplained
# (issue #306).
try:
obj = cls(device=device, **config)
except TypeError as exc:
raise ValueError(
f"malformed AMICATorchNG state: config does not match the "
f"AMICATorchNG constructor ({exc}); the payload may be "
"truncated or from an incompatible version."
) from exc
if version == _COLUMN_LAYOUT_FORMAT_VERSION:
state = _component_rows_state(state)
obj._load_params(state)
return obj
def _load_params(self, state: dict) -> None:
"""Restore fitted tensors/scalars from :meth:`state_dict` output onto
this instance's device/dtype."""
params = state["params"]
missing = [name for name in self._PARAM_TENSORS if name not in params]
if missing:
raise ValueError(f"malformed AMICATorchNG state: missing params {missing}")
# Guard against config/params drift: A and comp_list must match the
# dimensions the constructor just derived, or transform()/the E-step
# would fail later with a confusing matmul error far from load().
if tuple(params["A"].shape) != (self.n_comps, self.n_channels):
raise ValueError(
f"restored A has shape {tuple(params['A'].shape)}, expected "
f"{(self.n_comps, self.n_channels)} for n_channels="
f"{self.n_channels}, n_models={self.n_models}"
)
if tuple(params["comp_list"].shape) != (self.n_channels, self.n_models):
raise ValueError(
f"restored comp_list has shape {tuple(params['comp_list'].shape)}, "
f"expected {(self.n_channels, self.n_models)}"
)
for name in self._PARAM_TENSORS:
tensor = params[name]
# comp_list/pdtype hold integer indices/codes; preserve their dtype
# and only move devices. The float parameters follow self.dtype.
if name in self._INT_PARAM_TENSORS:
setattr(self, name, tensor.to(self.device))
else:
setattr(self, name, tensor.to(self.device, self.dtype))
# sphere was just replaced, so any cached back-map describes the old one.
self._sphere_pinv = None
# config's n_channels is the fitted (possibly reduced) rank, so the
# constructor above set the input count to it; the restored sphere's
# width is the true input count, which a refit validates X against.
assert self.sphere is not None # just set by the loop above
self._n_input_channels = int(self.sphere.shape[1])
extra = state["extra"]
self.sldet = extra["sldet"]
self.iteration = extra["iteration"]
self.ll_history = list(extra["ll_history"])
self.final_ll_ = extra["final_ll"]
self.stop_reason = extra["stop_reason"]
self.n_newton_fallbacks = extra["n_newton_fallbacks"]
self.n_kurt_done = extra["n_kurt_done"]
self.numrej = extra["numrej"]
good_idx = extra["good_idx"]
self.good_idx = None if good_idx is None else good_idx.to(self.device)
# The rates, validated first (a missing or non-finite one is a named
# ValueError); rholrate_cap is additive, see _saved_rates.
for name, value in _saved_rates(extra, "AMICATorchNG").items():
setattr(self, name, value)
# Additive-only, like the issue #207 config keys: a payload written
# before issue #198 simply has no restart records, and an empty search
# is the honest description of a single-fit model saved back then.
self.restart_seeds_ = list(extra.get("restart_seeds_", []))
self.restart_lls_ = list(extra.get("restart_lls_", []))
self.restart_stop_reasons_ = list(extra.get("restart_stop_reasons_", []))