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

scdesigner/minimal/loader.py

@@ -5,21 +5,45 @@ 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. This avoids calling `to_memory()` on
-    a per-row basis which is expensive for large backed files.
+    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)
@@ -28,6 +52,7 @@ class AnnDataDataset(Dataset):
 
         # Internal cache for the currently loaded chunk
         self._chunk: AnnData | None = None
+        self._chunk_X = None
         self._chunk_start = 0
 
     def __len__(self):
@@ -42,19 +67,12 @@ class AnnDataDataset(Dataset):
         """
         self._ensure_chunk_loaded(idx)
         local_idx = idx - self._chunk_start
-        adata_slice = self._chunk[local_idx]
-
-        # Get X data, accounting for potential sparse matrices
-        X = adata_slice.X
-        if hasattr(X, 'toarray'):
-            X = X.toarray()
 
-        # Get obs data
+        # Get obs data from GPU-cached matrices
         obs_dict = {}
         for key in self.formula.keys():
-            mat = self.obs_matrices.get(key)
-            obs_dict[key] = to_tensor(mat.values[local_idx: local_idx + 1])
-        return to_tensor(X), obs_dict
+            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."""
@@ -69,36 +87,45 @@ class AnnDataDataset(Dataset):
             self._chunk = chunk
             self._chunk_start = start
 
-            # Compute model matrices for this chunk's `obs` so we don't need
-            # to keep the full obs data model matrices in memory.
+            # 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():
-                self.obs_matrices[key] = model_matrix(f, obs_coded_chunk)
+                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 = {k: list(v.columns) for k, v in self.obs_matrices.items()}
-
-
-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).
-    """
+                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'])
-    if chunk_size is None:
-        if getattr(adata, 'isbacked', False):
-            chunk_size = 5000
-        else:
-            chunk_size = len(adata)
+    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)
 
-    dataset = AnnDataDataset(adata, formula, chunk_size)
     return DataLoader(
         dataset,
         batch_size=batch_size,
@@ -109,12 +136,30 @@ def adata_loader(adata: AnnData,
     )
 
 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
-        )
+    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

scdesigner/minimal/marginal.py

@@ -1,14 +1,14 @@
 from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
-from .loader import adata_loader
+from .loader import adata_loader, get_device
 from anndata import AnnData
 from typing import Union, Dict, Optional, Tuple
 import pandas as pd
-import pytorch_lightning as pl
 import torch
 import torch.nn as nn
+from abc import ABC, abstractmethod
 
 
-class Marginal:
+class Marginal(ABC):
     def __init__(self, formula: Union[Dict, str]):
         self.formula = formula
         self.feature_dims = None
@@ -17,6 +17,7 @@ class Marginal:
         self.predict = None
         self.predictor_names = None
         self.parameters = None
+        self.device = get_device()
 
     def setup_data(self, adata: AnnData, batch_size: int = 1024, **kwargs):
         """Set up the dataloader for the AnnData object."""
@@ -28,31 +29,31 @@ class Marginal:
         self.feature_dims = {k: v.shape[1] for k, v in obs_batch.items()}
         self.predictor_names = self.loader.dataset.predictor_names
 
-    def fit(self, **kwargs):
-        """Fit the marginal predictor"""
+    def fit(self, max_epochs: int = 100, **kwargs):
+        """Fit the marginal predictor using vanilla PyTorch training loop."""
         if self.predict is None:
             self.setup_optimizer(**kwargs)
-        trainer_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['trainer'])
-        trainer = pl.Trainer(**trainer_kwargs)
-        trainer.fit(self.predict, train_dataloaders=self.loader)
-        self.parameters = self.format_parameters()
 
-    def setup_optimizer(self, **kwargs):
-        raise NotImplementedError
+        for epoch in range(max_epochs):
+            epoch_loss, n_batches = 0.0, 0
 
-    def likelihood(self, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
-        """Compute the (negative) log-likelihood or loss for a batch.
-        """
-        raise NotImplementedError
+            for batch in self.loader:
+                y, x = batch
+                if y.device != self.device:
+                    y = y.to(self.device)
+                    x = {k: v.to(self.device) for k, v in x.items()}
 
-    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
-        """Invert pseudoobservations."""
-        raise NotImplementedError
+                self.predict.optimizer.zero_grad()
+                loss = self.predict.loss_fn((y, x))
+                loss.backward()
+                self.predict.optimizer.step()
 
-    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]):
-        """Uniformize using learned CDF.
-        """
-        raise NotImplementedError
+                epoch_loss += loss.item()
+                n_batches += 1
+
+            avg_loss = epoch_loss / n_batches
+            print(f"Epoch {epoch}/{max_epochs}, Loss: {avg_loss:.4f}", end='\r')
+        self.parameters = self.format_parameters()
 
     def format_parameters(self):
         """Convert fitted coefficient tensors into pandas DataFrames.
@@ -79,8 +80,29 @@ class Marginal:
             return 0
         return sum(p.numel() for p in self.predict.parameters() if p.requires_grad)
 
