API Reference

Transforms

mixture

Structure generation for amorphous solid and liquid mixtures.

Provides :func:mix_number and :func:mix_cell for building multicomponent mixtures using Packmol packing with crystal structure prototypes retrieved from the Materials Project.

mix_number(recipe: dict[str, int], density: float | None = None, tolerance: float = 2.0, rattle: float = 0.5, scale: float = 1.0, shuffle: bool = False, seed: int = 1, timeout: int = 30, log: bool = False, mp_api_key: str | None = MP_API_KEY, retry: int = 1000, retry_scale: float = 1.5) -> Atoms

Build a mixture structure by specifying formula unit counts.

Retrieves primitive crystal structures from Materials Project for each component, then uses Packmol to pack the specified number of formula units into a cubic simulation cell. The cell size is estimated from the sum of solid-state primitive cell volumes.

Parameters:

Name	Type	Description	Default
recipe	dict[str, int]	Mapping of chemical formula strings to the desired number of formula units (e.g., {"NaCl": 3, "KCl": 1}).	required
density	float | None	Target mass density in amu/ų. If provided, the cell is rescaled after packing. Defaults to None (use solid-state volume).	None
tolerance	float	Minimum distance between packed molecules in Å. Defaults to 2.0.	2.0
rattle	float	Standard deviation of Gaussian noise added to atomic positions in Å. Defaults to 0.5.	0.5
scale	float	Multiplicative factor for the estimated cubic cell edge. Defaults to 1.0.	1.0
shuffle	bool	If True, randomly permute atomic species numbers. Defaults to False.	False
seed	int	Random seed for Packmol and numpy. Defaults to 1.	1
timeout	int	Packmol timeout in seconds. Defaults to 30.	30
log	bool	If True, print diagnostic information. Defaults to False.	False
mp_api_key	str | None	Materials Project API key. Defaults to the MP_API_KEY environment variable.	MP_API_KEY
retry	int	Number of Packmol attempts before enlarging the box. Defaults to 1000.	1000
retry_scale	float	Factor by which to enlarge the box after retry failures. Defaults to 1.5.	1.5

Returns:

Type	Description
Atoms	Sorted ASE Atoms object with periodic boundary conditions.

Source code in muse/transforms/mixture.py

def mix_number(
    recipe: dict[str, int],
    density: float | None = None,
    tolerance: float = 2.0,
    rattle: float = 0.5,
    scale: float = 1.0,
    shuffle: bool = False,
    seed: int = 1,
    timeout: int = 30,
    log: bool = False,
    mp_api_key: str | None = MP_API_KEY,
    retry: int = 1000,
    retry_scale: float = 1.5,
) -> Atoms:
    """Build a mixture structure by specifying formula unit counts.

    Retrieves primitive crystal structures from Materials Project for each
    component, then uses Packmol to pack the specified number of formula
    units into a cubic simulation cell. The cell size is estimated from
    the sum of solid-state primitive cell volumes.

    Args:
        recipe: Mapping of chemical formula strings to the desired number
            of formula units (e.g., ``{"NaCl": 3, "KCl": 1}``).
        density: Target mass density in amu/ų. If provided, the cell is
            rescaled after packing. Defaults to None (use solid-state volume).
        tolerance: Minimum distance between packed molecules in Å.
            Defaults to 2.0.
        rattle: Standard deviation of Gaussian noise added to atomic
            positions in Å. Defaults to 0.5.
        scale: Multiplicative factor for the estimated cubic cell edge.
            Defaults to 1.0.
        shuffle: If True, randomly permute atomic species numbers.
            Defaults to False.
        seed: Random seed for Packmol and numpy. Defaults to 1.
        timeout: Packmol timeout in seconds. Defaults to 30.
        log: If True, print diagnostic information. Defaults to False.
        mp_api_key: Materials Project API key. Defaults to the
            ``MP_API_KEY`` environment variable.
        retry: Number of Packmol attempts before enlarging the box.
            Defaults to 1000.
        retry_scale: Factor by which to enlarge the box after ``retry``
            failures. Defaults to 1.5.

    Returns:
        Sorted ASE Atoms object with periodic boundary conditions.
    """
    mpr = MPRester(mp_api_key)
    rng = np.random.default_rng(seed)

    molecules = []

    for formula, units in recipe.items():
        if units == 0:
            continue

        reduced_formula, input_mult = Formula(formula).reduce()

        docs = mpr.materials.summary.search(
            formula=str(reduced_formula),
            is_stable=True,
            fields=["material_id", "formula_pretty", "structure"],
        )

        sga = SpacegroupAnalyzer(docs[0].structure)
        primitive_structure = sga.get_primitive_standard_structure()

        primitive_formula = Formula(primitive_structure.composition.to_pretty_string())

        molecule = Molecule(
            species=primitive_structure.species, coords=primitive_structure.cart_coords
        )
        _, primitive_mult = primitive_formula.reduce()

        number: float = 0.0
        count: int = 0
        while number == 0 or not math.isclose(number, round(number), rel_tol=1e-3):
            if count > 0:
                for key, value in recipe.items():
                    recipe[key] = value / count * (count + 1)

                for d in molecules:
                    d["number"] = d["number"] / count * (count + 1)

            number = input_mult * recipe[formula] / primitive_mult
            if log:
                print(recipe[formula], Formula(formula), number, primitive_formula)

            count += 1

        molecules.append(
            {
                "name": primitive_structure.composition.to_pretty_string(),
                "number": int(number),
                "coords": molecule,
                "volume": primitive_structure.volume,
            }
        )

    if log:
        print(molecules)

    total_volume = 0
    for molecule in molecules:
        total_volume += molecule["volume"] * molecule["number"]

    a = total_volume ** (1.0 / 3.0) * scale

    with ScratchDir("."):
        while True:
            try:
                input_gen = PackmolBoxGen(
                    tolerance=tolerance,
                    seed=seed,
                )

                margin = 0.5 * tolerance
                packmol_set = input_gen.get_input_set(
                    molecules=molecules,
                    box=[margin, margin, margin, a - margin, a - margin, a - margin],
                )
                packmol_set.write_input(".")
                packmol_set.run(".", timeout=timeout)

                atoms = read("packmol_out.xyz", format="xyz")
                break
            except Exception as e:
                if log:
                    print(e)
                seed += 1

                if a > total_volume ** (1.0 / 3.0) * scale * 2:
                    if log:
                        logger.warning(
                            "Box size was increased by more than 2x. "
                            "Generating random structure."
                        )

                    a = total_volume ** (1.0 / 3.0) * scale

                    symbols = ""
                    for molecule in molecules:
                        symbols += Formula(molecule["name"]) * int(molecule["number"])

                    atoms = Atoms(symbols=symbols, cell=[a, a, a], pbc=True)
                    atoms = sort(atoms)

                    atoms.set_scaled_positions(rng.random(size=atoms.positions.shape))
                    break

                if seed % retry == 0:
                    a *= retry_scale
                    if log:
                        logger.warning(
                            "Packmol failed %d times. Trying again with larger box. "
                            "New box size: %s",
                            retry,
                            a,
                        )

    atoms.set_cell([a, a, a])
    atoms.set_pbc(True)

    if a != total_volume ** (1.0 / 3.0) * scale:
        if log:
            logger.warning("Box size was increased. Shrinking to the designated size.")
        scaled_positions = atoms.get_scaled_positions()
        a = total_volume ** (1.0 / 3.0) * scale
        atoms.set_cell([a, a, a])
        atoms.set_scaled_positions(scaled_positions)

    if rattle > 0:
        atoms.positions += rng.normal(0, rattle, size=atoms.positions.shape)

    if shuffle:
        atoms.numbers = rng.permutation(atoms.numbers)

    if density is not None:
        cellpar = atoms.cell.cellpar()

        volume = atoms.get_masses().sum() / density

        factor = (volume / atoms.get_volume()) ** (1.0 / 3.0)
        atoms.set_cell(
            Cell.fromcellpar(
                [
                    cellpar[0] * factor,
                    cellpar[1] * factor,
                    cellpar[2] * factor,
                    cellpar[3],
                    cellpar[4],
                    cellpar[5],
                ]
            ),
            scale_atoms=True,
        )

    return sort(atoms)

