dextrademixer 0.0.1__tar.gz

dextrademixer-0.0.1/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, Benjamin Schubert
+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.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+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.

dextrademixer-0.0.1/PKG-INFO

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: dextrademixer
+Version: 0.0.1
+Summary: Probabilistic identification of antigen-specific T cells from single-cell pMHC multimer data
+Author-email: Yang an <yang.an@helmholtz-munich.de>, Benjamin Schubert <benjamin.schubert@helmholtz-munich.de>
+Maintainer-email: Yang an <yang.an@helmholtz-munich.de>, Benjamin Schubert <benjamin.schubert@helmholtz-munich.de>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/SchubertLab/DextraDemixer
+Project-URL: Issues, https://github.com/SchubertLab/DextraDemixer/issues
+Project-URL: Publication, https://doi.org/10.64898/2026.06.23.733339
+Project-URL: Repository, https://github.com/SchubertLab/DextraDemixer.git
+Keywords: antigen specificity,pMHC multimer,single-cell,T cell
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Programming Language :: Python :: 3 :: Only
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: anndata>=0.12.0
+Requires-Dist: arviz<1.0,>=0.22.0
+Requires-Dist: jax>=0.5.1
+Requires-Dist: matplotlib>=3.10.7
+Requires-Dist: mudata>=0.3.2
+Requires-Dist: muon>=0.1.7
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: numpyro>=0.16.0
+Requires-Dist: optax>=0.1.7
+Requires-Dist: pandas>=2.3.3
+Requires-Dist: scanpy>=1.10.0
+Requires-Dist: scikit-learn>=1.7.2
+Requires-Dist: scipy>=1.12.0
+Requires-Dist: scirpy>=0.13.0
+Requires-Dist: seaborn>=0.13.2
+Requires-Dist: statsmodels>=0.14.1
+Requires-Dist: tqdm>=4.68.3
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Dynamic: license-file
+
+# DextraDemixer
+
+[![bioRxiv](https://img.shields.io/badge/bioRxiv-10.64898%2F2026.06.23.733339-b31b1b)](https://doi.org/10.64898/2026.06.23.733339)
+[![License: BSD 3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://github.com/SchubertLab/DextraDemixer/blob/main/LICENSE)
+
+**Probabilistic identification of antigen-specific T cells from single-cell pMHC multimer experiments.**
+
+DextraDemixer models pMHC multimer UMI counts as a mixture of specific and nonspecific binding components. It returns a posterior binding probability and assignment for each cell, helping separate antigen-specific signal from background binding and assay noise.
+
+The method is described in [*DextraDemixer enables accurate identification of antigen-specific T cells from pMHC multimer experiments*](https://doi.org/10.64898/2026.06.23.733339).
+
+> [!NOTE]
+> DextraDemixer is research software under active development. Interfaces may change while the package matures.
+
+## Highlights
+
+- Infers cell-level posterior probabilities of antigen specificity from pMHC UMI counts.
+- Supports optional negative-control multimers.
+- Incorporates clonotype information through clone-level probability aggregation.
+- Provides fixed-threshold and Bayesian false discovery rate (FDR) assignments.
+- Works directly with [`MuData`](https://mudata.readthedocs.io/), [`AnnData`](https://anndata.readthedocs.io/), and [`Pandas DataFrames`](https://pandas.pydata.org/) objects and uses [JAX](https://jax.readthedocs.io/) and [NumPyro](https://num.pyro.ai/) for inference.
+
+## Installation
+
+```bash
+pip install dextrademixer
+```
+
+### Alternative: Conda environment
+
+For a conda environment.
+
+```bash
+git clone https://github.com/SchubertLab/DextraDemixer.git
+cd DextraDemixer
+conda env create -f environment.yaml
+conda activate dextrademixer
+pip install -e .  # or pip install dextrademixer
+```
+
+The final command installs the local package in editable mode. Runtime dependency constraints are aligned between `pyproject.toml` and `environment.yaml`. To reproduce the numbers from the manuscript, use `environment_reproducible.yaml` instead: it pins the exact versions, and results can shift slightly with other JAX/NumPyro releases.
+
+## Quick start
+
+### Single pMHC
+
+The example below fits one pMHC feature and classifies cells using a posterior-probability threshold of 0.5.
+
+```python
+import muon as mu
+
+from dextrademixer import DextraDemixer
+
+mdata = mu.read("data/example_data.h5mu")
+
+model = DextraDemixer().fit(mdata, pmhc_key="pmhc1", pmhc_modality_key="gex")
+p_binder, is_binder = model.predict(threshold=0.5)
+
+mdata.mod["gex"].obs["dextrademixer_probability"] = p_binder
+mdata.mod["gex"].obs["dextrademixer_assignment"] = is_binder
+```
+
+To include a negative-control multimer, pass `neg_ctrl_key`. To aggregate probabilities within clonotypes, pass the clonotype column as `ir_clone_key` and set `clonotype_median_p=True` when predicting:
+
+```python
+model = DextraDemixer().fit(
+    mdata,
+    pmhc_key="pmhc1",
+    pmhc_modality_key="gex",
+    neg_ctrl_key="neg_control",
+    ir_clone_key="clone_id",
+)
+
+p_pred, assignment = model.predict(threshold=0.5, clonotype_median_p=True)
+```
+
+Use either `threshold` for a fixed decision boundary or `target_fdr` for FDR-controlled assignments; do not specify both.
+
+### Multiple pMHCs
+
+`DextraDemixerMulti` fits one independent `DextraDemixer` per pMHC and returns cells x pMHC tables instead of vectors. Each fit stays reachable as `model.demixers[pmhc_key]`.
+
+```python
+from dextrademixer import DextraDemixerMulti
+
+mdata = mu.read("data/example_multi_pmhc.h5mu")
+
+model = DextraDemixerMulti().fit(
+    mdata,
+    pmhc_keys=["pmhc1", "pmhc2", "pmhc3"],
+    pmhc_modality_key="gex",
+    neg_ctrl_key="neg_control",
+)
+
+p_pred, assignment = model.predict(threshold=0.5, clonotype_median_p=True, max_prob=True)
+```
+
+Because the pMHCs are fitted independently, a cell can pass the threshold for several. `max_prob=True` keeps only the pMHC with the highest probability per cell, combined with `clonotype_median_p=True` the same choice is made per clonotype.
+
+## Input data
+
+DextraDemixer needs the pMHC UMI counts and, optionally, a cell-level clonotype identifier. It reads them from any of:
+
+- A [`MuData`](https://mudata.readthedocs.io/) object: counts in the `.X` of the feature modality (`pmhc_modality_key`, `gex` by default), clonotypes in the `.obs` of the AIRR modality (`ir_modality_key`, `airr` by default).
+- An [`AnnData`](https://anndata.readthedocs.io/) object: counts in `.X`, clonotypes in `.obs`; `pmhc_modality_key` and `ir_modality_key` are then unused.
+- A cells x features [`DataFrame`](https://pandas.pydata.org/): counts and annotation in the same table, so every key is a column name.
+
+`pmhc_key` (or `pmhc_keys`) and the optional `neg_ctrl_key` name the count columns, `ir_clone_key` the clonotype column, whose ids may be integers or strings. The bundled example dataset ships in all three formats: `data/example_data.h5mu`, `.h5ad` and `.csv`. `data/example_multi_pmhc.h5mu` is a simulated three-pMHC panel.
+
+See [`Tutorial.ipynb`](https://github.com/SchubertLab/DextraDemixer/blob/main/Tutorial.ipynb) for a complete, reproducible workflow using the bundled example dataset, including configuration, model fitting, evaluation, and visualization, and [`Tutorial DextraDemixerMulti.ipynb`](https://github.com/SchubertLab/DextraDemixer/blob/main/Tutorial%20DextraDemixerMulti.ipynb) for the same with multiple pMHCs:
+
+```bash
+jupyter lab Tutorial.ipynb
+```
+
+## Reproducing the manuscript
+
+The analyses and figure-generation code used in the manuscript are maintained in the [DextraDemixer reproducibility repository](https://github.com/SchubertLab/DextraDemixer_reproducibility).
+
+## Contributing
+
+Bug reports, feature requests, and pull requests are welcome through the [GitHub issue tracker](https://github.com/SchubertLab/DextraDemixer/issues). When reporting a problem, please include a minimal reproducible example and details of your operating system and Conda environment.
+
+## Citation
+
+If DextraDemixer supports your research, please cite:
+
+```bibtex
+@article{An2026DextraDemixer,
+  author       = {An, Yang and Drost, Felix and Bonafonte-Pard{\`a}s, Irene and Grotz, Myriam and Schober, Kilian and Schubert, Benjamin},
+  title        = {DextraDemixer enables accurate identification of antigen-specific T cells from pMHC multimer experiments},
+  journal      = {bioRxiv},
+  year         = {2026},
+  publisher    = {Cold Spring Harbor Laboratory},
+  doi          = {10.64898/2026.06.23.733339},
+  url          = {https://doi.org/10.64898/2026.06.23.733339},
+  elocation-id = {2026.06.23.733339}
+}
+```
+
+## License
+
+DextraDemixer is distributed under the [BSD 3-Clause License](https://github.com/SchubertLab/DextraDemixer/blob/main/LICENSE).

dextrademixer-0.0.1/README.md

@@ -0,0 +1,143 @@
+# DextraDemixer
+
+[![bioRxiv](https://img.shields.io/badge/bioRxiv-10.64898%2F2026.06.23.733339-b31b1b)](https://doi.org/10.64898/2026.06.23.733339)
+[![License: BSD 3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://github.com/SchubertLab/DextraDemixer/blob/main/LICENSE)
+
+**Probabilistic identification of antigen-specific T cells from single-cell pMHC multimer experiments.**
+
+DextraDemixer models pMHC multimer UMI counts as a mixture of specific and nonspecific binding components. It returns a posterior binding probability and assignment for each cell, helping separate antigen-specific signal from background binding and assay noise.
+
+The method is described in [*DextraDemixer enables accurate identification of antigen-specific T cells from pMHC multimer experiments*](https://doi.org/10.64898/2026.06.23.733339).
+
+> [!NOTE]
+> DextraDemixer is research software under active development. Interfaces may change while the package matures.
+
+## Highlights
+
+- Infers cell-level posterior probabilities of antigen specificity from pMHC UMI counts.
+- Supports optional negative-control multimers.
+- Incorporates clonotype information through clone-level probability aggregation.
+- Provides fixed-threshold and Bayesian false discovery rate (FDR) assignments.
+- Works directly with [`MuData`](https://mudata.readthedocs.io/), [`AnnData`](https://anndata.readthedocs.io/), and [`Pandas DataFrames`](https://pandas.pydata.org/) objects and uses [JAX](https://jax.readthedocs.io/) and [NumPyro](https://num.pyro.ai/) for inference.
+
+## Installation
+
+```bash
+pip install dextrademixer
+```
+
+### Alternative: Conda environment
+
+For a conda environment.
+
+```bash
+git clone https://github.com/SchubertLab/DextraDemixer.git
+cd DextraDemixer
+conda env create -f environment.yaml
+conda activate dextrademixer
+pip install -e .  # or pip install dextrademixer
+```
+
+The final command installs the local package in editable mode. Runtime dependency constraints are aligned between `pyproject.toml` and `environment.yaml`. To reproduce the numbers from the manuscript, use `environment_reproducible.yaml` instead: it pins the exact versions, and results can shift slightly with other JAX/NumPyro releases.
+
+## Quick start
+
+### Single pMHC
+
+The example below fits one pMHC feature and classifies cells using a posterior-probability threshold of 0.5.
+
+```python
+import muon as mu
+
+from dextrademixer import DextraDemixer
+
+mdata = mu.read("data/example_data.h5mu")
+
+model = DextraDemixer().fit(mdata, pmhc_key="pmhc1", pmhc_modality_key="gex")
+p_binder, is_binder = model.predict(threshold=0.5)
+
+mdata.mod["gex"].obs["dextrademixer_probability"] = p_binder
+mdata.mod["gex"].obs["dextrademixer_assignment"] = is_binder
+```
+
+To include a negative-control multimer, pass `neg_ctrl_key`. To aggregate probabilities within clonotypes, pass the clonotype column as `ir_clone_key` and set `clonotype_median_p=True` when predicting:
+
+```python
+model = DextraDemixer().fit(
+    mdata,
+    pmhc_key="pmhc1",
+    pmhc_modality_key="gex",
+    neg_ctrl_key="neg_control",
+    ir_clone_key="clone_id",
+)
+
+p_pred, assignment = model.predict(threshold=0.5, clonotype_median_p=True)
+```
+
+Use either `threshold` for a fixed decision boundary or `target_fdr` for FDR-controlled assignments; do not specify both.
+
+### Multiple pMHCs
+
+`DextraDemixerMulti` fits one independent `DextraDemixer` per pMHC and returns cells x pMHC tables instead of vectors. Each fit stays reachable as `model.demixers[pmhc_key]`.
+
+```python
+from dextrademixer import DextraDemixerMulti
+
+mdata = mu.read("data/example_multi_pmhc.h5mu")
+
+model = DextraDemixerMulti().fit(
+    mdata,
+    pmhc_keys=["pmhc1", "pmhc2", "pmhc3"],
+    pmhc_modality_key="gex",
+    neg_ctrl_key="neg_control",
+)
+
+p_pred, assignment = model.predict(threshold=0.5, clonotype_median_p=True, max_prob=True)
+```
+
+Because the pMHCs are fitted independently, a cell can pass the threshold for several. `max_prob=True` keeps only the pMHC with the highest probability per cell, combined with `clonotype_median_p=True` the same choice is made per clonotype.
+
+## Input data
+
+DextraDemixer needs the pMHC UMI counts and, optionally, a cell-level clonotype identifier. It reads them from any of:
+
+- A [`MuData`](https://mudata.readthedocs.io/) object: counts in the `.X` of the feature modality (`pmhc_modality_key`, `gex` by default), clonotypes in the `.obs` of the AIRR modality (`ir_modality_key`, `airr` by default).
+- An [`AnnData`](https://anndata.readthedocs.io/) object: counts in `.X`, clonotypes in `.obs`; `pmhc_modality_key` and `ir_modality_key` are then unused.
+- A cells x features [`DataFrame`](https://pandas.pydata.org/): counts and annotation in the same table, so every key is a column name.
+
+`pmhc_key` (or `pmhc_keys`) and the optional `neg_ctrl_key` name the count columns, `ir_clone_key` the clonotype column, whose ids may be integers or strings. The bundled example dataset ships in all three formats: `data/example_data.h5mu`, `.h5ad` and `.csv`. `data/example_multi_pmhc.h5mu` is a simulated three-pMHC panel.
+
+See [`Tutorial.ipynb`](https://github.com/SchubertLab/DextraDemixer/blob/main/Tutorial.ipynb) for a complete, reproducible workflow using the bundled example dataset, including configuration, model fitting, evaluation, and visualization, and [`Tutorial DextraDemixerMulti.ipynb`](https://github.com/SchubertLab/DextraDemixer/blob/main/Tutorial%20DextraDemixerMulti.ipynb) for the same with multiple pMHCs:
+
+```bash
+jupyter lab Tutorial.ipynb
+```
+
+## Reproducing the manuscript
+
+The analyses and figure-generation code used in the manuscript are maintained in the [DextraDemixer reproducibility repository](https://github.com/SchubertLab/DextraDemixer_reproducibility).
+
+## Contributing
+
+Bug reports, feature requests, and pull requests are welcome through the [GitHub issue tracker](https://github.com/SchubertLab/DextraDemixer/issues). When reporting a problem, please include a minimal reproducible example and details of your operating system and Conda environment.
+
+## Citation
+
+If DextraDemixer supports your research, please cite:
+
+```bibtex
+@article{An2026DextraDemixer,
+  author       = {An, Yang and Drost, Felix and Bonafonte-Pard{\`a}s, Irene and Grotz, Myriam and Schober, Kilian and Schubert, Benjamin},
+  title        = {DextraDemixer enables accurate identification of antigen-specific T cells from pMHC multimer experiments},
+  journal      = {bioRxiv},
+  year         = {2026},
+  publisher    = {Cold Spring Harbor Laboratory},
+  doi          = {10.64898/2026.06.23.733339},
+  url          = {https://doi.org/10.64898/2026.06.23.733339},
+  elocation-id = {2026.06.23.733339}
+}
+```
+
+## License
+
+DextraDemixer is distributed under the [BSD 3-Clause License](https://github.com/SchubertLab/DextraDemixer/blob/main/LICENSE).

dextrademixer-0.0.1/dextrademixer/__init__.py

@@ -0,0 +1,4 @@
+from dextrademixer.model import *
+from dextrademixer.model import __all__ as _model_all
+
+__all__ = list(_model_all)

dextrademixer-0.0.1/dextrademixer/model/ApMHCDeconvolution.py

@@ -0,0 +1,207 @@
+"""
+The interface shared by all pMHC deconvolution methods.
+
+`ApMHCDeconvolution` fixes the call order every method follows and holds the parts that do not
+depend on the model: resolving the supported input containers into counts and annotation
+(`as_counts`), validating them (`_check_parameters`), and turning posterior probabilities into class
+assignments by threshold or local-FDR control (`_predict_posterior_class`). DextraDemixer and the
+comparison baselines BEAM, ICON and ITRAP all build on it.
+"""
+from __future__ import annotations
+
+import abc
+
+from typing import TYPE_CHECKING, Tuple, Union
+
+import anndata as ad
+import mudata as md
+import pandas as pd
+
+from jax import lax
+import jax.numpy as jnp
+
+if TYPE_CHECKING:
+    from jax._src.typing import Array
+
+Data = Union[md.MuData, ad.AnnData, pd.DataFrame]
+
+
+class ApMHCDeconvolution:
+    """
+    Common interface of the pMHC deconvolution methods, i.e. DextraDemixer and the comparison
+    baselines. It fixes the call order `preprocess_model_data` -> `fit` -> `predict_posterior_class`
+    and provides the model-agnostic parts of it.
+    """
+
+    @abc.abstractmethod
+    def preprocess_model_data(self,
+                              data: Data,
+                              pmhc_key: str,
+                              pmhc_modality_key: str = "gex",
+                              neg_ctrl_key: str = None,
+                              ir_modality_key: str = "airr",
+                              ir_clone_key: str = None,
+                              **kwargs):
+        """
+        Extracts the counts and the annotation a method needs from `data`, see `as_counts`.
+
+        Args:
+            data: the pMHC counts, as a MuData, an AnnData or a cells x features DataFrame.
+            pmhc_key: the pMHC count column to deconvolve.
+            pmhc_modality_key: the MuData modality holding the counts.
+            neg_ctrl_key: (Optional) the negative control count column.
+            ir_modality_key: the MuData AIRR module key.
+            ir_clone_key: (Optional) the `obs` column holding clonotype ids.
+            kwargs: method-specific extras.
+        """
+
+    @abc.abstractmethod
+    def fit(self, *args, **kwargs):
+        """
+        Fits the method on the data prepared by `preprocess_model_data`.
+
+        Returns:
+            self, so that `predict_posterior_class` can be chained onto the call.
+        """
+
+    @abc.abstractmethod
+    def predict_posterior_class(self,
+                                threshold: float = None,
+                                target_fdr: float = None
+                                ) -> Tuple[Array, Array]:
+        """
+        Assigns each cell to the binding or non-binding class.
+
+        Args:
+            threshold: (Optional) probability in [0,1] above which a cell is called a binder.
+            target_fdr: (Optional) FDR to control instead of using a fixed threshold. Mutually
+                        exclusive with `threshold`.
+
+        Returns:
+            A tuple (p_pred, assignment) of per-cell binding probabilities and 0/1 assignments.
+        """
+
+    def predict(self, *args, **kwargs):
+        """Alias for `predict_posterior_class`, for the familiar scikit-learn call name."""
+        return self.predict_posterior_class(*args, **kwargs)
+
+    @staticmethod
+    def _predict_posterior_class(p_pred: Array,
+                                 threshold: float = None,
+                                 target_fdr: float = None
+                                 ) -> Array:
+        """
+        Turns per-cell binding probabilities into 0/1 assignments, either at a fixed threshold or
+        at the largest threshold whose estimated FDR stays below `target_fdr`.
+
+        Args:
+            p_pred: per-cell posterior probability of binding, shape (n_cells,).
+            threshold: (Optional) probability in [0,1] above which a cell is called a binder.
+            target_fdr: (Optional) FDR to control instead. Mutually exclusive with `threshold`;
+                        if neither is given, a threshold of 0.5 is used.
+
+        Returns:
+            The 0/1 class assignment, shape (n_cells,).
+
+        Raises:
+            ValueError: if both `threshold` and `target_fdr` are given, or either is outside [0,1].
+        """
+        if threshold is not None and target_fdr is not None:
+            raise ValueError("Please specify either a manual `threshold` or a `target_fdr` but not both.")
+
+        if threshold is not None and not (0 <= threshold <= 1):
+            raise ValueError(f"`threshold`must be in [0,1] but was {threshold}")
+
+        if target_fdr is not None and not (0 <= target_fdr <= 1):
+            raise ValueError(f"`target_fdr`must be in [0,1] but was {target_fdr}")
+
+        if threshold is None and target_fdr is None:
+            threshold = 0.5
+
+        # posterior probability of belonging to the binding class
+        if target_fdr is not None:
+            n_cells = p_pred.shape[0]
+
+            # Calculate the local FDR (1 - p)
+            lfdr = 1 - p_pred
+
+            sorted_indices = jnp.argsort(p_pred)[::-1]
+            sorted_p = p_pred[sorted_indices]
+            sorted_lfdr = lfdr[sorted_indices]
+
+            cumulative_lfdr = jnp.cumsum(sorted_lfdr)
+            cumulative_count = jnp.arange(1, n_cells + 1)
+
+            # Estimated FDR for each possible threshold
+            estimated_fdr = cumulative_lfdr / cumulative_count
+
+            # Find the largest index k such that estimated_fdr[k] <= target_fdr
+            valid_thresholds = estimated_fdr <= target_fdr
+            max_k = jnp.max(jnp.where(valid_thresholds, jnp.arange(n_cells), -1))
+            threshold = lax.cond(max_k >= 0, lambda: sorted_p[max_k], lambda: 1.0)
+
+        assignment = (p_pred >= threshold).astype("int32")
+        return assignment
+
+
+    @staticmethod
+    def as_counts(data: Data, pmhc_modality_key: str = "gex", ir_modality_key: str = "airr") -> Tuple[pd.DataFrame, pd.DataFrame]:
+        """
+        Normalizes any supported input into a (counts, obs) pair of DataFrames, so that count
+        columns are addressed by `counts[key]` and per-cell annotation by `obs[key]`:
+
+        - MuData: counts from the `pmhc_modality_key` modality, obs from the `ir_modality_key` modality
+        - AnnData: counts from `X`, obs from `obs`
+        - DataFrame (cells x features): obs is the frame itself, so clonotypes can be a column
+
+        Args:
+            data: a MuData, an AnnData or a cells x features DataFrame.
+            pmhc_modality_key: the MuData modality holding the counts, unused for the other types.
+            ir_modality_key: the MuData modality holding the annotation, unused for the other types.
+
+        Returns:
+            A tuple (counts, obs) of DataFrames sharing the cell order of `data`.
+
+        Raises:
+            TypeError: if `data` is of an unsupported type.
+        """
+        if isinstance(data, md.MuData):
+            counts, _ = ApMHCDeconvolution.as_counts(data.mod[pmhc_modality_key])
+            return counts, data.mod[ir_modality_key].obs if ir_modality_key in data.mod else data.obs
+        if isinstance(data, ad.AnnData):
+            return data.to_df(), data.obs
+        if isinstance(data, pd.DataFrame):
+            return data, data
+        raise TypeError(f"unsupported input type {type(data).__name__}, expected a MuData, an "
+                        f"AnnData or a cells x features DataFrame")
+
+    @staticmethod
+    def _check_parameters(x, x_neg, clone_id):
+        """
+        Checks consistency of the input data before initializing the model.
+
+        Args:
+            x: pMHC UMI counts, shape (n_cells,).
+            x_neg: (Optional) negative control counts, expected shape (n_cells,).
+            clone_id: (Optional) clonotype ids, expected shape (n_cells,).
+
+        Raises:
+            ValueError: if `x` contains NaNs or if `x_neg`/`clone_id` do not match its length.
+        """
+        n_cells = x.shape[0]
+
+        if jnp.isnan(x).any():
+            raise ValueError("Input data `x` contains NaN values. Please remove them before fitting the model.")
+
+        if clone_id is not None:
+            if clone_id.shape[0] != n_cells:
+                raise ValueError(f"`clone_id` and count data `x` require the same size but got "
+                                 f"{clone_id.shape[0]} and {n_cells}")
+
+        if x_neg is not None:
+            n_cells_neg = x_neg.shape[0]
+
+            if n_cells_neg != n_cells:
+                raise ValueError(f"`x_neg` must have the same size as `x` but got "
+                                 f"{n_cells_neg} vs {n_cells}.")
+