networkvi 0.1.0__py3-none-any.whl

networkvi/__init__.py

@@ -0,0 +1,32 @@
+"""scvi-tools."""
+
+# Set default logging handler to avoid logging with logging.lastResort logger.
+import logging
+import warnings
+
+from ._constants import REGISTRY_KEYS
+from ._settings import settings
+
+# this import needs to come after prior imports to prevent circular import
+from . import data, model, utils
+
+from importlib.metadata import version
+
+package_name = "networkvi"
+__version__ = "1.0.0"
+
+settings.verbosity = logging.INFO
+
+# Jax sets the root logger, this prevents double output.
+scvi_logger = logging.getLogger("networkvi")
+scvi_logger.propagate = False
+
+
+__all__ = [
+    "settings",
+    "REGISTRY_KEYS",
+    "data",
+    "model",
+    "utils",
+]
+

networkvi/_constants.py

@@ -0,0 +1,20 @@
+from typing import NamedTuple
+
+
+class _REGISTRY_KEYS_NT(NamedTuple):
+    X_KEY: str = "X"
+    BATCH_KEY: str = "batch"
+    PATIENT_KEY: str = "patient"
+    LABELS_KEY: str = "labels"
+    PROTEIN_EXP_KEY: str = "proteins"
+    CAT_COVS_KEY: str = "extra_categorical_covs"
+    CONT_COVS_KEY: str = "extra_continuous_covs"
+    INDICES_KEY: str = "ind_x"
+    SIZE_FACTOR_KEY: str = "size_factor"
+    MINIFY_TYPE_KEY: str = "minify_type"
+    LATENT_QZM_KEY: str = "latent_qzm"
+    LATENT_QZV_KEY: str = "latent_qzv"
+    OBSERVED_LIB_SIZE: str = "observed_lib_size"
+
+
+REGISTRY_KEYS = _REGISTRY_KEYS_NT()

networkvi/_settings.py

