Error with specifying the mass of each particle

[rdpdlv]

Hello,

I am trying to model the H2 molecule via netket. I define a 4 particle Hilbert space, and would like to ignore the kinetic energy of the nuclei. To do that, I try to set the mass of the nuclei to be infinity in defining the kinetic energy operator: mass = (1.0, 1.0, jnp.inf, jnp.inf). This is my code for testing:

import platform
import netket as nk

import jax
from jax import random, numpy as jnp
import flax
from flax import linen as nn
import math

# 4 particles in 3D space without periodic boundary conditions
hi = nk.hilbert.Particle(N=4, D=3)

sa = nk.sampler.MetropolisGaussian(hi, sigma=0.1)

ekin = nk.operator.KineticEnergy(hi, mass=(1.0, 1.0, jnp.inf, jnp.inf))

def pot(x):
  x = x.reshape(-1,3)
  dist = jnp.linalg.norm( (-x[jnp.newaxis,:,:]+x[:,jnp.newaxis,:])[jnp.triu_indices(4,1)], axis=1)
  return 1.0 / dist[0] + 1.0 / dist[-1] - jnp.sum(1.0 / dist[1:-1])

pot = nk.operator.PotentialEnergy(hi, pot)

ha = pot + ekin

class SimpleNN(nn.Module):
  features = (2, 1)

  def setup(self):
    self.dense1 = nn.Dense(self.features[0])
    self.dense2 = nn.Dense(self.features[1])

  @nn.compact
  def __call__(self,x):
    return jnp.squeeze(self.dense2(nn.relu(self.dense1(x))))

snn = SimpleNN()

vstate = nk.vqs.MCState(sa, snn, n_samples=128)

from flax.linen.initializers import normal

vstate.init_parameters(normal(stddev=1.0))

optimizer = nk.optimizer.Sgd(learning_rate=0.1)

gs = nk.driver.VMC(ha, optimizer, variational_state=vstate,
                   preconditioner=nk.optimizer.SR(diag_shift=0.1))

log = nk.logging.RuntimeLog()
gs.run(n_iter=10, out=log, timeit=True)

