embeddings

.gitignore

@@ -1,4 +1,4 @@
-dpacman/data_files
+dpacman/data_files/
 dpacman/preprocess/tfclust/*.log
 dpacman/preprocess/tfclust/temp.py
 bigBedToBed
@@ -25,4 +25,8 @@ dpacman/idmap_filt.csv
 dpacman/temp3.py
 dpacman/temp4.py
 dpacman/temp.ipynb
-dpacman/nohup.out
+dpacman/nohup.out
+dpacman/*/__pycache__/
+dpacman/data_tasks/split/__pycache__/
+dpacman/data_tasks/cluster/__pycache__/
+dpacman/data_tasks/embeddings/__pycache__/

configs/data_task/embeddings/dna.yaml

@@ -0,0 +1,9 @@
+name: dna
+type: embeddings
+
+genome_json_dir: null
+chrom_model: caduceus
+input_file: dpacman/data_files/processed/fimo/post_fimo/fimo_hits_only/maps/dna_seqid_to_dna_sequence.json
+out_dir: dpacman/data_files/processed/embeddings/fimo_hits_only
+
+device: gpu

dpacman/data_tasks/embeddings/__init__.py

@@ -0,0 +1,24 @@
+from .embedders import (
+    CaduceusEmbedder, 
+    DNABertEmbedder, 
+    NucleotideTransformerEmbedder, 
+    GPNEmbedder, 
+    SegmentNTEmbedder, 
+    ESMEmbedder,
+    ESMDBPEmbedder,
+    ProGenEmbedder
+)
+
+def get_embedder(name, device, for_dna=True):
+    name = name.lower()
+    if for_dna:
+        if name=="caduceus":   return CaduceusEmbedder(device)
+        if name=="dnabert":    return DNABertEmbedder(device)
+        if name=="nucleotide": return NucleotideTransformerEmbedder(device)
+        if name=="gpn":        return GPNEmbedder(device)
+        if name=="segmentnt":    return SegmentNTEmbedder(device)
+    else:
+        if name in ("esm",):    return ESMEmbedder(device)
+        if name in ("esm-dbp","esm_dbp"): return ESMDBPEmbedder(device)
+        if name=="progen":      return ProGenEmbedder(device)
+    raise ValueError(f"Unknown model {name} (for_dna={for_dna})")

dpacman/data_tasks/embeddings/dna.py

@@ -0,0 +1,52 @@
+from .utils import pad_token_embeddings, embed_and_save
+from dpacman.data_tasks.embeddings import get_embedder
+
+import logging
+import rootutils
+import os 
+import torch
+import json
+import pandas as pd
+from pathlib import Path
+from omegaconf import DictConfig
+
+root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+logger = logging.getLogger(__name__)
+
+def main(cfg: DictConfig):
+    logger.info(f"Making embeddings using {cfg.data_task.chrom_model} for dna sequences at {cfg.data_task.input_file}")
+    # make out dir if necessary
+    out_dir = Path(root) / cfg.data_task.out_dir
+    os.makedirs(out_dir, exist_ok=True)
+    
+    # set device
+    device = "cpu"
+    if cfg.data_task.device=="gpu":
+        if torch.cuda.is_available():
+            device = "cuda"
+    logger.info(f"Using device: {device}")
+
+    # read the input file 
+    input_file = Path(root) / cfg.data_task.input_file
+    if str(input_file).endswith(".json"):
+        # load the json and isolate the sequences and ids 
+        with open(input_file, "r") as f:
+            d = json.load(f)
+
+        df = pd.DataFrame.from_dict(d, orient="index").reset_index()
+        df.columns = ["seq_id","sequence"]
+        
+    # turn into list of sequences and IDs
+    peak_seqs = df["sequence"].tolist()
+    peak_ids = df["seq_id"].tolist()
+    logger.info(f"Embedding {len(peak_seqs)} binding peak sequences from processed remap data")
+    
+    # Get the DNA embedder
+    dna_embedder = get_embedder(cfg.data_task.chrom_model, device, for_dna=True)
+    out_peaks = out_dir/ f"peaks_{cfg.data_task.chrom_model}.npy"
+    embed_and_save(peak_seqs, peak_ids, dna_embedder, out_peaks)
+
+    logger.info("Finished embedding DNA sequences.")
+    
+if __name__=="__main__":
+    main()

dpacman/data_tasks/embeddings/embedders.py

@@ -22,10 +22,10 @@ import numpy as np
 from pathlib import Path
 import torch
 from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM, pipeline
