optimizers

@dataclass
class ASEOptimizer[InputType: Structure | Molecule, OutputType: Structure | Molecule](
    PymatGenMaker[InputType, OutputType], GeometryOptimization
):
    """Base class for geometry optimization using ASE optimizers.

    Combines geometry optimization workflows with ASE calculator interfaces.
    This class provides the framework for optimizing molecular structures
    using various ASE optimization algorithms (LBFGS, BFGS, FIRE, etc.) and
    different calculators (neural networks, machine learning, semi-empirical).

    Units:
        Pass a float in the listed unit or a pint Quantity (e.g. ``jfchemistry.ureg``
        or ``jfchemistry.Q_``):

        - fmax: [eV/Å]

    Subclasses should inherit from both a specific ASECalculator implementation
    and ASEOptimizer to create complete optimization workflows.

    Attributes:
        name: Name of the optimizer (default: "ASE Optimizer").
        optimizer: ASE optimization algorithm to use:
            - "LBFGS": Limited-memory BFGS (default, recommended)
            - "BFGS": Broyden-Fletcher-Goldfarb-Shanno
            - "GPMin": Conjugate gradient
            - "MDMin": Molecular dynamics minimization
            - "FIRE": Fast Inertial Relaxation Engine
            - "FIRE2": FIRE version 2
            - "QuasiNewton": Quasi-Newton method
        fmax: Maximum force convergence criterion [eV/Å] (default: 0.05).
            Accepts float in [eV/Å] or pint Quantity.
        steps: Maximum number of optimization steps (default: 250000).

    Examples:
        >>> from ase.build import molecule # doctest: +SKIP
        >>> from pymatgen.core import Molecule # doctest: +SKIP
        >>> from jfchemistry.optimizers import ASEOptimizer # doctest: +SKIP
        >>> from jfchemistry.calculators import TBLiteCalculator # doctest: +SKIP
        >>> molecule = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
        >>> # Create custom optimizer by inheriting
        >>> class MyOptimizer(ASEOptimizer, TBLiteCalculator): # doctest: +SKIP
        ...     pass # doctest: +SKIP
        >>> opt = MyOptimizer(optimizer="LBFGS", fmax=0.01) # doctest: +SKIP
        >>> job = opt.make(molecule) # doctest: +SKIP
    """

    name: str = "ASE Optimizer"
    calculator: ASECalculator = field(
        default_factory=lambda: ASECalculator,
        metadata={"description": "the calculator to use for the calculation"},
    )
    optimizer: Literal["LBFGS", "BFGS", "GPMin", "MDMin", "FIRE", "FIRE2", "QuasiNewton"] = field(
        default="LBFGS",
        metadata={"description": "the ASE optimizer to use for the calculation"},
    )
    unit_cell_optimizer: Optional[
        Literal["UnitCellFilter", "ExpCellFilter", "FrechetCellFilter"]
    ] = field(
        default=None,
        metadata={"description": "the ASE unit cell optimizer to use for the calculation"},
    )
    fmax: float | Quantity = field(
        default=0.05,
        metadata={
            "description": "the maximum force convergence criterion [eV/Å]. "
            "Accepts float in [eV/Å] or pint Quantity.",
            "unit": "eV/Å",
        },
    )
    steps: int = field(
        default=250000,
        metadata={"description": "the maximum number of optimization steps"},
    )
    trajectory: Optional[str] = field(
        default=None,
        metadata={"description": "the trajectory file to save the optimization"},
    )
    logfile: Optional[str] = field(
        default=None,
        metadata={"description": "the log file to save the optimization"},
    )

    def __post_init__(self):
        """Post-initialization hook."""
        if isinstance(self.fmax, Quantity):
            object.__setattr__(self, "fmax", to_magnitude(self.fmax, "eV/angstrom"))
        self.name = f"{self.name} with {self.calculator.name}"
        super().__post_init__()

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Optimize molecular structure using ASE.

        Performs geometry optimization by:
        1. Converting structure to ASE Atoms
        2. Setting up the calculator with charge and spin
        3. Running the specified ASE optimizer
        4. Converting back to Pymatgen Molecule
        5. Extracting properties from the calculation

        Args:
            input: Input molecular structure with 3D coordinates.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            Tuple containing:
                - Optimized Pymatgen Molecule
                - Dictionary of computed properties from calculator

        Examples:
            >>> from ase.build import molecule # doctest: +SKIP
            >>> from pymatgen.core import Molecule # doctest: +SKIP
            >>> from jfchemistry.optimizers import TBLiteOptimizer # doctest: +SKIP
            >>> ethane = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
            >>> opt = TBLiteOptimizer(optimizer="LBFGS", fmax=0.01) # doctest: +SKIP
            >>> structures, properties = opt.operation(ethane) # doctest: +SKIP
        """
        atoms = input.to_ase_atoms()
        charge = int(input.charge)
        if type(input) is Molecule:
            spin_multiplicity = int(input.spin_multiplicity)
        else:
            spin_multiplicity = 1
        atoms = self.calculator._set_calculator(
            atoms, charge=charge, spin_multiplicity=spin_multiplicity
        )

        if type(input) is Structure and self.unit_cell_optimizer is not None:
            opt_atoms = getattr(filters, self.unit_cell_optimizer)(atoms)
        else:
            opt_atoms = atoms

        opt_func = getattr(ase.optimize, self.optimizer)
        opt = opt_func(opt_atoms, logfile=self.logfile, trajectory=self.trajectory)
        opt.run(self.fmax, self.steps)
        if type(input) is Structure:
            if self.unit_cell_optimizer is not None and isinstance(opt_atoms, Filter):
                opt_atoms = opt_atoms.atoms

        properties = self.calculator._get_properties(opt_atoms)

        if isinstance(input, Structure):
            opt_structure = Structure.from_ase_atoms(opt_atoms)
        elif isinstance(input, Molecule):
            opt_structure = Molecule.from_ase_atoms(opt_atoms)
        else:
            raise ValueError(f"Unsupported structure type: {type(input)}")

        return cast("OutputType", opt_structure), properties

@dataclass
class ORCAOptimizer[InputType: Molecule, OutputType: Molecule](
    ORCACalculator, GeometryOptimization, PymatGenMaker[InputType, OutputType]
):
    """Optimize molecular structures using ORCA DFT calculator.

    Inherits all attributes from ORCACalculator.

    Attributes:
        name: Name of the optimizer (default: "Orca Optimizer").
        opt: The ORCA optimizer to use for the calculation (default: ["OPT"]).
    """

    name: str = "Orca Optimizer"
    opt: Optional[list[OptModelType]] = field(
        default_factory=lambda: ["OPT"],
        metadata={"description": "the ORCA optimizer to use for the calculation"},
    )
    steps: int = field(
        default=250000,
        metadata={
            "description": "Maximum optimization steps. Set to 0 for fixed-geometry evaluation."
        },
    )
    _basename: str = "orca_optimizer"

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Optimize a molecule using ORCA DFT calculator."""
        input.to("input.xyz", fmt="xyz")
        sk_list = super()._set_keywords()
        if self.steps != 0:
            for opt_kw in self.opt or []:
                sk_list.append(getattr(Opt, opt_kw))
        calc = super()._build_calculator(self._basename)
        calc.structure = Structure.from_xyz("input.xyz")
        super()._set_structure_charge_and_spin(calc, input.charge, input.spin_multiplicity)
        super()._configure_calculator_input(calc, sk_list)
        calc.write_input()
        calc.run()
        output = calc.get_output()
        properties = super()._parse_output(output)
        final_molecule = Molecule.from_file(output.get_file(".xyz"))
        final_molecule = cast("OutputType", final_molecule)
        return final_molecule, properties

