.. 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 43-47 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 47-66 .. code-block:: Python from __future__ import annotations import time import numpy as np import torch from nvalchemiops.interactions.electrostatics import ( ewald_real_space, ewald_summation, particle_mesh_ewald, pme_reciprocal_space, ) from nvalchemiops.interactions.electrostatics.parameters import ( estimate_pme_parameters, ) from nvalchemiops.neighborlist import neighbor_list as neighbor_list_fn .. GENERATED FROM PYTHON SOURCE LINES 67-69 Configure Device ---------------- .. GENERATED FROM PYTHON SOURCE LINES 69-77 .. 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") .. rst-class:: sphx-glr-script-out .. code-block:: none Using CUDA device: NVIDIA L4 .. GENERATED FROM PYTHON SOURCE LINES 78-81 Create a NaCl Crystal System ---------------------------- We define a helper function to create NaCl rock salt crystal supercells. .. GENERATED FROM PYTHON SOURCE LINES 81-120 .. 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 121-126 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 126-134 .. 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} (should be 0 for neutral)") .. rst-class:: sphx-glr-script-out .. code-block:: none System: 54 atoms NaCl crystal Cell size: 16.92 Å Total charge: 0.0 (should be 0 for neutral) .. GENERATED FROM PYTHON SOURCE LINES 135-136 Estimate optimal PME parameters: .. GENERATED FROM PYTHON SOURCE LINES 136-150 .. 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 151-152 Build neighbor list and run PME: .. GENERATED FROM PYTHON SOURCE LINES 152-181 .. 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() energies, forces = 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 ) 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: 15170.04 ms .. GENERATED FROM PYTHON SOURCE LINES 182-186 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 186-206 .. 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 207-208 Using neighbor list format: .. GENERATED FROM PYTHON SOURCE LINES 208-224 .. code-block:: Python t0 = time.time() energies_list, forces_list = 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=5.18 ms .. GENERATED FROM PYTHON SOURCE LINES 225-226 Using neighbor matrix format: .. GENERATED FROM PYTHON SOURCE LINES 226-246 .. code-block:: Python t0 = time.time() energies_matrix, forces_matrix = 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=10.89 ms Energy difference: 0.00e+00 Max force difference: 2.20e-17 .. GENERATED FROM PYTHON SOURCE LINES 247-252 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 252-260 .. 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 261-262 Run PME with different accuracy targets: .. GENERATED FROM PYTHON SOURCE LINES 262-295 .. 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 | 4.53 ms 1e-03 | 0.204 | (16, 16, 16) | 12.91 | -9.745599 | 5.26 ms 1e-04 | 0.204 | (16, 16, 16) | 14.90 | -9.743791 | 2.81 ms 1e-05 | 0.204 | (32, 32, 32) | 16.66 | -9.743703 | 4.47 ms 1e-06 | 0.204 | (64, 64, 64) | 18.25 | -9.743761 | 3.51 ms 1e-07 | 0.204 | (64, 64, 64) | 19.71 | -9.743762 | 3.24 ms .. GENERATED FROM PYTHON SOURCE LINES 296-297 Show convergence relative to highest accuracy: .. GENERATED FROM PYTHON SOURCE LINES 297-304 .. 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 305-309 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 309-317 .. 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 318-319 Sweep through accuracy values: .. GENERATED FROM PYTHON SOURCE LINES 319-349 .. 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 350-353 Accessing Real-Space and Reciprocal-Space Components ---------------------------------------------------- You can compute the components separately if needed. .. GENERATED FROM PYTHON SOURCE LINES 353-366 .. 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 367-368 Real-space component (uses same kernel as Ewald): .. GENERATED FROM PYTHON SOURCE LINES 368-381 .. 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 382-383 PME reciprocal-space component (FFT-based): .. GENERATED FROM PYTHON SOURCE LINES 383-395 .. 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 396-401 Charge Gradients for ML Potentials ---------------------------------- PME supports computing analytical charge gradients (∂E/∂q_i), which are useful for training machine learning potentials that predict atomic partial charges. The charge gradient represents the electrostatic potential at each atom. .. GENERATED FROM PYTHON SOURCE LINES 401-443 .. 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: PME reciprocal charge gradients shape: torch.Size([54]) PME reciprocal charge gradients range: [-0.2295, 0.2295] Real-space charge gradients range: [-0.1315, 0.1315] Total charge gradients range: [-0.3609, 0.3609] .. GENERATED FROM PYTHON SOURCE LINES 444-445 Full PME with charge gradients in one call: .. GENERATED FROM PYTHON SOURCE LINES 445-462 .. code-block:: Python energies_full, forces_full, charge_grads_full = 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 463-464 Verify charge gradients against autograd: .. GENERATED FROM PYTHON SOURCE LINES 464-485 .. 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 explicit 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: 2.22e-16 .. GENERATED FROM PYTHON SOURCE LINES 486-489 Batch Evaluation ---------------- Multiple systems can be evaluated simultaneously using batch_idx. .. GENERATED FROM PYTHON SOURCE LINES 489-509 .. 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 510-511 Concatenate all systems: .. GENERATED FROM PYTHON SOURCE LINES 511-528 .. 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"Per-system 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 Per-system alphas: [0.24943219038054099, 0.20366053061902537, 0.17637519326429443] Mesh dimensions: (32, 32, 32) Real-space cutoff: 19.24 Å .. GENERATED FROM PYTHON SOURCE LINES 529-530 Build batched neighbor list and run: .. GENERATED FROM PYTHON SOURCE LINES 530-565 .. 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 = 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: 32.91 ms Per-system results: System 0: 16 atoms, E=-2.8870, |F|_max=0.0000 System 1: 54 atoms, E=-9.7438, |F|_max=0.0000 System 2: 128 atoms, E=-23.0963, |F|_max=0.0000 .. GENERATED FROM PYTHON SOURCE LINES 566-567 Verify batch vs individual calculations: .. GENERATED FROM PYTHON SOURCE LINES 567-592 .. 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.7438 System 2: E=-23.0963 .. GENERATED FROM PYTHON SOURCE LINES 593-598 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 598-610 .. code-block:: Python from nvalchemiops.interactions.electrostatics.parameters import ( estimate_ewald_parameters, ) 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 611-612 Compare timing and accuracy: .. GENERATED FROM PYTHON SOURCE LINES 612-667 .. 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 | 2.29 | 2.94 | 2.53e-06 3 | 54 | 2.06 | 3.30 | 4.31e-08 4 | 128 | 3.38 | 3.17 | 2.47e-06 Note: PME becomes increasingly efficient for larger systems. .. GENERATED FROM PYTHON SOURCE LINES 668-689 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 689-691 .. 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 15.372 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 `_