add save_cav_trainer option

Author: yztxwd

tpcav/cavs.py

@@ -604,6 +604,7 @@ def run_tpcav(
     num_pc: Union[str,int]='full',
     p=1, 
     max_pending_jobs=4,
+    save_cav_trainer=True
 ):
     """
     One-stop function to compute CAVs on motif concepts and bed concepts, compute AUC of motif concept f-scores after correction
@@ -696,6 +697,8 @@ def run_tpcav(
                                         num_samples_for_cav,
                                         output_dir=str(output_path / f"cavs_{nm}_motifs/"),
                                         num_processes=p, max_pending=max_pending_jobs)
+        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:
         bed_cav_trainer = CavTrainer(tpcav_model, penalty="l2")
@@ -708,6 +711,8 @@ def run_tpcav(
             output_dir=str(output_path / f"cavs_bed_concepts/"),
             num_processes=p,
         )
+        if save_cav_trainer:
+            torch.save(bed_cav_trainer, str(output_path / f"cavs_bed_concepts/cav_trainer.pt"))
     else:
         bed_cav_trainer = None