From 2dfb08fc983293590dda37ad62337ee5bf0118f3 Mon Sep 17 00:00:00 2001
From: Sander Vandenhaute <svdenhau@gmail.com>
Date: Tue, 18 Jun 2024 05:12:08 -0400
Subject: [PATCH] support barostat mode=nst

---
 psiflow/sampling/sampling.py | 12 ++++++++++--
 psiflow/sampling/walker.py   |  3 +++
 psiflow/tools/server.py      | 20 ++++++++++++++++++++
 tests/test_sampling.py       | 19 +++++++++++++++++++
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py
index abd844b..c3e41d4 100644
--- a/psiflow/sampling/sampling.py
+++ b/psiflow/sampling/sampling.py
@@ -121,13 +121,21 @@ def setup_motion(walker: Walker, fix_com: bool) -> ET.Element:
         else:
             dynamics.append(thermostat)
     elif walker.npt:
-        dynamics = ET.Element("dynamics", mode="npt")
+        if walker.volume_constrained:
+            mode = "nst"
+        else:
+            mode = "npt"
+        dynamics = ET.Element("dynamics", mode=mode)
         dynamics.append(timestep_element)
         if walker.pimd:
             dynamics.append(thermostat_pimd)
         else:
             dynamics.append(thermostat)
-        barostat = ET.Element("barostat", mode="flexible")
+        if walker.volume_constrained:
+            mode = "anisotropic"
+        else:
+            mode = "flexible"
+        barostat = ET.Element("barostat", mode=mode)
         tau = ET.Element("tau", units="femtosecond")
         tau.text = "200"
         barostat.append(tau)
diff --git a/psiflow/sampling/walker.py b/psiflow/sampling/walker.py
index 7359d01..6a62e80 100644
--- a/psiflow/sampling/walker.py
+++ b/psiflow/sampling/walker.py
@@ -49,6 +49,7 @@ class Walker:
     state: Union[Geometry, AppFuture, None]
     temperature: Optional[float]
     pressure: Optional[float]
+    volume_constrained: bool
     nbeads: int
     timestep: float
     coupling: Optional[Coupling]
@@ -62,6 +63,7 @@ def __init__(
         state: Union[Geometry, AppFuture, None] = None,
         temperature: Optional[float] = 300,
         pressure: Optional[float] = None,
+        volume_constrained: bool = False,
         nbeads: int = 1,
         timestep: float = 0.5,
         metadynamics: Optional[Metadynamics] = None,
@@ -80,6 +82,7 @@ def __init__(
 
         self.temperature = temperature
         self.pressure = pressure
+        self.volume_constrained = volume_constrained
         self.nbeads = nbeads
         self.timestep = timestep
 
diff --git a/psiflow/tools/server.py b/psiflow/tools/server.py
index b6bf4a3..ce66f25 100644
--- a/psiflow/tools/server.py
+++ b/psiflow/tools/server.py
@@ -318,6 +318,25 @@ def insert_data_start(input_xml, data_start):
             initialize.text = "start_INDEX.xyz"
 
 
+def anisotropic_barostat_h0(input_xml, data_start):
+    path = "system_template/template/system/motion/dynamics/barostat"
+    barostat = input_xml.find(path)
+    if barostat is not None and (barostat.attrib["mode"] == "anisotropic"):
+        h0 = ET.SubElement(barostat, "h0", shape="(3, 3)", units="angstrom")
+        cell = np.array(data_start[0].cell)
+        cell = cell.flatten(order="F")
+        h0.text = " [ {} ] ".format(" , ".join([str(a) for a in cell]))
+
+        path = "system_template/template/system/ensemble"
+        ensemble = input_xml.find(path)
+        assert ensemble is not None
+        pressure = ensemble.find("pressure")
+        if pressure is not None:
+            ensemble.remove(pressure)
+            stress = ET.SubElement(ensemble, "stress", units="megapascal")
+            stress.text = " [ PRESSURE, 0, 0, 0, PRESSURE, 0, 0, 0, PRESSURE ] "
+
+
 def start(args):
     data_start = read(args.start_xyz, index=":")
     assert len(data_start) == args.nwalkers
@@ -333,6 +352,7 @@ def start(args):
 
     insert_data_start(input_xml, data_start)
     insert_addresses(input_xml)
+    anisotropic_barostat_h0(input_xml, data_start)
     with open("input.xml", "wb") as f:
         f.write(ET.tostring(input_xml, encoding="utf-8"))
 
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index 413ea88..327c700 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -300,6 +300,25 @@ def test_npt(dataset):
         output.state.result().cell,
     )
 
+    walker = Walker(
+        dataset[0],
+        einstein,
+        temperature=600,
+        pressure=0,
+        volume_constrained=True,
+    )
+    output = sample([walker], steps=30)[0]
+    # cell should have changed, but not the volume
+    assert not np.allclose(
+        walker.start.result().cell,
+        output.state.result().cell,
+    )
+    # volumea actually does change with barostat mode='nst' (?)
+    # assert np.allclose(
+    #     np.linalg.det(walker.start.result().cell),
+    #     np.linalg.det(output.state.result().cell),
+    # )
+
 
 def test_reset(dataset):
     einstein = EinsteinCrystal(dataset[0], force_constant=0.1)