mix_cell(recipe: dict[str, float], cell: Cell, tolerance: float = 2.0, rattle: float = 0.5, scale: float = 1.0, shuffle: bool = True, seed: int = 1, log: bool = False, mp_api_key: str | None = MP_API_KEY, retry_scale: float = 1.5) -> Atoms

Build a mixture structure to fill a given simulation cell.

Similar to :func:mix_number, but instead of specifying absolute formula unit counts, the recipe specifies molar fractions and the total number of atoms is determined by the target cell volume.

The function scales the number of molecules to fill the provided cell based on the ratio of cell volume to the total solid-state volume of the components.

Parameters:

Name	Type	Description	Default
recipe	dict[str, float]	Mapping of chemical formula strings to molar ratios (e.g., {"NaCl": 0.7, "KCl": 0.3}).	required
cell	Cell	Target ASE Cell object defining the simulation box shape.	required
tolerance	float	Minimum distance between packed molecules in Å. Defaults to 2.0.	2.0
rattle	float	Standard deviation of Gaussian noise added to atomic positions in Å. Defaults to 0.5.	0.5
scale	float	Multiplicative factor for cell dimensions during packing. Defaults to 1.0.	1.0
shuffle	bool	If True, randomly permute atomic species numbers. Defaults to True.	True
seed	int	Random seed for Packmol and numpy. Defaults to 1.	1
log	bool	If True, print diagnostic information. Defaults to False.	False
mp_api_key	str | None	Materials Project API key. Defaults to the MP_API_KEY environment variable.	MP_API_KEY
retry_scale	float	Factor by which to enlarge the box after 1000 Packmol failures. Defaults to 1.5.	1.5

Returns:

Type	Description
Atoms	Sorted ASE Atoms object with periodic boundary conditions
Atoms	matching the target cell.

Source code in muse/transforms/mixture.py

def mix_cell(
    recipe: dict[str, float],
    cell: Cell,
    tolerance: float = 2.0,
    rattle: float = 0.5,
    scale: float = 1.0,
    shuffle: bool = True,
    seed: int = 1,
    log: bool = False,
    mp_api_key: str | None = MP_API_KEY,
    retry_scale: float = 1.5,
) -> Atoms:
    """Build a mixture structure to fill a given simulation cell.

    Similar to :func:`mix_number`, but instead of specifying absolute
    formula unit counts, the recipe specifies molar fractions and the
    total number of atoms is determined by the target cell volume.

    The function scales the number of molecules to fill the provided
    cell based on the ratio of cell volume to the total solid-state
    volume of the components.

    Args:
        recipe: Mapping of chemical formula strings to molar ratios
            (e.g., ``{"NaCl": 0.7, "KCl": 0.3}``).
        cell: Target ASE Cell object defining the simulation box shape.
        tolerance: Minimum distance between packed molecules in Å.
            Defaults to 2.0.
        rattle: Standard deviation of Gaussian noise added to atomic
            positions in Å. Defaults to 0.5.
        scale: Multiplicative factor for cell dimensions during packing.
            Defaults to 1.0.
        shuffle: If True, randomly permute atomic species numbers.
            Defaults to True.
        seed: Random seed for Packmol and numpy. Defaults to 1.
        log: If True, print diagnostic information. Defaults to False.
        mp_api_key: Materials Project API key. Defaults to the
            ``MP_API_KEY`` environment variable.
        retry_scale: Factor by which to enlarge the box after 1000
            Packmol failures. Defaults to 1.5.

    Returns:
        Sorted ASE Atoms object with periodic boundary conditions
        matching the target cell.
    """
    mpr = MPRester(mp_api_key)
    rng = np.random.default_rng(seed)

    molecules = []

    for formula, units in recipe.items():
        if units == 0:
            continue

        reduced_formula, input_mult = Formula(formula).reduce()

        docs = mpr.materials.summary.search(
            formula=str(reduced_formula),
            is_stable=True,
            fields=["material_id", "formula_pretty", "structure"],
        )

        sga = SpacegroupAnalyzer(docs[0].structure)
        primitive_structure = sga.get_primitive_standard_structure()

        primitive_formula = Formula(primitive_structure.composition.to_pretty_string())

        molecule = Molecule(
            species=primitive_structure.species, coords=primitive_structure.cart_coords
        )
        _, primitive_mult = primitive_formula.reduce()

        number: float = 0.0
        count: int = 0
        while number == 0 or not number.is_integer():
            if count > 0:
                for key, value in recipe.items():
                    recipe[key] = value / count * (count + 1)

                for d in molecules:
                    d["number"] = d["number"] / count * (count + 1)

            number = input_mult * recipe[formula] / primitive_mult
            if log:
                print(recipe[formula], Formula(formula), number, primitive_formula)

            count += 1

        molecules.append(
            {
                "name": primitive_structure.composition.to_pretty_string(),
                "number": int(number),
                "coords": molecule,
                "volume": primitive_structure.volume,
            }
        )

    total_volume = 0
    for molecule in molecules:
        total_volume += molecule["volume"] * molecule["number"]

    nfactor = cell.volume / total_volume

    for molecule in molecules:
        molecule["number"] = int(molecule["number"] * nfactor)

    if log:
        print(molecules)

    a, b, c, alpha, beta, gamma = cell.cellpar()

    with ScratchDir("."):
        while True:
            try:
                input_gen = PackmolBoxGen(
                    tolerance=tolerance,
                    seed=seed,
                )

                margin = 0.5 * tolerance
                packmol_set = input_gen.get_input_set(
                    molecules=molecules,
                    box=[margin, margin, margin, a - margin, b - margin, c - margin],
                )
                packmol_set.write_input(".")
                packmol_set.run(".")

                atoms = read("packmol_out.xyz", format="xyz")
                break
            except Exception as e:
                if log:
                    print(e)
                seed += 1

                if a > cell.volume ** (1.0 / 3.0) * scale * 2:
                    if log:
                        logger.warning(
                            "Box size was increased by more than 2x. "
                            "Generating random structure."
                        )

                    a, b, c, alpha, beta, gamma = cell.cellpar()
                    atoms = Atoms(cell=cell, pbc=True)
                    atoms.set_scaled_positions(rng.random(size=atoms.positions.shape))
                    break

                if seed % 1000 == 0:
                    a *= retry_scale
                    if log:
                        logger.warning(
                            "Packmol failed 1000 times. Trying again with larger box. "
                            "New box size: %s",
                            a,
                        )

    atoms.set_cell(cell)
    atoms.set_pbc(True)

    if rattle > 0:
        atoms.positions += rng.normal(0, rattle, size=atoms.positions.shape)

    if shuffle:
        atoms.numbers = rng.permutation(atoms.numbers)

    return sort(atoms)

