Silicon

Dependencies

/// script requires-python = ">=3.10" dependencies = [ "a2c-ase @ git+https://github.com/abhijeetgangan/a2c_ase.git", "numpy", "pymatgen", "tqdm", "mace-torch", ] ///

Silicon Crystallization Example

This example demonstrates the complete a2c workflow for predicting the crystal structure of silicon from an amorphous precursor using the MACE machine learning potential.

Setup and Imports¶

We'll use MACE for accurate energy and force calculations.

In [1]:

import os
from collections import defaultdict

import numpy as np
from ase.build import bulk
from mace.calculators.foundations_models import mace_mp  # type: ignore
from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.core.composition import Composition
from pymatgen.core.structure import Structure
from tqdm import tqdm

from a2c_ase.runner import melt_quench_md, relax_unit_cell
from a2c_ase.utils import extract_crystallizable_subcells, random_packed_structure

import os from collections import defaultdict import numpy as np from ase.build import bulk from mace.calculators.foundations_models import mace_mp # type: ignore from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer from pymatgen.analysis.structure_matcher import StructureMatcher from pymatgen.core.composition import Composition from pymatgen.core.structure import Structure from tqdm import tqdm from a2c_ase.runner import melt_quench_md, relax_unit_cell from a2c_ase.utils import extract_crystallizable_subcells, random_packed_structure

/opt/hostedtoolcache/Python/3.10.18/x64/lib/python3.10/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))

cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.

Configuration¶

Define the system (Si64) and simulation parameters. In CI mode, we use reduced parameters for faster testing.

In [2]:

IS_CI = os.getenv("CI") is not None

comp = Composition("Si64")
cell = np.array([[11.1, 0.0, 0.0], [0.0, 11.1, 0.0], [0.0, 0.0, 11.1]])

# Optimization parameters
global_seed = 42
fmax = 0.01  # Force convergence criterion in eV/Å
max_iter = 2 if IS_CI else 100

# Molecular dynamics parameters
md_log_interval = 50
md_equi_steps = 10 if IS_CI else 2500  # High temperature equilibration steps
md_cool_steps = 10 if IS_CI else 2500  # Cooling steps
md_final_steps = 10 if IS_CI else 2500  # Low temperature equilibration steps
md_T_high = 2000.0  # Initial melting temperature (K)
md_T_low = 300.0  # Final temperature (K)
md_time_step = 2.0  # fs
md_friction = 0.01  # Langevin friction

if IS_CI:
    print("Running in CI mode with reduced parameters for fast testing")

IS_CI = os.getenv("CI") is not None comp = Composition("Si64") cell = np.array([[11.1, 0.0, 0.0], [0.0, 11.1, 0.0], [0.0, 0.0, 11.1]]) # Optimization parameters global_seed = 42 fmax = 0.01 # Force convergence criterion in eV/Å max_iter = 2 if IS_CI else 100 # Molecular dynamics parameters md_log_interval = 50 md_equi_steps = 10 if IS_CI else 2500 # High temperature equilibration steps md_cool_steps = 10 if IS_CI else 2500 # Cooling steps md_final_steps = 10 if IS_CI else 2500 # Low temperature equilibration steps md_T_high = 2000.0 # Initial melting temperature (K) md_T_low = 300.0 # Final temperature (K) md_time_step = 2.0 # fs md_friction = 0.01 # Langevin friction if IS_CI: print("Running in CI mode with reduced parameters for fast testing")

Running in CI mode with reduced parameters for fast testing

Step 1: Initialize MACE Calculator¶

MACE is a state-of-the-art machine learning potential trained on diverse materials data.

In [3]:

device = "cpu" if IS_CI else "cuda"
calculator = mace_mp(model="small-omat-0", device=device, dtype="float32")

device = "cpu" if IS_CI else "cuda" calculator = mace_mp(model="small-omat-0", device=device, dtype="float32")

Using model under Academic Software License (ASL) license, see https://github.com/gabor1/ASL 
 To use this model you accept the terms of the license.
Downloading MACE model from 'https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-small.model'

Cached MACE model to /home/runner/.cache/mace/maceomat0smallmodel
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/maceomat0smallmodel
Using float32 for MACECalculator, which is faster but less accurate. Recommended for MD. Use float64 for geometry optimization.
Using head omat_pbe out of ['omat_pbe']
Default dtype float32 does not match model dtype float64, converting models to float32.