However, I got the following error when running the code:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-43-6b0eb915d26d>](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in <cell line: 2>()
      1 log = nk.logging.RuntimeLog()
----> 2 gs.run(n_iter=10, out=log, timeit=True)

14 frames
[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in run(self, n_iter, out, obs, step_size, show_progress, save_params_every, write_every, callback, timeit)
    331                 first_step = True
    332 
--> 333                 for step in self.iter(n_iter, step_size):
    334                     with timing.timed_scope(name="observables"):
    335                         log_data = self.estimate(obs)

[/usr/local/lib/python3.10/dist-packages/netket/driver/abstract_variational_driver.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in iter(self, n_steps, step)
    222         for _ in range(0, n_steps, step):
    223             for i in range(0, step):
--> 224                 dp = self._forward_and_backward()
    225                 if i == 0:
    226                     yield self.step_count

[/usr/local/lib/python3.10/dist-packages/netket/driver/vmc.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in _forward_and_backward(self)
    124 
    125         # Compute the local energy estimator and average Energy
--> 126         self._loss_stats, self._loss_grad = self.state.expect_and_grad(self._ham)
    127 
    128         # if it's the identity it does

    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/netket/vqs/mc/mc_state/state.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in expect_and_grad(self, O, mutable, **kwargs)
    620             mutable = self.mutable
    621 
--> 622         return expect_and_grad(
    623             self,
    624             O,

    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/netket/vqs/mc/mc_state/expect_grad.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in expect_and_grad_default_formula(vstate, Ô, chunk_size, mutable, use_covariance, *args)
     56         # `expect_and_forces` and `force_to_grad` functions.
     57         # return expect_and_grad_from_covariance(vstate, Ô, *args, mutable=mutable)
---> 58         Ō, Ō_grad = expect_and_forces(vstate, Ô, chunk_size, *args, mutable=mutable)
     59         Ō_grad = force_to_grad(Ō_grad, vstate.parameters)
     60         return Ō, Ō_grad

    [... skipping hidden 1 frame]

[/usr/local/lib/python3.10/dist-packages/netket/vqs/mc/mc_state/expect_forces.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in expect_and_forces(vstate, Ô, chunk_size, mutable)
     50     local_estimator_fun = get_local_kernel(vstate, Ô)
     51 
---> 52     Ō, Ō_grad, new_model_state = forces_expect_hermitian(
     53         local_estimator_fun,
     54         vstate._apply_fun,

    [... skipping hidden 11 frame]

[/usr/local/lib/python3.10/dist-packages/netket/vqs/mc/mc_state/expect_forces.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in forces_expect_hermitian(local_value_kernel, model_apply_fun, mutable, parameters, model_state, σ, local_value_args)
     82     n_samples = σ.shape[0] * mpi.n_nodes
     83 
---> 84     O_loc = local_value_kernel(
     85         model_apply_fun,
     86         {"params": parameters, **model_state},

[/usr/local/lib/python3.10/dist-packages/netket/operator/_sumoperators.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in _expect_kernel(logpsi, params, x, data)
    120         logpsi: Callable, params: PyTree, x: Array, data: Optional[PyTree]
    121     ):
--> 122         result = [
    123             data.coeffs[i] * op._expect_kernel(logpsi, params, x, op_data)
    124             for i, (op, op_data) in enumerate(zip(data.ops, data.op_data))

[/usr/local/lib/python3.10/dist-packages/netket/operator/_sumoperators.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in <listcomp>(.0)
    121     ):
    122         result = [
--> 123             data.coeffs[i] * op._expect_kernel(logpsi, params, x, op_data)
    124             for i, (op, op_data) in enumerate(zip(data.ops, data.op_data))
    125         ]

    [... skipping hidden 3 frame]

[/usr/local/lib/python3.10/dist-packages/netket/operator/_kinetic.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in _expect_kernel(self, logpsi, params, x, coefficient)
    100         self, logpsi: Callable, params: PyTree, x: Array, coefficient: Optional[PyTree]
    101     ):
--> 102         return self._expect_kernel_single(logpsi, params, x, coefficient)
    103 
    104     def _pack_arguments(self) -> PyTree:

[/usr/local/lib/python3.10/dist-packages/netket/operator/_kinetic.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in _expect_kernel_single(self, logpsi, params, x, inverse_mass)
     94         dp_dx = dlogpsi_x(x)[0][0] ** 2
     95 
---> 96         return -0.5 * jnp.sum(inverse_mass * (dp_dx2 + dp_dx), axis=-1)
     97 
     98     @partial(jax.vmap, in_axes=(None, None, None, 0, None))

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in op(self, *args)
    734 def _forward_operator_to_aval(name):
    735   def op(self, *args):
--> 736     return getattr(self.aval, f"_{name}")(self, *args)
    737   return op
    738 

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in deferring_binary_op(self, other)
    262     args = (other, self) if swap else (self, other)
    263     if isinstance(other, _accepted_binop_types):
--> 264       return binary_op(*args)
    265     # Note: don't use isinstance here, because we don't want to raise for
    266     # subclasses, e.g. NamedTuple objects that may override operators.

    [... skipping hidden 11 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/ufuncs.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in fn(x1, x2)
     98   def fn(x1, x2, /):
     99     x1, x2 = promote_args(numpy_fn.__name__, x1, x2)
--> 100     return lax_fn(x1, x2) if x1.dtype != np.bool_ else bool_lax_fn(x1, x2)
    101   fn.__qualname__ = f"jax.numpy.{numpy_fn.__name__}"
    102   fn = jit(fn, inline=True)

    [... skipping hidden 7 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/lax/lax.py](https://l2x9jgvnp1-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240725-060215_RC00_655912111#) in broadcasting_shape_rule(name, *avals)
   1658         result_shape.append(non_1s[0])
   1659       else:
-> 1660         raise TypeError(f'{name} got incompatible shapes for broadcasting: '
   1661                         f'{", ".join(map(str, map(tuple, shapes)))}.')
   1662 

TypeError: mul got incompatible shapes for broadcasting: (4,), (12,).

It seems to me that the error is caused by a mismatch between the shapes of inverse_mass and dp_dx2 + dp_dx, where the former has shape (4,) while the latter has shape (12,). When I set the mass to be mass = 1.0, the code works fine. Would there be something wrong in defining the mass of the kinetic operator?

Besides, I am not sure if it is appropriate to set the mass to be infinity to ignore the kinetic term of that particle. Is there any better ways to do this?

Thank you so much in advance!

[gcarleo]

hi I think setting the mass to infinity is really not a good idea, because it will likely cause some numerical problem and will needlessly compute a laplacian of the wave function. it's best just to ignore the kinetic energy and to not add it to the Hamiltonian

[gcarleo]

in other words, you need to define the usual Born Oppenheimer Hamiltonian in which the nuclei are just seen as an external potential for the electrons

[PhilipVinc]

We do not support it easily.

Setting the mass to infinity is the right thing to do here. The inverse mass will be 0.

I suspect the error is because your squeeze is not keeping the right shapes that Jax expects.

try to specify that you only want to squeeze the axis=-1.
I fear you are sqeezing other axes as well and that confuses Jax

[gcarleo]

we shouldn't really be worried about supporting this, it's just wrong to use infinite mass for the nuclei here, the ground taste energy will depend on some random positions the nuclei have been initialized in by the monte Carlo !

[rdpdlv]

I try only to squeeze the axis=-1 but the same error comes out. But, yes I think I should just fix the nuclear positions. Thanks a lot!