-import esm
 from Bio import SeqIO
 import time
 import pandas as pd
+import esm
 from tqdm.auto import tqdm
 import logging, math
 
@@ -197,23 +197,6 @@ class NucleotideTransformerEmbedder:
         pooled = [ np.mean(x, axis=0) for x in all_embeddings ]
         return np.vstack(pooled)       
 
-# class ESMEmbedder:
-#     def __init__(self, device):
-#         self.model, self.alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
-#         self.batch_converter = self.alphabet.get_batch_converter()
-#         self.model.to(device).eval()
-#         self.device = device
-
-#     def embed(self, seqs):
-#         batch = [(str(i), seq) for i, seq in enumerate(seqs)]
-#         _, _, toks = self.batch_converter(batch)
-#         toks = toks.to(self.device)
-#         with torch.no_grad():
-#             results = self.model(toks, repr_layers=[33], return_contacts=False)
-#         reps = results["representations"][33]
-#         return reps[:, 1:-1].mean(1).cpu().numpy()
-
-
 class ESMEmbedder:
     def __init__(self, device, model_name="esm2_t33_650M_UR50D"):
         # Try to load the specified ESM-2 model; fallback to esm1b if missing
@@ -280,39 +263,6 @@ class ESMEmbedder:
             all_embeddings.append(seq_vec.cpu().numpy())
         return np.vstack(all_embeddings)  # (N, D)
 
-
-# class ESMDBPEmbedder:
-#     def __init__(self, device):
-#         base_model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()  
-#         model_path = (
-#             Path(__file__).resolve().parent.parent
-#             / "pretrained" / "ESM-DBP" / "ESM-DBP.model"
-#         )
-#         checkpoint = torch.load(model_path, map_location="cpu")
-#         clean_sd = {}
-#         for k, v in checkpoint.items():
-#             clean_sd[k.replace("module.", "")] = v
-#         result = base_model.load_state_dict(clean_sd, strict=False)
-#         if result.missing_keys:
-#             print(f"[ESMDBP] missing keys: {result.missing_keys}")
-#         if result.unexpected_keys:
-#             print(f"[ESMDBP] unexpected keys: {result.unexpected_keys}")
-
-#         self.model = base_model.to(device).eval()
-#         self.alphabet = alphabet
-#         self.batch_converter = alphabet.get_batch_converter()
-#         self.device = device
-
-#     def embed(self, seqs):
-#         batch = [(str(i), seq) for i, seq in enumerate(seqs)]
-#         _, _, toks = self.batch_converter(batch)
-#         toks = toks.to(self.device)
-#         with torch.no_grad():
-#             out = self.model(toks, repr_layers=[33], return_contacts=False)
-#         reps = out["representations"][33]
-#         # skip start/end tokens
-#         return reps[:, 1:-1].mean(1).cpu().numpy()
-
 class ESMDBPEmbedder:
     def __init__(self, device):
         base_model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
@@ -412,149 +362,4 @@ class ProGenEmbedder:
         ).to(self.device)
         with torch.no_grad():
             last_hidden = self.model(**inputs).last_hidden_state