@@ -0,0 +1,220 @@
+import logging
+import os
+from pathlib import Path
+from typing import Literal, Optional, Union
+
+import torch
+from lightning.pytorch import seed_everything
+from rich.console import Console
+from rich.logging import RichHandler
+
+scvi_logger = logging.getLogger("scvi")
+
+
+class ScviConfig:
+    """Config manager for scvi-tools.
+
+    Examples
+    --------
+    To set the seed
+
+    >>> networkvi.settings.seed = 1
+
+    To set the batch size for functions like `networkvi.get_latent_representation`
+
+    >>> networkvi.settings.batch_size = 1024
+
+    To set the progress bar style, choose one of "rich", "tqdm"
+
+    >>> networkvi.settings.progress_bar_style = "rich"
+
+    To set the verbosity
+
+    >>> import logging
+    >>> networkvi.settings.verbosity = logging.INFO
+
+    To set the number of threads PyTorch will use
+
+    >>> networkvi.settings.num_threads = 2
+
+    To prevent Jax from preallocating GPU memory on start (default)
+
+    >>> networkvi.settings.jax_preallocate_gpu_memory = False
+    """
+
+    def __init__(
+        self,
+        verbosity: int = logging.INFO,
+        progress_bar_style: Literal["rich", "tqdm"] = "tqdm",
+        batch_size: int = 128,
+        seed: Optional[int] = None,
+        logging_dir: str = "./scvi_log/",
+        dl_num_workers: int = 0,
+        jax_preallocate_gpu_memory: bool = False,
+        warnings_stacklevel: int = 2,
+    ):
+        self.warnings_stacklevel = warnings_stacklevel
+        self.seed = seed
+        self.batch_size = batch_size
+        if progress_bar_style not in ["rich", "tqdm"]:
+            raise ValueError("Progress bar style must be in ['rich', 'tqdm']")
+        self.progress_bar_style = progress_bar_style
+        self.logging_dir = logging_dir
+        self.dl_num_workers = dl_num_workers
+        self._num_threads = None
+        self.jax_preallocate_gpu_memory = jax_preallocate_gpu_memory
+        self.verbosity = verbosity
+
+    @property
+    def batch_size(self) -> int:
+        """Minibatch size for loading data into the model.
+
+        This is only used after a model is trained. Trainers have specific
+        `batch_size` parameters.
+        """
+        return self._batch_size
+
+    @batch_size.setter
+    def batch_size(self, batch_size: int):
+        """Minibatch size for loading data into the model.
+
+        This is only used after a model is trained. Trainers have specific
+        `batch_size` parameters.
+        """
+        self._batch_size = batch_size
+
+    @property
+    def dl_num_workers(self) -> int:
+        """Number of workers for PyTorch data loaders (Default is 0)."""
+        return self._dl_num_workers
+
+    @dl_num_workers.setter
+    def dl_num_workers(self, dl_num_workers: int):
+        """Number of workers for PyTorch data loaders (Default is 0)."""
+        self._dl_num_workers = dl_num_workers
+
+    @property
+    def logging_dir(self) -> Path:
+        """Directory for training logs (default `'./scvi_log/'`)."""
+        return self._logging_dir
+
+    @logging_dir.setter
+    def logging_dir(self, logging_dir: Union[str, Path]):
+        self._logging_dir = Path(logging_dir).resolve()
+
+    @property
+    def num_threads(self) -> None:
+        """Number of threads PyTorch will use."""
+        return self._num_threads
+
+    @num_threads.setter
+    def num_threads(self, num: int):
+        """Number of threads PyTorch will use."""
+        self._num_threads = num
+        torch.set_num_threads(num)
+
+    @property
+    def progress_bar_style(self) -> str:
+        """Library to use for progress bar."""
+        return self._pbar_style
+
+    @progress_bar_style.setter
+    def progress_bar_style(self, pbar_style: Literal["tqdm", "rich"]):
+        """Library to use for progress bar."""
+        self._pbar_style = pbar_style
+
+    @property
+    def seed(self) -> int:
+        """Random seed for torch and numpy."""
+        return self._seed
+
+    @seed.setter
+    def seed(self, seed: Union[int, None] = None):
+        """Random seed for torch and numpy."""
+        if seed is None:
+            self._seed = None
+        else:
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+            seed_everything(seed)
+            self._seed = seed
+
+    @property
+    def verbosity(self) -> int:
+        """Verbosity level (default `logging.INFO`)."""
+        return self._verbosity
+
+    @verbosity.setter
+    def verbosity(self, level: Union[str, int]):
+        """Sets logging configuration for scvi based on chosen level of verbosity.
+
+        If "scvi" logger has no StreamHandler, add one.
+        Else, set its level to `level`.
+
+        Parameters
+        ----------
+        level
+            Sets "scvi" logging level to `level`
+        force_terminal
+            Rich logging option, set to False if piping to file output.
+        """
+        self._verbosity = level
+        scvi_logger.setLevel(level)
+        if len(scvi_logger.handlers) == 0:
+            console = Console(force_terminal=True)
+            if console.is_jupyter is True:
+                console.is_jupyter = False
+            ch = RichHandler(level=level, show_path=False, console=console, show_time=False)
+            formatter = logging.Formatter("%(message)s")
+            ch.setFormatter(formatter)
+            scvi_logger.addHandler(ch)
+        else:
+            scvi_logger.setLevel(level)
+
+    @property
+    def warnings_stacklevel(self) -> int:
+        """Stacklevel for warnings."""
+        return self._warnings_stacklevel
+
+    @warnings_stacklevel.setter
+    def warnings_stacklevel(self, stacklevel: int):
+        """Stacklevel for warnings."""
+        self._warnings_stacklevel = stacklevel
+
+    def reset_logging_handler(self):
+        """Resets "scvi" log handler to a basic RichHandler().
+
+        This is useful if piping outputs to a file.
+        """
+        scvi_logger.removeHandler(scvi_logger.handlers[0])
+        ch = RichHandler(level=self._verbosity, show_path=False, show_time=False)
+        formatter = logging.Formatter("%(message)s")
+        ch.setFormatter(formatter)
+        scvi_logger.addHandler(ch)
+
+    @property
+    def jax_preallocate_gpu_memory(self):
+        """Jax GPU memory allocation settings.
+
+        If False, Jax will ony preallocate GPU memory it needs.
+        If float in (0, 1), Jax will preallocate GPU memory to that
+        fraction of the GPU memory.
+        """
+        return self._jax_gpu
+
+    @jax_preallocate_gpu_memory.setter
+    def jax_preallocate_gpu_memory(self, value: Union[float, bool]):
+        # see https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html#gpu-memory-allocation
+        if value is False:
+            os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+        elif isinstance(value, float):
+            if value >= 1 or value <= 0:
+                raise ValueError("Need to use a value between 0 and 1")
+            # format is ".XX"
+            os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = str(value)[1:4]
+        else:
+            raise ValueError("value not understood, need bool or float in (0, 1)")
+        self._jax_gpu = value
+
+
+settings = ScviConfig()

