From 5f9093040247b1fde28451cb73e86492621c7f25 Mon Sep 17 00:00:00 2001
From: Erick Matsen <ematsen@gmail.com>
Date: Thu, 6 Jun 2024 10:08:05 -0700
Subject: [PATCH] cleanup and update

---
 netam/dnsm.py      | 2 --
 tests/test_dnsm.py | 8 ++++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/netam/dnsm.py b/netam/dnsm.py
index dce3e1be..6cff9e7a 100644
--- a/netam/dnsm.py
+++ b/netam/dnsm.py
@@ -186,8 +186,6 @@ def load_branch_lengths(self, in_csv_path):
         self.branch_lengths = pd.read_csv(in_csv_path)["branch_length"].values
 
     def update_neutral_aa_mut_probs(self):
-        print("consolidating neutral rates into substitution probabilities...")
-
         neutral_aa_mut_prob_l = []
 
         for nt_parent, mask, rates, branch_length, subs_probs in zip(
diff --git a/tests/test_dnsm.py b/tests/test_dnsm.py
index 624dfa73..de4e76c5 100644
--- a/tests/test_dnsm.py
+++ b/tests/test_dnsm.py
@@ -6,7 +6,8 @@
 from netam.framework import (
     crepe_exists,
     load_crepe,
-    load_and_add_shm_model_outputs_to_pcp_df,
+    load_pcp_df,
+    add_shm_model_outputs_to_pcp_df,
 )
 from netam.common import aa_idx_tensor_of_str_ambig, MAX_AMBIG_AA_IDX
 from netam.models import TransformerBinarySelectionModelWiggleAct
@@ -22,8 +23,11 @@ def test_aa_idx_tensor_of_str_ambig():
 
 @pytest.fixture
 def pcp_df():
-    df = load_and_add_shm_model_outputs_to_pcp_df(
+    df = load_pcp_df(
         "data/wyatt-10x-1p5m_pcp_2023-11-30_NI.first100.csv.gz",
+    )
+    df = add_shm_model_outputs_to_pcp_df(
+        df,
         "data/cnn_joi_sml-shmoof_small",
     )
     return df