From ea7365f29ff21d17b94d475853da1e11e06949e6 Mon Sep 17 00:00:00 2001
From: Jayoung Ryu <jayoung_ryu@g.harvard.edu>
Date: Wed, 6 Dec 2023 09:59:54 -0500
Subject: [PATCH] fix duplicate a0 subset error

---
 bean/preprocessing/data_class.py | 5 ++---
 tests/test_run.py                | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/bean/preprocessing/data_class.py b/bean/preprocessing/data_class.py
index 961fc92..48f4347 100644
--- a/bean/preprocessing/data_class.py
+++ b/bean/preprocessing/data_class.py
@@ -200,7 +200,8 @@ def __getitem__(self, guide_idx):
         ndata.X_masked = ndata.X_masked[:, :, guide_idx]
         ndata.X_control = ndata.X_control[:, :, guide_idx]
         ndata.repguide_mask = ndata.repguide_mask[:, guide_idx]
-        ndata.a0 = ndata.a0[guide_idx]
+        if hasattr(ndata, "a0") and self.a0 is not None:
+            ndata.a0 = ndata.a0[guide_idx]
         return ndata
 
     def transform_data(self, X, n_bins=None):
@@ -369,8 +370,6 @@ def __getitem__(self, guide_idx):
             ndata.allele_counts = ndata.allele_counts[:, :, guide_idx, :]
         if hasattr(ndata, "a0_allele"):
             ndata.a0_allele = ndata.a0_allele[guide_idx, :]
-        if hasattr(ndata, "a0") and self.a0 is not None:
-            ndata.a0 = ndata.a0[guide_idx]
         if hasattr(ndata, "pi_a0") and self.pi_a0 is not None:
             ndata.pi_a0 = ndata.pi_a0[guide_idx]
         if hasattr(ndata, "a0_bcmatch") and self.a0_bcmatch is not None:
diff --git a/tests/test_run.py b/tests/test_run.py
index baa0a80..00bca37 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -69,7 +69,7 @@ def test_run_tiling_with_wo_negctrl_noacc():
 
 @pytest.mark.order(18)
 def test_run_tiling_with_wo_negctrl_uniform():
-    cmd = "bean-run sorting tiling tests/data/tiling_mini_screen_annotated.h5ad -o tests/test_res/tiling/ --uniform-edit --allele-df-key allele_counts_spacer_0_19_A.G_translated_prop0.1_0.3 --control-guide-tag None  --repguide-mask None"
+    cmd = "bean-run sorting tiling tests/data/tiling_mini_screen_annotated.h5ad -o tests/test_res/tiling/ --uniform-edit --allele-df-key allele_counts_spacer_0_19_A.G_translated_prop0.1_0.3 --control-guide-tag None --repguide-mask None"
     try:
         subprocess.check_output(
             cmd,