/opt/hostedtoolcache/Python/3.10.18/x64/lib/python3.10/site-packages/mace/calculators/mace.py:197: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  torch.load(f=model_path, map_location=device)

Step 2: Generate Random Packed Structure¶

Create an initial random configuration with no severe atomic overlaps.

In [4]:

packed_atoms, log_data = random_packed_structure(
    composition=comp,
    cell=cell,
    seed=global_seed,
    fmax=fmax,
    max_iter=max_iter,
    verbose=True,
    auto_diameter=True,
    trajectory_file=None,
)
print(f"Soft sphere packed structure is ready {packed_atoms}")

packed_atoms, log_data = random_packed_structure( composition=comp, cell=cell, seed=global_seed, fmax=fmax, max_iter=max_iter, verbose=True, auto_diameter=True, trajectory_file=None, ) print(f"Soft sphere packed structure is ready {packed_atoms}")

Using random pack diameter of 2.2
Reduce atom overlap using the soft_sphere potential
Initial energy: 2.6981
Step: 0, E: 2.6981, Fmax: 0.4073, Min dist: 0.6425
Step: 1, E: 2.6761, Fmax: 0.4043, Min dist: 0.6490
Step: 2, E: 2.6326, Fmax: 0.3984, Min dist: 0.6618
Final energy: 2.6326
Soft sphere packed structure is ready Atoms(symbols='Si64', pbc=True, cell=[11.1, 11.1, 11.1], calculator=SoftSphere(...))

Step 3: Melt-Quench MD Simulation¶

Heat to 2000K, quench to 300K to create an amorphous structure.

In [5]:

amorphous_atoms, md_log = melt_quench_md(
    atoms=packed_atoms,
    calculator=calculator,
    equi_steps=md_equi_steps,
    cool_steps=md_cool_steps,
    final_steps=md_final_steps,
    T_high=md_T_high,
    T_low=md_T_low,
    time_step=md_time_step,
    friction=md_friction,
    trajectory_file=None,
    seed=global_seed,
    verbose=True,
    log_interval=md_log_interval,
)
print(f"Final amorphous structure is ready {amorphous_atoms}")

amorphous_atoms, md_log = melt_quench_md( atoms=packed_atoms, calculator=calculator, equi_steps=md_equi_steps, cool_steps=md_cool_steps, final_steps=md_final_steps, T_high=md_T_high, T_low=md_T_low, time_step=md_time_step, friction=md_friction, trajectory_file=None, seed=global_seed, verbose=True, log_interval=md_log_interval, ) print(f"Final amorphous structure is ready {amorphous_atoms}")

Step 0/30: T = 2000.0 K, E_pot = 630.546 eV, E_kin = 16.545 eV

Melt-quench simulation completed:
Final temperature: 83754.3 K
Final energy: -182.198 eV
Final amorphous structure is ready Atoms(symbols='Si64', pbc=True, cell=[11.1, 11.1, 11.1], momenta=..., calculator=MACECalculator(...))

Step 4: Extract Crystallizable Subcells¶

Divide the amorphous structure into overlapping subcells.

In [6]:

crystallizable_cells = extract_crystallizable_subcells(
    atoms=amorphous_atoms,
    d_frac=0.2,  # Size of subcell as fraction of original cell
    n_min=2,  # Min atoms per subcell
    n_max=8,  # Max atoms per subcell
    cubic_only=True,
    allowed_atom_counts=[2, 4, 8],
)
print(f"Generated {len(crystallizable_cells)} candidate structures for crystallization")

