scdesigner 0.0.5__py3-none-any.whl → 0.0.10__py3-none-any.whl

scdesigner/minimal/composite.py

@@ -1,119 +0,0 @@
-from .loader import obs_loader
-from .scd3 import SCD3Simulator
-from .standard_copula import StandardCopula
-from anndata import AnnData
-from typing import Dict, Optional, List
-import numpy as np
-import torch
-
-class CompositeCopula(SCD3Simulator):
-    def __init__(self, marginals: List,
-                 copula_formula: Optional[str] = None) -> None:
-        self.marginals = marginals
-        self.copula = StandardCopula(copula_formula)
-        self.template = None
-        self.parameters = None
-        self.merged_formula = None
-
-    def fit(
-        self,
-        adata: AnnData,
-        **kwargs):
-        """Fit the simulator"""
-        self.template = adata
-        merged_formula = {}
-
-        # fit each marginal model
-        for m in range(len(self.marginals)):
-            self.marginals[m][1].setup_data(adata[:, self.marginals[m][0]], **kwargs)
-            self.marginals[m][1].setup_optimizer(**kwargs)
-            self.marginals[m][1].fit(**kwargs)
-
-            # prepare formula for copula loader
-            f = self.marginals[m][1].formula
-            prefixed_f = {f"group{m}_{k}": v for k, v in f.items()}
-            merged_formula = merged_formula | prefixed_f
-
-        # copula simulator
-        self.merged_formula = merged_formula
-        self.copula.setup_data(adata, merged_formula, **kwargs)
-        self.copula.fit(self.merged_uniformize, **kwargs)
-        self.parameters = {
-            "marginal": [m[1].parameters for m in self.marginals],
-            "copula": self.copula.parameters
-        }
-
-    def merged_uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]) -> torch.Tensor:
-        """Produce a merged uniformized matrix for all marginals.
-
-        Delegates to each marginal's `uniformize` method and places the
-        result into the columns of a full matrix according to the variable
-        selection given in `self.marginals[m][0]`.
-        """
-        y_np = y.detach().cpu().numpy()
-        u = np.empty_like(y_np, dtype=float)
-
-        for m in range(len(self.marginals)):
-            sel = self.marginals[m][0]
-            ix = _var_indices(sel, self.template)
-
-            # remove the `group{m}_` prefix we used to distinguish the marginals
-            prefix = f"group{m}_"
-            cur_x = {k.removeprefix(prefix): v if k.startswith(prefix) else v for k, v in x.items()}
-
-            # slice the subset of y for this marginal and call its uniformize
-            y_sub = torch.from_numpy(y_np[:, ix])
-            u[:, ix] = self.marginals[m][1].uniformize(y_sub, cur_x)
-        return torch.from_numpy(u)
-
-    def predict(self, obs=None, batch_size: int = 1000, **kwargs):
-        """Predict from an obs dataframe"""
-        # prepare an internal data loader for this obs
-        if obs is None:
-            obs = self.template.obs
-        loader = obs_loader(
-            obs,
-            self.merged_formula,
-            batch_size=batch_size,
-            **kwargs
-        )
-
-        # prepare per-marginal collectors
-        n_marginals = len(self.marginals)
-        local_pred = [[] for _ in range(n_marginals)]
-
-        # for each batch, call each marginal's predict on its subset of x
-        for _, x_dict in loader:
-            for m in range(n_marginals):
-                prefix = f"group{m}_"
-                # build cur_x where prefixed keys are unprefixed for the marginal
-                cur_x = {k.removeprefix(prefix): v for k, v in x_dict.items()}
-                params = self.marginals[m][1].predict(cur_x)
-                local_pred[m].append(params)
-
-        # merge batch-wise parameter dicts for each marginal and return
-        results = []
-        for m in range(n_marginals):
-            parts = local_pred[m]
-            keys = list(parts[0].keys())
-            results.append({k: torch.cat([d[k] for d in parts]).detach().cpu().numpy() for k in keys})
-
-        return results
-
-
-def _var_indices(sel, adata: AnnData) -> np.ndarray:
-    """Return integer indices of `sel` within `adata.var_names`.
-
-    Expected use: `sel` is a list (or tuple) of variable names (strings).
-    """
-    # If sel is a single string, make it a list so we return consistent shape
-    single_string = False
-    if isinstance(sel, str):
-        sel = [sel]
-        single_string = True
-
-    idx = np.asarray(adata.var_names.get_indexer(sel), dtype=int)
-    if (idx < 0).any():
-        missing = [s for s, i in zip(sel, idx) if i < 0]
-        raise KeyError(f"Variables not found in adata.var_names: {missing}")
-    return idx if not single_string else idx.reshape(-1)

