nltools 0.6.0.dev2__py3-none-any.whl → 0.6.0.dev4__py3-none-any.whl

nltools/algorithms/alignment/procrustes.py

@@ -151,8 +151,10 @@ def align(
         method (str): One of `'probabilistic_srm'`, `'deterministic_srm'`, or
             `'procrustes'`. Defaults to `'deterministic_srm'`.
         n_features (int | None): Number of features in the common space (SRM only).
-            None uses the number of voxels, and may not exceed any subject's
-            voxel count. Must be None for `'procrustes'`.
+            None uses the smallest subject's size along the aligned axis (voxels
+            for `axis=0`), the largest value every subject can support; an
+            explicit value may not exceed that size for any subject. Must be None
+            for `'procrustes'`.
         axis (int): Axis to align on: 0 aligns timepoints (ISC computed per voxel),
             1 aligns voxels (ISC computed per timepoint). Defaults to 0.
         n_iter (int): Number of `_SRM`/`_DetSRM` iterations; ignored by
@@ -235,7 +237,7 @@ def align(
     out = {}
     if method in ["deterministic_srm", "probabilistic_srm"]:
         if n_features is None:
-            n_features = int(data[0].shape[0])
+            n_features = int(min(x.shape[0] for x in data))
         if method == "deterministic_srm":
             srm = _DetSRM(
                 n_features=n_features, n_iter=n_iter, random_state=random_state

nltools/algorithms/signal.py

@@ -214,7 +214,9 @@ def upsample(
     return upsampled_df
 
 
-def make_cosine_basis(nsamples, sampling_freq, filter_length, unit_scale=True, drop=0):
+def make_cosine_basis(
+    nsamples, sampling_interval, filter_length, unit_scale=True, drop=0
+):
     """Create basis functions for a discrete cosine transform.
 
     Based on the implementation in ``spm_filter`` and ``spm_dctmtx`` because
@@ -224,11 +226,12 @@ def make_cosine_basis(nsamples, sampling_freq, filter_length, unit_scale=True, d
 
     Args:
         nsamples (int): Number of observations (e.g. TRs).
-        sampling_freq (float): Sampling *interval* in seconds (the TR), matching
-            SPM's `RT` — not a frequency, despite the name. The number of bases
-            is `trunc(2 * nsamples * sampling_freq / filter_length + 1)`, minus
-            the constant. `DesignMatrix.add_dct_basis` passes
-            `1 / DesignMatrix.sampling_freq` for this reason.
+        sampling_interval (float): Seconds between observations (the TR), SPM's
+            `RT`. The number of bases is
+            `trunc(2 * nsamples * sampling_interval / filter_length + 1)`, minus
+            the constant. `DesignMatrix` stores a sampling frequency, so
+            `DesignMatrix.add_dct_basis` inverts it and passes
+            `1 / DesignMatrix.sampling_freq`.
         filter_length (int): Filter length in seconds.
         unit_scale (bool): Scale the basis functions to the range [-1, 1]. Defaults to True.
         drop (int): Number of leading (slowest) bases to drop after the constant is
@@ -239,12 +242,13 @@ def make_cosine_basis(nsamples, sampling_freq, filter_length, unit_scale=True, d
         np.ndarray: Basis matrix of shape (nsamples, n_bases).
 
     Note:
-        The basis count follows `spm_dctmtx`'s `k = fix(2*(n*RT)/HParam + 1)`, so
-        100 TRs of 2 s with a 128 s filter give the same four bases SPM gives.
+        The basis count follows `spm_dctmtx`'s `k = fix(2*(n*RT)/HParam + 1)` less
+        the constant that `spm_filter` drops, so 100 TRs of 2 s under a 128 s
+        filter give three bases, the same three SPM uses.
     """
 
     # Figure out number of basis functions to create
-    order = int(np.trunc(2 * (nsamples * sampling_freq) / filter_length + 1))
+    order = int(np.trunc(2 * (nsamples * sampling_interval) / filter_length + 1))
 
     n = np.arange(nsamples)
 

nltools/data/designmatrix/regressors.py

@@ -427,7 +427,10 @@ def _add_dct_basis(
 
     # Create DCT basis matrix using stats function
     basis_mat = make_cosine_basis(
-        dm.shape[0], 1.0 / dm.sampling_freq, duration, drop=drop
+        dm.shape[0],
+        sampling_interval=1.0 / dm.sampling_freq,
+        filter_length=duration,
+        drop=drop,
     )
 
     # Generate column names (.nl_cosine_1, .nl_cosine_2, ...)

nltools/data/simulator/haxby.py

@@ -186,7 +186,9 @@ def _simulate_run(dm, rois, n_voxels, rng):
     return data
 
 
-def load_haxby_example(n_runs=1, *, space="mni", random_state=42):
+def load_haxby_example(
+    n_runs=1, *, space="mni", block_order="independent", random_state=42
+):
     """Load a synthetic Haxby-like dataset on the MNI grid, entirely in-memory.
 
     The quickest way to try nltools: nothing is downloaded beyond the MNI
@@ -204,7 +206,10 @@ def load_haxby_example(n_runs=1, *, space="mni", random_state=42):
     fusiform cortex, and scrambled pictures in early visual cortex — with a 3%
     peak response against 1% white noise. A condition also drives its
     category-mates' spheres at a third of that, so the eight response patterns
-    carry an animate / man-made / scene / control similarity structure.
+    carry an animate / man-made / scene / control similarity structure. Each run
+    draws its own block order by default; pass `block_order='shared'` when the
+    runs have to line up TR by TR, as they do for intersubject correlation,
+    alignment and anything else that compares runs timepoint against timepoint.
 
     Args:
         n_runs (int): Number of runs to generate. Default 1.
@@ -213,6 +218,10 @@ def load_haxby_example(n_runs=1, *, space="mni", random_state=42):
             `extract_roi` work on the result. `'grid'` is a 10 x 10 x 5
             synthetic volume with random condition clusters, for tests that need
             construction in milliseconds and never plot.
+        block_order (str): How the condition order varies across runs.
+            `'independent'` (default) draws a fresh order per run; `'shared'`
+            draws one order and gives it to every run, so the runs are
+            comparable TR by TR. Noise differs per run either way.
         random_state (int | None): Seed for reproducible output. Default 42.
 
     Returns:
@@ -223,7 +232,8 @@ def load_haxby_example(n_runs=1, *, space="mni", random_state=42):
             `_c0` (HRF-convolved boxcars).
 
     Raises:
-        ValueError: If `space` is not `'mni'` or `'grid'`.
+        ValueError: If `space` is not `'mni'` or `'grid'`, or `block_order` is not
+            `'independent'` or `'shared'`.
 
     Examples:
         ```python
@@ -258,16 +268,31 @@ def load_haxby_example(n_runs=1, *, space="mni", random_state=42):
 
     if space not in ("mni", "grid"):
         raise ValueError(f"space must be 'mni' or 'grid', got {space!r}")
+    if block_order not in ("independent", "shared"):
+        raise ValueError(
+            f"block_order must be 'independent' or 'shared', got {block_order!r}"
+        )
 
     rng = np.random.default_rng(random_state)
     mask_img = _mask_image(space)
     n_voxels = int((np.asarray(mask_img.dataobj) > 0).sum())
     rois = _condition_rois(space, mask_img, rng)
 
+    # Drawn once so every run gets the same sequence; 'independent' draws inside
+    # the loop instead, leaving the default run's RNG stream untouched.
+    shared_order = (
+        list(rng.permutation(_HAXBY_CONDITIONS)) if block_order == "shared" else None
+    )
+
     brain_data_list = []
     design_matrix_list = []
     for run in range(n_runs):
-        dm, labels = _run_design(list(rng.permutation(_HAXBY_CONDITIONS)))
+        order = (
+            shared_order
+            if shared_order is not None
+            else list(rng.permutation(_HAXBY_CONDITIONS))
+        )
+        dm, labels = _run_design(order)
         data = _simulate_run(dm, rois, n_voxels, rng)
         labels_and_run = pl.DataFrame(
             {

nltools/models/ridge.py

@@ -12,7 +12,7 @@ import numbers
 from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 
@@ -71,6 +71,31 @@ def _scoped_himalaya_backend(name: str):
         set_backend(previous, on_error="raise")
 
 
+def _solver_form(
+    n_samples: int, n_features: int, n_spaces: int = 1
+) -> Literal["primal", "kernel"]:
+    """Return `'kernel'` when the design is wide, otherwise `'primal'`.
+
+    Himalaya's flowchart routes designs with more features than samples to its
+    kernel-form solvers, whose cross-validation working set scales with the
+    sample count instead of the feature count. Both forms solve the same
+    problem; only the cost differs. A banded fit holds one `(n_samples,
+    n_samples)` kernel per space, so it is only cheaper in the kernel form when
+    the average space is wider than the sample count, not when the total is. A
+    tie stays primal, which for one space is also the threshold of Himalaya's
+    own "slower than kernel ridge" warning.
+
+    Args:
+        n_samples (int): Rows of the design.
+        n_features (int): Total columns across every feature space.
+        n_spaces (int): Number of feature spaces. Defaults to 1.
+
+    Returns:
+        str: `'kernel'` or `'primal'`.
+    """
+    return "kernel" if n_spaces * n_samples < n_features else "primal"
+
+
 def _himalaya_backend_name(backend) -> str:
     """Map a resolved nltools `_Backend` to its Himalaya backend name.
 
@@ -100,26 +125,62 @@ def _batch_sizes(
     n_targets: int,
     n_alphas: int,
     itemsize: int,
+    *,
+    solver_form: str = "primal",
+    n_spaces: int = 1,
 ) -> dict[str, int]:
     """Size the batches of a whole cross-validated or banded fit.
 
     Only the per-item working-set estimates live here; the budget itself and
     the batch arithmetic come from `nltools.algorithms.backends`. The estimates
-    follow Himalaya's dominant allocations: decomposition matrices of
-    `(n_alphas_batch, n_features, n_samples)`, cross-validated predictions of
-    `(n_alphas_batch, n_samples, n_targets_batch)`, and refit weights of
-    `(n_features, n_targets_batch_refit)`.
+    follow Himalaya's dominant allocations. In the primal form they are
+    decomposition matrices of `(n_alphas_batch, n_features, n_samples)`,
+    cross-validated predictions of `(n_alphas_batch, n_samples,
+    n_targets_batch)`, and refit weights of `(n_alphas_batch, n_features,
+    n_targets_batch_refit)`. In the kernel form every `n_features` above
+    becomes `n_samples`: the decomposition is of the `(n_samples, n_samples)`
+    kernel and the refit weights are dual.
+
+    The arrays that stay on the device for the whole fit are charged first,
+    and the batches are sized from what remains: the design and targets in
+    both forms, plus in the kernel form one `(n_samples, n_samples)` kernel per
+    space and the same-sized kernel Himalaya sums or slices per fold.
+
+    Args:
+        backend (Backend): Backend returned by `_resolve_backend`.
+        memory_budget_gb (float | None): Explicit budget, or None to measure.
+        n_samples (int): Rows of the design.
+        n_features (int): Total columns across every feature space.
+        n_targets (int): Columns of `y`.
+        n_alphas (int): Candidate alphas.
+        itemsize (int): Bytes per element of the working dtype.
+        solver_form (str): `'primal'` or `'kernel'`, from `_solver_form`.
+        n_spaces (int): Number of feature spaces. Defaults to 1.
 
     Returns:
         dict[str, int]: `n_targets_batch`, `n_targets_batch_refit`, and
             `n_alphas_batch`.
+
+    Raises:
+        ValueError: If the resident arrays alone exceed the budget, or one
+            batch item does.
     """
     budget_gb = _device_memory_budget(
         backend, max_gpu_memory_gb=memory_budget_gb, cap_for_batching=True
     )
+    resident_bytes = (n_samples * n_features + n_samples * n_targets) * itemsize
+    if solver_form == "kernel":
+        resident_bytes += (n_spaces + 1) * n_samples * n_samples * itemsize
+    budget_gb -= resident_bytes / 1e9
+    if budget_gb <= 0:
+        raise ValueError(
+            f"the resident design, targets and kernels need "
+            f"{resident_bytes / 1e9:.6g} GB, exceeding the budget"
+        )
+    weight_rows = n_samples if solver_form == "kernel" else n_features
     n_alphas_batch, _ = _auto_batch_size(
         n_alphas,
-        n_features * n_samples * itemsize,
+        weight_rows * n_samples * itemsize,
         budget_gb=budget_gb,
         overhead=_WORKING_SET_OVERHEAD,
     )
@@ -131,7 +192,7 @@ def _batch_sizes(
     )
     n_targets_batch_refit, _ = _auto_batch_size(
         n_targets,
-        n_alphas_batch * n_features * itemsize,
+        n_alphas_batch * weight_rows * itemsize,
         budget_gb=budget_gb,
         overhead=_WORKING_SET_OVERHEAD,
     )
@@ -156,8 +217,11 @@ def _refit_targets_batch(
 
     With one shared alpha Himalaya reuses a single shrinkage operator, so a
     target costs only its own columns of `Y` and of the weights. With a
-    per-target alpha it instead holds an `(n_targets_batch, n_samples,
-    n_samples)` block, which dominates everything else.
+    per-target alpha it instead holds one square matrix per target, which
+    dominates everything else: `(n_samples, n_samples)` in the kernel form a
+    wide design runs, and `(n_features, n_samples)` with `n_features <=
+    n_samples` in the primal form a tall one runs, so the sample-count square
+    bounds both.
 
     Returns:
         int: Target batch size in `[1, n_targets]`.
@@ -401,6 +465,12 @@ def _refit_fixed_hyperparameters(
     `sqrt(gamma[k])` before the solve and the resulting coefficients are scaled
     back, which is equivalent to the per-space penalty `alpha / gamma[k]`.
 
+    The scaled design is one ordinary ridge system, so it follows Himalaya's
+    flowchart with one space: a tall design runs `solve_ridge_svd`, a wide one
+    forms its linear kernel, runs `solve_kernel_ridge_eigenvalues`, and turns
+    the dual weights back into coefficients with Himalaya's own host-side
+    `primal_weights_kernel_ridge`.
+
     Targets that selected the same weight vector share one decomposition. The
     grouping is an implementation detail and does not change the result.
 
@@ -427,6 +497,11 @@ def _refit_fixed_hyperparameters(
         np.ndarray: Coefficients of shape `(n_features, n_targets)` in the
             original, unscaled feature coordinates, as CPU NumPy.
     """
+    from himalaya.kernel_ridge import (
+        linear_kernel,
+        primal_weights_kernel_ridge,
+        solve_kernel_ridge_eigenvalues,
+    )
     from himalaya.ridge import solve_ridge_svd
 
     if isinstance(feature_spaces, _ResidentDesign):
@@ -444,6 +519,7 @@ def _refit_fixed_hyperparameters(
 
     alphas = np.broadcast_to(np.asarray(alpha, dtype=np.float64), (n_targets,))
     coef = np.zeros((n_features, n_targets), dtype=np.float64)
+    solver_form = _solver_form(n_samples, n_features)
 
     with _scoped_himalaya_backend(resident.backend_name):
         stacked = _take_rows(resident.design, row_indices)
@@ -459,8 +535,8 @@ def _refit_fixed_hyperparameters(
                 ).astype(dtype)
                 design = stacked * _on_active_backend(scale)[0]
             # A shared alpha needs one shrinkage vector; a per-target alpha
-            # makes Himalaya hold an (n_targets_batch, n_samples, n_samples)
-            # block instead, so the two paths get different batch estimates.
+            # makes Himalaya hold a block of one square matrix per target
+            # instead, so the two paths get different batch estimates.
             group_alphas = alphas[columns]
             shared_alpha = bool(np.all(group_alphas == group_alphas[0]))
             batch = n_targets_batch
@@ -479,16 +555,29 @@ def _refit_fixed_hyperparameters(
                 if len(columns) == n_targets
                 else _take_columns(all_targets, columns)
             )
-            solved = solve_ridge_svd(
-                design,
-                targets,
-                alpha=dtype.type(group_alphas[0])
+            group_alpha = (
+                dtype.type(group_alphas[0])
                 if shared_alpha
-                else group_alphas.astype(dtype),
-                fit_intercept=False,
-                n_targets_batch=batch,
-                warn=False,
+                else group_alphas.astype(dtype)
             )
+            if solver_form == "kernel":
+                dual = solve_kernel_ridge_eigenvalues(
+                    linear_kernel(design),
+                    targets,
+                    alpha=group_alpha,
+                    fit_intercept=False,
+                    n_targets_batch=batch,
+                )
+                solved = primal_weights_kernel_ridge(dual, design)
+            else:
+                solved = solve_ridge_svd(
+                    design,
+                    targets,
+                    alpha=group_alpha,
+                    fit_intercept=False,
+                    n_targets_batch=batch,
+                    warn=False,
+                )
             solved = np.asarray(_to_cpu_numpy(solved), dtype=np.float64)
             if scale is not None:
                 solved = solved * scale[:, None]
@@ -516,6 +605,49 @@ def _on_active_backend(*arrays):
     return [backend.asarray(array) for array in arrays]
 
 
+def _to_backend_arrays(spaces, targets, alphas, dtype):
+    """Move a fit's inputs onto Himalaya's active backend in the working dtype.
+
+    The designs come back as a list even when there is one, because Himalaya's
+    banded solvers take a list of spaces: their `check_arrays` converts a list
+    element by element, which is what lets spaces of differing widths through,
+    whereas a tuple would be handed to `asarray` whole and fail as ragged.
+
+    Args:
+        spaces (Sequence[np.ndarray]): Feature matrices in coefficient order,
+            already contiguous in `dtype`.
+        targets (np.ndarray): `(n_samples, n_targets)` targets in `dtype`.
+        alphas (np.ndarray): Candidate alphas, float64.
+        dtype (np.dtype): Working dtype for the alpha grid.
+
+    Returns:
+        tuple: `(designs, targets, alphas)` on the active backend, with
+            `designs` a list holding one entry per space.
+    """
+    converted = _on_active_backend(*spaces, targets, alphas.astype(dtype))
+    return converted[:-2], converted[-2], converted[-1]
+
+
+def _linear_kernels(spaces):
+    """Stack the linear kernels of several feature spaces for the banded solver.
+
+    Each kernel is Himalaya's own `linear_kernel`, called inside the same
+    `_scoped_himalaya_backend` block as the solver that consumes it, so it is
+    built on the device and in the dtype the solve runs in.
+
+    Args:
+        spaces (Sequence): Feature matrices already on the active backend,
+            each `(n_samples, n_features_k)`.
+
+    Returns:
+        Array: Shape `(n_spaces, n_samples, n_samples)` on the active backend.
+    """
+    from himalaya.backend import get_backend
+    from himalaya.kernel_ridge import linear_kernel
+
+    return get_backend().stack([linear_kernel(space) for space in spaces])
+
+
 def _to_cpu_numpy(array) -> np.ndarray:
     """Return `array` as CPU NumPy, whatever backend produced it.
 
@@ -663,7 +795,10 @@ class _Ridge:
         cv_scores_ (float | np.ndarray | None): None for a fixed-alpha fit. For
             ordinary Ridge, the fold-averaged negative-MSE score at the selected
             alpha. For banded Ridge, `(search_iterations,)` or
-            `(search_iterations, n_targets)` fold-averaged scores.
+            `(search_iterations, n_targets)` fold-averaged scores. Stored as
+            Himalaya reports them: in the kernel form a candidate alpha below
+            the float32 rounding floor of the linear kernel scores `-1e5`, and
+            a target whose every candidate scored that way keeps the first.
         feature_space_weights_ (np.ndarray | None): None for ordinary Ridge.
             Strictly positive weights whose columns sum to one, shaped
             `(n_spaces,)` or `(n_spaces, n_targets)`.
@@ -672,6 +807,8 @@ class _Ridge:
         feature_space_sizes_ (tuple[int, ...] | None): Feature counts aligned
             with `feature_space_names_`; None for ordinary Ridge.
         backend_ (Backend): The resolved execution backend.
+        solver_form_ (str): `'primal'` or `'kernel'`, the Himalaya solver family
+            the fit ran. Chosen from the design shape by `_solver_form`.
         n_samples_ (int): Fitted sample count.
         n_features_in_ (int): Total fitted feature count across spaces.
         is_fitted_ (bool): True after a successful fit.
@@ -1010,12 +1147,15 @@ class _Ridge:
         n_targets = targets.shape[1]
         alphas = np.atleast_1d(np.asarray(alpha, dtype=np.float64))
 
-        self.backend_ = backend
+        # Every fitted attribute is assigned only once the solve has succeeded,
+        # so a fit that raises leaves the previous fitted state intact.
         if scalar_alpha:
             # The fixed-alpha refit sizes its own batch from the same budget; it
             # never runs the cross-validation or alpha loops the others measure.
-            self._fit_fixed_alpha(spaces, targets, float(alpha))
+            solver_form = _solver_form(n_samples, n_features)
+            self._fit_fixed_alpha(backend, spaces, targets, float(alpha))
         else:
+            solver_form = _solver_form(n_samples, n_features, len(spaces))
             try:
                 batches = _batch_sizes(
                     backend,
@@ -1025,6 +1165,8 @@ class _Ridge:
                     n_targets=n_targets,
                     n_alphas=alphas.size,
                     itemsize=dtype.itemsize,
+                    solver_form=solver_form,
+                    n_spaces=len(spaces),
                 )
             except ValueError as error:
                 raise ValueError(
@@ -1032,11 +1174,21 @@ class _Ridge:
                     f"fit with n_samples={n_samples}, n_features={n_features}, "
                     f"n_targets={n_targets}, n_alphas={alphas.size} ({error})"
                 ) from error
-            if is_banded:
-                self._fit_banded(spaces, targets, alphas, dtype, batches)
+            if is_banded and solver_form == "kernel":
+                self._fit_banded_kernel(
+                    backend, spaces, targets, alphas, dtype, batches
+                )
+            elif is_banded:
+                self._fit_banded(backend, spaces, targets, alphas, dtype, batches)
+            elif solver_form == "kernel":
+                self._fit_ordinary_cv_kernel(
+                    backend, spaces, targets, alphas, dtype, batches
+                )
             else:
-                self._fit_ordinary_cv(spaces, targets, alphas, dtype, batches)
+                self._fit_ordinary_cv(backend, spaces, targets, alphas, dtype, batches)
 
+        self.backend_ = backend
+        self.solver_form_ = solver_form
         self.feature_space_names_ = names
         self.feature_space_sizes_ = sizes if is_banded else None
         self.n_samples_ = int(n_samples)
@@ -1046,10 +1198,11 @@ class _Ridge:
         self.is_fitted_ = True
         return self
 
-    def _fit_fixed_alpha(self, spaces, targets, alpha) -> None:
+    def _fit_fixed_alpha(self, backend, spaces, targets, alpha) -> None:
         """Solve a fixed-alpha ordinary Ridge and store the fitted state.
 
         Args:
+            backend (Backend): Resolved execution backend.
             spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alpha (float): The fixed regularization strength.
@@ -1058,17 +1211,41 @@ class _Ridge:
             spaces,
             targets,
             alpha,
-            backend=self.backend_,
+            backend=backend,
             memory_budget_gb=self.memory_budget_gb,
         )
         self.alpha_ = float(alpha)
         self.cv_scores_ = None
         self.feature_space_weights_ = None
 
-    def _fit_ordinary_cv(self, spaces, targets, alphas, dtype, batches) -> None:
+    def _store_ordinary_state(self, coef, best_alphas, cv_scores, alphas) -> None:
+        """Normalize an ordinary cross-validated solve into the fitted state.
+
+        Shared by the primal and kernel forms, whose solvers differ only in how
+        `coef` is obtained. The selected alphas are snapped back onto the
+        candidate grid, and every array is stored as CPU float64.
+
+        Args:
+            coef: `(n_features, n_targets)` coefficients, on any backend.
+            best_alphas: `(n_targets,)` selected alphas, on any backend.
+            cv_scores: Fold-averaged scores at the selected alpha, on any backend.
+            alphas (np.ndarray): Candidate alphas.
+        """
+        self.coef_ = np.asarray(_to_cpu_numpy(coef), dtype=np.float64)
+        selected = _snap_to_grid(_to_cpu_numpy(best_alphas), alphas)
+        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+        self.cv_scores_ = np.asarray(
+            _to_cpu_numpy(cv_scores), dtype=np.float64
+        ).reshape(-1)
+        self.feature_space_weights_ = None
+
+    def _fit_ordinary_cv(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
         """Select an alpha by cross-validation and store the fitted state.
 
         Args:
+            backend (Backend): Resolved execution backend.
             spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alphas (np.ndarray): Candidate alphas.
@@ -1078,12 +1255,12 @@ class _Ridge:
         from himalaya.ridge import solve_ridge_cv_svd
         from himalaya.scoring import l2_neg_loss
 
-        with _scoped_himalaya_backend(_himalaya_backend_name(self.backend_)):
-            design, y_device, alpha_device = _on_active_backend(
-                spaces[0], targets, alphas.astype(dtype)
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
             )
             best_alphas, coefs, cv_scores = solve_ridge_cv_svd(
-                design,
+                designs[0],
                 y_device,
                 alphas=alpha_device,
                 fit_intercept=False,
@@ -1095,49 +1272,131 @@ class _Ridge:
                 **batches,
             )
 
-        self.coef_ = np.asarray(_to_cpu_numpy(coefs), dtype=np.float64)
-        selected = _snap_to_grid(_to_cpu_numpy(best_alphas), alphas)
-        self.alpha_ = float(selected[0]) if not self.per_target_alpha else selected
-        self.cv_scores_ = np.asarray(
-            _to_cpu_numpy(cv_scores), dtype=np.float64
-        ).reshape(-1)
-        self.feature_space_weights_ = None
+        self._store_ordinary_state(coefs, best_alphas, cv_scores, alphas)
 
-    def _fit_banded(self, spaces, targets, alphas, dtype, batches) -> None:
-        """Run the banded random search and store the fitted state.
+    def _fit_ordinary_cv_kernel(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
+        """Select an alpha by cross-validation in the kernel form and store the state.
+
+        The wide-design counterpart of `_fit_ordinary_cv`. Himalaya solves on the
+        `(n_samples, n_samples)` linear kernel and returns dual weights on the
+        host; its own `primal_weights_kernel_ridge` turns them back into `coef_`
+        on the host, which is where Himalaya keeps primal weights because they
+        can be large. The device copy of the design is released once the kernel
+        exists, so only the kernel stays resident through the solve.
 
         Args:
-            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alphas (np.ndarray): Candidate alphas.
             dtype (np.dtype): Working dtype.
             batches (dict[str, int]): Himalaya batch sizes.
         """
-        from himalaya.kernel_ridge import generate_dirichlet_samples
-        from himalaya.ridge import solve_group_ridge_random_search
+        from himalaya.kernel_ridge import (
+            linear_kernel,
+            primal_weights_kernel_ridge,
+            solve_kernel_ridge_cv_eigenvalues,
+        )
         from himalaya.scoring import l2_neg_loss
 
-        # Himalaya's sampler ends with `get_backend().asarray(gammas)`, so the
-        # candidates would otherwise take the dtype and device of whatever
-        # backend happened to be globally active. They are validated and clamped
-        # on the host, so draw them under an explicit numpy scope.
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
+            )
+            kernel = linear_kernel(designs[0])
+            del designs
+            best_alphas, dual_weights, cv_scores = solve_kernel_ridge_cv_eigenvalues(
+                kernel,
+                y_device,
+                alphas=alpha_device,
+                fit_intercept=False,
+                score_func=l2_neg_loss,
+                cv=self._resolved_cv(),
+                local_alpha=self.per_target_alpha,
+                conservative=self.prefer_conservative_alpha,
+                **batches,
+            )
+
+        dual = np.asarray(_to_cpu_numpy(dual_weights), dtype=dtype)
+        with _scoped_himalaya_backend("numpy"):
+            coef = primal_weights_kernel_ridge(dual, spaces[0])
+        self._store_ordinary_state(coef, best_alphas, cv_scores, alphas)
+
+    def _draw_feature_space_candidates(self, n_spaces: int, dtype) -> np.ndarray:
+        """Draw and condition the Dirichlet candidates for a banded search.
+
+        Himalaya's sampler ends with `get_backend().asarray(gammas)`, so the
+        candidates would otherwise take the dtype and device of whatever
+        backend happened to be globally active. They are validated and clamped
+        on the host, so they are drawn under an explicit numpy scope.
+
+        Args:
+            n_spaces (int): Number of feature spaces.
+            dtype (np.dtype): Working dtype, which sets the underflow floor.
+
+        Returns:
+            np.ndarray: `(search_iterations, n_spaces)` weights, each row on
+                the simplex and every entry at least `finfo(dtype).tiny`.
+        """
+        from himalaya.kernel_ridge import generate_dirichlet_samples
+
         with _scoped_himalaya_backend("numpy"):
             candidates = _to_cpu_numpy(
                 generate_dirichlet_samples(
                     n_samples=self.search_iterations,
-                    n_kernels=len(spaces),
+                    n_kernels=n_spaces,
                     concentration=self._concentration_for_himalaya(),
                     random_state=self.random_state,
                 )
             )
-        candidates = _prepare_feature_space_weights(candidates, dtype)
-
-        with _scoped_himalaya_backend(_himalaya_backend_name(self.backend_)):
-            converted = _on_active_backend(*spaces, targets, alphas.astype(dtype))
-            designs, y_device, alpha_device = (
-                converted[:-2],
-                converted[-2],
-                converted[-1],
+        return _prepare_feature_space_weights(candidates, dtype)
+
+    def _store_banded_state(self, deltas, refit_weights, cv_scores, alphas) -> None:
+        """Normalize a banded random search into the fitted state.
+
+        Shared by the primal and kernel forms. Himalaya reports the banded
+        solution as `deltas = log(gamma / alpha)` with each gamma column summing
+        to one, so the simplex weights and the selected alpha both fall out of a
+        log-sum-exp over the spaces. The recovered alpha is snapped back onto
+        the candidate grid, and every array is stored as CPU float64.
+
+        Args:
+            deltas: `(n_spaces, n_targets)` deltas, on any backend.
+            refit_weights: `(n_features, n_targets)` coefficients in original
+                feature coordinates, on any backend.
+            cv_scores: `(search_iterations, n_targets)` scores, on any backend.
+            alphas (np.ndarray): Candidate alphas.
+        """
+        from scipy.special import logsumexp, softmax
+
+        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
+        self.coef_ = np.asarray(_to_cpu_numpy(refit_weights), dtype=np.float64)
+        self.cv_scores_ = np.asarray(_to_cpu_numpy(cv_scores), dtype=np.float64)
+        self.feature_space_weights_ = softmax(deltas, axis=0)
+        selected = _snap_to_grid(np.exp(-logsumexp(deltas, axis=0)), alphas)
+        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+
+    def _fit_banded(self, backend, spaces, targets, alphas, dtype, batches) -> None:
+        """Run the banded random search in the primal form and store the fitted state.
+
+        Args:
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            targets (np.ndarray): `(n_samples, n_targets)` targets.
+            alphas (np.ndarray): Candidate alphas.
+            dtype (np.dtype): Working dtype.
+            batches (dict[str, int]): Himalaya batch sizes.
+        """
+        from himalaya.ridge import solve_group_ridge_random_search
+        from himalaya.scoring import l2_neg_loss
+
+        candidates = self._draw_feature_space_candidates(len(spaces), dtype)
+
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
             )
             deltas, refit_weights, cv_scores = solve_group_ridge_random_search(
                 designs,
@@ -1156,19 +1415,68 @@ class _Ridge:
                 **batches,
             )
 
-        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
-        self.coef_ = np.asarray(_to_cpu_numpy(refit_weights), dtype=np.float64)
-        self.cv_scores_ = np.asarray(_to_cpu_numpy(cv_scores), dtype=np.float64)
+        self._store_banded_state(deltas, refit_weights, cv_scores, alphas)
 
-        # deltas = log(gamma / alpha) with each gamma column summing to one, so
-        # the simplex weights and the selected alpha both fall out of a
-        # log-sum-exp over the spaces.
-        shifted = deltas - deltas.max(axis=0, keepdims=True)
-        weights = np.exp(shifted)
-        self.feature_space_weights_ = weights / weights.sum(axis=0, keepdims=True)
-        log_total = deltas.max(axis=0) + np.log(np.exp(shifted).sum(axis=0))
-        selected = _snap_to_grid(np.exp(-log_total), alphas)
-        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+    def _fit_banded_kernel(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
+        """Run the banded random search in the kernel form and store the fitted state.
+
+        The wide-design counterpart of `_fit_banded`. Himalaya searches over one
+        linear kernel per feature space and returns dual weights, as its own
+        `MultipleKernelRidgeCV` asks it to; its `primal_weights_weighted_kernel_ridge`
+        then recovers `coef_` once on the host from the raw spaces and the
+        deltas. Asking the solver for primal weights instead would rebuild the
+        gamma-scaled design on the device for every improving candidate. The
+        deltas carry the same meaning as in the primal search and go through
+        the same recovery.
+
+        Args:
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            targets (np.ndarray): `(n_samples, n_targets)` targets.
+            alphas (np.ndarray): Candidate alphas.
+            dtype (np.dtype): Working dtype.
+            batches (dict[str, int]): Himalaya batch sizes.
+        """
+        from himalaya.kernel_ridge import (
+            primal_weights_weighted_kernel_ridge,
+            solve_multiple_kernel_ridge_random_search,
+        )
+        from himalaya.scoring import l2_neg_loss
+
+        candidates = self._draw_feature_space_candidates(len(spaces), dtype)
+
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
+            )
+            kernels = _linear_kernels(designs)
+            deltas, refit_weights, cv_scores = (
+                solve_multiple_kernel_ridge_random_search(
+                    kernels,
+                    y_device,
+                    n_iter=candidates,
+                    alphas=alpha_device,
+                    fit_intercept=False,
+                    score_func=l2_neg_loss,
+                    cv=self._resolved_cv(),
+                    return_weights="dual",
+                    local_alpha=self.per_target_alpha,
+                    random_state=self.random_state,
+                    progress_bar=self.progress_bar,
+                    conservative=self.prefer_conservative_alpha,
+                    **batches,
+                )
+            )
+
+        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
+        dual = np.asarray(_to_cpu_numpy(refit_weights), dtype=dtype)
+        with _scoped_himalaya_backend("numpy"):
+            per_space = primal_weights_weighted_kernel_ridge(dual, deltas, spaces)
+        self._store_banded_state(
+            deltas, np.concatenate(per_space, axis=0), cv_scores, alphas
+        )
 
     def _concentration_for_himalaya(self):
         """Return `dirichlet_concentration` in the form Himalaya's sampler takes.

{nltools-0.6.0.dev2.dist-info → nltools-0.6.0.dev4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: nltools
-Version: 0.6.0.dev2
+Version: 0.6.0.dev4
 Summary: A Python package to analyze neuroimaging data
 Project-URL: Homepage, https://nltools.org
 Author-email: "Luke J. Chang" <luke.j.chang@dartmouth.edu>, Eshin Jolly <eshin.jolly@gmail.com>

{nltools-0.6.0.dev2.dist-info → nltools-0.6.0.dev4.dist-info}/RECORD

@@ -11,11 +11,11 @@ nltools/algorithms/decoding.py,sha256=W0GQJNPXRslwZ048K3V6_evHm0A_1x-KbKEFkZT1pg
 nltools/algorithms/neighborhoods.py,sha256=SjnIyZ94k8O4XrDoJJvizJW5nulyV4YofG6XDoCqk1U,6833
 nltools/algorithms/outliers.py,sha256=jysGIr-I22L2dvJsg82qcqXTTYqlcfDVBP4_qHlJbiI,14708
 nltools/algorithms/regression.py,sha256=X0Lssdt89m7torciZn-8PpcqZzuf3nD0kNQtEs6MtXc,4231
-nltools/algorithms/signal.py,sha256=HOcaUR16lNXjcLxu9Jw_sxbFs9y5Q9hmdmqtfdsa-aA,13500
+nltools/algorithms/signal.py,sha256=qZ28zWEeQQJDg59G2Pg4SyCGsi3S_AfNmoKrKjvwnu4,13589
 nltools/algorithms/similarity.py,sha256=Tuy644q16Gw9Xcwuc2Tio7dy-zQjKoCVyBE7AvR9uxM,8655
 nltools/algorithms/validation.py,sha256=XgGq9Ir5S3khLgn4zGVMkvlvhKOd6azHVLziMktJYy0,6091
 nltools/algorithms/alignment/__init__.py,sha256=WgAQFbOZ8Woy3SMhItHpqxVaiUlO2-RPEP0MHgwnE0Y,773
-nltools/algorithms/alignment/procrustes.py,sha256=ErZzr2d8PoAePQDOR9zXAlWlE6IHeQMEEOt-0_2Dh6c,23723
+nltools/algorithms/alignment/procrustes.py,sha256=JAAwZfnmFTg-1rNOAtHrP5jaNp6YVYxzrr0tdkX8AzY,23871
 nltools/algorithms/alignment/srm.py,sha256=en1gSFi2zyz8kqzd2S6i077cViZD6jNCBIEHjbjrAus,31445
 nltools/algorithms/inference/__init__.py,sha256=W0_HFffImX232hGljm5TYTkl3IUy2JbStZdHgtKVfbg,2070
 nltools/algorithms/inference/bootstrap.py,sha256=JHUZiFmI64g5NlBMS2hd_-9vkfUNkApxjoRu29dAayk,55223
@@ -63,19 +63,19 @@ nltools/data/designmatrix/append.py,sha256=Ax_NAnxGZ6UNbFi0D5c-t1vA76MemMTNyv8WS
 nltools/data/designmatrix/diagnostics.py,sha256=HdTTU6GyqwEiYCsDZHHkkA415w_0d2sdjH7m2ELrJ2w,9153
 nltools/data/designmatrix/io.py,sha256=Vp7ALBHuza88ZQClL4cZNTcJeW824sMkO3cJYNSGPZ8,14430
 nltools/data/designmatrix/plotting.py,sha256=Sdxy1H8x-TfAenyY2wqdO8auCnkWlZSyFXhCTVBhlXw,9252
-nltools/data/designmatrix/regressors.py,sha256=sfvh0Csnq4AV0zBOLNV4XcsMS-nb0yv2nOhbgZZbCbU,19692
+nltools/data/designmatrix/regressors.py,sha256=ZMsPBe_1p_tKKUyn2-82kchXwRYIk8uPZRqSjpPqe-8,19749
 nltools/data/designmatrix/transforms.py,sha256=0q3bHit3a9VkysA_LdFqsolu2WsmNCp3xDnGSBagSFY,7752
 nltools/data/designmatrix/utils.py,sha256=MatdEj_CqhYSuNQirez1lhHqiHNTU6A-rC0eH00nSV4,12831
 nltools/data/roc/__init__.py,sha256=kLJAAPwedzIAYl21V3I0XCZp1g_DbNlqtPDfBNEpYts,26804
 nltools/data/simulator/__init__.py,sha256=ODSlYqOCKLSzBzQylQaL-MTu5ioiBfv5p-PI98XvlCc,39084
-nltools/data/simulator/haxby.py,sha256=QkTodkaxVOjUM9KPNDtf_zrmK3RufDOlN5FBA1fwfW0,11004
+nltools/data/simulator/haxby.py,sha256=vopfvOFA32FSWJ3rwBU5BCUNq01AURvEZvhZkmkJHlo,12189
 nltools/io/__init__.py,sha256=oZ5oki0-YZwcTRU0kp9oc2eSCy38o5YRoDQdNi3LgZ8,305
 nltools/io/events.py,sha256=WEm8uLdDsEvzK4Cbb1lG7wBbe86BUJIlUAZHaNnZ8eo,2200
 nltools/io/h5.py,sha256=flN0dwEFfslLihIpl9_BE3nkpoPGoidccLylMkJX34g,13814
 nltools/models/__init__.py,sha256=-Vu2zF2s8UPRRzea9zOu-HJBjfzrHrAxf2knio2fdjA,394
 nltools/models/glm.py,sha256=E8e0OwZvkZBK08GQCO0R8OkJJ46KlzYZ0P62ePUg0ew,22939
 nltools/models/results.py,sha256=TNTNkJBZ5t-4ErsYrPC1k0n_cAvOOKvrIlv3uN20-ZI,3176
-nltools/models/ridge.py,sha256=rNSYAq8ifNqE-3a5uj5SkbX80lLh2OlrOCxYvXCBB4M,53769
+nltools/models/ridge.py,sha256=HMHHmeflYQDJM5VwCxQPmpcGNgGPrD5YbByhZdaaVL4,67893
 nltools/models/validation.py,sha256=RQ5xVzJZSWeIqqy47pNAEdtoYLdIbsK8wnOY0W-v_wc,856
 nltools/plotting/__init__.py,sha256=RXOKXHnHa-QmqdK4ZeApG3ErtPMVBdcWY2iCvUBCuf8,1005
 nltools/plotting/adjacency.py,sha256=X9qHM1GgO1dIkgxsziKyCFhl-EtYeMQICqaehFaOD_s,13077
@@ -90,7 +90,7 @@ nltools/templates/fetch.py,sha256=i3U7B7E-wCwEtL9P3uh4TZWnM6fG1vIANJs8cl0Nnlk,94
 nltools/templates/matching.py,sha256=0JlBLD6598BMprg1CgrW4bDzFSNqSQEKrFkvmxfkD_w,6679
 nltools/templates/paths.py,sha256=STnUDeLa7Nvs1A-k-B8RmiejOQPh-VJGNs-8RQMA0Do,3796
 nltools/templates/registry.py,sha256=zjDQMqUXEDzMj0YZ1-7Pfv6HnA7VSMAfwStSYEe6OsY,636
-nltools-0.6.0.dev2.dist-info/METADATA,sha256=exoa2BeXHR6w8PfTXlemk_gtxzrahNtfQPQmTTppdKc,3735
-nltools-0.6.0.dev2.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
-nltools-0.6.0.dev2.dist-info/licenses/LICENSE,sha256=ZOQ5Dr2VC_ZzSEVTTnkfqCA60sJ9LE2dfV6f-5uHXH0,1108
-nltools-0.6.0.dev2.dist-info/RECORD,,
+nltools-0.6.0.dev4.dist-info/METADATA,sha256=RVm9FlBuraDGGJ2zhs8z1dKzIup8L0TS7vd1g3Tn0_s,3735
+nltools-0.6.0.dev4.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+nltools-0.6.0.dev4.dist-info/licenses/LICENSE,sha256=ZOQ5Dr2VC_ZzSEVTTnkfqCA60sJ9LE2dfV6f-5uHXH0,1108
+nltools-0.6.0.dev4.dist-info/RECORD,,

{nltools-0.6.0.dev2.dist-info → nltools-0.6.0.dev4.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.32.0
+Generator: hatchling 1.32.4
 Root-Is-Purelib: true
 Tag: py3-none-any