diff --git a/setup.cfg b/setup.cfg index 6d6101d..23ef0c2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -56,7 +56,7 @@ install_requires = torch pyfaidx genomepy - grelu>=1.0.10,<=1.0.11 + grelu lightning torchmetrics bioframe diff --git a/src/decima/cli/attributions.py b/src/decima/cli/attributions.py index 783271a..b729210 100644 --- a/src/decima/cli/attributions.py +++ b/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, ) diff --git a/src/decima/cli/finetune.py b/src/decima/cli/finetune.py index 70105bb..64fbdf3 100755 --- a/src/decima/cli/finetune.py +++ b/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": diff --git a/src/decima/data/dataset.py b/src/decima/data/dataset.py index 3fad18d..6f3dbd7 100644 --- a/src/decima/data/dataset.py +++ b/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" ")" ) diff --git a/src/decima/data/write_hdf5.py b/src/decima/data/write_hdf5.py index ccd3c1d..0d01f45 100644 --- a/src/decima/data/write_hdf5.py +++ b/src/decima/data/write_hdf5.py @@ -1,10 +1,12 @@ 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: @@ -12,16 +14,22 @@ def write_hdf5(file, ad, pad=0, genome="hg38"): 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!") diff --git a/src/decima/interpret/attributions.py b/src/decima/interpret/attributions.py index 31b5922..e70f7c5 100644 --- a/src/decima/interpret/attributions.py +++ b/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 diff --git a/src/decima/model/lightning.py b/src/decima/model/lightning.py index 1c20ac7..7d32e2a 100644 --- a/src/decima/model/lightning.py +++ b/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"], diff --git a/src/decima/tools/inference.py b/src/decima/tools/inference.py index 448a8c8..aab9f07 100644 --- a/src/decima/tools/inference.py +++ b/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 diff --git a/src/decima/utils/io.py b/src/decima/utils/io.py index 72ad854..50613b5 100644 --- a/src/decima/utils/io.py +++ b/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) diff --git a/src/decima/utils/qc.py b/src/decima/utils/qc.py index b77b0ca..e375c4c 100644 --- a/src/decima/utils/qc.py +++ b/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(