Refactor/write h5 speedup numpy2 compat

setup.cfg

@@ -56,7 +56,7 @@ install_requires =
     torch
     pyfaidx
     genomepy
-    grelu>=1.0.10,<=1.0.11
+    grelu
     lightning
     torchmetrics
     bioframe

src/decima/cli/attributions.py

@@ -245,6 +245,9 @@ def cli_attributions_predict(
 @click.option(
     "--genome", type=str, show_default=True, default="hg38", help="Genome name or path to the genome fasta file."
 )
+@click.option(
+    "--no-bigwig", "disable_bigwig", is_flag=True, default=False, help="Disable bigwig output file generation."
+)
 def cli_attributions(
     output_prefix,
     genes,
@@ -267,6 +270,7 @@ def cli_attributions(
     meme_motif_db,
     device,
     genome,
+    disable_bigwig,
 ):
     """Generate and save attribution analysis results for a gene or a set of sequences and perform seqlet calling on the attributions.
 
@@ -312,6 +316,7 @@ def cli_attributions(
         pattern_type=pattern_type,
         meme_motif_db=meme_motif_db,
         genome=genome,
+        bigwig=not disable_bigwig,
     )
 
 

src/decima/cli/finetune.py

@@ -45,6 +45,7 @@
 @click.option("--logger", default="wandb", type=str, help="Logger.")
 @click.option("--num-workers", default=16, type=int, help="Number of workers.")
 @click.option("--seed", default=0, type=int, help="Random seed.")
+@click.option("--checkpoint", default=None, type=str, help="Path to a checkpoint to resume training from.")
 def cli_finetune(
     name,
     model,
@@ -63,6 +64,7 @@ def cli_finetune(
     logger,
     num_workers,
     seed,
+    checkpoint,
 ):
     """Finetune the Decima model.
 
@@ -102,8 +104,11 @@ def cli_finetune(
     )
     val_dataset = HDF5Dataset(h5_file=h5_file, ad=ad, key="val", max_seq_shift=0)
 
-    if isinstance(device, str) and device.isdigit():
-        device = int(device)
+    if isinstance(device, str):
+        if "," in device:
+            device = [int(d) for d in device.split(",")]
+        elif device.isdigit():
+            device = int(device)
 
     train_params = {
         "batch_size": batch_size,
@@ -137,7 +142,7 @@ def cli_finetune(
         run = wandb.init(project="decima", dir=name, name=name)
 
     logger.info("Training")
-    model.train_on_dataset(train_dataset, val_dataset)
+    model.train_on_dataset(train_dataset, val_dataset, checkpoint_path=checkpoint)
     train_dataset.close()
     val_dataset.close()
     if logger == "wandb":

src/decima/data/dataset.py

@@ -958,8 +958,8 @@ def __str__(self):
             "VariantDataset("
             f"{self.variants.shape[0]} variants "
             f"from {list(self.variants.chrom.unique())} "
-            f"between {self.variants.start.min()} "
-            f"and {self.variants.end.max()} bp from TSS"
+            f"between {self.variants.tss_dist.abs().min()} "
+            f"and {self.variants.tss_dist.abs().max()} bp from TSS"
             ")"
         )
 

src/decima/data/write_hdf5.py

@@ -1,27 +1,35 @@
 import h5py
 import numpy as np
-from grelu.sequence.format import convert_input_type
+from grelu.io.genome import get_genome
+from grelu.sequence.format import _BASE_LUT
 from grelu.sequence.utils import get_unique_length
+from tqdm import tqdm
 
 
-def write_hdf5(file, ad, pad=0, genome="hg38"):
+def write_hdf5(file, ad, pad=0, genome="hg38", batch_size=1000):
     """Write AnnData object to HDF5 file.
 
     Args:
         file: Path to the HDF5 file to write
         ad: AnnData object containing the data
         pad: Amount of padding to add. Defaults to 0
         genome: Genome name or path to the genome fasta file. Defaults to "hg38"
+        batch_size: Number of genes per write batch. Defaults to 1000
     """
-    # Calculate seq_len
     seq_len = get_unique_length(ad.var)
+    padded_seq_len = seq_len + 2 * pad
+    n_genes = ad.var.shape[0]
+    genome_obj = get_genome(genome)
+
+    intervals = ad.var[["chrom", "start", "end", "strand"]].copy()
+    intervals["start"] = intervals["start"] - pad
+    intervals["end"] = intervals["end"] + pad
 
     with h5py.File(file, "w") as f:
         # Metadata
         print("Writing metadata")
         f.create_dataset("pad", shape=(), data=pad)
         f.create_dataset("seq_len", shape=(), data=seq_len)
-        padded_seq_len = seq_len + 2 * pad
         f.create_dataset("padded_seq_len", shape=(), data=padded_seq_len)
 
         # Tasks
@@ -35,26 +43,47 @@ def write_hdf5(file, ad, pad=0, genome="hg38"):
         f.create_dataset("genes", shape=arr.shape, data=arr)
 
         # Labels
-        arr = np.expand_dims(ad.X.T.astype(np.float32), 2)
-        print(f"Writing labels array of shape: {arr.shape}")
+        print("Writing labels")
+        X = ad.X.toarray() if hasattr(ad.X, "toarray") else np.asarray(ad.X)
+        arr = np.expand_dims(X.T.astype(np.float32), 2)
+        print(f"  shape: {arr.shape}")
         f.create_dataset("labels", shape=arr.shape, dtype=np.float32, data=arr)
+        del X, arr
+
+        # Masks and sequences — written in batches to avoid OOM
+        print("Writing masks and sequences")
+        masks_ds = f.create_dataset(
+            "masks", shape=(n_genes, padded_seq_len), dtype=np.float32
+        )
+        seqs_ds = f.create_dataset(
+            "sequences", shape=(n_genes, padded_seq_len), dtype=np.int8
+        )
+
+        n_batches = (n_genes + batch_size - 1) // batch_size
+        for b in tqdm(range(n_batches), desc="Batches"):
+            start_i = b * batch_size
+            end_i = min(start_i + batch_size, n_genes)
+            batch_var = ad.var.iloc[start_i:end_i]
+            batch_iv = intervals.iloc[start_i:end_i]
+
+            masks = np.zeros((end_i - start_i, padded_seq_len), dtype=np.float32)
+            seqs = np.empty((end_i - start_i, padded_seq_len), dtype=np.int8)
+
+            for j, (row_var, row_iv) in enumerate(
+                zip(batch_var.itertuples(), batch_iv.itertuples())
+            ):
+                masks[j, row_var.gene_mask_start + pad : row_var.gene_mask_end + pad] = 1.0
+                seq = str(
+                    genome_obj.get_seq(
+                        row_iv.chrom,
+                        row_iv.start + 1,
+                        row_iv.end,
+                        rc=row_iv.strand == "-",
+                    )
+                ).upper()
+                seqs[j] = _BASE_LUT[np.frombuffer(seq.encode("ascii"), dtype=np.uint8)]
 
-        # Gene masks
-        print("Making gene masks")
-        shape = (ad.var.shape[0], padded_seq_len)
-        arr = np.zeros(shape=shape)
-        for i, row in enumerate(ad.var.itertuples()):
-            arr[i, row.gene_mask_start + pad : row.gene_mask_end + pad] += 1
-        print(f"Writing mask array of shape: {arr.shape}")
-        f.create_dataset("masks", shape=shape, dtype=np.float32, data=arr)
-
-        # Sequences
-        print("Encoding sequences")
-        arr = ad.var[["chrom", "start", "end", "strand"]].copy()
-        arr.start = arr.start - pad
-        arr.end = arr.end + pad
-        arr = convert_input_type(arr, "indices", genome=genome)
-        print(f"Writing sequence array of shape: {arr.shape}")
-        f.create_dataset("sequences", shape=arr.shape, dtype=np.int8, data=arr)
+            masks_ds[start_i:end_i] = masks
+            seqs_ds[start_i:end_i] = seqs
 
     print("Done!")

src/decima/interpret/attributions.py

@@ -164,7 +164,7 @@ def predict_save_attributions(
     attributer = DecimaAttributer(model, tasks, off_tasks, method, transform)
 
     with QCLogger(str(output_prefix) + ".warnings.qc.log", metadata_anndata=result) as qc:
-        if result.ground_truth is not None:
+        if result.ground_truth is not None and "preds" in result.anndata.layers:
             qc.log_correlation(tasks, off_tasks, plot=True)
 
         if (genes is not None) and (seqs is not None):
@@ -356,6 +356,7 @@ def predict_attributions_seqlet_calling(
     pattern_type: str = "both",
     meme_motif_db: str = "hocomoco_v13",
     genome: str = "hg38",
+    bigwig: bool = True,
 ):
     """Generate and save attribution analysis results for a gene.
     This function performs attribution analysis for a given gene and cell types, saving
@@ -413,6 +414,7 @@ def predict_attributions_seqlet_calling(
         top_n_markers=top_n_markers,
         device=device,
         genome=genome,
+        bigwig=bigwig,
     )
 
     custom_genome = False

src/decima/model/lightning.py

@@ -207,7 +207,7 @@ def count_params(self) -> int:
         """
         return sum(p.numel() for p in self.model.parameters() if p.requires_grad)
 
-    def parse_logger(self) -> str:
+    def parse_logger(self, checkpoint_path=None) -> str:
         """
         Parses the name of the logger supplied in train_params.
         """
@@ -218,7 +218,13 @@ def parse_logger(self) -> str:
                 save_dir=self.train_params["save_dir"],
             )
         elif self.train_params["logger"] == "csv":
-            logger = CSVLogger(name=self.name, save_dir=self.train_params["save_dir"])
+            import re
+            version = None
+            if checkpoint_path is not None:
+                match = re.search(r"version_(\d+)", checkpoint_path)
+                if match:
+                    version = int(match.group(1))
+            logger = CSVLogger(name=self.name, save_dir=self.train_params["save_dir"], version=version)
         else:
             raise NotImplementedError
         return logger
@@ -310,15 +316,17 @@ def train_on_dataset(
         torch.set_float32_matmul_precision("medium")
 
         # Set up logging
-        logger = self.parse_logger()
+        logger = self.parse_logger(checkpoint_path=checkpoint_path)
 
         # Set up trainer
+        devices = make_list(self.train_params["devices"])
         trainer = pl.Trainer(
             max_epochs=self.train_params["max_epochs"],
             accelerator="gpu",
-            devices=make_list(self.train_params["devices"]),
+            devices=devices,
+            strategy="ddp" if len(devices) > 1 else "auto",
             logger=logger,
-            callbacks=[ModelCheckpoint(monitor="val_loss", mode="min", save_top_k=self.train_params["save_top_k"])],
+            callbacks=[ModelCheckpoint(monitor="val_loss", mode="min", save_top_k=self.train_params["save_top_k"], save_last=True)],
             default_root_dir=self.train_params["save_dir"],
             accumulate_grad_batches=self.train_params["accumulate_grad_batches"],
             gradient_clip_val=self.train_params["clip"],

src/decima/tools/inference.py

@@ -86,7 +86,7 @@ def evaluate_gene_expression_predictions(ad):
 
     n_pbs = ad.shape[0]
     n_genes = ad.shape[1]
-    truth = ad.X
+    truth = ad.X.toarray() if hasattr(ad.X, "toarray") else np.asarray(ad.X)
     preds = ad.layers["preds"]
 
     # Compute Pearson correlation per gene

src/decima/utils/io.py

@@ -332,10 +332,10 @@ def add(
                 start = 0
                 end = DECIMA_CONTEXT_SIZE
             else:
-                gene_meta = self.result.get_gene_metadata(self.genes[idx])
-                chrom = gene_meta.chrom
-                start = gene_meta.start
-                end = gene_meta.end
+                gene_var = self.result.gene_metadata.loc[self.genes[idx]]
+                chrom = gene_var["chrom"]
+                start = int(gene_var["start"])
+                end = int(gene_var["end"])
 
             if self.correct_grad_bigwig:
                 attrs = attrs - attrs.mean(axis=0, keepdims=True)

src/decima/utils/qc.py

@@ -41,6 +41,8 @@ def log_gene(self, gene: str, threshold: float = 0.5):
             gene (str): Gene to log
             threshold (float, optional): Threshold for logging. Defaults to 0.5.
         """
+        if "pearson" not in self.result.gene_metadata.columns:
+            return
         gene_metadata = self.result.get_gene_metadata(gene)
         if gene_metadata.pearson < threshold:
             self.log(