Tutorials#

This page collects practical, copy-pasteable examples for common tasks.

Compute ESP and export a cube#

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 water_esp.pkl

# Map sampled ESP to rectilinear grid and write water.cube
# (See esp_cube_generation page for a full script.)

Train DCMNet on ESP datasets (CLI)#

python -m mmml.dcmnet.dcmnet.main \
  --data_files /path/set1.npz /path/set2.npz \
  --n_train 80000 --n_valid 2000 \
  --num_epochs 5000 --batch_size 1 --learning_rate 1e-3 \
  --esp_w 10000 --n_dcm 1 --type default \
  --log_dir ./runs/example --checkpoint_dir ./ckpts/example

Train DCMNet (programmatic)#

import jax
from tensorboardX import SummaryWriter
from mmml.dcmnet.dcmnet.data import prepare_datasets
from mmml.dcmnet.dcmnet.modules import MessagePassingModel
from mmml.dcmnet.dcmnet.training import train_model

key = jax.random.PRNGKey(0)
train_data, valid_data = prepare_datasets(
    key, 80000, 2000, ["/path/set1.npz", "/path/set2.npz"], clean=True
)
model = MessagePassingModel(16, 2, 2, 16, cutoff=4.0, n_dcm=1)
writer = SummaryWriter("./runs/programmatic")
train_model(key, model, train_data, valid_data, 5000, 1e-3, 1, writer, ndcm=1, esp_w=10000)

Dipole-augmented training#

from mmml.dcmnet.dcmnet.training import train_model_dipo
# Requires dataset with Dxyz, com, espMask
train_model_dipo(key, model, train_data, valid_data, 3000, 5e-4, 1, writer, ndcm=1, esp_w=10000)

Prepare a minimal NPZ from an ESP pickle#

import numpy as np, pickle
with open("water_esp.pkl", "rb") as f:
    d = pickle.load(f)
R = d["R"][None, ...]
Z = d["Z"][None, ...]
N = np.array([[len(d["Z"])]]);
mono = np.zeros((1, len(d["Z"])));
esp = d["esp"][None, ...]
vdw_surface = d["esp_grid"][None, ...]
n_grid = np.array([len(d["esp"])])
np.savez("water_esp_train.npz", R=R, Z=Z, N=N, mono=mono,
         esp=esp, vdw_surface=vdw_surface, n_grid=n_grid)

Tips#

  • Reduce memory by lowering batch size and model size.

  • Use TensorBoard to monitor Loss/train and Loss/valid.