Calculations

density

Calculator for density-related properties via NPT molecular dynamics.

DensityCalc(calculator: Calculator, optimizer: Optimizer | str = 'FIRE', steps: int = 500, interval: int = 1, fmax: float = 0.1, mask: list | np.ndarray | None = None, rtol: float = 0.0001, atol: float = 0.0001, out_stem: str | Path = '.')

Bases: PropCalc

Relax and run NPT simulations to compute the equilibrium density.

This calculator performs a three-stage molecular dynamics workflow:

1. 0 K relaxation — Minimize forces with the chosen optimizer.
2. NVT equilibration — Thermalize at the target temperature with fixed volume until energy converges.
3. NPT production — Allow both temperature and pressure to equilibrate, then compute density from the final volume.

Parameters:

Name	Type	Description	Default
calculator	Calculator	ASE calculator to use for energy/force evaluation.	required
optimizer	Optimizer | str	ASE optimizer class or name string. Defaults to "FIRE".	'FIRE'
steps	int	Number of MD steps per convergence window. Defaults to 500.	500
interval	int	Trajectory save interval in steps. Defaults to 1.	1
fmax	float	Maximum force for structural relaxation (eV/Å). Defaults to 0.1.	0.1
mask	list | ndarray | None	3×3 mask array controlling which cell degrees of freedom are relaxed in the NPT barostat. Defaults to None (all free).	None
rtol	float	Relative tolerance for energy convergence between windows. Defaults to 1e-4.	0.0001
atol	float	Absolute tolerance for stress convergence (eV/ų). Defaults to 1e-4.	0.0001
out_stem	str | Path	Path stem for saving trajectory and observer files. Defaults to ".".	'.'

Source code in muse/calcs/density.py

def __init__(
    self,
    calculator: Calculator,
    optimizer: Optimizer | str = "FIRE",
    steps: int = 500,
    interval: int = 1,
    fmax: float = 0.1,
    mask: list | np.ndarray | None = None,
    rtol: float = 1e-4,
    atol: float = 1e-4,
    out_stem: str | Path = ".",
):
    self.calculator = calculator

    # check str is valid optimizer key
    def is_ase_optimizer(key):
        return isclass(obj := getattr(optimize, key)) and issubclass(obj, Optimizer)

    valid_keys = [key for key in dir(optimize) if is_ase_optimizer(key)]
    if isinstance(optimizer, str) and optimizer not in valid_keys:
        raise ValueError(f"Unknown {optimizer=}, must be one of {valid_keys}")

    self.optimizer: Optimizer = (
        getattr(optimize, optimizer) if isinstance(optimizer, str) else optimizer
    )
    self.steps = steps
    self.interval = interval
    self.fmax = fmax
    self.mask = mask
    self.rtol = rtol
    self.atol = atol
    self.out_stem = out_stem

calc(atoms: Atoms, temperature: float, externalstress: float | np.ndarray, timestep: float = 2.0 * units.fs, ttime: float = 25.0 * units.fs, pfactor: float = (75 * units.fs) ** 2 * units.GPa, annealing: float = 1.0, momentum: float = 0.9) -> dict

Relax the structure and run NPT simulations to compute the density.

Parameters:

Name	Type	Description	Default
atoms	Atoms	Structure to relax and equilibrate.	required
temperature	float	Temperature of the simulation in Kelvin.	required
externalstress	float | ndarray	External pressure in eV/ų (scalar for isotropic, or Voigt 6-vector).	required
timestep	float	MD timestep in ASE internal units. Defaults to 2.0 fs.	2.0 * fs
ttime	float	Thermostat characteristic timescale in ASE internal units. Defaults to 25.0 fs.	25.0 * fs
pfactor	float	Barostat constant in ASE internal units. Defaults to (75 fs)² × 1 GPa.	(75 * fs) ** 2 * GPa
annealing	float	Temperature scaling factor for NVT initialization. Values > 1 start hotter to aid equilibration. Defaults to 1.0.	1.0
momentum	float	Exponential moving average factor for energy convergence tracking between NVT/NPT windows. Defaults to 0.9.	0.9

Returns:

Type	Description
dict	Dictionary with keys: - volume_avg: Mean cell volume (ų). - volume_std: Standard deviation of cell volume. - atomic_density: Number density (atoms/ų). - mass_density: Mass density (amu/ų). - energy_avg: Mean potential energy (eV). - energy_std: Standard deviation of potential energy.

Source code in muse/calcs/density.py