-        return last_hidden.mean(dim=1).cpu().numpy()
-
-# ---- main pipeline ----
-
-def get_embedder(name, device, for_dna=True):
-    name = name.lower()
-    if for_dna:
-        if name=="caduceus":   return CaduceusEmbedder(device)
-        if name=="dnabert":    return DNABertEmbedder(device)
-        if name=="nucleotide": return NucleotideTransformerEmbedder(device)
-        if name=="gpn":        return GPNEmbedder(device)
-        if name=="segmentnt":    return SegmentNTEmbedder(device)
-    else:
-        if name in ("esm",):    return ESMEmbedder(device)
-        if name in ("esm-dbp","esm_dbp"): return ESMDBPEmbedder(device)
-        if name=="progen":      return ProGenEmbedder(device)
-    raise ValueError(f"Unknown model {name} (for_dna={for_dna})")
-
-
-def pad_token_embeddings(list_of_arrays, pad_value=0.0):
-    """
-    list_of_arrays: list of (L_i, D) numpy arrays
-    Returns:
-      padded: (N, L_max, D) array
-      mask:   (N, L_max) boolean array where True = real token, False = padding
-    """
-    N = len(list_of_arrays)
-    D = list_of_arrays[0].shape[1]
-    L_max = max(arr.shape[0] for arr in list_of_arrays)
-    padded = np.full((N, L_max, D), pad_value, dtype=list_of_arrays[0].dtype)
-    mask = np.zeros((N, L_max), dtype=bool)
-    for i, arr in enumerate(list_of_arrays):
-        L = arr.shape[0]
-        padded[i, :L] = arr
-        mask[i, :L] = True
-    return padded, mask
-
-def embed_and_save(seqs, ids, embedder, out_path):
-    embs = embedder.embed(seqs)
-
-    # Decide whether we got variable-length per-token outputs (list of (L, D))
-    is_variable_token = isinstance(embs, (list, tuple)) and len(embs) > 0 and hasattr(embs[0], "shape") and embs[0].ndim == 2
-
-    if is_variable_token:
-        # pad to (N, L_max, D) + mask
-        padded, mask = pad_token_embeddings(embs)
-        # Save both embeddings and mask together in an .npz for convenience
-        np.savez_compressed(out_path.with_suffix(".caduceus.npz"),
-                            embeddings=padded,
-                            mask=mask,
-                            ids=np.array(ids, dtype=object))
-    else:
-        # fixed shape output, e.g., pooled (N, D)
-        array = np.vstack(embs) if isinstance(embs, list) else embs
-        np.save(out_path, array)
-        with open(out_path.with_suffix(".ids"), "w") as f:
-            f.write("\n".join(ids))
-
-
-if __name__=="__main__":
-
-    p = argparse.ArgumentParser()
-    #p.add_argument("--peak_fasta", default="binding_peaks_unique.fa", help="FASTA of deduplicated binding peak sequences; if present this is used for DNA embedding instead of genome JSONs")
-    p.add_argument("--genome-json-dir", default=None, help="(fallback) directory of UCSC JSONs for full chromosome embedding if peak FASTA is missing or you explicitly want chromosomes")
-    p.add_argument("--skip-dna", action="store_true", help="if set, skip the chromosome embedding step") #if glm embeddings successful but not plm embeddings
-    p.add_argument("--tf-fasta",      required=True, help="input TF FASTA file")
-    p.add_argument("--chrom-model",   default="caduceus")
-    p.add_argument("--tf-model",      default="esm-dbp")
-    p.add_argument("--out-dir",       default="dpacman/model/embeddings")
-    p.add_argument("--device",        default="cpu")
-    args = p.parse_args()
-
-    os.makedirs(args.out_dir, exist_ok=True)
-    device = args.device
-    print(device)
-
-    if not args.skip_dna:
-        if args.genome_json_dir == None:
-            dna_df = pd.read_parquet('/home/a03-akrishna/DPACMAN/dpacman/model/remap2022_crm_fimo_output_q_processed.parquet', engine='pyarrow')
-            #df.to_csv('/home/a03-akrishna/DPACMAN/dpacman/model/remap2022_crm_fimo_output_q_processed.csv', index=False)
-            peak_seqs = dna_df["dna_sequence"]
-            peak_ids = dna_df["ID"]
-            print(f"Embedding {len(peak_seqs)} binding peak sequences from processed remap data", flush=True)
-            dna_embedder = get_embedder(args.chrom_model, device, for_dna=True)
-            out_peaks = Path(args.out_dir) / f"peaks_{args.chrom_model}.npy"
-            embed_and_save(peak_seqs, peak_ids, dna_embedder, out_peaks)
-
-        # peak_fasta = Path(args.peak_fasta)
-        # if peak_fasta.exists():
-        #     # Load peak sequences from FASTA
-        #     from Bio import SeqIO
-
-        #     peak_seqs = []
-        #     peak_ids = []
-        #     for rec in SeqIO.parse(peak_fasta, "fasta"):
-        #         peak_ids.append(rec.id)
-        #         peak_seqs.append(str(rec.seq))
-        #     print(f"Embedding {len(peak_seqs)} binding peak sequences from {peak_fasta}", flush=True)
-        #     dna_embedder = get_embedder(args.chrom_model, device, for_dna=True)
-        #     out_peaks = Path(args.out_dir) / f"peaks_{args.chrom_model}.npy"
-        #     embed_and_save(peak_seqs, peak_ids, dna_embedder, out_peaks)
-        elif args.genome_json_dir:
-            # Legacy: load full chromosomes from JSONs (chr1–22, X, Y, M)
-            genome_dir = Path(args.genome_json_dir)
-            chrom_seqs, chrom_ids = [], []
-            primary_pattern = re.compile(r"^hg38_chr(?:[1-9]|1[0-9]|2[0-2]|X|Y|M)\.json$")
-            for j in sorted(genome_dir.iterdir()):
-                if not primary_pattern.match(j.name):
-                    continue
-                data = json.loads(j.read_text())
-                seq = data.get("dna") or data.get("sequence")
-                chrom = data.get("chrom") or j.stem.split("_")[-1]
-                chrom_seqs.append(seq)
-                chrom_ids.append(chrom)
-            cutoff = CaduceusEmbedder(device).chunk_size
-            long_chroms = [
-                (chrom, len(seq))
-                for chrom, seq in zip(chrom_ids, chrom_seqs)
-                if len(seq) > cutoff
-            ]
-            if long_chroms:
-                print("⚠️ Chromosomes exceeding Caduceus max tokens ({}):".format(cutoff))
-                for chrom, L in long_chroms:
-                    print(f"  {chrom}: {L} bases")
-            else:
-                print("All chromosomes ≤ Caduceus limit ({}).".format(cutoff))
-
-            chrom_embedder = get_embedder(args.chrom_model, device, for_dna=True)
-            out_chrom = Path(args.out_dir) / f"chrom_{args.chrom_model}.npy"
-            embed_and_save(chrom_seqs, chrom_ids, chrom_embedder, out_chrom)
-        else:
-            raise ValueError("No input for DNA embedding: provide a peak FASTA (default binding_peaks_unique.fa) or set --genome-json-dir for chromosome JSONs.")
-
-
-    #Load TF sequences
-    tf_seqs, tf_ids = [], []
-    for record in SeqIO.parse(args.tf_fasta, "fasta"):
-        tf_ids.append(record.id)
-        tf_seqs.append(str(record.seq))
-
-    # embed and save
-    tf_embedder = get_embedder(args.tf_model, device, for_dna=False)
-    out_tf = Path(args.out_dir) / f"tf_{args.tf_model}.npy"
-    embed_and_save(tf_seqs, tf_ids, tf_embedder, out_tf)
-
-    print("Done.")
+        return last_hidden.mean(dim=1).cpu().numpy()

