cellarr 0.2.2__tar.gz → 0.2.3__tar.gz

{cellarr-0.2.2 → cellarr-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cellarr
-Version: 0.2.2
+Version: 0.2.3
 Summary: TileDB-based array storage for genomics data collections.
 Home-page: https://github.com/BiocPy/cellarr
 Author: Jayaram Kancherla
@@ -165,6 +165,32 @@ datamodule = DataModule(
 )
 ```
 
+The package also includes a simple autoencoder in the `pytorch-lightning` which makes use of the dataloader. This can be used as a template to create custom architectures and models.
+
+```python
+import pytorch_lightning as pl
+from cellarr.autoencoder import AutoEncoder
+
+autoencoder = AutoEncoder(
+    n_genes=len(datamodule.gene_indices),
+    latent_dim=128,
+    hidden_dim=[1024, 1024, 1024],
+    dropout=0.5,
+    input_dropout=0.4,
+    residual=False,
+)
+
+model_path = "/path/to/model/mymodel/"
+params = {
+    "max_epochs": 500,
+    "logger": True,
+    "log_every_n_steps": 1,
+    "limit_train_batches": 100, # to specify number of batches per epoch
+}
+trainer = pl.Trainer(**params)
+trainer.fit(autoencoder, datamodule=datamodule)
+autoencoder.save_all(model_path=model_path)
+```
 <!-- pyscaffold-notes -->
 
 ## Note

{cellarr-0.2.2 → cellarr-0.2.3}/README.md

@@ -135,6 +135,32 @@ datamodule = DataModule(
 )
 ```
 
+The package also includes a simple autoencoder in the `pytorch-lightning` which makes use of the dataloader. This can be used as a template to create custom architectures and models.
+
+```python
+import pytorch_lightning as pl
+from cellarr.autoencoder import AutoEncoder
+
+autoencoder = AutoEncoder(
+    n_genes=len(datamodule.gene_indices),
+    latent_dim=128,
+    hidden_dim=[1024, 1024, 1024],
+    dropout=0.5,
+    input_dropout=0.4,
+    residual=False,
+)
+
+model_path = "/path/to/model/mymodel/"
+params = {
+    "max_epochs": 500,
+    "logger": True,
+    "log_every_n_steps": 1,
+    "limit_train_batches": 100, # to specify number of batches per epoch
+}
+trainer = pl.Trainer(**params)
+trainer.fit(autoencoder, datamodule=datamodule)
+autoencoder.save_all(model_path=model_path)
+```
 <!-- pyscaffold-notes -->
 
 ## Note

{cellarr-0.2.2 → cellarr-0.2.3}/docs/tutorial.md

@@ -237,6 +237,33 @@ batch = next(iter(dataloader))
 expression, labels, studies = batch
 ```
 
+The dataloader can then be used in training models. The package provides a simple autoencoder to serve as a template for the user's own models.
+
+```python
+import pytorch_lightning as pl
+from cellarr.autoencoder import AutoEncoder
+
+autoencoder = AutoEncoder(
+    n_genes=len(datamodule.gene_indices),
+    latent_dim=128,
+    hidden_dim=[1024, 1024, 1024],
+    dropout=0.5,
+    input_dropout=0.4,
+    residual=False,
+)
+
+model_path = "/path/to/model/mymodel/"
+params = {
+    "max_epochs": 500,
+    "logger": True,
+    "log_every_n_steps": 1,
+    "limit_train_batches": 100, # to specify number of batches per epoch
+}
+trainer = pl.Trainer(**params)
+trainer.fit(autoencoder, datamodule=datamodule)
+autoencoder.save_all(model_path=model_path)
+```
+
 ---
 
 Check out the [documentation](https://biocpy.github.io/cellarr/api/modules.html) for more details.

cellarr-0.2.3/src/cellarr/autoencoder.py

@@ -0,0 +1,444 @@
+import os
+import torch
+from torch import nn
+import torch.nn.functional as F
+from typing import List
+import pytorch_lightning as pl
+
+
+class Encoder(nn.Module):
+    """A class that encapsulates the encoder."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        latent_dim: int = 128,
+        hidden_dim: List[int] = [1024, 1024],
+        dropout: float = 0.5,
+        input_dropout: float = 0.4,
+        residual: bool = False,
+    ):
+        """Constructor.
+
+        Args:
+            n_genes:
+                The number of genes in the gene space, representing the input dimensions.
+
+            latent_dim:
+                The latent space dimensions
+
+            hidden_dim:
+                A list of hidden layer dimensions, describing the number of layers and their dimensions.
+                Hidden layers are constructed in the order of the list for the encoder and in reverse
+                for the decoder.
+
+            dropout:
+                The dropout rate for hidden layers
+
+            input_dropout:
+                The dropout rate for the input layer
+
+            residual:
+                Use residual connections.
+        """
+
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.network = nn.ModuleList()
+        self.residual = residual
+        if self.residual:
+            assert len(set(hidden_dim)) == 1
+        for i in range(len(hidden_dim)):
+            if i == 0:  # input layer
+                self.network.append(
+                    nn.Sequential(
+                        nn.Dropout(p=input_dropout),
+                        nn.Linear(n_genes, hidden_dim[i]),
+                        nn.BatchNorm1d(hidden_dim[i]),
+                        nn.PReLU(),
+                    )
+                )
+            else:  # hidden layers
+                self.network.append(
+                    nn.Sequential(
+                        nn.Dropout(p=dropout),
+                        nn.Linear(hidden_dim[i - 1], hidden_dim[i]),
+                        nn.BatchNorm1d(hidden_dim[i]),
+                        nn.PReLU(),
+                    )
+                )
+        # output layer
+        self.network.append(nn.Linear(hidden_dim[-1], latent_dim))
+
+    def forward(self, x) -> torch.Tensor:
+        """Forward.
+
+        Args:
+            x: torch.Tensor
+                Input tensor corresponding to input layer.
+
+        Returns:
+            Output tensor corresponding to output layer.
+        """
+
+        for i, layer in enumerate(self.network):
+            if self.residual and (0 < i < len(self.network) - 1):
+                x = layer(x) + x
+            else:
+                x = layer(x)
+        return F.normalize(x, p=2, dim=1)
+
+    def save_state(self, filename: str):
+        """Save model state.
+
+        Args:
+            filename:
+                Filename to save the model state.
+        """
+
+        torch.save({"state_dict": self.state_dict()}, filename)
+
+    def load_state(self, filename: str, use_gpu: bool = False):
+        """Load model state.
+
+        Args:
+            filename:
+                Filename containing the model state.
+
+            use_gpu:
+                Boolean indicating whether or not to use GPUs.
+        """
+
+        if not use_gpu:
+            ckpt = torch.load(filename, map_location=torch.device("cpu"))
+        else:
+            ckpt = torch.load(filename)
+        self.load_state_dict(ckpt["state_dict"])
+
+
+class Decoder(nn.Module):
+    """A class that encapsulates the decoder."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        latent_dim: int = 128,
+        hidden_dim: List[int] = [1024, 1024],
+        dropout: float = 0.5,
+        residual: bool = False,
+    ):
+        """Constructor.
+
+        Args:
+            n_genes:
+                The number of genes in the gene space, representing the input dimensions.
+
+            latent_dim:
+                The latent space dimensions
+
+            hidden_dim:
+                A list of hidden layer dimensions, describing the number of layers and their dimensions.
+                Hidden layers are constructed in the order of the list for the encoder and in reverse
+                for the decoder.
+
+            dropout:
+                The dropout rate for hidden layers
+
+            residual:
+                Use residual connections.
+        """
+
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.network = nn.ModuleList()
+        self.residual = residual
+        if self.residual:
+            assert len(set(hidden_dim)) == 1
+        for i in range(len(hidden_dim)):
+            if i == 0:  # first hidden layer
+                self.network.append(
+                    nn.Sequential(
+                        nn.Linear(latent_dim, hidden_dim[i]),
+                        nn.BatchNorm1d(hidden_dim[i]),
+                        nn.PReLU(),
+                    )
+                )
+            else:  # other hidden layers
+                self.network.append(
+                    nn.Sequential(
+                        nn.Dropout(p=dropout),
+                        nn.Linear(hidden_dim[i - 1], hidden_dim[i]),
+                        nn.BatchNorm1d(hidden_dim[i]),
+                        nn.PReLU(),
+                    )
+                )
+        # reconstruction layer
+        self.network.append(nn.Linear(hidden_dim[-1], n_genes))
+
+    def forward(self, x) -> torch.Tensor:
+        """Forward.
+
+        Args:
+            x:
+                Input tensor corresponding to input layer.
+
+        Returns:
+            Output tensor corresponding to output layer.
+        """
+        for i, layer in enumerate(self.network):
+            if self.residual and (0 < i < len(self.network) - 1):
+                x = layer(x) + x
+            else:
+                x = layer(x)
+        return x
+
+    def save_state(self, filename: str):
+        """Save model state.
+
+        Args:
+            filename:
+                Filename to save the model state.
+        """
+
+        torch.save({"state_dict": self.state_dict()}, filename)
+
+    def load_state(self, filename: str, use_gpu: bool = False):
+        """Load model state.
+
+        Args:
+            filename:
+                Filename containing the model state.
+
+            use_gpu:
+                Boolean indicating whether or not to use GPUs.
+        """
+
+        if not use_gpu:
+            ckpt = torch.load(filename, map_location=torch.device("cpu"))
+        else:
+            ckpt = torch.load(filename)
+        self.load_state_dict(ckpt["state_dict"])
+
+
+class AutoEncoder(pl.LightningModule):
+    """A class encapsulating training."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        latent_dim: int = 128,
+        hidden_dim: List[int] = [1024, 1024],
+        dropout: float = 0.5,
+        input_dropout: float = 0.4,
+        lr: float = 5e-3,
+        residual: bool = False,
+    ):
+        """Constructor.
+
+        Args:
+            n_genes:
+                The number of genes in the gene space, representing the input dimensions.
+
+            latent_dim:
+                The latent space dimensions. Defaults to 128.
+
+            hidden_dim:
+                A list of hidden layer dimensions, describing the number of layers and their dimensions.
+                Hidden layers are constructed in the order of the list for the encoder and in reverse
+                for the decoder.
+
+            dropout:
+                The dropout rate for hidden layers
+
+            input_dropout:
+                The dropout rate for the input layer
+
+            lr:
+                The initial learning rate
+
+            residual:
+                Use residual connections.
+        """
+
+        super().__init__()
+
+        # network architecture
+        self.n_genes = n_genes
+        self.latent_dim = latent_dim
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+        self.input_dropout = input_dropout
+        self.lr = lr
+        self.residual = residual
+
+        # networks
+        self.encoder = Encoder(
+            self.n_genes,
+            latent_dim=self.latent_dim,
+            hidden_dim=self.hidden_dim,
+            dropout=self.dropout,
+            input_dropout=self.input_dropout,
+            residual=self.residual,
+        )
+        self.decoder = Decoder(
+            self.n_genes,
+            latent_dim=self.latent_dim,
+            hidden_dim=list(reversed(self.hidden_dim)),
+            dropout=self.dropout,
+            residual=self.residual,
+        )
+
+        self.mse_loss_fn = nn.MSELoss()
+        self.scheduler = None
+        self.val_step_outputs = []
+
+    def configure_optimizers(self):
+        """Configure optimizers."""
+
+        optimizer = torch.optim.AdamW(self.parameters(), self.lr, weight_decay=0.01)
+        self.scheduler = {
+            "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, step_size=0.001),
+            "interval": "epoch",
+            "frequency": 1,
+        }
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": self.scheduler,
+        }  # pytorch-lightning required format
+
+    def forward(self, x):
+        """Forward.
+
+        Args:
+            x:
+                Input tensor corresponding to input layer.
+
+        Returns:
+            Output tensor corresponding to the last encoder layer.
+
+            Output tensor corresponding to the last decoder layer.
+        """
+
+        z = self.encoder(x)
+        x_hat = self.decoder(z)
+        return z, x_hat
+
+    def get_loss(self, batch):
+        """Calculate the loss.
+
+        Args:
+            batch:
+                A batch as defined by a pytorch DataLoader.
+
+        Returns:
+            The training loss
+        """
+
+        cells, labels, studies = batch
+        embedding, reconstruction = self(cells)
+        return self.mse_loss_fn(cells, reconstruction)
+
+    def training_step(self, batch, batch_idx):
+        """Pytorch-lightning training step."""
+
+        loss = self.get_loss(batch)
+        self.log("train loss", loss, prog_bar=True, logger=True)
+        return loss
+
+    def on_validation_epoch_start(self):
+        """Pytorch-lightning validation epoch start."""
+
+        super().on_validation_epoch_start()
+        self.val_step_outputs = []
+
+    def validation_step(self, batch, batch_idx):
+        """Pytorch-lightning validation step."""
+
+        if self.trainer.datamodule.val_dataset is None:
+            return {}
+        return self._eval_step(batch, prefix="val")
+
+    def on_validation_epoch_end(self):
+        """Pytorch-lightning validation epoch end evaluation."""
+
+        if self.trainer.datamodule.val_dataset is None:
+            return {}
+        return self._eval_epoch(prefix="val")
+
+    def _eval_step(self, batch, prefix: str):
+        """Evaluation of validation or test step.
+
+        Args:
+            batch:
+                A batch as defined by a pytorch DataLoader.
+            prefix:
+                A string prefix to label logs.
+
+        Returns:
+            A dictionary containing step evaluation metrics.
+        """
+
+        loss = self.get_loss(batch)
+
+        losses = {
+            f"{prefix}_loss": loss,
+        }
+
+        if prefix == "val":
+            self.val_step_outputs.append(losses)
+        return losses
+
+    def _eval_epoch(self, prefix: str):
+        """Evaluation of validation or test epoch.
+
+        Args:
+            prefix:
+                A string prefix to label logs.
+
+        Returns:
+            A dictionary containing epoch evaluation metrics.
+        """
+
+        if prefix == "val":
+            step_outputs = self.val_step_outputs
+
+        loss = torch.Tensor([step[f"{prefix}_loss"] for step in step_outputs]).mean()
+        self.log(f"{prefix} loss", loss, logger=True)
+
+        losses = {
+            f"{prefix}_loss": loss,
+        }
+        return losses
+
+    def save_all(
+        self,
+        model_path: str,
+    ):
+        if not os.path.isdir(model_path):
+            os.makedirs(model_path)
+
+        # save model
+        self.encoder.save_state(os.path.join(model_path, "encoder.ckpt"))
+        self.decoder.save_state(os.path.join(model_path, "decoder.ckpt"))
+
+    def load_state(
+        self,
+        encoder_filename: str,
+        decoder_filename: str,
+        use_gpu: bool = False,
+    ):
+        """Load model state.
+
+        Args:
+            encoder_filename:
+                Filename containing the encoder model state.
+
+            decoder_filename:
+                Filename containing the decoder model state.
+
+            use_gpu:
+                Boolean indicating whether or not to use GPUs.
+        """
+
+        self.encoder.load_state(encoder_filename, use_gpu)
+        self.decoder.load_state(decoder_filename, use_gpu)

{cellarr-0.2.2 → cellarr-0.2.3}/src/cellarr/build_cellarrdataset.py

@@ -291,9 +291,11 @@ def build_cellarrdataset(
         sample_metadata = pd.DataFrame({"cellarr_sample": _samples})
     elif isinstance(sample_metadata, str):
         sample_metadata = pd.read_csv(sample_metadata, header=0)
-        sample_metadata["cellarr_sample"] = _samples
+        if "cellarr_sample" not in sample_metadata.columns:
+            sample_metadata["cellarr_sample"] = _samples
     elif isinstance(sample_metadata, pd.DataFrame):
-        sample_metadata["cellarr_sample"] = _samples
+        if "cellarr_sample" not in sample_metadata.columns:
+            sample_metadata["cellarr_sample"] = _samples
     else:
         raise TypeError("'sample_metadata' is not an expected type.")
 
@@ -302,11 +304,13 @@ def build_cellarrdataset(
             "Scanning all files for feature ids (e.g. gene symbols), this may take long",
             UserWarning,
         )
-        gene_scan_set = uad.scan_for_features(files_cache, unique=False)
-        gene_set_str = [",".join(x) for x in gene_scan_set]
-        sample_metadata["cellarr_original_gene_set"] = gene_set_str
+        if "cellarr_original_gene_set" not in sample_metadata.columns:
+            gene_scan_set = uad.scan_for_features(files_cache, unique=False)
+            gene_set_str = [",".join(x) for x in gene_scan_set]
+            sample_metadata["cellarr_original_gene_set"] = gene_set_str
 
-        sample_metadata["cellarr_cell_counts"] = cell_counts
+        if "cellarr_cell_counts" not in sample_metadata.columns:
+            sample_metadata["cellarr_cell_counts"] = cell_counts
 
         _col_types = utf.infer_column_types(
             sample_metadata, sample_metadata_options.column_types
@@ -367,8 +371,11 @@ def build_cellarrdataset(
                 "Number of rows in 'cell_metadata' does not match the number of cells across files."
             )
 
-        cell_metadata["cellarr_sample"] = _sample_per_cell
-        cell_metadata["cellarr_cell_index_in_sample"] = _cell_index_in_sample
+        if "cellarr_sample" not in cell_metadata.columns:
+            cell_metadata["cellarr_sample"] = _sample_per_cell
+
+        if "cellarr_cell_index_in_sample" not in cell_metadata.columns:
+            cell_metadata["cellarr_cell_index_in_sample"] = _cell_index_in_sample
 
         cell_metadata.reset_index(drop=True, inplace=True)
 

{cellarr-0.2.2 → cellarr-0.2.3}/src/cellarr.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cellarr
-Version: 0.2.2
+Version: 0.2.3
 Summary: TileDB-based array storage for genomics data collections.
 Home-page: https://github.com/BiocPy/cellarr
 Author: Jayaram Kancherla
@@ -165,6 +165,32 @@ datamodule = DataModule(
 )
 ```
 
+The package also includes a simple autoencoder in the `pytorch-lightning` which makes use of the dataloader. This can be used as a template to create custom architectures and models.
+
+```python
+import pytorch_lightning as pl
+from cellarr.autoencoder import AutoEncoder
+
+autoencoder = AutoEncoder(
+    n_genes=len(datamodule.gene_indices),
+    latent_dim=128,
+    hidden_dim=[1024, 1024, 1024],
+    dropout=0.5,
+    input_dropout=0.4,
+    residual=False,
+)
+
+model_path = "/path/to/model/mymodel/"
+params = {
+    "max_epochs": 500,
+    "logger": True,
+    "log_every_n_steps": 1,
+    "limit_train_batches": 100, # to specify number of batches per epoch
+}
+trainer = pl.Trainer(**params)
+trainer.fit(autoencoder, datamodule=datamodule)
+autoencoder.save_all(model_path=model_path)
+```
 <!-- pyscaffold-notes -->
 
 ## Note

{cellarr-0.2.2 → cellarr-0.2.3}/src/cellarr.egg-info/SOURCES.txt

@@ -31,6 +31,7 @@ docs/_static/.gitignore
 src/cellarr/CellArrDataset.py
 src/cellarr/CellArrDatasetSlice.py
 src/cellarr/__init__.py
+src/cellarr/autoencoder.py
 src/cellarr/build_cellarrdataset.py
 src/cellarr/build_options.py
 src/cellarr/buildutils_tiledb_array.py

{cellarr-0.2.2 → cellarr-0.2.3}/tests/test_dataloader.py

@@ -2,8 +2,10 @@ import tempfile
 
 import anndata
 import numpy as np
+import os
 import pandas as pd
 import pytest
+import pytorch_lightning as pl
 import random
 from scipy.sparse import rand
 import tiledb
@@ -12,6 +14,7 @@ from cellarr import (
     MatrixOptions,
 )
 from cellarr.dataloader import DataModule
+from cellarr.autoencoder import AutoEncoder
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -78,3 +81,54 @@ def test_dataloader():
     assert len(labels) == 100
     assert len(studies) == 100
     assert data.shape == (100, len(datamodule.gene_indices))
+
+
+def test_autoencoder():
+    tempdir = tempfile.mkdtemp()
+
+    adata1 = generate_adata(1000, 100, study="test1")
+    adata2 = generate_adata(100, 1000, study="test2")
+    adata3 = generate_adata(100, 100, study="test3")
+    obs = pd.concat([adata1.obs, adata2.obs, adata3.obs])
+    obs = obs.reset_index(drop=True)
+
+    build_cellarrdataset(
+        output_path=tempdir,
+        files=[adata1, adata2, adata3],
+        cell_metadata=obs,
+        matrix_options=MatrixOptions(dtype=np.float32),
+    )
+
+    datamodule = DataModule(
+        dataset_path=tempdir,
+        cell_metadata_uri="cell_metadata",
+        gene_annotation_uri="gene_annotation",
+        matrix_uri="counts",
+        val_studies=["test3"],
+        label_column_name="label",
+        study_column_name="study",
+        batch_size=100,
+        lognorm=True,
+        target_sum=1e4,
+    )
+
+    autoencoder = AutoEncoder(
+        n_genes=len(datamodule.gene_indices),
+        latent_dim=128,
+        hidden_dim=[1024, 1024],
+        dropout=0.5,
+        input_dropout=0.4,
+        residual=False,
+    )
+
+    params = {
+        "max_epochs": 2,
+        "logger": True,
+        "log_every_n_steps": 1,
+        "limit_train_batches": 4,
+    }
+    trainer = pl.Trainer(**params)
+    trainer.fit(autoencoder, datamodule=datamodule)
+    autoencoder.save_all(model_path=os.path.join(tempdir, "test"))
+    assert os.path.isfile(os.path.join(tempdir, "test", "encoder.ckpt"))
+    assert os.path.isfile(os.path.join(tempdir, "test", "decoder.ckpt"))