.. 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-70 .. code-block:: Python from __future__ import annotations import time 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 71-73 Configure Device ---------------- .. GENERATED FROM PYTHON SOURCE LINES 73-81 .. 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 H100 PCIe .. GENERATED FROM PYTHON SOURCE LINES 82-85 Create a NaCl Crystal System ---------------------------- We define a helper function to create NaCl rock salt crystal supercells. .. GENERATED FROM PYTHON SOURCE LINES 85-124 .. 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 125-130 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 130-138 .. 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 /tmp/tmp_3o8x9b4/844a8844ea67e3784d47e94b0b0ad202396776b7/examples/electrostatics/03_pme_example.py:115: 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 (should be 0 for neutral) .. GENERATED FROM PYTHON SOURCE LINES 139-140 Estimate optimal PME parameters: .. GENERATED FROM PYTHON SOURCE LINES 140-154 .. 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 155-156 Build neighbor list and run PME: .. GENERATED FROM PYTHON SOURCE LINES 156-185 .. 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: 15696.76 ms .. GENERATED FROM PYTHON SOURCE LINES 186-190 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 190-210 .. 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 211-212 Using neighbor list format: .. GENERATED FROM PYTHON SOURCE LINES 212-228 .. 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=4.00 ms .. GENERATED FROM PYTHON SOURCE LINES 229-230 Using neighbor matrix format: .. GENERATED FROM PYTHON SOURCE LINES 230-250 .. 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=4.10 ms Energy difference: 0.00e+00 Max force difference: 2.52e-17 .. GENERATED FROM PYTHON SOURCE LINES 251-256 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 256-264 .. 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 265-266 Run PME with different accuracy targets: .. GENERATED FROM PYTHON SOURCE LINES 266-299 .. 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.63 ms 1e-03 | 0.204 | (16, 16, 16) | 12.91 | -9.745599 | 4.35 ms 1e-04 | 0.204 | (16, 16, 16) | 14.90 | -9.743791 | 2.92 ms 1e-05 | 0.204 | (32, 32, 32) | 16.66 | -9.743703 | 4.78 ms 1e-06 | 0.204 | (64, 64, 64) | 18.25 | -9.743761 | 2.90 ms 1e-07 | 0.204 | (64, 64, 64) | 19.71 | -9.743762 | 2.91 ms .. GENERATED FROM PYTHON SOURCE LINES 300-301 Show convergence relative to highest accuracy: .. GENERATED FROM PYTHON SOURCE LINES 301-308 .. 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 309-313 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 313-321 .. 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 322-323 Sweep through accuracy values: .. GENERATED FROM PYTHON SOURCE LINES 323-353 .. 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 354-357 Accessing Real-Space and Reciprocal-Space Components ---------------------------------------------------- You can compute the components separately if needed. .. GENERATED FROM PYTHON SOURCE LINES 357-370 .. 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 371-372 Real-space component (uses same kernel as Ewald): .. GENERATED FROM PYTHON SOURCE LINES 372-385 .. 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 386-387 PME reciprocal-space component (FFT-based): .. GENERATED FROM PYTHON SOURCE LINES 387-399 .. 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 400-405 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 405-447 .. 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 448-449 Full PME with charge gradients in one call: .. GENERATED FROM PYTHON SOURCE LINES 449-466 .. 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 467-468 Verify charge gradients against autograd: .. GENERATED FROM PYTHON SOURCE LINES 468-489 .. 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: 1.67e-16 .. GENERATED FROM PYTHON SOURCE LINES 490-493 Batch Evaluation ---------------- Multiple systems can be evaluated simultaneously using batch_idx. .. GENERATED FROM PYTHON SOURCE LINES 493-513 .. 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 514-515 Concatenate all systems: .. GENERATED FROM PYTHON SOURCE LINES 515-532 .. 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 533-534 Build batched neighbor list and run: .. GENERATED FROM PYTHON SOURCE LINES 534-569 .. 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: 7.94 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 570-571 Verify batch vs individual calculations: .. GENERATED FROM PYTHON SOURCE LINES 571-596 .. 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 597-602 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 602-611 .. 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 612-613 Compare timing and accuracy: .. GENERATED FROM PYTHON SOURCE LINES 613-668 .. 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 | 13.09 | 3.28 | 2.53e-06 3 | 54 | 2.11 | 2.93 | 4.31e-08 4 | 128 | 2.07 | 2.98 | 2.47e-06 Note: PME becomes increasingly efficient for larger systems. .. GENERATED FROM PYTHON SOURCE LINES 669-690 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 690-692 .. 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 21.052 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 `_