From 20c88cb72ac59ff06805534a3acc1bd41dc2eec0 Mon Sep 17 00:00:00 2001
From: OMalenfantThuot <malenfantthuotolivier@gmail.com>
Date: Wed, 4 Nov 2020 16:55:50 -0500
Subject: [PATCH] ase calculator with automatic differentiation capabilities -
 first version

---
 .../calculators/ase_calculators/__init__.py   |  1 +
 .../ase_calculators/asespkcalculator.py       | 32 +++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 mlcalcdriver/calculators/ase_calculators/__init__.py
 create mode 100644 mlcalcdriver/calculators/ase_calculators/asespkcalculator.py

diff --git a/mlcalcdriver/calculators/ase_calculators/__init__.py b/mlcalcdriver/calculators/ase_calculators/__init__.py
new file mode 100644
index 0000000..23e2dd6
--- /dev/null
+++ b/mlcalcdriver/calculators/ase_calculators/__init__.py
@@ -0,0 +1 @@
+from mlcalcdriver.calculators.ase_calculators.asespkcalculator import AseSpkCalculator
diff --git a/mlcalcdriver/calculators/ase_calculators/asespkcalculator.py b/mlcalcdriver/calculators/ase_calculators/asespkcalculator.py
new file mode 100644
index 0000000..2d5ee1d
--- /dev/null
+++ b/mlcalcdriver/calculators/ase_calculators/asespkcalculator.py
@@ -0,0 +1,32 @@
+from ase.calculators.calculator import Calculator, all_changes
+from mlcalcdriver.calculators import SchnetPackCalculator
+from mlcalcdriver.base import Posinp, Job
+from copy import deepcopy
+import numpy as np
+
+
+class AseSpkCalculator(Calculator):
+
+    def __init__(self, model_dir, available_properties=None, device="cpu", **kwargs):
+        Calculator.__init__(self, **kwargs)
+        self.schnetpackcalculator = SchnetPackCalculator(
+            model_dir=model_dir,
+            available_properties=available_properties,
+            device=device,
+        )
+        self.implemented_properties = self.schnetpackcalculator._get_available_properties()
+        if "energy" in self.implemented_properties and "forces" not in self.implemented_properties:
+            self.implemented_properties.append("forces")
+
+    def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
+        if self.calculation_required(atoms, properties):
+            Calculator.calculate(self, atoms)
+            posinp = Posinp.from_ase(atoms)
+
+            job = Job(posinp=posinp, calculator=self.schnetpackcalculator)
+            for prop in properties:
+                job.run(prop)
+            results = {}
+            for prop, result in zip(job.results.keys(), job.results.values()):
+                results[prop] = np.squeeze(result)
+            self.results = results