Merge branch 'optimize-training' of github.com:PMBio/deeprvat into op…

…timize-training
PMBio · Dec 14, 2023 · e0e67cd · e0e67cd
2 parents 2350f0c + b985e77
commit e0e67cd
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/deeprvat/deeprvat/train.py b/deeprvat/deeprvat/train.py
@@ -516,6 +516,18 @@ def __getitem__(self, index):
     #     slice_indices = np.nonzero(np.isin(arange, indices))
     #     return z[slice_indices]
 
+    def index_input_tensor_zarr(self, pheno: str, indices: np.ndarray):
+        # IMPORTANT!!! Never call this function after self.subset_samples()
+
+        x = self.data[pheno]["input_tensor_zarr"]
+        first_idx = indices[0]
+        last_idx = indices[-1]
+        slice_ = slice(first_idx, last_idx + 1)
+        arange = np.arange(first_idx, last_idx + 1)
+        z = x[slice_]
+        slice_indices = np.nonzero(np.isin(arange, indices))
+        return z[slice_indices]
+
 
 class MultiphenoBaggingData(pl.LightningDataModule):
     """