The Codes of Crystal Structure Optimisation with FAIRChem

Following our guide, Getting Started with FAIRChem, which covers the initial setup, this article delves into the practical code implementation for optimising a material's structure. We will walk through setting up your environment, defining atomic constraints, running the optimisation, and even parallelising the workload for efficiency.

Setting Up the Workspace

First, we need to define and create the necessary directories for our project. This helps keep our initial structures, optimised results, and calculation files organised.

from pathlib import Path

base_dir = Path(__file__).parent.parent
initial_dir = base_dir / "1_structures" / "initial_structures"
optimized_dir = base_dir / "1_structures" / "optimized_structures"
calc_dir = base_dir / "2_calculations"

# Create the necessary directories
optimized_dir.mkdir(parents=True, exist_ok=True)
calc_dir.mkdir(parents=True, exist_ok=True)

Identifying Atoms and Layers

For many simulations, we need to selectively manipulate or constrain parts of the structure. Here, we identify the aluminium (Al) atoms and then sort them by their z-coordinate to determine the different atomic layers in the crystal.

Identifying Specific Atoms

# Identify Al atoms and OH molecules
symbols = atoms.get_chemical_symbols()
al_indices = [i for i, sym in enumerate(symbols) if sym == 'Al']

Sorting Atoms into Layers

By rounding the z-coordinates, we can group atoms into distinct layers. This is crucial for applying layer-specific constraints, such as fixing the bottom of a slab.

import numpy as np

al_positions = atoms.positions[al_indices]
z_coords = al_positions[:, 2]
unique_z = np.unique(np.round(z_coords, decimals=3))
unique_z_sorted = np.sort(unique_z)
n_layers = len(unique_z_sorted)
n_fixed_layers = 3  # We will fix the bottom 3 layers

Applying Atomic Constraints

To simulate a surface or interface, it is common practice to fix the bottom layers of a crystal slab while allowing the top layers to relax. In this example, we fix the bottom three layers of Al atoms, while leaving all hydroxyl (OH) groups and the remaining Al atoms free to move.

from ase.constraints import FixAtoms

# Set a threshold for the layers to be fixed
z_threshold = unique_z_sorted[n_fixed_layers -1]

# Set the constraint to fix atoms (only Al atoms are fixed, OH groups are fully relaxed)
fixed_indices = [i for i in al_indices if atoms.positions[i, 2] <= z_threshold]
constraint = FixAtoms(indices=fixed_indices)
atoms.set_constraint(constraint)

Running the Geometry Optimisation

With the constraints in place, we can now set up and run the geometry optimisation using FAIRChem's calculator. We will use the Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm to relax the structure until the maximum force on any atom falls below a specified threshold.

from ase.optimize import BFGS
from fairchem.calculator import FAIRChemCalculator

# Configure the FAIRChem calculator
calc = FAIRChemCalculator(predictor, task_name="oc20")
atoms.calc = calc

# Set the path for the optimisation trajectory file
traj_file = calc_dir / f"{vasp_file.stem}_optimization.traj"

# Execute the geometry optimisation
optimizer = BFGS(atoms, trajectory=str(traj_file), logfile=None) # logfile=None prevents write conflicts in multiprocessing
optimizer.run(fmax=0.01, steps=200)  # Set max steps to 200

Parallelising the Calculations

Optimising multiple structures can be time-consuming. We can significantly speed up this process by running the optimisations in parallel using Python's multiprocessing module. This allows us to leverage all available CPU cores to handle several calculations simultaneously.

import os
import multiprocessing

# Use os.cpu_count() to get the number of CPU cores, adjust as needed
num_processes = os.cpu_count()
print(f"Starting parallel optimisation using {num_processes} CPU cores...")
print("-" * 80)

# Create a process pool and execute the tasks
with multiprocessing.Pool(processes=num_processes) as pool:
    results = pool.map(optimize_structure, tasks)

This concludes our walkthrough of a typical geometry optimisation workflow using FAIRChem. By following these steps, you can effectively prepare and run your own materials simulations.