dpacman/data_tasks/embeddings/utils.py

@@ -0,0 +1,47 @@
+"""
+Utility funcitons related to creating embeddings
+"""
+import numpy as np
+
+def pad_token_embeddings(list_of_arrays, pad_value=0.0):
+    """
+    list_of_arrays: list of (L_i, D) numpy arrays
+    Returns:
+      padded: (N, L_max, D) array
+      mask:   (N, L_max) boolean array where True = real token, False = padding
+    """
+    N = len(list_of_arrays)
+    D = list_of_arrays[0].shape[1]
+    L_max = max(arr.shape[0] for arr in list_of_arrays)
+    padded = np.full((N, L_max, D), pad_value, dtype=list_of_arrays[0].dtype)
+    mask = np.zeros((N, L_max), dtype=bool)
+    for i, arr in enumerate(list_of_arrays):
+        L = arr.shape[0]
+        padded[i, :L] = arr
+        mask[i, :L] = True
+    return padded, mask
+
+def embed_and_save(seqs, ids, embedder, out_path):
+    """
+    Using the passed embedder, make embeddings
+    """
+    embs = embedder.embed(seqs)
+
+    # Decide whether we got variable-length per-token outputs (list of (L, D))
+    is_variable_token = isinstance(embs, (list, tuple)) and len(embs) > 0 and hasattr(embs[0], "shape") and embs[0].ndim == 2
+
+    if is_variable_token:
+        # pad to (N, L_max, D) + mask
+        padded, mask = pad_token_embeddings(embs)
+        # Save both embeddings and mask together in an .npz for convenience
+        np.savez_compressed(out_path.with_suffix(".caduceus.npz"),
+                            embeddings=padded,
+                            mask=mask,
+                            ids=np.array(ids, dtype=object),
+                            seqs=np.array(seqs, dtype=object)) 
+    else:
+        # fixed shape output, e.g., pooled (N, D)
+        array = np.vstack(embs) if isinstance(embs, list) else embs
+        np.save(out_path, array)
+        with open(out_path.with_suffix(".ids"), "w") as f:
+            f.write("\n".join(ids))

dpacman/data_tasks/split/remap.py

@@ -450,7 +450,8 @@ def main(cfg: DictConfig):
         )
         dna_assign, kept_by_split = results
         