def calc(
    self,
    atoms: Atoms,
    temperature: float,
    externalstress: float | np.ndarray,
    timestep: float = 2.0 * units.fs,
    ttime: float = 25.0 * units.fs,
    pfactor: float = (75 * units.fs) ** 2 * units.GPa,
    annealing: float = 1.0,
    momentum: float = 0.9,
) -> dict:
    """Relax the structure and run NPT simulations to compute the density.

    Args:
        atoms: Structure to relax and equilibrate.
        temperature: Temperature of the simulation in Kelvin.
        externalstress: External pressure in eV/ų (scalar for isotropic,
            or Voigt 6-vector).
        timestep: MD timestep in ASE internal units. Defaults to 2.0 fs.
        ttime: Thermostat characteristic timescale in ASE internal units.
            Defaults to 25.0 fs.
        pfactor: Barostat constant in ASE internal units.
            Defaults to (75 fs)² × 1 GPa.
        annealing: Temperature scaling factor for NVT initialization.
            Values > 1 start hotter to aid equilibration. Defaults to 1.0.
        momentum: Exponential moving average factor for energy convergence
            tracking between NVT/NPT windows. Defaults to 0.9.

    Returns:
        Dictionary with keys:
            - ``volume_avg``: Mean cell volume (ų).
            - ``volume_std``: Standard deviation of cell volume.
            - ``atomic_density``: Number density (atoms/ų).
            - ``mass_density``: Mass density (amu/ų).
            - ``energy_avg``: Mean potential energy (eV).
            - ``energy_std``: Standard deviation of potential energy.
    """
    # relax the structure
    atoms.calc = self.calculator

    stream = io.StringIO()

    # step 0: relax at 0 K
    with contextlib.redirect_stdout(stream):
        optimizer = self.optimizer(atoms)

        if self.out_stem is not None:
            traj = Trajectory(f"{self.out_stem}-relax.traj", "w", atoms)
            optimizer.attach(traj.write, interval=self.interval)

        obs = TrajectoryObserver(atoms)
        optimizer.attach(obs, interval=self.interval)
        optimizer.run(fmax=self.fmax, steps=self.steps)

    if self.out_stem is not None:
        traj.close()
        obs()
        obs.save(f"{self.out_stem}-relax.pkl")
        del obs

    # step 1: run NVT simulation
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature * annealing)
    Stationary(atoms, preserve_temperature=True)

    nvt = NPT(
        atoms,
        timestep=timestep,
        temperature_K=temperature * annealing,
        externalstress=externalstress,
        ttime=ttime,
        pfactor=None,
    )
    nvt.set_fraction_traceless(0.0)

    converged, erg_converged, str_converged = False, False, False
    restart = 0
    last_erg_avg, first_erg_avg = None, None
    while not converged:
        if self.out_stem is not None:
            traj = Trajectory(f"{self.out_stem}-nvt-{restart}.traj", "w", atoms)
            nvt.attach(traj.write, interval=self.interval)

        obs = TrajectoryObserver(atoms)
        nvt.attach(obs, interval=self.interval)
        nvt.run(steps=self.steps)

        erg_avg, erg_std = np.mean(obs.energies), np.std(obs.energies)

        if last_erg_avg is None or first_erg_avg is None:
            last_erg_avg = erg_avg
            first_erg_avg = erg_avg
            erg_converged = False
        else:
            erg_converged = (
                abs(erg_avg - last_erg_avg) / last_erg_avg < self.rtol
                and np.sign(erg_avg - first_erg_avg) * (erg_avg - last_erg_avg) < 0
            )

        stress = np.mean(np.stack(obs.stresses, axis=0), axis=0)
        str_converged = np.allclose(
            nvt.externalstress, stress, atol=self.atol, rtol=self.rtol
        )

        converged = erg_converged  # and str_converged

        if self.out_stem is not None:
            traj.close()
            obs()
            obs.save(f"{self.out_stem}-nvt-{restart}.pkl")
            del obs

        if not converged:
            logger.info(
                "NVT - %d: Energy or stress not converged, restarting simulation.",
                restart,
            )
            if not erg_converged:
                logger.info(
                    "Current relative energy deviation: %.4f %%. Target: %.4f %%.",
                    (erg_avg - last_erg_avg) / last_erg_avg * 100,
                    self.rtol * 100,
                )
            else:
                logger.info("Energy converged.")
            if not str_converged:
                logger.info("Current pressure: %s eV/ų.", stress)
                logger.info("Target pressure: %s eV/ų.", nvt.externalstress)
            else:
                logger.info("Pressure converged.")

            nvt.observers.clear()
            last_erg_avg = momentum * erg_avg + (1 - momentum) * last_erg_avg
            restart += 1

    # step 2: run NPT simulation
    npt = NPT(
        atoms,
        timestep=timestep,
        temperature_K=temperature,
        externalstress=externalstress,
        ttime=ttime,
        pfactor=pfactor,
        mask=self.mask,
    )
    if np.array_equal(self.mask, np.eye(3)):
        npt.set_fraction_traceless(0.0)  # fix shape

    converged, erg_converged, str_converged = False, False, False
    restart = 0
    last_erg_avg, first_erg_avg = None, None
    while not converged:
        if self.out_stem is not None:
            self.final_traj_fpath = f"{self.out_stem}-npt-{restart}.traj"
            traj = Trajectory(self.final_traj_fpath, "w", atoms)
            npt.attach(traj.write, interval=self.interval)

        obs = TrajectoryObserver(atoms)
        npt.attach(obs, interval=self.interval)
        npt.run(steps=self.steps)

        erg_avg, erg_std = np.mean(obs.energies), np.std(obs.energies)

        if last_erg_avg is None or first_erg_avg is None:
            last_erg_avg = erg_avg
            first_erg_avg = erg_avg
            erg_converged = False
        else:
            erg_converged = (
                abs(erg_avg - last_erg_avg) / last_erg_avg < self.rtol
                and np.sign(erg_avg - first_erg_avg) * (erg_avg - last_erg_avg) < 0
            )

        stress = np.mean(np.stack(obs.stresses, axis=0), axis=0)
        str_converged = np.allclose(
            npt.externalstress, stress, atol=self.atol, rtol=self.rtol
        )

        converged = erg_converged and str_converged

        if self.out_stem is not None:
            traj.close()
            obs()
            obs.save(f"{self.out_stem}-npt-{restart}.pkl")

        if not converged:
            logger.info(
                "NPT - %d: Energy or stress not converged, restarting simulation.",
                restart,
            )
            if not erg_converged:
                logger.info(
                    "Current relative energy deviation: %.4f %%. Target: %.4f %%.",
                    (erg_avg - last_erg_avg) / last_erg_avg * 100,
                    self.rtol * 100,
                )
            else:
                logger.info("Energy converged.")
            if not str_converged:
                logger.info("Current pressure: %s eV/ų.", stress)
                logger.info("Target pressure: %s eV/ų.", npt.externalstress)
            else:
                logger.info("Pressure converged.")

            npt.observers.clear()
            last_erg_avg = momentum * erg_avg + (1 - momentum) * last_erg_avg
            restart += 1

    volumes = [Cell(matrix).volume for matrix in obs.cells]
    vol_avg, vol_std = np.mean(volumes), np.std(volumes)
    erg_avg, erg_std = np.mean(obs.energies), np.std(obs.energies)

    return {
        "volume_avg": vol_avg,
        "volume_std": vol_std,
        "atomic_density": atoms.get_global_number_of_atoms() / vol_avg,
        "mass_density": atoms.get_masses().sum() / vol_avg,
        "energy_avg": erg_avg,
        "energy_std": erg_std,
    }

utils

Trajectory observation utilities for MD simulations.

Provides :class:TrajectoryObserver, a callback hook that records energies, forces, stresses, positions, and cell parameters during ASE relaxations and molecular dynamics runs.

TrajectoryObserver(atoms: Atoms)

Trajectory observer that records simulation data at each step.

Attach this observer to an ASE optimizer or dynamics object to capture per-step energies, forces, stresses, positions, and cell matrices. Data can be serialized to a pickle file for post-processing.

This class is adapted from matcalc <https://github.com/materialsvirtuallab/matcalc>_.

Parameters:

Name	Type	Description	Default
atoms	Atoms	The ASE Atoms object to observe.	required

Source code in muse/calcs/utils.py

def __init__(self, atoms: Atoms) -> None:
    self.atoms = atoms
    self.energies: list[float] = []
    self.forces: list[np.ndarray] = []
    self.stresses: list[np.ndarray] = []
    self.atom_positions: list[np.ndarray] = []
    self.cells: list[np.ndarray] = []

__call__() -> None

Record the current state of the Atoms object.

Captures potential energy, forces, stress tensor (including ideal gas contribution when available), positions, and cell matrix.

Source code in muse/calcs/utils.py

