Ase

def md(
    self, dyn: MolecularDynamics, trajectory: list[Atoms] | None = None
) -> DynSchema:
    """
    Get tabulated results from an ASE MD run.

    Parameters
    ----------
    dyn
        ASE MolecularDynamics object.
    trajectory
        ASE Trajectory object or list[Atoms] from reading a trajectory file.

    Returns
    -------
    DynSchema
        Dictionary representation of the task document
    """
    # Check and set up variables
    base_task_doc = self.opt(dyn, trajectory=trajectory, check_convergence=False)
    del base_task_doc["converged"]

    # Clean up the opt parameters
    parameters_md = base_task_doc.pop("parameters_opt")
    parameters_md.pop("logfile", None)

    trajectory_log = []
    for t, atoms in enumerate(base_task_doc["trajectory"]):
        trajectory_log.append(
            {
                "kinetic_energy": atoms.get_kinetic_energy(),
                "temperature": atoms.get_temperature(),
                "time": t * parameters_md["timestep"],
            }
        )

    md_fields = {"parameters_md": parameters_md, "trajectory_log": trajectory_log}

    # Create a dictionary of the inputs/outputs
    unsorted_task_doc = base_task_doc | md_fields | self.additional_fields

    return clean_dict(unsorted_task_doc)

def neb(
    self,
    dyn: Optimizer,
    n_images: int,
    n_iter_return: int = -1,
    trajectory: TrajectoryWriter | list[Atoms] | None = None,
) -> OptSchema:
    """
    Summarize the NEB run results and store them in a database-friendly format.

    Parameters
    ----------
    dyn
        ASE Optimizer object used for the NEB run.
    n_images
        Number of images in the NEB run.
    n_iter_return
        Number of iterations to return. If -1, all iterations are returned.
    trajectory
        Trajectory of the NEB run, either as a Trajectory object or a list of Atoms objects.

    Returns
    -------
    OptSchema
        A dictionary containing the summarized NEB run results.
    """

    # Get trajectory
    atoms_trajectory = (
        trajectory or read(dyn.trajectory, index=":")
        if isinstance(dyn.trajectory, (str, Path))
        else read(dyn.trajectory.filename, index=":")
    )

    if n_iter_return == -1:
        atoms_trajectory = atoms_trajectory[-(n_images):]
    else:
        atoms_trajectory = _get_nth_iteration(
            atoms_trajectory,
            int(len(atoms_trajectory) / n_images),
            n_images,
            n_iter_return,
        )
    trajectory_results = [atoms.calc.results for atoms in atoms_trajectory]
    ts_index = (
        np.argmax(
            [
                result["energy"]
                for result in trajectory_results[-(n_images - 1) : -1]
            ]
        )
        + 1
    )
    ts_atoms = atoms_trajectory[ts_index]
    base_task_doc = atoms_to_metadata(atoms_trajectory[0])

    # Clean up the opt parameters
    parameters_opt = dyn.todict()
    parameters_opt.pop("logfile", None)
    parameters_opt.pop("restart", None)

    opt_fields = {
        "parameters_opt": parameters_opt,
        "trajectory": atoms_trajectory,
        "trajectory_results": trajectory_results,
        "ts_atoms": ts_atoms,
    }

    # Create a dictionary of the inputs/outputs
    unsorted_task_doc = base_task_doc | opt_fields | self.additional_fields

    return clean_dict(unsorted_task_doc)

