lineshape_tools 0.1.0__py3-none-any.whl

lineshape_tools/cli.py

@@ -0,0 +1,965 @@
+"""Contains functionality for the command line interface."""
+
+from __future__ import annotations
+
+import logging
+import sys
+import warnings
+from importlib.util import find_spec
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+import ase.io
+import numpy as np
+from cyclopts import Parameter
+from tqdm import tqdm
+
+from lineshape_tools.constants import omega2eV
+from lineshape_tools.lineshape import convert_A_to_L, gaussian, get_phonon_spec_func
+from lineshape_tools.phonon import get_disp_vect, get_ipr, get_phonons
+from lineshape_tools.plot import plot_spec_funcs
+
+if TYPE_CHECKING:
+    from ase.atoms import Atoms
+    from mace.calculators import MACECalculator
+
+logger = logging.getLogger(__name__)
+
+
+def collect(
+    files: Annotated[list[Path], Parameter(negative="")],
+    output_file: Path = Path("./database.extxyz"),
+    strategy: str = "none",
+    read_index: str = ":",
+    max_force: float = 2.0,
+    min_force: float = -np.inf,
+    dx_tol: float = 0.1,
+    rtol: float = 1e-5,
+    config_weight: float = 1.0,
+    force_weighting: bool = False,
+) -> None:
+    """Collect and process data for fine-tuning.
+
+    Collect files into an extxyz database that can be used for fine-tuning. An optional filtering
+    strategy can be applied. This is potentially useful if relaxation data is being included in the
+    dataset, as it can be noisey from having multiple closely spaced geometries close to the
+    equilibrium geomtry. By default, configurations with too large forces will be thrown away to
+    avoid potential anharmonic contributions to the PES.
+
+    Args:
+        files (list): list of paths to files that are parseable by ase.io. The files should contain
+            atomic geometries, total energies, and forces at a minimum (for example, vasprun.xml).
+        output_file (Path): optional path where data is written (output to stdout by default)
+        strategy (str): optional specification of strategy to be used for filtering. Available
+            options are 'none', 'qr', or 'dx'.
+        read_index (str): pythonic index passed to ase.io.read to determine which structures are
+            read from the input files (the same value is used for each file). The default ":" reads
+            all of the structures, while ":3" would read the first three for example.
+        max_force (float): remove structures where the maximum force acting on any atom is above
+            the specified value (in eV/Å)
+        min_force (float): remove structures where the maximum force acting on any atom is below
+            the specified value (in eV/Å)
+        dx_tol (float): tolerance (in Å) for how far atoms must move to accept configuration in
+            'dx' filtering strategy
+        rtol (float): tolerance ratio for determining rank of displacement vectors in 'qr' strategy
+        config_weight (float): set the configuration weight for training
+        force_weighting (bool): store a config_weight that's inversely proportional to the max
+            force that any atom feels in the configuration [min(0.02 / max_fpa, 1)]. Overwrites
+            the value specified by config_weight.
+    """
+    if strategy.lower() not in ("none", "qr", "dx"):
+        raise ValueError("invalid strategy choice")
+
+    if output_file.exists():
+        logger.info(f"{output_file} already exists, appending to it")
+
+    total_read = 0
+    all_atoms = []
+    for fname in tqdm(files, desc="[*] reading files", disable=len(files) < 2):
+        read_atoms = ase.io.read(fname, read_index)
+        total_read += len(read_atoms)
+        for atoms in tqdm(read_atoms, desc="[*] processing atoms", disable=len(read_atoms) < 10):
+            forces = atoms.get_forces()
+            if min_force < np.linalg.norm(forces, axis=1).max() < max_force:
+                atoms.info["REF_energy"] = atoms.get_potential_energy()
+                atoms.new_array("REF_forces", forces)
+                atoms.calc = None
+                atoms.info["config_weight"] = config_weight
+                if force_weighting:
+                    atoms.info["config_weight"] = min(
+                        0.02 / np.linalg.norm(forces, axis=1).max(), 1
+                    )
+                    if (cw := atoms.info["config_weight"]) < 0.02:
+                        logger.warning(f"small config weight found ({cw})")
+                all_atoms.append(atoms)
+
+    logger.info(
+        f"read {total_read} configurations and discarded {total_read - len(all_atoms)} based"
+        " on forces criteria"
+    )
+
+    filtered_atoms = []
+    if strategy.lower() == "qr":
+        from scipy.linalg import qr
+
+        logger.warning("qr strategy assumes the last read structure is the equilibrium geometry")
+        if np.linalg.norm(all_atoms[-1].arrays["REF_forces"], axis=1).max() > 0.02:
+            warnings.warn("large forces found in last structure", stacklevel=0)
+
+        sqrt_mass = np.repeat(np.sqrt(all_atoms[-1].get_masses()), 3)
+
+        dq = np.zeros((sqrt_mass.shape[0], len(all_atoms) - 1), dtype=np.float64)
+        for atoms_i, atoms in enumerate(tqdm(all_atoms[:-1], desc="[*] computing dq")):
+            dx = get_disp_vect(atoms, all_atoms[-1])
+            dq[:, atoms_i] = sqrt_mass * dx
+
+        logger.info("computing pivoted QR factorization")
+        _, R, P = qr(dq, mode="economic", pivoting=True)
+
+        rdiag = np.diag(R)
+        rank = np.sum(np.abs(rdiag) > rtol * rdiag.max())
+
+        filtered_atoms = [all_atoms[-1]] + [all_atoms[atoms_i] for atoms_i in np.sort(P[:rank])]
+    elif strategy.lower() == "dx":
+        filtered_atoms.append(all_atoms[0])
+        for atoms in tqdm(all_atoms[1:], desc="[*] dx filtering"):
+            dx = get_disp_vect(filtered_atoms[-1], atoms)
+            if np.linalg.norm(dx) > dx_tol:
+                filtered_atoms.append(atoms)
+    else:
+        filtered_atoms = all_atoms
+
+    logger.info(f"filtered {len(all_atoms) - len(filtered_atoms)} configurations from dataset")
+
+    with open(output_file, "a") as f:
+        for atoms in tqdm(filtered_atoms, desc="[*] writing"):
+            ase.io.write(f, atoms, format="extxyz")
+
+
+def get_force_opt_modes(
+    n: int,
+    omega2: np.ndarray,
+    U: np.ndarray,
+    sqrt_mass: np.ndarray,
+    F: float = 0.5,
+    tol: float = 1e-6,
+    seed: int = 897689932,
+    start_with_min_spread: bool = False,
+    save_plot: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Bin phonons by energy and find a vector in the subspace that minimizes the spread in forces.
+
+    Args:
+        n (int): number of modes to select
+        omega2 (np.ndarray): frequencies squared of the modes (directly from np.linalg.eigh)
+        U (np.ndarray): matrix where eigenvectors of modes are cols (directly from np.linalg.eigh)
+        sqrt_mass (np.ndarray): sqrt of the vector of atomic masses
+        F (float): target forces to optimize amplitudes for
+        tol (float): convergence tolerance for scipy minimize call
+        seed (int): seed value for random number generator
+        start_with_min_spread (bool): determines if mode with smallest force spread is used as the
+            starting point for optimization. Uses a random vector otherwise.
+        save_plot (bool): save a plot for analyzing resulting modes
+
+    Returns:
+        modes (np.ndarray): generated modes as columns of the matrix
+        mode_dqs (np.ndarray): optimized displacement amplitudes following above criteria
+    """
+    inv_sqrt_mass = 1 / sqrt_mass
+
+    # acoustic phonon filtering
+    omega2[omega2 < (0.0005 / omega2eV) ** 2] = 0.0
+
+    opt_path, opt_info = np.einsum_path("i,ij,j,j->i", sqrt_mass, U, omega2, np.ones(U.shape[0]))
+    logger.debug(opt_info)
+
+    def loss(x, ind=None):
+        if ind is not None:
+            tx = np.zeros(U.shape[0])
+            tx[ind] = x
+        else:
+            tx = x
+
+        Fx = np.einsum("i,ij,j,j->i", sqrt_mass, U, omega2, tx, optimize=opt_path)
+        Fx = np.linalg.norm(Fx.reshape((-1, 3)), axis=1)
+        return np.mean((Fx - Fx.mean()) ** 2)
+
+    def constr_f(x):
+        return np.sum(x**2)
+
+    def constr_J(x):
+        return 2 * x
+
+    def constr_H(x, v):
+        return np.diag(2 * v[0] * np.ones(x.shape[0]))
+
+    def log_callback(intermediate_result):
+        ir = intermediate_result
+        logger.debug(f"{ir.nit:04d} | {ir.fun:.06e} | {ir.optimality:.06e}")
+
+    from scipy.optimize import NonlinearConstraint, minimize
+
+    norm_constr = NonlinearConstraint(constr_f, 1, 1, jac=constr_J, hess=constr_H)
+
+    rng = np.random.default_rng(seed)
+
+    modes = np.zeros((U.shape[0], n), dtype=np.float64)
+    mode_dqs = np.zeros(n, dtype=np.float64)
+    freqs = np.empty(n)
+
+    subspace_inds = np.array_split(np.arange(3, U.shape[0]), n)
+    for imode, ind in enumerate(tqdm(subspace_inds, desc="[*] optimizing modes")):
+        subs_spread = np.empty(ind.shape[0])
+        for i, ind_i in enumerate(ind):
+            Fx = np.linalg.norm((sqrt_mass * omega2[ind_i] * U[:, ind_i]).reshape((-1, 3)), axis=1)
+            subs_spread[i] = np.mean((Fx - Fx.mean()) ** 2)
+
+        if start_with_min_spread:
+            imin = np.argmin(subs_spread)
+            logger.info(
+                f"starting from mode {ind[imin]} "
+                f"with frequency {1000 * omega2eV * np.sqrt(omega2[ind[imin]]):.02f} meV "
+                f"and spread {subs_spread[imin]:.06e}"
+            )
+            x0 = np.zeros(ind.shape[0])
+            x0[imin] = 1.0
+        else:
+            x0 = rng.random(ind.shape[0]) - 0.5
+            x0 /= np.linalg.norm(x0)
+
+        logger.debug(" nit |     loss     |  optimality")
+        res = minimize(
+            lambda x: loss(x, ind=ind),  # noqa: B023
+            x0,
+            tol=tol,
+            method="trust-constr",
+            constraints=[norm_constr],
+            callback=log_callback,
+        )
+
+        if not res.success:
+            logger.warning(f"optimization failed - {res.message}")
+
+        logger.debug(f"final spread {res.fun}, subspace spread {subs_spread}")
+        if not np.all(res.fun < subs_spread):
+            logger.warning("optimization failed to find a vector with smaller spread")
+
+        modes[ind, imode] = res.x
+        freqs[imode] = omega2eV * np.sqrt(modes[:, imode] @ np.diag(omega2) @ modes[:, imode])
+        logger.debug(f"found mode with frequency {1000 * freqs[imode]:.02f} meV")
+
+        Fx = np.linalg.norm(
+            np.einsum("i,ij,j,j->i", sqrt_mass, U, omega2, modes[:, imode]).reshape((-1, 3)),
+            axis=1,
+        )
+        mode_dqs[imode] = F / Fx.max()
+
+        dx = np.linalg.norm(
+            (inv_sqrt_mass * mode_dqs[imode] * (U @ modes[:, imode])).reshape((-1, 3)), axis=1
+        )
+        if dx.max() > 0.05:
+            new_dq = (0.05 / dx.max()) * mode_dqs[imode]
+            logger.warning(
+                f"max displacement too large ({dx.max()} > 0.05), "
+                f"resetting dq {mode_dqs[imode]} -> {new_dq}"
+            )
+            mode_dqs[imode] = new_dq
+        elif dx.max() < 0.005:
+            new_dq = (0.005 / dx.max()) * mode_dqs[imode]
+            logger.warning(
+                f"max displacement too small ({dx.max()} < 0.005), "
+                f"resetting dq {mode_dqs[imode]} -> {new_dq}"
+            )
+            mode_dqs[imode] = new_dq
+
+    if save_plot:
+        w = np.linspace(0, 0.1, 1000)
+        dos = np.zeros_like(w)
+        for i in range(U.shape[0]):
+            dos += gaussian(w - omega2eV * np.sqrt(omega2[i]), 0.001)
+        dos /= U.shape[0]
+
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots(figsize=(4, 3))
+        ax.plot(w, dos, color="k")
+        ax.fill_between(w, dos, color="k", alpha=0.2, lw=0)
+        for imode in range(n):
+            pdos = np.zeros_like(w)
+            for i in range(U.shape[0]):
+                pdos += modes[i, imode] ** 2 * gaussian(w - omega2eV * np.sqrt(omega2[i]), 0.001)
+            p = ax.plot(w, pdos / n, lw=1)
+            ax.axvline(x=freqs[imode], color=p[0].get_color(), alpha=0.2)
+        ax.set_xlabel("Energy [eV]")
+        ax.set_ylabel("Density of States")
+        plt.savefig("pdos.png", dpi=600, bbox_inches="tight")
+
+    return U @ modes, mode_dqs
+
+
+def get_random_phonons(
+    n: int,
+    omega2: np.ndarray,
+    U: np.ndarray,
+    sqrt_mass: np.ndarray,
+    F: float = 0.5,
+    seed: int = 897689932,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Bin phonons by energy and select one randomly from each bin.
+
+    The amplitude of each phonon is chosen to produce a max force per atom as close to F as
+    possible. The max displacement on a given atom is kept within a reasonable range (0.005, 0.05).
+
+    Args:
+        n (int): number of modes to select
+        omega2 (np.ndarray): frequencies squared of the modes (directly from np.linalg.eigh)
+        U (np.ndarray): matrix where eigenvectors of modes are cols (directly from np.linalg.eigh)
+        sqrt_mass (np.ndarray): sqrt of the vector of atomic masses
+        F (float): target forces to optimize amplitudes for
+        seed (int): seed value for random number generator
+
+    Returns:
+        modes (np.ndarray): generated modes as columns of the matrix
+        mode_dqs (np.ndarray): optimized displacement amplitudes following above criteria
+    """
+    inv_sqrt_mass = 1 / sqrt_mass
+
+    # acoustic phonon filtering
+    omega2[omega2 < (0.0005 / omega2eV) ** 2] = 0.0
+
+    modes = np.zeros((U.shape[0], n), dtype=np.float64)
+    mode_dqs = np.empty(n, dtype=np.float64)
+
+    rng = np.random.default_rng(seed)
+
+    # starting from 3 to skip acoustic phonons
+    subspace_inds = np.array_split(np.arange(3, U.shape[0]), n)
+    for i, inds in enumerate(tqdm(subspace_inds, desc="[*] selecting random modes")):
+        imode = inds[rng.integers(inds.shape[0])]
+        modes[:, i] = U[:, imode]
+
+        Fx = np.linalg.norm((sqrt_mass * omega2[imode] * U[:, imode]).reshape((-1, 3)), axis=1)
+        mode_dqs[i] = F / Fx.max()
+
+        dx = np.linalg.norm((mode_dqs[i] * inv_sqrt_mass * U[:, imode]).reshape((-1, 3)), axis=1)
+        if dx.max() > 0.05:
+            new_dq = (0.05 / dx.max()) * mode_dqs[i]
+            logger.warning(
+                f"max displacement too large ({dx.max()} > 0.05), "
+                f"resetting dq {mode_dqs[i]} -> {new_dq}"
+            )
+            mode_dqs[i] = new_dq
+        elif dx.max() < 0.005:
+            new_dq = (0.005 / dx.max()) * mode_dqs[i]
+            logger.warning(
+                f"max displacement too small ({dx.max()} < 0.005), "
+                f"resetting dq {mode_dqs[i]} -> {new_dq}"
+            )
+            mode_dqs[i] = new_dq
+    return modes, mode_dqs
+
+
+def gen_confs(
+    struct_path: Path,
+    num_conf: int,
+    strategy: str = "rand",
+    output_dir: Path = Path("./confs"),
+    accepting_mode: Path | None = None,
+    dynmat_file: Path | None = None,
+    orthogonalize: bool = False,
+    default_max_dx: float = 0.015,
+    start_with_min_spread: bool = False,
+    opt_tol: float = 1e-6,
+    seed: int = 897689932,
+) -> None:
+    """Generate additional configurations to enhance fine-tuning dataset.
+
+    Args:
+        struct_path (Path): path to file containing structure that will be displaced
+        num_conf (int): total number of additional configurations to generate
+        strategy (str): strategy used to generate the additional configurations. Available options
+            are 'rand', 'phon_rand', and 'phon_opt'.
+        output_dir (Path): output directory where the new configurations will be written to
+        accepting_mode (Path): path to file containing the structure that defines the accepting
+            mode. For example, if struct_path refers to the ground-state equilibrium geometry, then
+            accepting_mode should refer to the excited-state equilibrium geometry and vice versa.
+        dynmat_file (Path): path to the .npz file containing the dynamical matrix presumably
+            calculated using the "compute-dynmat" function.
+        orthogonalize (bool): perform Gram-Schmidt orthogonalization at the last step
+        default_max_dx (float): default value for max displacement of a given atom
+        start_with_min_spread (bool): determines if mode with smallest force spread is used as the
+            starting point for optimization in phon_opt strategy. Uses a random vector otherwise.
+        opt_tol (float): convergence tolerance for the call to scipy minimize in the phon_opt strat
+        seed (int): seed value for random number generator
+    """
+    if strategy.lower() not in ("rand", "phon_rand", "phon_opt"):
+        raise ValueError("invalid strategy choice")
+
+    if output_dir.exists():
+        raise ValueError("output directory already exists")
+
+    struct: Atoms = ase.io.read(struct_path)  # type: ignore[assignment]
+    sqrt_mass = np.repeat(np.sqrt(struct.get_masses()), 3)
+
+    # makes written poscars ugly and adds unnecessary data
+    if struct.has("momenta"):
+        del struct.arrays["momenta"]
+
+    imode = 0
+    modes = np.zeros((3 * len(struct), num_conf), dtype=np.float64)
+    mode_dqs = np.zeros(num_conf, dtype=np.float64)
+
+    if accepting_mode is not None:
+        am_struct: Atoms = ase.io.read(accepting_mode)  # type: ignore[assignment]
+        dx = get_disp_vect(struct, am_struct)
+        dq = sqrt_mass * dx
+        logger.info(f"accepting mode dQ={np.linalg.norm(dq)} amu^0.5 Å")
+
+        modes[:, 0] = dq / np.linalg.norm(dq)
+        mode_dqs[0] = (
+            default_max_dx
+            / np.linalg.norm((modes[:, 0] / sqrt_mass).reshape((-1, 3)), axis=1).max()
+        )
+
+        imode += 1
+
+    if strategy.lower()[:4] == "phon":
+        logger.info("working in phonon basis")
+
+        if dynmat_file is None:
+            raise ValueError("dynamical matrix file is needed to compute phonons")
+
+        logger.info(f"reading dynamical matrix from {dynmat_file}")
+        data = np.load(dynmat_file)
+        H = data["H"]
+
+        if not np.allclose(sqrt_mass, data["sqrt_mass"], rtol=1e-4):
+            logger.debug(sqrt_mass, data["sqrt_mass"])
+            raise ValueError("sqrt_mass from struct is not compatible with dynamical matrix file")
+
+        # diagonalize dynamical matrix
+        omega2, U = np.linalg.eigh(H)
+
+        if strategy.lower() == "phon_opt":
+            logger.info("generating force-optimized phonons")
+            modes[:, imode:], mode_dqs[imode:] = get_force_opt_modes(
+                num_conf - imode,
+                omega2,
+                U,
+                sqrt_mass,
+                tol=opt_tol,
+                start_with_min_spread=start_with_min_spread,
+                seed=seed,
+            )
+        else:
+            logger.info("selecting phonons randomly")
+            modes[:, imode:], mode_dqs[imode:] = get_random_phonons(
+                num_conf - imode,
+                omega2,
+                U,
+                sqrt_mass,
+                seed=seed,
+            )
+    else:
+        logger.info("generating random modes")
+
+        from scipy.linalg import qr
+
+        rng = np.random.default_rng(seed)
+        modes[:, imode:], _, _ = qr(
+            rng.random((modes.shape[0], num_conf - imode)) - 0.5, mode="economic", pivoting=True
+        )
+        for i in range(imode, num_conf):
+            mode_dqs[i] = (
+                default_max_dx
+                / np.linalg.norm((modes[:, i] / sqrt_mass).reshape((-1, 3)), axis=1).max()
+            )
+
+    if orthogonalize:
+        from scipy.linalg import qr
+
+        logger.info("performing final Gram-Schmidt orthogonalization of all modes")
+        modes, _ = qr(modes, mode="economic")
+
+        logger.info("recomputing displacement amplitude after orthogonalization")
+        for i in range(num_conf):
+            mode_dqs[i] = (
+                default_max_dx
+                / np.linalg.norm((modes[:, i] / sqrt_mass).reshape((-1, 3)), axis=1).max()
+            )
+
+    # write to output directory
+    for imode in tqdm(range(num_conf), desc="[*] writing structures"):
+        dx = mode_dqs[imode] * modes[:, imode] / sqrt_mass
+        logger.debug(
+            f"mode {imode} - max displacement  "
+            f"{np.linalg.norm(dx.reshape((-1, 3)), axis=1).max():.06f} Å  "
+            f"{np.abs(dx).max():.06f} Å"
+        )
+
+        atoms = struct.copy()
+        atoms.positions += dx.reshape((-1, 3))
+
+        fname = output_dir / f"{imode}/POSCAR"
+        fname.parent.mkdir(parents=True)
+        ase.io.write(fname, atoms, format="vasp", direct=True)
+
+    # store invoking command
+    with open(output_dir / "cmd.txt", "w") as f:
+        f.write(" ".join(sys.argv) + "\n")
+
+
+def reestimate_e0s_linear_system(
+    calculator: MACECalculator,
+    database_atoms: list[Atoms],
+    elements: list | None = None,
+    initial_e0s: dict | None = None,
+) -> dict:
+    """Estimate atomic reference energies (E0s) by solving a linear system.
+
+    Notes:
+        Slightly adapted from code by Noam Bernstein based on private communications
+        with Ilyes Batatia and Joe Hart.
+
+        This functionality will eventually be removed once merged into MACE.
+
+    Args:
+        calculator (MACECalculator): Calculator object for the MACE model.
+        database_atoms (list): List of ase Atoms objects with energy and atomic_numbers.
+        elements (list): List of element atomic numbers to consider, default to set present in
+            database_atoms.
+        initial_e0s (dict): Dictionary mapping element atomic numbers to E0 values, default to
+            values returned by foundation_model for isolated atom configs>
+
+    Returns:
+        Dictionary with re-estimated E0 values for each element
+    """
+    # filter configs without energy
+    database_atoms = [
+        atoms for atoms in database_atoms if atoms.info.get("REF_energy") is not None
+    ]
+
+    if len(database_atoms) == 0:
+        raise ValueError("database does not contain REF_energy tag")
+
+    if not elements:
+        elements = np.unique([Z for atoms in database_atoms for Z in atoms.numbers]).tolist()
+
+    if not initial_e0s:
+        initial_e0s = {Z: 0.0 for Z in elements}
+        try:
+            for Z in elements:
+                for i in range(len(calculator.models)):
+                    z_ind = calculator.z_table.z_to_index(Z)
+                    initial_e0s[Z] += float(
+                        calculator.models[i].atomic_energies_fn.atomic_energies[z_ind]
+                    )
+        except Exception as e:
+            logger.warning(f"unexpected exception in getting initial E0s: {e}")
+            logger.warning("falling back to explicit isolated atom calculations")
+
+            from ase.atoms import Atoms
+
+            for Z in elements:
+                calculator.calculate(
+                    atoms=Atoms(numbers=[Z], cell=[20] * 3, pbc=[True] * 3), properties=["energy"]
+                )
+                initial_e0s[Z] = calculator.results.get("energy")
+        logger.info(f"using initial E0s: {initial_e0s}")
+
+    # A matrix: each row contains atom counts for each element
+    # b vector: each entry is the prediction error for a configuration
+    A = np.zeros((len(database_atoms), len(elements)))
+    b = np.zeros(len(database_atoms))
+
+    logger.info(
+        f"solving linear system with {len(database_atoms)} equations and {len(elements)} unknowns"
+    )
+
+    # - A[i,j] is the count of element j in configuration i
+    # - b[i] is the error (true - predicted) for configuration i
+    # - x[j] will be the energy correction for element j
+    for i, atoms in enumerate(tqdm(database_atoms, desc="[*] foundation model predictions")):
+        calculator.calculate(atoms=atoms.copy(), properties=["energy"])
+        b[i] = atoms.info["REF_energy"] - calculator.results.get("energy")
+
+        # atom counts for each element
+        for j, element in enumerate(elements):
+            A[i, j] = np.sum(atoms.get_atomic_numbers() == element)
+
+    # solve with least squares
+    try:
+        corrections, _, rank, s = np.linalg.lstsq(A, b, rcond=None)
+    except np.linalg.LinAlgError as e:
+        logger.warning(f"error using lstsq to solve the linear system: {e}")
+        logger.warning("falling back to foundation model E0s")
+        return initial_e0s.copy()
+
+    if np.linalg.norm(corrections) > 1e6:
+        logger.critical(
+            f"abnormally large corrections found, rank determination may have failed: {s}"
+        )
+
+    new_e0s = {}
+    for i, element in enumerate(elements):
+        new_e0s[element] = initial_e0s[element] + corrections[i]
+        logger.debug(
+            f"element {element}: foundation E0 = {initial_e0s[element]:.4f}, "
+            f"correction = {corrections[i]:.4f}, new E0 = {new_e0s[element]:.4f}"
+        )
+
+    # statistics about the fit
+    b_after = b - A @ corrections
+    mse_before, mse_after = np.mean(b**2), np.mean(b_after**2)
+    improvement = (1 - mse_after / mse_before) * 100
+
+    logger.debug(f"mean squared error before correction: {mse_before:.4f} eV²")
+    logger.debug(f"mean squared error after correction: {mse_after:.4f} eV²")
+    logger.debug(f"improvement: {improvement:.1f}%")
+
+    if rank < len(elements):
+        logger.warning(f"system is rank deficient (rank {rank}/{len(elements)})")
+        logger.warning(
+            "some elements may be linearly dependent or not well represented in the dataset."
+        )
+
+    return new_e0s
+
+
+def gen_ft_config(
+    out: Path | str = "./config.default",
+    estimate_e0s: bool = False,
+    device: str = "cuda",
+    name: str = "fine-tuned",
+    mace_model: str = "medium-omat-0",
+    database: Path | str = "./database.extxyz",
+    head: str = "default",
+) -> None:
+    """Generate a configuration file for mace_run_train.
+
+    Args:
+        out (Path): path where the mace_run_train configuration file is written to.
+        estimate_e0s (bool): estimate the E0s for training of the foundation model.
+        device (str): device string passed to MACE to determine where calculation is performed.
+        name (str): name of the model.
+        mace_model (str): pre-trained MACE model that is used, can be a local path
+        database (Path): path to the training dataset file (likely generated with :func:`collect`)
+        head (str): which head from the model to use for prediction
+    """
+    e0s: str | dict
+
+    if Path(out).exists():
+        raise ValueError(f"Output path {out} already exists!")
+
+    if not estimate_e0s:
+        logger.warning(f'Please update the "E0s" tag in {out} prior to running mace_run_train.')
+        e0s = "TO_BE_REPLACED"
+    else:
+        from mace.calculators import mace_mp
+
+        calc = mace_mp(
+            model=mace_model,
+            dispersion=False,
+            default_dtype="float64",
+            device=device,
+            enable_cueq=(device == "cuda" and find_spec("cuequivariance") is not None),
+            head=head,
+        )
+
+        logger.info("running E0 estimation")
+        db_atoms: list[Atoms] = ase.io.read(database, ":")  # type: ignore[assignment]
+        e0s = reestimate_e0s_linear_system(calc, db_atoms)
+
+    with open(out, "w") as f:
+        f.write(
+            "\n".join(
+                [
+                    f"name: {name}",
+                    f"foundation_model: {mace_model}",
+                    f"train_file: {database}",
+                    f"valid_file: {database}",
+                    f'E0s: "{e0s}"',
+                    "multiheads_finetuning: false",
+                    "energy_weight: 1",
+                    "forces_weight: 10",
+                    "stress_weight: 0",
+                    "ema: true",
+                    "ema_decay: 0.999",
+                    "lr: 0.001",
+                    "max_num_epochs: 500",
+                    "default_dtype: float64",
+                    "batch_size: 1",
+                    "valid_batch_size: 1",
+                ]
+            )
+            + "\n"
+        )
+
+
+def parse_force_constants_file(fname: Path | str) -> np.ndarray:
+    """Parse a force constants file from phonopy."""
+    with open(fname) as f:
+        H = np.zeros([int(x) for x in next(f).split()] + [3, 3], dtype=np.float64)
+
+        for _ in range(H.shape[0] * H.shape[1]):
+            i, j = [int(x) - 1 for x in next(f).split()]
+
+            H[i, j, :, :] = np.fromstring(
+                " ".join([next(f).strip() for _ in range(3)]), sep=" "
+            ).reshape((3, 3))
+
+    H = H.swapaxes(1, 2).reshape((3 * H.shape[0], 3 * H.shape[0]))
+    return H
+
+
+def convert_from_phonopy(
+    fname: Path | str, atoms_file: Path | str, save_file: Path | str = "dynmat.npz"
+) -> None:
+    """Convert phonopy FORCE_CONSTANTS to dynmat.npz.
+
+    FORCE_CONSTANTS is written by phonopy when specifying the "--writefc" tag.
+
+    Args:
+        fname (Path): path to FORCE_CONSTANTS file.
+        atoms_file (Path): path to file containing the equilibrium structure that was used to
+            evaluate the force constants. (Only needed to extract sqrt masses.)
+        save_file (Path): path where dynamical matrix is saved to (should end in .npz)
+    """
+    logger.info(f"reading force constants from {fname}")
+    H = parse_force_constants_file(fname)
+
+    atoms: Atoms = ase.io.read(atoms_file)  # type: ignore[assignment]
+    sqrt_mass = np.repeat(np.sqrt(atoms.get_masses()), 3)
+    inv_sqrt_mass = 1 / sqrt_mass
+
+    H = np.einsum("i,ij,j->ij", inv_sqrt_mass, H, inv_sqrt_mass)
+
+    logger.info(f"saving dynamical matrix to {save_file}")
+    np.savez_compressed(save_file, H=H, sqrt_mass=sqrt_mass, cmd=" ".join(sys.argv))
+
+
+def compute_dynmat(
+    input_struct: Path,
+    save_file: Path = Path("./dynmat.npz"),
+    mace_model: str = "medium-omat-0",
+    device: str = "cuda",
+    head: str = "default",
+    relax_struct: bool = True,
+    analytical_hessian: bool = True,
+    relax_algo: str = "LBFGSLineSearch",
+    fmax: float = 0.001,
+) -> None:
+    """Calculate the dynamical matrix using MACE.
+
+    Args:
+        input_struct (Path): structure about which to compute the dynamical matrix
+        save_file (Path): path where dynamical matrix is saved to (should end in .npz)
+        mace_model (str): pre-trained MACE model that is used, can be a local path
+        device (str): device string passed to MACE to determine where calculation is performed
+        head (str): which head from the model to use for prediction
+        relax_struct (bool): determines if an atomic relaxation is performed prior to computing the
+            Hessian matrix. This is recommended if the model does not predict the same equilibrium
+            structure as your explicit DFT calculation, which is generally the case unless good
+            fine tuning has been performed.
+        analytical_hessian (bool): determines if the Hessian is computed analytically or
+            numerically using finite differences
+        relax_algo (str): name of algorithm from ase.optimize that is used for atomic relaxation.
+        fmax (float): force convergence criteria for atomic relaxation in eV/Ä
+    """
+    from mace.calculators import mace_mp
+
+    atoms: Atoms = ase.io.read(input_struct)  # type: ignore[assignment] # noqa: F823
+    atoms.calc = mace_mp(
+        model=mace_model,
+        dispersion=False,
+        default_dtype="float64",
+        device=device,
+        enable_cueq=(device == "cuda" and find_spec("cuequivariance") is not None),
+        head=head,
+    )
+
+    if relax_struct:
+        import ase.optimize as ase_optim
+
+        optim = getattr(ase_optim, relax_algo)
+        optim(atoms).run(fmax=fmax)
+
+    if np.linalg.norm(atoms.get_forces(), axis=1).max() > 0.02:
+        warnings.warn("large forces found", stacklevel=0)
+
+    if analytical_hessian:
+        H = atoms.calc.get_hessian(atoms).reshape((3 * len(atoms), 3 * len(atoms)))
+    else:
+        from ase.atoms import Atoms
+        from phonopy import Phonopy
+        from phonopy.structure.atoms import PhonopyAtoms
+
+        phonopy_atoms = PhonopyAtoms(
+            symbols=atoms.get_chemical_symbols(),
+            cell=atoms.get_cell(),
+            scaled_positions=atoms.get_scaled_positions(),
+        )
+        phonopy = Phonopy(phonopy_atoms, supercell_matrix=np.eye(3), log_level=2)
+        phonopy.generate_displacements(distance=0.02)
+
+        forces = []
+        for phonopy_atoms in tqdm(
+            phonopy.supercells_with_displacements, desc="[*] computing forces"
+        ):
+            atoms_dx = Atoms(
+                symbols=phonopy_atoms.symbols,
+                scaled_positions=phonopy_atoms.scaled_positions,
+                cell=phonopy_atoms.cell,
+                pbc=True,
+            )
+            atoms.calc.calculate(atoms=atoms_dx, properties=["forces"])
+            forces.append(atoms.calc.results["forces"])
+
+        phonopy.forces = np.array(forces)
+        phonopy.produce_force_constants()
+        phonopy.symmetrize_force_constants()
+
+        if phonopy.force_constants is None:
+            raise RuntimeError("phonopy failed to produce force constants")
+
+        H = phonopy.force_constants.swapaxes(1, 2).reshape((3 * len(atoms), 3 * len(atoms)))
+
+    if not np.allclose(H, H.T):
+        warnings.warn("Hessian matrix is not symmetric", stacklevel=0)
+
+    sqrt_mass = np.repeat(np.sqrt(atoms.get_masses()), 3)
+    inv_sqrt_mass = 1 / sqrt_mass
+    H = np.einsum("i,ij,j->ij", inv_sqrt_mass, H, inv_sqrt_mass)
+    np.savez_compressed(save_file, H=H, sqrt_mass=sqrt_mass, cmd=" ".join(sys.argv))
+
+
+def compute_lineshape(
+    ground: Path,
+    excited: Path,
+    dynmat_file: Path,
+    emission: Annotated[bool, Parameter(name="--luminescence", negative="--absorption")] = True,
+    dE: float | None = None,
+    gamma_zpl: float = 0.001,
+    sigma_zpl: float = 0.0,
+    sigma_psb: tuple[float, float] = (0.005, 0.001),
+    gamma_psb: tuple[float, float] | None = None,
+    omega_mult: float = 5.0,
+    norm: str = "area",
+    T: Annotated[float, Parameter(name=["--T", "-T"])] = 0.0,
+    plot: str | None = None,
+) -> None:
+    """Calculate the spectral density/function and lineshape for a given dynamical matrix.
+
+    Args:
+        ground (Path): path to structure containing the ground state equilibrium geometry.
+        excited (Path): path to structure containing the excited state equilibrium geometry.
+        dynmat_file (Path): path to dynamical matrix file produced by :func:`compute_dynmat` or by
+            phonopy and converted with :func:`convert_from_phonopy`.
+        emission (bool): write luminescence (True) or absorption (False) spectrum.
+        dE (float): zero-phonon line energy in eV, inferred from ground/excited if not provided.
+        gamma_zpl (float): Lorentzian broadening in the ZPL to capture homogeneous broadening.
+        sigma_zpl (float): Gaussian broadening in the ZPL to capture inhomogeneous broadening.
+        sigma_psb (float, float): Gaussian broadening used to broaden the partial Huang-Rhys
+            factors. The broadening factor is linearly interpolated from sigma_psb[0] at zero
+            frequency to sigma_psb[1] at the highest (non-LVM) frequency.
+        gamma_psb (float, float): Turns on Lorentzian broadening of local vibrational modes
+            identified by their inverse participation ratio. gamma_psb[0] is ipr_cut and
+            gamma_psb[1] is gamma_lvm. See :class:`Broadening`.
+        omega_mult (float): number of factors of maximum phonon frequency from ZPL to plot.
+        norm (str): normalization of luminescence (area or max).
+        T (float): Temperature in kelvin.
+        plot (str): if provide, specifies the type of plot to be generated and save in the current
+            working directory. Can be "subplot", "inset", "dos", "S", or "L". See
+            :func:`plot_spec_funcs` for more info.
+    """
+    initial, final = (excited, ground) if emission else (ground, excited)
+
+    ini_atoms: Atoms = ase.io.read(initial)  # type: ignore[assignment]
+    fin_atoms: Atoms = ase.io.read(final)  # type:ignore[assignment]
+
+    if dE is None:
+        dE = np.abs(fin_atoms.get_potential_energy() - ini_atoms.get_potential_energy())
+        logger.info(f"energy difference not provided, found dE = {dE} from input structures")
+
+    sqrt_mass = np.repeat(np.sqrt(fin_atoms.get_masses()), 3)
+
+    dx = get_disp_vect(fin_atoms, ini_atoms)
+    dq = sqrt_mass * dx
+    logger.info(f"dQ={np.linalg.norm(dq):.06f} amu^{{1/2}} Å")
+
+    logger.info("reading dynmat matrix")
+    data = np.load(dynmat_file)
+    H = data["H"]
+
+    # probably not needed, but doesn't hurt
+    if not np.allclose(sqrt_mass, data["sqrt_mass"]):
+        warnings.warn(
+            "sqrt_mass in dynmat file is not compatible with ground/excited conf",
+            stacklevel=0,
+        )
+
+    logger.info("diagonalizing")
+    omega, U = get_phonons(H)
+    dq_k = U.T @ dq
+
+    if gamma_psb is not None:
+        logger.info("computing inverse participation ratios")
+        ipr, ipr_cut, gamma_lvm = get_ipr(U), gamma_psb[0], gamma_psb[1]
+    else:
+        ipr, ipr_cut, gamma_lvm = None, None, None
+
+    logger.info("computing spectral functions")
+    w, dos, S, A = get_phonon_spec_func(
+        dq_k,
+        omega,
+        sigma_psb=sigma_psb,
+        gamma_zpl=gamma_zpl,
+        sigma_zpl=sigma_zpl,
+        gamma_lvm=gamma_lvm,
+        ipr=ipr,
+        ipr_cut=ipr_cut,
+        T=T,
+    )
+    tw, L = convert_A_to_L(w, A, dE, emission=emission, norm=norm)
+
+    logger.info("saving results to .txt files")
+    np.savetxt("spec_funcs.txt", np.array((w, S, A)).T)
+    np.savetxt("lineshape.txt", np.array((tw, L)).T)
+
+    if plot is not None:
+        import matplotlib.pyplot as plt
+
+        plot_spec_funcs(
+            (w, dos, S, tw, L),
+            None,
+            dE,
+            emission=emission,
+            omega_mult=omega_mult,
+            omega_max=(omega2eV * omega.max()),
+            plot_type=plot,
+        )
+        plt.savefig("lineshape.png", dpi=600, bbox_inches="tight")
+
+
+def analyze_dynmat(dynmat: Path, structure: Path) -> None:
+    """Produce analysis plots of the dynamical matrix."""
+    H = np.load(dynmat)["H"]
+
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(4, 3))
+    p = ax.imshow(np.log(np.abs(H)), vmin=-6)
+    plt.colorbar(p)
+    plt.savefig("H.png", dpi=600, bbox_inches="tight")
+
+    atoms: Atoms = ase.io.read(structure)  # type: ignore[assignment]
+    d = atoms.get_all_distances(mic=True)
+
+    fig, ax = plt.subplots(figsize=(4, 3))
+    for i in range(3):
+        for j in range(3):
+            ax.plot(d.flatten(), np.abs(H[i::3, j::3].flatten()), ".", ms=1)
+    ax.set_yscale("log")
+    ax.set_xlabel(r"$\vert {\bf R}_I - {\bf R}_J \vert$ [${\rm \AA}$]")
+    ax.set_ylabel(r"$\vert \Phi_{I\alpha,J\beta} \vert$ [eV/amu$^{1/2}$ ${\rm \AA}$]")
+    plt.savefig("radial_H.png", dpi=600, bbox_inches="tight")