From 3075a9b080937253d520962290654403d1c3cc2a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 18 Nov 2024 11:16:18 +0000
Subject: [PATCH] Add qm_charge kwarg to EMLECalculator constructor.

---
 emle/calculator.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/emle/calculator.py b/emle/calculator.py
index 10c8a7f..3af2db5 100644
--- a/emle/calculator.py
+++ b/emle/calculator.py
@@ -93,6 +93,7 @@ def __init__(
         method="electrostatic",
         alpha_mode="species",
         atomic_numbers=None,
+        qm_charge=0,
         backend="torchani",
         external_backend=None,
         plugin_path=".",
@@ -161,6 +162,12 @@ def __init__(
             you are using a fixed QM region, i.e. the same QM region for each
             call to the calculator.
 
+        qm_charge: int
+            The charge on the QM region. This is required when using an
+            EMLECalculator instance with the OpenMM interface. When using
+            the sander interface, the QM charge will be taken from the ORCA
+            input file.
+
         external_backend: str
             The name of an external backend to use to compute in vacuo energies.
             This should be a callback function formatted as 'module.function'.
@@ -415,6 +422,15 @@ def __init__(
         else:
             self._mm_charges = None
 
+        if qm_charge is not None:
+            try:
+                qm_charge = int(qm_charge)
+            except:
+                msg = "'qm_charge' must be of type 'int'"
+                _logger.error(msg)
+                raise TypeError(msg)
+            self._qm_charge = qm_charge
+
         # Create the EMLE model instance.
         self._emle = _EMLE(
             model=model,
@@ -422,6 +438,7 @@ def __init__(
             alpha_mode=alpha_mode,
             atomic_numbers=atomic_numbers,
             mm_charges=self._mm_charges,
+            qm_charge=self._qm_charge,
             device=self._device,
         )
 
@@ -962,6 +979,7 @@ def __init__(
             "method": self._method,
             "alpha_mode": self._alpha_mode,
             "atomic_numbers": None if atomic_numbers is None else atomic_numbers,
+            "qm_charge": self._qm_charge,
             "backend": self._backend,
             "external_backend": None if external_backend is None else external_backend,
             "mm_charges": None if mm_charges is None else self._mm_charges.tolist(),
@@ -1767,6 +1785,7 @@ def _sire_callback_optimised(
                 model_index=self._ani2x_model_index,
                 ani2x_model=self._torchani_model,
                 atomic_numbers=atomic_numbers,
+                qm_charge=self._qm_charge,
                 device=self._device,
             )