From 9c4d95d11a342302aa1167d204e2c77867f74f6b Mon Sep 17 00:00:00 2001
From: kaminow <benjamin.kaminow@choderalab.org>
Date: Thu, 31 Aug 2023 10:56:31 -0400
Subject: [PATCH] Add better repr for PIC50Readout so that substrate/Km values
 are shown.

---
 mtenn/model.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mtenn/model.py b/mtenn/model.py
index b23be30..ad8cef5 100644
--- 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.