Skip to content

Tasks Module

The tasks module contains the implementation of various simulation tasks.

Static Calculations

Static calculations: single-point energy and structure optimization

calculate_single_point

def calculate_single_point(atoms: Atoms, calculator: Calculator) -> dict

Calculate energy, forces, and stress for a structure.

Parameters:

  • atoms (Atoms): ASE Atoms object
  • calculator (Calculator): ASE calculator

Returns: dict with energy, forces, stress, etc.

Optimization

Static calculations: single-point energy and structure optimization

optimize_structure 

optimize_structure(atoms: Atoms, calculator, fmax: float = 0.05, steps: int = 500, optimizer: str = 'LBFGS', optimize_cell: bool = False, trajectory: Optional[str] = None, logfile: Optional[str] = None) -> Atoms

Optimize atomic structure.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object

 required
calculator

ASE calculator

 required
fmax float

Force convergence criterion (eV/Å)

 0.05
steps int

Maximum optimization steps

 500
optimizer str

Optimizer name ("LBFGS", "BFGS", "FIRE")

 'LBFGS'
optimize_cell bool

Whether to optimize cell parameters

 False
trajectory Optional[str]

Path to save optimization trajectory

 None
logfile Optional[str]

Path to save optimization log

 None

Returns:

Type Description
Atoms

Optimized Atoms object
Source code in src/mace_inference/tasks/static.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def optimize_structure(
    atoms: Atoms,
    calculator,
    fmax: float = 0.05,
    steps: int = 500,
    optimizer: str = "LBFGS",
    optimize_cell: bool = False,
    trajectory: Optional[str] = None,
    logfile: Optional[str] = None
) -> Atoms:
    """
    Optimize atomic structure.

    Args:
        atoms: Input ASE Atoms object
        calculator: ASE calculator
        fmax: Force convergence criterion (eV/Å)
        steps: Maximum optimization steps
        optimizer: Optimizer name ("LBFGS", "BFGS", "FIRE")
        optimize_cell: Whether to optimize cell parameters
        trajectory: Path to save optimization trajectory
        logfile: Path to save optimization log

    Returns:
        Optimized Atoms object
    """
    # Make a copy to avoid modifying original
    atoms = atoms.copy()
    atoms.calc = calculator

    # Select optimizer
    optimizer_map = {
        "LBFGS": LBFGS,
        "BFGS": BFGS,
        "FIRE": FIRE
    }

    if optimizer not in optimizer_map:
        raise ValueError(
            f"Unknown optimizer: {optimizer}. Choose from {list(optimizer_map.keys())}"
        )

    OptClass = optimizer_map[optimizer]

    # Apply cell filter if optimizing cell
    if optimize_cell:
        atoms_to_optimize = FrechetCellFilter(atoms)
    else:
        atoms_to_optimize = atoms

    # Create optimizer
    opt = OptClass(
        atoms_to_optimize,
        trajectory=trajectory,
        logfile=logfile
    )

    # Run optimization
    opt.run(fmax=fmax, steps=steps)

    # Return the original atoms object (which has been modified in-place)
    return atoms

optimize_structure

def optimize_structure(
    atoms: Atoms,
    calculator: Calculator,
    fmax: float = 0.05,
    steps: int = 500,
    optimizer: str = "BFGS",
    relax_cell: bool = False,
    fix_symmetry: bool = False,
    trajectory: Optional[str] = None
) -> dict

Optimize atomic positions and cell.

Molecular Dynamics

Molecular dynamics simulations

run_nvt_md 

run_nvt_md(atoms: Atoms, calculator, temperature_K: float = 300, timestep: float = 1.0, steps: int = 1000, trajectory: Optional[str] = None, logfile: Optional[str] = None, log_interval: int = 100, taut: Optional[float] = None, friction: Optional[float] = None, progress_callback: Optional[Callable[[int, int], None]] = None) -> Atoms

Run NVT molecular dynamics using Langevin thermostat.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object

 required
calculator

ASE calculator

 required
temperature_K float

Target temperature (K)

 300
timestep float

Time step (fs)

 1.0
steps int

Number of MD steps

 1000
trajectory Optional[str]

Trajectory file path

 None
logfile Optional[str]

Log file path

 None
log_interval int

Logging interval (steps)

 100
taut Optional[float]

Temperature coupling time (fs) - alternative to friction

 None
friction Optional[float]

Friction coefficient (1/fs)

 None
progress_callback Optional[Callable[[int, int], None]]

Optional callback function(current_step, total_steps)

 None

Returns:

Type Description
Atoms

Final Atoms object after MD

Examples:

