.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/electrostatics/03_pme_example.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_electrostatics_03_pme_example.py: Particle Mesh Ewald (PME) for Long-Range Electrostatics ======================================================= This example demonstrates how to compute long-range electrostatic interactions using the Particle Mesh Ewald (PME) method. PME achieves O(N log N) scaling by using FFT-based mesh interpolation for the reciprocal-space contribution. In this example you will learn: - How to set up and run PME with automatic parameter estimation - Using neighbor list and neighbor matrix formats - Understanding convergence with accuracy-based parameter estimation - Effect of the splitting parameter alpha and accuracy - Batch evaluation for multiple systems - Comparison between PME and standard Ewald summation - Computing charge gradients for ML potential training PME accelerates the reciprocal-space sum using B-spline interpolation: 1. Spread charges to mesh using B-splines 2. FFT to reciprocal space 3. Multiply by Green's function 4. Inverse FFT to get potentials 5. Interpolate forces back to atom positions .. important:: This script is intended as an API demonstration. Do not use this script for performance benchmarking; refer to the `benchmarks` folder instead. .. GENERATED FROM PYTHON SOURCE LINES 48-52 Setup and Imports ----------------- First, we import the necessary modules. The PME API provides unified functions that handle both single-system and batched calculations. .. GENERATED FROM PYTHON SOURCE LINES 52-71 .. code-block:: Python from __future__ import annotations import time import warnings import numpy as np import torch from nvalchemiops.torch.interactions.electrostatics import ( estimate_ewald_parameters, estimate_pme_parameters, ewald_real_space, ewald_summation, particle_mesh_ewald, pme_reciprocal_space, ) from nvalchemiops.torch.neighbors import neighbor_list as neighbor_list_fn .. GENERATED FROM PYTHON SOURCE LINES 72-74 Configure Device ---------------- .. GENERATED FROM PYTHON SOURCE LINES 74-94 .. code-block:: Python if torch.cuda.is_available(): device = torch.device("cuda:0") print(f"Using CUDA device: {torch.cuda.get_device_name(0)}") else: device = torch.device("cpu") print("Using CPU") def _legacy_direct_output_call(function, *args, **kwargs): """Call a deprecated direct-output path used for explicit migration checks.""" with warnings.catch_warnings(): warnings.filterwarnings( "ignore", category=DeprecationWarning, message="The direct-output flags.*", ) return function(*args, **kwargs) .. rst-class:: sphx-glr-script-out .. code-block:: none Using CUDA device: NVIDIA H100 PCIe .. GENERATED FROM PYTHON SOURCE LINES 95-98 Create a NaCl Crystal System ---------------------------- We define a helper function to create NaCl rock salt crystal supercells. .. GENERATED FROM PYTHON SOURCE LINES 98-137 .. code-block:: Python def create_nacl_system(n_cells: int = 2, lattice_constant: float = 5.64): """Create a NaCl crystal supercell. Parameters ---------- n_cells : int Number of unit cells in each direction. lattice_constant : float NaCl lattice constant in Angstroms. Returns ------- positions, charges, cell, pbc : torch.Tensor System tensors. """ base_positions = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]) base_charges = np.array([1.0, -1.0]) positions = [] charges = [] for i in range(n_cells): for j in range(n_cells): for k in range(n_cells): offset = np.array([i, j, k]) for pos, charge in zip(base_positions, base_charges): positions.append((pos + offset) * lattice_constant) charges.append(charge) positions = torch.tensor(positions, dtype=torch.float64, device=device) charges = torch.tensor(charges, dtype=torch.float64, device=device) cell = torch.eye(3, dtype=torch.float64, device=device) * lattice_constant * n_cells cell = cell.unsqueeze(0) pbc = torch.tensor([[True, True, True]], dtype=torch.bool, device=device) return positions, charges, cell, pbc .. GENERATED FROM PYTHON SOURCE LINES 138-143 Basic Usage with Automatic Parameters ------------------------------------- The simplest way to use PME is with automatic parameter estimation. Given an accuracy tolerance, the API estimates optimal alpha, mesh dimensions, and real-space cutoff. .. GENERATED FROM PYTHON SOURCE LINES 143-151 .. code-block:: Python # Create a NaCl crystal (3×3×3 unit cells = 54 atoms) positions, charges, cell, pbc = create_nacl_system(n_cells=3) print(f"System: {len(positions)} atoms NaCl crystal") print(f"Cell size: {cell[0, 0, 0]:.2f} Å") print(f"Total charge: {charges.sum().item():.1f} (constructed neutral system)") .. rst-class:: sphx-glr-script-out .. code-block:: none /tmp/tmp_3o8x9b4/c1e23460859a784e1d78043bcd1c8af0d1095fa2/examples/electrostatics/03_pme_example.py:128: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /pytorch/torch/csrc/utils/tensor_new.cpp:253.) positions = torch.tensor(positions, dtype=torch.float64, device=device) System: 54 atoms NaCl crystal Cell size: 16.92 Å Total charge: 0.0 (constructed neutral system) .. GENERATED FROM PYTHON SOURCE LINES 152-153 Estimate optimal PME parameters: .. GENERATED FROM PYTHON SOURCE LINES 153-167 .. code-block:: Python params = estimate_pme_parameters(positions.cpu(), cell.cpu(), accuracy=1e-6) print("\nEstimated parameters (accuracy=1e-6):") print(f" alpha = {params.alpha.item():.4f}") print(f" mesh_dimensions = {params.mesh_dimensions}") spacing = ( params.mesh_spacing[0] if params.mesh_spacing.dim() == 2 else params.mesh_spacing ) print( f" mesh_spacing = ({spacing[0].item():.2f}, {spacing[1].item():.2f}, {spacing[2].item():.2f}) Å" ) print(f" real_space_cutoff = {params.real_space_cutoff.item():.2f} Å") .. rst-class:: sphx-glr-script-out .. code-block:: none Estimated parameters (accuracy=1e-6): alpha = 0.2037 mesh_dimensions = (64, 64, 64) mesh_spacing = (0.26, 0.26, 0.26) Å real_space_cutoff = 18.25 Å .. GENERATED FROM PYTHON SOURCE LINES 168-169 Build neighbor list and run PME: .. GENERATED FROM PYTHON SOURCE LINES 169-199 .. code-block:: Python neighbor_list, neighbor_ptr, neighbor_shifts = neighbor_list_fn( positions, params.real_space_cutoff.item(), cell=cell, pbc=pbc, return_neighbor_list=True, ) t0 = time.time() positions_grad = positions.detach().clone().requires_grad_(True) energies = particle_mesh_ewald( positions=positions_grad, charges=charges, cell=cell, neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr, neighbor_shifts=neighbor_shifts, accuracy=1e-6, # Parameters estimated automatically ) forces = -torch.autograd.grad(energies.sum(), positions_grad)[0] t1 = time.time() total_energy = energies.sum().item() print("\nPME Results:") print(f" Total energy: {total_energy:.6f}") print(f" Energy per atom: {total_energy / len(positions):.6f}") print(f" Max force magnitude: {torch.norm(forces, dim=1).max().item():.6f}") print(f" Time: {(t1 - t0) * 1000:.2f} ms") .. rst-class:: sphx-glr-script-out .. code-block:: none PME Results: Total energy: -9.743761 Energy per atom: -0.180440 Max force magnitude: 0.000000 Time: 3271.68 ms .. GENERATED FROM PYTHON SOURCE LINES 200-204 Neighbor List vs Neighbor Matrix Format --------------------------------------- PME supports both neighbor formats, producing identical results. We use the estimated parameters from above for consistency. .. GENERATED FROM PYTHON SOURCE LINES 204-224 .. code-block:: Python # Build both formats using the estimated real-space cutoff neighbor_list, neighbor_ptr, neighbor_shifts = neighbor_list_fn( positions, params.real_space_cutoff.item(), cell=cell, pbc=pbc, return_neighbor_list=True, ) neighbor_matrix, _, neighbor_matrix_shifts = neighbor_list_fn( positions, params.real_space_cutoff.item(), cell=cell, pbc=pbc, return_neighbor_list=False, ) print("\nNeighbor format comparison (accuracy=1e-6):") print(f" Using alpha={params.alpha.item():.4f}, mesh_dims={params.mesh_dimensions}") .. rst-class:: sphx-glr-script-out .. code-block:: none Neighbor format comparison (accuracy=1e-6): Using alpha=0.2037, mesh_dims=(64, 64, 64) .. GENERATED FROM PYTHON SOURCE LINES 225-226 Using neighbor list format: .. GENERATED FROM PYTHON SOURCE LINES 226-243 .. code-block:: Python t0 = time.time() energies_list, forces_list = _legacy_direct_output_call( particle_mesh_ewald, positions=positions, charges=charges, cell=cell, neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr, neighbor_shifts=neighbor_shifts, compute_forces=True, accuracy=1e-6, # Parameters estimated automatically ) t_list = (time.time() - t0) * 1000 print(f" List format: E={energies_list.sum().item():.6f}, time={t_list:.2f} ms") .. rst-class:: sphx-glr-script-out .. code-block:: none List format: E=-9.743761, time=4.41 ms .. GENERATED FROM PYTHON SOURCE LINES 244-245 Using neighbor matrix format: .. GENERATED FROM PYTHON SOURCE LINES 245-266 .. code-block:: Python t0 = time.time() energies_matrix, forces_matrix = _legacy_direct_output_call( particle_mesh_ewald, positions=positions, charges=charges, cell=cell, neighbor_matrix=neighbor_matrix, neighbor_matrix_shifts=neighbor_matrix_shifts, compute_forces=True, accuracy=1e-6, # Same accuracy for comparison ) t_matrix = (time.time() - t0) * 1000 print(f" Matrix format: E={energies_matrix.sum().item():.6f}, time={t_matrix:.2f} ms") energy_diff = abs(energies_list.sum().item() - energies_matrix.sum().item()) force_diff = (forces_list - forces_matrix).abs().max().item() print(f"\nEnergy difference: {energy_diff:.2e}") print(f"Max force difference: {force_diff:.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none Matrix format: E=-9.743761, time=1202.81 ms Energy difference: 0.00e+00 Max force difference: 2.88e-17 .. GENERATED FROM PYTHON SOURCE LINES 267-272 Convergence with Accuracy Parameter ----------------------------------- The PME accuracy depends on the accuracy parameter, which controls both the mesh resolution and alpha parameter. The parameter estimation uses optimal formulas to balance computational cost between real and reciprocal space. .. GENERATED FROM PYTHON SOURCE LINES 272-280 .. code-block:: Python accuracies = [1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7] results_acc = [] print("\nConvergence with Accuracy Target:") print(" Accuracy | alpha | mesh_dims | r_cutoff | Energy | Time") print(" " + "-" * 75) .. rst-class:: sphx-glr-script-out .. code-block:: none Convergence with Accuracy Target: Accuracy | alpha | mesh_dims | r_cutoff | Energy | Time --------------------------------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 281-282 Run PME with different accuracy targets: .. GENERATED FROM PYTHON SOURCE LINES 282-315 .. code-block:: Python for acc in accuracies: # Estimate optimal parameters for this accuracy params_acc = estimate_pme_parameters(positions, cell, accuracy=acc) # Build neighbor list with appropriate cutoff nl_acc, nptr_acc, ns_acc = neighbor_list_fn( positions, params_acc.real_space_cutoff.item(), cell=cell, pbc=pbc, return_neighbor_list=True, ) t0 = time.time() energies_acc = particle_mesh_ewald( positions=positions, charges=charges, cell=cell, neighbor_list=nl_acc, neighbor_ptr=nptr_acc, neighbor_shifts=ns_acc, accuracy=acc, ) t_elapsed = (time.time() - t0) * 1000 total_e = energies_acc.sum().item() results_acc.append((acc, params_acc, total_e, t_elapsed)) print( f" {acc:.0e} | {params_acc.alpha.item():.3f} | {str(params_acc.mesh_dimensions):12s} |" f" {params_acc.real_space_cutoff.item():5.2f} | {total_e:10.6f} | {t_elapsed:.2f} ms" ) .. rst-class:: sphx-glr-script-out .. code-block:: none 1e-02 | 0.204 | (8, 8, 8) | 10.54 | -9.755594 | 9852.69 ms 1e-03 | 0.204 | (16, 16, 16) | 12.91 | -9.745599 | 62.20 ms 1e-04 | 0.204 | (16, 16, 16) | 14.90 | -9.743791 | 3.65 ms 1e-05 | 0.204 | (32, 32, 32) | 16.66 | -9.743703 | 108.42 ms 1e-06 | 0.204 | (64, 64, 64) | 18.25 | -9.743761 | 3.58 ms 1e-07 | 0.204 | (64, 64, 64) | 19.71 | -9.743762 | 3.40 ms .. GENERATED FROM PYTHON SOURCE LINES 316-317 Show convergence relative to highest accuracy: .. GENERATED FROM PYTHON SOURCE LINES 317-324 .. code-block:: Python ref_energy_acc = results_acc[-1][2] print("\nRelative error from reference (accuracy=1e-7):") for acc, _, e, _ in results_acc[:-1]: rel_err = abs((e - ref_energy_acc) / ref_energy_acc) print(f" accuracy={acc:.0e}: {rel_err:.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none Relative error from reference (accuracy=1e-7): accuracy=1e-02: 1.21e-03 accuracy=1e-03: 1.89e-04 accuracy=1e-04: 2.98e-06 accuracy=1e-05: 5.99e-06 accuracy=1e-06: 8.65e-08 .. GENERATED FROM PYTHON SOURCE LINES 325-329 Effect of Accuracy on PME Parameters ------------------------------------ The accuracy parameter controls how alpha and mesh dimensions are chosen. Lower accuracy targets require larger meshes and cutoffs. .. GENERATED FROM PYTHON SOURCE LINES 329-337 .. code-block:: Python positions_2, charges_2, cell_2, pbc_2 = create_nacl_system(n_cells=2) accuracies = [1e-2, 1e-3, 1e-4, 1e-5, 1e-6] print("\nAccuracy Effect on PME Parameters:") print(" Accuracy | alpha | mesh_dims | real_cutoff | Total Energy") print(" " + "-" * 65) .. rst-class:: sphx-glr-script-out .. code-block:: none Accuracy Effect on PME Parameters: Accuracy | alpha | mesh_dims | real_cutoff | Total Energy ----------------------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 338-339 Sweep through accuracy values: .. GENERATED FROM PYTHON SOURCE LINES 339-369 .. code-block:: Python for accuracy in accuracies: params_acc = estimate_pme_parameters( positions_2.cpu(), cell_2.cpu(), accuracy=accuracy ) nl_acc, nptr_acc, ns_acc = neighbor_list_fn( positions_2, params_acc.real_space_cutoff.item(), cell=cell_2, pbc=pbc_2, return_neighbor_list=True, ) energies_acc = particle_mesh_ewald( positions=positions_2, charges=charges_2, cell=cell_2, alpha=params_acc.alpha.item(), mesh_dimensions=tuple(params_acc.mesh_dimensions), neighbor_list=nl_acc, neighbor_ptr=nptr_acc, neighbor_shifts=ns_acc, ) print( f" {accuracy:.0e} | {params_acc.alpha.item():.2f} | {str(params_acc.mesh_dimensions):9s}" f" | {params_acc.real_space_cutoff.item():5.2f} | {energies_acc.sum().item():.6f}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none 1e-02 | 0.25 | (8, 8, 8) | 8.60 | -2.869979 1e-03 | 0.25 | (8, 8, 8) | 10.54 | -2.886155 1e-04 | 0.25 | (16, 16, 16) | 12.17 | -2.886909 1e-05 | 0.25 | (32, 32, 32) | 13.60 | -2.887047 1e-06 | 0.25 | (32, 32, 32) | 14.90 | -2.887036 .. GENERATED FROM PYTHON SOURCE LINES 370-373 Accessing Real-Space and Reciprocal-Space Components ---------------------------------------------------- You can compute the components separately if needed. .. GENERATED FROM PYTHON SOURCE LINES 373-386 .. code-block:: Python params_comp = estimate_pme_parameters(positions, cell, accuracy=1e-4) nl_comp, nptr_comp, ns_comp = neighbor_list_fn( positions, params_comp.real_space_cutoff.item(), cell=cell, pbc=pbc, return_neighbor_list=True, ) print("\nEnergy Components:") .. rst-class:: sphx-glr-script-out .. code-block:: none Energy Components: .. GENERATED FROM PYTHON SOURCE LINES 387-388 Real-space component (uses same kernel as Ewald): .. GENERATED FROM PYTHON SOURCE LINES 388-401 .. code-block:: Python real_energy = ewald_real_space( positions=positions, charges=charges, cell=cell, alpha=params_comp.alpha, neighbor_list=nl_comp, neighbor_ptr=nptr_comp, neighbor_shifts=ns_comp, ) print(f" Real-space: {real_energy.sum().item():.6f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Real-space: -3.549335 .. GENERATED FROM PYTHON SOURCE LINES 402-403 PME reciprocal-space component (FFT-based): .. GENERATED FROM PYTHON SOURCE LINES 403-415 .. code-block:: Python recip_energy = pme_reciprocal_space( positions=positions, charges=charges, cell=cell, alpha=params_comp.alpha, mesh_dimensions=tuple(params_comp.mesh_dimensions), ) print(f" Reciprocal-space (PME): {recip_energy.sum().item():.6f}") print(f" Total: {(real_energy.sum() + recip_energy.sum()).item():.6f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Reciprocal-space (PME): -6.194456 Total: -9.743791 .. GENERATED FROM PYTHON SOURCE LINES 416-421 Legacy Direct Charge-Gradient Outputs ------------------------------------- PME component charge-gradient outputs remain available for compatibility and migration checks. New differentiable training code should prefer the energy autograd recipe verified below. .. GENERATED FROM PYTHON SOURCE LINES 421-463 .. code-block:: Python print("\nCharge Gradients:") # Compute PME reciprocal-space with charge gradients recip_energies, recip_forces, recip_charge_grads = pme_reciprocal_space( positions=positions, charges=charges, cell=cell, alpha=params_comp.alpha, mesh_dimensions=tuple(params_comp.mesh_dimensions), compute_forces=True, compute_charge_gradients=True, ) print(f" PME reciprocal charge gradients shape: {recip_charge_grads.shape}") print( f" PME reciprocal charge gradients range: [{recip_charge_grads.min().item():.4f}, {recip_charge_grads.max().item():.4f}]" ) # Compute real-space with charge gradients real_energies, real_forces, real_charge_grads = ewald_real_space( positions=positions, charges=charges, cell=cell, alpha=params_comp.alpha, neighbor_list=nl_comp, neighbor_ptr=nptr_comp, neighbor_shifts=ns_comp, compute_forces=True, compute_charge_gradients=True, ) print( f" Real-space charge gradients range: [{real_charge_grads.min().item():.4f}, {real_charge_grads.max().item():.4f}]" ) # Total charge gradient is the sum of components total_charge_grads = real_charge_grads + recip_charge_grads print( f" Total charge gradients range: [{total_charge_grads.min().item():.4f}, {total_charge_grads.max().item():.4f}]" ) .. rst-class:: sphx-glr-script-out .. code-block:: none Charge Gradients: /tmp/tmp_3o8x9b4/c1e23460859a784e1d78043bcd1c8af0d1095fa2/examples/electrostatics/03_pme_example.py:425: DeprecationWarning: The component direct-output flag(s) compute_charge_gradients on pme_reciprocal_space are deprecated for differentiable training and will be removed in a future release. Component compute_forces=True remains supported for no-autograd MD/inference loops. For training, compute the energy and use torch.autograd.grad on the energy instead. recip_energies, recip_forces, recip_charge_grads = pme_reciprocal_space( PME reciprocal charge gradients shape: torch.Size([54]) PME reciprocal charge gradients range: [-0.2295, 0.2295] /tmp/tmp_3o8x9b4/c1e23460859a784e1d78043bcd1c8af0d1095fa2/examples/electrostatics/03_pme_example.py:441: DeprecationWarning: The component direct-output flag(s) compute_charge_gradients on ewald_real_space are deprecated for differentiable training and will be removed in a future release. Component compute_forces=True remains supported for no-autograd MD/inference loops. For training, compute the energy and use torch.autograd.grad on the energy instead. real_energies, real_forces, real_charge_grads = ewald_real_space( Real-space charge gradients range: [-0.1315, 0.1315] Total charge gradients range: [-0.3609, 0.3609] .. GENERATED FROM PYTHON SOURCE LINES 464-467 Full PME legacy charge-gradient output in one call: This direct-output flag remains functional for compatibility; use ``torch.autograd.grad`` on the scalar energy for new training code. .. GENERATED FROM PYTHON SOURCE LINES 467-485 .. code-block:: Python energies_full, forces_full, charge_grads_full = _legacy_direct_output_call( particle_mesh_ewald, positions=positions, charges=charges, cell=cell, neighbor_list=nl_comp, neighbor_ptr=nptr_comp, neighbor_shifts=ns_comp, compute_forces=True, compute_charge_gradients=True, accuracy=1e-4, ) print( f"\n Full PME charge gradients range: [{charge_grads_full.min().item():.4f}, {charge_grads_full.max().item():.4f}]" ) .. rst-class:: sphx-glr-script-out .. code-block:: none Full PME charge gradients range: [-0.3609, 0.3609] .. GENERATED FROM PYTHON SOURCE LINES 486-487 Verify legacy direct charge gradients against the training autograd recipe: .. GENERATED FROM PYTHON SOURCE LINES 487-508 .. code-block:: Python charges.requires_grad_(True) energies_total = particle_mesh_ewald( positions=positions, charges=charges, cell=cell, neighbor_list=nl_comp, neighbor_ptr=nptr_comp, neighbor_shifts=ns_comp, accuracy=1e-4, ).sum() energies_total.backward() autograd_charge_grads = charges.grad.clone() charges.requires_grad_(False) charges.grad = None # Compare legacy direct vs autograd charge gradients charge_grad_diff = (charge_grads_full - autograd_charge_grads).abs().max().item() print(f" Explicit vs Autograd charge gradient max diff: {charge_grad_diff:.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none Explicit vs Autograd charge gradient max diff: 5.55e-17 .. GENERATED FROM PYTHON SOURCE LINES 509-512 Batch Evaluation ---------------- Multiple systems can be evaluated simultaneously using batch_idx. .. GENERATED FROM PYTHON SOURCE LINES 512-532 .. code-block:: Python n_systems = 3 all_positions = [] all_charges = [] all_cells = [] all_pbc = [] batch_idx_list = [] print(f"\nBatch Evaluation: Creating {n_systems} systems...") for i in range(n_systems): n_cells = i + 2 # 2×2×2, 3×3×3, 4×4×4 pos, chrg, cell_i, pbc_i = create_nacl_system(n_cells=n_cells) batch_idx_list.extend([i] * len(pos)) all_positions.append(pos) all_charges.append(chrg) all_cells.append(cell_i) all_pbc.append(pbc_i) print(f" System {i}: {len(pos)} atoms ({n_cells}×{n_cells}×{n_cells})") .. rst-class:: sphx-glr-script-out .. code-block:: none Batch Evaluation: Creating 3 systems... System 0: 16 atoms (2×2×2) System 1: 54 atoms (3×3×3) System 2: 128 atoms (4×4×4) .. GENERATED FROM PYTHON SOURCE LINES 533-534 Concatenate all systems: .. GENERATED FROM PYTHON SOURCE LINES 534-551 .. code-block:: Python positions_batch = torch.cat(all_positions, dim=0) charges_batch = torch.cat(all_charges, dim=0) cells_batch = torch.cat(all_cells, dim=0) pbc_batch = torch.cat(all_pbc, dim=0) batch_idx = torch.tensor(batch_idx_list, dtype=torch.int32, device=device) # Estimate parameters for the batch with desired accuracy params_batch = estimate_pme_parameters( positions_batch, cells_batch, batch_idx=batch_idx, accuracy=1e-5 ) print(f"\nTotal atoms: {len(positions_batch)}") print(f"Estimated alphas: {params_batch.alpha.tolist()}") print(f"Mesh dimensions: {params_batch.mesh_dimensions}") print(f"Real-space cutoff: {params_batch.real_space_cutoff.max().item():.2f} Å") .. rst-class:: sphx-glr-script-out .. code-block:: none Total atoms: 198 Estimated alphas: [0.20366053061902537, 0.20366053061902537, 0.20366053061902537] Mesh dimensions: (32, 32, 32) Real-space cutoff: 16.66 Å .. GENERATED FROM PYTHON SOURCE LINES 552-553 Build batched neighbor list and run: .. GENERATED FROM PYTHON SOURCE LINES 553-589 .. code-block:: Python # Use the maximum real-space cutoff across all systems real_cutoff_batch = params_batch.real_space_cutoff.max().item() neighbor_matrix_batch, _, neighbor_matrix_shifts_batch = neighbor_list_fn( positions_batch, real_cutoff_batch, cell=cells_batch, pbc=pbc_batch, method="batch_naive", batch_idx=batch_idx, return_neighbor_list=False, ) t0 = time.time() energies_batch, forces_batch = _legacy_direct_output_call( particle_mesh_ewald, positions=positions_batch, charges=charges_batch, cell=cells_batch, batch_idx=batch_idx, neighbor_matrix=neighbor_matrix_batch, neighbor_matrix_shifts=neighbor_matrix_shifts_batch, compute_forces=True, accuracy=1e-5, # Parameters estimated automatically for batch ) t_batch = (time.time() - t0) * 1000 print(f"\nBatch evaluation time: {t_batch:.2f} ms") print("\nPer-system results:") for i in range(n_systems): mask = batch_idx == i n_atoms = mask.sum().item() sys_energy = energies_batch[mask].sum().item() max_force = torch.norm(forces_batch[mask], dim=1).max().item() print(f" System {i}: {n_atoms} atoms, E={sys_energy:.4f}, |F|_max={max_force:.4f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Batch evaluation time: 2406.78 ms Per-system results: System 0: 16 atoms, E=-2.8870, |F|_max=0.0000 System 1: 54 atoms, E=-9.7437, |F|_max=0.0000 System 2: 128 atoms, E=-23.0962, |F|_max=0.0000 .. GENERATED FROM PYTHON SOURCE LINES 590-591 Preferred training recipe: derive batched forces from the scalar energy. .. GENERATED FROM PYTHON SOURCE LINES 591-608 .. code-block:: Python positions_batch_ag = positions_batch.detach().clone().requires_grad_(True) energies_batch_ag = particle_mesh_ewald( positions=positions_batch_ag, charges=charges_batch, cell=cells_batch, batch_idx=batch_idx, neighbor_matrix=neighbor_matrix_batch, neighbor_matrix_shifts=neighbor_matrix_shifts_batch, accuracy=1e-5, ) forces_batch_ag = -torch.autograd.grad(energies_batch_ag.sum(), positions_batch_ag)[0] print("\nBatched energy-autograd forces:") print(f" Total energy: {energies_batch_ag.sum().item():.4f}") print(f" Max force magnitude: {torch.norm(forces_batch_ag, dim=1).max().item():.4f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Batched energy-autograd forces: Total energy: -35.7269 Max force magnitude: 0.0000 .. GENERATED FROM PYTHON SOURCE LINES 609-610 Verify batch vs individual calculations: .. GENERATED FROM PYTHON SOURCE LINES 610-635 .. code-block:: Python print("\nVerification (individual calculations with same accuracy):") for i in range(n_systems): mask = batch_idx == i pos_i = positions_batch[mask] chrg_i = charges_batch[mask] cell_i = cells_batch[i : i + 1] pbc_i = pbc_batch[i : i + 1] # Use same cutoff as batch for fair comparison nl_i, nptr_i, ns_i = neighbor_list_fn( pos_i, real_cutoff_batch, cell=cell_i, pbc=pbc_i, return_neighbor_list=True ) e_i = particle_mesh_ewald( positions=pos_i, charges=chrg_i, cell=cell_i, neighbor_list=nl_i, neighbor_ptr=nptr_i, neighbor_shifts=ns_i, accuracy=1e-5, # Same accuracy as batch ) print(f" System {i}: E={e_i.sum().item():.4f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Verification (individual calculations with same accuracy): System 0: E=-2.8870 System 1: E=-9.7437 System 2: E=-23.0939 .. GENERATED FROM PYTHON SOURCE LINES 636-641 PME vs Standard Ewald Comparison -------------------------------- PME becomes more efficient than standard Ewald for larger systems due to its O(N log N) scaling vs O(N²) for explicit k-vector summation. Both methods use the same accuracy target for fair comparison. .. GENERATED FROM PYTHON SOURCE LINES 641-650 .. code-block:: Python system_sizes = [2, 3, 4] accuracy_cmp = 1e-5 print(f"\nPME vs Ewald Performance (accuracy={accuracy_cmp:.0e}):") print(" N_cells | N_atoms | Ewald (ms) | PME (ms) | Energy diff") print(" " + "-" * 60) .. rst-class:: sphx-glr-script-out .. code-block:: none PME vs Ewald Performance (accuracy=1e-05): N_cells | N_atoms | Ewald (ms) | PME (ms) | Energy diff ------------------------------------------------------------ .. GENERATED FROM PYTHON SOURCE LINES 651-652 Compare timing and accuracy: .. GENERATED FROM PYTHON SOURCE LINES 652-707 .. code-block:: Python for n_cells in system_sizes: pos_cmp, chrg_cmp, cell_cmp, pbc_cmp = create_nacl_system(n_cells=n_cells) # Estimate parameters for both methods ewald_params = estimate_ewald_parameters(pos_cmp, cell_cmp, accuracy=accuracy_cmp) pme_params = estimate_pme_parameters(pos_cmp, cell_cmp, accuracy=accuracy_cmp) # Use the larger cutoff to ensure both methods have same neighbors real_cutoff_cmp = max( ewald_params.real_space_cutoff.item(), pme_params.real_space_cutoff.item(), ) nl_cmp, nptr_cmp, ns_cmp = neighbor_list_fn( pos_cmp, real_cutoff_cmp, cell=cell_cmp, pbc=pbc_cmp, return_neighbor_list=True ) # Standard Ewald with automatic parameter estimation t0 = time.time() energies_ewald = ewald_summation( positions=pos_cmp, charges=chrg_cmp, cell=cell_cmp, neighbor_list=nl_cmp, neighbor_ptr=nptr_cmp, neighbor_shifts=ns_cmp, accuracy=accuracy_cmp, ) if device.type == "cuda": torch.cuda.synchronize() t_ewald = (time.time() - t0) * 1000 # PME with automatic parameter estimation t0 = time.time() energies_pme = particle_mesh_ewald( positions=pos_cmp, charges=chrg_cmp, cell=cell_cmp, neighbor_list=nl_cmp, neighbor_ptr=nptr_cmp, neighbor_shifts=ns_cmp, accuracy=accuracy_cmp, ) if device.type == "cuda": torch.cuda.synchronize() t_pme = (time.time() - t0) * 1000 e_diff = abs(energies_ewald.sum().item() - energies_pme.sum().item()) print( f" {n_cells} | {len(pos_cmp):5d} | {t_ewald:9.2f} | {t_pme:8.2f} | {e_diff:.2e}" ) print("\nNote: PME becomes increasingly efficient for larger systems.") .. rst-class:: sphx-glr-script-out .. code-block:: none 2 | 16 | 384.11 | 4.28 | 2.53e-06 3 | 54 | 3.38 | 3.42 | 4.30e-08 4 | 128 | 2.98 | 3.35 | 2.47e-06 Note: PME becomes increasingly efficient for larger systems. .. GENERATED FROM PYTHON SOURCE LINES 708-729 Summary ------- This example demonstrated: 1. **Automatic parameter estimation** for alpha and mesh dimensions using ``estimate_pme_parameters`` with target accuracy 2. **Neighbor format flexibility** with list and matrix formats 3. **Accuracy-based convergence** showing how the accuracy parameter controls both mesh resolution and real-space cutoff 4. **Accuracy-parameter relationships** for PME 5. **Component access** for real-space and reciprocal-space 6. **Charge gradients** (∂E/∂q_i) for ML potential training 7. **Batch evaluation** for multiple systems with automatic per-system alpha 8. **PME vs Ewald** performance comparison with same accuracy Key PME steps: - Charge spreading: :math:`Q(\\mathbf{x}) = \\sum_i q_i M_p(\\mathbf{x} - \\mathbf{r}_i)` - FFT convolution: :math:`\\tilde{\\Phi}(\\mathbf{k}) = G(\\mathbf{k}) \\tilde{Q}(\\mathbf{k})` - Force interpolation from mesh gradients - Charge gradient: :math:`\\frac{\\partial E}{\\partial q_i} = \\phi_i` (electrostatic potential) .. GENERATED FROM PYTHON SOURCE LINES 729-731 .. code-block:: Python print("\nPME example complete!") .. rst-class:: sphx-glr-script-out .. code-block:: none PME example complete! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 29.714 seconds) .. _sphx_glr_download_examples_electrostatics_03_pme_example.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 03_pme_example.ipynb <03_pme_example.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 03_pme_example.py <03_pme_example.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 03_pme_example.zip <03_pme_example.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_