def opt(
    self,
    dyn: Optimizer,
    trajectory: list[Atoms] | None = None,
    check_convergence: bool | DefaultSetting = QuaccDefault,
) -> OptSchema:
    """
    Get tabulated results from an ASE optimization.

    Parameters
    ----------
    dyn
        ASE Optimizer object.
    trajectory
        ASE Trajectory object or list[Atoms] from reading a trajectory file.
    check_convergence
        Whether to check the convergence of the calculation. Defaults to True in
        settings.

    Returns
    -------
    OptSchema
        Dictionary representation of the task document
    """

    # Check and set up variables
    check_convergence = (
        self._settings.CHECK_CONVERGENCE
        if check_convergence == QuaccDefault
        else check_convergence
    )

    # Get trajectory
    atoms_trajectory = (
        trajectory or read(dyn.trajectory, index=":")
        if isinstance(dyn.trajectory, (str, Path))
        else read(dyn.trajectory.filename, index=":")
    )
    trajectory_results = [atoms.calc.results for atoms in atoms_trajectory]

    initial_atoms = atoms_trajectory[0]
    final_atoms = get_final_atoms_from_dynamics(dyn)
    directory = self.directory or final_atoms.calc.directory

    # Check convergence
    gradient = (
        dyn.optimizable.get_gradient() if hasattr(dyn, "optimizable") else None
    )
    is_converged = dyn.converged(gradient)
    if check_convergence and not is_converged:
        msg = f"Optimization did not converge. Refer to {directory}"
        raise RuntimeError(msg)

    # Base task doc
    base_task_doc = self.run(final_atoms, initial_atoms)

    # Clean up the opt parameters
    parameters_opt = dyn.todict()
    parameters_opt.pop("logfile", None)
    parameters_opt.pop("restart", None)

    opt_fields = {
        "parameters_opt": parameters_opt,
        "converged": is_converged,
        "trajectory": atoms_trajectory,
        "trajectory_results": trajectory_results,
    }

    # Create a dictionary of the inputs/outputs
    unsorted_task_doc = base_task_doc | opt_fields | self.additional_fields

    return clean_dict(unsorted_task_doc)

def run(self, final_atoms: Atoms, input_atoms: Atoms) -> RunSchema:
    """
    Get tabulated results from a standard ASE run.

    Parameters
    ----------
    final_atoms
        ASE Atoms following a calculation. A calculator must be attached.
    input_atoms
        Input ASE Atoms object to store.

    Returns
    -------
    RunSchema
        Dictionary representation of the task document
    """

    # Check and set up variables
    if not final_atoms.calc:
        msg = "ASE Atoms object has no attached calculator."
        raise ValueError(msg)
    if not final_atoms.calc.results:
        msg = "ASE Atoms object's calculator has no results."
        raise ValueError(msg)

    directory = self.directory or final_atoms.calc.directory

    # Generate input atoms metadata
    input_atoms_metadata = atoms_to_metadata(input_atoms) if input_atoms else {}

    # Generate the base of the task document
    inputs = {
        "parameters": final_atoms.calc.parameters,
        "nid": get_uri(directory).split(":")[0],
        "dir_name": directory,
        "input_atoms": input_atoms_metadata,
        "quacc_version": __version__,
    }
    results = {"results": final_atoms.calc.results}

    # Prepare atoms for the next run
    atoms_to_store = prep_next_run(final_atoms, move_magmoms=self.move_magmoms)

    # Generate final atoms metadata
    final_atoms_metadata = atoms_to_metadata(atoms_to_store) if final_atoms else {}

    # Create a dictionary of the inputs/outputs
    unsorted_task_doc = (
        final_atoms_metadata | inputs | results | self.additional_fields
    )

    return clean_dict(unsorted_task_doc)

