Synchronise with devel

bin/emle-analyze

@@ -36,7 +36,7 @@ parser.add_argument(
 parser.add_argument(
     "--backend",
     type=str,
-    choices=["torchani", "mace", "deepmd"],
+    choices=["torchani", "mace", "emle-mace", "deepmd"],
     help="Gas phase ML backend",
 )
 parser.add_argument(
@@ -48,7 +48,7 @@ parser.add_argument(
 parser.add_argument(
     "--mace-model",
     type=str,
-    help="MACE model file (for backend='mace')",
+    help="MACE model file (for 'mace' and 'emle-mace' backends)",
 )
 parser.add_argument(
     "--qm-xyz", type=str, metavar="name.xyz", required=True, help="QM xyz file"
@@ -73,8 +73,14 @@ parser.add_argument(
 parser.add_argument(
     "--q-total",
     type=int,
+    default=0,
     help="Total charge of the QM region",
 )
+parser.add_argument(
+    "--use-dipoles",
+    action="store_true",
+    help="Whether to include static atomic dipoles",
+)
 parser.add_argument(
     "--start",
     type=int,
@@ -97,8 +103,8 @@ from emle._orca_parser import ORCAParser
 
 if args.backend == "deepmd" and not args.deepmd_model:
     parser.error("--deepmd-model is required when backend='deepmd'")
-if args.backend == "mace" and not args.mace_model:
-    parser.error("--mace-model is required when backend='mace'")
+if args.backend in ("mace", "emle-mace") and not args.mace_model:
+    parser.error("--mace-model is required when backend is 'mace' or 'emle-mace'")
 
 backend = None
 if args.backend == "torchani":
@@ -110,10 +116,14 @@ elif args.backend == "mace":
 
     backend = MACEEMLE(emle_model=args.emle_model, mace_model=args.mace_model)
 
-elif args.backend == "maceemle":
+elif args.backend == "emle-mace":
     from emle.models import MACEEMLEJoint
 
-    backend = MACEEMLEJoint(emle_model=args.emle_model, mace_model=args.mace_model)
+    backend = MACEEMLEJoint(
+        emle_model=args.emle_model,
+        mace_model=args.mace_model,
+        qm_charge=args.q_total,
+    )
 
 elif args.backend == "deepmd":
     from emle._backends import DeePMD
@@ -133,6 +143,7 @@ analyzer = EMLEAnalyzer(
     backend,
     orca_parser,
     args.q_total,
+    args.use_dipoles,
     start=args.start,
     end=args.end,
 )
@@ -148,6 +159,9 @@ result = {
     "atomic_alpha_emle": analyzer.atomic_alpha,
     "alpha_emle": analyzer.alpha,
 }
+if args.use_dipoles:
+    result["mu_emle"] = analyzer.mu
+    result["E_static_mu_emle"] = analyzer.e_static_mu
 
 if orca_parser:
     result.update(
@@ -163,9 +177,12 @@ if orca_parser:
             "E_static_mbis": analyzer.e_static_mbis,
         }
     )
+    if args.use_dipoles:
+        result["mu_qm"] = orca_parser.mbis["mu"]
 
 if args.backend:
     result["E_vac_emle"] = analyzer.e_backend
+    result["grad_vac_emle"] = analyzer.grad_backend
 if args.alpha:
     result["alpha_qm"] = orca_parser.alpha
 

bin/emle-server

@@ -61,6 +61,10 @@ except:
 model = os.getenv("EMLE_MODEL")
 method = os.getenv("EMLE_METHOD")
 alpha_mode = os.getenv("EMLE_ALPHA_MODE")
+try:
+    use_dipoles = strtobool(os.getenv("EMLE_USE_DIPOLES"))
+except:
+    use_dipoles = False
 atomic_numbers = os.getenv("EMLE_ATOMIC_NUMBERS")
 mm_charges = os.getenv("EMLE_MM_CHARGES")
 try:
@@ -144,6 +148,7 @@ env = {
     "model": model,
     "method": method,
     "alpha_mode": alpha_mode,
+    "use_dipoles": use_dipoles,
     "atomic_numbers": atomic_numbers,
     "mm_charges": mm_charges,
     "num_clients": num_clients,

bin/tarball-to-extxyz

@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+
+from emle._tarball_to_extXYZ import main
+
+if __name__ == "__main__":
+    main()

emle/__init__.py

@@ -36,6 +36,7 @@
 _supported_backends = [
     "torchani",
     "mace",
+    "emle-mace",
     "ace",
     "deepmd",
     "orca",

emle/_analyzer.py

@@ -33,6 +33,7 @@
 
 from ._units import _HARTREE_TO_KCAL_MOL, _ANGSTROM_TO_BOHR
 from ._utils import pad_to_max as _pad_to_max
+from .models import MACEEMLEJoint as _MACEEMLEJoint
 
 
 class EMLEAnalyzer:
@@ -48,6 +49,7 @@ def __init__(
         backend=None,
         parser=None,
         q_total=None,
+        use_dipoles=False,
         start=None,
         end=None,
     ):
@@ -145,13 +147,18 @@ def __init__(
             if isinstance(backend, _torch.nn.Module):
                 backend = backend.to(device).to(dtype)
                 atomic_numbers = _torch.tensor(atomic_numbers, device=device)
-                qm_xyz = _torch.tensor(qm_xyz, dtype=dtype, device=device)
+                qm_xyz = _torch.tensor(
+                    qm_xyz, dtype=dtype, device=device, requires_grad=True
+                )
                 charges_mm = _torch.empty((len(qm_xyz), 0), dtype=dtype, device=device)
                 mm_xyz = _torch.empty((len(qm_xyz), 0, 3), dtype=dtype, device=device)
             self.e_backend = (
-                backend(atomic_numbers, charges_mm, qm_xyz, mm_xyz).T
+                backend(atomic_numbers, charges_mm, qm_xyz, mm_xyz, qm_charge=q_total).T
                 * _HARTREE_TO_KCAL_MOL
             )
+            self.grad_backend = _torch.autograd.grad(self.e_backend.T[0].sum(), qm_xyz)[
+                0
+            ]
 
         self.atomic_numbers = _torch.tensor(
             atomic_numbers, dtype=_torch.int, device=device
@@ -163,11 +170,25 @@ def __init__(
         qm_xyz_bohr = self.qm_xyz * _ANGSTROM_TO_BOHR
         pc_xyz_bohr = self.pc_xyz * _ANGSTROM_TO_BOHR
 
-        self.s, self.q_core, self.q_val, self.A_thole = emle_base(
-            self.atomic_numbers,
-            self.qm_xyz,
-            self.q_total,
-        )
+        if isinstance(backend, _MACEEMLEJoint):
+            self.s = _torch.stack(backend.emle_values["s"])
+            self.q_core = _torch.stack(backend.emle_values["q_core"])
+            self.q_val = _torch.stack(backend.emle_values["q_val"])
+            self.mu = _torch.stack(backend.emle_values["mu"])
+
+            a_Thole = backend._mace.a_Thole
+            species_id = emle_base._species_map[self.atomic_numbers]
+            k = backend._mace.elements_alpha_v_ratios[species_id]
+            r_data = emle_base._get_r_data(qm_xyz_bohr, atomic_numbers > 0)
+            self.A_thole = emle_base._get_A_thole(
+                r_data, self.s, self.q_val, k, a_Thole
+            )
+        else:
+            self.s, self.q_core, self.q_val, self.A_thole = emle_base(
+                self.atomic_numbers,
+                self.qm_xyz,
+                self.q_total,
+            )
         self.atomic_alpha = 1.0 / _torch.diagonal(self.A_thole, dim1=1, dim2=2)[:, ::3]
         self.alpha = self._get_mol_alpha(self.A_thole, self.atomic_numbers)
 
@@ -179,6 +200,13 @@ def __init__(
             )
             * _HARTREE_TO_KCAL_MOL
         )
+        if use_dipoles:
+            self.e_static_mu = (
+                emle_base.get_static_energy(
+                    self.q_core, self.q_val, self.pc_charges, mesh_data, self.mu
+                )
+                * _HARTREE_TO_KCAL_MOL
+            )
         self.e_induced = (
             emle_base.get_induced_energy(
                 self.A_thole, self.pc_charges, self.s, mesh_data, mask
@@ -203,13 +231,17 @@ def __init__(
             "s",
             "q_core",
             "q_val",
+            "q",
+            "mu",
             "q_total",
             "atomic_alpha",
             "alpha",
             "e_backend",
             "e_static",
+            "e_static_mu",
             "e_induced",
             "e_static_mbis",
+            "grad_backend",
         ):
             if attr in self.__dict__:
                 setattr(self, attr, getattr(self, attr).detach().cpu().numpy())

emle/_orca_parser.py

@@ -33,8 +33,7 @@
 import tarfile as _tarfile
 
 from ._utils import pad_to_max
-
-_HARTREE_TO_KCALMOL = 627.509
+from ._units import _HARTREE_TO_KCAL_MOL
 
 
 class ORCAParser:
@@ -132,6 +131,9 @@ def __init__(self, filename, decompose=False, alpha=False):
                 self.mbis = self._parse_horton()
                 self.z, self.xyz = self._get_z_xyz()
 
+                self.E_vac = self.get_E_vac()
+                self.forces = self.get_forces()
+
                 if decompose:
                     self.vac_E, self.pc_E = self._get_E()
                     self.E = self.pc_E - self.vac_E
@@ -165,13 +167,13 @@ def _get_E(self):
     def _get_E_from_out(self, f):
         E_prefix = b"FINAL SINGLE POINT ENERGY"
         E_line = next(line for line in f if line.startswith(E_prefix))
-        return float(E_line.split()[-1]) * _HARTREE_TO_KCALMOL
+        return float(E_line.split()[-1]) * _HARTREE_TO_KCAL_MOL
 
     def _get_E_static(self):
         vpot_all = self._get_vpot()
         pc_all = self._get_pc()
         result = _np.array([(vpot @ pc) for vpot, pc in zip(vpot_all, pc_all)])
-        return result * _HARTREE_TO_KCALMOL
+        return result * _HARTREE_TO_KCAL_MOL
 
     def _get_vpot(self):
         return [
@@ -251,3 +253,70 @@ def _parse_horton_out(self, f):
 
     def _get_file(self, name, suffix):
         return self._tar.extractfile(f"{name}.{suffix}")
+
+    def get_E_vac(self):
+        """
+        Parse the gas-phase ORCA output for each configuration and return the
+        total energy in Hartree.
+
+        Returns
+        -------
+
+        E_tot: numpy.ndarray (N_CONFIGS,)
+            Total gas-phase energy per configuration, in Hartree.
+        """
+        E_tot = [
+            self._get_E_vac_from_out(self._get_file(name, "vac.orca"))
+            for name in self.names
+        ]
+        return _np.array(E_tot)
+
+    def _get_E_vac_from_out(self, f):
+        """
+        Extract the total energy from an ORCA gas-phase output file.
+        """
+        E_prefix = b"Total Energy       :"
+        E_line = [line for line in f if line.startswith(E_prefix)]
+        return float(E_line[0].split()[-2])
+
+    def get_forces(self):
+        """
+        Parse the gas-phase ORCA output for each configuration and return the
+        atomic forces in Hartree/Bohr.
+
+        Returns
+        -------
+
+        forces: List[numpy.ndarray]
+            Per-configuration atomic forces in Hartree/Bohr, each of shape
+            (N_ATOMS, 3).
+        """
+        forces = [
+            self._get_forces_from_out(self._get_file(name, "vac.orca"))
+            for name in self.names
+        ]
+        return forces
+
+    def _get_forces_from_out(self, f):
+        """
+        Extract the Cartesian forces from an ORCA gas-phase output file.
+        ORCA reports the gradient in Hartree/Bohr; we negate (forces =
+        -gradient) and return in the same units.
+        """
+        while next(f) != b"CARTESIAN GRADIENT\n":
+            pass
+        next(f)
+        next(f)
+        forces = []
+        try:
+            while True:
+                line_elements = next(f).split()
+                if len(line_elements) == 0:
+                    break
+                # Extract gradient (last 3 elements) and negate to get forces.
+                force_values = [-float(x) for x in line_elements[-3:]]
+                forces.append(force_values)
+
+        except ValueError:
+            pass
+        return _np.asarray(forces)