nltools 0.6.0.dev0__py3-none-any.whl

nltools/data/designmatrix/utils.py

@@ -0,0 +1,350 @@
+"""Shared helpers for DesignMatrix submodules.
+
+These are internal utilities used by the facade and submodules — not part of the
+public API.
+"""
+
+from __future__ import annotations
+
+import functools
+from copy import deepcopy
+import re
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+from nltools.data.ownership import _copy_frame
+
+
+if TYPE_CHECKING:
+    from nltools.data.designmatrix import DesignMatrix
+
+
+RESERVED_PREFIX = ".nl_"
+"""Prefix marking a column name generated by nltools rather than the user.
+
+Every column nltools invents — polynomial drift (`.nl_poly_0`), DCT cosine
+bases (`.nl_cosine_1`), spike indicators (`.nl_global_spike1`), and the
+run-separated variants produced by a multi-run `DesignMatrix.append` — carries
+this prefix. Machinery that needs to recognize its own columns keys on the
+prefix, never on heuristics over user-controlled names (underscore counts,
+substring matches), so users may name their own regressors anything without
+colliding with nltools internals.
+"""
+
+_RUN_SEPARATED_RE = re.compile(re.escape(RESERVED_PREFIX) + r"r(\d+)_(.+)")
+
+
+def _reserved_name(base: str) -> str:
+    """Build a generated column name inside the reserved namespace.
+
+    Args:
+        base: Name without the reserved prefix, e.g. ``'poly_0'``.
+
+    Returns:
+        str: `base` prefixed with `RESERVED_PREFIX`, idempotently — a name
+            that already carries the prefix is returned unchanged.
+    """
+    return base if _is_reserved_name(base) else f"{RESERVED_PREFIX}{base}"
+
+
+def _is_reserved_name(name: str) -> bool:
+    """Return True if ``name`` is in the nltools-generated column namespace."""
+    return name.startswith(RESERVED_PREFIX)
+
+
+def _strip_reserved_prefix(name: str) -> str:
+    """Return ``name`` without its reserved prefix (a no-op if it has none)."""
+    return name[len(RESERVED_PREFIX) :] if _is_reserved_name(name) else name
+
+
+def _run_separated_name(run_idx: int, name: str) -> str:
+    """Build the run-separated variant of a column name.
+
+    Run separation is an nltools-generated naming decision, so the result
+    always lands in the reserved namespace regardless of whether the source
+    column was user-named (``motion_x`` → ``.nl_r0_motion_x``) or already
+    generated (``.nl_poly_0`` → ``.nl_r0_poly_0``; prefixes never stack).
+
+    Args:
+        run_idx: Zero-based run index.
+        name: Column name to separate.
+
+    Returns:
+        str: ``.nl_r{run_idx}_{base}``.
+    """
+    return f"{RESERVED_PREFIX}r{run_idx}_{_strip_reserved_prefix(name)}"
+
+
+def _parse_run_separated(name: str) -> tuple[int, str] | None:
+    """Split a run-separated column name into its run index and base name.
+
+    Args:
+        name: Column name to parse.
+
+    Returns:
+        tuple[int, str] | None: `(run_idx, base)` for a run-separated name (e.g.
+            `'.nl_r1_poly_0'` → `(1, 'poly_0')`), else None.
+    """
+    match = _RUN_SEPARATED_RE.fullmatch(name)
+    return (int(match.group(1)), match.group(2)) if match else None
+
+
+# Row selectors operate on a temporary column for zero-column designs.
+_ROW_SELECTION = frozenset({"head", "tail", "slice", "filter", "limit"})
+_MUTATORS = frozenset({"insert_column", "replace_column", "drop_in_place", "extend"})
+
+
+def _design_from_generated(
+    frame: pl.DataFrame,
+    *,
+    sampling_freq: float | None = None,
+    n_rows: int | None = None,
+) -> DesignMatrix:
+    """Build a DesignMatrix whose every column is an nltools-generated confound.
+
+    Callers outside this package (`find_spikes`) name their columns plainly and
+    hand the frame here; the reserved prefix is applied in the one package that
+    owns the namespace. Every renamed column is marked a confound.
+
+    Args:
+        frame (pl.DataFrame): Generated columns under their plain names.
+        sampling_freq (float | None): Sampling frequency in Hz, or None.
+        n_rows (int | None): Row count to record when `frame` has no columns.
+
+    Returns:
+        DesignMatrix: The design, with every column in the reserved namespace
+            and listed in `confounds`.
+    """
+    from nltools.data.designmatrix import DesignMatrix
+
+    renamed = frame.rename({name: _reserved_name(name) for name in frame.columns})
+    return DesignMatrix(
+        renamed,
+        sampling_freq=sampling_freq,
+        confounds=list(renamed.columns),
+        n_rows=n_rows,
+    )
+
+
+def _effective_frame(dm: DesignMatrix) -> pl.DataFrame:
+    """Represent recorded observations during operations on a column-less frame."""
+    if dm.data.width == 0 and dm._n_rows is not None:
+        return pl.DataFrame({"": pl.repeat(None, dm.shape[0], eager=True)})
+    return dm.data
+
+
+def _replacement_names(frame: pl.DataFrame, exprs, named_exprs) -> list[str]:
+    """Ask Polars which columns the supplied expressions produce."""
+    return frame.lazy().select(*exprs, **named_exprs).collect_schema().names()
+
+
+# Drift terms generated by add_poly / add_dct_basis, matched on the base name
+# left after the run prefix is stripped (e.g. '.nl_r1_poly_0' -> 'poly_0').
+_DRIFT_BASE_RE = re.compile(r"(?:poly|cosine)_\d+")
+_INTERCEPT_BASE_RE = re.compile(r"(?:poly|cosine)_0")
+
+
+def _is_generated_intercept(name: str) -> bool:
+    """Return True if ``name`` is an intercept column nltools generated.
+
+    Covers the zeroth-order drift terms from `add_poly` / `add_dct_basis`
+    (``.nl_poly_0`` / ``.nl_cosine_0``) and their run-separated variants. Both
+    are all-ones columns, so anything computing a correlation matrix has to
+    drop them. Keyed on the reserved namespace: a user column is never an
+    intercept by this definition, however it happens to be named.
+    """
+    if not _is_reserved_name(name):
+        return False
+    parsed = _parse_run_separated(name)
+    base = parsed[1] if parsed is not None else _strip_reserved_prefix(name)
+    return _INTERCEPT_BASE_RE.fullmatch(base) is not None
+
+
+def _has_run_separated_drift(dm: DesignMatrix) -> bool:
+    """Return True if ``dm`` carries per-run polynomial or cosine drift terms.
+
+    Adding a global drift term to a design that already models drift per run
+    is ambiguous, so both `add_poly` and `add_dct_basis` refuse it. Detection
+    keys on the reserved namespace nltools controls (``.nl_r{run}_poly_{i}`` /
+    ``.nl_r{run}_cosine_{i}``), so a user confound is never mistaken for one
+    however it is named.
+    """
+    for col in dm.confounds or []:
+        parsed = _parse_run_separated(col)
+        if parsed is not None and _DRIFT_BASE_RE.fullmatch(parsed[1]):
+            return True
+    return False
+
+
+def _is_column_selection(value) -> bool:
+    """Recognize plain Polars column selectors without interpreting expressions."""
+    if isinstance(value, str):
+        return True
+    if isinstance(value, pl.Expr):
+        return value.meta.is_column_selection()
+    if isinstance(value, (list, tuple)):
+        return all(_is_column_selection(item) for item in value)
+    return False
+
+
+def _df_passthrough(dm: DesignMatrix, name: str):
+    """Forward Polars operations with explicit row, column and mutation context."""
+    attr = getattr(dm.data, name)
+    if not callable(attr):
+        return attr
+
+    @functools.wraps(attr)  # nosemgrep: kwargs-internal-forwarding  # Polars adapter
+    def wrapped(*args, **kwargs):
+        mutation = name in _MUTATORS or (
+            name in {"hstack", "vstack", "shrink_to_fit"}
+            and kwargs.get("in_place", False)
+        )
+        frame = _copy_frame(dm.data) if mutation else dm.data
+        if name in _ROW_SELECTION:
+            frame = _effective_frame(dm)
+        result = getattr(frame, name)(*args, **kwargs)
+        if (
+            name in {"insert_column", "hstack"}
+            and dm.data.width == 0
+            and dm._n_rows is not None
+        ):
+            populated = frame if mutation else result
+            if populated.height != dm.shape[0]:
+                raise ValueError(
+                    "Added columns must match the recorded number of rows."
+                )
+        operation = "unknown"
+        rename = None
+        replaced = None
+        if name in _ROW_SELECTION:
+            operation = "preserve"
+        elif (
+            name == "select"
+            and not kwargs
+            and all(_is_column_selection(arg) for arg in args)
+        ):
+            operation = "preserve"
+        elif name == "rename":
+            operation = "rename"
+            rename = dict(zip(dm.columns, result.columns))
+        elif name == "replace_column":
+            operation = "replace"
+            index = args[0] if args else kwargs["index"]
+            replaced = [dm.columns[index], frame.columns[index]]
+        elif name in {"insert_column", "hstack"}:
+            operation = "preserve"
+        elif name in {"drop_in_place", "shrink_to_fit"}:
+            operation = "preserve"
+        if mutation:
+            updated = _copy_with(dm, frame, operation=operation, replaced=replaced)
+            dm.__dict__.update(updated.__dict__)
+            return dm if result is frame else result
+        if isinstance(result, pl.DataFrame):
+            n_rows = result.height
+            if name in _ROW_SELECTION and dm.data.width == 0:
+                result = pl.DataFrame()
+            elif name == "select" and operation == "preserve" and result.width == 0:
+                n_rows = dm.shape[0]
+            return _copy_with(
+                dm, result, operation=operation, rename=rename, n_rows=n_rows
+            )
+        if isinstance(result, pl.Series):
+            return _copy_frame(result.to_frame()).to_series()
+        return result
+
+    return wrapped
+
+
+def _copy_with(
+    dm: DesignMatrix,
+    new_df: pl.DataFrame,
+    *,
+    operation: str = "preserve",
+    rename: dict | None = None,
+    replaced: list[str] | None = None,
+    sampling_freq=...,
+    convolved: list[str] | None = None,
+    confounds: list[str] | None = None,
+    multi: bool | None = None,
+    n_rows: int | None = None,
+    run_count: int | None = None,
+) -> DesignMatrix:
+    """Own a transformed frame and apply its caller-established metadata policy."""
+    from nltools.data.designmatrix import DesignMatrix
+
+    metadata = _get_metadata(dm)
+    if operation == "unknown":
+        metadata.update(
+            sampling_freq=None, convolved=[], confounds=[], multi=False, run_count=0
+        )
+    elif operation == "rename":
+        for key in ("convolved", "confounds"):
+            metadata[key] = [(rename or {}).get(c, c) for c in metadata[key]]
+    elif operation == "replace":
+        metadata["convolved"] = [
+            c for c in metadata["convolved"] if c not in (replaced or [])
+        ]
+    if sampling_freq is not ...:
+        metadata["sampling_freq"] = sampling_freq
+    if convolved is not None:
+        metadata["convolved"] = convolved
+    if confounds is not None:
+        metadata["confounds"] = confounds
+    if multi is not None:
+        metadata["multi"] = multi
+    if run_count is not None:
+        metadata["run_count"] = run_count
+    for key in ("convolved", "confounds"):
+        metadata[key] = [c for c in metadata[key] if c in new_df.columns]
+    new = DesignMatrix.__new__(DesignMatrix)
+    memo = {id(dm): new}
+    new.data = _copy_frame(new_df, memo)
+    new.sampling_freq = metadata["sampling_freq"]
+    new._convolved = deepcopy(metadata["convolved"], memo)
+    new._confounds = deepcopy(metadata["confounds"], memo)
+    new.multi = metadata["multi"]
+    new._run_count = metadata["run_count"]
+    new._n_rows = (
+        (dm.shape[0] if n_rows is None else n_rows) if new_df.width == 0 else None
+    )
+    return new
+
+
+def _get_metadata(dm: DesignMatrix) -> dict:
+    """Extract metadata as a dict (for copying).
+
+    Args:
+        dm (DesignMatrix): DesignMatrix instance.
+
+    Returns:
+        dict: Dictionary with keys 'sampling_freq', 'convolved', 'confounds',
+            'multi', 'n_rows'.
+    """
+    return {
+        "sampling_freq": dm.sampling_freq,
+        "convolved": dm.convolved.copy(),
+        "confounds": dm.confounds.copy(),
+        "multi": dm.multi,
+        "n_rows": dm._n_rows,
+        "run_count": dm._run_count,
+    }
+
+
+def _get_data_columns(dm: DesignMatrix, exclude_confounds: bool = True) -> list[str]:
+    """Get column names, optionally excluding confound regressors.
+
+    Used wherever experimental regressors must be distinguished from
+    nuisance/confound columns (polynomial drift, DCT cosines, motion, etc.).
+
+    Args:
+        dm (DesignMatrix): DesignMatrix instance.
+        exclude_confounds (bool): If True, exclude nuisance columns tracked in
+            ``dm.confounds`` from the result. Default: True.
+
+    Returns:
+        list[str]: Column names (excluding confounds if requested).
+    """
+    if exclude_confounds and dm.confounds:
+        return [col for col in dm.columns if col not in dm.confounds]
+    return list(dm.columns)

