Merge 2822bd6 into c48bd22

mlcalcdriver/base/job.py

@@ -158,39 +158,32 @@ def run(self, property, device="cpu", batch_size=128, finite_difference=False):
             if not torch.cuda.is_available():
                 warnings.warn("CUDA was asked for, but is not available.", UserWarning)
 
-        if property not in self.calculator.available_properties:
-            if not (
-                property == "forces"
-                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(
-                    property="energy", posinp=self.posinp
-                )
-                pred_idx = 0
-                predictions = {}
-                predictions["energy"], predictions["forces"] = [], []
-                for struct_idx in range(self.num_struct):
-                    predictions["energy"].append(raw_predictions["energy"][pred_idx][0])
-                    pred_idx += 1
-                    predictions["forces"].append(
-                        self._calculate_forces(
-                            raw_predictions["energy"][
-                                pred_idx : pred_idx
-                                + 12 * len(self._init_posinp[struct_idx])
-                            ]
-                        )
-                    )
-                    pred_idx += 12 * len(self._init_posinp[struct_idx])
-                self.posinp = deepcopy(self._init_posinp)
+        if not finite_difference:
+            predictions = self.calculator.run(
+                property=property, posinp=self.posinp, batch_size=batch_size
+            )
         else:
-            predictions = self.calculator.run(property=property, posinp=self.posinp)
+            self._create_additional_structures()
+            raw_predictions = self.calculator.run(
+                property="energy", posinp=self.posinp, batch_size=batch_size
+            )
+            pred_idx = 0
+            predictions = {}
+            predictions["energy"], predictions["forces"] = [], []
+            for struct_idx in range(self.num_struct):
+                predictions["energy"].append(raw_predictions["energy"][pred_idx][0])
+                pred_idx += 1
+                predictions["forces"].append(
+                    self._calculate_forces(
+                        raw_predictions["energy"][
+                            pred_idx : pred_idx
+                            + 12 * len(self._init_posinp[struct_idx])
+                        ]
+                    )
+                )
+                pred_idx += 12 * len(self._init_posinp[struct_idx])
+            self.posinp = deepcopy(self._init_posinp)
+
         for pred in predictions.keys():
             # Future proofing, will probably need some work
             if pred in ["energy", "gap"]:
@@ -201,6 +194,11 @@ def run(self, property, device="cpu", batch_size=128, finite_difference=False):
                     predictions[pred] *= HA_TO_EV
                 if self.calculator.units["positions"] == "atomic":
                     predictions[pred] *= ANG_TO_B
+            elif pred == "hessian":
+                if self.calculator.units["energy"] == "hartree":
+                    predictions[pred] *= HA_TO_EV
+                if self.calculator.units["positions"] == "atomic":
+                    predictions[pred] *= ANG_TO_B ** 2
             elif pred in ["dipole_moment", "mu"]:
                 if self.calculator.units["dipole_moment"] == "Debye":
                     pass

mlcalcdriver/calculators/schnetpack.py

@@ -9,8 +9,10 @@
 from schnetpack import AtomsLoader
 from mlcalcdriver.globals import eVA
 from mlcalcdriver.calculators import Calculator
+from mlcalcdriver.calculators.utils import torch_derivative, get_derivative_names
 from mlcalcdriver.interfaces import posinp_to_ase_atoms, SchnetPackData
 from schnetpack.environment import SimpleEnvironmentProvider, AseEnvironmentProvider
+from mlcalcdriver.globals import EV_TO_HA, B_TO_ANG
 
 
 class SchnetPackCalculator(Calculator):
@@ -43,39 +45,18 @@ def __init__(self, model_dir, available_properties=None, device="cpu", units=eVA
         self._get_representation_type()
 
     def run(
-        self, property, derivative=False, posinp=None, device="cpu", batch_size=128,
+        self, property, 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 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
+        init_property, out_name, derivative, wrt = get_derivative_names(
+            property, self.available_properties
+        )
+
+        if len(posinp) > 1 and derivative:
+            batch_size = 1
 
         data = [posinp_to_ase_atoms(pos) for pos in posinp]
         pbc = True if any(pos.pbc.any() for pos in data) else False
@@ -92,40 +73,52 @@ def run(
         data_loader = AtomsLoader(data, batch_size=batch_size)
 
         pred = []
-        if self.model.output_modules[0].derivative is not None:
+        if derivative == 0:
+            if self.model.output_modules[0].derivative is not None:
+                for batch in data_loader:
+                    batch = {k: v.to(device) for k, v in batch.items()}
+                    pred.append(self.model(batch))
+            else:
+                with torch.no_grad():
+                    for batch in data_loader:
+                        batch = {k: v.to(device) for k, v in batch.items()}
+                        pred.append(self.model(batch))
+        if abs(derivative) == 1:
             for batch in data_loader:
                 batch = {k: v.to(device) for k, v in batch.items()}
-                pred.append(self.model(batch))
-        elif derivative:
+                batch[wrt[0]].requires_grad_()
+                results = self.model(batch)
+                deriv1 = torch.unsqueeze(
+                    torch_derivative(results[init_property], batch[wrt[0]]), 0
+                )
+                if derivative < 0:
+                    deriv1 = -1.0 * deriv1
+                pred.append({out_name: deriv1})
+        if abs(derivative) == 2:
             for batch in data_loader:
                 batch = {k: v.to(device) for k, v in batch.items()}
-                batch["_positions"].requires_grad_()
+                for inp in set(wrt):
+                    batch[inp].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]
+                deriv2 = torch.unsqueeze(
+                    torch_derivative(
+                        torch_derivative(
+                            results[init_property], batch[wrt[0]], create_graph=True,
+                        ),
+                        batch[wrt[0]],
+                    ),
+                    0,
                 )
-                pred.append({deriv_name: drdx})
-        else:
-            with torch.no_grad():
-                pred = []
-                for batch in data_loader:
-                    batch = {k: v.to(device) for k, v in batch.items()}
-                    pred.append(self.model(batch))
-
+                if derivative < 0:
+                    deriv2 = -1.0 * deriv2
+                pred.append({out_name: deriv2})
         predictions = {}
         if derivative:
-            predictions[out_name] = np.concatenate(
-                [batch[deriv_name].cpu().detach().numpy() for batch in pred]
+            predictions[property] = np.concatenate(
+                [batch[out_name].cpu().detach().numpy() for batch in pred]
             )
         else:
-            predictions[out_name] = np.concatenate(
+            predictions[property] = np.concatenate(
                 [batch[init_property].cpu().detach().numpy() for batch in pred]
             )
         return predictions
@@ -149,6 +142,7 @@ def _get_available_properties(self):
 
     def _get_representation_type(self):
         r"""
+        Private method to determine representation type (schnet or wcasf).
         """
         if "representation.cutoff.cutoff" in self.model.state_dict().keys():
             self.model_type = "wacsf"

mlcalcdriver/calculators/utils.py

@@ -0,0 +1,42 @@
+import torch
+
+
+def torch_derivative(fx, x, create_graph=False):
+    if len(fx.shape) == 2:
+        fx = fx.unsqueeze(0)
+    dfdx = []
+    flat_fx = fx.reshape(-1)
+    for i in range(len(flat_fx)):
+        (grad_x,) = torch.autograd.grad(
+            flat_fx[i],
+            x,
+            torch.ones_like(flat_fx[i]),
+            retain_graph=True,
+            create_graph=create_graph,
+        )
+        dfdx.append(grad_x)
+    return torch.stack(dfdx).reshape(fx.shape[2] * x.shape[1], fx.shape[1] * x.shape[2])
+
+
+def get_derivative_names(property, avail):
+    if property not in avail:
+        if property == "forces" and "energy" in avail:
+            init_property, out_name, derivative = "energy", "forces", -1
+            wrt = ["_positions"]
+        elif property == "hessian" and "forces" in avail:
+            init_property, out_name, derivative = "forces", "hessian", -1
+            wrt = ["_positions"]
+        elif property == "hessian" and "energy" in avail:
+            init_property, out_name, derivative = "energy", "hessian", 2
+            wrt = ["_positions", "_positions"]
+        else:
+            raise ValueError(
+                "The property {} is not in the available properties of the model : {}.".format(
+                    property, avail
+                )
+            )
+    elif property == "energy" and "energy_U0" in avail:
+        init_property, out_name, derivative, wrt = "energy_U0", "", 0, []
+    else:
+        init_property, out_name, derivative, wrt = property, "", 0, []
+    return init_property, out_name, derivative, wrt

mlcalcdriver/workflows/phonon.py

@@ -25,7 +25,14 @@ class Phonon:
     translated by a small amount around the equilibrium positions.
     """
 
-    def __init__(self, posinp, calculator, relax=True, translation_amplitudes=None):
+    def __init__(
+        self,
+        posinp,
+        calculator,
+        relax=True,
+        finite_difference=False,
+        translation_amplitudes=None,
+    ):
         r"""
         The initial position fo the atoms are taken from the `init_state`
         Posinp instance. If they are not part of a relaxed geometry, the
@@ -52,16 +59,19 @@ def __init__(self, posinp, calculator, relax=True, translation_amplitudes=None):
         relax : bool
             Wether the initial positions need to be relaxed or not.
             Default is `True`.
+        finite_difference: bool
+            If True, the hessian matrix is calculated using finite
+            displacements of atoms. Default is False. Mostly there for
+            legacy reasons.
         translation_amplitudes: list of length 3
             Amplitudes of the translations to be applied to each atom
             along each of the three space coordinates (in angstroms).
-        order : int
-            Order of the numerical differentiation used to compute the
-            dynamical matrix. Can be either 1, 2 or 3.
+            Only relevant if finite_difference is True.
         """
         self.posinp = posinp
         self.calculator = calculator
         self.relax = relax
+        self.finite_difference = finite_difference
         self.translation_amplitudes = translation_amplitudes
 
         if self.relax:
@@ -151,6 +161,21 @@ def relax(self, relax):
         relax = bool(relax)
         self._relax = relax
 
+    @property
+    def finite_difference(self):
+        r"""
+        Returns
+        -------
+        finite_difference : bool
+            If `True`, the hessian matrix is calculated using small finite
+            movements on the atoms. Default is `False`.
+        """
+        return self._finite_difference
+
+    @finite_difference.setter
+    def finite_difference(self, finite_difference):
+        self._finite_difference = bool(finite_difference)
+
     @property
     def energies(self):
         r"""
@@ -212,15 +237,20 @@ def run(self, device="cpu", batch_size=128, **kwargs):
             geopt.run(device=device, batch_size=batch_size)
             self._ground_state = deepcopy(geopt.final_posinp)
 
-        job = Job(posinp=self._create_displacements(), calculator=self.calculator)
-        job.run(property="forces", device=device, batch_size=batch_size)
+        if self.finite_difference:
+            job = Job(posinp=self._create_displacements(), calculator=self.calculator)
+            job.run(property="forces", device=device, batch_size=batch_size)
+        else:
+            job = Job(posinp=self._ground_state, calculator=self.calculator)
+            job.run(property="hessian", device=device, batch_size=batch_size)
         self._post_proc(job)
 
     def _create_displacements(self):
         r"""
         Set the displacements each atom must undergo from the amplitudes
         of displacement in each direction. The numerical derivatives are obtained
-        with the five-point stencil method.
+        with the five-point stencil method. Only used if the phonons are
+        calculated using finite_displacements.
         """
         structs = []
         for i in range(len(self._ground_state)):
@@ -263,14 +293,22 @@ def _compute_hessian(self, job):
         Computes the hessian matrix from the forces
         """
         n_at = len(self.posinp)
-        hessian = np.zeros((3 * n_at, 3 * n_at))
-        forces = np.array(job.results["forces"]) * EV_TO_HA * B_TO_ANG
-        for i in range(3 * n_at):
-            hessian[i, :] = (
-                -forces[4 * i].flatten()
-                + forces[4 * i + 3].flatten()
-                + 8 * (forces[4 * i + 1].flatten() - forces[4 * i + 2].flatten())
-            ) / (12 * self.translation_amplitudes * ANG_TO_B)
+        if "hessian" in job.results.keys():
+            h = (
+                job.results["hessian"].reshape(3 * n_at, 3 * n_at)
+                * EV_TO_HA
+                * B_TO_ANG ** 2
+            )
+            return (h + h.T) / 2.0
+        else:
+            hessian = np.zeros((3 * n_at, 3 * n_at))
+            forces = np.array(job.results["forces"]) * EV_TO_HA * B_TO_ANG
+            for i in range(3 * n_at):
+                hessian[i, :] = (
+                    -forces[4 * i].flatten()
+                    + forces[4 * i + 3].flatten()
+                    + 8 * (forces[4 * i + 1].flatten() - forces[4 * i + 2].flatten())
+                ) / (12 * self.translation_amplitudes * ANG_TO_B)
         return -(hessian + hessian.T) / 2.0
 
     def _solve_dyn_mat(self):

requirements_dev.txt

@@ -1,15 +1,15 @@
 pip==20.0.2
 wheel==0.34.2
 black==19.10b0
-numpy==1.18.1
+numpy==1.18.2
 torch==1.4.0
 torchvision==0.5.0
 twine==3.1.1
 pytest==5.3.5
 pytest-cov==2.8.1
 pytest-sugar==0.9.2
-coveralls==1.10.0
+coveralls==1.11.1
 schnetpack==0.3
-sphinx==2.4.2
+sphinx==2.4.4
 sphinx-rtd-theme==0.4.3
 m2r==0.2.1