fix multigrid v2 stress spike

[puzhichen]

_rks_exc_strain_deriv reused the base cell's cached image_indices and block mappings for strained cells while eplacing neighboring_images with a new list from get_lattice_Ls(cell_strained). Because get_lattice_Ls applies distance < rcut cutoff, a boundary image can enter or leave the list under a 1e-4 strain (946 vs 947 images on t-HfO2 -xx/-yy), shifting all subsequent indices and causing pairs to reference wrong translations. A demo example is added as the unit test, and listed here

import numpy as np
from pyscf.data.nist import BOHR
from pyscf.pbc import gto


# Lattice and atomic coordinates of t-HfO2; all elements are changed to H for fast Cell construction.
LATTICE = np.array([
    [3.535871610000000, 0.000000000293102, -0.000000000005598],
    [0.000000000293102, 3.535871610000000, -0.000000000129270],
    [-0.000000000007998, -0.000000000189320, 5.190846310000000],
])
COORDS = [
    (1.767935803700000, 1.767935803726000, 2.595423152457000),
    (-0.000000000170000, -0.000000000413000, -0.000000000009000),
    (-0.000000000217000, 1.767935803226000, 3.539517351764000),
    (1.767935803466000, -0.000000000847000, 4.246760404410000),
    (1.767935803449000, -0.000000000675000, 1.651328953277000),
    (-0.000000000265000, 1.767935803396000, 0.944085900686000),
]


def strained_cell(cell, strain):
    lattice = cell.lattice_vectors()
    coords = cell.atom_coords()
    lattice *= BOHR
    coords *= BOHR
    transform = np.eye(3)
    transform[0, 0] += strain
    new_cell = cell.set_geom_(coords @ transform.T, inplace=False)
    new_cell.a = lattice @ transform.T
    return new_cell


def integer_images(cell):
    images = gto.eval_gto.get_lattice_Ls(cell)
    return np.rint(images @ np.linalg.inv(cell.lattice_vectors())).astype(int)


cell = gto.Cell()
cell.atom = [('H', xyz) for xyz in COORDS]
cell.a = LATTICE
cell.unit = 'Angstrom'
cell.basis = [[0, [1.0, 1.0]]]
cell.rcut = 34.71833632109211
cell.verbose = 0
cell.build()

base = integer_images(cell)
plus = integer_images(strained_cell(cell, +1e-4))
minus = integer_images(strained_cell(cell, -1e-4))

first_mismatch = next(
    i for i, (base_image, minus_image) in enumerate(zip(base, minus))
    if not np.array_equal(base_image, minus_image)
)

assert len(base) == len(plus) == 946
assert len(minus) == 947
assert first_mismatch == 15