networkvi/_types.py

@@ -0,0 +1,13 @@
+from typing import Literal, Union
+
+import anndata
+import jax.numpy as jnp
+import mudata
+import torch
+
+Number = Union[int, float]
+AnnOrMuData = Union[anndata.AnnData, mudata.MuData]
+Tensor = Union[torch.Tensor, jnp.ndarray]
+LossRecord = Union[dict[str, Tensor], Tensor]
+# TODO(adamgayoso): Add constants for minified data types.
+MinifiedDataType = Literal["latent_posterior_parameters"]

networkvi/data/__init__.py

@@ -0,0 +1,31 @@
+from anndata import read_csv, read_h5ad, read_loom, read_text
+
+from ._anntorchdataset import AnnTorchDataset
+from ._manager import AnnDataManager, AnnDataManagerValidationCheck
+from ._preprocessing import (
+    add_dna_sequence,
+    organize_cite_seq_10x,
+    organize_multiome_anndatas,
+    poisson_gene_selection,
+    reads_to_fragments,
+)
+from ._read import read_10x_atac, read_10x_multiome
+
+__all__ = [
+    "AnnTorchDataset",
+    "AnnDataManagerValidationCheck",
+    "AnnDataManager",
+    "poisson_gene_selection",
+    "organize_cite_seq_10x",
+    "read_h5ad",
+    "read_csv",
+    "read_loom",
+    "read_text",
+    "read_10x_atac",
+    "read_10x_multiome",
+    "organize_multiome_anndatas",
+    "add_dna_sequence",
+    "reads_to_fragments",
+]
+
+

