.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/dynamics/13_lbfgs_variable_cell.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_dynamics_13_lbfgs_variable_cell.py: Variable-Cell L-BFGS Optimization (FCC Argon) ============================================== Joint relaxation of atomic coordinates and the simulation cell with the **L-BFGS** optimizer, using the PyTorch binding :func:`nvalchemiops.torch.lbfgs.lbfgs_step_coord_cell`. Compared with the FIRE2 variable-cell example (``11_fire2_variable_cell.py``), the L-BFGS path is a single call and takes the **stress** directly: - You pass ``stress``; the optimizer converts it to a cell driving force internally, so there is no separate ``stress_to_cell_force`` step - The same ``stress`` is used for the stress convergence criterion, which is compared against a real stress rather than against a packed cell force - There is no timestep or cell mass to tune How the coupling works ---------------------- Positions and cell are mapped into one packed coordinate vector, so the quasi-Newton recursion couples them with no special handling. The coordinates follow ASE's ``UnitCellFilter`` convention: with a reference cell ``H0`` fixed at the start and lattice vectors held as columns, the deformation gradient is ``Phi = H H0^-1``, atoms are stored as ``u = Phi^-1 r`` and the cell as ``kappa * Phi``. Two consequences are worth knowing: - The cell must be aligned with ``align_cell`` before the first step, exactly as ``fire2_step_coord_cell`` requires, and the six packed cell components are the same six FIRE2 packs. - The reference cell is captured **once**, by ``lbfgs_prepare_cell_state``. Rebuilding the chart mid-run invalidates every stored curvature pair, because they compare cell coordinates across steps. This is the one rule FIRE2 does not share: it keeps no history. - The optimizer buffers must be sized for ``num_atoms + 2 * num_systems`` degrees of freedom, because the cell contributes two packed entries per system. The preparation functions allocate and initialize both states, but the packed topology stays yours: this example builds it with the generic batch utilities, which is what makes ragged batches (systems with different atom counts) expressible. Apply your thresholds to the **Cartesian** forces and the stress, never to the packed norms, so they keep their physical meaning however far the cell deforms. The full set of cell rules -- alignment, the six-component convention, the fixed reference, topology, and ragged and empty systems -- is stated once in the :mod:`nvalchemiops.dynamics.optimizers.lbfgs` module documentation. This example follows it rather than restating it. .. GENERATED FROM PYTHON SOURCE LINES 68-99 .. code-block:: Python from __future__ import annotations import matplotlib.pyplot as plt import numpy as np import torch import warp as wp from _dynamics_utils import ( EPSILON_AR, SIGMA_AR, MDSystem, create_fcc_lattice, pressure_ev_per_a3_to_gpa, pressure_gpa_to_ev_per_a3, virial_to_stress, ) from nvalchemiops.dynamics.utils import align_cell, wrap_positions_to_cell from nvalchemiops.torch.lbfgs import ( lbfgs_prepare_cell_state, lbfgs_prepare_state, lbfgs_step_coord_cell, ) CUTOFF = 2.5 * SIGMA_AR # ~8.5 Å wp.init() device = "cuda:0" if wp.is_cuda_available() else "cpu" torch_device = torch.device(device) print(f"Using device: {device}") .. rst-class:: sphx-glr-script-out .. code-block:: none Warp UserWarning: Kernel cache artifacts from a previous Warp version were found in '/tmp/warp-cache'. These will be ignored. You can safely delete them. Using device: cuda:0 .. GENERATED FROM PYTHON SOURCE LINES 100-105 Create Initial System --------------------- An FCC argon lattice, deliberately expanded from its equilibrium spacing so there is something for the cell to relax. .. GENERATED FROM PYTHON SOURCE LINES 105-134 .. code-block:: Python n_cells = 3 # 3x3x3 = 108 atoms a_initial = 5.5 # Å (equilibrium is near 5.26 Å) positions_np, cell_np = create_fcc_lattice(n_cells, a_initial) num_atoms = len(positions_np) num_systems = 1 target_pressure_gpa = 0.01 target_pressure = pressure_gpa_to_ev_per_a3(target_pressure_gpa) print(f"System: {num_atoms} atoms in {n_cells}³ FCC lattice") print(f"Initial lattice constant: {a_initial:.3f} Å") print(f"Target external pressure: {target_pressure_gpa:.3f} GPa") positions = wp.array(positions_np, dtype=wp.vec3d, device=device) cell = wp.array(cell_np.reshape(1, 3, 3), dtype=wp.mat33d, device=device) md_system = MDSystem( positions=positions_np, cell=cell_np, epsilon=EPSILON_AR, sigma=SIGMA_AR, cutoff=CUTOFF, skin=0.5, switch_width=1.0, # smooth cutoff, so the energy is differentiable device=device, ) .. rst-class:: sphx-glr-script-out .. code-block:: none System: 108 atoms in 3³ FCC lattice Initial lattice constant: 5.500 Å Target external pressure: 0.010 GPa Initialized MD system with 108 atoms Cell: 16.50 x 16.50 x 16.50 Å Cutoff: 8.50 Å (+ 0.50 Å skin) LJ: ε = 0.0104 eV, σ = 3.40 Å Device: cuda:0, dtype: Units: x [Å], t [fs], E [eV], m [eV·fs²/Ų] (from amu), v [Å/fs] .. GENERATED FROM PYTHON SOURCE LINES 135-141 Align the Cell, Once -------------------- The optimizer keeps the cell in its aligned form, stopping it drifting into a rotation. Align *before* capturing the reference cell: re-aligning mid-run redefines the chart and invalidates the history. .. GENERATED FROM PYTHON SOURCE LINES 141-159 .. code-block:: Python transform = wp.empty(1, dtype=wp.mat33d, device=device) positions, cell = align_cell(positions, cell, transform=transform, device=device) wp.copy(md_system.wp_positions, positions) md_system.update_cell(cell) # Wrap once, before the first step. Wrapping *during* the run teleports atoms # across the boundary, and that discontinuity corrupts the curvature pairs. # FIRE2 tolerates per-step wrapping because it carries no history. wrap_positions_to_cell( positions=md_system.wp_positions, cells=md_system.wp_cell, cells_inv=md_system.wp_cell_inv, device=device, ) wp.synchronize() print(f"\nAligned cell:\n{cell.numpy()[0]}") .. rst-class:: sphx-glr-script-out .. code-block:: none Aligned cell: [[16.5 0. 0. ] [ 0. 16.5 0. ] [ 0. 0. 16.5]] .. GENERATED FROM PYTHON SOURCE LINES 160-169 Prepare the Two States ---------------------- Two objects, both yours: an :class:`~nvalchemiops.dynamics.optimizers.lbfgs.LBFGSState` as on the coordinate-only path, and an :class:`~nvalchemiops.dynamics.optimizers.lbfgs.LBFGSCellState` carrying the chart and its scratch. The preparation functions allocate, initialize and validate both; every field stays reachable by name. .. GENERATED FROM PYTHON SOURCE LINES 169-193 .. code-block:: Python # Work on the system's own arrays, so in-place writes are visible to the force # engine. Only the cell needs handing back: its inverse and the neighbor list # depend on it. positions_t = wp.to_torch(md_system.wp_positions) cell_t = wp.to_torch(cell).reshape(num_systems, 3, 3) batch_idx = torch.zeros(num_atoms, dtype=torch.int32, device=torch_device) dtype, i32 = torch.float64, torch.int32 history_size = 6 # The cell contributes two packed entries per system, so the optimizer state is # sized for num_atoms + 2 * num_systems degrees of freedom, not num_atoms. num_dofs = num_atoms + 2 * num_systems # Every per-system scalar follows the coordinate dtype, so this is an fp64 # state end to end; passing float32 would give an fp32 one. state = lbfgs_prepare_state( num_dofs, num_systems, dtype=dtype, history_size=history_size, device=torch_device, ) .. GENERATED FROM PYTHON SOURCE LINES 194-206 Prepare the Variable-Cell Chart ------------------------------- One call. Hand it the ordinary ``atom_ptr`` you already have and the aligned cells, and it derives the packed topology, captures the reference cell that defines the chart, and computes ``kappa`` -- the state comes back ready to step, with nothing left to repair. Deriving the topology costs no generality: ``atom_ptr`` already carries each system's atom count, so ragged batches work exactly as before. ``kappa`` scales the cell coordinate against the atomic ones; raise ``cell_force_scale`` to make the cell move less per step. .. GENERATED FROM PYTHON SOURCE LINES 206-219 .. code-block:: Python atom_ptr = torch.arange(num_systems + 1, dtype=i32, device=torch_device) * num_atoms cell_state = lbfgs_prepare_cell_state( atom_ptr, cell_t, cell_force_scale=1.0 / num_atoms, dtype=dtype, device=torch_device, ) print(f"\nOptimizer degrees of freedom: {num_atoms} atoms + {2 * num_systems} cell") .. rst-class:: sphx-glr-script-out .. code-block:: none Optimizer degrees of freedom: 108 atoms + 2 cell .. GENERATED FROM PYTHON SOURCE LINES 220-228 Optimization Loop ----------------- One force/stress evaluation per call. The optimizer owns neither tolerance and reports no terminal status, so the loop below applies both thresholds itself. Both are physical, and both are applied to the **Cartesian** forces and the stress rather than to the packed norms: ``force_tol`` is eV/Å on the largest per-atom force, ``stress_tol`` is a stress. .. GENERATED FROM PYTHON SOURCE LINES 228-313 .. code-block:: Python max_evals = 200 force_tol = 1e-4 # eV/Å stress_tol = pressure_gpa_to_ev_per_a3(0.03) # 0.03 GPa maxstep = 0.2 # Å energy_hist: list[float] = [] max_force_hist: list[float] = [] volume_hist: list[float] = [] pressure_hist: list[float] = [] print("\n" + "=" * 95) print("VARIABLE-CELL L-BFGS OPTIMIZATION (FCC argon)") print("=" * 95) print(f"{'eval':>6} {'PE (eV)':>14} {'max|F|':>11} {'|stress| GPa':>13} {'a (Å)':>9}") n_evals = 0 converged = False for step in range(max_evals): energies, forces, virial = md_system.compute_forces_virial() stress = virial_to_stress(virial, md_system.wp_cell, target_pressure, device) n_evals += 1 # Both criteria are the caller's, and both are physical: they are applied # to the *Cartesian* forces and the stress, never to packed norms, so they # keep their meaning however far the cell deforms. Tested before stepping. forces_t = wp.to_torch(forces) stress_t = wp.to_torch(stress).reshape(num_systems, 3, 3) fmax_now = float(forces_t.norm(dim=1).max()) smax_now = float(np.linalg.svd(stress.numpy()[0], compute_uv=False).max()) converged = fmax_now < force_tol and smax_now < stress_tol if step < 20 or step % 25 == 0 or converged: volume = float(np.linalg.det(cell_t.detach().cpu().numpy()[0])) stress_np = stress.numpy()[0] stress_gpa = pressure_ev_per_a3_to_gpa(0.5 * (stress_np + stress_np.T)) stress_residual = float(np.linalg.svd(stress_gpa, compute_uv=False).max()) pe = float(energies.numpy().sum()) energy_hist.append(pe) max_force_hist.append(fmax_now) volume_hist.append(volume) pressure_hist.append(stress_residual) print( f"{step:>6d} {pe:>14.6f} {fmax_now:>11.3e} {stress_residual:>13.3e} " f"{(volume / (n_cells**3)) ** (1 / 3):>9.4f}" ) if converged: break lbfgs_step_coord_cell( positions_t, cell_t, forces_t, stress_t, state, cell_state, batch_idx, maxstep=maxstep, ) # Positions were written in place; the cell needs handing back so its # inverse is recomputed and the neighbor list is rebuilt. md_system.update_cell(cell) # The loop tests before stepping, so on the converged path the last values # logged already describe the geometry and cell we are keeping. If the budget # ran out instead, the loop's final act was a *step*, and those values describe # the point before it -- while `cell_t` below is the post-step cell. Pairing # them would report a lattice constant and a force from two different # geometries. Re-evaluate once so everything below agrees. This is a reporting # evaluation, so it is deliberately not counted in `n_evals`. if not converged: energies, forces, virial = md_system.compute_forces_virial() stress = virial_to_stress(virial, md_system.wp_cell, target_pressure, device) fmax_now = float(wp.to_torch(forces).norm(dim=1).max()) stress_np = stress.numpy()[0] stress_gpa = pressure_ev_per_a3_to_gpa(0.5 * (stress_np + stress_np.T)) energy_hist.append(float(energies.numpy().sum())) max_force_hist.append(fmax_now) volume_hist.append(float(np.linalg.det(cell_t.detach().cpu().numpy()[0]))) pressure_hist.append(float(np.linalg.svd(stress_gpa, compute_uv=False).max())) .. rst-class:: sphx-glr-script-out .. code-block:: none =============================================================================================== VARIABLE-CELL L-BFGS OPTIMIZATION (FCC argon) =============================================================================================== eval PE (eV) max|F| |stress| GPa a (Å) 0 -8.410375 4.297e-16 2.434e-01 5.5000 1 -8.572777 7.297e-16 2.100e-01 5.4508 2 -8.703961 1.810e-15 1.660e-01 5.4015 3 -8.797196 3.380e-15 1.117e-01 5.3523 4 -8.848598 5.744e-15 5.293e-02 5.3031 5 -8.857636 8.593e-15 4.860e-03 5.2602 .. GENERATED FROM PYTHON SOURCE LINES 314-316 Result ------ .. GENERATED FROM PYTHON SOURCE LINES 316-330 .. code-block:: Python outcome = "CONVERGED" if converged else "ran out of evaluations" final_volume = float(np.linalg.det(cell_t.detach().cpu().numpy()[0])) final_a = (final_volume / (n_cells**3)) ** (1 / 3) print(f"\nFinished after {n_evals} evaluations: {outcome}") print(f" lattice constant: {a_initial:.4f} Å -> {final_a:.4f} Å") # All four of these describe the same, final geometry on both paths. print(f" final max|F| : {max_force_hist[-1]:.3e} eV/Å") print(f" final volume : {final_volume:.2f} ų") print(f" final |stress| : {pressure_hist[-1]:.4f} GPa") print(" (textbook FCC argon equilibrium is near 5.26 Å)") .. rst-class:: sphx-glr-script-out .. code-block:: none Finished after 6 evaluations: CONVERGED lattice constant: 5.5000 Å -> 5.2602 Å final max|F| : 8.593e-15 eV/Å final volume : 3929.91 ų final |stress| : 0.0049 GPa (textbook FCC argon equilibrium is near 5.26 Å) .. GENERATED FROM PYTHON SOURCE LINES 331-336 Plot convergence ---------------- Energy is plotted for interest, not as a convergence signal: with no Armijo test forcing it down it may rise on a step. The *force* is what converges. .. GENERATED FROM PYTHON SOURCE LINES 336-354 .. code-block:: Python points = np.arange(len(energy_hist)) fig, ax = plt.subplots(3, 1, figsize=(7.0, 7.5), sharex=True, constrained_layout=True) ax[0].plot(points, energy_hist, lw=1.5) ax[0].set_ylabel("Potential Energy (eV)") ax[0].set_title("Variable-Cell L-BFGS Convergence") ax[1].semilogy(points, max_force_hist, lw=1.5, label=r"max$|F|$") ax[1].axhline(force_tol, color="k", ls="--", lw=1.0, label="force tolerance") ax[1].set_ylabel(r"max$|F|$ (eV/$\AA$)") ax[1].legend(frameon=False, loc="best") ax[2].plot(points, [(v / n_cells**3) ** (1 / 3) for v in volume_hist], lw=1.5) ax[2].set_xlabel("Log point index") ax[2].set_ylabel("Lattice constant (Å)") plt.show() .. image-sg:: /examples/dynamics/images/sphx_glr_13_lbfgs_variable_cell_001.png :alt: Variable-Cell L-BFGS Convergence :srcset: /examples/dynamics/images/sphx_glr_13_lbfgs_variable_cell_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.390 seconds) .. _sphx_glr_download_examples_dynamics_13_lbfgs_variable_cell.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 13_lbfgs_variable_cell.ipynb <13_lbfgs_variable_cell.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 13_lbfgs_variable_cell.py <13_lbfgs_variable_cell.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 13_lbfgs_variable_cell.zip <13_lbfgs_variable_cell.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_