Revert v3 update for netam.multihit (#149)

Author: willdumm

netam/multihit.py

@@ -8,6 +8,8 @@
 apply the correction to existing codon probability predictions, we multiply the
 probability of each child codon by the correction factor for its hit class, then
 renormalize. The correction factor for hit class 0 is fixed at 1.
+
+NOTE: Unlike the rest of netam, this module is not updated to the v3 data format, since Thrifty isn't either.
 """
 
 import torch
@@ -436,10 +438,10 @@ def to_crepe(self):
 def hit_class_dataset_from_pcp_df(
     pcp_df: pd.DataFrame, branch_length_multiplier: int = 1.0
 ) -> HitClassDataset:
-    nt_parents = pcp_df["parent_heavy"].reset_index(drop=True)
-    nt_children = pcp_df["child_heavy"].reset_index(drop=True)
-    nt_rates = pcp_df["nt_rates_heavy"].reset_index(drop=True)
-    nt_csps = pcp_df["nt_csps_heavy"].reset_index(drop=True)
+    nt_parents = pcp_df["parent"].reset_index(drop=True)
+    nt_children = pcp_df["child"].reset_index(drop=True)
+    nt_rates = pcp_df["nt_rates"].reset_index(drop=True)
+    nt_csps = pcp_df["nt_csps"].reset_index(drop=True)
 
     return HitClassDataset(
         nt_parents,
@@ -455,10 +457,10 @@ def train_test_datasets_of_pcp_df(
 ) -> Tuple[HitClassDataset, HitClassDataset]:
     """Splits a pcp_df prepared by `prepare_pcp_df` into a training and testing
     HitClassDataset."""
-    nt_parents = pcp_df["parent_heavy"].reset_index(drop=True)
-    nt_children = pcp_df["child_heavy"].reset_index(drop=True)
-    nt_rates = pcp_df["nt_rates_heavy"].reset_index(drop=True)
-    nt_csps = pcp_df["nt_csps_heavy"].reset_index(drop=True)
+    nt_parents = pcp_df["parent"].reset_index(drop=True)
+    nt_children = pcp_df["child"].reset_index(drop=True)
+    nt_rates = pcp_df["nt_rates"].reset_index(drop=True)
+    nt_csps = pcp_df["nt_csps"].reset_index(drop=True)
 
     train_len = int(train_frac * len(nt_parents))
     train_parents, val_parents = nt_parents[:train_len], nt_parents[train_len:]
@@ -496,18 +498,12 @@ def prepare_pcp_df(
 
     Returns the modified dataframe, which is the input dataframe modified in-place.
     """
-    pcp_df["parent_heavy"] = _trim_to_codon_boundary_and_max_len(
-        pcp_df["parent_heavy"], site_count
-    )
-    pcp_df["child_heavy"] = _trim_to_codon_boundary_and_max_len(
-        pcp_df["child_heavy"], site_count
-    )
-    pcp_df = pcp_df[pcp_df["parent_heavy"] != pcp_df["child_heavy"]].reset_index(
-        drop=True
-    )
+    pcp_df["parent"] = _trim_to_codon_boundary_and_max_len(pcp_df["parent"], site_count)
+    pcp_df["child"] = _trim_to_codon_boundary_and_max_len(pcp_df["child"], site_count)
+    pcp_df = pcp_df[pcp_df["parent"] != pcp_df["child"]].reset_index(drop=True)
     ratess, cspss = framework.trimmed_shm_model_outputs_of_crepe(
-        crepe, pcp_df["parent_heavy"]
+        crepe, pcp_df["parent"]
     )
-    pcp_df["nt_rates_heavy"] = ratess
-    pcp_df["nt_csps_heavy"] = cspss
+    pcp_df["nt_rates"] = ratess
+    pcp_df["nt_csps"] = cspss
     return pcp_df

tests/test_multihit.py

@@ -87,6 +87,14 @@
 @pytest.fixture
 def mini_multihit_train_val_datasets():
     df = pd.read_csv("data/wyatt-10x-1p5m_pcp_2023-11-30_NI.first100.csv.gz")
+    # Rename _heavy columns to drop _heavy and drop all _light columns
+    # (multihit training is Thrifty territory, and not yet v3 format)
+    df = df.rename(
+        columns={
+            col: col[: -len("_heavy")] for col in df.columns if col.endswith("_heavy")
+        }
+    )
+    df = df.drop(columns=[col for col in df.columns if col.endswith("_light")])
     crepe = pretrained.load("ThriftyHumV0.2-45")
     df = multihit.prepare_pcp_df(df, crepe, 500)
     return multihit.train_test_datasets_of_pcp_df(df)