>>> def progress(current, total):
...     print(f"MD progress: {current}/{total} steps ({100*current/total:.1f}%)")
>>> atoms = run_nvt_md(atoms, calc, steps=1000, progress_callback=progress)
Source code in src/mace_inference/tasks/dynamics.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def run_nvt_md(
    atoms: Atoms,
    calculator,
    temperature_K: float = 300,
    timestep: float = 1.0,
    steps: int = 1000,
    trajectory: Optional[str] = None,
    logfile: Optional[str] = None,
    log_interval: int = 100,
    taut: Optional[float] = None,
    friction: Optional[float] = None,
    progress_callback: Optional[Callable[[int, int], None]] = None
) -> Atoms:
    """
    Run NVT molecular dynamics using Langevin thermostat.

    Args:
        atoms: Input ASE Atoms object
        calculator: ASE calculator
        temperature_K: Target temperature (K)
        timestep: Time step (fs)
        steps: Number of MD steps
        trajectory: Trajectory file path
        logfile: Log file path
        log_interval: Logging interval (steps)
        taut: Temperature coupling time (fs) - alternative to friction
        friction: Friction coefficient (1/fs)
        progress_callback: Optional callback function(current_step, total_steps)

    Returns:
        Final Atoms object after MD

    Examples:
        >>> def progress(current, total):
        ...     print(f"MD progress: {current}/{total} steps ({100*current/total:.1f}%)")
        >>> atoms = run_nvt_md(atoms, calc, steps=1000, progress_callback=progress)
    """
    atoms = atoms.copy()
    atoms.calc = calculator

    # Initialize velocities
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)

    # Calculate friction from taut if provided
    if taut is not None and friction is None:
        friction = 1.0 / taut  # Convert coupling time to friction
    elif friction is None:
        friction = 0.002  # Default friction coefficient

    # Create Langevin dynamics
    dyn = Langevin(
        atoms,
        timestep=timestep * units.fs,
        temperature_K=temperature_K,
        friction=friction / units.fs
    )

    # Attach trajectory writer
    if trajectory:
        traj = Trajectory(trajectory, 'w', atoms)
        dyn.attach(traj.write, interval=log_interval)

    # Attach logger
    if logfile:
        from ase.md import MDLogger
        logger = MDLogger(
            dyn,
            atoms,
            logfile,
            header=True,
            stress=False,
            peratom=False,
            mode='w'
        )
        dyn.attach(logger, interval=log_interval)

    # Attach progress callback
    if progress_callback:
        def _progress_wrapper():
            progress_callback(dyn.nsteps, steps)
        dyn.attach(_progress_wrapper, interval=log_interval)

    # Run MD
    dyn.run(steps)

    # Close trajectory
    if trajectory:
        traj.close()

    return atoms

run_npt_md 

run_npt_md(atoms: Atoms, calculator, temperature_K: float = 300, pressure_GPa: float = 0.0, timestep: float = 1.0, steps: int = 1000, trajectory: Optional[str] = None, logfile: Optional[str] = None, log_interval: int = 100, taut: Optional[float] = None, taup: Optional[float] = None, compressibility_au: Optional[float] = None, progress_callback: Optional[Callable[[int, int], None]] = None) -> Atoms

Run NPT molecular dynamics using NPT ensemble.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object

 required
calculator

ASE calculator

 required
temperature_K float

Target temperature (K)

 300
pressure_GPa float

Target pressure (GPa)

 0.0
timestep float

Time step (fs)

 1.0
steps int

Number of MD steps

 1000
trajectory Optional[str]

Trajectory file path

 None
logfile Optional[str]

Log file path

 None
log_interval int

Logging interval (steps)

 100
taut Optional[float]

Temperature coupling time (fs)

 None
taup Optional[float]

Pressure coupling time (fs)

 None
compressibility_au Optional[float]

Isothermal compressibility (atomic units)

 None
progress_callback Optional[Callable[[int, int], None]]

Optional callback function(current_step, total_steps)

 None

Returns:

Type Description
Atoms

Final Atoms object after MD

Examples:

