#!/usr/bin/env amspython from __future__ import annotations from pathlib import Path from typing import Iterable import warnings import matplotlib.pyplot as plt import numpy as np from scm.base import Units from scm.input_classes import AMS from scm.plams import AMSJob, Molecule, Settings, finish, init, packmol # ============================================================================= # User settings # ============================================================================= WORKDIR = Path("PolymerResin_workdir") STRUCTURE_DIR = Path(".") SYSTEM = { "polymer": STRUCTURE_DIR / "PI_30mer.xyz", "resin": STRUCTURE_DIR / "HDCPD_trimer.xyz", } N_PI_CHAINS = 20 N_RESIN_MOLECULES = 46 PACKING_DENSITY_G_CM3 = 0.05 PACKMOL_SEED = 2026 PACKED_STRUCTURE = Path("polymer_resin_initial.in") UFF_RELAX_MAX_ITERATIONS = 1000 PRESHRINK_STEPS = 20_000 PRESHRINK_TARGET_DENSITY_G_CM3 = 1.0 PRESHRINK_LAMBDA = 1.5 REAXFF_FORCE_FIELD = "dispersion/CHONSSi-lg.ff" REAXFF_RELAX_MAX_ITERATIONS = 1000 REAXFF_TIMESTEP_FS = 0.5 PRESSURE_PA = 101325.0 ANNEALING_STEPS = 25_000 ANNEALING_SAMPLING_FREQ = 500 ANNEALING_LOW_TEMPERATURE_K = 298.15 ANNEALING_HIGH_TEMPERATURE_K = 413.0 MAX_ANNEALING_ITERATIONS = 7 DENSITY_CONVERGENCE_PERCENT = 0.5 PRODUCTION_TEMPERATURE_K = 413.0 PRODUCTION_STEPS = 50_000 PRODUCTION_SAMPLING_FREQ = 100 RDF_PAIRS = (("resin", "resin"), ("resin", "polymer"), ("polymer", "polymer")) RDF_RMAX = 40.0 RDF_DR = 0.1 ANALYSIS_START_FRACTION = 0.5 DENSITY_MAP_BINS = 80 # ============================================================================= # General helpers # ============================================================================= KG_M3_TO_G_CM3 = 1.0 / 1000.0 AU_DENSITY_TO_G_CM3 = ( Units.conversion_factor("au", "kg") / Units.conversion_factor("bohr", "m") ** 3 * KG_M3_TO_G_CM3 ) def validate_settings(settings: Settings) -> Settings: AMS.from_settings(settings) return settings def molecule_density_g_cm3(molecule: Molecule) -> float: return molecule.get_density() * KG_M3_TO_G_CM3 def density_trace_g_cm3(job: AMSJob) -> np.ndarray: densities_au = job.results.get_history_property("Density", "MDHistory") return np.asarray(densities_au, dtype=float) * AU_DENSITY_TO_G_CM3 def time_trace_ps(job: AMSJob) -> np.ndarray: times_fs = job.results.get_history_property("Time", "MDHistory") return np.asarray(times_fs, dtype=float) / 1000.0 def last_third_average(values: Iterable[float]) -> float: values = np.asarray(list(values), dtype=float) if len(values) == 0: raise ValueError("Cannot average an empty trace") return float(np.mean(values[-max(1, len(values) // 3) :])) def guess_bonds(molecule: Molecule) -> Molecule: if hasattr(molecule, "guess_bonds"): molecule.guess_bonds() return molecule def run_or_load(molecule: Molecule, settings: Settings, name: str) -> tuple[AMSJob, Molecule]: job_path = WORKDIR / name if job_path.exists(): print(f"Loading {name}", flush=True) job = AMSJob.load_external(job_path) else: print(f"Running {name}", flush=True) job = AMSJob(molecule=molecule, settings=settings, name=name) job.run() output_molecule = job.results.get_main_molecule() print(f"{name}: density {molecule_density_g_cm3(output_molecule):.3f} g/cm^3", flush=True) return job, output_molecule # ============================================================================= # Structure setup # ============================================================================= def build_packed_blend() -> tuple[Molecule, Molecule, Molecule]: pi_chain = guess_bonds(Molecule(SYSTEM["polymer"])) resin = guess_bonds(Molecule(SYSTEM["resin"])) blend, _ = packmol( molecules=[pi_chain, resin], n_molecules=[N_PI_CHAINS, N_RESIN_MOLECULES], density=PACKING_DENSITY_G_CM3, region_names=["polymer", "resin"], tolerance=2.0, seed=PACKMOL_SEED, return_details=True, ) blend.write(PACKED_STRUCTURE) print(f"Packed structure: {PACKED_STRUCTURE}", flush=True) print(f"Atoms: {len(blend)}", flush=True) print(f"Initial density: {molecule_density_g_cm3(blend):.3f} g/cm^3", flush=True) return blend, pi_chain, resin # ============================================================================= # AMS settings builders # ============================================================================= def uff_relax_settings() -> Settings: settings = Settings() settings.input.ams.task = "GeometryOptimization" settings.input.ams.GeometryOptimization.MaxIterations = UFF_RELAX_MAX_ITERATIONS settings.input.ams.GeometryOptimization.PretendConverged = "Yes" settings.input.ForceField.Type = "UFF" return validate_settings(settings) def uff_shrink_settings(target_lattice: tuple[float, float, float]) -> Settings: a, b, c = target_lattice settings = Settings() settings.input.ams.task = "MolecularDynamics" settings.input.ams.MolecularDynamics.NSteps = PRESHRINK_STEPS settings.input.ams.MolecularDynamics.TimeStep = 0.25 settings.input.ams.MolecularDynamics.Trajectory.SamplingFreq = 1000 settings.input.ams.MolecularDynamics.InitialVelocities.Temperature = 5.0 settings.input.ams.MolecularDynamics.Thermostat.Type = "Berendsen" settings.input.ams.MolecularDynamics.Thermostat.Temperature = 5.0 settings.input.ams.MolecularDynamics.Thermostat.Tau = 10.0 settings.input.ams.MolecularDynamics.Deformation.TargetLattice._1 = f"{a:.3f} 0 0" settings.input.ams.MolecularDynamics.Deformation.TargetLattice._2 = f"0 {b:.3f} 0" settings.input.ams.MolecularDynamics.Deformation.TargetLattice._3 = f"0 0 {c:.3f}" settings.input.ForceField.Type = "UFF" return validate_settings(settings) def reaxff_relax_settings() -> Settings: settings = Settings() settings.input.ams.task = "GeometryOptimization" settings.input.ams.GeometryOptimization.MaxIterations = REAXFF_RELAX_MAX_ITERATIONS settings.input.ams.GeometryOptimization.PretendConverged = "Yes" settings.input.ReaxFF.ForceField = REAXFF_FORCE_FIELD return validate_settings(settings) def reaxff_annealing_settings() -> Settings: heating_steps = ANNEALING_STEPS // 2 cooling_steps = ANNEALING_STEPS - heating_steps settings = Settings() settings.input.ams.task = "MolecularDynamics" settings.input.ams.MolecularDynamics.NSteps = ANNEALING_STEPS settings.input.ams.MolecularDynamics.TimeStep = REAXFF_TIMESTEP_FS settings.input.ams.MolecularDynamics.Trajectory.SamplingFreq = ANNEALING_SAMPLING_FREQ settings.input.ams.MolecularDynamics.InitialVelocities.Temperature = ANNEALING_LOW_TEMPERATURE_K settings.input.ams.MolecularDynamics.Thermostat.Type = "Berendsen" settings.input.ams.MolecularDynamics.Thermostat.Temperature = ( f"{ANNEALING_LOW_TEMPERATURE_K} " f"{ANNEALING_HIGH_TEMPERATURE_K} " f"{ANNEALING_LOW_TEMPERATURE_K}" ) settings.input.ams.MolecularDynamics.Thermostat.Duration = f"{heating_steps} {cooling_steps}" settings.input.ams.MolecularDynamics.Thermostat.Tau = 100.0 settings.input.ams.MolecularDynamics.Barostat.Type = "Berendsen" settings.input.ams.MolecularDynamics.Barostat.Pressure = PRESSURE_PA settings.input.ams.MolecularDynamics.Barostat.Tau = 500.0 settings.input.ReaxFF.ForceField = REAXFF_FORCE_FIELD return validate_settings(settings) def reaxff_production_settings() -> Settings: settings = Settings() settings.input.ams.task = "MolecularDynamics" settings.input.ams.MolecularDynamics.NSteps = PRODUCTION_STEPS settings.input.ams.MolecularDynamics.TimeStep = REAXFF_TIMESTEP_FS settings.input.ams.MolecularDynamics.Trajectory.SamplingFreq = PRODUCTION_SAMPLING_FREQ settings.input.ams.MolecularDynamics.InitialVelocities.Temperature = PRODUCTION_TEMPERATURE_K settings.input.ams.MolecularDynamics.Thermostat.Type = "Berendsen" settings.input.ams.MolecularDynamics.Thermostat.Temperature = PRODUCTION_TEMPERATURE_K settings.input.ams.MolecularDynamics.Thermostat.Tau = 100.0 settings.input.ReaxFF.ForceField = REAXFF_FORCE_FIELD return validate_settings(settings) def preshrink_lattice(molecule: Molecule) -> tuple[float, float, float, float, float]: current_density = molecule_density_g_cm3(molecule) current_volume = molecule.unit_cell_volume() target_volume = current_density * current_volume / PRESHRINK_TARGET_DENSITY_G_CM3 new_a = (target_volume * PRESHRINK_LAMBDA) ** (1.0 / 3.0) new_b = new_a new_c = new_a / PRESHRINK_LAMBDA return current_density, target_volume, new_a, new_b, new_c # ============================================================================= # Analysis helpers # ============================================================================= def atom_region(atom) -> str | None: properties = atom.properties return properties.get("region") if hasattr(properties, "get") else getattr(properties, "region", None) def region_molecules(molecule: Molecule, region: str, atoms_per_molecule: int) -> list[list[int]]: atom_indices = [i for i, atom in enumerate(molecule, start=1) if atom_region(atom) == region] if not atom_indices: raise ValueError(f"No atoms found for region '{region}'") if len(atom_indices) % atoms_per_molecule: raise ValueError( f"Region '{region}' has {len(atom_indices)} atoms, which is not divisible by " f"{atoms_per_molecule} atoms per molecule" ) return [ atom_indices[i : i + atoms_per_molecule] for i in range(0, len(atom_indices), atoms_per_molecule) ] def region_indices(molecule: Molecule, region: str) -> list[int]: return [i for i, atom in enumerate(molecule, start=1) if atom_region(atom) == region] def lattice_matrix(molecule: Molecule) -> np.ndarray: return np.asarray(molecule.lattice, dtype=float) def minimum_image_vectors(vectors: np.ndarray, lattice: np.ndarray) -> np.ndarray: fractional = vectors @ np.linalg.inv(lattice) return vectors - np.round(fractional) @ lattice def center_of_mass(molecule: Molecule, atom_indices: list[int]) -> np.ndarray: molecule_atoms = list(molecule) atoms = [molecule_atoms[i - 1] for i in atom_indices] masses = np.asarray([atom.mass for atom in atoms], dtype=float) coords = np.asarray([atom.coords for atom in atoms], dtype=float) lattice = lattice_matrix(molecule) unwrapped = coords[0] + minimum_image_vectors(coords - coords[0], lattice) com = np.average(unwrapped, axis=0, weights=masses) return com - np.floor(com @ np.linalg.inv(lattice)) @ lattice def pair_distances( points_a: np.ndarray, points_b: np.ndarray, lattice: np.ndarray, same_region: bool, ) -> np.ndarray: if same_region: vectors = [ points_a[j] - points_a[i] for i in range(len(points_a)) for j in range(i + 1, len(points_a)) ] else: vectors = [point_b - point_a for point_a in points_a for point_b in points_b] if not vectors: return np.empty(0, dtype=float) return np.linalg.norm(minimum_image_vectors(np.asarray(vectors), lattice), axis=1) def calculate_rdfs( job: AMSJob, pairs: tuple[tuple[str, str], ...], atoms_per_molecule: dict[str, int], rmax: float, dr: float, start_fraction: float, ) -> tuple[np.ndarray, dict[str, np.ndarray]]: n_frames = job.results.readrkf("History", "nEntries") start_frame = max(1, min(n_frames, int(n_frames * start_fraction) + 1)) first_molecule = job.results.get_history_molecule(start_frame) max_r = min(np.linalg.norm(vector) for vector in lattice_matrix(first_molecule)) / 2.0 if rmax > max_r: warnings.warn(f"Capping rmax from {rmax:.3f} to {max_r:.3f} A", RuntimeWarning) rmax = max_r bins = np.arange(0.0, rmax + dr, dr) histograms = {f"{a}-{b}": np.zeros(len(bins) - 1, dtype=float) for a, b in pairs} volumes: list[float] = [] molecule_counts: dict[str, int] | None = None print(f"Using RDF frames {start_frame}-{n_frames} of {n_frames}", flush=True) for frame in range(start_frame, n_frames + 1): molecule = job.results.get_history_molecule(frame) centers = { region: np.asarray( [ center_of_mass(molecule, group) for group in region_molecules(molecule, region, atoms_per_molecule[region]) ] ) for region in atoms_per_molecule } counts = {region: len(points) for region, points in centers.items()} if molecule_counts is None: molecule_counts = counts print("Molecule counts:", molecule_counts, flush=True) elif counts != molecule_counts: raise ValueError(f"Frame {frame} has a different number of molecules") lattice = lattice_matrix(molecule) for region_a, region_b in pairs: distances = pair_distances( centers[region_a], centers[region_b], lattice, region_a == region_b, ) histograms[f"{region_a}-{region_b}"] += np.histogram(distances, bins=bins)[0] volumes.append(molecule.unit_cell_volume()) if molecule_counts is None or not volumes: raise ValueError("No trajectory frames were used") radii = 0.5 * (bins[:-1] + bins[1:]) shell_volumes = 4.0 * np.pi / 3.0 * (bins[1:] ** 3 - bins[:-1] ** 3) volume = float(np.mean(volumes)) rdf_data: dict[str, np.ndarray] = {} for region_a, region_b in pairs: n_a = molecule_counts[region_a] n_b = molecule_counts[region_b] pair_count = n_a * (n_a - 1) / 2.0 if region_a == region_b else n_a * n_b normalization = len(volumes) * shell_volumes * pair_count / volume rdf_data[f"{region_a}-{region_b}"] = histograms[f"{region_a}-{region_b}"] / normalization return radii, rdf_data def density_map_xy( job: AMSJob, region: str = "resin", bins: int = DENSITY_MAP_BINS, start_fraction: float = ANALYSIS_START_FRACTION, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: n_frames = job.results.readrkf("History", "nEntries") start_frame = max(1, min(n_frames, int(n_frames * start_fraction) + 1)) final_molecule = job.results.get_main_molecule() lx = final_molecule.lattice[0][0] ly = final_molecule.lattice[1][1] histogram = np.zeros((bins, bins), dtype=float) print(f"Using density-map frames {start_frame}-{n_frames} of {n_frames}", flush=True) for used_frames, frame in enumerate(range(start_frame, n_frames + 1), start=1): molecule = job.results.get_history_molecule(frame) atoms = list(molecule) indices = region_indices(molecule, region) if not indices: raise ValueError(f"No atoms found for region '{region}'") coords = np.asarray([atoms[i - 1].coords[:2] for i in indices], dtype=float) coords[:, 0] %= lx coords[:, 1] %= ly histogram += np.histogram2d( coords[:, 0], coords[:, 1], bins=bins, range=[[0.0, lx], [0.0, ly]], )[0] dx_nm = (lx / bins) / 10.0 dy_nm = (ly / bins) / 10.0 density = histogram / used_frames / (dx_nm * dy_nm) x_edges = np.linspace(0.0, lx / 10.0, bins + 1) y_edges = np.linspace(0.0, ly / 10.0, bins + 1) return x_edges, y_edges, density.T def write_annealing_density_plot(annealing_jobs: list[AMSJob]) -> None: fig, ax = plt.subplots(figsize=(6.0, 3.2)) for job in annealing_jobs: ax.plot(time_trace_ps(job), density_trace_g_cm3(job), label=job.name) ax.set_xlabel("Time (ps)") ax.set_ylabel("Density (g/cm^3)") ax.legend(frameon=False) fig.tight_layout() fig.savefig("density_per_annealing_step.png") plt.close(fig) def write_rdf_outputs(production_job: AMSJob, atoms_per_molecule: dict[str, int]) -> None: radii, rdf_data = calculate_rdfs( production_job, RDF_PAIRS, atoms_per_molecule, RDF_RMAX, RDF_DR, ANALYSIS_START_FRACTION, ) fig, ax = plt.subplots(figsize=(5.2, 3.4)) for label, values in rdf_data.items(): ax.plot(radii, values, label=label) ax.set_xlabel("r (A)") ax.set_ylabel("g(r)") ax.legend(frameon=False) fig.tight_layout() fig.savefig("RDF_com_regions.png") plt.close(fig) rdf_table = np.column_stack([radii, *rdf_data.values()]) np.savetxt("RDF_com_regions.txt", rdf_table, header="r_A " + " ".join(rdf_data)) def write_density_map_outputs(production_job: AMSJob) -> None: x_edges, y_edges, resin_density = density_map_xy(production_job) np.savetxt("DENSMAP_resin.txt", resin_density, header="resin atoms / nm^2 / frame") fig, ax = plt.subplots(figsize=(5.0, 4.0)) mesh = ax.pcolormesh(x_edges, y_edges, resin_density, shading="auto") fig.colorbar(mesh, ax=ax, label="resin atoms / nm2 / frame") ax.set_xlabel("x (nm)") ax.set_ylabel("y (nm)") ax.set_title("Resin density map") fig.tight_layout() fig.savefig("DENSMAP_resin.png") plt.close(fig) # ============================================================================= # Main workflow # ============================================================================= def main() -> None: init(folder=str(WORKDIR)) try: blend, pi_chain, resin = build_packed_blend() _, relaxed_blend = run_or_load(blend, uff_relax_settings(), "UFF_InitialRelax") current_density, target_volume, new_a, new_b, new_c = preshrink_lattice(relaxed_blend) print(f"Current density: {current_density:.3f} g/cm^3", flush=True) print(f"Target density: {PRESHRINK_TARGET_DENSITY_G_CM3:.3f} g/cm^3", flush=True) print(f"Target volume: {target_volume:.1f} A^3", flush=True) print(f"Target lattice: a={new_a:.3f}, b={new_b:.3f}, c={new_c:.3f} A", flush=True) _, shrunken_blend = run_or_load( relaxed_blend, uff_shrink_settings((new_a, new_b, new_c)), "UFF_Shrink", ) _, reaxff_relaxed_blend = run_or_load( shrunken_blend, reaxff_relax_settings(), "ReaxFF_InitialRelax", ) annealing_settings = reaxff_annealing_settings() annealing_jobs: list[AMSJob] = [] density_summary: list[dict[str, float]] = [] current_molecule = reaxff_relaxed_blend previous_last_third_density: float | None = None for iteration in range(1, MAX_ANNEALING_ITERATIONS + 1): job_name = f"ReaxFF_Anneal_{iteration}" initial_density = molecule_density_g_cm3(current_molecule) job, current_molecule = run_or_load(current_molecule, annealing_settings, job_name) annealing_jobs.append(job) densities = density_trace_g_cm3(job) last_third_density = last_third_average(densities) if previous_last_third_density is None: density_change = np.nan else: density_change = abs(last_third_density / previous_last_third_density - 1.0) * 100.0 density_summary.append( { "annealing_step": float(iteration), "initial_density_g_cm3": initial_density, "final_density_g_cm3": float(densities[-1]), "last_third_density_g_cm3": last_third_density, "change_percent": float(density_change), } ) print(density_summary[-1], flush=True) if previous_last_third_density is not None and density_change < DENSITY_CONVERGENCE_PERCENT: break previous_last_third_density = last_third_density write_annealing_density_plot(annealing_jobs) production_job, production_blend = run_or_load( current_molecule, reaxff_production_settings(), "ReaxFF_Production", ) print( f"Final production density: {molecule_density_g_cm3(production_blend):.3f} g/cm^3", flush=True, ) atoms_per_molecule = { "polymer": len(pi_chain), "resin": len(resin), } write_rdf_outputs(production_job, atoms_per_molecule) write_density_map_outputs(production_job) finally: finish() if __name__ == "__main__": main()