-        edge_df["split"] = edge_df["dna_seqid"].map(dna_assign)
+        # assign datapoints to cluster by their DNA cluster rep
+        edge_df["split"] = edge_df["dna_cluster_rep"].map(dna_assign)
     else:
         results = split_bipartite_by_components(
             edges,
@@ -498,9 +499,11 @@ def main(cfg: DictConfig):
     # ensure there is no overlap
     check_validity(train, val, test, split_by=cfg.data_task.split_by)
     
-    logger.info(f"Length of train dataset: {len(train)} ({100*len(train)/sum([len(train),len(val),len(test)]):.2f}%)")
-    logger.info(f"Length of val dataset: {len(val)} ({100*len(val)/sum([len(train),len(val),len(test)]):.2f}%)")
-    logger.info(f"Length of test dataset: {len(test)} ({100*len(test)/sum([len(train),len(val),len(test)]):.2f}%)")
+    total = sum([len(train),len(val),len(test)])
+    logger.info(f"Length of train dataset: {len(train)} ({100*len(train)/total:.2f}%)")
+    logger.info(f"Length of val dataset: {len(val)} ({100*len(val)/total:.2f}%)")
+    logger.info(f"Length of test dataset: {len(test)} ({100*len(test)/total:.2f}%)")
+    logger.info(f"Total sequences = {total}. Same as edges size? {total==len(edge_df)}")
     
     # create the output dir
     split_out_dir = Path(root)/cfg.data_task.split_out_dir

dpacman/scripts/preprocess.py

@@ -15,6 +15,7 @@ from dpacman.data_tasks.fimo.run_fimo import main as run_fimo_main
 from dpacman.data_tasks.fimo.post_fimo import main as post_fimo_main
 from dpacman.data_tasks.cluster.remap import main as cluster_remap_main
 from dpacman.data_tasks.split.remap import main as split_remap_main
+from dpacman.data_tasks.embeddings.dna import main as embed_dna_main
 
 @hydra.main(
     config_path=str(root / "configs"), config_name="preprocess", version_base="1.3"
@@ -59,12 +60,19 @@ def main(cfg: DictConfig):
         else:
             raise ValueError(f"No clean pipeline defined for: {task_name}")
     
+    # Split
     elif task_type == "split":
         if task_name == "remap":
             split_remap_main(cfg)
         else:
             raise ValueError(f"No clean pipeline defined for: {task_name}")
 
+    # Embed
+    elif task_type=="embeddings":
+        if task_name == "dna":
+            embed_dna_main(cfg)
+        else:
+            raise ValueError(f"No clean pipeline defined for: {task_name}")
     # Unknown - error
     else:
         raise ValueError(f"Unknown task type: {task_type}")

dpacman/scripts/run_embeddings.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+# Manually specify values used in the config
+main_task="preprocess"
+data_task_type="embeddings"
+timestamp=$(date "+%Y-%m-%d_%H-%M-%S")
+
+run_dir="$HOME/DPACMAN/logs/${main_task}/${data_task_type}/runs/${timestamp}"
+mkdir -p "$run_dir"
+
+nohup python -u -m scripts.preprocess \
+  hydra.run.dir="${run_dir}" \
+  data_task="${data_task_type}/dna" \
+  > "${run_dir}/run.log" 2>&1 &
+
+echo $! > "${run_dir}/pid.txt"

dpacman/scripts/run_split.sh

@@ -5,7 +5,7 @@ main_task="preprocess"
 data_task_type="split"
 timestamp=$(date "+%Y-%m-%d_%H-%M-%S")
 
-run_dir="/vast/projects/pranam/lab/sophie/DPACMAN/logs/${main_task}/${data_task_type}/runs/${timestamp}"
+run_dir="$HOME/DPACMAN/logs/${main_task}/${data_task_type}/runs/${timestamp}"
 mkdir -p "$run_dir"
 
 nohup python -u -m scripts.preprocess \

environment.yaml

@@ -19,12 +19,12 @@ dependencies:
   - matplotlib=3.10.*
   - pip:
       # Pull GPU wheels (CUDA 12.8) from PyTorch's cu128 index; fall back to PyPI for others
-      - --index-url https://download.pytorch.org/whl/cu128
+      - --index-url https://download.pytorch.org/whl/cu129
       - --extra-index-url https://pypi.org/simple
 
-      # PyTorch + CUDA 12.8
-      - torch==2.7.1
-      - torchvision==0.22.1
+      # PyTorch + CUDA 12.9
+      - torch==2.8
+      - torchvision==0.23
       # - torchaudio==2.7.1   # optional, if you need it
 
       # Lightning (classic)
@@ -41,5 +41,6 @@ dependencies:
       - scikit-learn==1.7.1
       - biopython==1.85
       - ortools==9.14.6206
+      - esm==3.2.1.post1
       # Your package in editable mode
       - -e .