>>> def progress(current, total):
...     print(f"NPT MD: {current}/{total} steps")
>>> atoms = run_npt_md(atoms, calc, temperature_K=300, pressure_GPa=0.1,
...                    steps=1000, progress_callback=progress)
Source code in src/mace_inference/tasks/dynamics.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def run_npt_md(
    atoms: Atoms,
    calculator,
    temperature_K: float = 300,
    pressure_GPa: float = 0.0,
    timestep: float = 1.0,
    steps: int = 1000,
    trajectory: Optional[str] = None,
    logfile: Optional[str] = None,
    log_interval: int = 100,
    taut: Optional[float] = None,
    taup: Optional[float] = None,
    compressibility_au: Optional[float] = None,
    progress_callback: Optional[Callable[[int, int], None]] = None
) -> Atoms:
    """
    Run NPT molecular dynamics using NPT ensemble.

    Args:
        atoms: Input ASE Atoms object
        calculator: ASE calculator
        temperature_K: Target temperature (K)
        pressure_GPa: Target pressure (GPa)
        timestep: Time step (fs)
        steps: Number of MD steps
        trajectory: Trajectory file path
        logfile: Log file path
        log_interval: Logging interval (steps)
        taut: Temperature coupling time (fs)
        taup: Pressure coupling time (fs)
        compressibility_au: Isothermal compressibility (atomic units)
        progress_callback: Optional callback function(current_step, total_steps)

    Returns:
        Final Atoms object after MD

    Examples:
        >>> def progress(current, total):
        ...     print(f"NPT MD: {current}/{total} steps")
        >>> atoms = run_npt_md(atoms, calc, temperature_K=300, pressure_GPa=0.1,
        ...                    steps=1000, progress_callback=progress)
    """
    atoms = atoms.copy()
    atoms.calc = calculator

    # Initialize velocities
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)

    # Default coupling times
    if taut is None:
        taut = 100.0  # fs
    if taup is None:
        taup = 1000.0  # fs

    # Calculate pfactor from taup if compressibility not provided
    if compressibility_au is None:
        # Estimate bulk modulus (typical for MOFs: ~10 GPa = 0.062 eV/ų)
        B_estimate = 10.0 * units.GPa  # Convert to eV/ų
        pfactor = (taup * units.fs)**2 * B_estimate
    else:
        pfactor = (taup * units.fs)**2 / compressibility_au

    # Convert pressure to atomic units (stress tensor)
    # Note: NPT uses stress (positive in tension), so negate pressure
    externalstress = -pressure_GPa * units.GPa

    # Create NPT dynamics
    dyn = NPT(
        atoms,
        timestep=timestep * units.fs,
        temperature_K=temperature_K,
        externalstress=externalstress,
        ttime=taut * units.fs,
        pfactor=pfactor
    )

    # Attach trajectory writer
    if trajectory:
        traj = Trajectory(trajectory, 'w', atoms)
        dyn.attach(traj.write, interval=log_interval)

    # Attach logger
    if logfile:
        from ase.md import MDLogger
        logger = MDLogger(
            dyn,
            atoms,
            logfile,
            header=True,
            stress=True,
            peratom=False,
            mode='w'
        )
        dyn.attach(logger, interval=log_interval)

    # Attach volume monitor
    def log_volume():
        if dyn.nsteps % log_interval == 0:
            vol = atoms.get_volume()
            if hasattr(dyn, '_initial_volume'):
                vol_change = (vol - dyn._initial_volume) / dyn._initial_volume * 100
            else:
                dyn._initial_volume = vol
                vol_change = 0.0
            print(f"Step {dyn.nsteps}: Volume = {vol:.2f} ų ({vol_change:+.2f}%)")

    dyn.attach(log_volume, interval=log_interval)

    # Attach progress callback
    if progress_callback:
        def _progress_wrapper():
            progress_callback(dyn.nsteps, steps)
        dyn.attach(_progress_wrapper, interval=log_interval)

    # Run MD
    dyn.run(steps)

    # Close trajectory
    if trajectory:
        traj.close()

    return atoms

run_nvt

def run_nvt(
    atoms: Atoms,
    calculator: Calculator,
    temperature: float,
    timestep: float = 1.0,
    steps: int = 1000,
    friction: float = 0.01,
    save_interval: int = 10,
    log_interval: int = 100
) -> dict

Run NVT molecular dynamics with Langevin thermostat.

run_npt

def run_npt(
    atoms: Atoms,
    calculator: Calculator,
    temperature: float,
    pressure: float,
    timestep: float = 1.0,
    steps: int = 1000,
    save_interval: int = 10,
    log_interval: int = 100
) -> dict

Run NPT molecular dynamics with Berendsen barostat.

Phonon Calculations

Phonon calculations and thermal properties

ase_to_phonopy 

ase_to_phonopy(atoms: Atoms) -> PhonopyAtoms

Convert ASE Atoms to Phonopy Atoms.

Source code in src/mace_inference/tasks/phonon.py 
13
14
15
16
17
18
19
def ase_to_phonopy(atoms: Atoms) -> PhonopyAtoms:
    """Convert ASE Atoms to Phonopy Atoms."""
    return PhonopyAtoms(
        symbols=atoms.get_chemical_symbols(),
        scaled_positions=atoms.get_scaled_positions(),
        cell=atoms.get_cell()
    )

calculate_phonon 

calculate_phonon(atoms: Atoms, calculator: Calculator, supercell_matrix: Union[List[int], ndarray, int] = 2, displacement: float = 0.01, mesh: Optional[List[int]] = None, output_dir: Optional[str] = None) -> PhononResult

Calculate phonon properties.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object

 required
calculator Calculator

ASE calculator

 required
supercell_matrix Union[List[int], ndarray, int]

Supercell size (e.g., [2, 2, 2] or 2)

 2
displacement float

Atomic displacement distance (Å)

 0.01
mesh Optional[List[int]]

k-point mesh for phonon DOS (default: [20, 20, 20])

 None
output_dir Optional[str]

Output directory for phonon files

 None

Returns:

Type Description
PhononResult