crystallizable_cells = extract_crystallizable_subcells( atoms=amorphous_atoms, d_frac=0.2, # Size of subcell as fraction of original cell n_min=2, # Min atoms per subcell n_max=8, # Max atoms per subcell cubic_only=True, allowed_atom_counts=[2, 4, 8], ) print(f"Generated {len(crystallizable_cells)} candidate structures for crystallization")

Created 1915 subcells from amorphous structure
Subcells kept after filtering: 39
Generated 39 candidate structures for crystallization

Step 5: Optimize Subcells¶

Relax each subcell to find stable crystalline structures.

In [7]:

relaxed_atoms_list = []
print("Relaxing candidate structures...")
for atoms in tqdm(crystallizable_cells):
    # Relax the structure
    relaxed_atoms, logger = relax_unit_cell(
        atoms=atoms, calculator=calculator, max_iter=max_iter, fmax=fmax, verbose=False
    )

    # Get final energy and pressure
    final_energy = relaxed_atoms.get_potential_energy()
    final_pressure = -np.trace(relaxed_atoms.get_stress(voigt=False)) / 3.0

    # Store the relaxed structure and its properties
    relaxed_atoms_list.append((relaxed_atoms, logger, final_energy, final_pressure))

relaxed_atoms_list = [] print("Relaxing candidate structures...") for atoms in tqdm(crystallizable_cells): # Relax the structure relaxed_atoms, logger = relax_unit_cell( atoms=atoms, calculator=calculator, max_iter=max_iter, fmax=fmax, verbose=False ) # Get final energy and pressure final_energy = relaxed_atoms.get_potential_energy() final_pressure = -np.trace(relaxed_atoms.get_stress(voigt=False)) / 3.0 # Store the relaxed structure and its properties relaxed_atoms_list.append((relaxed_atoms, logger, final_energy, final_pressure))

Relaxing candidate structures...

  0%|          | 0/39 [00:00<?, ?it/s]

  3%|▎         | 1/39 [00:00<00:17,  2.16it/s]

  5%|▌         | 2/39 [00:00<00:18,  1.99it/s]

  8%|▊         | 3/39 [00:01<00:17,  2.06it/s]

 10%|█         | 4/39 [00:01<00:16,  2.07it/s]

 13%|█▎        | 5/39 [00:02<00:15,  2.19it/s]

 15%|█▌        | 6/39 [00:02<00:14,  2.24it/s]

 18%|█▊        | 7/39 [00:03<00:16,  1.97it/s]

 21%|██        | 8/39 [00:04<00:18,  1.64it/s]

 23%|██▎       | 9/39 [00:04<00:16,  1.86it/s]

 26%|██▌       | 10/39 [00:05<00:17,  1.66it/s]

 28%|██▊       | 11/39 [00:05<00:15,  1.77it/s]

 31%|███       | 12/39 [00:06<00:13,  2.01it/s]

 33%|███▎      | 13/39 [00:06<00:12,  2.12it/s]

 36%|███▌      | 14/39 [00:07<00:13,  1.91it/s]

 38%|███▊      | 15/39 [00:07<00:11,  2.00it/s]

 41%|████      | 16/39 [00:08<00:11,  1.96it/s]

 44%|████▎     | 17/39 [00:08<00:11,  1.98it/s]

 46%|████▌     | 18/39 [00:09<00:11,  1.90it/s]

 49%|████▊     | 19/39 [00:09<00:09,  2.02it/s]

 51%|█████▏    | 20/39 [00:10<00:08,  2.30it/s]

 54%|█████▍    | 21/39 [00:10<00:08,  2.10it/s]

 56%|█████▋    | 22/39 [00:11<00:08,  2.02it/s]

 59%|█████▉    | 23/39 [00:11<00:07,  2.12it/s]

 62%|██████▏   | 24/39 [00:12<00:07,  1.89it/s]

 64%|██████▍   | 25/39 [00:12<00:06,  2.05it/s]

 67%|██████▋   | 26/39 [00:12<00:05,  2.21it/s]

 69%|██████▉   | 27/39 [00:13<00:05,  2.27it/s]

 72%|███████▏  | 28/39 [00:13<00:05,  2.03it/s]

 74%|███████▍  | 29/39 [00:14<00:05,  1.88it/s]

 77%|███████▋  | 30/39 [00:15<00:05,  1.73it/s]

 79%|███████▉  | 31/39 [00:15<00:04,  1.87it/s]

 82%|████████▏ | 32/39 [00:16<00:03,  2.09it/s]

 85%|████████▍ | 33/39 [00:16<00:02,  2.29it/s]

 87%|████████▋ | 34/39 [00:16<00:02,  2.23it/s]

 90%|████████▉ | 35/39 [00:17<00:01,  2.47it/s]

 92%|█████████▏| 36/39 [00:17<00:01,  2.63it/s]

 95%|█████████▍| 37/39 [00:17<00:00,  2.48it/s]

 97%|█████████▋| 38/39 [00:18<00:00,  2.48it/s]

100%|██████████| 39/39 [00:18<00:00,  2.24it/s]

100%|██████████| 39/39 [00:18<00:00,  2.06it/s]

Step 6: Analyze Results¶

Find the lowest energy structure and determine its space group.

In [8]:

lowest_e_candidate = min(relaxed_atoms_list, key=lambda x: x[-2] / len(x[0]))
lowest_e_atoms, lowest_e_logger, lowest_e_energy, lowest_e_pressure = lowest_e_candidate

# Convert to pymatgen structure for space group analysis
pymatgen_struct = Structure(
    lattice=lowest_e_atoms.get_cell(),
    species=lowest_e_atoms.get_chemical_symbols(),
    coords=lowest_e_atoms.get_positions(),
    coords_are_cartesian=True,
)

# Analyze space group
spg = SpacegroupAnalyzer(pymatgen_struct)
print("Space group of predicted crystallization product:", spg.get_space_group_symbol())
print(
    f"Final energy: {lowest_e_energy:.4f} eV, "
    f"Energy per atom: {lowest_e_energy / len(lowest_e_atoms):.4f} eV/atom"
)
print(f"Final pressure: {lowest_e_pressure:.6f} eV/ų")

# Count frequency of space groups across all candidates
spg_counter = defaultdict(lambda: 0)

for s in relaxed_atoms_list:
    pymatgen_struct = Structure(
        lattice=s[0].get_cell(),
        species=s[0].get_chemical_symbols(),
        coords=s[0].get_positions(),
        coords_are_cartesian=True,
    )
    try:
        sp = SpacegroupAnalyzer(pymatgen_struct).get_space_group_symbol()
        spg_counter[sp] += 1
    except TypeError:
        continue


print("All space groups encountered:", dict(spg_counter))

# Compare to reference diamond structure
si_diamond = bulk("Si", "diamond", a=5.43)
pymatgen_ref_struct = Structure(
    lattice=si_diamond.get_cell(),
    species=si_diamond.get_chemical_symbols(),
    coords=si_diamond.get_positions(),
    coords_are_cartesian=True,
)
print(
    "Prediction matches diamond-cubic Si?",
    StructureMatcher().fit(pymatgen_struct, pymatgen_ref_struct),
)

# Save the lowest energy structure
lowest_e_atoms.write("final_crystal_structure.cif")
print("Lowest energy structure saved to final_crystal_structure.cif")

lowest_e_candidate = min(relaxed_atoms_list, key=lambda x: x[-2] / len(x[0])) lowest_e_atoms, lowest_e_logger, lowest_e_energy, lowest_e_pressure = lowest_e_candidate # Convert to pymatgen structure for space group analysis pymatgen_struct = Structure( lattice=lowest_e_atoms.get_cell(), species=lowest_e_atoms.get_chemical_symbols(), coords=lowest_e_atoms.get_positions(), coords_are_cartesian=True, ) # Analyze space group spg = SpacegroupAnalyzer(pymatgen_struct) print("Space group of predicted crystallization product:", spg.get_space_group_symbol()) print( f"Final energy: {lowest_e_energy:.4f} eV, " f"Energy per atom: {lowest_e_energy / len(lowest_e_atoms):.4f} eV/atom" ) print(f"Final pressure: {lowest_e_pressure:.6f} eV/ų") # Count frequency of space groups across all candidates spg_counter = defaultdict(lambda: 0) for s in relaxed_atoms_list: pymatgen_struct = Structure( lattice=s[0].get_cell(), species=s[0].get_chemical_symbols(), coords=s[0].get_positions(), coords_are_cartesian=True, ) try: sp = SpacegroupAnalyzer(pymatgen_struct).get_space_group_symbol() spg_counter[sp] += 1 except TypeError: continue print("All space groups encountered:", dict(spg_counter)) # Compare to reference diamond structure si_diamond = bulk("Si", "diamond", a=5.43) pymatgen_ref_struct = Structure( lattice=si_diamond.get_cell(), species=si_diamond.get_chemical_symbols(), coords=si_diamond.get_positions(), coords_are_cartesian=True, ) print( "Prediction matches diamond-cubic Si?", StructureMatcher().fit(pymatgen_struct, pymatgen_ref_struct), ) # Save the lowest energy structure lowest_e_atoms.write("final_crystal_structure.cif") print("Lowest energy structure saved to final_crystal_structure.cif")

Space group of predicted crystallization product: P1
Final energy: -18.5615 eV, Energy per atom: -4.6404 eV/atom
Final pressure: -0.018813 eV/ų
All space groups encountered: {'P1': 25, 'P-1': 13, 'P2_1/m': 1}
Prediction matches diamond-cubic Si? False
Lowest energy structure saved to final_crystal_structure.cif