feat(num_feat): add NumericalFeature.feature_matrix (#337)

aaanalysis/feature_engineering/_backend/num_feat/feature_matrix.py

@@ -0,0 +1,138 @@
+"""
+This is a script for the backend of ``NumericalFeature.feature_matrix``.
+
+``get_feature_matrix_num_`` reconstructs the ``(n_samples, n_features)`` value
+matrix ``X`` from pre-sliced per-part numerical tensors (``dict_num_parts``),
+the numerical-mode analog of :func:`utils_feature.get_feature_matrix_`.
+
+The value of feature ``PART-SPLIT-SCALE`` for sample ``i`` is:
+
+    ``round(nanmean( arr[part][i, split_positions_i, scale_idx] ), 5)``
+
+where ``split_positions_i`` are the **0-based residue indices** obtained by
+applying the ``SPLIT`` to ``arange(L_part_i)`` (``L_part_i`` = the sample's real
+residue count for that part), and ``scale_idx`` is the D-axis index of ``SCALE``
+in ``df_scales.columns``. This is the SAME reconstruction ``CPP.run_num`` uses in
+``recompute_feature_matrix`` (position buffer built from ``arange`` + per-split
+``np.nanmean`` + ``np.round(_, 5)``), so the columns of ``X`` are byte-identical
+to the values ``run_num`` computed for the same feature ids — verified against
+``recompute_feature_matrix`` for uniform and variable-length parts.
+
+Note that the ``positions`` column of ``run_num``'s ``df_feat`` is a *display*
+numbering in TMD-JMD coordinate space (JMD-offset, e.g. ``21..30`` for a TMD),
+NOT an index into the ``(L, D)`` tensor; value reconstruction therefore re-applies
+the SPLIT to the residue axis rather than reading ``positions``.
+"""
+import numpy as np
+from joblib import Parallel, delayed
+
+import aaanalysis.utils as ut
+from ..cpp._split import SplitVec
+from ..cpp.utils_feature import _get_split_info
+
+
+# I Helper Functions
+def _feature_column_num(feature=None, dict_num_parts=None, part_lens=None,
+                        scale_to_idx=None, sp_vec=None):
+    """Compute the ``(n_samples,)`` value column for one ``PART-SPLIT-SCALE`` feature."""
+    part_upper, split_str, scale = ut.split_feat_id(feat_id=feature)
+    part = part_upper.lower()
+    split_type, split_kwargs = _get_split_info(split=split_str)
+    f_split = getattr(sp_vec, split_type.lower())
+    arr_2d = dict_num_parts[part][:, :, scale_to_idx[scale]]  # (n, L_max) float64
+    lens = part_lens[part]
+    n_samples = arr_2d.shape[0]
+
+    # Fast path: every sample shares the same real length for this part, so the
+    # split positions are identical across samples -> one vectorized gather +
+    # ``nanmean(axis=1)`` (element-wise identical to the per-sample nanmean).
+    if n_samples > 0 and bool(np.all(lens == lens[0])):
+        positions = np.asarray(
+            f_split(seq=np.arange(int(lens[0]), dtype=np.int64), **split_kwargs),
+            dtype=np.int64,
+        )
+        if positions.size == 0:
+            return np.round(np.full(n_samples, np.nan, dtype=np.float64), 5)
+        col = np.nanmean(arr_2d[:, positions], axis=1)
+        return np.round(col, 5)
+
+    # General path: positions depend on each sample's real length.
+    col = np.empty(n_samples, dtype=np.float64)
+    for i in range(n_samples):
+        positions = np.asarray(
+            f_split(seq=np.arange(int(lens[i]), dtype=np.int64), **split_kwargs),
+            dtype=np.int64,
+        )
+        col[i] = np.nanmean(arr_2d[i, positions]) if positions.size else np.nan
+    return np.round(col, 5)
+
+
+def _compute_features_num(features=None, dict_num_parts=None, part_lens=None,
+                          scale_to_idx=None):
+    """Fill an ``(n_samples, n_features)`` matrix for a chunk of features."""
+    sp_vec = SplitVec(type_str=False)  # one shared instance, hoisted out of the loop
+    n_samples = next(iter(dict_num_parts.values())).shape[0]
+    X = np.empty((n_samples, len(features)), dtype=np.float64)
+    for j, feature in enumerate(features):
+        X[:, j] = _feature_column_num(
+            feature=feature, dict_num_parts=dict_num_parts, part_lens=part_lens,
+            scale_to_idx=scale_to_idx, sp_vec=sp_vec,
+        )
+    return X
+
+
+# II Main Functions
+def get_feature_matrix_num_(features=None, dict_num_parts=None, part_lens=None,
+                            df_scales=None, n_jobs=None):
+    """Build the ``(n_samples, n_features)`` numerical-mode feature matrix ``X``.
+
+    Numerical-mode analog of ``utils_feature.get_feature_matrix_``: instead of an
+    AA -> scale lookup, per-residue values come from the pre-sliced
+    ``dict_num_parts`` tensors. Byte-identical to the values ``CPP.run_num``
+    computes for the same feature ids (verified against ``recompute_feature_matrix``).
+
+    Parameters
+    ----------
+    features : list of str
+        Feature ids (``PART-SPLIT-SCALE``); output column order matches.
+    dict_num_parts : dict[str, np.ndarray]
+        Per-part ``(n_samples, L_part_max, D)`` NaN-padded numerical tensors, as
+        produced by ``NumericalFeature.get_parts``.
+    part_lens : dict[str, np.ndarray]
+        Per-part ``(n_samples,)`` int64 real residue counts, from
+        ``_derive_dict_part_lens(df_parts)`` — the SAME length source ``CPP.run_num``
+        uses (non-gap character count of the string ``df_parts``). Applying each
+        split to ``arange(part_lens[part][i])`` therefore lands on exactly the
+        residues ``run_num`` selected.
+    df_scales : pd.DataFrame
+        Names the D dimensions; ``df_scales.columns`` order defines the SCALE ->
+        D-index mapping.
+    n_jobs : int, None, or -1
+        CPU cores for the feature loop (resolved via ``ut.resolve_n_jobs``).
+
+    Returns
+    -------
+    X : np.ndarray, shape (n_samples, n_features) float64
+        Per-sample per-feature values, rounded to 5 decimals.
+    """
+    features = [features] if isinstance(features, str) else list(features)
+    scale_to_idx = {s: i for i, s in enumerate(list(df_scales.columns))}
+
+    n_jobs = ut.resolve_n_jobs(n_jobs=n_jobs, n_work=len(features))
+    if n_jobs == 1 or len(features) <= 1:
+        return _compute_features_num(
+            features=features, dict_num_parts=dict_num_parts,
+            part_lens=part_lens, scale_to_idx=scale_to_idx,
+        )
+
+    feature_chunks = np.array_split(features, n_jobs)
+
+    def _mp(chunk):
+        return _compute_features_num(
+            features=list(chunk), dict_num_parts=dict_num_parts,
+            part_lens=part_lens, scale_to_idx=scale_to_idx,
+        )
+
+    with Parallel(n_jobs=n_jobs) as parallel:
+        results = parallel(delayed(_mp)(chunk) for chunk in feature_chunks)
+    return np.concatenate(results, axis=1)

aaanalysis/feature_engineering/_numerical_feature.py

@@ -9,10 +9,13 @@
 
 import aaanalysis.utils as ut
 
-from ._backend.check_feature import check_df_scales, expand_pos_anchors_
+from ._backend.check_feature import (check_df_scales, expand_pos_anchors_,
+                                     check_match_df_scales_features)
 from ._backend.feature_filter import filter_correlation_
 from ._backend.num_feat.extend_alphabet import extend_alphabet_
+from ._backend.num_feat.feature_matrix import get_feature_matrix_num_
 from ._backend.cpp._filters._assign import assign_dict_num_to_parts
+from ._cpp import _derive_dict_part_lens
 
 
 # I Helper Functions
@@ -66,6 +69,37 @@ def check_match_df_seq_dict_num(df_seq=None, dict_num=None) -> None:
         raise ValueError("'dict_num[*]' has D=0; should be >= 1.")
 
 
+def check_dict_num_parts(dict_num_parts=None) -> Tuple[int, int]:
+    """Validate ``dict_num_parts`` (the value source for ``feature_matrix``) and return ``(n_samples, D)``.
+
+    Each value must be a 3-D ``(n_samples, L_part_max, D)`` ndarray, with a
+    consistent ``n_samples`` (axis 0) and ``D`` (axis 2) across all parts — the
+    contract emitted by :meth:`NumericalFeature.get_parts`.
+    """
+    ut.check_dict(name="dict_num_parts", val=dict_num_parts, accept_none=False)
+    if len(dict_num_parts) == 0:
+        raise ValueError("'dict_num_parts' should not be empty; it must map each part name "
+                         "to a (n_samples, L_part_max, D) tensor (see NumericalFeature.get_parts).")
+    n_seen, d_seen = set(), set()
+    for part, arr in dict_num_parts.items():
+        if not hasattr(arr, "shape") or getattr(arr, "ndim", None) != 3:
+            raise ValueError(
+                f"'dict_num_parts[{part!r}]' should be a 3-D ndarray (n_samples, L_part_max, D); "
+                f"got ndim={getattr(arr, 'ndim', None)}."
+            )
+        n_seen.add(arr.shape[0])
+        d_seen.add(arr.shape[2])
+    if len(n_seen) > 1:
+        raise ValueError(f"'dict_num_parts' has inconsistent n_samples (axis 0) across parts: "
+                         f"{sorted(n_seen)} — all parts must share the same number of samples.")
+    if len(d_seen) > 1:
+        raise ValueError(f"'dict_num_parts' has inconsistent D (axis 2) across parts: "
+                         f"{sorted(d_seen)} — all parts must share the same dimensionality.")
+    if d_seen and 0 in d_seen:
+        raise ValueError("'dict_num_parts[*]' has D=0; should be >= 1.")
+    return n_seen.pop(), d_seen.pop()
+
+
 # II Main Functions
 class NumericalFeature:
     """
@@ -75,7 +109,8 @@ class NumericalFeature:
     It provides numeric helpers for the :class:`CPP` feature engineering pipeline:
     extending the amino acid alphabet of a scale DataFrame, slicing per-residue tensors
     into sequence parts (numerical analog of :meth:`SequenceFeature.get_df_parts`),
-    and removing redundant features by Pearson correlation.
+    reconstructing the model matrix ``X`` from :meth:`CPP.run_num`-selected features
+    (:meth:`feature_matrix`), and removing redundant features by Pearson correlation.
 
     .. versionadded:: 0.1.3
     """
@@ -277,6 +312,160 @@ def get_parts(df_seq: pd.DataFrame,
         # no need to expose them as a separate output (keeps the API one-shape).
         return df_parts, dict_part_vals
 
+    @staticmethod
+    def feature_matrix(features: Union[ut.ArrayLike1D, pd.DataFrame],
+                       dict_num_parts: Dict[str, np.ndarray],
+                       df_parts: pd.DataFrame,
+                       df_scales: Optional[pd.DataFrame] = None,
+                       n_jobs: Union[int, None] = 1,
+                       ) -> np.ndarray:
+        """
+        Create the numerical-mode feature matrix ``X`` for given feature ids and per-part tensors.
+
+        Numerical analog of :meth:`SequenceFeature.feature_matrix`: for each sample and each
+        feature id, reconstructs the feature value from the pre-sliced per-residue tensors in
+        ``dict_num_parts`` (PLM embeddings, structure, annotations, ...) instead of an
+        amino-acid-to-scale lookup, so per-residue context is preserved. The result is the
+        numerical input ``X`` for a downstream model or for
+        :meth:`NumericalFeature.filter_correlation`. This is the missing step that turns
+        :meth:`CPP.run_num`-selected features back into a model matrix.
+
+        .. versionadded:: 1.1.0
+
+        Parameters
+        ----------
+        features : array-like, shape (n_features,) or pd.DataFrame
+            Ids of features (``'PART-SPLIT-SCALE'``) for which the matrix of feature values should
+            be created. Alternatively, a ``df_feat`` DataFrame (e.g. from :meth:`CPP.run_num`), in
+            which case its ``'feature'`` column is used.
+        dict_num_parts : dict[str, np.ndarray]
+            Per-part NaN-padded numerical tensors, as produced by :meth:`NumericalFeature.get_parts`
+            and consumed by :meth:`CPP.run_num`. Each value has shape ``(n_samples, L_part_max, D)``,
+            row-aligned across parts; must cover every part referenced in ``features``.
+        df_parts : pd.DataFrame, shape (n_samples, n_parts)
+            The string ``df_parts`` returned **alongside** ``dict_num_parts`` by
+            :meth:`NumericalFeature.get_parts` (row-aligned with it). Supplies each part's real
+            residue length via the exact same helper :meth:`CPP.run_num` uses internally (non-gap
+            character count), so every split lands on the residues ``run_num`` selected. Its columns
+            must cover every part in ``dict_num_parts``.
+        df_scales : pd.DataFrame, shape (n_letters, n_scales), optional
+            DataFrame whose columns **name the D dimensions** of ``dict_num_parts`` (the same
+            ``df_scales`` used to construct the :class:`CPP` for :meth:`CPP.run_num`); its column
+            order defines the ``SCALE`` -> D-index mapping. Its row (amino acid) values are unused in
+            numerical mode. Default from :meth:`load_scales` unless specified in
+            ``options['df_scales']`` — pass your own when the D axis is a custom (e.g. embedding)
+            space, since the default AA-scale set will not match a custom ``D``.
+        n_jobs : int, None, or -1, default=1
+            Number of CPU cores (>=1) used for multiprocessing. If ``None``, the number is optimized
+            automatically. If ``-1``, the number is set to all available cores. Overridden by
+            ``options['n_jobs']`` when set.
+
+        Returns
+        -------
+        X : array-like, shape (n_samples, n_features)
+            Feature matrix containing feature values for samples. Column ``i`` corresponds to
+            feature ``i`` in ``features``; values are byte-identical to those :meth:`CPP.run_num`
+            computed for the same feature ids — both take per-part lengths from the same ``df_parts``,
+            so they select identical residues for every split.
+
+        Raises
+        ------
+        ValueError
+            If ``dict_num_parts`` is empty or has inconsistent / non-3D tensors, if ``df_parts`` is
+            missing a part column or its row count / real lengths do not match ``dict_num_parts``, if
+            the ``D`` axis does not match ``len(df_scales.columns)``, if a feature id is malformed or
+            references a part/scale absent from ``dict_num_parts`` / ``df_scales``, or if a feature's
+            split selects only padded (all-NaN) residues (producing a ``NaN`` value).
+
+        Notes
+        -----
+        * **Mapping ``positions`` back to ``dict_num_parts`` — do not index with it.** The
+          ``'positions'`` column of a :meth:`CPP.run_num` ``df_feat`` is a **display numbering** in
+          TMD-JMD coordinate space: positions are offset by the JMD length and use the ``start`` /
+          ``tmd_len`` / ``jmd_n_len`` / ``jmd_c_len`` shown at run time (e.g. a ``jmd_n_len=10`` TMD
+          is numbered ``21..30``, decoupled from the actual per-sample TMD length). These numbers do
+          **not** directly index the ``(L, D)`` per-part array. ``feature_matrix`` therefore does not
+          read ``positions``; it reconstructs each value the same way :meth:`CPP.run_num` does — by
+          re-applying the ``SPLIT`` encoded in the feature id to the part's **0-based residue axis**
+          (``arange(L_part)``, where ``L_part`` is the sample's real residue count taken from
+          ``df_parts``), selecting the ``SCALE`` column along ``D``, and averaging the selected
+          residues (``nanmean``, rounded to 5 decimals). Because the split is applied to the residue
+          axis and not to the display numbering, ``X`` lines up with ``run_num`` by construction.
+        * **Real per-part lengths come from ``df_parts``**, via the very same length rule
+          (non-gap character count of ``df_parts``) that :meth:`CPP.run_num` applies
+          internally — not inferred from the tensor's NaN padding. Passing the ``df_parts`` that
+          :meth:`NumericalFeature.get_parts` returns alongside ``dict_num_parts`` therefore makes
+          ``X`` identical to ``run_num`` in every case, including when a *real* residue is all-NaN
+          across ``D`` (an unresolved structure position or a masked embedding): its length is still
+          counted from the string, exactly as ``run_num`` counts it.
+        * A feature whose split selects only padded (all-NaN) residues yields a ``NaN`` value; this
+          raises a ``ValueError`` rather than returning a silently-``NaN`` column.
+        * Unlike :meth:`SequenceFeature.feature_matrix` there is no ``accept_gaps`` option: a numeric
+          tensor carries no gap characters, and NaN padding is ignored via ``nanmean`` automatically.
+
+        See Also
+        --------
+        * :meth:`SequenceFeature.feature_matrix`: sequence-mode (per-AA scale) equivalent.
+        * :meth:`NumericalFeature.get_parts`: produces the ``dict_num_parts`` consumed here.
+        * :meth:`CPP.run_num`: selects the feature ids whose values this method reconstructs.
+
+        Examples
+        --------
+        .. include:: examples/nf_feature_matrix.rst
+        """
+        # Load defaults
+        if df_scales is None:
+            df_scales = ut.load_default_scales()
+        # Check input
+        _n_samples, D = check_dict_num_parts(dict_num_parts=dict_num_parts)
+        ut.check_df(name="df_parts", df=df_parts, accept_none=False,
+                    cols_required=list(dict_num_parts.keys()))
+        if len(df_parts) != _n_samples:
+            raise ValueError(
+                f"'df_parts' has {len(df_parts)} rows but 'dict_num_parts' has {_n_samples} samples. "
+                f"Pass the 'df_parts' returned alongside 'dict_num_parts' by NumericalFeature.get_parts."
+            )
+        check_df_scales(df_scales=df_scales)
+        n_scales = len(df_scales.columns)
+        if D != n_scales:
+            raise ValueError(
+                f"'dict_num_parts' D={D} should equal len(df_scales.columns)={n_scales}. "
+                f"'df_scales' names the D dimensions in numerical mode — pass the same 'df_scales' "
+                f"used to construct the CPP for 'run_num'."
+            )
+        features = ut.check_features(features=features, list_parts=list(dict_num_parts.keys()),
+                                     list_scales=list(df_scales.columns))
+        check_match_df_scales_features(df_scales=df_scales, features=features)
+        n_jobs = ut.check_n_jobs(n_jobs=n_jobs)
+        # Real per-part lengths from the SAME source CPP.run_num uses (the string 'df_parts',
+        # non-gap character count) rather than the tensor's NaN padding, so each split lands on
+        # exactly the residues run_num selected — identical X even when a real residue is all-NaN.
+        part_lens = _derive_dict_part_lens(df_parts=df_parts)
+        for part, arr in dict_num_parts.items():
+            if part_lens[part].size and int(part_lens[part].max()) > arr.shape[1]:
+                raise ValueError(
+                    f"'df_parts[{part!r}]' has a real residue length exceeding the padded tensor "
+                    f"length L_part_max={arr.shape[1]} in 'dict_num_parts'. Pass 'df_parts' and "
+                    f"'dict_num_parts' from the same NumericalFeature.get_parts call."
+                )
+        # Build the feature matrix (byte-identical to CPP.run_num's value reconstruction)
+        try:
+            X = get_feature_matrix_num_(features=features, dict_num_parts=dict_num_parts,
+                                        part_lens=part_lens, df_scales=df_scales, n_jobs=n_jobs)
+        except IndexError as error:
+            raise ValueError(
+                "A feature's 'SPLIT' references a residue position beyond a part's length in "
+                "'dict_num_parts'. Ensure 'features' were generated for these parts (e.g. via "
+                "CPP.run_num on the same 'dict_num_parts')."
+            ) from error
+        if not np.isfinite(X).all():
+            raise ValueError(
+                "'feature_matrix' produced NaN feature values: at least one feature's split selects "
+                "only padded (all-NaN) residues in 'dict_num_parts'. Drop that feature or widen the "
+                "part/split so it covers real residues."
+            )
+        return X
+
     @staticmethod
     def extend_alphabet(df_scales: pd.DataFrame,
                         new_letter: str,