diff --git a/README.md b/README.md index 2c50427..f87a65c 100644 --- a/README.md +++ b/README.md @@ -9,28 +9,35 @@ [![PyPI version](https://img.shields.io/pypi/v/kalelinear?color=blue)](https://pypi.org/project/kalelinear/) [![PyPI downloads](https://pepy.tech/badge/kalelinear)](https://pepy.tech/project/kalelinear) -`kalelinear` is a Python library for learning harmonized or individualized models from multi-source/multi-view data in linear or reproducing kernel Hilbert spaces (RKHS). It provides NumPy-based methods for leveraging related data distributions and structural assumptions, including transfer learning, domain adaptation, manifold regularization, and group-aware learning, through a [`scikit-learn`](https://github.com/scikit-learn/scikit-learn) style API. +KaleLinear is a Python library for non-deep machine learning that learns transferable, shared, or group-specific models from data across multiple sources, groups, blocks, or views. It provides NumPy-based methods in linear or reproducing kernel Hilbert spaces (RKHS), including transfer learning, domain adaptation, manifold regularization, and group-aware learning, through a [`scikit-learn`](https://github.com/scikit-learn/scikit-learn) style API. The package is part of the [PyKale](https://github.com/pykale/pykale) ecosystem and focuses on linear and kernel methods for data characterized by covariates (e.g., domain labels, group labels, side information), unlabeled target samples, or tensor structures. -## What's included - -- Transformer models for learning feature embeddings: - - Multilinear Principal Component Analysis (MPCA) [[1](#references)] - - Transfer Component Analysis (TCA) [[2](#references)] - - Joint Distribution Adaptation (JDA) [[3](#references)] - - Balanced Distribution Adaptation (BDA) [[4](#references)] - - Maximum Independence Domain Adaptation (MIDA) [[5](#references)] -- Estimator models for classification: - - Manifold Regularization Learning Framework (LapSVM, LapRLS) [[6](#references)] - - Adaptation Regularization Learning Framework (ARSVM, ARRLS) [[7](#references)] - - Covariate Independence Regularized Learning Framework (CoIRSVM, CoIRLS) [[8](#references)][[9](#references)] - - Group-specific Discriminant Analysis (GSDA) [[9](#references)][[10](#references)] -- NumPy-compatible inputs and outputs. +## Key features + +- Feature transformation models for data embedding via `kalelinear.transformer` (PyKale-style alias: `kalelinear.embed`): + - Dimension reduction for multiview tensor data: + - Multilinear Principal Component Analysis (`MPCA`) [[1](#references)] + - Transferable / generalizable feature extraction across domains or groups: + - Transfer Component Analysis (`TCA`) [[2](#references)] + - Joint Distribution Adaptation (`JDA`) [[3](#references)] + - Balanced Distribution Adaptation (`BDA`) [[4](#references)] + - Maximum Independence Domain Adaptation (`MIDA`) [[5](#references)] + - Common (or shared or joint) and individual feature separation / extraction across groups or blocks: + - Common and Individual Feature Extraction (`CIFE`) [[11](#references)] + - Angle-based Joint and Individual Variation Explained (`AJIVE`) [[12](#references)] +- Estimator models for prediction via `kalelinear.estimator` (PyKale-style alias: `kalelinear.predict`): + - Predictive models that generalize across domains or groups: + - Manifold Regularization Learning Framework (`LapSVM`, `LapRLS`) [[6](#references)] + - Adaptation Regularization Learning Framework (`ARSVM`, `ARRLS`) [[7](#references)] + - Covariate Independence Regularized Learning Framework (`CoIRSVM`, `CoIRLS`) [[8](#references)][[9](#references)] + - Group-specific predictive models: + - Group-specific Discriminant Analysis (`GSDA`) [[9](#references)][[10](#references)] +- Lightweight: plain NumPy array inputs and outputs — no deep-learning framework or GPU required. - scikit-learn style `fit`, `transform`, `predict`, `fit_transform`, and `fit_predict` workflows where applicable. -- Optional covariate encoding for categorical domain or group labels. +- Most methods accept additional `covariates` — e.g., domain or group labels — alongside `X` and `y`, with optional one-hot encoding for categorical values; multiblock transformers (CIFE, AJIVE) take `groups` to specify block membership. -`kalelinear` requires Python 3.10 or later. Core dependencies include: +KaleLinear requires Python 3.10 or later. Core dependencies include: - [NumPy](http://www.numpy.org/) - [SciPy](https://www.scipy.org/) @@ -99,6 +106,7 @@ Worked examples for the main transformers and estimators are collected in - Learn a domain-invariant embedding with TCA - Use MIDA with categorical covariates +- Extract common and individual features across groups with CIFE or AJIVE - Train a domain adaptation classifier (ARSVM, ARRLS) - Train a manifold-regularized classifier (LapSVM, LapRLS) @@ -124,6 +132,10 @@ Worked examples for the main transformers and estimators are collected in [10] Zhou, S., Luo, J., Jiang, Y., Wang, H., Lu, H. and Gong, G., 2025. [Group-specific discriminant analysis enhances detection of sex differences in brain functional network lateralization](https://academic.oup.com/gigascience/article/doi/10.1093/gigascience/giaf082/8244707). _GigaScience_, 14, p.giaf082. +[11] Zhou, G., Cichocki, A., Zhang, Y. and Mandic, D., 2016. [Group component analysis for multiblock data: Common and individual feature extraction](https://ieeexplore.ieee.org/abstract/document/7310871). _IEEE Transactions on Neural Networks and Learning Systems_, 27(11), pp.2426-2439. + +[12] Feng, Q., Jiang, M., Hannig, J. and Marron, J.S., 2018. [Angle-based joint and individual variation explained](https://www.sciencedirect.com/science/article/pii/S0047259X1730204X). _Journal of Multivariate Analysis_, 166, pp.241-265. + ## Other open domain adaptation repositories - [POT: Python Optimal Transport](https://github.com/rflamary/POT) @@ -136,4 +148,4 @@ Worked examples for the main transformers and estimators are collected in ## License -`kalelinear` is released under the MIT License. See [LICENSE](LICENSE) for details. +KaleLinear is released under the MIT License. See [LICENSE](LICENSE) for details. diff --git a/TUTORIALS.md b/TUTORIALS.md index 320e125..8941c36 100644 --- a/TUTORIALS.md +++ b/TUTORIALS.md @@ -101,3 +101,47 @@ clf = LapSVM(kernel="linear") clf.fit(X_train, ys) y_pred = clf.predict(X_target) ``` + +## Learn Common and Individual Features Across Groups + +CIFE and AJIVE decompose multiblock data into a common feature subspace shared +by all blocks plus block-specific individual subspaces. Each block below is +generated from a shared two-dimensional common basis plus a block-specific +three-dimensional individual basis, so every 50x10 block is rank-deficient and +the common structure is identifiable. Blocks can be passed either as a stacked +matrix with a `groups` array of block ids or as a list of block matrices +sharing the same feature space. + +```python +import numpy as np +from kalelinear.transformer import CIFE, AJIVE + +rng = np.random.default_rng(0) +n_common, n_individual = 2, 3 + +# Shared common basis plus block-specific individual bases, so each block only +# spans 2 + 3 = 5 of the 10 feature dimensions. +common_basis, _ = np.linalg.qr(rng.normal(size=(10, n_common))) +blocks = [] +for _ in range(3): + individual_basis, _ = np.linalg.qr(rng.normal(size=(10, n_individual))) + individual_basis -= common_basis @ (common_basis.T @ individual_basis) + individual_basis, _ = np.linalg.qr(individual_basis) + block = rng.normal(size=(50, n_common)) @ common_basis.T + block += rng.normal(size=(50, n_individual)) @ individual_basis.T + blocks.append(block) + +X = np.vstack(blocks) +groups = np.repeat([0, 1, 2], 50) + +cife = CIFE(n_individual_components=n_individual, random_state=0) +z_common = cife.fit_transform(X, groups=groups) +z_individual = cife.transform_individual(X, groups=groups) + +ajive = AJIVE(n_resamples=50, random_state=0) +z_common_ajive = ajive.fit_transform(X, groups=groups) +``` + +`transform` returns the common feature scores of the samples (two common +components are recovered here), and `transform_individual` returns a list with +the block-specific individual scores (three per block). diff --git a/docs/images/kalelinear.jpg b/docs/images/kalelinear.jpg index ffb3519..9f76423 100644 Binary files a/docs/images/kalelinear.jpg and b/docs/images/kalelinear.jpg differ diff --git a/docs/source/api.rst b/docs/source/api.rst index 29a5189..31f6df5 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -6,8 +6,5 @@ API Reference This page is kept for existing links. The API reference is now split into the following pages: -* :doc:`api_embed` * :doc:`api_transformers` -* :doc:`api_predict` * :doc:`api_estimators` -* :doc:`api_utilities` diff --git a/docs/source/api_embed.rst b/docs/source/api_embed.rst deleted file mode 100644 index a3dd0dd..0000000 --- a/docs/source/api_embed.rst +++ /dev/null @@ -1,37 +0,0 @@ -Embed -===== - -.. currentmodule:: kalelinear.embed - -This compatibility module mirrors the PyKale-style import path while exposing -the same transformer classes as :mod:`kalelinear.transformer`. - -.. autoclass:: MPCA - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: TCA - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: JDA - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: BDA - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: MIDA - :members: - :undoc-members: - :show-inheritance: - :no-index: diff --git a/docs/source/api_estimators.rst b/docs/source/api_estimators.rst index 3f27682..f66cace 100644 --- a/docs/source/api_estimators.rst +++ b/docs/source/api_estimators.rst @@ -3,6 +3,9 @@ Estimators .. currentmodule:: kalelinear.estimator +The estimator classes are also available through the PyKale-compatible alias +:mod:`kalelinear.predict`. + .. autoclass:: LapSVM :members: :undoc-members: diff --git a/docs/source/api_predict.rst b/docs/source/api_predict.rst deleted file mode 100644 index 8483b17..0000000 --- a/docs/source/api_predict.rst +++ /dev/null @@ -1,49 +0,0 @@ -Predict -======= - -.. currentmodule:: kalelinear.predict - -This compatibility module mirrors the PyKale-style import path while exposing -the same estimator classes as :mod:`kalelinear.estimator`. - -.. autoclass:: LapSVM - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: LapRLS - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: ARSVM - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: ARRLS - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: CoIRSVM - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: CoIRLS - :members: - :undoc-members: - :show-inheritance: - :no-index: - -.. autoclass:: GSDA - :members: - :undoc-members: - :show-inheritance: - :no-index: diff --git a/docs/source/api_transformers.rst b/docs/source/api_transformers.rst index b7db20e..d20ac72 100644 --- a/docs/source/api_transformers.rst +++ b/docs/source/api_transformers.rst @@ -3,6 +3,9 @@ Transformers .. currentmodule:: kalelinear.transformer +The transformer classes are also available through the PyKale-compatible alias +:mod:`kalelinear.embed`. + .. autoclass:: MPCA :members: :undoc-members: @@ -27,3 +30,13 @@ Transformers :members: :undoc-members: :show-inheritance: + +.. autoclass:: CIFE + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: AJIVE + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_utilities.rst b/docs/source/api_utilities.rst deleted file mode 100644 index 0a17198..0000000 --- a/docs/source/api_utilities.rst +++ /dev/null @@ -1,13 +0,0 @@ -Utilities -========= - -.. currentmodule:: kalelinear.utils - -.. autofunction:: lap_norm -.. autofunction:: mmd_coef -.. autofunction:: centering_matrix -.. autofunction:: centered_kernel_matrix -.. autofunction:: hsic_grad_term -.. autofunction:: kernel_fit_matrices -.. autofunction:: base_init -.. autofunction:: to_numpy diff --git a/docs/source/index.rst b/docs/source/index.rst index 343bc78..b94aeb7 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,4 +1,4 @@ -kalelinear Documentation +KaleLinear Documentation ========================= Getting Started @@ -11,21 +11,19 @@ Getting Started installation tutorial -kalelinear API +KaleLinear API --------------- .. toctree:: :maxdepth: 2 - api_embed api_transformers - api_predict api_estimators - api_utilities -kalelinear APIs above are ordered following the machine learning pipeline, -i.e., feature embedding transformers, predictive estimators, and reusable -utilities, rather than alphabetically. +The API pages above are ordered following the machine learning pipeline, +i.e., feature embedding transformers and predictive estimators, rather than +alphabetically. `kalelinear.embed` and `kalelinear.predict` are PyKale-compatible +aliases of `kalelinear.transformer` and `kalelinear.estimator`, respectively. Project ------- diff --git a/docs/source/installation.rst b/docs/source/installation.rst index aad43fd..f16e79d 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -13,7 +13,7 @@ Install a local checkout for development: pip install -e ".[dev]" -Kale-Linear requires Python 3.10 or later. Core dependencies include NumPy, +KaleLinear requires Python 3.10 or later. Core dependencies include NumPy, SciPy, scikit-learn, pandas, tensorly, cvxopt, and osqp. To build the documentation locally, install the documentation requirements and diff --git a/docs/source/introduction.rst b/docs/source/introduction.rst index 1611d71..a21aa52 100644 --- a/docs/source/introduction.rst +++ b/docs/source/introduction.rst @@ -1,7 +1,7 @@ Introduction ============ -Kale-Linear is a Python library for non-deep, knowledge-aware machine learning +KaleLinear is a Python library for non-deep, knowledge-aware machine learning from multiple sources, domains, or views. It provides NumPy-based implementations of transfer learning, domain adaptation, manifold regularization, and group-aware linear learning methods with a scikit-learn @@ -14,8 +14,8 @@ covariates, side information, or unlabeled target samples. Main Features ------------- -* Transformer models for learning feature embeddings: MPCA, TCA, JDA, BDA, and - MIDA. +* Transformer models for learning feature embeddings: MPCA, TCA, JDA, BDA, + MIDA, CIFE, and AJIVE. * Estimator models for classification and adaptation: LapSVM, LapRLS, ARSVM, ARRLS, CoIRSVM, CoIRLS, and GSDA. * NumPy-compatible inputs and outputs. diff --git a/kalelinear/__init__.py b/kalelinear/__init__.py index 73855ce..e54af09 100644 --- a/kalelinear/__init__.py +++ b/kalelinear/__init__.py @@ -1,7 +1,7 @@ """ kalelinear. -Learning harmonized or individualized models from multi-source/multi-view data in linear or reproducing kernel Hilbert spaces (RKHS). +Non-deep machine learning that learns transferable, shared, or group-specific models from data across multiple sources, groups, blocks, or views. """ from importlib import import_module diff --git a/kalelinear/embed.py b/kalelinear/embed.py index a9f59b3..3cb98b5 100644 --- a/kalelinear/embed.py +++ b/kalelinear/embed.py @@ -1,5 +1,5 @@ """Embedding models exposed with a PyKale-style API.""" -from kalelinear.transformer import BDA, JDA, MIDA, MPCA, TCA +from kalelinear.transformer import AJIVE, BDA, CIFE, JDA, MIDA, MPCA, TCA -__all__ = ["TCA", "JDA", "BDA", "MIDA", "MPCA"] +__all__ = ["TCA", "JDA", "BDA", "MIDA", "MPCA", "CIFE", "AJIVE"] diff --git a/kalelinear/transformer/__init__.py b/kalelinear/transformer/__init__.py index 7a68902..b0c44bb 100644 --- a/kalelinear/transformer/__init__.py +++ b/kalelinear/transformer/__init__.py @@ -1,3 +1,5 @@ +from kalelinear.transformer._ajive import AJIVE +from kalelinear.transformer._cife import CIFE from kalelinear.transformer._jda import BDA, JDA from kalelinear.transformer._mida import MIDA from kalelinear.transformer._mpca import MPCA @@ -9,4 +11,6 @@ "BDA", "MIDA", "MPCA", + "CIFE", + "AJIVE", ] diff --git a/kalelinear/transformer/_ajive.py b/kalelinear/transformer/_ajive.py new file mode 100644 index 0000000..4c29c69 --- /dev/null +++ b/kalelinear/transformer/_ajive.py @@ -0,0 +1,285 @@ +# ============================================================================= +# @author: Shuo Zhou, The University of Sheffield +# @contact: shuo.zhou@sheffield.ac.uk +# ============================================================================= +"""Angle-based Joint and Individual Variation Explained (AJIVE) implementation. + +References +---------- +Feng, Q., Jiang, M., Hannig, J. and Marron, J.S., 2018. Angle-based joint and +individual variation explained. Journal of Multivariate Analysis, 166, +pp.241-265. + +The algorithm follows the authors' reference MATLAB implementation +(MeileiJiang/AJIVE_Project on GitHub) with the same perturbation-bound and rank +selection steps. +""" + +from numbers import Integral, Real + +import numpy as np +from sklearn.utils._param_validation import Interval + +from kalelinear.transformer._multiblock import _check_per_block_ranks, BaseCommonIndividualTransformer + +_FERROR = 1e-10 + + +def _jive_rand_null_norm(data, basis, n_sim, random_state): + """Operator norm of data on random directions orthogonal to ``basis``.""" + n_ambient_dims = basis.shape[0] + n_null_dims = n_ambient_dims - basis.shape[1] + if n_null_dims <= 0: + return np.zeros(n_sim) + # ``initial_ranks`` may legally exceed the null-space dimension, in which + # case requesting one orthogonal direction per basis vector would run out + # of genuine null directions: once ``current`` spans the ambient space the + # retry normalises round-off noise (or never terminates in exact + # arithmetic). Only ask for as many directions as the null space can hold. + n_directions = min(basis.shape[1], n_null_dims) + null_norms = np.empty(n_sim) + for i in range(n_sim): + current = basis.copy() + directions = [] + for _ in range(n_directions): + direction = random_state.randn(n_ambient_dims) + direction = direction - current @ (current.T @ direction) + norm = np.linalg.norm(direction) + while norm == 0: + direction = random_state.randn(n_ambient_dims) + direction = direction - current @ (current.T @ direction) + norm = np.linalg.norm(direction) + direction /= norm + directions.append(direction) + current = np.column_stack((current, direction)) + directions = np.column_stack(directions) + # The Wedin bound needs the spectral (2-)norm, matching the reference + # implementation's MATLAB ``norm(data * nulldir)``; the default + # Frobenius norm grows with the number of sampled directions and would + # systematically inflate the angle bound. + null_norms[i] = np.linalg.norm(data @ directions, ord=2) + return null_norms + + +def _wedin_angle_bound(block, n_sim, U, S, V, random_state): + """Resampled Wedin perturbation-angle bound for one data block.""" + delta = S[-1] + if delta <= S[0] * np.finfo(float).eps: + raise ValueError( + "`initial_ranks` exceeds the numerical rank of a data block: the " + "smallest retained singular value is zero up to machine precision. " + "Lower `initial_ranks` for that block." + ) + row_bound = _jive_rand_null_norm(block, V, n_sim, random_state) + column_bound = _jive_rand_null_norm(block.T, U, n_sim, random_state) + ratio = np.maximum(row_bound, column_bound) / delta + ratio = np.clip(ratio, 0.0, 1.0) + return np.rad2deg(np.arcsin(ratio)) + + +def _random_direction_ssv(n, ranks, n_sim, random_state): + """Largest squared singular values of random stacked subspaces.""" + stacked = np.zeros((int(np.sum(ranks)), n)) + values = np.empty(n_sim) + for i in range(n_sim): + row = 0 + for rank in ranks: + Q, _ = np.linalg.qr(random_state.randn(n, rank)) + stacked[row : row + rank] = Q.T + row += rank + values[i] = np.linalg.norm(stacked, 2) ** 2 + return values + + +class AJIVE(BaseCommonIndividualTransformer): + """Angle-based Joint and Individual Variation Explained (AJIVE). + + AJIVE decomposes multiblock data into a joint (common) subspace shared by + all blocks and block-specific individual subspaces. It estimates the joint + rank from the Wedin perturbation bound of the stacked per-block row spaces + and reconstructs the joint and individual structures with low-rank SVDs. + + ``X`` can be a single matrix with a ``groups`` array of block ids, or a + list of block matrices that all share the same feature space. + :meth:`transform` returns the common feature scores of the samples while + :meth:`transform_individual` returns the block-specific individual scores. + + Parameters + ---------- + n_common_components : int or None, default=None + Number of joint components. When None, the joint rank is selected from + the Wedin perturbation bound. + n_individual_components : int, array-like or None, default=None + Number of individual components per block. When None, components above + the per-block singular-value threshold are kept. + initial_ranks : int, array-like or None, default=None + Initial signal rank of each block. When None, ranks are estimated from + the fraction of variance explained by ``variance_threshold``. + variance_threshold : float, default=0.95 + Cumulative explained-variance fraction used to estimate ``initial_ranks`` + when it is None. + n_resamples : int, default=1000 + Number of re-samples for the Wedin perturbation bound. + percentile : float, default=5 + Percentile of the Wedin bounds used for the joint rank selection. + random_state : int, RandomState or None, default=None + Random seed for the perturbation-bound re-sampling. + + Attributes + ---------- + common_components_ : ndarray of shape (n_features, n_common_components_) + Orthonormal joint feature-space basis. + individual_components_ : list of ndarray of shape (n_features, rank) + Per-block orthonormal individual feature-space bases. + common_scores_ : list of ndarray of shape (n_samples_in_block, n_common_components_) + Common scores of the training samples in each block. + individual_scores_ : list of ndarray of shape (n_samples_in_block, rank) + Individual scores of the training samples in each block. + individual_ranks_ : ndarray of shape (n_blocks,) + Number of individual components retained per block. + n_common_components_ : int + Number of joint components retained. + """ + + _parameter_constraints: dict = { + **BaseCommonIndividualTransformer._parameter_constraints, + "initial_ranks": ["array-like", Interval(Integral, 1, None, closed="left"), None], + "variance_threshold": [Interval(Real, 0, 1, closed="both")], + "n_resamples": [Interval(Integral, 1, None, closed="left")], + "percentile": [Interval(Real, 0, 100, closed="both")], + } + + def __init__( + self, + n_common_components=None, + n_individual_components=None, + initial_ranks=None, + variance_threshold=0.95, + n_resamples=1000, + percentile=5, + random_state=None, + ): + self.initial_ranks = initial_ranks + self.variance_threshold = variance_threshold + self.n_resamples = n_resamples + self.percentile = percentile + super().__init__( + n_common_components=n_common_components, + n_individual_components=n_individual_components, + random_state=random_state, + ) + + def _resolve_initial_ranks(self, blocks): + if self.initial_ranks is None: + ranks = [] + for block in blocks: + singular_values = np.linalg.svd(block, compute_uv=False) + if singular_values.size == 0 or singular_values[0] == 0: + # A zero-energy block has no signal rank; the positive-rank + # validation below then rejects it with a clear message. + ranks.append(0) + continue + explained = np.cumsum(singular_values**2) / np.sum(singular_values**2) + rank = int(np.searchsorted(explained, self.variance_threshold) + 1) + ranks.append(int(min(rank, singular_values.size))) + return np.asarray(ranks, dtype=int) + ranks = _check_per_block_ranks(self.initial_ranks, self.n_blocks_, "initial_ranks") + for n, (rank, block) in enumerate(zip(ranks, blocks)): + if rank < 1 or rank > min(block.shape): + raise ValueError( + f"`initial_ranks[{n}]` must be between 1 and min(samples, features) " + f"= {min(block.shape)}, got {rank}." + ) + if not np.any(block): + # A zero-energy block has no row or column space to share: the + # perturbation bound would divide by a zero singular value and + # every residual direction would be misclassified as individual + # signal. Reject it even when the ranks were supplied explicitly. + raise ValueError(f"Block {n} has zero energy, so `initial_ranks[{n}]` must not be positive.") + return ranks + + def _fit_blocks(self, blocks): + D = self.n_features_in_ + ranks = self._resolve_initial_ranks(blocks) + if np.any(ranks < 1): + raise ValueError("`initial_ranks` must contain positive values for every block.") + + stacked = np.zeros((int(np.sum(ranks)), D)) + thresholds = np.empty(self.n_blocks_) + angle_bounds = [] + row = 0 + for n, block in enumerate(blocks): + rank = ranks[n] + U, s, Vt = np.linalg.svd(block, full_matrices=False) + if rank + 1 <= len(s): + thresholds[n] = 0.5 * (s[rank - 1] + s[rank]) + else: + thresholds[n] = 0.5 * s[-1] + U0, S0, V0 = U[:, :rank], s[:rank], Vt[:rank].T + stacked[row : row + rank] = V0.T + row += rank + angle_bounds.append(_wedin_angle_bound(block, self.n_resamples, U0, S0, V0, self.random_state_)) + angle_bounds = np.vstack(angle_bounds) + + _, s_stacked, Vt_stacked = np.linalg.svd(stacked, full_matrices=False) + max_joint_rank = int(np.min(ranks)) + # The joint space is a subspace of every block's row space, so it can + # have at most ``max_joint_rank`` dimensions. Keep only the leading + # singular values and vectors, as the reference implementation computes + # only ``min(vecr)`` singular vectors; directions supported by only a + # subset of blocks are then never candidates for the common subspace. + s_stacked = s_stacked[:max_joint_rank] + Vt_stacked = Vt_stacked[:max_joint_rank] + wedin_ssv_bounds = np.maximum(np.sum(np.cos(np.deg2rad(angle_bounds)) ** 2, axis=0), 1.0) + wedin_ssv_bound = np.percentile(wedin_ssv_bounds, self.percentile) + random_ssvs = _random_direction_ssv(D, ranks, 100, self.random_state_) + random_ssv_bound = np.percentile(random_ssvs, 95) + # Take the more conservative (larger) of the two perturbation bounds, + # following the reference implementation: max(wedin, random). + joint_threshold = max(wedin_ssv_bound, random_ssv_bound) + joint_rank = int(np.sum(s_stacked**2 + _FERROR > joint_threshold)) + if self.n_common_components is not None: + joint_rank = min(int(self.n_common_components), len(s_stacked)) + + row_joint = Vt_stacked[:joint_rank] + drop_rows = set() + for n, block in enumerate(blocks): + projected = block @ row_joint.T + low_variance = np.flatnonzero(np.sqrt(np.sum(projected**2, axis=0)) <= thresholds[n] + _FERROR) + drop_rows.update(low_variance.tolist()) + if drop_rows: + keep_rows = [j for j in range(row_joint.shape[0]) if j not in drop_rows] + row_joint = row_joint[keep_rows] + joint_rank = row_joint.shape[0] + + common_components = row_joint.T + common_scores = [block @ common_components for block in blocks] + + ranks_spec = _check_per_block_ranks(self.n_individual_components, self.n_blocks_, "n_individual_components") + individual_components = [] + individual_scores = [] + individual_ranks = [] + for n, block in enumerate(blocks): + individual = block - block @ common_components @ common_components.T + s_individual = np.linalg.svd(individual, compute_uv=False) + if ranks_spec is None: + rank = int(np.sum(s_individual + _FERROR > thresholds[n])) + else: + rank = min(int(ranks_spec[n]), len(s_individual)) + rank = max(rank, 0) + if rank > 0: + U_i, S_i, Vt_i = np.linalg.svd(individual, full_matrices=False) + U_i, S_i, Vt_i = U_i[:, :rank], S_i[:rank], Vt_i[:rank] + individual_components.append(Vt_i.T) + individual_scores.append(U_i * S_i) + else: + individual_components.append(np.zeros((D, 0))) + individual_scores.append(np.zeros((block.shape[0], 0))) + individual_ranks.append(rank) + + self.common_components_ = common_components + self.common_scores_ = common_scores + self.individual_components_ = individual_components + self.individual_scores_ = individual_scores + self.individual_ranks_ = np.asarray(individual_ranks, dtype=int) + self.n_common_components_ = joint_rank diff --git a/kalelinear/transformer/_cife.py b/kalelinear/transformer/_cife.py new file mode 100644 index 0000000..d6d6c8c --- /dev/null +++ b/kalelinear/transformer/_cife.py @@ -0,0 +1,290 @@ +# ============================================================================= +# @author: Shuo Zhou, The University of Sheffield +# @contact: shuo.zhou@sheffield.ac.uk +# ============================================================================= +"""Common and Individual Feature Extraction (CIFE) implementation. + +References +---------- +Zhou, G., Cichocki, A., Zhang, Y. and Mandic, D., 2016. Group component +analysis for multiblock data: Common and individual feature extraction. +IEEE Transactions on Neural Networks and Learning Systems, 27(11), +pp.2426-2439. + +The common orthogonal basis extraction (COBE) steps follow the pycifa and the +accompanying MATLAB code implementations at https://github.com/kharyuk/pycifa. +""" + +from numbers import Integral, Real + +import numpy as np +from sklearn.utils._param_validation import Interval + +from kalelinear.transformer._multiblock import _check_per_block_ranks, BaseCommonIndividualTransformer + + +def _column_space_basis(Y, pca_dim=None): + """Return an orthonormal basis of the column space of ``Y`` (D x J).""" + D, J = Y.shape + U, s, _ = np.linalg.svd(Y, full_matrices=False) + if s.size == 0 or s[0] == 0: + raise ValueError("Each block must have a non-zero column space.") + rank = int(np.sum(s > s[0] * max(D, J) * np.finfo(s.dtype).eps)) + if rank >= D: + if pca_dim is None: + raise ValueError( + "A block spans the whole feature space, so common and individual " + "subspaces cannot be separated. Reduce the dimensionality first " + "or set `pca_dim` to truncate the per-block column spaces." + ) + if 0 < pca_dim < 1: + rank = min(int(np.floor(D * pca_dim)), D - 1) + else: + rank = min(int(pca_dim), D - 1) + rank = max(rank, 1) + return U[:, :rank], rank + + +def _cobe_common_basis(blocks, c, max_iter, tol, epsilon, pca_dim, random_state): + """Extract a common orthogonal basis shared by all blocks. + + Parameters + ---------- + blocks : list of ndarray of shape (D, J_n) + Data blocks sharing the same feature dimension ``D``. + c : int or None + Number of common components. When None, the number is estimated from + the residual ``epsilon`` criterion. + max_iter : int + Maximum power iterations per common direction. + tol : float + Convergence tolerance for the power iterations. + epsilon : float + Residual threshold below which a direction counts as common. + pca_dim : int, float or None + Optional per-block dimensionality truncation for blocks spanning the + whole feature space. + random_state : RandomState + Random number generator for initializing the power iterations. + + Returns + ------- + common_basis : ndarray of shape (D, n_common) + Orthonormal common basis shared by all blocks. + """ + n_blocks = len(blocks) + D = blocks[0].shape[0] + if c is not None and c <= 0: + # No common subspace is requested, so skip the per-block column-space + # validation entirely; the individual components are computed from the + # untouched blocks. + return np.zeros((D, 0)) + bases = [] + ranks = [] + for Y in blocks: + basis, rank = _column_space_basis(Y, pca_dim=pca_dim) + bases.append(basis) + ranks.append(rank) + min_rank = min(ranks) + if min_rank == 0: + return np.zeros((D, 0)) + + order = np.argsort(ranks) + projections = [np.zeros((bases[n].shape[1], min_rank)) for n in range(n_blocks)] + common_basis = np.zeros((D, min_rank)) + residuals = [] + + def _power_iteration(initial, column): + direction = initial / np.linalg.norm(initial) + for _ in range(max_iter): + previous = direction + update = np.zeros(D) + for n in range(n_blocks): + projections[n][:, column] = bases[n].T @ direction + update += bases[n] @ projections[n][:, column] + update_norm = np.linalg.norm(update) + if update_norm == 0: + break + direction = update / update_norm + if abs(previous @ direction) > 1 - tol: + break + return direction + + # Seek the first common direction. + initial = bases[order[0]] @ random_state.randn(bases[order[0]].shape[1]) + first = _power_iteration(initial, 0) + residual = 0.0 + for n in range(n_blocks): + projection = bases[n].T @ first + residual += 1 - projection @ projection + residual /= n_blocks + residuals.append(residual) + + if c is None and residual > epsilon: + return np.zeros((D, 0)) + + if c is not None: + c = min(c, min_rank) + common_basis = np.zeros((D, c)) + common_basis[:, 0] = first + residuals.extend([np.inf] * (c - 1)) + else: + common_basis[:, 0] = first + residuals.extend([np.inf] * (min_rank - 1)) + + # Seek the remaining common directions with deflation. + for j in range(1, min_rank): + if c is not None and j >= c: + break + for n in range(n_blocks): + basis = bases[n] + bases[n] = basis - np.outer(basis @ projections[n][:, j - 1], projections[n][:, j - 1]) + initial = bases[order[0]] @ random_state.randn(bases[order[0]].shape[1]) + direction = _power_iteration(initial, j) + residual = 0.0 + for n in range(n_blocks): + projection = bases[n].T @ direction + residual += 1 - projection @ projection + residual /= n_blocks + residuals[j] = residual + if c is None and residual > epsilon: + residuals[j] = np.inf + break + common_basis[:, j] = direction + + common_basis = common_basis[:, ~np.isinf(np.asarray(residuals))] + if common_basis.shape[1] > 0: + u, _, vt = np.linalg.svd(common_basis, full_matrices=False) + common_basis = u @ vt + return common_basis + + +class CIFE(BaseCommonIndividualTransformer): + """Common and Individual Feature Extraction (CIFE). + + CIFE decomposes multiblock data into a common feature subspace shared by + all blocks and block-specific individual subspaces. The common subspace is + extracted with the common orthogonal basis extraction (COBE) algorithm and + the individual subspaces are obtained from the residual of each block after + removing its common part. + + ``X`` can be a single matrix with a ``groups`` array of block ids, or a + list of block matrices that all share the same feature space. + :meth:`transform` returns the common feature scores of the samples while + :meth:`transform_individual` returns the block-specific individual scores. + + Parameters + ---------- + n_common_components : int or None, default=None + Number of common components to extract. When None, the number is + estimated automatically from the residual threshold ``epsilon``. + n_individual_components : int, array-like or None, default=None + Number of individual components per block. When None, all numerically + non-zero residual directions are kept. + max_iter : int, default=200 + Maximum power iterations for each common direction. + tol : float, default=1e-6 + Convergence tolerance for the power iterations. + epsilon : float, default=0.01 + Residual threshold used to decide whether a direction is common when + ``n_common_components`` is None. + pca_dim : int, float or None, default=None + Optional truncation of per-block column spaces, either as a relative + fraction in (0, 1) or an absolute number of components. Required when + a block spans the whole feature space. + random_state : int, RandomState or None, default=None + Random seed for initializing the power iterations. + + Attributes + ---------- + common_components_ : ndarray of shape (n_features, n_common_components_) + Orthonormal common feature-space basis. + individual_components_ : list of ndarray of shape (n_features, rank) + Per-block orthonormal individual feature-space bases. + common_scores_ : list of ndarray of shape (n_samples_in_block, n_common_components_) + Common scores of the training samples in each block. + individual_scores_ : list of ndarray of shape (n_samples_in_block, rank) + Individual scores of the training samples in each block. + individual_ranks_ : ndarray of shape (n_blocks,) + Number of individual components retained per block. + n_common_components_ : int + Number of common components retained. + """ + + _parameter_constraints: dict = { + **BaseCommonIndividualTransformer._parameter_constraints, + "max_iter": [Interval(Integral, 1, None, closed="left")], + "tol": [Interval(Real, 0, None, closed="left")], + "epsilon": [Interval(Real, 0, None, closed="left")], + "pca_dim": [ + Interval(Real, 0, 1, closed="right"), + Interval(Integral, 2, None, closed="left"), + None, + ], + } + + def __init__( + self, + n_common_components=None, + n_individual_components=None, + max_iter=200, + tol=1e-6, + epsilon=0.01, + pca_dim=None, + random_state=None, + ): + self.max_iter = max_iter + self.tol = tol + self.epsilon = epsilon + self.pca_dim = pca_dim + super().__init__( + n_common_components=n_common_components, + n_individual_components=n_individual_components, + random_state=random_state, + ) + + def _fit_blocks(self, blocks): + D = self.n_features_in_ + transposed = [block.T for block in blocks] + common_components = _cobe_common_basis( + transposed, + c=self.n_common_components, + max_iter=self.max_iter, + tol=self.tol, + epsilon=self.epsilon, + pca_dim=self.pca_dim, + random_state=self.random_state_, + ) + n_common = common_components.shape[1] + common_scores = [block @ common_components for block in blocks] + + ranks_spec = _check_per_block_ranks(self.n_individual_components, self.n_blocks_, "n_individual_components") + individual_components = [] + individual_scores = [] + individual_ranks = [] + for n, (X_n, Y_n) in enumerate(zip(blocks, transposed)): + residual = Y_n - common_components @ (common_components.T @ Y_n) + U_i, s_i, _ = np.linalg.svd(residual, full_matrices=False) + if ranks_spec is None: + if s_i.size == 0: + rank = 0 + else: + noise_tol = s_i[0] * max(residual.shape) * np.finfo(s_i.dtype).eps + rank = int(np.sum(s_i > noise_tol)) + else: + rank = min(int(ranks_spec[n]), U_i.shape[1]) + if rank > 0: + U_i = U_i[:, :rank] + individual_components.append(U_i) + individual_scores.append(X_n @ U_i) + else: + individual_components.append(np.zeros((D, 0))) + individual_scores.append(np.zeros((X_n.shape[0], 0))) + individual_ranks.append(rank) + + self.common_components_ = common_components + self.common_scores_ = common_scores + self.individual_components_ = individual_components + self.individual_scores_ = individual_scores + self.individual_ranks_ = np.asarray(individual_ranks, dtype=int) + self.n_common_components_ = n_common diff --git a/kalelinear/transformer/_multiblock.py b/kalelinear/transformer/_multiblock.py new file mode 100644 index 0000000..51a538b --- /dev/null +++ b/kalelinear/transformer/_multiblock.py @@ -0,0 +1,241 @@ +# ============================================================================= +# @author: Shuo Zhou, The University of Sheffield +# @contact: shuo.zhou@sheffield.ac.uk +# ============================================================================= +"""Shared base classes for multiblock common and individual feature transformers.""" + +from abc import abstractmethod +from numbers import Integral + +import numpy as np +from sklearn.base import BaseEstimator, ClassNamePrefixFeaturesOutMixin, TransformerMixin +from sklearn.utils._param_validation import Interval +from sklearn.utils.validation import check_is_fitted, check_random_state + + +def _check_multiblock_input(X, groups=None, min_blocks=2): + """Validate a multiblock input and return one matrix per block. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) or list of array-like + When ``X`` is a single matrix, ``groups`` must give the block id of + each row. When ``X`` is a list, each element is one block and all + blocks must share the same feature space (columns). + groups : array-like of shape (n_samples,), default=None + Block id for each sample when ``X`` is a single stacked matrix. + min_blocks : int, default=2 + Minimum number of blocks required. Set to 1 when the caller only + needs to stack the blocks, e.g., the common projection in + :meth:`transform` only needs one block, so callers that project new + samples can pass ``min_blocks=1``. + + + Returns + ------- + blocks : list of ndarray of shape (n_samples_in_block, n_features) + groups : ndarray of shape (n_samples,) or None + block_ids : ndarray + Unique block ids in the order the blocks are returned: first-appearance + order for stacked input, positional ``np.arange(len(blocks))`` for a + list of blocks. + """ + if isinstance(X, (list, tuple)): + if groups is not None: + raise ValueError("`groups` must be None when `X` is a list of blocks.") + if len(X) == 0: + raise ValueError("`X` must contain at least one block.") + blocks = [] + n_features = None + for block in X: + block = np.asarray(block, dtype=float) + if block.ndim != 2: + raise ValueError("Each block in `X` must be a 2D array.") + if block.shape[0] == 0: + raise ValueError("Each block in `X` must contain at least one sample.") + if n_features is None: + n_features = block.shape[1] + elif block.shape[1] != n_features: + raise ValueError("All blocks must share the same number of features.") + blocks.append(block) + if len(blocks) < min_blocks: + raise ValueError("At least two blocks are required for common and individual feature extraction.") + return blocks, None, np.arange(len(blocks)) + + X = np.asarray(X, dtype=float) + if X.ndim != 2: + raise ValueError("`X` must be a 2D array or a list of 2D block arrays.") + if groups is None: + raise ValueError("`groups` must be provided when `X` is a single stacked matrix.") + groups = np.asarray(groups) + if groups.ndim != 1 or groups.shape[0] != X.shape[0]: + raise ValueError("`groups` must be a 1D array aligned with the rows of `X`.") + block_ids, first_idx = np.unique(groups, return_index=True) + block_ids = block_ids[np.argsort(first_idx)] + blocks = [X[groups == block_id] for block_id in block_ids] + if len(blocks) < min_blocks: + raise ValueError(f"At least {min_blocks} blocks are required for common and individual feature extraction.") + if any(block.shape[0] == 0 for block in blocks): + raise ValueError("Each block must contain at least one sample.") + return blocks, groups, block_ids + + +def _align_blocks_to_fit_ids(blocks, incoming_ids, fit_ids): + """Validate incoming stacked block ids and return blocks in fit-time order. + + ``_check_multiblock_input`` builds blocks in first-appearance order, so the + incoming id order can differ from the fit-time order even when the id sets + match. The sets are validated for equality and the blocks are reordered to + the fit-time order, preventing a stacked ``groups`` array with different + ids from silently being paired with the wrong per-block bases. + """ + incoming_ids = np.asarray(incoming_ids) + fit_ids = np.asarray(fit_ids) + if incoming_ids.shape != fit_ids.shape or not np.array_equal(np.sort(incoming_ids), np.sort(fit_ids)): + raise ValueError( + "`groups` must contain the same block ids used at fit time " + f"({fit_ids.tolist()}), got {incoming_ids.tolist()}." + ) + positions = {block_id: index for index, block_id in enumerate(incoming_ids)} + return [blocks[positions[block_id]] for block_id in fit_ids] + + +def _check_per_block_ranks(n_components, n_blocks, name): + """Validate an integer or a per-block sequence of component counts.""" + if n_components is None: + return None + if isinstance(n_components, (Integral, np.integer)): + return np.full(n_blocks, int(n_components), dtype=int) + ranks = np.asarray(n_components) + if ranks.ndim != 1 or ranks.shape[0] != n_blocks: + raise ValueError(f"{name} must be an integer or a sequence with one value per block.") + if not np.issubdtype(ranks.dtype, np.number): + raise ValueError(f"{name} must contain numeric values.") + if np.any(np.isnan(ranks)): + raise ValueError(f"{name} must not contain NaN values.") + if np.any(np.isinf(ranks)): + raise ValueError(f"{name} must not contain infinite values.") + if np.any(ranks < 0): + raise ValueError(f"{name} must contain non-negative values.") + if not np.all(np.equal(ranks, np.floor(ranks))): + raise ValueError(f"{name} must contain integer values.") + return ranks.astype(int) + + +class BaseCommonIndividualTransformer(ClassNamePrefixFeaturesOutMixin, TransformerMixin, BaseEstimator): + """Base class for multiblock common and individual feature transformers. + + Subclasses learn a common feature subspace shared by all data blocks plus + block-specific individual subspaces, following the common and individual + feature extraction (CIFE) framework of Zhou and Cichocki (2016). + + The input convention follows the rest of ``kalelinear``: ``X`` is a sample + matrix whose rows are partitioned into blocks by ``groups``, or a list of + block matrices that all share the same feature space. + + Attributes + ---------- + block_ids_ : ndarray + Unique block ids in fit-time block order, used to align stacked + ``groups`` in :meth:`transform_individual`. + """ + + _parameter_constraints: dict = { + "n_common_components": [Interval(Integral, 0, None, closed="left"), None], + "n_individual_components": ["array-like", Interval(Integral, 0, None, closed="left"), None], + "random_state": ["random_state"], + } + + def __init__(self, n_common_components=None, n_individual_components=None, random_state=None): + self.n_common_components = n_common_components + self.n_individual_components = n_individual_components + self.random_state = random_state + + def fit(self, X, y=None, groups=None, **fit_params): + """Fit the transformer on multiblock data. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) or list of array-like + Stacked samples from all blocks, or a list of block matrices + sharing the same feature space. + y : array-like of shape (n_samples,), default=None + Ignored. Present for scikit-learn API consistency. + groups : array-like of shape (n_samples,), default=None + Block id for each sample, required when ``X`` is a single stacked + matrix. + + Returns + ------- + self : object + Fitted transformer. + """ + self._validate_params() + blocks, _, block_ids = _check_multiblock_input(X, groups) + self.n_features_in_ = blocks[0].shape[1] + self.n_blocks_ = len(blocks) + self.block_ids_ = block_ids + self.block_sizes_ = np.array([block.shape[0] for block in blocks]) + self.random_state_ = check_random_state(self.random_state) + self._fit_blocks(blocks) + self._n_features_out = self.n_common_components_ + return self + + @abstractmethod + def _fit_blocks(self, blocks): + """Run the algorithm on validated per-block matrices.""" + + def transform(self, X, groups=None): + """Project samples onto the learned common feature subspace. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) or list of array-like + New samples, either stacked or given as a list of blocks. A + single block is accepted for the common projection. + groups : array-like of shape (n_samples,), default=None + Ignored for the common projection. Present for API consistency. + + Returns + ------- + X_new : ndarray of shape (n_samples, n_common_components_) + Common feature scores shared by all blocks. + """ + check_is_fitted(self, "common_components_") + if isinstance(X, (list, tuple)): + blocks, _, _ = _check_multiblock_input(X, min_blocks=1) + X_stacked = np.vstack(blocks) + else: + X_stacked = np.asarray(X, dtype=float) + if X_stacked.ndim != 2: + raise ValueError("`X` must be a 2D array or a list of 2D block arrays.") + if X_stacked.shape[1] != self.n_features_in_: + raise ValueError(f"Expected {self.n_features_in_} features, got {X_stacked.shape[1]}.") + return X_stacked @ self.common_components_ + + def transform_individual(self, X, groups=None): + """Project samples onto the individual feature subspaces. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) or list of array-like + New samples, either stacked or given as a list of blocks. + groups : array-like of shape (n_samples,), default=None + Block id for each sample when ``X`` is a single stacked matrix. + Must contain the same block ids used at fit time; blocks are + aligned to the fit-time block order. + + Returns + ------- + scores : list of ndarray + One array per block, of shape (n_samples_in_block, individual_ranks_[i]). + """ + check_is_fitted(self, "individual_components_") + blocks, groups, block_ids = _check_multiblock_input(X, groups) + if len(blocks) != self.n_blocks_: + raise ValueError(f"Expected {self.n_blocks_} blocks, got {len(blocks)}.") + if any(block.shape[1] != self.n_features_in_ for block in blocks): + raise ValueError(f"Expected {self.n_features_in_} features in every block, got mismatched blocks.") + if groups is not None: + blocks = _align_blocks_to_fit_ids(blocks, block_ids, self.block_ids_) + return [block @ components for block, components in zip(blocks, self.individual_components_)] diff --git a/tests/test_public_api.py b/tests/test_public_api.py index b431f3b..708c32a 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -10,6 +10,8 @@ def test_embed_module_exposes_transformers(): assert embed.BDA is transformer.BDA assert embed.MIDA is transformer.MIDA assert embed.MPCA is transformer.MPCA + assert embed.CIFE is transformer.CIFE + assert embed.AJIVE is transformer.AJIVE def test_predict_module_exposes_estimators(): diff --git a/tests/transformer/test_ajive.py b/tests/transformer/test_ajive.py new file mode 100644 index 0000000..d50b535 --- /dev/null +++ b/tests/transformer/test_ajive.py @@ -0,0 +1,236 @@ +import numpy as np +import pytest +from numpy import testing +from sklearn.base import clone + +from kalelinear.transformer import AJIVE +from kalelinear.transformer._ajive import _jive_rand_null_norm +from tests.utils.test_utils import make_common_individual_dataset + + +@pytest.fixture(scope="module") +def multiblock_data(): + return make_common_individual_dataset(random_state=0) + + +def _subspace_error(estimated, planted): + projection = planted @ planted.T + return np.linalg.norm(estimated - projection @ estimated) / np.sqrt(estimated.shape[1]) + + +def test_ajive_recovers_joint_and_individual_subspaces(multiblock_data): + X, groups, _, common_basis = multiblock_data + ajive = AJIVE(initial_ranks=[5, 6, 4], n_resamples=50, random_state=0) + ajive.fit(X, groups=groups) + + assert ajive.n_common_components_ == 2 + assert ajive.common_components_.shape == (X.shape[1], 2) + assert _subspace_error(ajive.common_components_, common_basis) < 1e-6 + testing.assert_array_equal(ajive.individual_ranks_, [3, 4, 2]) + + +def test_ajive_automatic_initial_ranks(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(variance_threshold=1.0, n_resamples=50, random_state=0) + ajive.fit(X, groups=groups) + assert ajive.n_common_components_ == 2 + + +def test_ajive_percentile_uses_larger_perturbation_bound(multiblock_data): + X, groups, _, _ = multiblock_data + ranks = [ + AJIVE(percentile=percentile, n_resamples=50, random_state=0).fit(X, groups=groups).n_common_components_ + for percentile in (5, 50, 95) + ] + # A larger percentile gives a non-decreasing Wedin threshold, so the joint + # rank selected from max(wedin, random-direction) bound cannot increase. + assert ranks[0] >= ranks[1] >= ranks[2] + + +def test_ajive_manual_joint_rank(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE( + n_common_components=1, + initial_ranks=[5, 6, 4], + n_resamples=50, + random_state=0, + ) + ajive.fit(X, groups=groups) + assert ajive.n_common_components_ == 1 + + +def test_ajive_individual_rank_override(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE( + initial_ranks=[5, 6, 4], + n_individual_components=[2, 2, 2], + n_resamples=50, + random_state=0, + ) + ajive.fit(X, groups=groups) + testing.assert_array_equal(ajive.individual_ranks_, [2, 2, 2]) + + +def test_ajive_list_input_matches_stacked(multiblock_data): + X, groups, blocks, _ = multiblock_data + from_stacked = AJIVE(n_resamples=50, random_state=0).fit(X, groups=groups) + from_list = AJIVE(n_resamples=50, random_state=0).fit(blocks) + testing.assert_allclose(from_stacked.common_components_, from_list.common_components_) + + +def test_ajive_transform_consistency(multiblock_data): + X, groups, blocks, _ = multiblock_data + ajive = AJIVE(n_resamples=50, random_state=0) + ajive.fit(X, groups=groups) + + z = ajive.transform(X) + testing.assert_allclose(z, X @ ajive.common_components_) + testing.assert_allclose(ajive.transform(blocks), z) + single_block = ajive.transform([blocks[0]]) + testing.assert_allclose(single_block, blocks[0] @ ajive.common_components_) + testing.assert_allclose(ajive.transform((blocks[0],)), single_block) + + individual = ajive.transform_individual(X, groups=groups) + assert [scores.shape for scores in individual] == [(60, 3), (50, 4), (70, 2)] + individual_from_blocks = ajive.transform_individual(blocks) + for expected, actual in zip(individual, individual_from_blocks): + testing.assert_allclose(expected, actual) + + +def test_ajive_transform_individual_rejects_unknown_block_ids(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(n_resamples=50, random_state=0) + ajive.fit(X, groups=groups) + with pytest.raises(ValueError, match="block ids"): + ajive.transform_individual(X, groups=groups + 10) + + +def test_ajive_transform_individual_preserves_block_assignment(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(n_resamples=50, random_state=0) + ajive.fit(X, groups=groups) + expected = ajive.transform_individual(X, groups=groups) + rng = np.random.default_rng(0) + perm = rng.permutation(len(X)) + actual = ajive.transform_individual(X[perm], groups=groups[perm]) + position = np.empty(len(perm), dtype=int) + position[perm] = np.arange(len(perm)) + for k, block_id in enumerate(np.unique(groups)): + row_indices = np.flatnonzero(groups == block_id) + row_order = np.argsort(position[row_indices]) + testing.assert_allclose(actual[k], expected[k][row_order]) + + +def test_ajive_fit_transform_and_clone(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(n_resamples=50, random_state=0) + z = ajive.fit_transform(X, groups=groups) + assert z.shape == (X.shape[0], 2) + + copied = clone(ajive) + copied.fit(X, groups=groups) + testing.assert_allclose(copied.common_components_, ajive.common_components_) + + +def test_ajive_feature_names(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(n_resamples=50, random_state=0).fit(X, groups=groups) + testing.assert_array_equal(ajive.get_feature_names_out(), np.array(["ajive0", "ajive1"])) + + +def test_ajive_input_validation(): + X = np.ones((10, 4)) + with pytest.raises(ValueError, match="`groups` must be provided"): + AJIVE().fit(X) + with pytest.raises(ValueError, match="same number of features"): + AJIVE().fit([np.ones((5, 3)), np.ones((5, 4))]) + with pytest.raises(ValueError, match="initial_ranks"): + AJIVE(initial_ranks=[5, 6]).fit([np.ones((4, 3)), np.ones((4, 3))]) + + +def test_ajive_zero_energy_block_is_rejected(): + blocks = [np.zeros((5, 3)), np.ones((5, 3))] + with pytest.raises(ValueError, match="positive values"): + AJIVE(n_resamples=50, random_state=0).fit(blocks) + + +def test_ajive_rejects_zero_energy_block_with_explicit_initial_ranks(): + blocks = [np.zeros((5, 3)), np.ones((5, 3))] + with pytest.raises(ValueError, match="zero energy"): + AJIVE(initial_ranks=[2, 1], n_resamples=20, random_state=0).fit(blocks) + + +def test_jive_rand_null_norm_limits_directions_to_null_space_dimension(): + rng = np.random.RandomState(0) + n_ambient, rank = 30, 20 # the null space has only 10 dimensions + data = rng.randn(60, n_ambient) + basis, _ = np.linalg.qr(rng.randn(n_ambient, rank)) + + null_norms = _jive_rand_null_norm(data, basis, 5, rng) + # With the direction count capped at the null-space dimension, the sampled + # directions span the entire null space and the spectral norm becomes the + # exact operator norm of the data restricted to that null space. + expected = np.linalg.norm(data - data @ basis @ basis.T, ord=2) + testing.assert_allclose(null_norms, expected, rtol=1e-8) + + +def test_ajive_accepts_initial_ranks_larger_than_null_space(multiblock_data): + X, groups, _, _ = multiblock_data + ajive = AJIVE(initial_ranks=[20, 18, 17], n_resamples=20, random_state=0) + ajive.fit(X, groups=groups) + assert ajive.common_components_.shape == (X.shape[1], ajive.n_common_components_) + assert np.all(np.isfinite(ajive.common_components_)) + + +def test_ajive_common_rank_is_capped_by_smallest_initial_rank(): + rng = np.random.RandomState(0) + D = 10 + # e2 is shared by blocks 1-2 only; the joint space shared by every block + # is span{e1}, so at most min(initial_ranks) = 1 component can be common. + V1 = np.eye(D)[[0, 1]].T + V2 = np.eye(D)[[0, 1]].T + V3 = np.eye(D)[[0]].T + blocks = [rng.randn(40, V.shape[1]) @ V.T for V in (V1, V2, V3)] + ajive = AJIVE( + n_common_components=5, # more than the theoretical joint-space dimension + initial_ranks=[2, 2, 1], + n_resamples=50, + random_state=0, + ) + ajive.fit(blocks) + assert ajive.n_common_components_ <= 1 + assert ajive.common_components_.shape[1] <= 1 + + +def test_ajive_rejects_initial_ranks_exceeding_numerical_rank(): + rng = np.random.RandomState(0) + # The first block has an exactly-zero column, so its numerical rank is 4 + # even though it is 40 x 5. Requesting rank 5 makes the smallest retained + # singular value zero and the Wedin angle bound undefined. + blocks = [np.column_stack([rng.randn(40, 4), np.zeros(40)]), rng.randn(40, 5), rng.randn(40, 5)] + with pytest.raises(ValueError, match="numerical rank"): + AJIVE(initial_ranks=[5, 2, 2], n_resamples=20, random_state=0).fit(blocks) + + +def test_ajive_rejects_invalid_per_block_rank_specs(multiblock_data): + X, groups, _, _ = multiblock_data + invalid_specs = [ + ([np.nan, 2, 2], "NaN"), + ([np.inf, 2, 2], "infinite"), + ([1.5, 2, 2], "integer"), + ([-1, 2, 2], "non-negative"), + ] + for spec, message in invalid_specs: + with pytest.raises(ValueError, match=message): + AJIVE(n_individual_components=spec, n_resamples=20, random_state=0).fit(X, groups=groups) + with pytest.raises(ValueError, match="NaN"): + AJIVE(initial_ranks=[np.nan, 2, 2], n_resamples=20, random_state=0).fit(X, groups=groups) + + +def test_ajive_transform_individual_rejects_wrong_number_of_blocks(multiblock_data): + X, groups, blocks, _ = multiblock_data + ajive = AJIVE(n_resamples=20, random_state=0).fit(X, groups=groups) + with pytest.raises(ValueError, match="Expected 3 blocks"): + ajive.transform_individual(blocks[:2]) + with pytest.raises(ValueError, match="Expected 3 blocks"): + ajive.transform_individual(blocks + [blocks[0]]) diff --git a/tests/transformer/test_cife.py b/tests/transformer/test_cife.py new file mode 100644 index 0000000..d87068c --- /dev/null +++ b/tests/transformer/test_cife.py @@ -0,0 +1,173 @@ +import numpy as np +import pytest +from numpy import testing +from sklearn.base import clone + +from kalelinear.transformer import CIFE +from tests.utils.test_utils import make_common_individual_dataset + + +@pytest.fixture(scope="module") +def multiblock_data(): + return make_common_individual_dataset(random_state=0) + + +def _subspace_error(estimated, planted): + projection = planted @ planted.T + return np.linalg.norm(estimated - projection @ estimated) / np.sqrt(estimated.shape[1]) + + +def test_cife_recovers_common_and_individual_subspaces(multiblock_data): + X, groups, _, common_basis = multiblock_data + cife = CIFE( + n_common_components=2, + n_individual_components=[3, 4, 2], + tol=1e-10, + random_state=0, + ) + cife.fit(X, groups=groups) + + assert cife.n_common_components_ == 2 + assert cife.common_components_.shape == (X.shape[1], 2) + assert _subspace_error(cife.common_components_, common_basis) < 1e-4 + testing.assert_array_equal(cife.individual_ranks_, [3, 4, 2]) + assert [components.shape[1] for components in cife.individual_components_] == [3, 4, 2] + + +def test_cife_automatic_common_rank(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(tol=1e-10, random_state=0) + cife.fit(X, groups=groups) + assert cife.n_common_components_ == 2 + + +def test_cife_detects_no_common_structure(): + random_state = np.random.RandomState(1) + blocks = [random_state.randn(50, 4) @ random_state.randn(30, 4).T for _ in range(3)] + cife = CIFE(tol=1e-10, random_state=0) + cife.fit(blocks) + assert cife.n_common_components_ == 0 + assert cife.common_components_.shape == (30, 0) + + +def test_cife_list_input_matches_stacked(multiblock_data): + X, groups, blocks, _ = multiblock_data + from_stacked = CIFE(tol=1e-10, random_state=0).fit(X, groups=groups) + from_list = CIFE(tol=1e-10, random_state=0).fit(blocks) + testing.assert_allclose(from_stacked.common_components_, from_list.common_components_) + + +def test_cife_transform_consistency(multiblock_data): + X, groups, blocks, _ = multiblock_data + cife = CIFE(n_common_components=2, n_individual_components=[3, 4, 2], random_state=0) + cife.fit(X, groups=groups) + + z = cife.transform(X) + testing.assert_allclose(z, X @ cife.common_components_) + testing.assert_allclose(cife.transform(blocks), z) + single_block = cife.transform([blocks[0]]) + testing.assert_allclose(single_block, blocks[0] @ cife.common_components_) + testing.assert_allclose(cife.transform((blocks[0],)), single_block) + + individual = cife.transform_individual(X, groups=groups) + assert [scores.shape for scores in individual] == [(60, 3), (50, 4), (70, 2)] + individual_from_blocks = cife.transform_individual(blocks) + for expected, actual in zip(individual, individual_from_blocks): + testing.assert_allclose(expected, actual) + + +def test_cife_transform_individual_rejects_unknown_block_ids(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, n_individual_components=[3, 4, 2], random_state=0) + cife.fit(X, groups=groups) + with pytest.raises(ValueError, match="block ids"): + cife.transform_individual(X, groups=groups + 10) + + +def test_cife_transform_individual_preserves_block_assignment(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, n_individual_components=[3, 4, 2], random_state=0) + cife.fit(X, groups=groups) + expected = cife.transform_individual(X, groups=groups) + rng = np.random.default_rng(0) + perm = rng.permutation(len(X)) + actual = cife.transform_individual(X[perm], groups=groups[perm]) + position = np.empty(len(perm), dtype=int) + position[perm] = np.arange(len(perm)) + for k, block_id in enumerate(np.unique(groups)): + row_indices = np.flatnonzero(groups == block_id) + row_order = np.argsort(position[row_indices]) + testing.assert_allclose(actual[k], expected[k][row_order]) + + +def test_cife_fit_transform_and_clone(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0) + z = cife.fit_transform(X, groups=groups) + assert z.shape == (X.shape[0], 2) + + copied = clone(cife) + copied.fit(X, groups=groups) + testing.assert_allclose(copied.common_components_, cife.common_components_) + + +def test_cife_feature_names(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0).fit(X, groups=groups) + testing.assert_array_equal(cife.get_feature_names_out(), np.array(["cife0", "cife1"])) + + +def test_cife_full_rank_block_requires_pca_dim(): + random_state = np.random.RandomState(2) + blocks = [random_state.randn(40, 10) for _ in range(3)] + with pytest.raises(ValueError, match="whole feature space"): + CIFE(random_state=0).fit(blocks) + cife = CIFE(pca_dim=0.5, random_state=0) + cife.fit(blocks) + assert cife.n_common_components_ == 0 + + +def test_cife_zero_common_components_skips_full_rank_validation(): + random_state = np.random.RandomState(2) + blocks = [random_state.randn(40, 10) for _ in range(3)] + cife = CIFE(n_common_components=0, random_state=0) + cife.fit(blocks) + assert cife.n_common_components_ == 0 + assert cife.common_components_.shape == (10, 0) + testing.assert_array_equal(cife.individual_ranks_, [10, 10, 10]) + + +def test_cife_rejects_infinite_individual_ranks(multiblock_data): + X, groups, _, _ = multiblock_data + with pytest.raises(ValueError, match="infinite"): + CIFE(n_individual_components=[np.inf, 3, 2], random_state=0).fit(X, groups=groups) + + +def test_cife_rejects_invalid_per_block_rank_specs(multiblock_data): + X, groups, _, _ = multiblock_data + invalid_specs = [ + ([np.nan, 3, 2], "NaN"), + ([1.5, 3, 2], "integer"), + ([-1, 3, 2], "non-negative"), + ] + for spec, message in invalid_specs: + with pytest.raises(ValueError, match=message): + CIFE(n_individual_components=spec, random_state=0).fit(X, groups=groups) + + +def test_cife_transform_individual_rejects_wrong_number_of_blocks(multiblock_data): + X, groups, blocks, _ = multiblock_data + cife = CIFE(n_common_components=2, n_individual_components=[3, 4, 2], random_state=0) + cife.fit(X, groups=groups) + with pytest.raises(ValueError, match="Expected 3 blocks"): + cife.transform_individual(blocks[:2]) + + +def test_cife_input_validation(): + X = np.ones((10, 4)) + with pytest.raises(ValueError, match="`groups` must be provided"): + CIFE().fit(X) + with pytest.raises(ValueError, match="same number of features"): + CIFE().fit([np.ones((5, 3)), np.ones((5, 4))]) + with pytest.raises(ValueError, match="(?i)at least two blocks"): + CIFE().fit([np.ones((5, 3))]) diff --git a/tests/transformer/test_multiblock.py b/tests/transformer/test_multiblock.py new file mode 100644 index 0000000..479ee69 --- /dev/null +++ b/tests/transformer/test_multiblock.py @@ -0,0 +1,154 @@ +import numpy as np +import pytest +from numpy import testing + +from kalelinear.transformer import CIFE +from kalelinear.transformer._multiblock import _check_multiblock_input, _check_per_block_ranks +from tests.utils.test_utils import make_common_individual_dataset + + +@pytest.fixture(scope="module") +def multiblock_data(): + return make_common_individual_dataset(random_state=0) + + +def test_check_multiblock_input_rejects_groups_with_list(): + with pytest.raises(ValueError, match="`groups` must be None"): + _check_multiblock_input([np.ones((5, 3)), np.ones((5, 3))], groups=np.zeros(10)) + + +def test_check_multiblock_input_rejects_empty_list(): + with pytest.raises(ValueError, match="at least one block"): + _check_multiblock_input([]) + + +def test_check_multiblock_input_rejects_1d_block(): + with pytest.raises(ValueError, match="2D array"): + _check_multiblock_input([np.ones(5), np.ones((5, 3))]) + + +def test_check_multiblock_input_rejects_empty_block(): + with pytest.raises(ValueError, match="at least one sample"): + _check_multiblock_input([np.ones((0, 3)), np.ones((5, 3))]) + + +def test_check_multiblock_input_rejects_mismatched_features(): + with pytest.raises(ValueError, match="same number of features"): + _check_multiblock_input([np.ones((5, 3)), np.ones((5, 4))]) + + +def test_check_multiblock_input_rejects_single_block_list(): + with pytest.raises(ValueError, match="(?i)at least two blocks"): + _check_multiblock_input([np.ones((5, 3))]) + + +def test_check_multiblock_input_accepts_single_block_with_min_blocks_one(): + blocks, groups, block_ids = _check_multiblock_input([np.ones((5, 3))], min_blocks=1) + assert len(blocks) == 1 + assert groups is None + testing.assert_array_equal(block_ids, [0]) + + +def test_check_multiblock_input_rejects_1d_stacked(): + with pytest.raises(ValueError, match="2D array or a list"): + _check_multiblock_input(np.ones(5)) + + +def test_check_multiblock_input_requires_groups(): + with pytest.raises(ValueError, match="`groups` must be provided"): + _check_multiblock_input(np.ones((10, 3))) + + +def test_check_multiblock_input_validates_groups_shape(): + X = np.ones((10, 3)) + with pytest.raises(ValueError, match="1D array aligned"): + _check_multiblock_input(X, groups=np.ones((10, 1))) + with pytest.raises(ValueError, match="1D array aligned"): + _check_multiblock_input(X, groups=np.zeros(5)) + + +def test_check_multiblock_input_rejects_single_group(): + with pytest.raises(ValueError, match="(?i)at least 2 blocks"): + _check_multiblock_input(np.ones((10, 3)), groups=np.zeros(10, dtype=int)) + + +def test_check_multiblock_input_rejects_nan_group_block(): + X = np.ones((6, 3)) + groups = np.array([0.0, 0.0, 1.0, 1.0, np.nan, np.nan]) + with pytest.raises(ValueError, match="at least one sample"): + _check_multiblock_input(X, groups) + + +def test_check_multiblock_input_orders_blocks_by_first_appearance(): + X = np.arange(6).reshape(6, 1) + groups = np.array([2, 2, 0, 0, 1, 1]) + blocks, _, block_ids = _check_multiblock_input(X, groups) + testing.assert_array_equal([block[0, 0] for block in blocks], [0.0, 2.0, 4.0]) + testing.assert_array_equal(block_ids, [2, 0, 1]) + + +def test_check_per_block_ranks_accepts_integer(): + testing.assert_array_equal(_check_per_block_ranks(3, 4, "ranks"), np.full(4, 3, dtype=int)) + + +def test_check_per_block_ranks_accepts_sequence(): + testing.assert_array_equal(_check_per_block_ranks([1, 2], 2, "ranks"), np.array([1, 2])) + + +def test_check_per_block_ranks_accepts_none(): + assert _check_per_block_ranks(None, 3, "ranks") is None + + +def test_check_per_block_ranks_rejects_wrong_length(): + with pytest.raises(ValueError, match="one value per block"): + _check_per_block_ranks([1, 2], 3, "ranks") + + +def test_check_per_block_ranks_rejects_non_numeric(): + with pytest.raises(ValueError, match="numeric"): + _check_per_block_ranks(["a", "b"], 2, "ranks") + + +def test_check_per_block_ranks_rejects_nan(): + with pytest.raises(ValueError, match="NaN"): + _check_per_block_ranks([1.0, np.nan], 2, "ranks") + + +def test_check_per_block_ranks_rejects_negative(): + with pytest.raises(ValueError, match="non-negative"): + _check_per_block_ranks([1, -1], 2, "ranks") + + +def test_check_per_block_ranks_rejects_non_integer(): + with pytest.raises(ValueError, match="integer"): + _check_per_block_ranks([1.5, 2.0], 2, "ranks") + + +def test_transform_rejects_1d_input(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0).fit(X, groups=groups) + with pytest.raises(ValueError, match="2D array or a list"): + cife.transform(np.ones(X.shape[1])) + + +def test_transform_rejects_wrong_feature_count(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0).fit(X, groups=groups) + with pytest.raises(ValueError, match="Expected .* features, got"): + cife.transform(np.ones((5, X.shape[1] + 1))) + + +def test_transform_individual_rejects_wrong_block_count(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0).fit(X, groups=groups) + with pytest.raises(ValueError, match="Expected 3 blocks, got 2"): + cife.transform_individual([X[:5], X[5:10]]) + + +def test_transform_individual_rejects_wrong_feature_count(multiblock_data): + X, groups, _, _ = multiblock_data + cife = CIFE(n_common_components=2, random_state=0).fit(X, groups=groups) + n_features = X.shape[1] + blocks = [np.ones((5, n_features + 1)) for _ in range(3)] + with pytest.raises(ValueError, match="features in every block"): + cife.transform_individual(blocks) diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index c1776a7..614c373 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -108,3 +108,65 @@ def make_domain_shifted_dataset( domains = domains[idx] return X, y, domains + + +def make_common_individual_dataset( + n_blocks=3, + n_features=30, + n_common=2, + individual_ranks=(3, 4, 2), + n_samples=(60, 50, 70), + noise=0.0, + random_state=None, +): + """Create multiblock data with planted common and individual subspaces. + + Each block is generated as ``B_c A_c^T + B_i A_i^T (+ noise)`` where + ``A_c`` is a common feature-space basis shared by all blocks and ``A_i`` + is a block-specific basis orthogonal to ``A_c``. + """ + individual_ranks = np.asarray(individual_ranks) + n_samples = np.asarray(n_samples) + if individual_ranks.ndim != 1 or individual_ranks.shape[0] != n_blocks: + raise ValueError( + f"`individual_ranks` must be a sequence with one rank per block: " + f"expected {n_blocks} values, got {individual_ranks.size}." + ) + if n_samples.ndim != 1 or n_samples.shape[0] != n_blocks: + raise ValueError( + f"`n_samples` must be a sequence with one sample count per block: " + f"expected {n_blocks} values, got {n_samples.size}." + ) + random_state = check_random_state(random_state) + common_basis, _ = np.linalg.qr(random_state.randn(n_features, n_common)) + blocks = [] + group_lists = [] + for k in range(n_blocks): + individual_basis, _ = np.linalg.qr(random_state.randn(n_features, individual_ranks[k])) + individual_basis -= common_basis @ (common_basis.T @ individual_basis) + individual_basis, _ = np.linalg.qr(individual_basis) + block = random_state.randn(n_samples[k], n_common) @ common_basis.T + block += random_state.randn(n_samples[k], individual_ranks[k]) @ individual_basis.T + if noise: + block += noise * random_state.randn(n_samples[k], n_features) + blocks.append(block) + group_lists.append(np.full(n_samples[k], k)) + X = np.vstack(blocks) + groups = np.concatenate(group_lists) + return X, groups, blocks, common_basis + + +def test_make_common_individual_dataset_validates_block_parameters(): + with pytest.raises(ValueError, match="individual_ranks"): + make_common_individual_dataset(n_blocks=4) + with pytest.raises(ValueError, match="n_samples"): + make_common_individual_dataset(n_blocks=4, individual_ranks=(1, 2, 3, 4)) + + +def test_make_common_individual_dataset_accepts_custom_n_blocks(): + X, groups, blocks, _ = make_common_individual_dataset( + n_blocks=2, individual_ranks=(1, 2), n_samples=(20, 30), random_state=0 + ) + assert len(blocks) == 2 + assert [block.shape[0] for block in blocks] == [20, 30] + assert groups.shape == (50,)