Merge Claudiu Craciun and Alan Kuznicki's semester project work into main

.gitignore

@@ -3,12 +3,14 @@ generated/**
 wandb/**
 */.DS_Store
 .idea
+*.png
 
 /data
 /checkpoint
 /logs
 /plots
 /archive/imgs/
+/tmp
 
 
 # Created by https://www.gitignore.io/api/macos,python,R,linux,vim,emacs,visualstudiocode,intellij

src/asap/__init__.py

@@ -1,2 +1,5 @@
 from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset 
-from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac
+
+
+from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model, train_multiheaded_model_progressively, extend_multiheaded_model_progressively
+

src/asap/dataloader/base.py

@@ -31,10 +31,12 @@ def __init__(
     ):
         super().__init__()
         self.genome = genome
-        if signal_files is not None:
-            self.signal_files = [signal_files]
-        else:
+        if signal_files is None:
             self.signal_files = None
+        elif type(signal_files) is list:
+            self.signal_files = signal_files
+        else:
+            self.signal_files = [signal_files]
 
         # if we randomly shift to augment the data
         self.random_shift = random_shift

src/asap/models/convnext_dcnn.py

@@ -20,7 +20,8 @@
     'kernel2': 1,
     'dropout': 0.3,
     'final_dropout': 0.05,
-    'use_map': False
+    'use_map': False,
+    'num_heads': 1,
 }
 
 
@@ -50,7 +51,8 @@ def __init__(self, **kwargs) -> None:
                       kernel1=config['kernel1'],
                       kernel2=config['kernel2'],
                       dropout=config['dropout'],
-                      final_dropout=config['final_dropout'])
+                      final_dropout=config['final_dropout'],
+                      num_heads=config['num_heads'])
 
     def forward(self, x:torch.Tensor, return_unmap=False) -> torch.Tensor:
         x = torch.transpose(x, dim0=-1, dim1=-2)
@@ -70,7 +72,7 @@ class BasenjiCoreBlock(nn.Module):
     def __init__(self, nr_tracks: int, window: int, filters_in,
                   nr_res_blocks: int = 11, rate_mult: float = 1.5, bin_size: int = 100, filters1: int = 128,
                   filters3: int = None, kernel1: int = 3, kernel2: int = 1, dropout: float = 0.3,
-                  final_dropout: float = 0.05):
+                  final_dropout: float = 0.05, num_heads=1):
         super().__init__()
         if not filters3:
             filters3 = window
@@ -93,10 +95,11 @@ def __init__(self, nr_tracks: int, window: int, filters_in,
         pool_size = bin_size // 2
         self.pool = nn.AvgPool1d(kernel_size=pool_size, stride=pool_size)
 
-        self.linear_out = nn.Linear(
-            in_features=filters3,
-            out_features=nr_tracks,
-        )
+        print("creating output layer with tracks:", num_heads)
+        self.heads = nn.ModuleList([
+            nn.Linear(filters3, nr_tracks)
+            for _ in range(num_heads)
+        ])
         self.activation = nn.Softplus()
 
     def forward(self, x:torch.Tensor) -> torch.Tensor:
@@ -115,9 +118,12 @@ def forward(self, x:torch.Tensor) -> torch.Tensor:
         x = self.final_dropout(x)
         x = self.pool(x)
         x = torch.transpose(x, -2, -1)
-        x = self.linear_out(x)
-        x = self.activation(x)
-        return x
+        outputs = []
+        for head in self.heads:
+            out = head(x)
+            out = self.activation(out)
+            outputs.append(out)
+        return outputs
 
 
 class ConvBlock(nn.Module):

src/asap/snv/predict.py

@@ -20,7 +20,7 @@ def _idx_to_ohe(idx: np.ndarray) -> np.ndarray:
     one_hot_encoded = eye[idx]
     return one_hot_encoded
 
-def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device) -> None:
+def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device, target_head=0) -> None:
     # Ensure output columns exist
     output_columns = ['signal_true', 'signal_pred_ref', 'signal_pred_alt']
     for col in output_columns:
@@ -56,8 +56,8 @@ def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: n
             x = _idx_to_ohe(x).astype(np.float32)
             x_var = _idx_to_ohe(x_var).astype(np.float32)
 
-            ref_pred = model(torch.from_numpy(x).to(device))
-            var_pred = model(torch.from_numpy(x_var).to(device))
+            ref_pred = model(torch.from_numpy(x).to(device))[target_head]
+            var_pred = model(torch.from_numpy(x_var).to(device))[target_head]
             for (i, _), true, ref, var in zip(chr_df_i.iterrows(), y, ref_pred.cpu().detach().numpy(), var_pred.cpu().detach().numpy()):
                 snv.loc[i, 'signal_true'] = json.dumps(true.flatten().tolist())
                 snv.loc[i, 'signal_pred_ref'] = json.dumps(ref.tolist())