Dictionary containing phonon object and basic properties
Source code in src/mace_inference/tasks/phonon.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def calculate_phonon(
    atoms: Atoms,
    calculator: "Calculator",
    supercell_matrix: Union[List[int], np.ndarray, int] = 2,
    displacement: float = 0.01,
    mesh: Optional[List[int]] = None,
    output_dir: Optional[str] = None
) -> "PhononResult":
    """
    Calculate phonon properties.

    Args:
        atoms: Input ASE Atoms object
        calculator: ASE calculator
        supercell_matrix: Supercell size (e.g., [2, 2, 2] or 2)
        displacement: Atomic displacement distance (Å)
        mesh: k-point mesh for phonon DOS (default: [20, 20, 20])
        output_dir: Output directory for phonon files

    Returns:
        Dictionary containing phonon object and basic properties
    """
    # Convert supercell_matrix to list
    if isinstance(supercell_matrix, int):
        supercell_matrix = [[supercell_matrix, 0, 0],
                           [0, supercell_matrix, 0],
                           [0, 0, supercell_matrix]]
    elif isinstance(supercell_matrix, list) and len(supercell_matrix) == 3:
        if all(isinstance(x, int) for x in supercell_matrix):
            supercell_matrix = [[supercell_matrix[0], 0, 0],
                               [0, supercell_matrix[1], 0],
                               [0, 0, supercell_matrix[2]]]

    # Create Phonopy object
    phonopy_atoms = ase_to_phonopy(atoms)
    phonon = Phonopy(phonopy_atoms, supercell_matrix=supercell_matrix)

    # Generate displacements
    phonon.generate_displacements(distance=displacement)

    # Calculate forces for displaced structures
    supercells = phonon.supercells_with_displacements
    forces_list = []

    for i, scell in enumerate(supercells):
        # Convert Phonopy atoms to ASE atoms
        # PhonopyAtoms uses .symbols instead of .get_chemical_symbols()
        ase_scell = Atoms(
            symbols=scell.symbols,
            positions=scell.positions,
            cell=scell.cell,
            pbc=True
        )
        ase_scell.calc = calculator
        forces = ase_scell.get_forces()
        forces_list.append(forces)

    # Set forces to phonon
    phonon.forces = forces_list

    # Produce force constants
    phonon.produce_force_constants()

    # Calculate phonon DOS if mesh specified
    if mesh is None:
        mesh = [20, 20, 20]

    phonon.run_mesh(mesh)
    phonon.run_total_dos()

    # Save phonon files if output directory specified
    if output_dir:
        import os
        os.makedirs(output_dir, exist_ok=True)
        phonon.save(f"{output_dir}/phonopy.yaml")
        phonon.write_yaml_force_constants(f"{output_dir}/force_constants.yaml")

    result = {
        "phonon": phonon,
        "supercell_matrix": supercell_matrix,
        "displacement": displacement,
        "mesh": mesh,
    }

    return result

calculate_thermal_properties 

calculate_thermal_properties(phonon: Phonopy, t_min: float = 0, t_max: float = 1000, t_step: float = 10) -> ThermalPropertiesResult

Calculate thermal properties from phonon object.

Parameters:

Name Type Description Default
phonon Phonopy

Phonopy object with force constants

 required
t_min float

Minimum temperature (K)

 0
t_max float

Maximum temperature (K)

 1000
t_step float

Temperature step (K)

 10

Returns:

Type Description
ThermalPropertiesResult

Dictionary with thermal properties
Source code in src/mace_inference/tasks/phonon.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def calculate_thermal_properties(
    phonon: Phonopy,
    t_min: float = 0,
    t_max: float = 1000,
    t_step: float = 10
) -> "ThermalPropertiesResult":
    """
    Calculate thermal properties from phonon object.

    Args:
        phonon: Phonopy object with force constants
        t_min: Minimum temperature (K)
        t_max: Maximum temperature (K)
        t_step: Temperature step (K)

    Returns:
        Dictionary with thermal properties
    """
    # Run thermal properties calculation
    phonon.run_thermal_properties(t_step=t_step, t_max=t_max, t_min=t_min)

    # Get thermal properties dictionary
    tp_dict = phonon.get_thermal_properties_dict()

    result = {
        "temperatures": tp_dict["temperatures"],  # K
        "free_energy": tp_dict["free_energy"],    # kJ/mol
        "entropy": tp_dict["entropy"],            # J/(mol·K)
        "heat_capacity": tp_dict["heat_capacity"], # J/(mol·K)
    }

    return result

calculate_phonon

def calculate_phonon(
    atoms: Atoms,
    calculator: Calculator,
    supercell: Tuple[int, int, int] = (2, 2, 2),
    delta: float = 0.01,
    mesh: Tuple[int, int, int] = (20, 20, 20)
) -> dict

Calculate phonon frequencies and density of states.

calculate_thermal_properties

def calculate_thermal_properties(
    phonon: Phonopy,
    t_min: float = 0,
    t_max: float = 1000,
    t_step: float = 10
) -> dict

Calculate thermal properties from phonon data.

Mechanical Properties

Mechanical properties calculations

calculate_bulk_modulus 

calculate_bulk_modulus(atoms: Atoms, calculator: Calculator, n_points: int = 11, scale_range: Tuple[float, float] = (0.95, 1.05), eos_type: str = 'birchmurnaghan') -> BulkModulusResult

Calculate bulk modulus using equation of state.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object

 required
calculator Calculator

ASE calculator

 required
n_points int

Number of volume points

 11
scale_range Tuple[float, float]

Volume scaling range (min_scale, max_scale)

 (0.95, 1.05)
eos_type str

Equation of state type ("birchmurnaghan", "murnaghan", etc.)

 'birchmurnaghan'

Returns:

Type Description
BulkModulusResult

