create_surface_slab.py

import numpy as np
from ase import io
from ase.build import surface, sort
from ase.constraints import FixAtoms
from ase.visualize import view

# ------------------- USER SETTINGS ---------------------
input_file = "TiN.cif"                  # Input bulk structure
surface_output_file_name = "Out"        # Output CIF base name
surface_output_file_name_POSCAR = "POSCAR_235"  # Output POSCAR file name

indices = (2, 3, 5)      # Miller indices
layers = 20               # Number of layers
vacuum = 10.0             # Vacuum thickness in Angstrom
supercell = (1, 1, 1)     # Supercell repetition
thickness_to_fix = 4.0    # Thickness from bottom to fix atoms (Angstrom)

show_structure = True     # Set False on clusters/headless servers
# ------------------------------------------------------


def get_top_bottom_info(atoms):
    """Return info about the top-most and bottom-most atoms."""
    positions = atoms.get_positions()
    symbols = atoms.get_chemical_symbols()

    min_z_idx = np.argmin(positions[:, 2])
    max_z_idx = np.argmax(positions[:, 2])

    return {
        "bottom_element": symbols[min_z_idx],
        "bottom_z": positions[min_z_idx, 2],
        "top_element": symbols[max_z_idx],
        "top_z": positions[max_z_idx, 2],
    }


def get_fixed_indices_by_thickness(atoms, thickness):
    """Return atom indices within 'thickness' Angstrom from the bottom."""
    positions = atoms.get_positions()
    min_z = np.min(positions[:, 2])
    return [i for i, pos in enumerate(positions) if pos[2] < min_z + thickness]


# 1) Read input structure
atoms_input = io.read(input_file)

# 2) Build slab
atoms = surface(atoms_input, indices, layers, vacuum)
atoms = atoms.repeat(supercell)

print("\n--------------------------------------------------------------------\n")
print(
    f"Successfully created surface with Miller indices {indices}, "
    f"{layers} layers, {vacuum} Å vacuum, and supercell {supercell}\n"
    f"from input structure: {input_file}\n"
)

# 3) Report top/bottom atoms
info = get_top_bottom_info(atoms)
print(f"The bottom-most element is {info['bottom_element']} with z = {info['bottom_z']:.6f} Å")
print(f"The top-most element is {info['top_element']} with z = {info['top_z']:.6f} Å")

# 4) Save slab to CIF
cif_filename = surface_output_file_name + ".cif"
io.write(cif_filename, atoms, format="cif")

# 5) Print unique species in original order of appearance
symbols = atoms.get_chemical_symbols()
unique_symbols = list(dict.fromkeys(symbols))
print(unique_symbols)
print(f"Elements in structure: {unique_symbols}")
print("\n--------------------------------------------------------------------\n")

# 6) Prepare POSCAR structure
# Sort atoms by chemical symbol so POSCAR groups species correctly
atoms_poscar = atoms.copy()
atoms_poscar = sort(atoms_poscar)

# 7) Apply constraints for bottom region
fixed_indices = get_fixed_indices_by_thickness(atoms_poscar, thickness_to_fix)
constraint = FixAtoms(indices=fixed_indices)
atoms_poscar.set_constraint(constraint)

print(f"Number of fixed atoms: {len(fixed_indices)}")
print(f"Atoms with z < (z_min + {thickness_to_fix} Å) are fixed.")

# 8) Save POSCAR
io.write(surface_output_file_name_POSCAR, atoms_poscar, format="vasp", direct=True)

# 9) Save PNG snapshot
atoms_png = atoms.copy()
atoms_png.set_tags(np.zeros(len(atoms_png), dtype=int))
io.write(surface_output_file_name_POSCAR + ".png", atoms_png, format="png")

# 10) Optional visualization
if show_structure:
    view(atoms)