networkvi/data/_anntorchdataset.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import h5py
+import numpy as np
+import pandas as pd
+import torch
+
+try:
+    # anndata >= 0.10
+    from anndata.experimental import CSCDataset, CSRDataset
+
+    SparseDataset = (CSRDataset, CSCDataset)
+except ImportError:
+    from anndata._core.sparse_dataset import SparseDataset
+
+from scipy.sparse import issparse
+from torch.utils.data import Dataset
+
+from networkvi._constants import REGISTRY_KEYS
+
+if TYPE_CHECKING:
+    from ._manager import AnnDataManager
+from ._utils import registry_key_to_default_dtype, scipy_to_torch_sparse
+
+logger = logging.getLogger(__name__)
+
+class AnnTorchDataset(Dataset):
+    """Extension of :class:`~torch.utils.data.Dataset` for :class:`~anndata.AnnData` objects.
+
+    Parameters
+    ----------
+    adata_manager
+        :class:`~networkvi.data.AnnDataManager` object with a registered AnnData object.
+    getitem_tensors
+        Specifies the keys in the data registry (``adata_manager.data_registry``) to return in
+        ``__getitem__``. One of the following:
+
+        * ``dict``: Keys correspond to keys in the data registry and values correspond to the
+        desired :class:`~np.dtype` of the returned data.
+        * ``list``: Elements correspond to keys in the data registry. Continuous data will be
+        returned as :class:`~np.float32` and discrete data will be returned as :class:`~np.int64`.
+        * ``None``: All registered data will be returned. Continuous data will be returned as
+        :class:`~np.float32` and discrete data will be returned as :class:`~np.int64`.
+    load_sparse_tensor
+        ``EXPERIMENTAL`` If ``True``, loads data with sparse CSR or CSC layout as a
+        :class:`~torch.Tensor` with the same layout. Can lead to speedups in data transfers to
+        GPUs, depending on the sparsity of the data.
+    """
+
+    def __init__(
+        self,
+        adata_manager: AnnDataManager,
+        getitem_tensors: list | dict[str, type] | None = None,
+        load_sparse_tensor: bool = False,
+    ):
+        super().__init__()
+
+        if adata_manager.adata is None:
+            raise ValueError("Please run ``register_fields`` on ``adata_manager`` first.")
+        self.adata_manager = adata_manager
+        self.keys_and_dtypes = getitem_tensors
+        self.load_sparse_tensor = load_sparse_tensor
+
+    @property
+    def registered_keys(self):
+        """Keys in the data registry."""
+        return self.adata_manager.data_registry.keys()
+
+    @property
+    def keys_and_dtypes(self):
+        """Keys and corresponding :class:`~np.dtype` of data to fetch in ``__getitem__``."""
+        return self._keys_and_dtypes
+
+    @keys_and_dtypes.setter
+    def keys_and_dtypes(self, getitem_tensors: list | dict[str, type] | None):
+        """Set keys and corresponding :class:`~np.dtype` of data to fetch in ``__getitem__``.
+
+        Raises an error if any of the keys are not in the data registry.
+        """
+        if isinstance(getitem_tensors, list):
+            keys_to_dtypes = {key: registry_key_to_default_dtype(key) for key in getitem_tensors}
+        elif isinstance(getitem_tensors, dict):
+            keys_to_dtypes = getitem_tensors
+        elif getitem_tensors is None:
+            keys_to_dtypes = {
+                key: registry_key_to_default_dtype(key) for key in self.registered_keys
+            }
+        else:
+            raise ValueError("`getitem_tensors` must be a `list`, `dict`, or `None`")
+
+        for key in keys_to_dtypes:
+            if key not in self.registered_keys:
+                raise KeyError(f"{key} not found in the data registry.")
+
+        self._keys_and_dtypes = keys_to_dtypes
+
+    @property
+    def data(self):
+        """Dictionary of data tensors.
+
+        First time this is accessed, data is fetched from the underlying
+        :class:`~anndata.AnnData` object. Subsequent accesses will return the
+        cached dictionary.
+        """
+        if not hasattr(self, "_data"):
+            self._data = {
+                key: self.adata_manager.get_from_registry(key) for key in self.keys_and_dtypes
+            }
+        return self._data
+
+    def __len__(self):
+        return self.adata_manager.adata.shape[0]
+
+    def __getitem__(
+        self, indexes: int | list[int] | slice
+    ) -> dict[str, np.ndarray | torch.Tensor]:
+        """Fetch data from the :class:`~anndata.AnnData` object.
+
+        Parameters
+        ----------
+        indexes
+            Indexes of the observations to fetch. Can be a single index, a list of indexes, or a
+            slice.
+
+        Returns
+        -------
+        Mapping of data registry keys to arrays of shape ``(n_obs, ...)``.
+        """
+        if isinstance(indexes, int):
+            indexes = [indexes]  # force batched single observations
+
+        if self.adata_manager.adata.isbacked and isinstance(indexes, (list, np.ndarray)):
+            # need to sort indexes for h5py datasets
+            indexes = np.sort(indexes)
+
+        data_map = {}
+
+        for key, dtype in self.keys_and_dtypes.items():
+            data = self.data[key]
+
+            if isinstance(data, (np.ndarray, h5py.Dataset)):
+                sliced_data = data[indexes].astype(dtype, copy=False)
+            elif isinstance(data, pd.DataFrame):
+                sliced_data = data.iloc[indexes, :].to_numpy().astype(dtype, copy=False)
+            elif issparse(data) or isinstance(data, SparseDataset):
+                sliced_data = data[indexes].astype(dtype, copy=False)
+                if self.load_sparse_tensor:
+                    sliced_data = scipy_to_torch_sparse(sliced_data)
+                else:
+                    sliced_data = sliced_data.toarray()
+            elif issparse(data) or isinstance(data, da.Array):
+                sliced_data = data[indexes].compute()
+            elif isinstance(data, str) and key == REGISTRY_KEYS.MINIFY_TYPE_KEY:
+                # for minified  anndata, we need this because we can have a string
+                # for `data``, which is the value of the MINIFY_TYPE_KEY in adata.uns,
+                # used to record the type data minification
+                # TODO: Adata manager should have a list of which fields it will load
+                continue
+            else:
+                raise TypeError(f"{key} is not a supported type")
+
+            data_map[key] = sliced_data
+
+        return data_map

networkvi/data/_compat.py

