From b56a9f48de5d109d01a8dc60556881fc2f8d459a Mon Sep 17 00:00:00 2001
From: kaminow <benjamin.kaminow@choderalab.org>
Date: Fri, 13 Sep 2024 11:31:06 -0400
Subject: [PATCH] Adjust pKiReadout math for numerical stability.

---
 mtenn/readout.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/mtenn/readout.py b/mtenn/readout.py
index 5db20cb..82f097d 100644
--- a/mtenn/readout.py
+++ b/mtenn/readout.py
@@ -148,6 +148,10 @@ class PKiReadout(Readout):
         \mathrm{pK_i} &= \mathrm{-log10(K_i)}
 
         \mathrm{pK_i} &= \mathrm{-log10(exp(\Delta G))}
+
+        \mathrm{pK_i} &= \\frac{\mathrm{-ln(exp(\Delta G))}}{\mathrm{ln(10)}}
+
+        \mathrm{pK_i} &= \\frac{\mathrm{-\Delta G}}{\mathrm{ln(10)}}
     """
 
     def __repr__(self):
@@ -168,6 +172,6 @@ def forward(self, delta_g):
         torch.Tensor
             Calculated :math:`\mathrm{pK_i}` value.
         """
-        pki = -torch.log10(torch.exp(delta_g))
+        pki = -delta_g / torch.log(torch.tensor(10, dtype=delta_g.dtype))
 
         return pki