def vib(self, *, is_molecule: bool) -> VibSchema:
    """
    Get tabulated results from an ASE Vibrations object and store them in a database-
    friendly format.

    Parameters
    ----------
    is_molecule
        Whether the Atoms object is a molecule. If True, the vibrational modes are
        sorted by their absolute value and the 3N-5 or 3N-6 modes are taken. If False,
        all vibrational modes are taken.

    Returns
    -------
    VibSchema
        Dictionary representation of the task document
    """
    # Tabulate input parameters
    vib_freqs_raw = self.vib_object.get_frequencies().tolist()
    vib_energies_raw = self.vib_object.get_energies().tolist()
    if isinstance(self.vib_object, VibrationsData):
        atoms = self.vib_object._atoms
        directory = self.directory
        inputs = {"nid": get_uri(directory).split(":")[0], "dir_name": directory}
    else:
        atoms = self.vib_object.atoms
        directory = self.directory or atoms.calc.directory
        inputs = {
            "parameters": atoms.calc.parameters,
            "parameters_vib": {
                "delta": self.vib_object.delta,
                "direction": self.vib_object.direction,
                "method": self.vib_object.method,
                "ndof": self.vib_object.ndof,
                "nfree": self.vib_object.nfree,
            },
            "nid": get_uri(directory).split(":")[0],
            "dir_name": directory,
            "quacc_version": __version__,
        }

    # Convert imaginary modes to negative values for DB storage
    for i, f in enumerate(vib_freqs_raw):
        if np.imag(f) > 0:
            vib_freqs_raw[i] = -np.abs(f)
            vib_energies_raw[i] = -np.abs(vib_energies_raw[i])
        else:
            vib_freqs_raw[i] = np.abs(f)
            vib_energies_raw[i] = np.abs(vib_energies_raw[i])

    # Get the true vibrational modes
    atoms_metadata = atoms_to_metadata(atoms)

    natoms = len(atoms)
    if natoms == 1:
        vib_freqs = []
        vib_energies = []
    elif is_molecule:
        is_linear = (
            PointGroupData()
            .from_molecule(AseAtomsAdaptor().get_molecule(atoms))
            .linear
            if atoms.pbc.any()
            else atoms_metadata["molecule_metadata"]["symmetry"]["linear"]
        )

        # Sort by absolute value
        vib_freqs_raw_sorted = vib_freqs_raw.copy()
        vib_energies_raw_sorted = vib_energies_raw.copy()
        vib_freqs_raw_sorted.sort(key=np.abs)
        vib_energies_raw_sorted.sort(key=np.abs)

        # Cut the 3N-5 or 3N-6 modes based on their absolute value
        n_modes = 3 * natoms - 5 if is_linear else 3 * natoms - 6
        vib_freqs = vib_freqs_raw_sorted[-n_modes:]
        vib_energies = vib_energies_raw_sorted[-n_modes:]
    else:
        vib_freqs = vib_freqs_raw
        vib_energies = vib_energies_raw

    imag_vib_freqs = [f for f in vib_freqs if f < 0]

    vib_results = {
        "results": {
            "imag_vib_freqs": imag_vib_freqs,
            "n_imag": len(imag_vib_freqs),
            "vib_energies": vib_energies,
            "vib_freqs": vib_freqs,
            "vib_energies_raw": vib_energies_raw,
            "vib_freqs_raw": vib_freqs_raw,
        }
    }
    unsorted_task_doc = (
        atoms_metadata | inputs | vib_results | self.additional_fields
    )

    return clean_dict(unsorted_task_doc)

def vib_and_thermo(
    self,
    thermo_method: Literal["ideal_gas", "harmonic"],
    energy: float = 0.0,
    temperature: float = 298.15,
    pressure: float = 1.0,
) -> VibThermoSchema:
    """
    Get tabulated results from an ASE Vibrations object and thermochemistry.

    Parameters
    ----------
    thermo_method
        Method to use for thermochemistry calculations. If None, no thermochemistry
        calculations are performed.
    energy
        Potential energy in eV used as the reference point for thermochemistry calculations.
    temperature
        Temperature in K for thermochemistry calculations.
    pressure
        Pressure in atm for thermochemistry calculations

    Returns
    -------
    VibThermoSchema
        Dictionary representation of the task document
    """
    atoms = (
        self.vib_object._atoms
        if isinstance(self.vib_object, VibrationsData)
        else self.vib_object.atoms
    )
    is_molecule = bool(thermo_method == "ideal_gas")

    # Generate vib data
    vib_schema = self.vib(is_molecule=is_molecule)

    # Generate thermo data
    thermo_summary = ThermoSummarize(
        atoms,
        vib_schema["results"]["vib_freqs_raw"],
        energy=energy,
        additional_fields=self.additional_fields,
    )
    if thermo_method == "ideal_gas":
        thermo_schema = thermo_summary.ideal_gas(
            temperature=temperature, pressure=pressure
        )
    elif thermo_method == "harmonic":
        thermo_schema = thermo_summary.harmonic(
            temperature=temperature, pressure=pressure
        )
    else:
        raise ValueError(f"Unsupported thermo_method: {thermo_method}.")

    # Merge the vib and thermo data
    unsorted_task_doc = recursive_dict_merge(vib_schema, thermo_schema)

    return clean_dict(unsorted_task_doc)