def __call__(self) -> None:
    """Record the current state of the Atoms object.

    Captures potential energy, forces, stress tensor (including ideal
    gas contribution when available), positions, and cell matrix.
    """
    self.energies.append(float(self.atoms.get_potential_energy()))
    self.forces.append(self.atoms.get_forces())
    # Stress tensor should include the contribution from the momenta, otherwise
    # during MD simulation the stress tensor ignores the effect of kinetic part,
    # leading to the discrepancy between applied pressure and the stress tensor.
    # For more details, see: https://gitlab.com/ase/ase/-/merge_requests/1500
    try:
        stress = self.atoms.get_stress(include_ideal_gas=True)
    except TypeError:
        stress = self.atoms.get_stress()
    self.stresses.append(stress)

    self.atom_positions.append(self.atoms.get_positions())
    self.cells.append(self.atoms.get_cell()[:])

save(filename: str | Path) -> None

Save the recorded trajectory data to a pickle file.

The output dictionary contains

• energy: List of potential energies (eV).
• forces: List of force arrays (eV/Å).
• stresses: List of stress tensors (eV/ų Voigt).
• atom_positions: List of position arrays (Å).
• cell: List of cell matrices (Å).
• atomic_number: Array of atomic numbers.

Parameters:

Name	Type	Description	Default
filename	str | Path	Path to the output pickle file.	required

Source code in muse/calcs/utils.py

def save(self, filename: str | Path) -> None:
    """Save the recorded trajectory data to a pickle file.

    The output dictionary contains:
        - ``energy``: List of potential energies (eV).
        - ``forces``: List of force arrays (eV/Å).
        - ``stresses``: List of stress tensors (eV/ų Voigt).
        - ``atom_positions``: List of position arrays (Å).
        - ``cell``: List of cell matrices (Å).
        - ``atomic_number``: Array of atomic numbers.

    Args:
        filename: Path to the output pickle file.
    """
    out = {
        "energy": self.energies,
        "forces": self.forces,
        "stresses": self.stresses,
        "atom_positions": self.atom_positions,
        "cell": self.cells,
        "atomic_number": self.atoms.get_atomic_numbers(),
    }
    with open(filename, "wb") as file:
        pickle.dump(out, file)

Plotting

density

Binary density–composition diagram with Redlich–Kister curve fitting.

BinaryDXDiagram(fig: Figure, *args: Any, facecolor=None, frameon: bool = True, sharex: Axes | None = None, sharey: Axes | None = None, label: str = '', xscale: float | None = None, yscale: float | None = None, box_aspect: float | None = None, **kwargs)

Bases: Axes

Custom Matplotlib Axes for binary density–composition (D–x) diagrams.

Processes MD trajectories at various compositions to compute density and molar volume, then plots the results with an optional Redlich–Kister polynomial fit for the excess property.

Source code in muse/plots/density.py

def __init__(
    self,
    fig: Figure,
    *args: Any,
    facecolor=None,
    frameon: bool = True,
    sharex: Axes | None = None,
    sharey: Axes | None = None,
    label: str = "",
    xscale: float | None = None,
    yscale: float | None = None,
    box_aspect: float | None = None,
    **kwargs,
) -> None:
    super().__init__(
        fig,
        args,
        facecolor=facecolor,
        frameon=frameon,
        sharex=sharex,
        sharey=sharey,
        label=label,
        xscale=xscale,
        yscale=yscale,
        box_aspect=box_aspect,
        **kwargs,
    )

process(trajectories: Sequence[Sequence[Atoms]], phases: Sequence[str | Formula]) -> None

Process MD trajectories to extract density and volume statistics.

Computes mass density (g/cm³) and molar volume (ų/formula unit) for each trajectory, storing the results sorted by composition.

Parameters:

Name	Type	Description	Default
trajectories	Sequence[Sequence[Atoms]]	List of MD trajectories, each a sequence of Atoms.	required
phases	Sequence[str | Formula]	Two-element list of phase formulas defining the binary system.	required

Source code in muse/plots/density.py

def process(
    self,
    trajectories: Sequence[Sequence[Atoms]],
    phases: Sequence[str | Formula],
) -> None:
    """Process MD trajectories to extract density and volume statistics.

    Computes mass density (g/cm³) and molar volume (ų/formula unit)
    for each trajectory, storing the results sorted by composition.

    Args:
        trajectories: List of MD trajectories, each a sequence of Atoms.
        phases: Two-element list of phase formulas defining the binary system.
    """
    # Change strings to Formula objects and sort symbols
    for phase in phases:
        phase = Formula.from_list(phase) if isinstance(phase, str) else sort(phase)

    x, denavgs, denstds, volavgs, volstds = [], [], [], [], []
    for traj in trajectories:
        atoms = traj[0]
        formula = atoms.symbols.formula

        portions = {}
        for phase in phases:
            portions[phase] = formula // phase

        total_units = sum(portions.values())

        fractions = {}
        for phase in phases:
            fractions[phase] = portions[phase] / total_units

        x.append(fractions[phases[-1]])

        densities = []
        volumes = []
        for atoms in traj:
            densities.append(
                atoms.get_masses().sum()
                * 1.66054e-24
                / (atoms.get_volume() * 1e-24)
            )
            volumes.append(atoms.get_volume() / total_units)

        denavgs.append(np.mean(densities))
        denstds.append(np.std(densities))
        volavgs.append(np.mean(volumes))
        volstds.append(np.std(volumes))

    x = np.array(x)
    idx = np.argsort(x)
    self.x = x[idx]

    self.y = {} if getattr(self, "y", None) is None else self.y
    self.y["density.avg"] = np.array(denavgs)[idx]
    self.y["density.std"] = np.array(denstds)[idx]
    self.y["volume.avg"] = np.array(volavgs)[idx]
    self.y["volume.std"] = np.array(volstds)[idx]

from_trajectories(trajectories: Sequence[Sequence[Atoms]], phases: Sequence[str | Formula], temperature: float | None = None, label: str | None = None, rk: int = 2, **kwargs) -> None

Plot a binary density–composition diagram from MD trajectories.

Computes densities from the trajectories, plots them with error bars, and overlays a Redlich–Kister polynomial fit.

Parameters:

Name	Type	Description	Default
trajectories	Sequence[Sequence[Atoms]]	List of MD trajectories at different compositions.	required
phases	Sequence[str | Formula]	Two-element list of phase formulas defining the binary system.	required
temperature	float | None	Temperature in Kelvin for the Redlich–Kister fit. Defaults to 1000 K.	None
label	str | None	Label prefix for the legend entries.	None
rk	int	Number of Redlich–Kister terms to use in the fit. Defaults to 2.	2
**kwargs		Additional keyword arguments passed to errorbar.	{}

Source code in muse/plots/density.py