scdesigner/minimal/copula.py

@@ -1,205 +0,0 @@
-from typing import Dict, Callable, Tuple
-import torch
-from anndata import AnnData
-from .loader import adata_loader
-from abc import ABC, abstractmethod
-import numpy as np
-import pandas as pd
-from typing import Optional, Union
-class Copula(ABC):
-    def __init__(self, formula: str, **kwargs):
-        self.formula = formula
-        self.loader = None
-        self.n_outcomes = None
-        self.parameters = None # Should be a dictionary of CovarianceStructure objects
-
-    def setup_data(self, adata: AnnData, marginal_formula: Dict[str, str], batch_size: int = 1024, **kwargs):
-        self.adata = adata
-        self.formula = self.formula | marginal_formula
-        self.loader = adata_loader(adata, self.formula, batch_size=batch_size, **kwargs)
-        X_batch, _ = next(iter(self.loader))
-        self.n_outcomes = X_batch.shape[1]
-    
-    def decorrelate(self, row_pattern: str, col_pattern: str, group: Union[str, list, None] = None):
-        """Decorrelate the covariance matrix for the given row and column patterns.
-        
-        Args:
-            row_pattern (str): The regex pattern for the row names to match.
-            col_pattern (str): The regex pattern for the column names to match.
-            group (Union[str, list, None]): The group or groups to apply the transformation to. If None, the transformation is applied to all groups.
-        """
-        if group is None:
-            for g in self.groups:
-                self.parameters[g].decorrelate(row_pattern, col_pattern)
-        elif isinstance(group, str):
-            self.parameters[group].decorrelate(row_pattern, col_pattern)
-        else:
-            for g in group:
-                self.parameters[g].decorrelate(row_pattern, col_pattern)
-                
-    def correlate(self, factor: float, row_pattern: str, col_pattern: str, group: Union[str, list, None] = None):
-        """Multiply selected off-diagonal entries by factor.
-        
-        Args:
-            row_pattern (str): The regex pattern for the row names to match.
-            col_pattern (str): The regex pattern for the column names to match.
-            factor (float): The factor to multiply the off-diagonal entries by.
-            group (Union[str, list, None]): The group or groups to apply the transformation to. If None, the transformation is applied to all groups.
-        """
-        if group is None:
-            for g in self.groups:
-                self.parameters[g].correlate(row_pattern, col_pattern, factor)
-        elif isinstance(group, str):
-            self.parameters[group].correlate(row_pattern, col_pattern, factor)
-        else:
-            for g in group:
-                self.parameters[g].correlate(row_pattern, col_pattern, factor)
-                
-    @abstractmethod
-    def fit(self, uniformizer: Callable, **kwargs):
-        raise NotImplementedError
-
-    @abstractmethod
-    def pseudo_obs(self, x_dict: Dict):
-        raise NotImplementedError
-
-    @abstractmethod
-    def likelihood(self, uniformizer: Callable, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
-        raise NotImplementedError
-
-    @abstractmethod
-    def num_params(self, **kwargs):
-        raise NotImplementedError
-
-    # @abstractmethod
-    # def format_parameters(self):
-    #     raise NotImplementedError
-    
-class CovarianceStructure:
-    """
-    Efficient storage for covariance matrices in copula-based gene expression modeling.
-    
-    This class provides memory-efficient storage for covariance information by storing
-    either a full covariance matrix or a block matrix with diagonal variances for
-    remaining genes. This enables fast copula estimation and sampling for large
-    gene expression datasets.
-    
-
-    
-    Attributes
-    ----------
-    cov : pd.DataFrame
-        Covariance matrix for modeled genes with gene names as index/columns
-    modeled_indices : np.ndarray
-        Indices of modeled genes in original ordering
-    remaining_var : pd.Series or None
-        Diagonal variances for remaining genes, None if full matrix stored
-    remaining_indices : np.ndarray or None
-        Indices of remaining genes in original ordering
-    num_modeled_genes : int
-        Number of modeled genes
-    num_remaining_genes : int
-        Number of remaining genes (0 if full matrix stored)
-    total_genes : int
-        Total number of genes
-    """
-    
-    def __init__(self, cov: np.ndarray, 
-                 modeled_names: pd.Index, 
-                 modeled_indices: Optional[np.ndarray] = None,
-                 remaining_var: Optional[np.ndarray] = None, 
-                 remaining_indices: Optional[np.ndarray] = None, 
-                 remaining_names: Optional[pd.Index] = None):
-        """initialize a CovarianceStructure object.
-
-        Args:
-            cov (np.ndarray): Covariance matrix for modeled genes, shape (n_modeled_genes, n_modeled_genes)
-            modeled_names (pd.Index): Gene names for the modeled genes
-            modeled_indices (Optional[np.ndarray], optional): Indices of modeled genes in original ordering. Defaults to sequential indices.
-            remaining_var (Optional[np.ndarray], optional): Diagonal variances for remaining genes, shape (n_remaining_genes,)
-            remaining_indices (Optional[np.ndarray], optional): Indices of remaining genes in original ordering
-            remaining_names (Optional[pd.Index], optional): Gene names for remaining genes
-        """
-        self.cov = pd.DataFrame(cov, index=modeled_names, columns=modeled_names)
-        
-        if modeled_indices is not None:
-            self.modeled_indices = modeled_indices
-        else:
-            self.modeled_indices = np.arange(len(modeled_names))
-        
-        if remaining_var is not None:
-            self.remaining_var = pd.Series(remaining_var, index=remaining_names)
-        else: 
-            self.remaining_var = None
-        
-        self.remaining_indices = remaining_indices
-        self.num_modeled_genes = len(modeled_names)
-        self.num_remaining_genes = len(remaining_indices) if remaining_indices is not None else 0
-        self.total_genes = self.num_modeled_genes + self.num_remaining_genes
-        
-    def __repr__(self):
-        if self.remaining_var is None:
-            return self.cov.__repr__()
-        else:
-            return f"CovarianceStructure(modeled_genes={self.num_modeled_genes}, \
-                total_genes={self.total_genes})"
-    
-    def _repr_html_(self):
-        """Jupyter Notebook display"""
-        if self.remaining_var is None:
-            return self.cov._repr_html_()
-        else:
-            html = f"<b>CovarianceStructure:</b> {self.num_modeled_genes} modeled genes, {self.total_genes} total<br>"
-            html += "<h4>Modeled Covariance Matrix</h4>" + self.cov._repr_html_()
-            html += "<h4>Remaining Gene Variances</h4>" + self.remaining_var.to_frame("variance").T._repr_html_()
-            return html
-    
-    def decorrelate(self, row_pattern: str, col_pattern: str):
-        """Decorrelate the covariance matrix for the given row and column patterns.
-        """
-        from .transform import data_frame_mask
-        m1 = data_frame_mask(self.cov, ".", col_pattern)
-        m2 = data_frame_mask(self.cov, row_pattern, ".")
-        mask = (m1 | m2)
-        np.fill_diagonal(mask, False)
-        self.cov.values[mask] = 0
-        
-    def correlate(self, row_pattern: str, col_pattern: str, factor: float):
-        """Multiply selected off-diagonal entries by factor.
-
-        Args:
-            row_pattern (str): The regex pattern for the row names to match.
-            col_pattern (str): The regex pattern for the column names to match.
-            factor (float): The factor to multiply the off-diagonal entries by.
-        """
-        from .transform import data_frame_mask
-        m1 = data_frame_mask(self.cov, ".", col_pattern)
-        m2 = data_frame_mask(self.cov, row_pattern, ".")
-        mask = (m1 | m2)
-        np.fill_diagonal(mask, False)
-        self.cov.values[mask] = self.cov.values[mask] * factor
-    
-    @property
-    def shape(self):
-        return (self.total_genes, self.total_genes)
-    
-    def to_full_matrix(self):
-        """
-        Convert to full covariance matrix for compatibility/debugging.
-        Returns:
-        --------
-        np.ndarray : Full covariance matrix with shape (total_genes, total_genes)
-        """
-        if self.remaining_var is None:
-            return self.cov.values
-        else:
-            full_cov = np.zeros((self.total_genes, self.total_genes))
-            
-            # Fill in top-k block
-            ix_modeled = np.ix_(self.modeled_indices, self.modeled_indices)
-            full_cov[ix_modeled] = self.cov.values
-            
-            # Fill in diagonal for remaining genes
-            full_cov[self.remaining_indices, self.remaining_indices] = self.remaining_var.values
-        
-        return full_cov 

scdesigner/minimal/formula.py

@@ -1,23 +0,0 @@
-from typing import Union
-import warnings
-
-def standardize_formula(formula: Union[str, dict], allowed_keys = None):
-    # The first element of allowed_keys should be the name of default parameter
-    if allowed_keys is None:
-        raise ValueError("Internal error: allowed_keys must be specified")
-    formula = {allowed_keys[0]: formula} if isinstance(formula, str) else formula
-
-    formula_keys = set(formula.keys())
-    allowed_keys = set(allowed_keys)
-
-    if not formula_keys & allowed_keys:
-        raise ValueError(f"formula must have at least one of the following keys: {allowed_keys}")
-    
-    if extra_keys := formula_keys - allowed_keys:
-        warnings.warn(
-            f"Invalid formulas in dictionary will not be used: {extra_keys}",
-            UserWarning,
-        )
-    
-    formula.update({k: '~ 1' for k in allowed_keys - formula_keys})
-    return formula

scdesigner/minimal/gaussian.py

@@ -1,65 +0,0 @@
-from .formula import standardize_formula
-from .marginal import GLMPredictor, Marginal
-from .loader import _to_numpy
-from typing import Union, Dict, Optional
-import torch
-import numpy as np
-from scipy.stats import norm
-
-class Gaussian(Marginal):
-    """Gaussian marginal estimator"""
-    def __init__(self, formula: Union[Dict, str]):
-        formula = standardize_formula(formula, allowed_keys=['mean', 'sdev'])
-        super().__init__(formula)
-
-    def setup_optimizer(
-            self,
-            optimizer_class: Optional[callable] = torch.optim.Adam,
-            **optimizer_kwargs,
-    ):
-        if self.loader is None:
-            raise RuntimeError("self.loader is not set (call setup_data first)")
-
-        nll = lambda batch: -self.likelihood(batch).sum()
-        link_fns = {"mean": lambda x: x}
-        self.predict = GLMPredictor(
-            n_outcomes=self.n_outcomes,
-            feature_dims=self.feature_dims,
-            link_fns=link_fns,
-            loss_fn=nll,
-            optimizer_class=optimizer_class,
-            optimizer_kwargs=optimizer_kwargs
-        )
-
-    def likelihood(self, batch):
-        """Compute the negative log-likelihood"""
-        y, x = batch
-        params = self.predict(x)
-        mu = params.get("mean")
-        sigma = params.get("sdev")
-
-        # log likelihood for Gaussian
-        log_likelihood = -0.5 * (torch.log(2 * torch.pi * sigma ** 2) + ((y - mu) ** 2) / (sigma ** 2))
-        return log_likelihood
-
-    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
-        """Invert pseudoobservations."""
-        mu, sdev, u = self._local_params(x, u)
-        y = norm(loc=mu, scale=sdev).ppf(u)
-        return torch.from_numpy(y).float()
-
-    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor], epsilon=1e-6):
-        """Return uniformized pseudo-observations for counts y given covariates x."""
-        # cdf values using scipy's parameterization
-        mu, sdev, y = self._local_params(x, y)
-        u = norm.cdf(y, loc=mu, scale=sdev)
-        u = np.clip(u, epsilon, 1 - epsilon)
-        return torch.from_numpy(u).float()
-
-    def _local_params(self, x, y=None):
-        params = self.predict(x)
-        mu = params.get('mean')
-        sdev = params.get('sdev')
-        if y is None:
-            return _to_numpy(mu, sdev)
-        return _to_numpy(mu, sdev, y)

