From 6af98ad800e25d4728cb75ddd91fbb81c5b51592 Mon Sep 17 00:00:00 2001
From: Louis Vanduyfhuys <louis.vanduyfhuys@ugent.be>
Date: Fri, 10 Jul 2020 14:53:26 +0200
Subject: [PATCH] Compatibility patch for Yaff 1.6.0 regarding pickling of the
 perturbation trajectories

---
 quickff/valence.py | 39 +++++++++++++++++++++++++++++++++++++++
 setup.py           |  2 +-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/quickff/valence.py b/quickff/valence.py
index a55b902..b7856f4 100644
--- a/quickff/valence.py
+++ b/quickff/valence.py
@@ -61,6 +61,45 @@ def __init__(self, index, basename, kind, ics, tasks, units,master=None, slaves=
         self.slaves = slaves
         self.diag_term_indexes=diag_term_indexes
 
+    def __getstate__(self):
+        state = {
+            'index': self.index,
+            'basename': self.basename,
+            'kind': self.kind,
+            'ics': [(ic.kind, ic.index_pairs) for ic in self.ics],
+            'tasks': self.tasks,
+            'units': self.units,
+            'master': self.master,
+            'slaves': self.slaves,
+            'diag_term_indexes': self.diag_term_indexes,
+        }
+        return state
+    
+    def __setstate__(self, state):
+        self.index = state['index']
+        self.basename = state['basename']
+        self.kind = state['kind']
+        self.ics = []
+        for (kind, index_pairs) in state['ics']:
+            if kind==0:
+                ic = Bond(*index_pairs[0])
+            elif kind==1:
+                ic = BendCos(index_pairs[0][1],index_pairs[1][0],index_pairs[1][1])
+            elif kind==2:
+                ic = BendAngle(index_pairs[0][1],index_pairs[1][0],index_pairs[1][1])
+            elif kind==10:
+                ic = OopDist(index_pairs[0][0],index_pairs[1][0],index_pairs[2][0],index_pairs[2][1])
+            elif kind==11:
+                ic = SqOopDist(index_pairs[0][0],index_pairs[1][0],index_pairs[2][0],index_pairs[2][1])
+            else:
+                raise NotImplementedError('Internal coordinate of kind %i not implemented yet in __setstate__ routine of Term class.' %kind)
+            self.ics.append(ic)
+        self.tasks = state['tasks']
+        self.units = state['units']
+        self.master = state['master']
+        self.slaves = state['slaves']
+        self.diag_term_indexes = state['diag_term_indexes']
+
     def is_master(self):
         return self.master==self.index
 
diff --git a/setup.py b/setup.py
index b2fefda..4e93360 100644
--- a/setup.py
+++ b/setup.py
@@ -64,7 +64,7 @@
 
 setup(
     name='QuickFF',
-    version='2.2.6',
+    version='2.2.6_bugfixpickle',
     description='Python library to quickly derive force fields from ab initio training data.',
     author='Louis Vanduyfhuys',
     author_email='Louis.Vanduyfhuys@UGent.be',