Source code for atomate2.turbomole.sets.core

"""Module defining core atomate2-turbomole input set generators."""

from __future__ import annotations

import json
from dataclasses import dataclass, field

import numpy as np
from pymatgen.core import Molecule, Structure
from turbomoleio.core.control import cpc
from turbomoleio.core.molecule import MoleculeSystem
from turbomoleio.core.periodic import PeriodicSystem

from atomate2.turbomole.sets.base import (
    BaseTurbomoleInputGenerator,
    DefineInputSet,
    TurbomoleInputSet,
)
from atomate2.turbomole.utils import get_define_parameters




@dataclass
class TurbomoleDefineInputGenerator(BaseTurbomoleInputGenerator):
    """A class to generate Turbomole input sets."""

    define_template: str | None = None
    define_parameters: dict = field(default_factory=dict)
    user_define_parameters_settings: dict = field(default_factory=dict)
    user_datagroups_settings: dict = field(default_factory=dict)
    use_system_charge: bool = True
    use_system_multiplicity: bool = False



    def get_input_set(
        self,
        system: MoleculeSystem | PeriodicSystem | Molecule | Structure,
        charge=None,
        unpaired_electrons=None,
    ) -> DefineInputSet:
        """
        Generate a DefineInputSet object.

        Parameters
        ----------
        system
            either a pymatgen Molecule or Structure or a
            turbomoleio MoleculeSystem or PeriodicSystem.
        charge
            charge of the system.
        unpaired_electrons
            number of unpaired electrons.

        Returns
        -------
            a DefineInputSet object

        TODO: can we use pydantic to make sure system is of the correct kind?
        """
        # Define the system
        if isinstance(system, Molecule):
            system = MoleculeSystem(system)
        elif isinstance(system, Structure):
            system = PeriodicSystem(system)
        elif not isinstance(  # pragma: no branch (trivial)
            system, (MoleculeSystem, PeriodicSystem)
        ):
            raise RuntimeError(
                'The "molecule_or_structure" input should be a pymatgen '
                "Molecule/Structure or a turbomoleio MoleculeSystem/PeriodicSystem."
            )

        # Get the define parameters
        dp = get_define_parameters(self.define_template, self.define_parameters)
        if unpaired_electrons is not None:
            dp["unpaired_electrons"] = unpaired_electrons
        elif self.use_system_multiplicity:  # pragma: no branch (trivial)
            raise NotImplementedError(
                "multiplicity from the system not yet implemented."
            )
        if charge is not None:
            dp["charge"] = charge
        elif dp.get("charge", None) is None and self.use_system_charge:
            dp["charge"] = system._molecule_or_structure.charge
        # Charges should be integers
        if dp.get("charge", None) is not None:
            int_charge = int(np.around(dp["charge"]))
            if not np.isclose(int_charge, dp["charge"]):
                raise ValueError("Charge is not close to an integer.")
            dp["charge"] = int_charge

        # The exchange-correlation functional should be set outside of define as it
        # does not know about libxc functionals
        func = dp.pop("functional", None)

        return DefineInputSet(
            inputs={
                "coord": system.to_coord_string(),
                "define_parameters.json": json.dumps(dp),
                "datagroups.json": json.dumps(self.user_datagroups_settings),
            },
            xc_func=func,
            datagroups=self.user_datagroups_settings,
            define_parameters=dp,
            system=system,
        )






@dataclass
class TurbomoleInputGenerator(BaseTurbomoleInputGenerator):
    """A class to generate Turbomole input sets."""



    def get_input_set(self, prev_output=None) -> TurbomoleInputSet:
        """
        Generate a TurbomoleInputSet object.

        TODO: what kind of file is prev_output?
        """
        if prev_output is not None:
            cpc(".", control_dir=prev_output)
        return TurbomoleInputSet()






@dataclass
class DscfInputGenerator(TurbomoleInputGenerator):
    """A class to generate Turbomole input sets."""



    def get_input_set(self, prev_output=None) -> TurbomoleInputSet:
        """
        Generate a TurbomoleInputSet object.

        TODO: what kind of file is prev_output?
        """
        if prev_output is not None:
            cpc(".", control_dir=prev_output)
        return TurbomoleInputSet()