Generating ESP cube files
=========================

This guide shows how to compute electrostatic potentials (ESP) on a grid using GPU-accelerated PySCF and how to export those values to a Gaussian ``.cube`` file for visualization in tools like VMD or Avogadro.

Prerequisites
-------------

- A working CUDA setup if using GPU acceleration
- Python dependencies from the docs requirements
- Optional: `VMD`, `Avogadro`, or similar for visualization

Compute ESP with GPU4PySCF
--------------------------

The function ``compute_dft`` in ``mmml.pyscf4gpuInterface.calcs`` supports an ESP-and-density workflow when the ``--dens_esp`` flag is set. It returns arrays for the ESP values and their grid positions alongside molecule data.

Command-line usage
~~~~~~~~~~~~~~~~~~

Example: compute ESP for water (neutral singlet) at PBE0/def2-TZVP and write a pickle with results.

.. code-block:: bash

   python -m mmml.pyscf4gpuInterface.calcs \
       --mol "O 0.000 0.000 0.000; H 0.000 0.757 0.586; H 0.000 -0.757 0.586" \
       --basis def2-tzvp \
       --xc wB97m-v \
       --spin 0 \
       --charge 0 \
       --dens_esp \
       --energy \
       --output esp_output.pkl

The output pickle contains (at least):

- ``esp``: ESP values at the selected grid points (array of shape ``(N,)``)
- ``esp_grid``: grid point positions in Ångström (array of shape ``(N, 3)``)
- ``R``: nuclear coordinates in Ångström (array of shape ``(M, 3)``)
- ``Z``: nuclear charges (array of shape ``(M,)``)

Note: The grid points are selected by density thresholds inside ``compute_dft`` so that the potential is sampled on a physically meaningful shell. You can adapt those thresholds in code if desired.

Export ESP to a Gaussian .cube file
-----------------------------------

Gaussian cube files require a rectilinear grid. If you prefer to visualize the non-rectilinear sampled potential directly, many tools can interpolate scattered points. Otherwise, the snippet below shows how to map the sampled ESP to a coarse rectilinear grid and write a cube.

.. code-block:: python

   import numpy as np
   import pickle

   # Load data produced by mmml.pyscf4gpuInterface.calcs
   with open("esp_output.pkl", "rb") as f:
       data = pickle.load(f)

   esp = data["esp"]            # (N,)
   grid = data["esp_grid"]      # (N, 3) in Angstrom
   R = data["R"]                # (M, 3)
   Z = data["Z"]                # (M,)

   # Define a cube (rectilinear) grid around the molecule
   pad = 4.0  # Angstrom padding
   mins = np.min(R, axis=0) - pad
   maxs = np.max(R, axis=0) + pad
   nx, ny, nz = 40, 40, 40
   xs = np.linspace(mins[0], maxs[0], nx)
   ys = np.linspace(mins[1], maxs[1], ny)
   zs = np.linspace(mins[2], maxs[2], nz)
   X, Y, Zz = np.meshgrid(xs, ys, zs, indexing="ij")
   cube_points = np.stack([X, Y, Zz], axis=-1).reshape(-1, 3)

   # Interpolate scattered ESP onto the rectilinear grid (nearest-neighbor)
   from scipy.spatial import cKDTree
   tree = cKDTree(grid)
   _, idx = tree.query(cube_points, k=1)
   cube_values = esp[idx].reshape(nx, ny, nz)

   def write_cube(path, atoms_Z, atoms_R, origin, axes, values):
       # values must be shaped (nx, ny, nz)
       nx, ny, nz = values.shape
       ax, ay, az = axes  # each is (3,) box vector in Angstrom / count
       with open(path, "w") as f:
           f.write("ESP cube generated by mmml\n")
           f.write("ESP on a rectilinear grid\n")
           f.write(f"{len(atoms_Z):4d} {origin[0]:12.6f} {origin[1]:12.6f} {origin[2]:12.6f}\n")
           f.write(f"{nx:4d} {ax[0]:12.6f} {ax[1]:12.6f} {ax[2]:12.6f}\n")
           f.write(f"{ny:4d} {ay[0]:12.6f} {ay[1]:12.6f} {ay[2]:12.6f}\n")
           f.write(f"{nz:4d} {az[0]:12.6f} {az[1]:12.6f} {az[2]:12.6f}\n")
           for Zq, Rq in zip(atoms_Z, atoms_R):
               f.write(f"{int(Zq):4d} {float(Zq):12.6f} {Rq[0]:12.6f} {Rq[1]:12.6f} {Rq[2]:12.6f}\n")
           # Write voxel data, 6 values per line
           count = 0
           for i in range(nx):
               for j in range(ny):
                   line_vals = []
                   for k in range(nz):
                       line_vals.append(f"{values[i, j, k]:13.5e}")
                       count += 1
                       if len(line_vals) == 6:
                           f.write(" ".join(line_vals) + "\n")
                           line_vals = []
                   if line_vals:
                       f.write(" ".join(line_vals) + "\n")

   # Build cube axes from box lengths
   ax = np.array([xs[1]-xs[0], 0.0, 0.0])
   ay = np.array([0.0, ys[1]-ys[0], 0.0])
   az = np.array([0.0, 0.0, zs[1]-zs[0]])
   write_cube(
       "esp.cube",
       atoms_Z=Z,
       atoms_R=R,
       origin=np.array([xs[0], ys[0], zs[0]]),
       axes=(ax, ay, az),
       values=cube_values,
   )

Visualize the cube
------------------

- VMD: File → New Molecule → Browse to ``esp.cube`` → Load
- Avogadro: File → Open → select ``esp.cube``; enable Surfaces → ESP

Troubleshooting
---------------

- If the ESP looks noisy, increase the rectilinear grid resolution or apply a smoother interpolant.
- To compute ESP exactly on a desired rectilinear grid rather than mapping later, adapt the grid generation in ``compute_dft`` and evaluate ESP directly at those points.