Dictionary with equilibrium volume, energy, and bulk modulus
Source code in src/mace_inference/tasks/mechanics.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def calculate_bulk_modulus(
    atoms: Atoms,
    calculator: "Calculator",
    n_points: int = 11,
    scale_range: Tuple[float, float] = (0.95, 1.05),
    eos_type: str = "birchmurnaghan"
) -> "BulkModulusResult":
    """
    Calculate bulk modulus using equation of state.

    Args:
        atoms: Input ASE Atoms object
        calculator: ASE calculator
        n_points: Number of volume points
        scale_range: Volume scaling range (min_scale, max_scale)
        eos_type: Equation of state type ("birchmurnaghan", "murnaghan", etc.)

    Returns:
        Dictionary with equilibrium volume, energy, and bulk modulus
    """
    atoms = atoms.copy()
    atoms.calc = calculator

    # Generate volume points
    scale_factors = np.linspace(scale_range[0], scale_range[1], n_points)

    volumes = []
    energies = []

    # Calculate energy at each volume
    original_cell = atoms.get_cell()

    for scale in scale_factors:
        # Scale cell uniformly
        scaled_atoms = atoms.copy()
        scaled_atoms.set_cell(original_cell * scale, scale_atoms=True)
        scaled_atoms.calc = calculator

        volumes.append(scaled_atoms.get_volume())
        energies.append(scaled_atoms.get_potential_energy())

    # Fit equation of state
    eos = EquationOfState(volumes, energies, eos=eos_type)
    v0, e0, B = eos.fit()

    # Convert bulk modulus from eV/ų to GPa
    B_GPa = B * 160.21766208

    result = {
        "v0": v0,           # Equilibrium volume (ų)
        "e0": e0,           # Equilibrium energy (eV)
        "B": B,             # Bulk modulus (eV/ų)
        "B_GPa": B_GPa,     # Bulk modulus (GPa)
        "volumes": volumes,
        "energies": energies,
        "eos_type": eos_type
    }

    return result

calculate_elastic_constants 

calculate_elastic_constants(atoms: Atoms, calculator: Calculator, delta: float = 0.01, symmetry: str = 'auto') -> Dict[str, Any]

Calculate elastic constants using stress-strain relationships.

Uses finite differences to compute the elastic tensor from stress responses to applied strains.

Parameters:

Name Type Description Default
atoms Atoms

Input ASE Atoms object (should be relaxed)

 required
calculator Calculator

ASE calculator

 required
delta float

Strain magnitude for finite differences

 0.01
symmetry str

Crystal symmetry ("auto", "cubic", "hexagonal", "orthorhombic", "full")

 'auto'

Returns:

Type Description
Dict[str, Any]

Dictionary with elastic constants in GPa
Example

result = calculate_elastic_constants(atoms, calculator) print(f"C11 = {result['C'][0,0]:.1f} GPa") print(f"Bulk modulus = {result['K_V']:.1f} GPa")

