amica-python 0.1.0__tar.gz

amica_python-0.1.0/LICENSE

@@ -0,0 +1,25 @@
+BSD 2-Clause License
+
+Copyright (c) 2015-2020, Jason Palmer and contributors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 

amica_python-0.1.0/PKG-INFO

@@ -0,0 +1,196 @@
+Metadata-Version: 2.4
+Name: amica-python
+Version: 0.1.0
+Summary: Adaptive Mixture ICA in Python
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: loguru
+Requires-Dist: rich
+Requires-Dist: pooch>=1.5
+Requires-Dist: psutil
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: scikit-learn>=1.7.0
+Provides-Extra: torch-cpu
+Requires-Dist: torch; extra == "torch-cpu"
+Provides-Extra: torch-cuda
+Requires-Dist: torch; extra == "torch-cuda"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-timeout; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: mne; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Provides-Extra: doc
+Requires-Dist: sphinx<8.2; extra == "doc"
+Requires-Dist: shibuya; extra == "doc"
+Requires-Dist: sphinx-gallery; extra == "doc"
+Requires-Dist: numpydoc; extra == "doc"
+Requires-Dist: sphinx-design; extra == "doc"
+Requires-Dist: sphinxcontrib-bibtex; extra == "doc"
+Requires-Dist: sphinx-copybutton; extra == "doc"
+Requires-Dist: healpy; extra == "doc"
+Requires-Dist: pandas; extra == "doc"
+Dynamic: license-file
+
+[![codecov](https://codecov.io/github/scott-huberty/amica-python/graph/badge.svg?token=Gt7dvyE9mL)](https://codecov.io/github/scott-huberty/amica-python)
+[![tests](https://github.com/scott-huberty/amica-python/actions/workflows/ci.yaml/badge.svg?branch=main)](https://github.com/scott-huberty/amica-python/actions/workflows/ci.yaml)
+[![docs](https://img.shields.io/github/actions/workflow/status/scott-huberty/amica-python/circleci_redirect.yml?label=Docs)](https://dl.circleci.com/status-badge/redirect/gh/scott-huberty/amica-python/tree/main)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+# AMICA-Python
+### Yes, it's fast.
+
+A Python implementation of the [AMICA](https://sccn.ucsd.edu/~jason/amica_a.pdf) (Adaptive Mixture Independent Component Analysis) algorithm for blind source separation, that was originally [developed in FORTRAN](https://github.com/sccn/amica) by Jason Palmer at the Swartz Center for Computational Neuroscience (SCCN).
+
+AMICA-Python is pre-alpha but is tested against the Fortran implementation and is ready for test driving.
+
+| Python | Fortran |
+|--------|---------|
+| <img src="https://raw.githubusercontent.com/scott-huberty/amica-python/main/docs/source/_static/amica-python.gif" width=400px /> | <img src="https://raw.githubusercontent.com/scott-huberty/amica-python/main/docs/source/_static/amica-fortran.gif" width=400px /> |
+
+
+## Installation
+
+For now, AMICA-Python should be installed from source, and you will have to manually install
+PyTorch (see below) yourself:
+
+```bash
+git clone https://github.com/scott-huberty/amica-python.git
+cd amica-python
+pip install -e .
+```
+
+> [!IMPORTANT]
+> You must install PyTorch before using AMICA-Python.
+
+### Installing PyTorch
+
+Depending on your system and preferences, you can install PyTorch with or without GPU support.
+
+To install the standard version of PyTorch, run:
+
+```bash
+python -m pip install torch
+```
+
+To install the CPU-only version of PyTorch, run:
+
+```bash
+python -m pip install torch --index-url https://download.pytorch.org/whl/cu113
+```
+
+Or for Conda users:
+
+```bash
+conda install -c conda-forge pytorch-cpu
+```
+
+>[!WARNING]
+> If you are using an Intel Mac, you cannot install Pytorch via pip, because there are no precompiled wheels for that platform. Instead, you must install PyTorch via Conda, e.g.:
+
+```bash
+conda install pytorch -c conda-forge
+```
+
+If you use UV, you can also just install torch while installing AMICA-Python:
+
+```bash
+uv pip install -e ".[torch-cpu]"
+```
+
+```bash
+uv pip install -e ".[torch-cuda]"
+```
+
+## Usage
+
+AMICA-Python exposes a scikit-learn style interface. Here is an example of how to use it:
+
+```python
+import numpy as np
+from scipy import signal
+from amica import AMICA
+
+
+rng = np.random.default_rng(0)
+n_samples = 2000
+time = np.linspace(0, 8, n_samples)
+
+s1 = np.sin(2 * time)                     # Sinusoidal
+s2 = np.sign(np.sin(3 * time))            # Square wave
+s3 = signal.sawtooth(2 * np.pi * time)    # Sawtooth
+
+S = np.c_[s1, s2, s3]
+S += 0.2 * rng.standard_normal(S.shape)   # Add noise
+S /= S.std(axis=0)                        # Standardize
+
+A = np.array([[1, 1, 1],
+              [0.5, 2, 1.0],
+              [1.5, 1.0, 2.0]])           # Mixing matrix
+
+X = S @ A.T                               # Observed mixtures
+
+ica = AMICA(random_state=0)
+X_new = ica.fit_transform(X)
+```
+
+<img src="https://scott-huberty.github.io/amica-python/_images/sphx_glr_plot_ica_blind_source_separation_001.png" alt="AMICA-Python vs FastICA outputs" width="50%" style="display: block; margin: 0 auto;"/>
+
+### GPU acceleration
+
+If PyTorch was installed with CUDA support, you can fit AMICA on GPU:
+
+```python
+ica = AMICA(device='cuda', random_state=0)
+```
+
+<br/>
+
+For more examples and documentation, please see the [documentation](https://scott-huberty.github.io/amica-python/).
+
+## What is AMICA?
+
+AMICA is composed of two main ideas, which are hinted at by the name and the title of the original paper:
+*AMICA: An Adaptive Mixture of Independent Component Analyzers with Shared Components*.
+
+#### 1. *Adaptive Mixture* ICA
+
+Standard ICA assumes each source is independent and *non-Gaussian*. Extended Infomax ICA
+improves on this by handling both *sub-Gaussian* and *super-Gaussian* sources. AMICA goes
+further by modeling each source as a *mixture of multiple Gaussians*. This flexibility
+lets AMICA represent virtually any source shape - super-Gaussian, sub-Gaussian,
+or even some funky bimodal distribution:
+
+<img src="docs/source/_static/GMM.png" alt="Source distributions modeled by AMICA" width="25%"/>
+
+In practice, the authors argue that this leads to a more accurate
+approximation of the source signals.
+
+#### 2. *Shared Components*
+
+AMICA can learn multiple ICA decompositions (i.e. models). This is a work around to the assumption of ICA that the sources are
+stationary (they do not change over time). AMICA will
+decide which model best explains the data at each sample, effectively allowing
+the sources to change over time. The "shared components" part of the paper title refers
+to AMICA's ability to allow the various ICA models to share some components (i.e. sources)
+between them, to reduce computational load.
+
+# What does AMICA-Python implement?
+
+In short, AMICA-Python implements point 1 above (Adaptive Mixture ICA),
+but does not implement point 2 (running multiple ICA models simultaneously).
+
+AMICA-Python is powered by [Torch](https://pytorch.org/) and wrapped in an easy-to-use [scikit-learn](https://scikit-learn.org/stable/) style interface.
+
+The outputs are numerically tested against the original FORTRAN implementation to ensure correctness and minimize bugs.
+
+# What wasn't implemented?
+
+  - The ability to model multiple ICA decompositions simultaneously.
+  - The ability to reject unlikely samples based on a thresholded log-likelihood (in the
+    FORTRAN implementation, this is a strategy to deal with artifacts in the data).
+  - AMICA-Python does not expose all the hyper-parameters available in the original FORTRAN implementation.
+    Instead I have tried to pick sensible defaults that should work well in most cases,
+    thus reducing the complexity of the interface.

amica_python-0.1.0/README.md

@@ -0,0 +1,160 @@
+[![codecov](https://codecov.io/github/scott-huberty/amica-python/graph/badge.svg?token=Gt7dvyE9mL)](https://codecov.io/github/scott-huberty/amica-python)
+[![tests](https://github.com/scott-huberty/amica-python/actions/workflows/ci.yaml/badge.svg?branch=main)](https://github.com/scott-huberty/amica-python/actions/workflows/ci.yaml)
+[![docs](https://img.shields.io/github/actions/workflow/status/scott-huberty/amica-python/circleci_redirect.yml?label=Docs)](https://dl.circleci.com/status-badge/redirect/gh/scott-huberty/amica-python/tree/main)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+# AMICA-Python
+### Yes, it's fast.
+
+A Python implementation of the [AMICA](https://sccn.ucsd.edu/~jason/amica_a.pdf) (Adaptive Mixture Independent Component Analysis) algorithm for blind source separation, that was originally [developed in FORTRAN](https://github.com/sccn/amica) by Jason Palmer at the Swartz Center for Computational Neuroscience (SCCN).
+
+AMICA-Python is pre-alpha but is tested against the Fortran implementation and is ready for test driving.
+
+| Python | Fortran |
+|--------|---------|
+| <img src="https://raw.githubusercontent.com/scott-huberty/amica-python/main/docs/source/_static/amica-python.gif" width=400px /> | <img src="https://raw.githubusercontent.com/scott-huberty/amica-python/main/docs/source/_static/amica-fortran.gif" width=400px /> |
+
+
+## Installation
+
+For now, AMICA-Python should be installed from source, and you will have to manually install
+PyTorch (see below) yourself:
+
+```bash
+git clone https://github.com/scott-huberty/amica-python.git
+cd amica-python
+pip install -e .
+```
+
+> [!IMPORTANT]
+> You must install PyTorch before using AMICA-Python.
+
+### Installing PyTorch
+
+Depending on your system and preferences, you can install PyTorch with or without GPU support.
+
+To install the standard version of PyTorch, run:
+
+```bash
+python -m pip install torch
+```
+
+To install the CPU-only version of PyTorch, run:
+
+```bash
+python -m pip install torch --index-url https://download.pytorch.org/whl/cu113
+```
+
+Or for Conda users:
+
+```bash
+conda install -c conda-forge pytorch-cpu
+```
+
+>[!WARNING]
+> If you are using an Intel Mac, you cannot install Pytorch via pip, because there are no precompiled wheels for that platform. Instead, you must install PyTorch via Conda, e.g.:
+
+```bash
+conda install pytorch -c conda-forge
+```
+
+If you use UV, you can also just install torch while installing AMICA-Python:
+
+```bash
+uv pip install -e ".[torch-cpu]"
+```
+
+```bash
+uv pip install -e ".[torch-cuda]"
+```
+
+## Usage
+
+AMICA-Python exposes a scikit-learn style interface. Here is an example of how to use it:
+
+```python
+import numpy as np
+from scipy import signal
+from amica import AMICA
+
+
+rng = np.random.default_rng(0)
+n_samples = 2000
+time = np.linspace(0, 8, n_samples)
+
+s1 = np.sin(2 * time)                     # Sinusoidal
+s2 = np.sign(np.sin(3 * time))            # Square wave
+s3 = signal.sawtooth(2 * np.pi * time)    # Sawtooth
+
+S = np.c_[s1, s2, s3]
+S += 0.2 * rng.standard_normal(S.shape)   # Add noise
+S /= S.std(axis=0)                        # Standardize
+
+A = np.array([[1, 1, 1],
+              [0.5, 2, 1.0],
+              [1.5, 1.0, 2.0]])           # Mixing matrix
+
+X = S @ A.T                               # Observed mixtures
+
+ica = AMICA(random_state=0)
+X_new = ica.fit_transform(X)
+```
+
+<img src="https://scott-huberty.github.io/amica-python/_images/sphx_glr_plot_ica_blind_source_separation_001.png" alt="AMICA-Python vs FastICA outputs" width="50%" style="display: block; margin: 0 auto;"/>
+
+### GPU acceleration
+
+If PyTorch was installed with CUDA support, you can fit AMICA on GPU:
+
+```python
+ica = AMICA(device='cuda', random_state=0)
+```
+
+<br/>
+
+For more examples and documentation, please see the [documentation](https://scott-huberty.github.io/amica-python/).
+
+## What is AMICA?
+
+AMICA is composed of two main ideas, which are hinted at by the name and the title of the original paper:
+*AMICA: An Adaptive Mixture of Independent Component Analyzers with Shared Components*.
+
+#### 1. *Adaptive Mixture* ICA
+
+Standard ICA assumes each source is independent and *non-Gaussian*. Extended Infomax ICA
+improves on this by handling both *sub-Gaussian* and *super-Gaussian* sources. AMICA goes
+further by modeling each source as a *mixture of multiple Gaussians*. This flexibility
+lets AMICA represent virtually any source shape - super-Gaussian, sub-Gaussian,
+or even some funky bimodal distribution:
+
+<img src="docs/source/_static/GMM.png" alt="Source distributions modeled by AMICA" width="25%"/>
+
+In practice, the authors argue that this leads to a more accurate
+approximation of the source signals.
+
+#### 2. *Shared Components*
+
+AMICA can learn multiple ICA decompositions (i.e. models). This is a work around to the assumption of ICA that the sources are
+stationary (they do not change over time). AMICA will
+decide which model best explains the data at each sample, effectively allowing
+the sources to change over time. The "shared components" part of the paper title refers
+to AMICA's ability to allow the various ICA models to share some components (i.e. sources)
+between them, to reduce computational load.
+
+# What does AMICA-Python implement?
+
+In short, AMICA-Python implements point 1 above (Adaptive Mixture ICA),
+but does not implement point 2 (running multiple ICA models simultaneously).
+
+AMICA-Python is powered by [Torch](https://pytorch.org/) and wrapped in an easy-to-use [scikit-learn](https://scikit-learn.org/stable/) style interface.
+
+The outputs are numerically tested against the original FORTRAN implementation to ensure correctness and minimize bugs.
+
+# What wasn't implemented?
+
+  - The ability to model multiple ICA decompositions simultaneously.
+  - The ability to reject unlikely samples based on a thresholded log-likelihood (in the
+    FORTRAN implementation, this is a strategy to deal with artifacts in the data).
+  - AMICA-Python does not expose all the hyper-parameters available in the original FORTRAN implementation.
+    Instead I have tried to pick sensible defaults that should work well in most cases,
+    thus reducing the complexity of the interface.

amica_python-0.1.0/pyproject.toml

@@ -0,0 +1,97 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "amica-python"
+version = "0.1.0"
+description = "Adaptive Mixture ICA in Python"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "loguru",
+    "rich",
+    "pooch >= 1.5", 
+    "psutil",
+    "numpy>=2.2.6",
+    "scikit-learn>=1.7.0",
+]
+
+[project.optional-dependencies]
+torch-cpu = ["torch"]
+torch-cuda = ["torch"]
+dev = ["pytest", "pytest-cov", "pytest-timeout", "matplotlib", "mne", "ruff"]
+doc = [
+    "sphinx<8.2",
+    "shibuya",
+    "sphinx-gallery",
+    "numpydoc",
+    "sphinx-design",
+    "sphinxcontrib-bibtex",
+    "sphinx-copybutton",
+    # For Tutorials
+    "healpy",
+    "pandas",  # Needed to load the MNIST dataset example
+    # "smica @ git+https://github.com/scott-huberty/smica.git",
+    ]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["amica*"]
+exclude = ["amica.tests*"]
+
+[tool.ruff.lint]
+select = ["A", "B006", "D", "E", "F", "I", "UP", "UP031", "W"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.ruff.lint.per-file-ignores]
+"src/amica/_types.py" = ["E501"]  # Line too long
+"src/amica/**/__init__.py" = ["D104"]  # Missing docstring in public package
+"src/amica/**/tests/*.py" = ["D100"]  # Missing docstring in public module
+
+
+[[tool.uv.index]]
+name = "pytorch_cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true # only fetch from this index if we explicitly map a package there.
+
+[tool.uv.sources]
+torch = { index = "pytorch_cpu" }
+markupsafe = { index = "pytorch_cpu" }
+
+[tool.pytest.ini_options]
+addopts = [
+    "--cov=amica",
+    "--cov-branch",
+    "--cov-report=xml",
+    "--cov-report=term",
+    "--ignore=src/amica/tests/test_kernels.py",
+]
+markers = [
+    "sklearn_api: Tests that validate Scikit-Learn API conformance",
+    "slow: Marks tests as slow (deselect with `pytest -m not slow`)"
+]
+
+[tool.coverage.run]
+source = ["amica"] # Source files to measure.
+branch = true # Add branch coverage to the analysis.
+omit = [
+    "*/tests/*",
+    ".venv/*",
+]
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+]
+show_missing = true
+skip_covered = true
+
+[tool.coverage.xml]
+output = "coverage.xml"

amica_python-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

amica_python-0.1.0/src/amica/__init__.py

@@ -0,0 +1,5 @@
+from . import datasets, utils
+from ._sklearn_interface import AMICA
+from .core import fit_amica
+
+__all__ = ['fit_amica', 'AMICA', 'datasets', 'utils']

amica_python-0.1.0/src/amica/_batching.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Union
+from warnings import warn
+
+import numpy as np
+import psutil
+import torch
+
+ArrayLike2D = Union[np.ndarray, "np.typing.NDArray[np.floating]"]
+
+
+class BatchLoader:
+    """Iterate over an array in fixed-size batches of data along a chosen axis.
+
+    We hand rolled this instead of using DataLoader because 1) we want to yield
+    slices of input array (i.e. a view), and 2) return the indices as
+    a slice object. DataLoader would internally convert the slice into a tensor
+    of indices.
+
+    Example (AMICA shape):
+        X: (n_samples, n_features)
+        it = BatchLoader(X, axis=0, batch_size=4096)
+        for X_blk, sl in it:
+            # X_blk is X[sl, :] where sl is slice(start, end)
+            ...
+    """
+
+    def __init__(self, X: ArrayLike2D, axis: int, batch_size: int | None = None):
+        # Validate inputs
+        cls_name = self.__class__.__name__
+        if not isinstance(X, torch.Tensor):
+            raise TypeError(f"{cls_name} expects a torch.Tensor")  # pragma: no cover
+        if X.ndim < 1:
+            raise ValueError(
+                f"{cls_name} expects an array with at least 1 dimension"
+            )  # pragma: no cover
+        self.X = X
+        self.axis = axis
+
+        if self.axis < 0:
+            self.axis += X.ndim
+        if not (0 <= self.axis < X.ndim):
+            raise ValueError(
+                f"axis {self.axis} out of bounds for array with ndim={X.ndim}"
+                )
+
+        # Determine batching parameters
+        n = X.shape[self.axis]
+        start = 0
+        stop = n
+        if batch_size is None:
+            # Treat as single chunk spanning [start:stop]
+            batch_size = stop
+
+        # Validate parameters
+        assert (0 <= start <= n), f"start {start} out of range [0, {n}]"
+        assert (0 <= stop <= n), f"stop {stop} out of range [0, {n}]"
+        assert start <= stop, f"start {start} must be <= stop {stop}"
+        if batch_size < 0:
+            raise ValueError(f"batch_size must be positive. Got {batch_size}.")
+        if batch_size > X.shape[self.axis]:
+            raise ValueError(
+                f"batch_size {batch_size} exceeds data size {X.shape[self.axis]} "
+                f"along axis {self.axis}."
+            )
+
+        # Store parameters
+        self.start = start
+        self.stop = stop
+        self.batch_size = int(batch_size)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        start = self.start + idx * self.batch_size
+        stop = min(start + self.batch_size, self.stop)
+
+
+        idx = [slice(None)] * self.X.ndim
+        idx[self.axis] = slice(start, stop)
+        return self.X[tuple(idx)]
+
+    def __iter__(self) -> Iterator[tuple[np.ndarray, slice]]:
+        axis = self.axis
+        start = self.start
+        stop = self.stop
+        step = self.batch_size
+
+        idx = [slice(None)] * self.X.ndim
+        assert -((stop - start) // -step) == len(self)  # sanity check
+        for s in range(start, stop, step):
+            e = min(s + step, stop)
+            batch_slice = slice(s, e)
+            idx[axis] = batch_slice
+            yield self.X[tuple(idx)], batch_slice
+
+    def __len__(self) -> int:
+        return (self.X.shape[self.axis] + self.batch_size - 1) // self.batch_size
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(Data shape: {self.X.shape}, "
+            f"Batched axis: {self.axis}, batch_size: {self.batch_size}, "
+            f"n_batches: {len(self)})"
+        )
+
+def choose_batch_size(
+        *,
+        N: int,
+        n_comps: int,
+        n_mix: int,
+        n_models: int = 1,
+        dtype: np.dtype = np.float64,
+        memory_fraction: float = 0.25,      # use up to 25% of available memory
+        memory_cap: float = 1.5 * 1024**3,  # 1.5 GB absolute ceiling
+        ) -> int:
+    """
+    Choose batch size for processing data in chunks.
+
+    Parameters
+    ----------
+    N : int
+        Total number of samples.
+    n_comps : int
+        Number of components to be learned in the model, e.g. size of the n_components
+        dimension of the data.
+    n_mix : int
+        Number of mixture components per source/component to be learned in the model.
+    dtype : np.dtype, optional
+        Data type of the input data, by default np.float64.
+    memory_cap : float, optional
+        Maximum memory (in bytes) to be used for processing, by default
+        ``1.5 * 1024**3`` (1.5 GB).
+
+    Notes
+    -----
+    The batch size is primarily determined by the estimated size of hot buffers (e.g.
+    y, z, fp, ufp), which scale with the size of n_samples:
+    - One array of shape (N,):
+        - loglik
+    - Two arrays of shape (N, n_models):
+        - modloglik
+        - v (model responsibilities)
+    - Two arrays of shape (N, n_comps)
+        - b
+        - g
+    - Five arrays of shape (N, n_comps, n_mix): u, y, z, fp, ufp
+        - u (mixture responsibilities)
+        - y
+        - z
+        - fp
+        - ufp
+    """
+    dtype_size = np.dtype(dtype).itemsize
+    # per-sample cost across pre-allocated buffers
+    bytes_per_sample = (
+        1                       # loglik
+        + 2 * n_models          # modloglik, v
+        + 2 * n_comps           # b, g
+        + 5 * n_comps * n_mix   # fp, u, ufp, y, z,
+        ) * dtype_size
+    # Plus small headroom for intermediates
+    bytes_per_sample = int(bytes_per_sample * 1.2)
+
+    # Pick memory budget
+    try:
+        hard_cap = 4 * 1024**3  # 4 GiB (avoid runaway memory use)
+        avail_mem = psutil.virtual_memory().available
+        mem_cap = min(avail_mem * memory_fraction, hard_cap)
+    except Exception:
+        mem_cap = memory_cap  # fallback to user-specified cap
+
+    max_batch_size = mem_cap // bytes_per_sample
+
+    # Ensure at least 1 sample. This should only trigger if n_comps and n_mix are huge.
+    if max_batch_size < 1:
+        raise MemoryError(
+            f"Cannot fit even 1 sample within memory cap of "
+            f"{mem_cap / 1024**3:.2f} GiB. "
+            f"Per-sample memory cost is {bytes_per_sample / 1024**3:.2f} GB."
+        )
+    batch_size = int(min(N, max_batch_size))
+
+    # Heuristic floor, we don't want absurdly small chunks or chunks that are too
+    # small relative to the model complexity (n_comps)
+    # This heuristic works well for typical ICA regimes, where n_comps is < 256
+    min_batch_size = max(8192, n_comps * 32)  # at least 32 samples per component
+    min_batch_size = min(min_batch_size, N)  # Cannot exceed N
+    if batch_size < min_batch_size:
+        warn(
+            f"Warning: To stay within the memory cap, batch size is {batch_size} "
+            f"samples, which is below the recommended minimum of {min_batch_size}."
+        )
+    return batch_size