add synthetic gc content concept & test

Author: yztxwd

test/test_cav_trainer.py

@@ -115,6 +115,7 @@ def test_run_tpcav_random_control(self):
             genome_fasta=genome_fasta,
             num_motif_insertions=[4, 8],
             bed_seq_file="data/hg38_rmsk.sample.bed",
+            synthetic_gc_concept_step=0.1,
             output_dir="data/test_run_tpcav_output/",
         )
 
@@ -234,13 +235,37 @@ def test_motif_concepts_against_permute_control(self):
 
         cav_trainer.train_concepts_pairs(concepts_pairs, 200, output_dir="data/cavs_permute/", num_processes=2)
 
+    def test_synthetic_gc_content_concept(self):
+
+        builder = ConceptBuilder(
+            genome_fasta="data/hg38.analysisSet.fa",
+            input_window_length=1024,
+            bws=None,
+            num_motifs=12,
+            include_reverse_complement=True,
+            min_samples=1000,
+            batch_size=8,
+        )
+
+        builder.build_control()
+
+        builder.add_synthetic_gc_content_concepts(0.1)
+
+        for c in builder.concepts:
+            gc_content = float(c.name.split('_')[-1])
+            seq, _ = next(iter(c.data_iter))
+            total_c = 0.; gc_c = 0.
+            for s in seq:
+                total_c += len(s)
+                gc_c += s.count('G') + s.count('C')
+            assert (gc_c/total_c) - gc_content < 0.05
 
     def test_all(self):
-        lp = LineProfiler()
-                # Add installed-package functions you care about
-        lp.add_function(utils.iterate_seq_df_chunk)
-        lp.add_function(CavTrainer.train_concepts)
-        lp.enable_by_count()
+        #lp = LineProfiler()
+        #        # Add installed-package functions you care about
+        #lp.add_function(utils.iterate_seq_df_chunk)
+        #lp.add_function(CavTrainer.train_concepts)
+        #lp.enable_by_count()
 
         motif_path = Path("data") / "motif-clustering-v2.1beta_consensus_pwms.test.meme"
         self.assertTrue(motif_path.exists(), "Motif file is missing")
@@ -259,6 +284,8 @@ def test_all(self):
 
         builder.add_meme_motif_concepts(str(motif_path))
 
+        builder.add_synthetic_gc_content_concepts(0.1)
+
         builder.apply_transform(transform_fasta_to_one_hot_seq)
 
         batch = next(iter(builder.all_concepts()[0].data_iter))
@@ -370,7 +397,7 @@ def forward_from_layer_1_embeddings(tm, avs_residual, avs_projected):
             f"Attributions do not match, max difference is {torch.abs(attributions - attributions_old).max()}",
         )
 
-        lp.disable_by_count()
+        #lp.disable_by_count()
         #lp.print_stats()
 
         # test save and restore states

tpcav/cavs.py

