From 5ae2f2c83b71b2a9a209c633306b99db15dc3077 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Salom=C3=A9=20Rieder?= <salome.rieder@solnet.ch>
Date: Wed, 21 Apr 2021 16:47:53 +0200
Subject: [PATCH 1/2] add NUM to MULTIBATH block

add NUM to MULTIBATH block for Nose-Hoover chain thermostat
---
 pygromos/files/blocks/imd_blocks.py | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/pygromos/files/blocks/imd_blocks.py b/pygromos/files/blocks/imd_blocks.py
index b3868dfd..4ea5691f 100644
--- a/pygromos/files/blocks/imd_blocks.py
+++ b/pygromos/files/blocks/imd_blocks.py
@@ -476,6 +476,8 @@ class MULTIBATH(_generic_imd_block):
     Attributes
     ----------
     ALGORITHM:  int
+    NUM: int, optional
+        Mumber of chains in Nosé Hoover chains scheme [only specify when needed]
     NBATHS: int
         Number of temperature baths
     TEMP0:  List[float]
@@ -499,6 +501,7 @@ class MULTIBATH(_generic_imd_block):
     name: str = "MULTIBATH"
 
     ALGORITHM: int
+    NUM: int
     NBATHS: int
     TEMP0: List[float]
     TAU: List[float]
@@ -507,14 +510,15 @@ class MULTIBATH(_generic_imd_block):
     COMBATH: List[int]
     IRBATH: List[int]
 
-    _order: List[List[str]] = [[["ALGORITHM"], ["NBATHS"], ["TEMP0(1 ... NBATHS)", "TAU(1 ... NBATHS)"],
+    _order: List[List[str]] = [[["ALGORITHM"], ["NUM"], ["NBATHS"], ["TEMP0(1 ... NBATHS)", "TAU(1 ... NBATHS)"],
                                 ["DOFSET"], ["LAST(1 ... DOFSET)", "COMBATH(1 ... DOFSET)", "IRBATH(1 ... DOFSET)"]]]
 
     def __init__(self, ALGORITHM: int, NBATHS: int, TEMP0: List[float], TAU: List[float], DOFSET: int, LAST: List[int], COMBATH: List[int],
-                 IRBATH: List[int]):
+                 IRBATH: List[int], NUM: int = None):
 
         super().__init__(used=True)
         self.ALGORITHM = ALGORITHM
+        self.NUM = NUM
         self.NBATHS = NBATHS
         self.TEMP0 = TEMP0
         self.TAU = TAU
@@ -536,7 +540,7 @@ def __init__(self, ALGORITHM: int, NBATHS: int, TEMP0: List[float], TAU: List[fl
         if not len(LAST) == len(IRBATH):
             warnings.warn("Warning in MULTIBATH block. There must be the same number of IRBATH and LAST parameters")
 
-    def adapt_multibath(self, last_atoms_bath: Dict[int, int], algorithm: int = None, T: (float, List[float]) = None, TAU: float = None) -> None:
+    def adapt_multibath(self, last_atoms_bath: Dict[int, int], algorithm: int = None, num: int = None, T: (float, List[float]) = None, TAU: float = None) -> None:
         """ adapt_multibath
                 This function is adding each atom set into a single multibath.
                 #TODO implementation not correct with com_bath and irbath! Works for super simple cases though
@@ -577,6 +581,9 @@ def adapt_multibath(self, last_atoms_bath: Dict[int, int], algorithm: int = None
         else:
             self.ALGORITHM = algorithm
 
+        if (num != None):
+            self.NUM = num
+
         # TODO implementation not correct with com_bath and irbath! Works for super simple cases though
         print(last_atoms_bath)
         print("\n")
@@ -598,14 +605,19 @@ def block_to_string(self) -> str:
         result += str(self.name) + "\n"
         result += "# " + "\t".join(self._order[0][0]) + "\n"
         result += "  " + str(self.ALGORITHM) + "\n"
-        result += "# " + "\t".join(self._order[0][1]) + "\n"
-        result += "  " + str(self.NBATHS) + "\n"
+
+        if(self.ALGORITHM == "2"):
+            result += "# " + "\t".join(self._order[0][1]) + "\n"
+            result += "  " + str(self.NUM) + "\n"
+
         result += "# " + "\t".join(self._order[0][2]) + "\n"
+        result += "  " + str(self.NBATHS) + "\n"
+        result += "# " + "\t".join(self._order[0][3]) + "\n"
         for index in range(len(self.TEMP0)):
             result += "  " + str(self.TEMP0[index]) + "\t" + str(self.TAU[index]) + "\n"
-        result += "# " + "\t".join(map(str, self._order[0][3])) + "\n"
-        result += "  " + str(self.DOFSET) + "\n"
         result += "# " + "\t".join(map(str, self._order[0][4])) + "\n"
+        result += "  " + str(self.DOFSET) + "\n"
+        result += "# " + "\t".join(map(str, self._order[0][5])) + "\n"
         for index in range(len(self.LAST)):
             result += "  " + str(self.LAST[index]) + "\t\t" + str(self.COMBATH[index]) + "\t" + str(self.IRBATH[index]) + "\n"
         result += "END\n"

From d61f30b62c605e35e8b46e08f1ace6c16e708466 Mon Sep 17 00:00:00 2001
From: B_Ries <benjamin-ries@outlook.com>
Date: Thu, 22 Apr 2021 11:45:06 +0200
Subject: [PATCH 2/2] Update imd_blocks.py

just a minor save
---
 pygromos/files/blocks/imd_blocks.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pygromos/files/blocks/imd_blocks.py b/pygromos/files/blocks/imd_blocks.py
index 4ea5691f..95d24855 100644
--- a/pygromos/files/blocks/imd_blocks.py
+++ b/pygromos/files/blocks/imd_blocks.py
@@ -607,6 +607,9 @@ def block_to_string(self) -> str:
         result += "  " + str(self.ALGORITHM) + "\n"
 
         if(self.ALGORITHM == "2"):
+            if(self.NUM is None):
+                raise Exception("You need to specify the NUM parameter for MULTIBATH if ALGORITHM is 2!")
+               
             result += "# " + "\t".join(self._order[0][1]) + "\n"
             result += "  " + str(self.NUM) + "\n"