4-site water models

# Several water models (starting from the venerable TIP4P model of # Abascal and C. Vega, JCP (2005)) # use a center of negative charge that is displaced from the O position. # This is easily implemented, yielding the forces on the O and H positions # generated by the displaced charge.

import ase
import torch

import torchpme

structure = ase.Atoms(
    positions=[
        [0, 0, 0],
        [0, 1, 0],
        [1, -0.2, 0],
    ],
    cell=[6, 6, 6],
    symbols="OHH",
)

cell = torch.from_numpy(structure.cell.array)
positions = torch.from_numpy(structure.positions)

The key step is to create a “fourth site” based on the oxygen positions and use it in the interpolate step.

charges = torch.tensor([[-1.0], [0.5], [0.5]])

positions.requires_grad_(True)
charges.requires_grad_(True)
cell.requires_grad_(True)

positions_4site = torch.vstack(
    [
        ((positions[1::3] + positions[2::3]) * 0.5 + positions[0::3] * 3) / 4,
        positions[1::3],
        positions[2::3],
    ]
)

Important

For the automatic differentiation to work it is important to make a new tensor as positions_4site and do not “overwrite” the original tensor.

ns = torch.tensor([5, 5, 5])
interpolator = torchpme.lib.MeshInterpolator(
    cell=cell, ns_mesh=ns, interpolation_nodes=3, method="Lagrange"
)
interpolator.compute_weights(positions_4site)
mesh = interpolator.points_to_mesh(charges)

value = (mesh**2).sum()

The gradients can be computed by just running backward on the end result. Gradients are computed on the H and O positions.

value.backward()

print(
    f"""
Position gradients:
{positions.grad.T}

Cell gradients:
{cell.grad}

Charges gradients:
{charges.grad.T}
"""
)

Position gradients:
tensor([[-0.5106, -0.0857,  0.5133],
        [-0.4691,  0.5263, -0.0164],
        [ 0.0000,  0.0000,  0.0000]], dtype=torch.float64)

Cell gradients:
tensor([[-0.0855,  0.0027, -0.0000],
        [ 0.0314, -0.0883, -0.0000],
        [-0.0000, -0.0000, -0.0000]], dtype=torch.float64)

Charges gradients:
tensor([[-1.6628,  0.6843,  0.6317]])