def from_trajectories(
    self,
    trajectories: Sequence[Sequence[Atoms]],
    phases: Sequence[str | Formula],
    temperature: float | None = None,
    label: str | None = None,
    rk: int = 2,
    **kwargs,
) -> None:
    """Plot a binary density–composition diagram from MD trajectories.

    Computes densities from the trajectories, plots them with error bars,
    and overlays a Redlich–Kister polynomial fit.

    Args:
        trajectories: List of MD trajectories at different compositions.
        phases: Two-element list of phase formulas defining the binary system.
        temperature: Temperature in Kelvin for the Redlich–Kister fit. Defaults to 1000 K.
        label: Label prefix for the legend entries.
        rk: Number of Redlich–Kister terms to use in the fit. Defaults to 2.
        **kwargs: Additional keyword arguments passed to ``errorbar``.
    """
    assert len(phases) == 2

    self.process(trajectories, phases)

    color = kwargs.pop("color", "k")

    self.errorbar(
        self.x,
        self.y["density.avg"],
        yerr=self.y["density.std"],
        label=f"{label}: $\\rho_m$" if label else "$\\rho_m$",
        color=color,
        fmt="o",
        **kwargs,
    )

    T = temperature or 1000

    # Fitting the Redlich-Kister expansion model
    initial_guess = [0.0] * (2 * rk)

    y = self.y["density.avg"]

    params_opt, params_cov = curve_fit(
        lambda x_T, *params: redlich_kister_model(x_T[0], x_T[1], *params),
        (self.x, np.ones_like(self.x) * T),
        y - (y[0] + self.x * (y[-1] - y[0])),
        p0=initial_guess,
    )

    fitted_params = params_opt.reshape(-1, 2)
    logger.info("Fitted Redlich-Kister parameters (A_n, B_n): %s", fitted_params)

    xs = np.linspace(self.x.min(), self.x.max(), int(1e3))
    ys = redlich_kister_model(xs, T, *params_opt)
    ys += y[0] + xs * (y[-1] - y[0])

    self.plot(
        xs,
        ys,
        label=f"{label}: Redlich-Kister Fit" if label else "Redlich-Kister Fit",
        linestyle="--",
        lw=kwargs.get("lw", 1),
        color=color,
    )

plot_volume(label: str | None = None, **kwargs)

Plot molar volume on a secondary y-axis.

Must be called after from_trajectories or process so that self.x and self.y are populated.

Parameters:

Name	Type	Description	Default
label	str | None	Legend label for the volume curve.	None
**kwargs		Additional keyword arguments passed to errorbar.	{}

Returns:

Name	Type	Description
Axes		The secondary y-axis Axes.

Source code in muse/plots/density.py

def plot_volume(
    self,
    label: str | None = None,
    **kwargs,
):
    """Plot molar volume on a secondary y-axis.

    Must be called after ``from_trajectories`` or ``process`` so that
    ``self.x`` and ``self.y`` are populated.

    Args:
        label: Legend label for the volume curve.
        **kwargs: Additional keyword arguments passed to ``errorbar``.

    Returns:
        Axes: The secondary y-axis Axes.
    """
    self.vol_ax = self.twinx()
    self.vol_ax.errorbar(
        self.x,
        self.y["volume.avg"],
        yerr=self.y["volume.std"],
        label="$\\bar{V}$" if label else "$\\bar{V}$",
        **kwargs,
    )
    return self.vol_ax

pd

Binary Gibbs energy–composition (G–x) diagram with Redlich–Kister curve fitting.

BinaryGXDiagram(fig: Figure, *args: Any, facecolor=None, frameon: bool = True, sharex: Axes | None = None, sharey: Axes | None = None, label: str = '', xscale: float | None = None, yscale: float | None = None, box_aspect: float | None = None, **kwargs)

Bases: Axes

Custom Matplotlib Axes for binary Gibbs energy–composition (G–x) diagrams.

Computes mixing enthalpy ΔH and ideal entropy of mixing ΔS from MD trajectories, then fits ΔH with a Redlich–Kister polynomial.

Source code in muse/plots/pd.py

def __init__(
    self,
    fig: Figure,
    *args: Any,
    facecolor=None,
    frameon: bool = True,
    sharex: Axes | None = None,
    sharey: Axes | None = None,
    label: str = "",
    xscale: float | None = None,
    yscale: float | None = None,
    box_aspect: float | None = None,
    **kwargs,
) -> None:
    super().__init__(
        fig,
        args,
        facecolor=facecolor,
        frameon=frameon,
        sharex=sharex,
        sharey=sharey,
        label=label,
        xscale=xscale,
        yscale=yscale,
        box_aspect=box_aspect,
        **kwargs,
    )

from_trajectories(trajectories: Sequence[Sequence[Atoms]], phases: Sequence[str | Formula], temperature: float | None = None, label: str | None = None, rk: int = 2, **kwargs) -> None

Plot a binary G–x diagram from MD trajectories.

Computes mixing enthalpy ΔH = E(x) - [E(0) + x*(E(1) - E(0))] and ideal entropy of mixing, then fits ΔH with a Redlich–Kister polynomial expansion.

Parameters:

Name	Type	Description	Default
trajectories	Sequence[Sequence[Atoms]]	List of MD trajectories at different compositions.	required
phases	Sequence[str | Formula]	Two-element list of phase formulas defining the binary system.	required
temperature	float | None	Temperature in Kelvin for the Redlich–Kister fit. Defaults to 1000 K if not provided.	None
label	str | None	Label prefix for the legend entries.	None
rk	int	Number of Redlich–Kister terms. Defaults to 2.	2
**kwargs		Additional keyword arguments passed to errorbar.	{}

Source code in muse/plots/pd.py

def from_trajectories(
    self,
    trajectories: Sequence[Sequence[Atoms]],
    phases: Sequence[str | Formula],
    temperature: float | None = None,
    label: str | None = None,
    rk: int = 2,
    **kwargs,
) -> None:
    """Plot a binary G–x diagram from MD trajectories.

    Computes mixing enthalpy ΔH = E(x) - [E(0) + x*(E(1) - E(0))]
    and ideal entropy of mixing, then fits ΔH with a Redlich–Kister
    polynomial expansion.

    Args:
        trajectories: List of MD trajectories at different compositions.
        phases: Two-element list of phase formulas defining the binary system.
        temperature: Temperature in Kelvin for the Redlich–Kister fit.
            Defaults to 1000 K if not provided.
        label: Label prefix for the legend entries.
        rk: Number of Redlich–Kister terms. Defaults to 2.
        **kwargs: Additional keyword arguments passed to ``errorbar``.
    """
    assert len(phases) == 2

    # Change strings to Formula objects and sort symbols
    for phase in phases:
        phase = Formula.from_list(phase) if isinstance(phase, str) else sort(phase)

    x, ergavgs, ergstds = [], [], []
    for traj in trajectories:
        atoms = traj[0]
        formula = atoms.symbols.formula

        portions = {}
        for phase in phases:
            portions[phase] = formula // phase

        total_units = sum(portions.values())

        fractions = {}
        for phase in phases:
            fractions[phase] = portions[phase] / total_units

        x.append(fractions[phases[-1]])

        energies = []
        for atoms in traj:
            energies.append(atoms.get_potential_energy() / total_units)
        ergavgs.append(np.mean(energies))
        ergstds.append(np.std(energies))

    x = np.array(x)
    idx = np.argsort(x)
    x = x[idx]

    ergavgs = np.array(ergavgs)[idx]
    ergstds = np.array(ergstds)[idx]

    dH = ergavgs - (ergavgs[0] + x * (ergavgs[-1] - ergavgs[0]))
    _dS = -units.kB * (x * np.log(x + EPS) + (1 - x) * np.log(1 - x + EPS))  # noqa: F841

    color = kwargs.pop("color", "k")

    self.errorbar(
        x,
        dH,
        yerr=ergstds,
        label=f"{label}: $\\Delta H$" if label else "$\\Delta H$",
        color=color,
        fmt="o",
        **kwargs,
    )

    # Fitting the Redlich-Kister expansion model to Delta H
    initial_guess = [0.0] * (2 * rk)
    params_opt, params_cov = curve_fit(
        lambda x_T, *params: redlich_kister_model(x_T[0], x_T[1], *params),
        (x, np.ones_like(x) * (temperature or 1000)),
        dH,
        p0=initial_guess,
    )

    fitted_params = params_opt.reshape(-1, 2)
    logger.info("Fitted Redlich-Kister parameters (A_n, B_n): %s", fitted_params)

    xs = np.linspace(x.min(), x.max(), int(1e3))
    dH_fitted = redlich_kister_model(xs, temperature, *params_opt)

    self.plot(
        xs,
        dH_fitted,
        label=f"{label}: Redlich-Kister Fit" if label else "Redlich-Kister Fit",
        linestyle="--",
        lw=kwargs.get("lw", 1),
        color=color,
    )

