From 484a89f82e4e0d58720997e35b99f2d357f842c6 Mon Sep 17 00:00:00 2001
From: Kyle Zhang <lzhang43@uw.edu>
Date: Tue, 30 Jun 2026 20:26:40 -0700
Subject: [PATCH 1/3] Create partition function attributes for RotamerEnsemble
 and dRotamerEnsemble classes.

---
 src/chilife/RotamerEnsemble.py  |  9 +++++++++
 src/chilife/dRotamerEnsemble.py | 12 +++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/chilife/RotamerEnsemble.py b/src/chilife/RotamerEnsemble.py
index 6d681f8..3a42935 100644
--- a/src/chilife/RotamerEnsemble.py
+++ b/src/chilife/RotamerEnsemble.py
@@ -1362,6 +1362,15 @@ def set_dihedral_sampling_sigmas(self, value):
             raise ValueError('`dihedral_sigmas` must be a scalar, an array the length of the `self.dihedral atoms` or '
                              'an array with the shape of (len(self.weights), len(self.dihedral_atoms))')
 
+    @property
+    def partition_function(self):
+        """The partition function of the rotamer ensemble after reweighting."""
+        # Calculate external energies
+        energies = self.energy_func(self)
+        self.rot_clash_energy = energies
+        # Calculate total weights (combining internal and external)
+        self.weights, self.partition = scoring.reweight_rotamers(energies, self.temp, self.weights)
+        return self.partition
 
 def assign_defaults(kwargs):
     """
diff --git a/src/chilife/dRotamerEnsemble.py b/src/chilife/dRotamerEnsemble.py
index 2ffe8e5..e8886bf 100644
--- a/src/chilife/dRotamerEnsemble.py
+++ b/src/chilife/dRotamerEnsemble.py
@@ -718,6 +718,16 @@ def copy(self):
             else:
                 new_copy.__dict__[item] = deepcopy(self.__dict__[item])
         return new_copy
+    
+    @property
+    def partition_function(self):
+        """The partition function of the rotamer ensemble after reweighting."""
+        # Calculate external energies
+        energies = self.energy_func(self)
+        self.rot_clash_energy = energies
+        # Calculate total weights (combining internal and external)
+        self.weights, self.partition = scoring.reweight_rotamers(energies, self.temp, self.weights)
+        return self.partition
 
 def proc_sites(sites):
     sites = sorted(sites)
@@ -792,4 +802,4 @@ def dassign_defaults(kwargs):
             f'{", ".join(key for key in kwargs if key not in defaults)}'
         )
 
-    return kwargs_filled.items()
\ No newline at end of file
+    return kwargs_filled.items()

From 5e648b4d7b629ec12ff1ae74f9869a869d5f7842 Mon Sep 17 00:00:00 2001
From: Kyle Zhang <lzhang43@uw.edu>
Date: Wed, 1 Jul 2026 10:00:55 -0700
Subject: [PATCH 2/3] Refactor partition_function to avoid energy recalculation

Pull partition function calculation from existing energy calculation.
---
 src/chilife/dRotamerEnsemble.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/chilife/dRotamerEnsemble.py b/src/chilife/dRotamerEnsemble.py
index e8886bf..9a2b419 100644
--- a/src/chilife/dRotamerEnsemble.py
+++ b/src/chilife/dRotamerEnsemble.py
@@ -722,11 +722,8 @@ def copy(self):
     @property
     def partition_function(self):
         """The partition function of the rotamer ensemble after reweighting."""
-        # Calculate external energies
-        energies = self.energy_func(self)
-        self.rot_clash_energy = energies
-        # Calculate total weights (combining internal and external)
-        self.weights, self.partition = scoring.reweight_rotamers(energies, self.temp, self.weights)
+        if self.partition == 1 and self.protein is not None:
+            self.evaluate()
         return self.partition
 
 def proc_sites(sites):

From 08844f227783e39af2f5e29f4fef42a76068a81a Mon Sep 17 00:00:00 2001
From: Kyle Zhang <lzhang43@uw.edu>
Date: Wed, 1 Jul 2026 10:03:39 -0700
Subject: [PATCH 3/3] Refactor partition function to avoid double calculations

Pull partition values directly from previous energy calculation to avoid double calculations.
---
 src/chilife/RotamerEnsemble.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/chilife/RotamerEnsemble.py b/src/chilife/RotamerEnsemble.py
index 3a42935..5e383cc 100644
--- a/src/chilife/RotamerEnsemble.py
+++ b/src/chilife/RotamerEnsemble.py
@@ -1365,11 +1365,8 @@ def set_dihedral_sampling_sigmas(self, value):
     @property
     def partition_function(self):
         """The partition function of the rotamer ensemble after reweighting."""
-        # Calculate external energies
-        energies = self.energy_func(self)
-        self.rot_clash_energy = energies
-        # Calculate total weights (combining internal and external)
-        self.weights, self.partition = scoring.reweight_rotamers(energies, self.temp, self.weights)
+        if self.partition == 1 and self.protein is not None:
+            self.evaluate ()
         return self.partition
 
 def assign_defaults(kwargs):