nltools/data/ownership.py

@@ -0,0 +1,129 @@
+"""Copying that leaves the clone owning its own buffers.
+
+A data class holds polars frames whose storage may be a view onto a NumPy
+array the user still holds, and frames whose `pl.Object` cells are Python
+objects shared with the original. Copying either naively hands the clone a
+buffer or a cell somebody else can mutate. `_copy_frame` detaches one frame;
+`_copy_graph` walks a whole object's `__dict__` and detaches every frame it
+finds, preserving the aliases inside that graph through a shared memo.
+
+The two frame copiers are deliberately not the same function: `_copy_frame`
+re-`gather`s every non-Object series so a `DesignMatrix` clone owns its
+numeric buffers outright, while `_copy_object_frames` clones the frame and
+rewrites only its `pl.Object` columns, which is what `BrainData` and
+`Adjacency` metadata need.
+"""
+
+from copy import deepcopy
+
+import polars as pl
+
+
+def _copy_frame(frame: pl.DataFrame, memo: dict | None = None) -> pl.DataFrame:
+    """Detach frame storage and Python Object cells with a shared copy memo."""
+    if memo is None:
+        memo = {}
+    if id(frame) in memo:
+        return memo[id(frame)]
+    frames = []
+    visited = set()
+
+    def discover(value):
+        if id(value) in visited or id(value) in memo:
+            return
+        visited.add(id(value))
+        if isinstance(value, pl.DataFrame):
+            frames.append(value)
+            for series in value:
+                if series.dtype == pl.Object:
+                    for cell in series:
+                        discover(cell)
+        elif isinstance(value, dict):
+            for key, item in value.items():
+                discover(key)
+                discover(item)
+        elif isinstance(value, (list, tuple)):
+            for item in value:
+                discover(item)
+
+    discover(frame)
+    for source in frames:
+        memo[id(source)] = source.clone()
+    for source in frames:
+        for index, series in enumerate(source):
+            if series.dtype == pl.Object:
+                detached = pl.Series(
+                    series.name,
+                    [deepcopy(cell, memo) for cell in series],
+                    dtype=pl.Object,
+                )
+            else:
+                # Gather owns buffers even when the frame wraps a NumPy view.
+                detached = series.gather(pl.int_range(0, len(series), eager=True))
+            memo[id(source)].replace_column(index, detached)
+    return memo[id(frame)]
+
+
+def _copy_graph(source, *, memo=None, exclude=(), replacements=None):
+    """Copy one retained object graph, preserving its internal aliases."""
+    if memo is None:
+        memo = {}
+    if id(source) in memo:
+        return memo[id(source)]
+    new = type(source).__new__(type(source))
+    memo[id(source)] = new
+    values = {
+        key: value for key, value in source.__dict__.items() if key not in exclude
+    }
+    if replacements is not None:
+        values.update(replacements)
+    _copy_object_frames(values, memo)
+    for key, value in values.items():
+        setattr(new, key, deepcopy(value, memo))
+    return new
+
+
+def _copy_object_frames(values, memo):
+    """Prepare Polars Object cells for deepcopy without sharing Python objects."""
+    frames = []
+    seen = set()
+
+    def discover(value):
+        if id(value) in seen or id(value) in memo:
+            return
+        seen.add(id(value))
+        if isinstance(value, pl.DataFrame):
+            frames.append(value)
+            for series in value:
+                if series.dtype == pl.Object:
+                    for cell in series:
+                        discover(cell)
+        elif isinstance(value, dict):
+            for key, item in value.items():
+                discover(key)
+                discover(item)
+        elif isinstance(value, (list, tuple)):
+            for item in value:
+                discover(item)
+
+    discover(values)
+    # Register all frames first, including frames referred to by Object cells.
+    # The common memo preserves cycles and cell aliases across metadata frames.
+    for frame in frames:
+        memo[id(frame)] = frame.clone()
+    for frame in frames:
+        for index, series in enumerate(frame):
+            if series.dtype == pl.Object:
+                memo[id(frame)].replace_column(
+                    index,
+                    pl.Series(
+                        series.name,
+                        [deepcopy(cell, memo) for cell in series],
+                        dtype=pl.Object,
+                    ),
+                )
+
+
+def _copy_complete(source, memo=None):
+    """Return a complete independently owned snapshot."""
+    return _copy_graph(source, memo=memo)