From 57b9b0a2bb1be937476b49a99a4622fb89444dc7 Mon Sep 17 00:00:00 2001
From: hsulab <ahcigar@foxmail.com>
Date: Fri, 19 Apr 2024 20:27:25 +0100
Subject: [PATCH] update harmonic

---
 src/gdpx/bias/__init__.py          |  3 +-
 src/gdpx/bias/harmonic/__init__.py |  3 +-
 src/gdpx/bias/harmonic/distance.py | 99 ++++++++++++++++++++++++++++++
 src/gdpx/bias/timeio.py            | 70 +++++++++++++++++++++
 4 files changed, 173 insertions(+), 2 deletions(-)
 create mode 100644 src/gdpx/bias/harmonic/distance.py
 create mode 100644 src/gdpx/bias/timeio.py

diff --git a/src/gdpx/bias/__init__.py b/src/gdpx/bias/__init__.py
index fdc2695e..01670646 100644
--- a/src/gdpx/bias/__init__.py
+++ b/src/gdpx/bias/__init__.py
@@ -24,7 +24,8 @@
 # from .harmonic import HarmonicBias
 # bias_register.register("harmonic")(HarmonicBias)
 
-from .harmonic import PlaneHarmonicCalculator
+from .harmonic import DistanceHarmonicCalculator, PlaneHarmonicCalculator
+bias_register.register("distance_harmonic")(DistanceHarmonicCalculator)
 bias_register.register("plane_harmonic")(PlaneHarmonicCalculator)
 
 # - gaussian
diff --git a/src/gdpx/bias/harmonic/__init__.py b/src/gdpx/bias/harmonic/__init__.py
index 3997e783..5842e891 100644
--- a/src/gdpx/bias/harmonic/__init__.py
+++ b/src/gdpx/bias/harmonic/__init__.py
@@ -2,10 +2,11 @@
 # -*- coding: utf-8 -*-
 
 
+from .distance import DistanceHarmonicCalculator
 from .plane import PlaneHarmonicCalculator
 
 
-__all__ = ["PlaneHarmonicCalculator"]
+__all__ = ["DistanceHarmonicCalculator", "PlaneHarmonicCalculator"]
 
 
 if __name__ == "__main__":
diff --git a/src/gdpx/bias/harmonic/distance.py b/src/gdpx/bias/harmonic/distance.py
new file mode 100644
index 00000000..32f7ecd0
--- /dev/null
+++ b/src/gdpx/bias/harmonic/distance.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+
+import copy
+from typing import Optional, Tuple, List
+
+import numpy as np
+
+from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.geometry import find_mic
+
+from ..timeio import TimeIOCalculator
+
+
+def compute_distance(cell, positions, pbc: bool = True):
+    """"""
+    # compute colvar
+    assert positions.shape[0] == 2
+    vec, dis = find_mic(positions[0] - positions[1], cell=cell)
+
+    return vec, dis
+
+
+def compute_distance_harmonic_energy_and_forces(
+    vec, dis: float, center: float, kspring: float
+):
+    """"""
+    # compute energy
+    energy = np.sum(0.5 * kspring * (dis - center) ** 2)
+
+    # compute forces
+    forces = np.zeros((2, 3))
+    frc_i = -kspring * vec / dis
+    forces[0] += frc_i
+    forces[1] += -frc_i
+
+    return energy, forces
+
+
+class DistanceHarmonicCalculator(TimeIOCalculator):
+
+    implemented_properties = ["energy", "free_energy", "forces"]
+
+    def __init__(
+        self, group: List[int], center: float, kspring: float = 0.1, *args, **kwargs
+    ):
+        """"""
+        super().__init__(*args, **kwargs)
+
+        num_group_atoms = len(group)
+        assert num_group_atoms == 2
+        self.group = group
+
+        self.center = center
+
+        self.kspring = kspring
+
+        return
+
+    def _icalculate(self, atoms, properties, system_changes) -> Tuple[dict, list]:
+        """"""
+        vec, dis = compute_distance(atoms.cell, atoms.positions[self.group], pbc=True)
+
+        energy, ext_forces = compute_distance_harmonic_energy_and_forces(
+            vec, dis, self.center, self.kspring
+        )
+        forces = np.zeros(atoms.positions.shape)
+        forces[self.group] = ext_forces
+
+        results = {}
+        results["energy"] = energy
+        results["free_energy"] = energy
+        results["forces"] = forces
+
+        step_info = (self.num_steps, dis, energy)
+
+        return results, step_info
+
+    def _write_first_step(self):
+        """"""
+        content = "# {:>10s}  {:>12s}  {:>12s}\n".format("step", "distance", "energy")
+        with open(self.log_fpath, "w") as fopen:
+            fopen.write(content)
+
+        return
+    
+    def _write_step(self):
+        """"""
+        content = "{:>12d}  {:>12.4f}  {:>12.4f}\n".format(*self.step_info)
+        with open(self.log_fpath, "a") as fopen:
+            fopen.write(content)
+
+        return
+
+
+if __name__ == "__main__":
+    ...
diff --git a/src/gdpx/bias/timeio.py b/src/gdpx/bias/timeio.py
new file mode 100644
index 00000000..bba08c45
--- /dev/null
+++ b/src/gdpx/bias/timeio.py
@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+
+import pathlib
+
+from ase.calculators.calculator import Calculator
+
+
+class TimeIOCalculator(Calculator):
+
+    def __init__(self, pace: int=1, *args, **kwargs):
+        """"""
+        super().__init__(*args, **kwargs)
+
+        self.pace = pace
+
+        self._num_steps = 0
+
+        return
+
+    @property
+    def num_steps(self) -> int:
+        """Finished steps that match the host driver e.g. MD."""
+
+        return self._num_steps
+
+    @property
+    def log_fpath(self):
+        """"""
+
+        return pathlib.Path(self.directory) / "calc.log"
+
+    def calculate(
+        self,
+        atoms=None,
+        properties=["energy"],
+        system_changes=["positions", "numbers", "cell"],
+    ):
+        super().calculate(atoms, properties, system_changes)
+
+        if self.num_steps == 0:
+            self._write_first_step()
+
+        self.results, self.step_info = self._icalculate(atoms, properties, system_changes)
+
+        if self.num_steps % self.pace == 0:
+            self._write_step()
+
+        self._num_steps += 1
+
+        return
+
+    def _icalculate(self, atoms, properties, system_changes):
+        """"""
+
+        raise NotImplementedError()
+
+    def _write_first_step(self):
+        """"""
+
+        raise NotImplementedError()
+
+    def _write_step(self):
+        """"""
+
+        raise NotImplementedError()
+
+if __name__ == "__main__":
+    ...