volume

Binary mixing volume diagram with Redlich–Kister curve fitting.

MixingVolumeDiagram(fig: Figure, *args: Any, facecolor=None, frameon: bool = True, sharex: Axes | None = None, sharey: Axes | None = None, label: str = '', xscale: float | None = None, yscale: float | None = None, box_aspect: float | None = None, **kwargs)

Bases: Axes

Custom Matplotlib Axes for binary excess mixing volume diagrams.

Computes the deviation of molar volume from ideal mixing (Vegard's law) and fits the excess volume with a Redlich–Kister polynomial.

Source code in muse/plots/volume.py

def __init__(
    self,
    fig: Figure,
    *args: Any,
    facecolor=None,
    frameon: bool = True,
    sharex: Axes | None = None,
    sharey: Axes | None = None,
    label: str = "",
    xscale: float | None = None,
    yscale: float | None = None,
    box_aspect: float | None = None,
    **kwargs,
) -> None:
    super().__init__(
        fig,
        args,
        facecolor=facecolor,
        frameon=frameon,
        sharex=sharex,
        sharey=sharey,
        label=label,
        xscale=xscale,
        yscale=yscale,
        box_aspect=box_aspect,
        **kwargs,
    )

process(trajectories: Sequence[Sequence[Atoms]], phases: Sequence[str | Formula]) -> None

Process MD trajectories to extract density, volume, and excess volume.

Computes mass density (g/cm³), molar volume (ų/formula unit), and volume deviation from ideal (Vegard's law) mixing for each trajectory, storing results sorted by composition.

Parameters:

Name	Type	Description	Default
trajectories	Sequence[Sequence[Atoms]]	List of MD trajectories, each a sequence of Atoms.	required
phases	Sequence[str | Formula]	Two-element list of phase formulas defining the binary system.	required

Source code in muse/plots/volume.py

def process(
    self,
    trajectories: Sequence[Sequence[Atoms]],
    phases: Sequence[str | Formula],
) -> None:
    """Process MD trajectories to extract density, volume, and excess volume.

    Computes mass density (g/cm³), molar volume (ų/formula unit),
    and volume deviation from ideal (Vegard's law) mixing for each
    trajectory, storing results sorted by composition.

    Args:
        trajectories: List of MD trajectories, each a sequence of Atoms.
        phases: Two-element list of phase formulas defining the binary system.
    """
    # Change strings to Formula objects and sort symbols
    for phase in phases:
        phase = Formula.from_list(phase) if isinstance(phase, str) else sort(phase)

    x, denavgs, denstds, volavgs, volstds = [], [], [], [], []
    for traj in trajectories:
        atoms = traj[0]
        formula = atoms.symbols.formula

        portions = {}
        for phase in phases:
            portions[phase] = formula // phase

        total_units = sum(portions.values())

        fractions = {}
        for phase in phases:
            fractions[phase] = portions[phase] / total_units

        x.append(fractions[phases[-1]])

        densities = []
        volumes = []
        for atoms in traj:
            densities.append(
                atoms.get_masses().sum()
                * 1.66054e-24
                / (atoms.get_volume() * 1e-24)
            )
            volumes.append(atoms.get_volume() / total_units)

        denavgs.append(np.mean(densities))
        denstds.append(np.std(densities))
        volavgs.append(np.mean(volumes))
        volstds.append(np.std(volumes))

    x = np.array(x)
    idx = np.argsort(x)
    self.x = x[idx]

    self.y = {} if getattr(self, "y", None) is None else self.y
    self.y["density.avg"] = np.array(denavgs)[idx]
    self.y["density.std"] = np.array(denstds)[idx]
    self.y["volume.avg"] = np.array(volavgs)[idx]
    self.y["volume.std"] = np.array(volstds)[idx]
    self.y["volume.deviation"] = self.y["volume.avg"] - (
        self.y["volume.avg"][0]
        + self.x * (self.y["volume.avg"][-1] - self.y["volume.avg"][0])
    )

from_trajectories(trajectories: Sequence[Sequence[Atoms]], phases: Sequence[str | Formula], temperature: float | None = None, label: str | None = None, rk: int = 2, **kwargs) -> None

Plot a binary excess volume diagram from MD trajectories.

Computes the volume deviation from ideal mixing, plots it with error bars, and overlays a Redlich–Kister polynomial fit.

Parameters:

Name	Type	Description	Default
trajectories	Sequence[Sequence[Atoms]]	List of MD trajectories at different compositions.	required
phases	Sequence[str | Formula]	Two-element list of phase formulas defining the binary system.	required
temperature	float | None	Temperature in Kelvin for the Redlich–Kister fit. Defaults to 1000 K if not provided.	None
label	str | None	Label prefix for the legend entries.	None
rk	int	Number of Redlich–Kister terms. Defaults to 2.	2
**kwargs		Additional keyword arguments passed to errorbar.	{}

Source code in muse/plots/volume.py

def from_trajectories(
    self,
    trajectories: Sequence[Sequence[Atoms]],
    phases: Sequence[str | Formula],
    temperature: float | None = None,
    label: str | None = None,
    rk: int = 2,
    **kwargs,
) -> None:
    """Plot a binary excess volume diagram from MD trajectories.

    Computes the volume deviation from ideal mixing, plots it with
    error bars, and overlays a Redlich–Kister polynomial fit.

    Args:
        trajectories: List of MD trajectories at different compositions.
        phases: Two-element list of phase formulas defining the binary system.
        temperature: Temperature in Kelvin for the Redlich–Kister fit.
            Defaults to 1000 K if not provided.
        label: Label prefix for the legend entries.
        rk: Number of Redlich–Kister terms. Defaults to 2.
        **kwargs: Additional keyword arguments passed to ``errorbar``.
    """
    assert len(phases) == 2

    self.process(trajectories, phases)

    color = kwargs.pop("color", "k")

    self.errorbar(
        self.x,
        self.y["volume.deviation"],
        yerr=self.y["volume.std"],
        label="$\\Delta \\overline{V}$ " + label
        if label
        else "$\\Delta \\overline{V}$",
        color=color,
        fmt="o",
        **kwargs,
    )

    # Fitting the Redlich-Kister expansion model to excess volume
    initial_guess = [0.0] * (2 * rk)
    params_opt, params_cov = curve_fit(
        lambda x_T, *params: redlich_kister_model(x_T[0], x_T[1], *params),
        (self.x, np.ones_like(self.x) * (temperature or 1000)),
        self.y["volume.deviation"],
        p0=initial_guess,
    )

    fitted_params = params_opt.reshape(-1, 2)
    logger.info("Fitted Redlich-Kister parameters (A_n, B_n): %s", fitted_params)

    xs = np.linspace(self.x.min(), self.x.max(), int(1e3))
    dV_fitted = redlich_kister_model(xs, temperature, *params_opt)

    self.plot(
        xs,
        dV_fitted,
        label=f"{label}: Redlich-Kister Fit" if label else "Redlich-Kister Fit",
        linestyle="--",
        lw=kwargs.get("lw", 1),
        color=color,
    )

I/O

mptrj

Convert pymatgen trajectory data to extended XYZ format.

Supports writing energies, forces, stresses, charges, magnetic moments, and dipoles from Materials Project trajectory (MPtrj) data.

pmgtraj_to_extxyz(pmgtraj: Trajectory, fname: str | Path) -> None

Convert a pymatgen Trajectory to an extended XYZ file.

Writes each frame of the trajectory in the extended XYZ format, including lattice vectors, per-atom properties (species, positions, forces, charges, magnetic moments, dipoles), and frame-level properties (energies, stresses).

Parameters:

Name	Type	Description	Default
pmgtraj	Trajectory	A pymatgen Trajectory object, typically from MPtrj data.	required
fname	str | Path	Output file path for the extended XYZ file.	required

Source code in muse/io/mptrj.py

def pmgtraj_to_extxyz(pmgtraj: Trajectory, fname: str | Path) -> None:
    """Convert a pymatgen Trajectory to an extended XYZ file.

    Writes each frame of the trajectory in the extended XYZ format,
    including lattice vectors, per-atom properties (species, positions,
    forces, charges, magnetic moments, dipoles), and frame-level
    properties (energies, stresses).

    Args:
        pmgtraj: A pymatgen Trajectory object, typically from MPtrj data.
        fname: Output file path for the extended XYZ file.
    """
    with open(fname, "w") as f:
        for iframe, structure in enumerate(pmgtraj):
            # Write the extxyz format for each snapshot
            f.write(f"{len(structure)}\n")
            f.write(
                f'Lattice="{" ".join(map(str, structure.lattice.matrix.ravel()))}" '
            )
            f.write("Properties=species:S:1:pos:R:3")

            properties = pmgtraj.frame_properties[iframe]
            if "forces" in properties:
                f.write(":forces:R:3")
            if "charges" in properties:
                f.write(":charges:R:1")
            if "magmoms" in properties:
                f.write(":magmoms:R:3")
            if "dipoles" in properties:
                f.write(":dipoles:R:3")

            if "e_0_energy" in properties:
                f.write(f" energy={properties['e_0_energy']}")
            if "e_fr_energy" in properties:
                f.write(f" free_energy={properties['e_fr_energy']}")
            # e_wo_entrp: energy without electronic entropy (0K electron)
            if "e_wo_entrp" in properties:
                f.write(f" e_wo_entrp={properties['e_wo_entrp']}")

            if "stresses" in properties:
                f.write(
                    f' stresses="{" ".join(map(str, np.array(properties["stresses"]).ravel()))}"'
                )

            f.write("\n")

            for idx, site in enumerate(structure):
                line = f"{site.species_string} {' '.join(map(str, site.coords))}"

                if "forces" in properties:
                    line += f" {' '.join(map(str, properties['forces'][idx]))}"
                if "charges" in properties:
                    line += f" {' '.join(map(str, properties['charges'][idx]))}"
                if "magmoms" in properties:
                    line += f" {' '.join(map(str, properties['magmoms'][idx]))}"
                if "dipoles" in properties:
                    line += f" {' '.join(map(str, properties['dipoles'][idx]))}"

                line += "\n"
                f.write(line)

Jobs

slurm

SLURM job submission utilities.

submit_job(cmd: str, time: str, partition: str, nodes: int, ntasks_per_node: int, job_name: str, account: str | None = None, **kwargs: str | None) -> subprocess.CompletedProcess

Submit a SLURM batch job using sbatch --wrap.

Constructs and executes an sbatch command with the given resource parameters, wrapping the provided command string.

Parameters:

Name	Type	Description	Default
cmd	str	The command to run inside the SLURM job.	required
time	str	Wall time limit (e.g., "2:00:00").	required
partition	str	SLURM partition/QOS name.	required
nodes	int	Number of nodes to request.	required
ntasks_per_node	int	Number of tasks per node.	required
job_name	str	Name for the SLURM job.	required
account	str | None	SLURM account/project for billing. Defaults to None.	None
**kwargs	str | None	Additional sbatch options. Supported keys: - constraint: Node constraint string. - exclude: Nodes to exclude.	{}

Returns:

Type	Description
CompletedProcess	The completed process result from subprocess.run.

Raises:

Type	Description
CalledProcessError	If sbatch returns a non-zero exit code.

Source code in muse/jobs/slurm.py

def submit_job(
    cmd: str,
    time: str,
    partition: str,
    nodes: int,
    ntasks_per_node: int,
    job_name: str,
    account: str | None = None,
    **kwargs: str | None,
) -> subprocess.CompletedProcess:
    """Submit a SLURM batch job using ``sbatch --wrap``.

    Constructs and executes an ``sbatch`` command with the given
    resource parameters, wrapping the provided command string.

    Args:
        cmd: The command to run inside the SLURM job.
        time: Wall time limit (e.g., ``"2:00:00"``).
        partition: SLURM partition/QOS name.
        nodes: Number of nodes to request.
        ntasks_per_node: Number of tasks per node.
        job_name: Name for the SLURM job.
        account: SLURM account/project for billing. Defaults to None.
        **kwargs: Additional sbatch options. Supported keys:
            - ``constraint``: Node constraint string.
            - ``exclude``: Nodes to exclude.

    Returns:
        The completed process result from ``subprocess.run``.

    Raises:
        subprocess.CalledProcessError: If sbatch returns a non-zero exit code.
    """
    scmd = [
        *f"""sbatch
        --time={time}
        --qos={partition}
        --job-name={job_name}
        --output={job_name}.out
        --error={job_name}.err
        --nodes={nodes}
        --ntasks-per-node={ntasks_per_node}""".replace("'", "").split(),
    ]

    if account is not None:
        scmd += [f"--account={account}"]

    if kwargs.get("constraint"):
        scmd += [f"--constraint={kwargs['constraint']}"]

    if kwargs.get("exclude"):
        scmd += ["--exclude"]

    scmd += ["--wrap", cmd]

    return subprocess.run(scmd, check=True)