Merge pull request #27 from choderalab/improve-pic50readout-repr

Improve PIC50Readout __str__ and __repr__
choderalab · Aug 31, 2023 · 94abc63 · 94abc63
2 parents aafd56b + 9c4d95d
commit 94abc63
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/mtenn/model.py b/mtenn/model.py
@@ -417,11 +417,20 @@ def __init__(self, substrate: Optional[float] = None, Km: Optional[float] = None
         """
         super(PIC50Readout, self).__init__()
 
+        self.substrate = substrate
+        self.Km = Km
+
         if substrate and Km:
             self.cp_val = 1 + substrate / Km
         else:
             self.cp_val = None
 
+    def __repr__(self):
+        return f"PIC50Readout(substrate={self.substrate}, Km={self.Km})"
+
+    def __str__(self):
+        return repr(self)
+
     def forward(self, delta_g):
         """
         Method to convert a predicted delta G value into a pIC50 value.