@@ -877,6 +877,7 @@ def run_tpcav(
     motif_control_type="random",
     bed_seq_file: Optional[str] = None,
     bed_chrom_file: Optional[str] = None,
+    synthetic_gc_concept_step: Optional[float] = None,
     layer_name: Optional[str]=None,
     layer=None,
     output_dir: str = "tpcav/",
@@ -941,8 +942,8 @@ def run_tpcav(
         motif_concept_builders[nm] = builder
 
     ## bed concepts (optional)
-    if bed_seq_file is not None or bed_chrom_file is not None:
-        bed_builder = ConceptBuilder(
+    if bed_seq_file is not None or bed_chrom_file is not None or synthetic_gc_concept_step is not None:
+        non_motif_concept_builder = ConceptBuilder(
             genome_fasta=genome_fasta,
             input_window_length=input_window_length,
             bws=bws,
@@ -953,23 +954,26 @@ def run_tpcav(
             rng_seed = seed,
         )
         # use random regions as control
-        bed_builder.build_control()
+        non_motif_concept_builder.build_control()
         if bed_seq_file is not None:
             # build concepts from fasta sequences in bed file
-            bed_builder.add_bed_sequence_concepts(bed_seq_file)
+            non_motif_concept_builder.add_bed_sequence_concepts(bed_seq_file)
         if bed_chrom_file is not None:
             # build concepts from chromatin tracks in bed file
-            bed_builder.add_bed_chrom_concepts(bed_chrom_file)
+            non_motif_concept_builder.add_bed_chrom_concepts(bed_chrom_file)
+        if synthetic_gc_concept_step is not None:
+            # build synthetic gc content concepts
+            non_motif_concept_builder.add_synthetic_gc_content_concepts(synthetic_gc_concept_step)
         # apply transform to convert fasta sequences to one-hot encoded sequences
-        bed_builder.apply_transform(input_transform_func)
+        non_motif_concept_builder.apply_transform(input_transform_func)
     else:
-        bed_builder = None
+        non_motif_concept_builder = None
 
     # create TPCAV model on top of the given model
     tpcav_model = TPCAV(model, layer_name=layer_name, layer=layer)
     # fit PCA on sampled all concept activations of the last builder (should have the most motifs)
     tpcav_model.fit_pca(
-        concepts=motif_concept_builders[num_motif_insertions[-1]].concepts_for_pca() + bed_builder.concepts_for_pca() if  bed_builder is not None else motif_concept_builders[num_motif_insertions[-1]].concepts_for_pca(),
+        concepts=motif_concept_builders[num_motif_insertions[-1]].concepts_for_pca() + non_motif_concept_builder.concepts_for_pca() if non_motif_concept_builder is not None else motif_concept_builders[num_motif_insertions[-1]].concepts_for_pca(),
         num_samples_per_concept=num_samples_for_pca,
         num_pc=num_pc,
     )
@@ -993,13 +997,14 @@ def run_tpcav(
         if save_cav_trainer:
             torch.save(cav_trainer, str(output_path / f"cavs_{nm}_motifs/cav_trainer.pt"))
         motif_cav_trainers[nm] = cav_trainer
-    if bed_builder is not None:
+
+    if non_motif_concept_builder is not None:
         bed_cav_trainer = CavTrainer(tpcav_model, penalty="l2")
         bed_cav_trainer.set_control(
-            bed_builder.control_concepts[0], num_samples=num_samples_for_cav
+            non_motif_concept_builder.control_concepts[0], num_samples=num_samples_for_cav
         )
         bed_cav_trainer.train_concepts(
-            bed_builder.concepts,
+            non_motif_concept_builder.concepts,
             num_samples_for_cav,
             output_dir=str(output_path / f"cavs_bed_concepts/"),
             num_processes=p,

tpcav/concepts.py

@@ -33,6 +33,26 @@ def __iter__(self):
             yield inputs
 
 
+class _SyntheticGCSeqIterator:
+    def __init__(self, seq_len: int, n: int, batch_size: int, gc: float, seed: int):
+        self.seq_len = int(seq_len)
+        self.n = int(n)
+        self.batch_size = int(batch_size)
+        self.gc = float(gc)
+        self.seed = int(seed)
+
+    def __iter__(self):
+        rng = np.random.RandomState(self.seed)
+        bases = np.array(["A", "C", "G", "T"], dtype="<U1")
+        p_at = (1.0 - self.gc) / 2.0
+        p_gc = self.gc / 2.0
+        p = [p_at, p_gc, p_gc, p_at]
+        for start in range(0, self.n, self.batch_size):
+            bs = min(self.batch_size, self.n - start)
+            arr = rng.choice(4, size=(bs, self.seq_len), p=p)
+            seqs = ["".join(bases[row]) for row in arr]
+            yield seqs
+
 def _construct_motif_concept_dataloader_from_control(
     control_seq_df: pd.DataFrame,
     genome_fasta: str,
@@ -145,6 +165,43 @@ def _control_chrom_dl(self):
         )
         return chrom_iter
 
+    def add_synthetic_gc_content_concepts(self, gc_content_step=0.1):
+        """
+        Add a list of GC content concepts,
+        according to gc_content_step, example 0.1, GC content of each concept increases from 0.0 to 1.0 by the step
+        concept iter is basically the same as other add concept function, batch the generated input
+        """
+        step = float(gc_content_step)
+        if step <= 0 or step > 1:
+            raise ValueError("gc_content_step must be in (0, 1].")
+
+        # Include 1.0 endpoint (within floating tolerance).
+        gc_values = np.arange(0.0, 1.0 + 1e-9, step, dtype=float)
+
+        added: List[Concept] = []
+        for gc in gc_values:
+            gc = float(np.clip(gc, 0.0, 1.0))
+            concept_name = f"synthetic_gc_{gc:.2f}" + self.concept_name_suffix
+            seed = self.rng_seed + int(round(gc * 10000))
+            seq_iter = _SyntheticGCSeqIterator(
+                seq_len=self.input_window_length,
+                n=self.min_samples,
+                batch_size=self.batch_size,
+                gc=gc,
+                seed=seed,
+            )
+            concept = Concept(
+                id=self._reserve_id(),
+                name=concept_name,
+                data_iter=_PairedLoader(seq_iter, self._control_chrom_dl()),
+            )
+            self.concepts.append(concept)
+            added.append(concept)
+
+        self.metadata["synthetic_gc_content_step"] = step
+        self.metadata["synthetic_gc_content_values"] = gc_values.tolist()
+        return added
+
     def add_custom_motif_concepts(
         self, motif_table: str, control_regions: Optional[pd.DataFrame] = None, build_permute_control=True
     ) -> Union[List[Concept], List[Tuple[Concept]]]: