From cda9ca3efb7b0e52ba013c20aa3484ade67ecacd Mon Sep 17 00:00:00 2001
From: Samuel Musson <wppj21@ncc1.clients.dur.ac.uk>
Date: Mon, 15 May 2023 14:10:30 +0100
Subject: [PATCH] molearn bug fixes to avoid recreating dope_score_class in
 molearn_analysis

---
 src/molearn/analysis/analyser.py  |  4 ++++
 src/molearn/scoring/dope_score.py | 13 +++++++++----
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/molearn/analysis/analyser.py b/src/molearn/analysis/analyser.py
index 4ca560e..839f03f 100644
--- a/src/molearn/analysis/analyser.py
+++ b/src/molearn/analysis/analyser.py
@@ -466,3 +466,7 @@ def generate(self, crd):
 
         return s*self.stdval + self.meanval
 
+
+    def __getstate__(self):
+        return {key:value for key, value in dict(self.__dict__).items() if key not in ['dope_score_class', 'ramachandran_score_class']}
+
diff --git a/src/molearn/scoring/dope_score.py b/src/molearn/scoring/dope_score.py
index 21f6737..c7a665f 100644
--- a/src/molearn/scoring/dope_score.py
+++ b/src/molearn/scoring/dope_score.py
@@ -78,7 +78,12 @@ def get_dope(self, frame, refine=False):
                 if i+1<frame.shape[0]:
                     j.x, j.y, j.z = frame[self.fast_atom_order[i], :]
             self.fast_mdl.build(build_method='INTERNAL_COORDINATES', initialize_xyz=False)
-            
+            if refine == 'both':
+                with ShutUp():
+                    dope_unrefined = self.fast_fs.assess_dope()
+                    self.cg.optimize(self.fast_fs, max_iterations=50)
+                    dope_refined = self.fast_fs.assess_dope()
+                    return dope_unrefined, dope_refined
             with ShutUp():
                 if refine:
                     self.cg.optimize(self.fast_fs, max_iterations=50)
@@ -155,7 +160,7 @@ def __init__(self, mol, processes=-1, **kwargs):
             processes = min(processes, cpu_count())
         else:
             processes = cpu_count()
-            
+        self.processes = processes
         self.mol = deepcopy(mol)
         score = DOPE_Score
         ctx = get_context('spawn')
@@ -166,11 +171,11 @@ def __init__(self, mol, processes=-1, **kwargs):
         self.process_function = process_dope
 
     def __reduce__(self):
-        return (self.__class__, (self.mol,))
+        return (self.__class__, (self.mol, processes=self.processes))
 
     def get_score(self, coords, **kwargs):
         '''
         :param np.array coords: # shape (N, 3) numpy array
         '''
         #is copy necessary?
-        return self.pool.apply_async(self.process_function, (coords.copy(), kwargs))
\ No newline at end of file
+        return self.pool.apply_async(self.process_function, (coords.copy(), kwargs))