@@ -0,0 +1,174 @@
+from copy import deepcopy
+from typing import Optional
+
+import numpy as np
+
+from networkvi import REGISTRY_KEYS
+
+from . import _constants
+from .fields import (
+    CategoricalJointObsField,
+    CategoricalObsField,
+    LabelsWithUnlabeledObsField,
+    LayerField,
+    NumericalJointObsField,
+    ProteinObsmField,
+)
+
+LEGACY_REGISTRY_KEY_MAP = {
+    "X": REGISTRY_KEYS.X_KEY,
+    "batch_indices": REGISTRY_KEYS.BATCH_KEY,
+    "labels": REGISTRY_KEYS.LABELS_KEY,
+    "cat_covs": REGISTRY_KEYS.CAT_COVS_KEY,
+    "cont_covs": REGISTRY_KEYS.CONT_COVS_KEY,
+    "protein_expression": REGISTRY_KEYS.PROTEIN_EXP_KEY,
+    "ind_x": REGISTRY_KEYS.INDICES_KEY,
+}
+
+
+def _infer_setup_args(model_cls, setup_dict: dict, unlabeled_category: Optional[str]) -> dict:
+    setup_args = {}
+    data_registry = setup_dict[_constants._DATA_REGISTRY_KEY]
+    categorical_mappings = setup_dict["categorical_mappings"]
+    for registry_key, adata_mapping in data_registry.items():
+        if registry_key not in LEGACY_REGISTRY_KEY_MAP:
+            continue
+        new_registry_key = LEGACY_REGISTRY_KEY_MAP[registry_key]
+
+        attr_name = adata_mapping[_constants._DR_ATTR_NAME]
+        attr_key = adata_mapping[_constants._DR_ATTR_KEY]
+        if attr_name == _constants._ADATA_ATTRS.X:
+            setup_args["layer"] = None
+        elif attr_name == _constants._ADATA_ATTRS.LAYERS:
+            setup_args["layer"] = attr_key
+        elif attr_name == _constants._ADATA_ATTRS.OBS:
+            if new_registry_key in {REGISTRY_KEYS.BATCH_KEY, REGISTRY_KEYS.LABELS_KEY}:
+                # Omit labels field for TOTALVI
+                if (
+                    model_cls.__name__ == "TOTALVI"
+                    and new_registry_key == REGISTRY_KEYS.LABELS_KEY
+                ):
+                    continue
+                setup_args[f"{new_registry_key}_key"] = categorical_mappings[attr_key][
+                    "original_key"
+                ]
+        elif attr_name == _constants._ADATA_ATTRS.OBSM:
+            if new_registry_key == REGISTRY_KEYS.CONT_COVS_KEY:
+                obs_keys = setup_dict["extra_continuous_keys"]
+                setup_args["continuous_covariate_keys"] = obs_keys
+            elif new_registry_key == REGISTRY_KEYS.CAT_COVS_KEY:
+                obs_keys = setup_dict["extra_categoricals"]["keys"]
+                setup_args["categorical_covariate_keys"] = obs_keys
+            elif new_registry_key == REGISTRY_KEYS.PROTEIN_EXP_KEY:
+                setup_args["protein_expression_obsm_key"] = attr_key
+                setup_args["protein_names_uns_key"] = "_protein_names"
+
+    if unlabeled_category is not None:
+        setup_args["unlabeled_category"] = unlabeled_category
+
+    return {
+        _constants._MODEL_NAME_KEY: model_cls.__name__,
+        _constants._SETUP_ARGS_KEY: setup_args,
+    }
+
+
+def registry_from_setup_dict(
+    model_cls, setup_dict: dict, unlabeled_category: Optional[str] = None
+) -> dict:
+    """Converts old setup dict format to new registry dict format.
+
+    Only to be used for backwards compatibility when loading setup dictionaries for models.
+    Takes old hard-coded setup dictionary structure and fills in the analogous registry structure.
+
+    Parameters
+    ----------
+    model_cls
+        Model class used for setup.
+    setup_dict
+        Setup dictionary created after registering an AnnData with former ``setup_anndata``
+        implementation.
+    unlabeled_category
+        Unlabeled category value used in :class:`~networkvi.model.SCANVI`.
+    """
+    registry = {
+        _constants._SCVI_VERSION_KEY: setup_dict[_constants._SCVI_VERSION_KEY],
+        _constants._FIELD_REGISTRIES_KEY: {},
+    }
+    data_registry = setup_dict[_constants._DATA_REGISTRY_KEY]
+    categorical_mappings = setup_dict["categorical_mappings"]
+    summary_stats = setup_dict[_constants._SUMMARY_STATS_KEY]
+    field_registries = registry[_constants._FIELD_REGISTRIES_KEY]
+    for (
+        registry_key,
+        adata_mapping,
+    ) in data_registry.items():  # Note: this does not work for empty fields.
+        if registry_key not in LEGACY_REGISTRY_KEY_MAP:
+            continue
+        new_registry_key = LEGACY_REGISTRY_KEY_MAP[registry_key]
+
+        attr_name = adata_mapping[_constants._DR_ATTR_NAME]
+        attr_key = adata_mapping[_constants._DR_ATTR_KEY]
+        if attr_key is not None and attr_key.startswith("_scvi"):
+            adata_mapping[_constants._DR_ATTR_KEY] = f"_scvi_{new_registry_key}"
+
+        field_registries[new_registry_key] = {
+            _constants._DATA_REGISTRY_KEY: adata_mapping,
+            _constants._STATE_REGISTRY_KEY: {},
+            _constants._SUMMARY_STATS_KEY: {},
+        }
+        field_registry = field_registries[new_registry_key]
+        field_state_registry = field_registry[_constants._STATE_REGISTRY_KEY]
+        field_summary_stats = field_registry[_constants._SUMMARY_STATS_KEY]
+
+        if attr_name in (_constants._ADATA_ATTRS.X, _constants._ADATA_ATTRS.LAYERS):
+            field_state_registry[LayerField.N_VARS_KEY] = summary_stats["n_vars"]
+            field_summary_stats.update(field_state_registry)
+        elif attr_name == _constants._ADATA_ATTRS.OBS:
+            categorical_mapping = categorical_mappings[attr_key]
+            # Default labels field for TOTALVI
+            if model_cls.__name__ == "TOTALVI" and new_registry_key == REGISTRY_KEYS.LABELS_KEY:
+                field_state_registry[CategoricalObsField.CATEGORICAL_MAPPING_KEY] = np.zeros(
+                    1, dtype=np.int64
+                )
+            else:
+                field_state_registry[CategoricalObsField.CATEGORICAL_MAPPING_KEY] = (
+                    categorical_mapping["mapping"]
+                )
+            field_state_registry[CategoricalObsField.ORIGINAL_ATTR_KEY] = categorical_mapping[
+                "original_key"
+            ]
+            if new_registry_key == REGISTRY_KEYS.BATCH_KEY:
+                field_summary_stats[f"n_{new_registry_key}"] = summary_stats["n_batch"]
+            elif new_registry_key == REGISTRY_KEYS.LABELS_KEY:
+                field_summary_stats[f"n_{new_registry_key}"] = summary_stats["n_labels"]
+                if unlabeled_category is not None:
+                    field_state_registry[LabelsWithUnlabeledObsField.UNLABELED_CATEGORY] = (
+                        unlabeled_category
+                    )
+        elif attr_name == _constants._ADATA_ATTRS.OBSM:
+            if new_registry_key == REGISTRY_KEYS.CONT_COVS_KEY:
+                columns = setup_dict["extra_continuous_keys"].copy()
+                field_state_registry[NumericalJointObsField.COLUMNS_KEY] = columns
+                field_summary_stats[f"n_{new_registry_key}"] = columns.shape[0]
+            elif new_registry_key == REGISTRY_KEYS.CAT_COVS_KEY:
+                extra_categoricals_mapping = deepcopy(setup_dict["extra_categoricals"])
+                field_state_registry.update(extra_categoricals_mapping)
+                field_state_registry[CategoricalJointObsField.FIELD_KEYS_KEY] = (
+                    field_state_registry.pop("keys")
+                )
+                field_summary_stats[f"n_{new_registry_key}"] = len(
+                    extra_categoricals_mapping["keys"]
+                )
+            elif new_registry_key == REGISTRY_KEYS.PROTEIN_EXP_KEY:
+                field_state_registry[ProteinObsmField.COLUMN_NAMES_KEY] = setup_dict[
+                    "protein_names"
+                ].copy()
+                if "totalvi_batch_mask" in setup_dict:
+                    field_state_registry[ProteinObsmField.PROTEIN_BATCH_MASK] = setup_dict[
+                        "totalvi_batch_mask"
+                    ].copy()
+                field_summary_stats[f"n_{new_registry_key}"] = len(setup_dict["protein_names"])
+
+    registry.update(_infer_setup_args(model_cls, setup_dict, unlabeled_category))
+
+    return registry