Merge bbec00d into 871187b

mlcalcdriver/base/job.py

@@ -133,7 +133,7 @@ def calculator(self, calculator):
                 """
             )
 
-    def run(self, property, device="cpu", batch_size=128):
+    def run(self, property, device="cpu", batch_size=128, finite_difference=False):
         r"""
         Main method to call to obtain results for a Job
 
@@ -164,6 +164,10 @@ def run(self, property, device="cpu", batch_size=128):
                 and "energy" in self.calculator.available_properties
             ):
                 raise ValueError("The property {} is not available".format(property))
+            elif not finite_difference:
+                predictions = self.calculator.run(
+                    property="forces", posinp=self.posinp, derivative=True
+                )
             else:
                 self._create_additional_structures()
                 raw_predictions = self.calculator.run(

mlcalcdriver/base/posinp.py

@@ -172,8 +172,10 @@ def _from_lines(cls, lines):
         # Decode the second line
         line2 = lines.pop(0)
         boundary_conditions = line2[0].lower()
-        if boundary_conditions != "free":
-            cell = [0.0 if dim in [".inf", "inf"] else dim for dim in line2[1:4]]
+        if boundary_conditions == "periodic":
+            cell = line2[1:4]
+        elif boundary_conditions == "surface":
+            cell = [line2[1], 0.0, line2[3]]
         else:
             cell = None
         # Remove the lines about the forces, if there are some

mlcalcdriver/calculators/schnetpack.py

@@ -42,20 +42,40 @@ def __init__(self, model_dir, available_properties=None, device="cpu", units=eVA
         super(SchnetPackCalculator, self).__init__(units=units)
         self._get_representation_type()
 
-    def run(self, property, posinp=None, device="cpu", batch_size=128):
+    def run(
+        self, property, derivative=False, posinp=None, device="cpu", batch_size=128,
+    ):
         r"""
         Main method to use when making a calculation with
         the calculator.
         """
         if property not in self.available_properties:
-            if property == "energy" and "energy_U0" in self.available_properties:
-                pass
+            if derivative:
+                if property == "forces":
+                    if "energy" in self.available_properties:
+                        init_property, deriv_name, out_name = (
+                            "energy",
+                            "forces",
+                            "forces",
+                        )
+                    else:
+                        raise ValueError(
+                            "This model can't be used for forces predictions."
+                        )
+                else:
+                    raise NotImplementedError(
+                        "Derivatives of other quantities than the energy are not implemented yet."
+                    )
             else:
                 raise ValueError(
                     "The property {} is not in the available properties of the model : {}.".format(
                         property, self.available_properties
                     )
                 )
+        elif property == "energy" and "energy_U0" in self.available_properties:
+            init_property, out_name = "energy_U0", "energy"
+        else:
+            init_property, out_name = property, property
 
         data = [posinp_to_ase_atoms(pos) for pos in posinp]
         pbc = True if any(pos.pbc.any() for pos in data) else False
@@ -76,6 +96,22 @@ def run(self, property, posinp=None, device="cpu", batch_size=128):
             for batch in data_loader:
                 batch = {k: v.to(device) for k, v in batch.items()}
                 pred.append(self.model(batch))
+        elif derivative:
+            for batch in data_loader:
+                batch = {k: v.to(device) for k, v in batch.items()}
+                batch["_positions"].requires_grad_()
+                results = self.model(batch)
+                drdx = (
+                    -1.0
+                    * torch.autograd.grad(
+                        results[init_property],
+                        batch["_positions"],
+                        grad_outputs=torch.ones_like(results[init_property]),
+                        create_graph=True,
+                        retain_graph=True,
+                    )[0]
+                )
+                pred.append({deriv_name: drdx})
         else:
             with torch.no_grad():
                 pred = []
@@ -84,15 +120,15 @@ def run(self, property, posinp=None, device="cpu", batch_size=128):
                     pred.append(self.model(batch))
 
         predictions = {}
-        if "energy_U0" not in self.available_properties:
-            predictions[property] = np.concatenate(
-                [batch[property].cpu().detach().numpy() for batch in pred]
+        if derivative:
+            predictions[out_name] = np.concatenate(
+                [batch[deriv_name].cpu().detach().numpy() for batch in pred]
             )
         else:
-            predictions[property] = np.concatenate(
-                [batch["energy_U0"].cpu().detach().numpy() for batch in pred]
+            print(init_property)
+            predictions[out_name] = np.concatenate(
+                [batch[init_property].cpu().detach().numpy() for batch in pred]
             )
-
         return predictions
 
     def _get_available_properties(self):

tests/test_phonon.py

@@ -17,7 +17,7 @@ class TestPhonon:
     def test_ph_N2(self):
         ph1 = Phonon(posinp=self.posN2, calculator=self.calcN2)
         ph1.run()
-        assert np.isclose(ph1.energies.max(), 2339.57, atol=0.01)
+        assert np.isclose(ph1.energies.max(), 2339.53, atol=0.01)
         assert all(np.abs(np.delete(ph1.energies, np.argmax(ph1.energies))) < 30)
 
         ph2 = Phonon(

tests/test_schnetpack.py

@@ -32,6 +32,13 @@ def test_forces_from_deriv(self):
         assert self.job.calculator.available_properties == ["energy"]
         self.job.run("forces")
         assert np.float32(-2979.6067) == self.job.results["energy"][0]
+        ref_forces = np.array([[-0.0, -0.0, -0.31532133], [-0.0, -0.0, 0.31532133]])
+        assert np.isclose(self.job.results["forces"][0], ref_forces).all()
+
+    def test_forces_from_finite_difference(self):
+        assert self.job.calculator.available_properties == ["energy"]
+        self.job.run("forces", finite_difference=True)
+        assert np.float32(-2979.6067) == self.job.results["energy"][0]
         ref_forces = np.array([[-0.0, -0.0, -0.32416448], [-0.0, -0.0, 0.32416448]])
         assert np.isclose(self.job.results["forces"][0], ref_forces).all()