#!/usr/bin/env amspython """Run a parallel ReaxFF tribology case with AMS/PLAMS. Run from the directory containing the slab files with: "$AMSBIN/amspython" -u tribology_case.py All scientific and resource controls are in the USER SETTINGS section below. The input slabs must already be terminated or passivated and periodic in x and y. Relative slab paths are searched from the working directory, the script directory, and the script directory's parent. """ from __future__ import annotations import csv from dataclasses import dataclass from math import ceil from pathlib import Path import matplotlib import numpy as np matplotlib.use("Agg") import matplotlib.pyplot as plt from scm.base import Units from scm.plams import ( AMSJob, AMSNVTJob, JobRunner, KFHistory, Molecule, Settings, config, finish, from_smiles, init, ) from scm.plams.interfaces.molecule.packmol import packmol # ============================== USER SETTINGS =============================== TOP_SLAB_FILE = Path("Fe2O3_0001_H2O_rectangular.xyz") BOTTOM_SLAB_FILE = Path("Fe2O3_0001_H2O_rectangular.xyz") FLIP_TOP_SLAB = True # Put the opposite face of the supplied top slab at the contact. CASE = "stearic" # dry, hexadecane, oleic, stearic TARGET_XY_A = (20.0, 35.0) # Approximate target size; slabs are repeated by whole unit cells. INITIAL_GAP_A = 20.0 FILM_DENSITY_G_CM3 = 0.77 PACKMOL_TOLERANCE_A = 2.0 FLUID_SURFACE_CLEARANCE_A = 1.0 # A scalar is accepted for a one-load or one-seed debug run. For several # values use, for example, ``(2.5, 5.0, 10.0)`` or ``(101, 202)``. LOADS_NN = (7.5, 12.5, 15.0) SEEDS = (101, 202) TEMPERATURE_K = 300.0 # +y sliding; change together with SHEAR_FORCE_AXIS if needed. VELOCITY_A_PER_FS = (0.0, 0.0010, 0.0) TIMESTEP_FS = 0.25 EQUILIBRATION_STEPS = 400_000 PRODUCTION_STEPS = 800_000 EQUILIBRATION_TRAJECTORY_SAMPLING_STEPS = 1000 PRODUCTION_TRAJECTORY_SAMPLING_STEPS = 1000 ANALYSIS_BLOCK_STEPS = 10_000 # Production gaps are averaged after this initial fraction of the NEMD run. GAP_ANALYSIS_DISCARD_FRACTION = 0.20 # Geometry optimizations are deliberately bounded pre-optimizations. Their # output is used to remove bad contacts before MD, not as a convergence claim. SLAB_PREOPT_MAX_ITERATIONS = 100 INTERFACE_PREOPT_MAX_ITERATIONS = 100 # Fractions of the thickness of each prepared slab, measured from its outside # face. The unassigned remainder is the contact layer. FIXED_LAYER_FRACTION = 0.20 THERMOSTAT_LAYER_FRACTION = 0.25 DRIVE_LAYER_FRACTION = 0.20 # Gap definition. The gap is the average z position of the lowest Fe layer in # the top pad minus the average z position of the highest Fe layer in the # bottom pad. For another material, set the two element symbols accordingly. # GAP_LAYER_SELECTION_TOLERANCE_A selects all matching atoms within this z # distance of the appropriate outermost metal layer in the initial structure. BOTTOM_PAD_GAP_ELEMENT = "Fe" TOP_PAD_GAP_ELEMENT = "Fe" GAP_LAYER_SELECTION_TOLERANCE_A = 0.50 # The y component is the current sliding direction. This setting selects the # tangential force used in both the plots and friction analysis. The CSV always # contains Fx, Fy, and Fz. SHEAR_FORCE_AXIS = "y" THERMOSTAT = "NHC" THERMOSTAT_TAU_FS = 100.0 FORCE_FIELD = "FeOCHCl-ox.ff" # Every independent equilibration and every load x seed NEMD run is submitted # through PLAMS' parallel JobRunner. The maximum CPU demand is # MAX_PARALLEL_JOBS * NPROC_PER_JOB. if isinstance(LOADS_NN, (int, float)): LOADS_NN = (float(LOADS_NN),) else: LOADS_NN = tuple(float(load) for load in LOADS_NN) if isinstance(SEEDS, int): SEEDS = (SEEDS,) else: SEEDS = tuple(int(seed) for seed in SEEDS) MAX_PARALLEL_JOBS = len(LOADS_NN) * len(SEEDS) NPROC_PER_JOB = 8 # ============================================================================= @dataclass(frozen=True) class Component: """A component in a Packmol mixture, specified by its mole fraction.""" label: str smiles: str mole_fraction: float CASE_COMPONENTS: dict[str, tuple[Component, ...]] = { "dry": (), "hexadecane": (Component("hexadecane", "CCCCCCCCCCCCCCCC", 1.0),), "oleic": ( Component("hexadecane", "CCCCCCCCCCCCCCCC", 0.75), Component("oleic_acid", "CCCCCCCCC/C=C\\CCCCCCCC(=O)O", 0.25), ), "stearic": ( Component("hexadecane", "CCCCCCCCCCCCCCCC", 0.75), Component("stearic_acid", "CCCCCCCCCCCCCCCCCC(=O)O", 0.25), ), } def validate_user_settings() -> None: """Validate user-editable settings before starting any AMS jobs.""" components = CASE_COMPONENTS.get(CASE) if components is None: raise ValueError(f"Unknown CASE={CASE!r}; choose from {', '.join(CASE_COMPONENTS)}.") if components: fractions = [item.mole_fraction for item in components] if any(fraction <= 0.0 for fraction in fractions) or not np.isclose( sum(fractions), 1.0 ): raise ValueError(f"Mole fractions for CASE={CASE!r} must be positive and sum to 1.") if not LOADS_NN or any(load <= 0.0 for load in LOADS_NN): raise ValueError("LOADS_NN must contain at least one positive load.") if len(set(LOADS_NN)) != len(LOADS_NN): raise ValueError("LOADS_NN must not contain duplicate values.") if not SEEDS or len(set(SEEDS)) != len(SEEDS): raise ValueError("SEEDS must contain at least one value and no duplicates.") if TIMESTEP_FS <= 0.0 or EQUILIBRATION_STEPS <= 0 or PRODUCTION_STEPS <= 0: raise ValueError("The MD timestep and step counts must be positive.") if TEMPERATURE_K <= 0.0 or NPROC_PER_JOB <= 0: raise ValueError("TEMPERATURE_K and NPROC_PER_JOB must be positive.") if any(length <= 0.0 for length in TARGET_XY_A): raise ValueError("Both TARGET_XY_A dimensions must be positive.") if INITIAL_GAP_A <= 0.0 or FILM_DENSITY_G_CM3 <= 0.0 or PACKMOL_TOLERANCE_A <= 0.0: raise ValueError("The gap, fluid density, and Packmol tolerance must be positive.") if GAP_LAYER_SELECTION_TOLERANCE_A <= 0.0 or THERMOSTAT_TAU_FS <= 0.0: raise ValueError("The gap-layer tolerance and thermostat time constant must be positive.") if SLAB_PREOPT_MAX_ITERATIONS <= 0 or INTERFACE_PREOPT_MAX_ITERATIONS <= 0: raise ValueError("Geometry-optimization iteration limits must be positive.") if EQUILIBRATION_TRAJECTORY_SAMPLING_STEPS <= 0 or PRODUCTION_TRAJECTORY_SAMPLING_STEPS <= 0: raise ValueError("Trajectory sampling intervals must be positive.") if ANALYSIS_BLOCK_STEPS <= 0: raise ValueError("ANALYSIS_BLOCK_STEPS must be positive.") if not 0.0 <= GAP_ANALYSIS_DISCARD_FRACTION < 1.0: raise ValueError("GAP_ANALYSIS_DISCARD_FRACTION must be in [0, 1).") if not 0.0 < 2.0 * FLUID_SURFACE_CLEARANCE_A < INITIAL_GAP_A: raise ValueError("FLUID_SURFACE_CLEARANCE_A must leave a positive packing interval.") layer_fractions = ( FIXED_LAYER_FRACTION, THERMOSTAT_LAYER_FRACTION, DRIVE_LAYER_FRACTION, ) if any(not 0.0 < fraction < 1.0 for fraction in layer_fractions): raise ValueError("All solid-layer fractions must be between 0 and 1.") if FIXED_LAYER_FRACTION + THERMOSTAT_LAYER_FRACTION >= 1.0: raise ValueError("The bottom fixed and thermostat fractions must sum to less than 1.") if DRIVE_LAYER_FRACTION + THERMOSTAT_LAYER_FRACTION >= 1.0: raise ValueError("The top drive and thermostat fractions must sum to less than 1.") shear_axis = SHEAR_FORCE_AXIS.lower() if shear_axis not in {"x", "y"}: raise ValueError("SHEAR_FORCE_AXIS must be 'x' or 'y'.") shear_index = {"x": 0, "y": 1}[shear_axis] if len(VELOCITY_A_PER_FS) != 3 or VELOCITY_A_PER_FS[shear_index] == 0.0: raise ValueError("VELOCITY_A_PER_FS must be nonzero along SHEAR_FORCE_AXIS.") def locate_slab(path: Path) -> Path: """Resolve a slab path from common website-download layouts.""" if path.is_absolute(): candidates = [path] else: script_directory = Path(__file__).resolve().parent candidates = [ Path.cwd() / path, script_directory / path, script_directory.parent / path, ] checked: list[Path] = [] for candidate in candidates: resolved = candidate.resolve() if resolved not in checked: checked.append(resolved) if resolved.is_file(): return resolved locations = "\n - ".join(str(candidate) for candidate in checked) raise FileNotFoundError(f"Slab file {path!s} was not found. Checked:\n - {locations}") def runscript_settings() -> Settings: """Return shared AMS resource settings.""" settings = Settings() settings.runscript.nproc = NPROC_PER_JOB settings.runscript.preamble_lines = ["export OMP_NUM_THREADS=1"] return settings def reaxff_settings() -> Settings: """Return the common ReaxFF engine settings.""" settings = runscript_settings() settings.input.ReaxFF.ForceField = FORCE_FIELD return settings def add_region(atom, region: str) -> None: """Add one AMS region without removing an atom's existing regions.""" previous = atom.properties.get("region", set()) if isinstance(previous, str): previous = {previous} else: previous = set(previous) atom.properties.region = previous | {region} def slab_extents(slab: Molecule) -> tuple[float, float]: """Return the lowest and highest z coordinates in Angstrom.""" z = [atom.z for atom in slab] return min(z), max(z) def in_plane_lengths(slab: Molecule) -> tuple[float, float]: """Return x/y lengths for the rectangular two-dimensional lattice.""" if len(slab.lattice) != 2: raise ValueError( "Input slabs must have exactly two lattice vectors (x and y periodicity)." ) first, second = (np.asarray(vector, dtype=float) for vector in slab.lattice) if not np.allclose(first[1:], 0.0, atol=1.0e-6) or not np.allclose( second[[0, 2]], 0.0, atol=1.0e-6 ): raise ValueError("This workflow requires positive, axis-aligned rectangular x/y vectors.") lx, ly = float(first[0]), float(second[1]) if lx <= 0.0 or ly <= 0.0: raise ValueError("The x/y lattice-vector lengths must be positive.") return lx, ly def make_slab_relaxation_job(name: str, slab: Molecule) -> AMSJob: """Relax a prepared slab, including its two periodic in-plane lattice vectors.""" settings = reaxff_settings() settings.input.ams.Task = "GeometryOptimization" settings.input.ams.GeometryOptimization.OptimizeLattice = True settings.input.ams.GeometryOptimization.MaxIterations = SLAB_PREOPT_MAX_ITERATIONS settings.input.ams.GeometryOptimization.PretendConverged = True # The supplied slabs are rectangular. Relax their lengths, but retain a # rectangular in-plane cell so the subsequent 2D replication and Packmol # box are well defined. settings.input.ams.Constraints.FreezeStrain = "xy" return AMSJob(name=name, molecule=slab, settings=settings) def repeat_to_target(slab: Molecule) -> Molecule: """Repeat the relaxed slab by whole unit cells to reach TARGET_XY_A.""" lx, ly = in_plane_lengths(slab) return slab.supercell(ceil(TARGET_XY_A[0] / lx), ceil(TARGET_XY_A[1] / ly)) def reflect_z(slab: Molecule) -> None: """Reflect a slab through its z midpoint, retaining the lattice.""" zmin, zmax = slab_extents(slab) for atom in slab: atom.coords = (atom.x, atom.y, zmin + zmax - atom.z) def assign_solid_regions(bottom: Molecule, top: Molecule) -> None: """Assign regions from fractions of each slab's z thickness.""" bottom_min, bottom_max = slab_extents(bottom) bottom_thickness = bottom_max - bottom_min bottom_fixed_end = bottom_min + FIXED_LAYER_FRACTION * bottom_thickness bottom_thermo_end = bottom_fixed_end + THERMOSTAT_LAYER_FRACTION * bottom_thickness for atom in bottom: if atom.z <= bottom_fixed_end: add_region(atom, "bottom_fixed") elif atom.z <= bottom_thermo_end: add_region(atom, "bottom_thermo") else: add_region(atom, "bottom_contact") top_min, top_max = slab_extents(top) top_thickness = top_max - top_min top_drive_start = top_max - DRIVE_LAYER_FRACTION * top_thickness top_thermo_start = top_drive_start - THERMOSTAT_LAYER_FRACTION * top_thickness for atom in top: if atom.z >= top_drive_start: add_region(atom, "top_drive") elif atom.z >= top_thermo_start: add_region(atom, "top_thermo") else: add_region(atom, "top_contact") def stack_slabs(bottom: Molecule, top: Molecule) -> Molecule: """Place independently supplied slabs around INITIAL_GAP_A and assign regions.""" bottom = bottom.copy() top = top.copy() if FLIP_TOP_SLAB: reflect_z(top) bottom_lx, bottom_ly = in_plane_lengths(bottom) top_lx, top_ly = in_plane_lengths(top) if not np.allclose((bottom_lx, bottom_ly), (top_lx, top_ly), atol=1.0e-5): raise ValueError("Top and bottom slabs must have matching in-plane supercells.") bottom_min, bottom_max = slab_extents(bottom) bottom.translate((0.0, 0.0, -bottom_min)) top_min, _ = slab_extents(top) top.translate((0.0, 0.0, bottom_max - bottom_min + INITIAL_GAP_A - top_min)) assign_solid_regions(bottom, top) interface = Molecule() interface += bottom interface += top interface.lattice = bottom.lattice return interface def pack_fluid(interface: Molecule, seed: int) -> Molecule: """Pack the selected liquid mixture only in the initial inter-slab gap.""" components = CASE_COMPONENTS.get(CASE) if components is None: raise ValueError(f"Unknown CASE={CASE!r}; choose from {', '.join(CASE_COMPONENTS)}.") if not components: return interface lx, ly = in_plane_lengths(interface) molecules = [from_smiles(component.smiles) for component in components] mole_fractions = [component.mole_fraction for component in components] fluid, details = packmol( molecules=molecules, mole_fractions=mole_fractions, density=FILM_DENSITY_G_CM3, box_bounds=[ 0.0, 0.0, FLUID_SURFACE_CLEARANCE_A, lx, ly, INITIAL_GAP_A - FLUID_SURFACE_CLEARANCE_A, ], tolerance=PACKMOL_TOLERANCE_A, seed=seed, region_names=["fluid"] * len(molecules), return_details=True, ) # Align Packmol's lower atom with a small clearance above the bottom # contact surface. The matching upper clearance avoids initial overlaps # with either slab. fluid_zmin, _ = slab_extents(fluid) fluid.translate( (0.0, 0.0, region_z_max(interface, "bottom_") + FLUID_SURFACE_CLEARANCE_A - fluid_zmin) ) interface += fluid interface.lattice = interface.lattice[:2] component_counts = ", ".join( f"{component.label}={count}" for component, count in zip(components, details["n_molecules"]) ) print( f"Packed {sum(details['n_molecules'])} molecules ({component_counts}) " f"at {details['density']:.3f} g/cm^3 for {CASE}." ) return interface def constrained_settings() -> Settings: """Return settings that keep the bottom support fixed.""" settings = reaxff_settings() settings.input.ams.Constraints.FixedRegion = ["bottom_fixed"] return settings def make_interface_relaxation_job(name: str, molecule: Molecule) -> AMSJob: """Relax Packmol contacts while retaining the fixed bottom support.""" settings = constrained_settings() settings.input.ams.Task = "GeometryOptimization" settings.input.ams.GeometryOptimization.MaxIterations = INTERFACE_PREOPT_MAX_ITERATIONS settings.input.ams.GeometryOptimization.PretendConverged = True return AMSJob(name=name, molecule=molecule, settings=settings) def apply_normal_load(md: Settings, load_nn: float) -> None: """Apply one total compressive normal load to the top drive layer.""" md.ApplyForce.Force = f"0.0 0.0 {-load_nn:.12g}e-9 [Newton]" md.ApplyForce.Region = "top_drive" md.ApplyForce.PerAtom = "No" def set_restart_velocities(md: Settings, previous_job: AMSNVTJob) -> None: """Use the final coordinates and velocities from a preceding MD stage.""" md.InitialVelocities.Type = "FromFile" del md.InitialVelocities.Temperature md.InitialVelocities.File = str((Path(previous_job.path).resolve() / "ams.rkf")) def make_equilibration_job(name: str, molecule: Molecule, load_nn: float, seed: int) -> AMSNVTJob: """Equilibrate one load/seed replica directly at its target normal load.""" settings = constrained_settings() settings.input.ams.RNGSeed = seed job = AMSNVTJob( name=name, molecule=molecule, settings=settings, nsteps=EQUILIBRATION_STEPS, timestep=TIMESTEP_FS, temperature=TEMPERATURE_K, thermostat=THERMOSTAT, tau=THERMOSTAT_TAU_FS, thermostat_region="bottom_thermo+top_thermo", samplingfreq=EQUILIBRATION_TRAJECTORY_SAMPLING_STEPS, writevelocities=True, ) apply_normal_load(job.settings.input.ams.MolecularDynamics, load_nn) return job def make_nemd_job(name: str, equilibration: AMSNVTJob, load_nn: float, seed: int) -> AMSNVTJob: """Create one load/seed NEMD job from a seed-equilibrated geometry and velocities.""" settings = constrained_settings() settings.input.ams.RNGSeed = seed + 1_000_000 job = AMSNVTJob( name=name, molecule=equilibration.results.get_main_molecule(), settings=settings, nsteps=PRODUCTION_STEPS, timestep=TIMESTEP_FS, temperature=TEMPERATURE_K, thermostat=THERMOSTAT, tau=THERMOSTAT_TAU_FS, thermostat_region="bottom_thermo+top_thermo", samplingfreq=PRODUCTION_TRAJECTORY_SAMPLING_STEPS, writevelocities=True, writeenginegradients=True, ) # Validate the reusable MD base before adding the restart-file reference. # The typed parser is intentionally not used for the latter: it treats an # absolute Unix filename as an input expression, whereas AMS accepts it. md = job.settings.input.ams.MolecularDynamics set_restart_velocities(md, equilibration) md.ApplyVelocity.Velocity = " ".join(str(component) for component in VELOCITY_A_PER_FS) md.ApplyVelocity.Components = "XY" md.ApplyVelocity.Region = "top_drive" # Keep the same normal load that established this load-specific gap; the # lateral velocity is the only new driving condition in production. apply_normal_load(md, load_nn) return job def region_atom_indices(molecule: Molecule, region: str) -> list[int]: """Return zero-based indices for atoms belonging to an AMS region.""" indices: list[int] = [] for index, atom in enumerate(molecule): regions = atom.properties.get("region", set()) if isinstance(regions, str): regions = {regions} if region in regions: indices.append(index) return indices def region_prefix_atom_indices(molecule: Molecule, prefix: str) -> list[int]: """Return zero-based indices for atoms in regions with a given prefix.""" indices: list[int] = [] for index, atom in enumerate(molecule): regions = atom.properties.get("region", set()) if isinstance(regions, str): regions = {regions} if any(region.startswith(prefix) for region in regions): indices.append(index) return indices def region_z_max(molecule: Molecule, region_prefix: str) -> float: """Return the upper z extent of all regions with the given prefix.""" z_values = [] for atom in molecule: regions = atom.properties.get("region", set()) if isinstance(regions, str): regions = {regions} if any(region.startswith(region_prefix) for region in regions): z_values.append(atom.z) if not z_values: raise RuntimeError(f"No atoms found for region prefix {region_prefix!r}.") return max(z_values) def gap_layer_atom_indices(molecule: Molecule) -> tuple[list[int], list[int]]: """Select fixed inner-facing metal layers used for the slab-plane gap.""" bottom_region = set(region_prefix_atom_indices(molecule, "bottom_")) top_region = set(region_prefix_atom_indices(molecule, "top_")) bottom_candidates = [ index for index, atom in enumerate(molecule) if atom.symbol == BOTTOM_PAD_GAP_ELEMENT and index in bottom_region ] top_candidates = [ index for index, atom in enumerate(molecule) if atom.symbol == TOP_PAD_GAP_ELEMENT and index in top_region ] if not bottom_candidates or not top_candidates: raise RuntimeError( "Gap-layer atom selection failed. Check BOTTOM_PAD_GAP_ELEMENT and " "TOP_PAD_GAP_ELEMENT." ) bottom_surface_z = max(molecule[index + 1].z for index in bottom_candidates) top_surface_z = min(molecule[index + 1].z for index in top_candidates) bottom_layer = [ index for index in bottom_candidates if molecule[index + 1].z >= bottom_surface_z - GAP_LAYER_SELECTION_TOLERANCE_A ] top_layer = [ index for index in top_candidates if molecule[index + 1].z <= top_surface_z + GAP_LAYER_SELECTION_TOLERANCE_A ] return bottom_layer, top_layer def gap_from_coordinates( coordinates: np.ndarray, bottom_layer: list[int], top_layer: list[int] ) -> np.ndarray: """Return slab-plane gaps from coordinates in Angstrom.""" return coordinates[:, top_layer, 2].mean(axis=1) - coordinates[:, bottom_layer, 2].mean(axis=1) def gap_trajectory(job: AMSNVTJob) -> tuple[np.ndarray, np.ndarray]: """Return trajectory time (ps) and fixed-layer gap (Angstrom).""" bottom_layer, top_layer = gap_layer_atom_indices(job.molecule) history = KFHistory(job.results.rkfs["ams"], "History") steps = np.asarray(list(history.iter("Step")), dtype=float) coordinates = np.asarray(list(history.iter("Coords")), dtype=float).reshape( -1, len(job.molecule), 3 ) if len(steps) == 0 or len(coordinates) != len(steps): raise RuntimeError(f"Incomplete trajectory history for {job.name}.") coordinates *= Units.conversion_factor("bohr", "angstrom") return steps * TIMESTEP_FS * 1.0e-3, gap_from_coordinates(coordinates, bottom_layer, top_layer) def gap_from_molecule(molecule: Molecule) -> float: """Return the fixed-layer slab-plane gap for one geometry in Angstrom.""" bottom_layer, top_layer = gap_layer_atom_indices(molecule) bottom_z = np.mean([molecule[index + 1].z for index in bottom_layer]) top_z = np.mean([molecule[index + 1].z for index in top_layer]) return float(top_z - bottom_z) def production_gap_statistics(job: AMSNVTJob) -> tuple[float, float]: """Return mean and standard deviation of the late-production gap in Angstrom.""" time_ps, gaps = gap_trajectory(job) discard_time_ps = ( GAP_ANALYSIS_DISCARD_FRACTION * PRODUCTION_STEPS * TIMESTEP_FS * 1.0e-3 ) selected_gaps = gaps[time_ps >= discard_time_ps] if len(selected_gaps) == 0: selected_gaps = gaps return float(selected_gaps.mean()), float(selected_gaps.std(ddof=0)) def write_csv(path: Path, columns: dict[str, np.ndarray]) -> None: """Write equally sized numeric columns to a CSV file.""" names = list(columns) if not names or len({len(columns[name]) for name in names}) != 1: raise ValueError(f"CSV columns for {path} must be present and equally sized.") with path.open("w", encoding="utf-8", newline="") as handle: writer = csv.writer(handle) writer.writerow(names) writer.writerows(zip(*(columns[name] for name in names))) def save_equilibration_diagnostics(job: AMSNVTJob) -> None: """Write gap(t) CSV and PNG next to one completed equilibration job.""" time_ps, gaps = gap_trajectory(job) directory = Path(job.path) write_csv(directory / "equilibration_gap.csv", {"time_ps": time_ps, "gap_A": gaps}) figure, axis = plt.subplots(figsize=(6, 3.5), constrained_layout=True) axis.plot(time_ps, gaps, color="tab:blue") axis.set(xlabel="Time (ps)", ylabel="Fe-layer gap (Angstrom)", title="Load equilibration") figure.savefig(directory / "equilibration_gap.png", dpi=180) plt.close(figure) def save_production_diagnostics(job: AMSNVTJob) -> None: """Write force/gap time series CSV and PNG next to one NEMD job.""" time_ps, gaps = gap_trajectory(job) top_drive = region_atom_indices(job.molecule, "top_drive") if not top_drive: raise RuntimeError(f"No atoms found in top_drive for {job.name}.") history = KFHistory(job.results.rkfs["ams"], "History") gradients = np.asarray(list(history.iter("EngineGradients")), dtype=float).reshape( -1, len(job.molecule), 3 ) forces_nn = ( -gradients[:, top_drive, :].sum(axis=1) * Units.conversion_factor("hartree/bohr", "newton") * 1.0e9 ) if len(forces_nn) != len(time_ps): raise RuntimeError(f"Trajectory length mismatch while writing diagnostics for {job.name}.") directory = Path(job.path) write_csv( directory / "production_timeseries.csv", { "time_ps": time_ps, "force_x_nN": forces_nn[:, 0], "force_y_nN": forces_nn[:, 1], "force_z_nN": forces_nn[:, 2], "gap_A": gaps, }, ) force_index = shear_force_index() force_name = SHEAR_FORCE_AXIS.lower() figure, axes = plt.subplots(3, 1, sharex=True, figsize=(6, 7), constrained_layout=True) axes[0].plot(time_ps, forces_nn[:, force_index], color="tab:red") axes[0].set(ylabel=f"F{force_name} (nN)") axes[1].plot(time_ps, forces_nn[:, 2], color="tab:purple") axes[1].set(ylabel="Fz (nN)") axes[2].plot(time_ps, gaps, color="tab:blue") axes[2].set(xlabel="Time (ps)", ylabel="Fe-layer gap (Angstrom)") figure.suptitle("NEMD force and gap diagnostics") figure.savefig(directory / "production_timeseries.png", dpi=180) plt.close(figure) def block_means(values: np.ndarray, block_size: int) -> np.ndarray: """Average complete consecutive blocks only.""" if len(values) == 0: raise ValueError("Cannot calculate block means from an empty array.") nblocks = len(values) // block_size if nblocks == 0: return values.mean(axis=0, keepdims=True) return values[: nblocks * block_size].reshape(nblocks, block_size, -1).mean(axis=1) def shear_force_index() -> int: """Return the Cartesian index of the configured in-plane shear axis.""" return {"x": 0, "y": 1}[SHEAR_FORCE_AXIS.lower()] def analyze_nemd(job: AMSNVTJob) -> tuple[float, float, float]: """Return block-averaged shear, normal, and friction magnitudes.""" top_drive = region_atom_indices(job.molecule, "top_drive") if not top_drive: raise RuntimeError(f"No atoms found in top_drive for {job.name}.") history = KFHistory(job.results.rkfs["ams"], "History") gradients = np.asarray(list(history.iter("EngineGradients")), dtype=float).reshape( -1, len(job.molecule), 3 ) # EngineGradients are the gradients returned by the engine (the negative # of forces). Project them directly as in AMS' friction-coefficient # analysis: |Gshear| is the shear magnitude and -Gz the resisting normal part. gradients_on_drive = gradients[:, top_drive, :].sum(axis=1) shear_normal = np.column_stack( (gradients_on_drive[:, shear_force_index()], -gradients_on_drive[:, 2]) ) frames_per_block = max( 1, round(ANALYSIS_BLOCK_STEPS / PRODUCTION_TRAJECTORY_SAMPLING_STEPS) ) blocks = block_means(shear_normal, frames_per_block) if np.any(blocks[:, 1] == 0.0): raise RuntimeError(f"A zero normal-force block prevents friction analysis for {job.name}.") conversion = Units.conversion_factor("hartree/bohr", "newton") * 1.0e9 # The signs are set by the chosen driving directions. Report the physical # magnitudes after block averaging so short-time force fluctuations do not # bias the result. mean_shear_nn, mean_normal_nn = np.mean(np.abs(blocks), axis=0) * conversion friction = float(np.mean(np.abs(blocks[:, 0] / blocks[:, 1]))) return float(mean_shear_nn), float(mean_normal_nn), friction def load_label(load_nn: float) -> str: """Make a deterministic, filesystem-safe representation of a load.""" return f"{load_nn:g}".replace("-", "m").replace(".", "p") def check_results(results: list, stage: str) -> None: """Stop at a stage barrier if an AMS job did not complete.""" failed = [result.job.path for result in results if not result.ok()] if failed: raise RuntimeError(f"{stage} failed: " + ", ".join(map(str, failed))) def main() -> int: """Build the selected case and run slab relaxation, equilibration, and NEMD.""" validate_user_settings() init() config.default_jobrunner = JobRunner(parallel=True, maxjobs=MAX_PARALLEL_JOBS) try: top_input = Molecule(str(locate_slab(TOP_SLAB_FILE))) bottom_input = Molecule(str(locate_slab(BOTTOM_SLAB_FILE))) slab_jobs = [ make_slab_relaxation_job(f"{CASE}_top_slab_relax", top_input), make_slab_relaxation_job(f"{CASE}_bottom_slab_relax", bottom_input), ] slab_results = [job.run() for job in slab_jobs] check_results(slab_results, "slab relaxation") top = repeat_to_target(slab_jobs[0].results.get_main_molecule()) bottom = repeat_to_target(slab_jobs[1].results.get_main_molecule()) # Packing is deterministic for the selected case. Every load/seed # replica starts from this same preoptimized interface. interface = pack_fluid(stack_slabs(bottom, top), seed=SEEDS[0]) relaxation = make_interface_relaxation_job(f"{CASE}_relax", interface) relaxation_result = relaxation.run() check_results([relaxation_result], "interface relaxation") equilibrations: list[tuple[float, int, AMSNVTJob]] = [] relaxed = relaxation.results.get_main_molecule() for load_nn in LOADS_NN: for seed in SEEDS: name = f"{CASE}_equil_load{load_label(load_nn)}_seed{seed}" equilibrations.append( (load_nn, seed, make_equilibration_job(name, relaxed, load_nn, seed)) ) equilibration_results = [job.run() for _, _, job in equilibrations] check_results(equilibration_results, "load equilibration") equilibration_gaps: dict[tuple[float, int], float] = {} for load_nn, seed, job in equilibrations: _, gaps = gap_trajectory(job) equilibration_gaps[load_nn, seed] = float(gaps[-1]) save_equilibration_diagnostics(job) print("case | load (nN) | seed | load-equilibrated gap (Angstrom)") for load_nn, seed, _ in equilibrations: print(f"{CASE} | {load_nn:g} | {seed} | {equilibration_gaps[load_nn, seed]:.4f}") nemd_jobs: list[tuple[float, int, AMSNVTJob]] = [] for load_nn, seed, equilibration in equilibrations: name = f"{CASE}_load{load_label(load_nn)}_seed{seed}" nemd_jobs.append((load_nn, seed, make_nemd_job(name, equilibration, load_nn, seed))) nemd_results = [job.run() for _, _, job in nemd_jobs] check_results(nemd_results, "NEMD") for _, _, job in nemd_jobs: save_production_diagnostics(job) print( "case | load (nN) | seed | equilibrated gap (Angstrom) | " "production gap (Angstrom) | mean shear (nN) | mean normal (nN) | " "mean blockwise friction ratio" ) for load_nn, seed, job in nemd_jobs: shear_nn, normal_nn, friction = analyze_nemd(job) gap_mean, gap_std = production_gap_statistics(job) print( f"{CASE} | {load_nn:g} | {seed} | {equilibration_gaps[load_nn, seed]:.4f} | " f"{gap_mean:.4f} +/- {gap_std:.4f} | {shear_nn:.5g} | " f"{normal_nn:.5g} | {friction:.5g}" ) finally: finish() return 0 if __name__ == "__main__": raise SystemExit(main())