Source code in src/mace_inference/tasks/mechanics.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def calculate_elastic_constants(
    atoms: Atoms,
    calculator: "Calculator",
    delta: float = 0.01,
    symmetry: str = "auto"
) -> Dict[str, Any]:
    """
    Calculate elastic constants using stress-strain relationships.

    Uses finite differences to compute the elastic tensor from
    stress responses to applied strains.

    Args:
        atoms: Input ASE Atoms object (should be relaxed)
        calculator: ASE calculator
        delta: Strain magnitude for finite differences
        symmetry: Crystal symmetry ("auto", "cubic", "hexagonal", "orthorhombic", "full")

    Returns:
        Dictionary with elastic constants in GPa

    Example:
        >>> result = calculate_elastic_constants(atoms, calculator)
        >>> print(f"C11 = {result['C'][0,0]:.1f} GPa")
        >>> print(f"Bulk modulus = {result['K_V']:.1f} GPa")
    """
    atoms = atoms.copy()
    atoms.calc = calculator

    # Get reference stress (should be ~0 for relaxed structure)
    _ref_stress = atoms.get_stress(voigt=False)  # 3x3 tensor, used to verify relaxed state

    # Voigt notation indices: 
    # 0=xx, 1=yy, 2=zz, 3=yz, 4=xz, 5=xy
    voigt_pairs = [(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)]

    # Elastic tensor in Voigt notation (6x6)
    C = np.zeros((6, 6))

    # Get original cell
    cell0 = atoms.get_cell().copy()

    for i, (p, q) in enumerate(voigt_pairs):
        # Apply positive strain
        strained_atoms = atoms.copy()
        strain_matrix = np.eye(3)
        if p == q:
            # Normal strain
            strain_matrix[p, p] = 1 + delta
        else:
            # Shear strain (symmetric)
            strain_matrix[p, q] = delta / 2
            strain_matrix[q, p] = delta / 2

        new_cell = cell0 @ strain_matrix
        strained_atoms.set_cell(new_cell, scale_atoms=True)
        strained_atoms.calc = calculator
        stress_plus = strained_atoms.get_stress(voigt=False)

        # Apply negative strain
        strained_atoms = atoms.copy()
        strain_matrix = np.eye(3)
        if p == q:
            strain_matrix[p, p] = 1 - delta
        else:
            strain_matrix[p, q] = -delta / 2
            strain_matrix[q, p] = -delta / 2

        new_cell = cell0 @ strain_matrix
        strained_atoms.set_cell(new_cell, scale_atoms=True)
        strained_atoms.calc = calculator
        stress_minus = strained_atoms.get_stress(voigt=False)

        # Compute stress derivative (central difference)
        dstress = (stress_plus - stress_minus) / (2 * delta)

        # Fill elastic tensor (convert to Voigt notation)
        for j, (r, s) in enumerate(voigt_pairs):
            C[j, i] = -dstress[r, s]  # Negative because stress = -C * strain

    # Convert from eV/ų to GPa
    eV_per_A3_to_GPa = 160.21766208
    C_GPa = C * eV_per_A3_to_GPa

    # Symmetrize
    C_GPa = (C_GPa + C_GPa.T) / 2

    # Calculate derived properties using Voigt-Reuss-Hill averages
    # Compliance tensor
    try:
        S = np.linalg.inv(C_GPa)
    except np.linalg.LinAlgError:
        S = np.zeros((6, 6))

    # Voigt averages (upper bound)
    K_V = (C_GPa[0,0] + C_GPa[1,1] + C_GPa[2,2] + 
           2*(C_GPa[0,1] + C_GPa[1,2] + C_GPa[0,2])) / 9
    G_V = ((C_GPa[0,0] + C_GPa[1,1] + C_GPa[2,2]) - 
           (C_GPa[0,1] + C_GPa[1,2] + C_GPa[0,2]) +
           3*(C_GPa[3,3] + C_GPa[4,4] + C_GPa[5,5])) / 15

    # Reuss averages (lower bound)
    K_R = 1 / (S[0,0] + S[1,1] + S[2,2] + 2*(S[0,1] + S[1,2] + S[0,2]))
    G_R = 15 / (4*(S[0,0] + S[1,1] + S[2,2]) - 
                4*(S[0,1] + S[1,2] + S[0,2]) +
                3*(S[3,3] + S[4,4] + S[5,5]))

    # Hill averages
    K_H = (K_V + K_R) / 2
    G_H = (G_V + G_R) / 2

    # Young's modulus and Poisson's ratio from Hill averages
    E_H = 9 * K_H * G_H / (3 * K_H + G_H)
    nu_H = (3 * K_H - 2 * G_H) / (2 * (3 * K_H + G_H))

    result = {
        "C": C_GPa,                    # Full elastic tensor (GPa)
        "S": S,                        # Compliance tensor (1/GPa)
        "K_V": float(K_V),             # Voigt bulk modulus (GPa)
        "K_R": float(K_R),             # Reuss bulk modulus (GPa)
        "K_H": float(K_H),             # Hill bulk modulus (GPa)
        "G_V": float(G_V),             # Voigt shear modulus (GPa)
        "G_R": float(G_R),             # Reuss shear modulus (GPa)
        "G_H": float(G_H),             # Hill shear modulus (GPa)
        "E_H": float(E_H),             # Hill Young's modulus (GPa)
        "nu_H": float(nu_H),           # Hill Poisson's ratio
        "delta": delta,
        "symmetry": symmetry,
    }

    # Add cubic-specific constants if applicable
    if symmetry in ["auto", "cubic"]:
        result["C11"] = float(C_GPa[0, 0])
        result["C12"] = float(C_GPa[0, 1])
        result["C44"] = float(C_GPa[3, 3])
        # Zener anisotropy ratio
        if abs(C_GPa[0, 0] - C_GPa[0, 1]) > 1e-6:
            result["A"] = float(2 * C_GPa[3, 3] / (C_GPa[0, 0] - C_GPa[0, 1]))

    return result

calculate_bulk_modulus

def calculate_bulk_modulus(
    atoms: Atoms,
    calculator: Calculator,
    scale_range: Tuple[float, float] = (0.95, 1.05),
    n_points: int = 7,
    eos: str = "birchmurnaghan"
) -> dict

Calculate bulk modulus by fitting equation of state.

Adsorption

Adsorption energy and coordination analysis

calculate_adsorption_energy 

calculate_adsorption_energy(framework: Atoms, adsorbate: Union[str, Atoms], site_position: List[float], calculator: Calculator, optimize: bool = True, fmax: float = 0.05, fix_framework: bool = True) -> AdsorptionResult

Calculate gas adsorption energy in framework (MOF, zeolite, etc.).

E_ads = E(framework+adsorbate) - E(framework) - E(adsorbate)

Parameters:

Name Type Description Default
framework Atoms

Framework structure (MOF, zeolite, etc.) as ASE Atoms

 required
adsorbate Union[str, Atoms]

Adsorbate molecule name (e.g., "CO2") or Atoms object

 required
site_position List[float]

Adsorption site position [x, y, z]

 required
calculator Calculator

ASE calculator

 required
optimize bool

Whether to optimize the adsorption complex

 True
fmax float

Force convergence for optimization

 0.05
fix_framework bool

Whether to fix framework atoms during optimization

 True

Returns:

Type Description
AdsorptionResult

