Issue with Convergence in Jastrow Method for Ground State Energy Calculation

[MrAAafa99]

Hello,

I am using the Jastrow method to calculate the ground state energy of a multi-particle system with a specified interaction, as detailed in the code below. However, I do not observe any convergence in energy calculation.

Could the issue be related to the implementation of the Jastrow method, or might the problem lie elsewhere? Specifically, how should I adjust the use of the Jastrow method in this code to achieve better energy convergence?

I would appreciate any guidance or suggestions you can provide.

Here is the code I'm using:

import jax
import jax.numpy as jnp
import flax.linen as nn
import flax
from flax.linen.dtypes import promote_dtype
from flax.linen.initializers import normal
import matplotlib.pyplot as plt
import numpy as np
import netket as nk
from netket.utils.types import DType, Array, NNInitFunc


# System parameters
N = 4
sdim = 1
g = 10.0  # Interaction strength

# Define the Hilbert space
hilb = nk.hilbert.Particle(N=N, L=jnp.inf, pbc=False)

def minimum_distance(x, sdim):
    """
    Computes distances between particles
    """
    n_particles = x.shape[0] // sdim
    x = x.reshape(-1, sdim)

    above_diagonal = (-x[jnp.newaxis, :, :] + x[:, jnp.newaxis, :])[jnp.triu_indices(N, 1)]

    above_diagonal_norm = jnp.linalg.norm(above_diagonal, axis=1)

    return above_diagonal_norm

# Function to compute the harmonic oscillator
def potential1(x):
    potential_energy = 0.5 * jnp.linalg.norm(x) ** 2
    return potential_energy

# Function to compute interaction potential
def potential2(x, sdim):
    dis = minimum_distance(x, sdim=sdim)
    potential_energy = jnp.sum(g / ((dis) ** 2 + 1))
    return potential_energy

# Define the operators
ekin = nk.operator.KineticEnergy(hilb, mass=1.0)
pot1 = nk.operator.PotentialEnergy(hilb, potential1)
pot2 = nk.operator.PotentialEnergy(hilb, lambda x: potential2(x, 1))
ha = ekin + pot1 + pot2

# Increase number of chains and decrease discard
sa = nk.sampler.MetropolisGaussian(hilb, sigma=0.1, n_chains=32)

class Jastrow(nn.Module):
    r"""
    Jastrow wave function :math:`\Psi(s) = \exp(\sum_{i \neq j} s_i W_{ij} s_j)`,
    where W is a symmetric matrix.
    """

    param_dtype: DType = jnp.complex128
    kernel_init: NNInitFunc = normal()

    @nn.compact
    def __call__(self, x_in: Array):
        nv = x_in.shape[-1]
        il = jnp.tril_indices(nv, k=-1)
        kernel = self.param(
            "kernel", self.kernel_init, (nv * (nv - 1) // 2,), self.param_dtype
        )
        W = jnp.zeros((nv, nv), dtype=self.param_dtype).at[il].set(kernel)
        W, x_in = promote_dtype(W, x_in, dtype=None)
        y = jnp.einsum("...i,ij,...j", x_in, W, x_in)
        return y

# Create an instance of the model
Jastrow = Jastrow()

# Construct the variational state using the model and the sampler above
vs = nk.vqs.MCState(sa, Jastrow, n_samples=10**5, n_discard_per_chain=100)  

# First we reset the parameters to run the optimisation again
vs.init_parameters(normal(stddev=1.0))

# Then we create an optimiser from the standard library
optimizer = nk.optimizer.Sgd(learning_rate=0.05)

# Build the optimisation driver
gs = nk.driver.VMC(ha, optimizer, variational_state=vs)

# Run the driver for more iterations
log = nk.logging.RuntimeLog()
gs.run(n_iter=5000, out=log)  # Increased number of iteration

[PhilipVinc]

The jastrow you are using makes sense for a lattice system. But for a continuous space system you need to build a jastrow out of some single particle orbitals which you should chose reasonably.

A simple choice you could make, though not very good, would be to use Gaussian Orbitals.

You can read the first chapter of the Becca Sorella book to learn about it.

[MrAAafa99]

After reading the first chapter of the bookthe Becca Sorella book, I rewrote my Jastrow model from scratch and implemented it in the following way:

class  Gaussian(nn.Module):

    @nn.compact
    def __call__(self, x_in: Array):
        nv = x_in.shape[-1]

        kernel = self.param("kernel", normal(), (nv, nv), jnp.float64)
        kernel = jnp.dot(kernel.T, kernel)

        kernel, x_in = promote_dtype(kernel, x_in, dtype=None)
        y = -0.5 * jnp.einsum("...i,ij,...j", x_in, kernel, x_in)

        return y



class interaction(nn.Module):
    param_dtype: DType = jnp.complex128
    kernel_init: NNInitFunc = normal()


    @nn.compact
    def __call__(self, x_in: Array):
        batch_size, n_particles = x_in.shape[0], x_in.shape[1]

        # Calculate pairwise distances
        pairwise_distances = jnp.abs(x_in[:, :, jnp.newaxis] - x_in[:, jnp.newaxis, :])

        # Initialize parameters for the potential function
        params = {}
        param_names =[f'a_{i}' for i in range(0, 10)]
        for name in param_names:
            params[name] = self.param(name, self.kernel_init, (), self.param_dtype)

        # Calculate the potential U(xi-xj)
        potential = 0
        for i in range(0, 10):
          potential = (params1[f'a_{i}'] *(pairwise_distances)**(i))
        # Sum only the lower triangle of the distance matrix (excluding the diagonal)
        y = jnp.sum(jnp.tril(potential, k=-1), axis=(1, 2))

        # Output log of wave function
        log_wave_function = jnp.log(y)
        return log_wave_function



class fullmodel(nn.Module):

    def setup(self):
        self.m1 = Gaussian()
        self.m2 = interaction()

    @nn.compact
    def __call__(self, x):
        result = self.m1(x) + self.m2(x)
        return result

one class for Gaussian orbitals and another for the terms arising from the distances between particles with variational parameters. However, the Jastrow model still lacks the required accuracy, and in some cases, increasing the number of variational parameters causes the energy to diverge.

[gcarleo]

hi, I am not sure this is a good choice for a Jastrow. I would suggest you take the standard Jastrow that includes cusp conditions and that maybe has only one or two variational parameters. Depending on the application it will have a different form. I suggest you read this review for an introduction https://journals.aps.org/rmp/abstract/10.1103/RevModPhys.73.33. Becca and Sorella is an excellent book, but to my knowledge it does not cover very extensively this topic

[gcarleo]

in general, for most used interaction potentials analytical, "simple" jastrow form is known in the literature. for example for the Lennard Jones potential it's the McMillan form of the Jastrow found in the first paper on VMC, for Coulomb interactions can be found in the RMP I linked earlier etc. If you are using an interaction potential that has not been studied, you need to solve the Schroedinger equation for two particles exactly, for example with finite differences. Then from the log of the two-particle wave function you can infer the form of the jastrow.