+    @abstractmethod
+    def setup_optimizer(self, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def likelihood(self, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
+        """Compute the (negative) log-likelihood or loss for a batch.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Invert pseudoobservations."""
+        raise NotImplementedError
 
-class GLMPredictor(pl.LightningModule):
+    @abstractmethod
+    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Uniformize using learned CDF.
+        """
+        raise NotImplementedError
+
+
+class GLMPredictor(nn.Module):
     """GLM-style predictor with arbitrary named parameters.
 
     Args:
@@ -106,21 +128,22 @@ class GLMPredictor(pl.LightningModule):
         self.feature_dims = dict(feature_dims)
         self.param_names = list(self.feature_dims.keys())
 
-        # create default link functions and parameter matrices
         self.link_fns = link_fns or {k: torch.exp for k in self.param_names}
         self.coefs = nn.ParameterDict()
         for key, dim in self.feature_dims.items():
             self.coefs[key] = nn.Parameter(torch.zeros(dim, self.n_outcomes))
-
-        # optimization parameters
         self.reset_parameters()
+
         self.loss_fn = loss_fn
-        self.optimizer_class = optimizer_class
-        self.optimizer_kwargs = optimizer_kwargs
+        self.to(get_device())
+
+        optimizer_kwargs = optimizer_kwargs or {}
+        filtered_kwargs = _filter_kwargs(optimizer_kwargs, DEFAULT_ALLOWED_KWARGS['optimizer'])
+        self.optimizer = optimizer_class(self.parameters(), **filtered_kwargs)
 
     def reset_parameters(self):
         for p in self.coefs.values():
-            nn.init.normal_(p, mean=0.0, std=1e-2)
+            nn.init.normal_(p, mean=0.0, std=1e-4)
 
     def forward(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         out = {}
@@ -128,13 +151,4 @@ class GLMPredictor(pl.LightningModule):
             x_beta = obs_dict[name] @ self.coefs[name]
             link = self.link_fns.get(name, torch.exp)
             out[name] = link(x_beta)
-        return out
-
-    def training_step(self, batch):
-        loss = self.loss_fn(batch)
-        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
-        return loss
-
-    def configure_optimizers(self, **kwargs):
-        optimizer_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['optimizer'])
-        return self.optimizer_class(self.parameters(), **optimizer_kwargs)
+        return out

scdesigner/minimal/negbin.py

@@ -30,7 +30,7 @@ class NegBin(Marginal):
         )
 
     def likelihood(self, batch):
-        """Compute the negative log-likelihood"""
+        """Compute the log-likelihood"""
         y, x = batch
         params = self.predict(x)
         mu = params.get('mean')

scdesigner/minimal/scd3.py

@@ -6,6 +6,7 @@ from anndata import AnnData
 from tqdm import tqdm
 import torch
 import numpy as np
+from abc import ABC, abstractmethod
 
 class SCD3Simulator(Simulator):
     """Simulation wrapper"""

scdesigner/minimal/scd3_instances.py

@@ -2,7 +2,7 @@ from .scd3 import SCD3Simulator
 from .negbin import NegBin
 from .zero_inflated_negbin import ZeroInflatedNegBin
 from .gaussian import Gaussian
-from .standard_covariance import StandardCovariance
+from .standard_copula import StandardCopula
 from typing import Optional
 
 
@@ -12,7 +12,7 @@ class NegBinCopula(SCD3Simulator):
                  dispersion_formula: Optional[str] = None,
                  copula_formula: Optional[str] = None) -> None:
         marginal = NegBin({"mean": mean_formula, "dispersion": dispersion_formula})
-        covariance = StandardCovariance(copula_formula)
+        covariance = StandardCopula(copula_formula)
         super().__init__(marginal, covariance)
 
 
@@ -27,7 +27,7 @@ class ZeroInflatedNegBinCopula(SCD3Simulator):
             "dispersion": dispersion_formula,
             "zero_inflation_formula": zero_inflation_formula
         })
-        covariance = StandardCovariance(copula_formula)
+        covariance = StandardCopula(copula_formula)
         super().__init__(marginal, covariance)
 
 
@@ -36,7 +36,7 @@ class BernoulliCopula(SCD3Simulator):
                  mean_formula: Optional[str] = None,
                  copula_formula: Optional[str] = None) -> None:
         marginal = NegBin({"mean": mean_formula})
-        covariance = StandardCovariance(copula_formula)
+        covariance = StandardCopula(copula_formula)
         super().__init__(marginal, covariance)
 
 
@@ -46,5 +46,5 @@ class GaussianCopula(SCD3Simulator):
                  sdev_formula: Optional[str] = None,
                  copula_formula: Optional[str] = None) -> None:
         marginal = Gaussian({"mean": mean_formula, "sdev": sdev_formula})
-        covariance = StandardCovariance(copula_formula)
+        covariance = StandardCopula(copula_formula)
         super().__init__(marginal, covariance)