Dictionary with adsorption energy and structures
Source code in src/mace_inference/tasks/adsorption.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def calculate_adsorption_energy(
    framework: Atoms,
    adsorbate: Union[str, Atoms],
    site_position: List[float],
    calculator: "Calculator",
    optimize: bool = True,
    fmax: float = 0.05,
    fix_framework: bool = True
) -> "AdsorptionResult":
    """
    Calculate gas adsorption energy in framework (MOF, zeolite, etc.).

    E_ads = E(framework+adsorbate) - E(framework) - E(adsorbate)

    Args:
        framework: Framework structure (MOF, zeolite, etc.) as ASE Atoms
        adsorbate: Adsorbate molecule name (e.g., "CO2") or Atoms object
        site_position: Adsorption site position [x, y, z]
        calculator: ASE calculator
        optimize: Whether to optimize the adsorption complex
        fmax: Force convergence for optimization
        fix_framework: Whether to fix framework atoms during optimization

    Returns:
        Dictionary with adsorption energy and structures
    """
    # 1. Calculate framework energy
    framework_copy = framework.copy()
    framework_copy.calc = calculator
    E_framework = framework_copy.get_potential_energy()

    # 2. Get adsorbate molecule
    if isinstance(adsorbate, str):
        ads = molecule(adsorbate)
    elif isinstance(adsorbate, Atoms):
        ads = adsorbate.copy()
    else:
        raise TypeError("adsorbate must be str or Atoms object")

    # Calculate adsorbate energy (in vacuum)
    ads.calc = calculator
    ads.center(vacuum=10.0)  # Add vacuum around molecule
    E_ads_mol = ads.get_potential_energy()

    # 3. Create adsorption complex
    complex_atoms = framework_copy.copy()

    # Position adsorbate molecule at adsorption site
    ads_centered = ads.copy()
    ads_com = ads_centered.get_center_of_mass()
    translation = np.array(site_position) - ads_com
    ads_centered.translate(translation)

    # Combine structures
    complex_atoms += ads_centered
    complex_atoms.calc = calculator

    # 4. Optimize complex if requested
    if optimize:
        # Fix framework atoms if requested, only optimize adsorbate
        if fix_framework:
            from ase.constraints import FixAtoms
            n_framework_atoms = len(framework)
            constraint = FixAtoms(indices=range(n_framework_atoms))
            complex_atoms.set_constraint(constraint)

        opt = LBFGS(complex_atoms)
        opt.run(fmax=fmax)

    E_complex = complex_atoms.get_potential_energy()

    # 5. Calculate adsorption energy
    E_ads = E_complex - E_framework - E_ads_mol

    result = {
        "E_ads": E_ads,              # Adsorption energy (eV)
        "E_mof": E_framework,        # Framework energy (eV) - kept as E_mof for backward compatibility
        "E_gas": E_ads_mol,          # Adsorbate energy (eV) - kept as E_gas for backward compatibility
        "E_complex": E_complex,      # Complex energy (eV)
        "complex_structure": complex_atoms,
        "optimized": optimize
    }

    return result

analyze_coordination 

analyze_coordination(atoms: Atoms, metal_indices: Optional[List[int]] = None, cutoff_multiplier: float = 1.2) -> CoordinationResult

Analyze coordination environment around metal centers.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object

 required
metal_indices Optional[List[int]]

Indices of metal atoms (auto-detect if None)

 None
cutoff_multiplier float

Multiplier for natural cutoff radii

 1.2

Returns:

Type Description
CoordinationResult

Dictionary with coordination analysis
Source code in src/mace_inference/tasks/adsorption.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def analyze_coordination(
    atoms: Atoms,
    metal_indices: Optional[List[int]] = None,
    cutoff_multiplier: float = 1.2
) -> "CoordinationResult":
    """
    Analyze coordination environment around metal centers.

    Args:
        atoms: ASE Atoms object
        metal_indices: Indices of metal atoms (auto-detect if None)
        cutoff_multiplier: Multiplier for natural cutoff radii

    Returns:
        Dictionary with coordination analysis
    """
    # Auto-detect metal atoms if not provided
    if metal_indices is None:
        # Common metal elements in MOFs
        metal_symbols = ['Cu', 'Zn', 'Zr', 'Fe', 'Ni', 'Co', 'Mn', 'Cr', 'Ti', 'V',
                        'Al', 'Mg', 'Ca', 'Sr', 'Ba', 'Cd', 'Hg', 'Pd', 'Pt']
        metal_indices = [i for i, sym in enumerate(atoms.get_chemical_symbols()) 
                        if sym in metal_symbols]

    if not metal_indices:
        raise ValueError("No metal atoms found. Specify metal_indices manually.")

    # Create neighbor list with natural cutoffs
    cutoffs = natural_cutoffs(atoms, mult=cutoff_multiplier)
    nl = NeighborList(cutoffs, skin=0.3, self_interaction=False, bothways=False)
    nl.update(atoms)

    coordination_data = {}

    for metal_idx in metal_indices:
        indices, offsets = nl.get_neighbors(metal_idx)

        # Calculate distances
        metal_pos = atoms.positions[metal_idx]
        neighbor_data = []

        for neighbor_idx, offset in zip(indices, offsets):
            neighbor_pos = atoms.positions[neighbor_idx] + offset @ atoms.cell
            distance = np.linalg.norm(neighbor_pos - metal_pos)

            neighbor_data.append({
                "index": int(neighbor_idx),
                "symbol": atoms[neighbor_idx].symbol,
                "distance": float(distance)
            })

        # Sort by distance
        neighbor_data.sort(key=lambda x: x["distance"])

        coordination_data[int(metal_idx)] = {
            "metal_symbol": atoms[metal_idx].symbol,
            "coordination_number": len(neighbor_data),
            "neighbors": neighbor_data,
            "average_distance": float(np.mean([n["distance"] for n in neighbor_data])) if neighbor_data else 0.0
        }

    result = {
        "coordination": coordination_data,
        "n_metal_centers": len(metal_indices),
        "metal_indices": metal_indices
    }

    return result

