nltools 0.6.0.dev1__tar.gz → 0.6.0.dev3__tar.gz

{nltools-0.6.0.dev1 → nltools-0.6.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: nltools
-Version: 0.6.0.dev1
+Version: 0.6.0.dev3
 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.dev1 → nltools-0.6.0.dev3}/nltools/algorithms/alignment/procrustes.py

@@ -109,6 +109,23 @@ def _hyperalign(data, n_iter):
     return aligned, transformation_matrix, template, disparity, scale
 
 
+def _procrustes_similarity(mat1, mat2):
+    """One minus scipy's Procrustes disparity, the only value inference needs.
+
+    Module level so `loky` can pickle it, and returning the scalar rather than
+    scipy's full tuple keeps the two transformed matrices out of the payload
+    every permutation ships back to the parent process.
+
+    Args:
+        mat1 (np.ndarray): Reference matrix, `(n_rows, n_cols)`.
+        mat2 (np.ndarray): Matrix fit to `mat1`, same shape.
+
+    Returns:
+        float: `1 - disparity`, higher meaning more similar.
+    """
+    return 1 - procrust(mat1, mat2)[2]
+
+
 def align(
     data,
     method="deterministic_srm",
@@ -134,7 +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. 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
@@ -153,9 +173,14 @@ def align(
 
     Raises:
         ValueError: If `data` is not a same-typed list, `method` or `axis` is
-            unknown, or `method='procrustes'` is combined with `axis=1` on
+            unknown, `method='procrustes'` is combined with `axis=1` on
             `BrainData` input — that transform spans images on both axes and has
-            no voxel axis to be returned on.
+            no voxel axis to be returned on — or `method='procrustes'` is given
+            `BrainData` subjects with different voxel counts, whose zero-padded
+            results would not fit their own masks. Pass the subjects' `.data`
+            arrays to get the zero-padded result instead; it has no mask that
+            could describe it. An SRM `n_features` above any subject's voxel
+            count also raises.
 
     Examples:
         ```python
@@ -173,7 +198,7 @@ def align(
         ```
     """
 
-    from nltools.data import BrainData, Adjacency
+    from nltools.data import BrainData
 
     if not isinstance(data, list):
         raise ValueError("Make sure you are inputting data is a list.")
@@ -185,6 +210,7 @@ def align(
         )
 
     if isinstance(data[0], BrainData):
+        from nltools.data.braindata.analysis import _brain_result
         from nltools.data.braindata.utils import _result_from_array
 
         data_type = "BrainData"
@@ -211,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
@@ -251,18 +277,22 @@ def align(
 
     if data_type == "BrainData":
         if method == "procrustes":
+            # `_hyperalign` zero-pads every subject's feature axis up to the
+            # widest subject, so a narrower subject's result is wider than its
+            # own mask. `_brain_result` refuses that rather than returning an
+            # object whose `to_nifti` fails later.
             out["transformed"] = [
-                _result_from_array(source, values.T, rows="preserve")
+                _brain_result(source, values.T, "transformed", rows="preserve")
                 for source, values in zip(sources, out["transformed"])
             ]
-            out["common_model"] = _result_from_array(
-                sources[0], out["common_model"], rows="clear"
+            out["common_model"] = _brain_result(
+                sources[0], out["common_model"], "common_model", rows="clear"
             )
             # `_hyperalign` already returns these in the
             # `transformed = original @ T` orientation, and they are square on
             # the voxel axis, so unlike the SRM matrices they are wrapped as-is.
             out["transformation_matrix"] = [
-                _result_from_array(source, values, rows="clear")
+                _brain_result(source, values, "transformation_matrix", rows="clear")
                 for source, values in zip(sources, out["transformation_matrix"])
             ]
         else:
@@ -281,67 +311,28 @@ def align(
     #   BrainData: (timepoints, voxels)
     #   numpy: (voxels, timepoints)
 
-    a = Adjacency()
+    # For procrustes, transformed holds BrainData objects; extract .data.
+    # For SRM methods it already holds numpy arrays.
+    transformed_arrays = [
+        x.data if isinstance(x, BrainData) else x for x in out["transformed"]
+    ]
+
+    # Put every case in one orientation, (aligned units, observations), so the
+    # correlation below reads the same way whatever came in. BrainData results
+    # are (timepoints, voxels) and numpy results are (voxels, timepoints), so
+    # exactly one of the two needs a transpose for a given axis.
+    if (data_type == "BrainData") == (axis == 0):
+        units = [x.T for x in transformed_arrays]
+    else:
+        units = transformed_arrays
 
-    if data_type == "BrainData":
-        # BrainData transformed shape: (timepoints, voxels)
-        # For procrustes, transformed contains BrainData objects; extract .data
-        # For SRM methods, transformed contains numpy arrays after the .T
-        transformed_arrays = [
-            x.data if isinstance(x, BrainData) else x for x in out["transformed"]
-        ]
-        if axis == 0:
-            # Aligned timepoints → ISC per voxel (correlation over time)
-            n_isc = transformed_arrays[0].shape[1]  # n_voxels
-            for v in range(n_isc):
-                # Extract timecourse for voxel v from each subject
-                isc_data = np.array([x[:, v] for x in transformed_arrays])
-                a = a.append(
-                    Adjacency(
-                        1 - pairwise_distances(isc_data, metric="correlation"),
-                        matrix_type="similarity",
-                    )
-                )
-        else:  # axis == 1
-            # Aligned voxels → ISC per timepoint (spatial correlation)
-            n_isc = transformed_arrays[0].shape[0]  # n_timepoints
-            for t in range(n_isc):
-                # Extract spatial pattern at timepoint t from each subject
-                isc_data = np.array([x[t, :] for x in transformed_arrays])
-                a = a.append(
-                    Adjacency(
-                        1 - pairwise_distances(isc_data, metric="correlation"),
-                        matrix_type="similarity",
-                    )
-                )
-    else:  # numpy
-        # numpy transformed shape: (voxels, timepoints)
-        if axis == 0:
-            # Aligned timepoints → ISC per voxel (correlation over time)
-            n_isc = out["transformed"][0].shape[0]  # n_voxels
-            for v in range(n_isc):
-                # Extract timecourse for voxel v from each subject
-                isc_data = np.array([x[v, :] for x in out["transformed"]])
-                a = a.append(
-                    Adjacency(
-                        1 - pairwise_distances(isc_data, metric="correlation"),
-                        matrix_type="similarity",
-                    )
-                )
-        else:  # axis == 1
-            # Aligned voxels → ISC per timepoint (spatial correlation)
-            n_isc = out["transformed"][0].shape[1]  # n_timepoints
-            for t in range(n_isc):
-                # Extract spatial pattern at timepoint t from each subject
-                isc_data = np.array([x[:, t] for x in out["transformed"]])
-                a = a.append(
-                    Adjacency(
-                        1 - pairwise_distances(isc_data, metric="correlation"),
-                        matrix_type="similarity",
-                    )
-                )
-
-    out["isc"] = dict(zip(np.arange(n_isc), a.mean(axis=1)))
+    upper_triangle = np.triu_indices(len(units), k=1)
+    out["isc"] = {}
+    for unit in range(units[0].shape[0]):
+        similarity = 1 - pairwise_distances(
+            np.array([x[unit] for x in units]), metric="correlation"
+        )
+        out["isc"][unit] = float(np.nanmean(similarity[upper_triangle]))
 
     return out
 
@@ -478,14 +469,12 @@ def procrustes_distance(
     # the SAME scale. Previously the observed disparity was compared against a
     # null of similarities, inverting the scales and yielding p ~ 1 for
     # near-identical matrices.
-    _, _, disparity = procrust(mat1, mat2)
-    observed_similarity = 1 - disparity
+    observed_similarity = _procrustes_similarity(mat1, mat2)
 
-    null_disparities = Parallel(n_jobs=n_jobs)(
-        delayed(procrust)(random_state.permutation(mat1), mat2)
+    null_similarity = Parallel(n_jobs=n_jobs)(
+        delayed(_procrustes_similarity)(random_state.permutation(mat1), mat2)
         for _ in range(n_permute)
     )
-    null_similarity = [1 - x[2] for x in null_disparities]
 
     # Use _compute_pvalue from inference module (signature: obs_stat, null_dist, tail)
     stats = {"similarity": float(observed_similarity)}
@@ -521,7 +510,8 @@ def align_states(
             reordered data. Defaults to False.
         replace_zero_variance (bool): Replace zero-variance columns with uniform
             random numbers before computing distances; avoids NaNs with the
-            correlation metric. Defaults to False.
+            correlation metric. Integer inputs are converted to float so the
+            replacement noise survives. Defaults to False.
 
     Returns:
         np.ndarray: If `return_index=False` (default), `target[:, remapping]` — the
@@ -541,12 +531,18 @@ def align_states(
         Prevents NaN values when correlation-based distance metrics encounter
         constant columns.
 
+        The array is converted to float first: writing U(0, 1) draws into an
+        integer array truncates every one of them to zero, leaving the constant
+        column constant and the correlation distance NaN.
+
         Args:
             data (np.ndarray): 2-D array whose columns are checked for zero variance.
 
         Returns:
-            np.ndarray: Array with zero-variance columns replaced by U(0, 1) values.
+            np.ndarray: Float array with zero-variance columns replaced by
+                U(0, 1) values.
         """
+        data = np.asarray(data, dtype=float)
         if np.any(data.std(axis=0) == 0):
             for i in np.where(data.std(axis=0) == 0)[0]:
                 data[:, i] = np.random.uniform(low=0, high=1, size=data.shape[0])

{nltools-0.6.0.dev1 → nltools-0.6.0.dev3}/nltools/algorithms/alignment/srm.py

@@ -104,6 +104,62 @@ def _init_w_transforms(
     return w, voxels
 
 
+def _check_n_features(X: list[np.ndarray], n_features: int) -> None:
+    """Reject a feature count no subject's voxel dimension can support.
+
+    `_init_w_transforms` takes a reduced QR, which silently returns
+    `min(voxels, n_features)` columns, so an oversized request would otherwise
+    produce a model of a dimension the caller never asked for.
+
+    Note:
+        "Voxels" is this module's name for the first axis throughout, and the
+        message follows it. `align(axis=1)` transposes before fitting, so on
+        that path the first axis is timepoints and the message still says
+        voxels.
+
+    Args:
+        X (list[np.ndarray]): One (voxels_i, samples) array per subject.
+        n_features (int): Requested number of shared features.
+
+    Raises:
+        ValueError: If `n_features` is not a positive integer, or exceeds any
+            subject's voxel count.
+    """
+    if not isinstance(n_features, (int, np.integer)) or n_features < 1:
+        raise ValueError(f"n_features must be a positive integer, got {n_features!r}.")
+    for subject, data in enumerate(X):
+        if data is None:
+            continue
+        if data.shape[0] < n_features:
+            raise ValueError(
+                f"subject {subject} has {data.shape[0]} voxels, too few to "
+                f"support {n_features} features. Lower n_features to at most "
+                "the smallest subject's voxel count."
+            )
+
+
+def _update_transform_subject(Xi: np.ndarray, S: np.ndarray) -> np.ndarray:
+    """Update the mapping $W_i$ for one subject.
+
+    Solves the orthogonal Procrustes problem
+    $\\min ||X_i - W_i S||_F^2$ subject to $W_i^T W_i = I$: with the SVD
+    $U \\Sigma V^T = X_i S^T$, the optimum is $W_i = U V^T$.
+
+    Args:
+        Xi (np.ndarray): The subject's data $X_i$, shape (voxels, timepoints).
+        S (np.ndarray): The shared response, shape (n_features, timepoints).
+
+    Returns:
+        np.ndarray: The orthogonal transform $W_i$, shape (voxels, n_features).
+    """
+    # Compute cross-covariance: X_i S^T
+    A = Xi.dot(S.T)
+    # Solve the Procrustes problem via SVD
+    # Optimal orthogonal transform: W_i = U V^T where A = U Σ V^T
+    U, _, V = np.linalg.svd(A, full_matrices=False)
+    return U.dot(V)
+
+
 class _SRM(BaseEstimator, TransformerMixin):
     """Probabilistic Shared Response Model (SRM).
 
@@ -142,7 +198,7 @@ class _SRM(BaseEstimator, TransformerMixin):
     Examples:
         ```python
         import numpy as np
-        from nltools.algorithms import SRM
+        from nltools.algorithms.alignment import _SRM
 
         data = [np.random.randn(100, 50) for _ in range(3)]  # 3 subjects
 
@@ -173,6 +229,11 @@ class _SRM(BaseEstimator, TransformerMixin):
 
         Returns:
             _SRM: Fitted model (`self`).
+
+        Raises:
+            ValueError: If there are fewer than two subjects, the subjects
+                disagree on sample count, or `n_features` exceeds any subject's
+                voxel count.
         """
         logger.info("Starting Probabilistic SRM")
 
@@ -197,6 +258,10 @@ class _SRM(BaseEstimator, TransformerMixin):
                 f"Different number of samples between subjects: {sample_counts}."
             )
 
+        # After the sample-count check, so input wrong on both axes reports the
+        # mismatched samples first.
+        _check_n_features(X, self.n_features)
+
         # Validate all data is finite
         for subject in range(number_subjects):
             if X[subject] is not None:
@@ -328,28 +393,6 @@ class _SRM(BaseEstimator, TransformerMixin):
 
         return loglikehood
 
-    @staticmethod
-    def _update_transform_subject(Xi, S):
-        """Update the mapping $W_i$ for one subject.
-
-        Solves the orthogonal Procrustes problem
-        $\\min ||X_i - W_i S||_F^2$ subject to $W_i^T W_i = I$: with the SVD
-        $U \\Sigma V^T = X_i S^T$, the optimum is $W_i = U V^T$.
-
-        Args:
-            Xi (np.ndarray): The subject's data $X_i$, shape (voxels, timepoints).
-            S (np.ndarray): The shared response, shape (n_features, timepoints).
-
-        Returns:
-            np.ndarray: The orthogonal transform $W_i$, shape (voxels, n_features).
-        """
-        # Compute cross-covariance: X_i S^T
-        A = Xi.dot(S.T)
-        # Solve the Procrustes problem via SVD
-        # Optimal orthogonal transform: W_i = U V^T where A = U Σ V^T
-        U, _, V = np.linalg.svd(A, full_matrices=False)
-        return U.dot(V)
-
     def transform_subject(self, X: np.ndarray) -> np.ndarray:
         """Transform a new subject using the existing model.
 
@@ -373,7 +416,7 @@ class _SRM(BaseEstimator, TransformerMixin):
                 "The number of timepoints(TRs) does not match the one in the model."
             )
 
-        w = self._update_transform_subject(X, self.s_)
+        w = _update_transform_subject(X, self.s_)
 
         return w
 
@@ -533,7 +576,7 @@ class _DetSRM(BaseEstimator, TransformerMixin):
     Examples:
         ```python
         import numpy as np
-        from nltools.algorithms import DetSRM
+        from nltools.algorithms.alignment import _DetSRM
 
         data = [np.random.randn(100, 50) for _ in range(3)]  # 3 subjects
 
@@ -563,6 +606,11 @@ class _DetSRM(BaseEstimator, TransformerMixin):
 
         Returns:
             _DetSRM: Fitted model (`self`).
+
+        Raises:
+            ValueError: If there are fewer than two subjects, the subjects
+                disagree on sample count, or `n_features` exceeds any subject's
+                voxel count.
         """
         logger.info("Starting Deterministic SRM")
 
@@ -587,6 +635,10 @@ class _DetSRM(BaseEstimator, TransformerMixin):
             if X[subject].shape[1] != number_trs:
                 raise ValueError("Different number of samples between subjects.")
 
+        # After the sample-count check, so input wrong on both axes reports the
+        # mismatched samples first.
+        _check_n_features(X, self.n_features)
+
         # Run SRM
         self.w_, self.s_ = self._srm(X)
 
@@ -655,28 +707,6 @@ class _DetSRM(BaseEstimator, TransformerMixin):
 
         return s
 
-    @staticmethod
-    def _update_transform_subject(Xi, S):
-        """Update the mapping $W_i$ for one subject.
-
-        Solves the orthogonal Procrustes problem
-        $\\min ||X_i - W_i S||_F^2$ subject to $W_i^T W_i = I$: with the SVD
-        $U \\Sigma V^T = X_i S^T$, the optimum is $W_i = U V^T$.
-
-        Args:
-            Xi (np.ndarray): The subject's data $X_i$, shape (voxels, timepoints).
-            S (np.ndarray): The shared response, shape (n_features, timepoints).
-
-        Returns:
-            np.ndarray: The orthogonal transform $W_i$, shape (voxels, n_features).
-        """
-        # Compute cross-covariance: X_i S^T
-        A = Xi.dot(S.T)
-        # Solve the Procrustes problem via SVD
-        # Optimal orthogonal transform: W_i = U V^T where A = U Σ V^T
-        U, _, V = np.linalg.svd(A, full_matrices=False)
-        return U.dot(V)
-
     def transform_subject(self, X: np.ndarray) -> np.ndarray:
         """Transform a new subject using the existing model.
 
@@ -700,7 +730,7 @@ class _DetSRM(BaseEstimator, TransformerMixin):
                 "The number of timepoints(TRs) does not match the one in the model."
             )
 
-        w = self._update_transform_subject(X, self.s_)
+        w = _update_transform_subject(X, self.s_)
 
         return w
 

{nltools-0.6.0.dev1 → nltools-0.6.0.dev3}/nltools/algorithms/backends.py

@@ -654,6 +654,10 @@ def _ridge_bootstrap_batch_size(
 #: every summary payload is converted to CPU float64 before it is retained.
 _BOOTSTRAP_OUTPUT_ITEMSIZE = 8
 
+#: Bytes per resampling index. `_generate_bootstrap_indices` returns an int64
+#: matrix, and every CPU worker's closure captures it.
+_BOOTSTRAP_INDEX_ITEMSIZE = 8
+
 #: Output-sized arrays a bootstrap run always holds beyond its retained
 #: replicates: the two Welford accumulators (running mean and running sum of
 #: squared deviations) and the four `BootstrapResult` summary payloads.
@@ -734,15 +738,19 @@ def _bootstrap_output_bytes(
     *,
     confidence_level: float,
     return_samples: bool,
+    n_obs: int,
     n_workers: int = 1,
 ) -> int:
-    """Bytes a bootstrap run must hold for its retained output.
+    """Bytes a bootstrap run must hold for its retained output and its indices.
 
     Charges eight bytes for every output-sized array a run holds at once: the
     two bounded tails, the replicates buffered before the next flush and the
     two temporaries that flush builds, one dispatch window of in-flight
     replicates, every replicate when `return_samples=True`, and the two Welford
-    accumulators plus the four summary payloads.
+    accumulators plus the four summary payloads. On top of that it charges the
+    resampling index matrix twice: the run retains one `(n_samples, n_obs)`
+    int64 array, and building it holds the per-draw vectors alongside the
+    stacked result.
 
     Args:
         output_shape (tuple[int, ...]): Shape of one replicate's output.
@@ -750,6 +758,8 @@ def _bootstrap_output_bytes(
         confidence_level (float): Interval confidence level, which sets the
             retained tail size.
         return_samples (bool): Whether the complete distribution is retained.
+        n_obs (int): Observations resampled per replicate, which sizes the
+            index matrix.
         n_workers (int): Planned CPU worker count, which sets the dispatch
             window. Defaults to 1 (the GPU driver budgets its own batch through
             `_ridge_bootstrap_batch_size` instead).
@@ -770,7 +780,8 @@ def _bootstrap_output_bytes(
         + (int(n_samples) if return_samples else 0)
         + _BOOTSTRAP_FIXED_OUTPUT_ARRAYS
     )
-    return output_size * arrays * _BOOTSTRAP_OUTPUT_ITEMSIZE
+    index_bytes = 2 * int(n_obs) * int(n_samples) * _BOOTSTRAP_INDEX_ITEMSIZE
+    return output_size * arrays * _BOOTSTRAP_OUTPUT_ITEMSIZE + index_bytes
 
 
 def _bootstrap_memory_preflight(
@@ -779,6 +790,7 @@ def _bootstrap_memory_preflight(
     *,
     confidence_level: float,
     return_samples: bool,
+    n_obs: int,
     n_workers: int = 1,
     memory_budget_gb: float | None = None,
     backend: "_Backend | None" = None,
@@ -795,6 +807,8 @@ def _bootstrap_memory_preflight(
         n_samples (int): Number of bootstrap replicates.
         confidence_level (float): Interval confidence level.
         return_samples (bool): Whether the complete distribution is retained.
+        n_obs (int): Observations resampled per replicate, which sizes the
+            index matrix.
         n_workers (int): Planned CPU worker count, which sets the dispatch
             window. Defaults to 1.
         memory_budget_gb (float | None): Explicit budget in GB, or None to
@@ -813,6 +827,7 @@ def _bootstrap_memory_preflight(
         n_samples,
         confidence_level=confidence_level,
         return_samples=return_samples,
+        n_obs=n_obs,
         n_workers=n_workers,
     )
     required_gb = required_bytes / 1e9
@@ -925,7 +940,9 @@ def _compute_oom_safe(fn, *arrays, min_chunk: int = 1):
     a numpy array whose axis 0 corresponds row-for-row to its inputs. On a
     device OOM the cache is emptied, the arrays are split in half along
     axis 0, and the halves are retried recursively; partial results are
-    concatenated along axis 0.
+    concatenated along axis 0. Recovery happens after leaving the exception
+    handler, so the failed call's allocations are released before the smaller
+    retries ask for them.
 
     Because splitting reuses the *already generated* inputs rather than
     re-drawing them, recovery never changes which permutations a seeded
@@ -954,17 +971,22 @@ def _compute_oom_safe(fn, *arrays, min_chunk: int = 1):
     except Exception as exc:
         if not _is_oom_error(exc):
             raise
-        _empty_device_cache()
         if n <= min_chunk:
+            _empty_device_cache()
             raise MemoryError(
                 f"Device out of memory even for a single item (chunk of {n}). "
                 "Reduce the problem size, lower max_gpu_memory_gb elsewhere on "
                 "the device, or use device='cpu'."
             ) from exc
-        mid = n // 2
-        left = _compute_oom_safe(fn, *(a[:mid] for a in arrays), min_chunk=min_chunk)
-        right = _compute_oom_safe(fn, *(a[mid:] for a in arrays), min_chunk=min_chunk)
-        return np.concatenate([left, right], axis=0)
+
+    # Outside the handler, where Python's implicit `del exc` has dropped the
+    # traceback: the failed call's frame — and the device allocations its
+    # locals held — are gone before the cache is emptied and the halves retried.
+    _empty_device_cache()
+    mid = n // 2
+    left = _compute_oom_safe(fn, *(a[:mid] for a in arrays), min_chunk=min_chunk)
+    right = _compute_oom_safe(fn, *(a[mid:] for a in arrays), min_chunk=min_chunk)
+    return np.concatenate([left, right], axis=0)
 
 
 # ----------------------------------------------------------------------