@dataclass
class TorchSimOptimizer[InputType: Molecule | Structure, OutputType: Molecule | Structure](
    PymatGenMaker[InputType, OutputType], GeometryOptimization
):
    """Base class for geometry optimization using TorchSim calculators.

    Combines geometry optimization with TorchSim calculator interfaces.
    This class provides the framework for optimizing structures
    using various TorchSim calculators (neural networks, machine learning,
    semi-empirical, etc.).

    Attributes:
        name: Name of the optimizer (default: "TorchSim Optimizer").

    Examples:
        >>> from ase.build import molecule # doctest: +SKIP
        >>> from pymatgen.core import Molecule # doctest: +SKIP
        >>> from jfchemistry.optimizers import TorchSimOptimizer # doctest: +SKIP
        >>> from jfchemistry.calculators.torchsim import OrbCalculator # doctest: +SKIP
        >>> molecule = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
        >>> # Create custom optimizer by inheriting
        >>> class MyOptimizer(TorchSimOptimizer, OrbCalculator): # doctest: +SKIP
        ...     pass # doctest: +SKIP
        >>> opt = MyOptimizer(optimizer="FIRE") # doctest: +SKIP
        >>> job = opt.make(molecule) # doctest: +SKIP
    """

    name: str = "TorchSim Optimizer"
    calculator: TorchSimCalculator = field(
        default_factory=lambda: TorchSimCalculator,
        metadata={"description": "the calculator to use for the calculation"},
    )
    optimizer: Literal["FIRE", "Gradient Descent"] = field(
        default="FIRE", metadata={"description": "The optimizer to use"}
    )
    autobatcher: bool = field(
        default=True, metadata={"description": "Whether to use the autobatcher"}
    )
    steps: int = field(
        default=10_000,
        metadata={
            "description": "The maximum number of steps to take. Set to 0 for fixed geometry."
        },
    )
    max_steps: Optional[int] = field(
        default=None,
        metadata={"description": "Deprecated alias for steps; if set, overrides steps."},
    )
    steps_between_swaps: int = field(
        default=5,
        metadata={
            "description": "Number of steps to take before checking convergence\
                 and swapping out states."
        },
    )

    def __post_init__(self):
        """Post-initialization hook."""
        if self.max_steps is not None:
            self.steps = self.max_steps
        self.max_steps = self.steps
        self.name = f"{self.name} with {self.calculator.name}"
        super().__post_init__()

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Optimize molecular structure using TorchSim.

        Performs geometry optimization by:
        1. Converting structure to TorchSim state
        2. Setting up the calculator with charge and spin
        3. Running the optimizer
        4. Converting back to Pymatgen Molecule
        5. Extracting properties from the calculation

        Args:
            input: Input molecular structure with 3D coordinates.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            Tuple containing:
                - Optimized Pymatgen Molecule
                - Dictionary of computed properties from calculator

        Examples:
            >>> from ase.build import molecule # doctest: +SKIP
            >>> from pymatgen.core import Molecule # doctest: +SKIP
            >>> from jfchemistry.optimizers import TorchSimOptimizer # doctest: +SKIP
            >>> ethane = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
            >>> opt = TorchSimOptimizer(optimizer="FIRE") # doctest: +SKIP
            >>> structures, properties = opt.operation(ethane) # doctest: +SKIP
        """
        optimizer = getattr(ts.Optimizer, self.optimizer.lower().replace(" ", "_"))
        model = self.calculator._get_model()
        input.to_ase_atoms().write("initial_structure.xyz")
        if self.steps == 0:
            final_structure = input
            properties = self.calculator._get_properties(final_structure)
        else:
            final_state = ts.optimize(
                system=input.to_ase_atoms(),
                model=model,
                optimizer=optimizer,
                max_steps=self.steps,
                steps_between_swaps=self.steps_between_swaps,
                autobatcher=self.autobatcher,
                pbar=True,
            )
            final_atoms = final_state.to_atoms()[0]
            final_structure = type(input).from_ase_atoms(final_atoms)
            properties = self.calculator._get_properties(final_structure)
        final_structure.to_ase_atoms().write("final_structure.xyz")
        return cast("OutputType", final_structure), properties