Fix the transofrmation

xiki-tempula · Nov 16, 2020 · d4f5ec1 · d4f5ec1
1 parent 787e2c5
commit d4f5ec1
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/alchemicalitp/entry.py b/alchemicalitp/entry.py
@@ -372,13 +372,14 @@ def to_str(self):
                                                                                                  self.th0B, self.cthB,
                                                                                                  self.comment)
         elif self.func == 5:
-            return
-            '{: <7} {: <7} {: <7} {: <6} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} ; {}'.format(
+            return '{: <7} {: <7} {: <7} {: <6} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} {: <10} ; {}'.format(
                 self.i, self.j,
                 self.k, self.func,
                 self.C0, self.C1, self.C2, self.C3,
                 self.C0B, self.C1B, self.C2B, self.C3B,
                 self.comment)
+        else:
+            raise NotImplementedError
     def __eq__(self, other):
         if (self.i == other.i) and (self.j == other.j) and (self.k == other.k) and \
                 (self.func == other.func):
@@ -395,8 +396,10 @@ def __eq__(self, other):
         else:
             return False
     def __repr__(self): # pragma: no cover
-        return str((self.i, self.j, self.k, self.func, self.th0, self.cth))
-
+        if self.func == 1:
+            return str((self.i, self.j, self.k, self.func, self.th0, self.cth))
+        elif self.func == 5:
+            return str((self.i, self.j, self.k, self.func, self.C0, self.C1, self.C2, self.C3))
     def update_idx(self, mapping, sort=True):
         i = mapping[self.i]
         j = mapping[self.j]

diff --git a/alchemicalitp/field.py b/alchemicalitp/field.py
@@ -1,6 +1,7 @@
 from .entry import Comment, Atomtype, Atom, Bond, Pair, Angle, Dihedral, Cmap
 import copy
 import numpy as np
+from collections import defaultdict
 
 class Field():
     def __init__(self, name):
@@ -365,8 +366,28 @@ def __init__(self):
 
     def sort(self):
         self.merge_comment()
-        # Make sure that it is sorted by func, phase and pn
-        self.content.sort(key = lambda x: (x.i, x.j, x.k, x.l, x.func, x.phase, x.pn))
+        def sort_149(dihedral_list):
+            # Sort dihedral type 1 4 9
+            # Make sure that it is sorted by func, phase and pn
+            # As multiple lines can exist for the same dihedral angle
+            dihedral_list.sort(key = lambda x: (x.i, x.j, x.k, x.l, x.func, x.phase, x.pn))
+        def sort(dihedral_list):
+            dihedral_list.sort(
+                key=lambda x: (x.i, x.j, x.k, x.l, x.func))
+
+        dihedral_dict = defaultdict(list)
+        for dihedral in self.content:
+            dihedral_dict[dihedral.func].append(dihedral)
+        new_dihedral_list = []
+        for func in sorted(list(dihedral_dict.keys())):
+            if int(func) in [1, 4, 9]:
+                sort_149(dihedral_dict[func])
+                new_dihedral_list.extend(dihedral_dict[func])
+            else:
+                sort(dihedral_dict[func])
+                new_dihedral_list.extend(dihedral_dict[func])
+
+        self.content = new_dihedral_list
 
     def remove_zero(self):
         new_content = []