find_adsorption_sites 

find_adsorption_sites(atoms: Atoms, grid_spacing: float = 0.5, probe_radius: float = 1.2, min_distance: float = 2.0, energy_cutoff: float = None) -> List[np.ndarray]

Find potential adsorption sites in porous structure using grid-based search.

This function identifies void spaces in the structure that could accommodate adsorbate molecules by searching for grid points that are sufficiently far from all atoms.

Parameters:

Name Type Description Default
atoms Atoms

ASE Atoms object (porous structure like MOF, zeolite)

 required
grid_spacing float

Grid spacing for site search (Å)

 0.5
probe_radius float

Minimum distance from any atom to be considered a void (Å)

 1.2
min_distance float

Minimum distance between identified sites (Å)

 2.0
energy_cutoff float

Optional energy cutoff for filtering sites (not implemented)

 None

Returns:

Type Description
List[ndarray]

List of potential adsorption site positions [x, y, z]
Example

sites = find_adsorption_sites(mof, probe_radius=1.5) print(f"Found {len(sites)} potential sites") for site in sites[:5]: ... print(f" Site at {site}")

Source code in src/mace_inference/tasks/adsorption.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def find_adsorption_sites(
    atoms: Atoms,
    grid_spacing: float = 0.5,
    probe_radius: float = 1.2,
    min_distance: float = 2.0,
    energy_cutoff: float = None
) -> List[np.ndarray]:
    """
    Find potential adsorption sites in porous structure using grid-based search.

    This function identifies void spaces in the structure that could
    accommodate adsorbate molecules by searching for grid points that
    are sufficiently far from all atoms.

    Args:
        atoms: ASE Atoms object (porous structure like MOF, zeolite)
        grid_spacing: Grid spacing for site search (Å)
        probe_radius: Minimum distance from any atom to be considered a void (Å)
        min_distance: Minimum distance between identified sites (Å)
        energy_cutoff: Optional energy cutoff for filtering sites (not implemented)

    Returns:
        List of potential adsorption site positions [x, y, z]

    Example:
        >>> sites = find_adsorption_sites(mof, probe_radius=1.5)
        >>> print(f"Found {len(sites)} potential sites")
        >>> for site in sites[:5]:
        ...     print(f"  Site at {site}")
    """
    # Get cell parameters
    cell = atoms.get_cell()

    # Create grid points within the unit cell
    n_grid = [max(1, int(np.linalg.norm(cell[i]) / grid_spacing)) for i in range(3)]

    # Generate fractional coordinates
    grid_points = []
    for i in range(n_grid[0]):
        for j in range(n_grid[1]):
            for k in range(n_grid[2]):
                frac = np.array([i / n_grid[0], j / n_grid[1], k / n_grid[2]])
                grid_points.append(frac)

    grid_points = np.array(grid_points)

    # Convert to Cartesian coordinates
    cart_points = grid_points @ cell

    # Get atomic positions
    positions = atoms.get_positions()

    # Find points that are far enough from all atoms
    void_sites = []

    for point in cart_points:
        # Calculate minimum distance to any atom (considering PBC)
        min_dist = float('inf')

        for pos in positions:
            # Simple distance (could be improved with proper PBC handling)
            diff = point - pos
            # Apply minimum image convention
            frac_diff = np.linalg.solve(cell.T, diff)
            frac_diff = frac_diff - np.round(frac_diff)
            diff_pbc = frac_diff @ cell
            dist = np.linalg.norm(diff_pbc)
            min_dist = min(min_dist, dist)

        if min_dist >= probe_radius:
            void_sites.append(point)

    if not void_sites:
        return []

    void_sites = np.array(void_sites)

    # Cluster nearby sites to avoid redundancy
    selected_sites = []
    used = np.zeros(len(void_sites), dtype=bool)

    for i, site in enumerate(void_sites):
        if used[i]:
            continue

        # Mark nearby sites as used
        for j in range(i + 1, len(void_sites)):
            if used[j]:
                continue
            diff = site - void_sites[j]
            # Apply PBC
            frac_diff = np.linalg.solve(cell.T, diff)
            frac_diff = frac_diff - np.round(frac_diff)
            diff_pbc = frac_diff @ cell
            if np.linalg.norm(diff_pbc) < min_distance:
                used[j] = True

        selected_sites.append(site)
        used[i] = True

    return selected_sites

calculate_adsorption_energy

def calculate_adsorption_energy(
    framework: Atoms,
    adsorbate: Atoms,
    site_position: List[float],
    calculator: Calculator,
    optimize: bool = True,
    fmax: float = 0.05,
    fix_framework: bool = True
) -> dict

Calculate adsorption energy of a gas molecule in a framework.

The adsorption energy is calculated as:

\[E_{ads} = E_{complex} - E_{framework} - E_{gas}\]