scdesigner/minimal/loader.py

@@ -1,211 +0,0 @@
-from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
-from anndata import AnnData
-from formulaic import model_matrix
-from torch.utils.data import Dataset, DataLoader
-from typing import Dict
-import numpy as np
-import pandas as pd
-import scipy.sparse
-import torch
-
-def get_device():
-    """Detect and return the best available device (MPS, CUDA, or CPU)."""
-    if torch.backends.mps.is_available():
-        return torch.device("mps")
-    elif torch.cuda.is_available():
-        return torch.device("cuda")
-    else:
-        return torch.device("cpu")
-
-
-class PreloadedDataset(Dataset):
-    """Dataset that assumes x and y are both fully in memory."""
-    def __init__(self, y_tensor, x_tensors, predictor_names):
-        self.y = y_tensor
-        self.x = x_tensors
-        self.predictor_names = predictor_names
-
-    def __len__(self):
-        return len(self.y)
-
-    def __getitem__(self, idx):
-        return self.y[idx], {k: v[idx] for k, v in self.x.items()}
-
-class AnnDataDataset(Dataset):
-    """Simple PyTorch Dataset for AnnData objects.
-
-    Supports optional chunked loading for backed AnnData objects. When
-    `chunk_size` is provided, the dataset will load contiguous slices
-    of rows (of size `chunk_size`) into memory once and serve individual
-    rows from that cached chunk. Chunks are moved to device for faster access.
-    """
-    def __init__(self, adata: AnnData, formula: Dict[str, str], chunk_size: int):
-        self.adata = adata
-        self.formula = formula
-        self.chunk_size = chunk_size
-        self.device = get_device()
-
-        # keeping track of covariate-related information
-        self.obs_levels = categories(self.adata.obs)
-        self.obs_matrices = {}
-        self.predictor_names = None
-
-        # Internal cache for the currently loaded chunk
-        self._chunk: AnnData | None = None
-        self._chunk_X = None
-        self._chunk_start = 0
-
-    def __len__(self):
-        return len(self.adata)
-
-    def __getitem__(self, idx):
-        """Returns (X, obs) for the given index.
-
-        If `chunk_size` was specified the dataset will load a chunk
-        containing `idx` into memory (if not already cached) and
-        index into that chunk.
-        """
-        self._ensure_chunk_loaded(idx)
-        local_idx = idx - self._chunk_start
-
-        # Get obs data from GPU-cached matrices
-        obs_dict = {}
-        for key in self.formula.keys():
-            obs_dict[key] = self.obs_matrices[key][local_idx: local_idx + 1]
-        return self._chunk_X[local_idx], obs_dict
-
-    def _ensure_chunk_loaded(self, idx: int) -> None:
-        """Load the chunk that contains `idx` into the internal cache."""
-        start = (idx // self.chunk_size) * self.chunk_size
-        end = min(start + self.chunk_size, len(self.adata))
-
-        if (self._chunk is None) or not (self._chunk_start <= idx < self._chunk_start + len(self._chunk)):
-            # load the next chunk into memory
-            chunk = self.adata[start:end]
-            if getattr(chunk, 'isbacked', False):
-                chunk = chunk.to_memory()
-            self._chunk = chunk
-            self._chunk_start = start
-
-            # Move chunk to GPU
-            X = chunk.X
-            if hasattr(X, 'toarray'):
-                X = X.toarray()
-            self._chunk_X = torch.tensor(X, dtype=torch.float32).to(self.device)
-
-            # Compute model matrices for this chunk's `obs` and move to GPU
-            obs_coded_chunk = code_levels(self._chunk.obs.copy(), self.obs_levels)
-            self.obs_matrices = {}
-            predictor_names = {}
-            for key, f in self.formula.items():
-                mat = model_matrix(f, obs_coded_chunk)
-                predictor_names [key] = list(mat.columns)
-                self.obs_matrices[key] = torch.tensor(mat.values, dtype=torch.float32).to(self.device)
-
-            # Capture predictor (column) names from the model matrices once.
-            if self.predictor_names is None:
-                self.predictor_names = predictor_names
-
-
-def adata_loader(
-    adata: AnnData,
-    formula: Dict[str, str],
-    chunk_size: int = None,
-    batch_size: int = 1024,
-    shuffle: bool = False,
-    num_workers: int = 0,
-    **kwargs
-) -> DataLoader:
-    """Create a DataLoader from AnnData that returns batches of (X, obs)."""
-    data_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['data'])
-    device = get_device()
-
-    # separate chunked from non-chunked cases
-    if not getattr(adata, 'isbacked', False):
-        dataset = _preloaded_adata(adata, formula, device)
-    else:
-        dataset = AnnDataDataset(adata, formula, chunk_size or 5000)
-
-    return DataLoader(
-        dataset,
-        batch_size=batch_size,
-        shuffle=shuffle,
-        num_workers=num_workers,
-        collate_fn=dict_collate_fn,
-        **data_kwargs
-    )
-
-def obs_loader(obs: pd.DataFrame, marginal_formula, **kwargs):
-    adata = AnnData(X=np.zeros((len(obs), 1)), obs=obs)
-    return adata_loader(
-        adata,
-        marginal_formula,
-        **kwargs
-    )
-
-################################################################################
-## Extraction of in-memory AnnData to PreloadedDataset
-################################################################################
-
-def _preloaded_adata(adata: AnnData, formula: Dict[str, str], device: torch.device) -> PreloadedDataset:
-    X = adata.X
-    if scipy.sparse.issparse(X):
-        X = X.toarray()
-    y = torch.tensor(X, dtype=torch.float32).to(device)
-
-    obs = code_levels(adata.obs.copy(), categories(adata.obs))
-    x = {
-        k: torch.tensor(model_matrix(f, obs).values, dtype=torch.float32).to(device)
-        for k, f in formula.items()
-    }
-    predictor_names = {k: list(model_matrix(f, obs).columns) for k, f in formula.items()}
-    return PreloadedDataset(y, x, predictor_names)
-
-################################################################################
-## Helper functions
-################################################################################
-
-def dict_collate_fn(batch):
-    """
-    Custom collate function for handling dictionary obs tensors.
-    """
-    X_batch = torch.stack([item[0] for item in batch])
-    obs_batch = [item[1] for item in batch]
-
-    obs_dict = {}
-    for key in obs_batch[0].keys():
-        obs_dict[key] = torch.stack([obs[key] for obs in obs_batch])
-    return X_batch, obs_dict
-
-def to_tensor(X):
-    # If the tensor is 2D with second dim == 1, squeeze only the first
-    # dim when appropriate (e.g. converting a single-row X to 1D samples)
-    t = torch.tensor(X, dtype=torch.float32)
-    if t.dim() == 2 and t.size(1) == 1:
-        if t.size(0) == 1:
-            return t.view(1)
-        return t
-    return t.squeeze()
-
-def categories(obs):
-    levels = {}
-    for k in obs.columns:
-        obs_type = str(obs[k].dtype)
-        if obs_type in ["category", "object"]:
-            levels[k] = obs[k].unique()
-    return levels
-
-
-def code_levels(obs, categories):
-    for k in obs.columns:
-        if str(obs[k].dtype) == "category":
-            obs[k] = obs[k].astype(pd.CategoricalDtype(categories[k]))
-    return obs
-
-###############################################################################
-## Misc. Helper functions
-###############################################################################
-
-def _to_numpy(*tensors):
-    """Convenience helper: detach, move to CPU, and convert tensors to numpy arrays."""
-    return tuple(t.detach().cpu().numpy() for t in tensors)