Support for spin-polarized calculations

.vscode/settings.json

@@ -14,5 +14,7 @@
   "python.defaultInterpreterPath": "${env:CONDA_PREFIX}/envs/dxtb/bin/python",
   "python.testing.pytestArgs": [],
   "python.testing.unittestEnabled": false,
-  "python.testing.pytestEnabled": true
+  "python.testing.pytestEnabled": true,
+  "python-envs.defaultEnvManager": "ms-python.python:system",
+  "python-envs.defaultPackageManager": "ms-python.python:pip"
 }

docs/source/01_quickstart/getting_started.rst

@@ -153,6 +153,47 @@ the example `here <https://github.com/grimme-lab/dxtb/blob/main/examples/forces.
 
            assert torch.allclose(g1, g2)
 
+Restricted, Unrestricted, and Spin Polarisation
+------------------------------------------------
+
+By default, *dxtb* uses restricted SCF (one spin channel).
+If you want to run unrestricted SCF (UHF) explicitly, enable ``uhf_mode`` in
+the calculator options.
+
+.. code-block:: python
+
+    import dxtb
+
+    calc = dxtb.Calculator(
+        numbers,
+        dxtb.GFN1_XTB,
+        opts={"uhf_mode": True},
+        **dd,
+    )
+
+    # Closed-shell run in unrestricted mode (useful for RHF/UHF comparisons)
+    energy = calc.energy(positions, chrg=0, spin=0)
+
+To include spin-polarisation, add the spin-polarisation interaction component.
+
+.. code-block:: python
+
+    import dxtb
+    from dxtb.components.spinpolarisation import new_spinpolarisation
+
+    spinpol = new_spinpolarisation(numbers, **dd)
+    calc = dxtb.Calculator(
+        numbers,
+        dxtb.GFN1_XTB,
+        interaction=[spinpol],
+        **dd,
+    )
+
+    # Open-shell example with one unpaired electron
+    energy = calc.energy(positions, chrg=0, spin=1)
+
+When a spin-polarisation interaction is present, *dxtb* automatically runs in
+unrestricted mode (two spin channels).
 
 More Properties
 ---------------

docs/source/02_indepth/components_interactions.rst

@@ -17,4 +17,5 @@ Below, you can find the detailed documentation.
 
    Coulomb <../_autosummary/dxtb.components.coulomb>
    Fields <../_autosummary/dxtb.components.field>
+   Spin Polarisation <../_autosummary/dxtb.components.spinpolarisation>
    Solvation <../_autosummary/dxtb.components.solvation>

src/dxtb/__main__.py

@@ -18,7 +18,7 @@
 Entry point for command line interface via `python -m <prog>`.
 """
 
-from .cli import console_entry_point
+from ._src.cli import console_entry_point
 
 if __name__ == "__main__":
     raise SystemExit(console_entry_point())

src/dxtb/_src/calculators/config/main.py

@@ -123,6 +123,7 @@ def __init__(
         fermi_maxiter: int = defaults.FERMI_MAXITER,
         fermi_thresh: float | int | None = defaults.FERMI_THRESH,
         fermi_partition: str | int = defaults.FERMI_PARTITION,
+        uhf_mode: bool = defaults.UHF_MODE,
         # cache
         cache_enabled: bool = defaults.CACHE_ENABLED,
         cache_hcore: bool = defaults.CACHE_STORE_HCORE,
@@ -227,6 +228,8 @@ def __init__(
             # SCF: PyTorch
             device=device,
             dtype=dtype,
+            # SCF: Mode
+            uhf_mode=uhf_mode,
         )
 
         # compatibility checks (only need to be skipped for some tests)
@@ -286,6 +289,8 @@ def from_args(cls, args: Namespace) -> Self:
             fermi_maxiter=args.fermi_maxiter,
             fermi_thresh=args.fermi_thresh,
             fermi_partition=args.fermi_partition,
+            # SCF: UHF
+            uhf_mode=args.uhf_mode or getattr(args, "spinpol", False),
             # Cache
             cache_enabled=args.cache_enabled,
             cache_hcore=args.cache_hcore,

src/dxtb/_src/calculators/config/scf.py

@@ -141,6 +141,10 @@ class ConfigSCF:
     dtype: torch.dtype
     """Data type for calculations."""
 
+    # UHF
+    uhf_mode: bool
+    """Restricted (false) or unrestricted (true) SCF"""
+
     def __init__(
         self,
         *,
@@ -172,6 +176,8 @@ def __init__(
         # PyTorch
         device: torch.device = get_default_device(),
         dtype: torch.dtype = get_default_dtype(),
+        # UHF
+        uhf_mode: bool = defaults.UHF_MODE,
     ) -> None:
         self.strict = strict
         self.method = method
@@ -343,6 +349,8 @@ def __init__(
         self.device = device
         self.dtype = dtype
 
+        self.uhf_mode = uhf_mode
+
         self.x_atol = check_tols(x_atol, dtype)
         self.x_atol_max = check_tols(x_atol_max, dtype)
         self.f_atol = check_tols(f_atol, dtype)
@@ -378,6 +386,7 @@ def info(self) -> dict[str, Any]:
                 "x tolerance": self.x_atol,
                 "f(x) tolerance": self.f_atol,
                 **self.fermi.info(),
+                "UHF Mode for SCF": self.uhf_mode,
             }
         }
 
@@ -397,6 +406,7 @@ def __str__(self):  # pragma: no cover
             f"  xitorch absolute Tolerance: {self.x_atol}",
             f"  xitorch Functional Tolerance: {self.f_atol}",
             f"  Fermi Configuration: {self.fermi}",
+            f"  Unrestricted SCF: {self.uhf_mode}",
         ]
         return "\n".join(config_str)
 

src/dxtb/_src/calculators/types/energy.py

@@ -119,6 +119,10 @@ def singlepoint(
             else:
                 hashed_key += f"{sep}{arg}"
 
+        # Normalize coordinates once so all downstream components use the
+        # calculator dtype/device consistently.
+        positions = positions.to(**self.dd)
+
         is_batched = self.numbers.ndim == 2
 
         if is_batched:

src/dxtb/_src/cli/argparser.py

@@ -399,6 +399,16 @@ def parser(name: str = "dxtb", **kwargs: Any) -> argparse.ArgumentParser:
             "in the SCF. Note that 'charge' and 'charges' is identical."
         ),
     )
+    p.add_argument(
+        "--uhf-mode",
+        "--uhf_mode",
+        default=defaults.UHF_MODE,
+        action="store_true",
+        help=(
+            "R|Use unrestricted SCF mode\n"
+            "note this takes roughly twice as long as restricted SCF"
+        ),
+    )
 
     # Integrals
     p.add_argument(
@@ -480,6 +490,16 @@ def parser(name: str = "dxtb", **kwargs: Any) -> argparse.ArgumentParser:
         action="store_true",
         help="R|Whether to compute gradients for positions w.r.t. energy.",
     )
+    p.add_argument(
+        "--spinpol",
+        action="store_true",
+        default=defaults.SPINPOL,
+        help=(
+            "R|Activates Spin Polarisation\n"
+            "Recommended for open-shell system.\n"
+            "Requires unrestricted SCF and thus activates uhf_mode"
+        ),
+    )
 
     # Cache
 

src/dxtb/_src/cli/driver.py

@@ -32,6 +32,7 @@
 from dxtb._src.calculators.config import Config
 from dxtb._src.calculators.result import Result
 from dxtb._src.components.interactions.field import new_efield
+from dxtb._src.components.interactions.spin import new_spinpolarisation
 from dxtb._src.constants import labels
 from dxtb._src.timing import timer
 from dxtb._src.typing import Tensor
@@ -223,6 +224,9 @@ def singlepoint(self) -> Result | Tensor:
             )
             interactions.append(new_efield(field, **dd))
 
+        if args.spinpol is True:
+            interactions.append(new_spinpolarisation(numbers, **dd))
+
         # setup calculator
         calc = Calculator(
             numbers,
@@ -232,6 +236,7 @@ def singlepoint(self) -> Result | Tensor:
             **dd,
             timer=args.timer,
         )
+
         timer.stop("Setup")
 
         ####################################################

src/dxtb/_src/components/classicals/dispersion/d4.py

@@ -130,8 +130,9 @@ def get_cache(
             model = d4.model.D4Model(
                 numbers, ref_charges=self.ref_charges, **self.dd
             )
-        else:
-            model = model.type(self.dtype).to(self.device)
+
+        # Ensure externally provided models follow this component settings.
+        model = model.type(self.dtype).to(self.device)
 
         rcov = kwargs.pop("rcov", None)
         if rcov is not None and not isinstance(rcov, Tensor):
@@ -197,19 +198,32 @@ def get_energy(
 
         # FIXME: Charge should be REQUIRED for D4!
         if self.charge is None and "charge" not in kwargs:
+            charge_input: Tensor | float | int | None = torch.tensor(
+                0.0, **self.dd
+            )
+        else:
+            charge_input = kwargs.pop("charge", self.charge)
+
+        if charge_input is None:
             charge = torch.tensor(0.0, **self.dd)
+        elif isinstance(charge_input, Tensor):
+            charge = charge_input.to(**self.dd)
         else:
-            charge = kwargs.pop("charge", self.charge)
+            charge = torch.tensor(charge_input, **self.dd)
+
+        q_d4 = cache.q if q is None else q
+        if q_d4 is not None:
+            q_d4 = q_d4.to(**self.dd)
 
         return d4.dftd4(
             self.numbers,
-            positions,
+            positions.to(**self.dd),
             charge,
             self.param,
             model=cache.model,
             rcov=cache.rcov,
             r4r2=cache.r4r2,
-            q=cache.q if q is None else q,
+            q=q_d4,
             cutoff=cache.cutoff,
             counting_function=cache.counting_function,
             damping_function=cache.damping_function,

src/dxtb/_src/components/interactions/__init__.py

@@ -27,3 +27,4 @@
 from .field import *
 from .list import InteractionList, InteractionListCache
 from .solvation import *
+from .spin import *

src/dxtb/_src/components/interactions/base.py

@@ -93,6 +93,17 @@ class Interaction(Component):
     label: str
     """Label for the interaction."""
 
+    spin_channel: int | None = None
+    """
+    Which spin channel this interaction reads from and writes to
+    when charges carry an ``nspin`` dimension.
+
+    * ``None`` (default): charge-type interaction, reads from channel 0
+      (total charges).
+    * ``1``: magnetization-type interaction (e.g. spin polarisation),
+      reads from channel 1 (magnetization charges).
+    """
+
     def __init__(
         self,
         device: torch.device | None = None,
@@ -131,6 +142,22 @@ def get_cache(
         """
         return InteractionCache()
 
+    def _extract_mono_charges(self, charges: Charges) -> Tensor:
+        """
+        Extract the relevant monopole charges for this interaction.
+
+        For spin-polarized calculations (``charges.nspin > 1``), the
+        appropriate spin channel is selected according to
+        :attr:`spin_channel`.  Charge-type interactions (``spin_channel
+        is None``) use channel 0 (total charges); the spin interaction
+        uses channel 1 (magnetization).
+        """
+        nspin = getattr(charges, "nspin", 1)
+        if nspin > 1:
+            ch = self.spin_channel if self.spin_channel is not None else 0
+            return charges.mono[..., ch, :]
+        return charges.mono
+
     @final
     def get_potential(
         self,
@@ -155,9 +182,11 @@ def get_potential(
         Tensor
             Potential vector for each orbital partial charge.
         """
+        nspin = getattr(charges, "nspin", 1)
+        qat_ch = self._extract_mono_charges(charges)
 
         # monopole potential: shell-resolved
-        qsh = ihelp.reduce_orbital_to_shell(charges.mono)
+        qsh = ihelp.reduce_orbital_to_shell(qat_ch)
         vsh = self.get_monopole_shell_potential(cache, qsh)
 
         # monopole potential: atom-resolved
@@ -170,6 +199,19 @@ def get_potential(
         vsh += ihelp.spread_atom_to_shell(vat)
         vmono = ihelp.spread_shell_to_orbital(vsh)
 
+        # Route potential into the correct spin channel
+        if nspin > 1:
+            ch = self.spin_channel if self.spin_channel is not None else 0
+            vmono_full = torch.zeros(
+                *vmono.shape[:-1],
+                nspin,
+                vmono.shape[-1],
+                device=vmono.device,
+                dtype=vmono.dtype,
+            )
+            vmono_full[..., ch, :] = vmono
+            vmono = vmono_full
+
         # multipole potentials
         vdipole = self.get_dipole_atom_potential(
             cache, qat, charges.dipole, charges.quad
@@ -337,7 +379,9 @@ def get_energy(
                 "charges are required."
             )
 
-        qsh = ihelp.reduce_orbital_to_shell(charges.mono)
+        qat_ch = self._extract_mono_charges(charges)
+
+        qsh = ihelp.reduce_orbital_to_shell(qat_ch)
         esh = self.get_monopole_shell_energy(cache, qsh)
 
         qat = ihelp.reduce_shell_to_atom(qsh)
@@ -381,7 +425,10 @@ def get_monopole_atom_energy(
         return torch.zeros_like(qat)
 
     def get_monopole_shell_energy(
-        self, cache: InteractionCache, qat: Tensor, **_: Any
+        self,
+        cache: InteractionCache,
+        qat: Tensor,
+        **_: Any,
     ) -> Tensor:
         """
         Compute the energy from the charges, all quantities are shell-resolved.

src/dxtb/_src/components/interactions/container.py

@@ -312,6 +312,21 @@ class Charges(Container):
     Container for the charges used in the SCF.
     """
 
+    nspin: int
+    """Number of spin channels (1 = RHF, 2 = UHF)."""
+
+    def __init__(
+        self,
+        mono: Tensor | None = None,
+        dipole: Tensor | None = None,
+        quad: Tensor | None = None,
+        label: str | list[str] | None = None,
+        batch_mode: int = 0,
+        nspin: int = 1,
+    ) -> None:
+        super().__init__(mono, dipole, quad, label, batch_mode)
+        self.nspin = nspin
+
     @property
     def mono(self) -> Tensor:
         if self._mono is None:

src/dxtb/_src/components/interactions/coulomb/secondorder.py

@@ -373,7 +373,10 @@ def get_monopole_atom_energy(
 
     @override
     def get_monopole_shell_energy(
-        self, cache: ES2Cache, qat: Tensor, **_: Any
+        self,
+        cache: ES2Cache,
+        qat: Tensor,
+        **_: Any,
     ) -> Tensor:
         return (
             0.5 * qat * self.get_monopole_shell_potential(cache, qat)