About configs

[qetuo2]

Hi all,

I try to see how FermiNet can be run by following some examples in the configs. One of the examples I tried is the following

import sys
from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train
from ferminet.pbc import envelopes
import numpy as np
import jax.numpy as jnp
from jax.config import config

def _sc_lattice_vecs(rs: float, nelec: int) -> np.nd
```array:
  """Returns simple cubic lattice vectors with Wigner-Seitz radius rs."""
  volume = (4 / 3) * np.pi * (rs**3) * nelec
  length = volume**(1 / 3)
  return length * np.eye(3)

logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)

cfg = base_config.default()
cfg.system.electrons = (27, 0)
cfg.system.molecule = [system.Atom("X", (0., 0., 0.))]
cfg.pretrain.method = None
lattice = _sc_lattice_vecs(1.0, sum(cfg.system.electrons))
kpoints = envelopes.make_kpoints(lattice, cfg.system.electrons)
cfg.system.make_local_energy_fn = "ferminet.pbc.hamiltonian.local_energy"
cfg.system.make_local_energy_kwargs = {"lattice": lattice, "heg": True}
cfg.network.make_feature_layer_fn = (
    "ferminet.pbc.feature_layer.make_pbc_feature_layer")
cfg.network.make_feature_layer_kwargs = {
    "lattice": lattice,
    "include_r_ae": False
}
cfg.network.make_envelope_fn = (
    "ferminet.pbc.envelopes.make_multiwave_envelope")
cfg.network.make_envelope_kwargs = {"kpoints": kpoints}
cfg.network.full_det = False
cfg.batch_size = 512 
cfg.pretrain.iterations = 0
train.train(cfg)

However, it seems to run this code, one needs to first modify line line396 in train.py "charges" to "charges.shape[0]" to avoid error.

Even though the code can be run after that, it reports energy 28.6515, which seems to be much lower than the number mentioned in arxiv. 2202.05183 (which should be 1.2615x27=34.0605).

In addition, if one tries to run lattice = _sc_lattice_vecs(30, sum(cfg.system.electrons)), it prints out "Initial Energy: 5296.6323 E_h" and then after two steps just becomes nan.

May I know whether the above setup is proper for the related tasks?

[jsspencer]

A small batch size makes the estimate of the gradients and Fisher information matrix noisier and hence NaNs more likely.

You are comparing calculations using a simple cubic cell and a body-centred cubic cell. It is not surprising that they give very different energies for small system sizes.

[qetuo2]

Hi James,

Thanks for the reply. Actually I also tried the following:

import sys
from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train
from ferminet.pbc import envelopes
import numpy as np
import jax.numpy as jnp
from jax.config import config

def _sc_lattice_vecs(rs: float, nelec: int) -> np.ndarray:
  """Returns simple cubic lattice vectors with Wigner-Seitz radius rs."""
  volume = (4 / 3) * np.pi * (rs**3) * nelec
  length = volume**(1 / 3)
  return length * np.eye(3)

# Optional, for also printing training progress to STDOUT.
# If running a script, you can also just use the --alsologtostderr flag.
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)

# Get default options.
cfg = base_config.default()
cfg.system.electrons = (7, 7)
# A ghost atom at the origin defines one-electron coordinate system.
# Element 'X' is a dummy nucleus with zero charge
cfg.system.molecule = [system.Atom("X", (0., 0., 0.))]
# Pretraining is not currently implemented for systems in PBC
cfg.pretrain.method = None

lattice = _sc_lattice_vecs(1.0, sum(cfg.system.electrons))
kpoints = envelopes.make_kpoints(lattice, cfg.system.electrons)

cfg.system.make_local_energy_fn = "ferminet.pbc.hamiltonian.local_energy"
cfg.system.make_local_energy_kwargs = {"lattice": lattice, "heg": True}
cfg.network.make_feature_layer_fn = (
    "ferminet.pbc.feature_layer.make_pbc_feature_layer")
cfg.network.make_feature_layer_kwargs = {
    "lattice": lattice,
    "include_r_ae": False
}
cfg.network.make_envelope_fn = (
    "ferminet.pbc.envelopes.make_multiwave_envelope")
cfg.network.make_envelope_kwargs = {"kpoints": kpoints}
cfg.network.full_det = False 
train.train(cfg)

If I understand correctly, this should use the default 4096 examples and have the same 14 electrons simple cubic setup. However, it also reports NaN after 100ish steps optimization. May I know whether one should use more samples?

In addition, may I know whether it is also proper to use cfg.network.network_type='psiformer' in the above code. It seems that the optimization then reports energy of 7.173, which is also quite lower than the number reported in the paper (around 8.0). Is this setup correct?

[jsspencer]

1. More likely caused by the initalisation. The code defaults to initialising based upon a Gaussian centred at the origin; the HEG calculations used a uniform initialisation over the unit cell (easily implemented but not included in the pull request implementing the HEG Hamiltonian.
2. Yes, though note the defaults do not match exactly the network configurations used in the ICLR paper.
3. What do you mean by 'reports an energy of'? After how many iterations? What is the variance? Did you do any analysis to calculate the expected energy from an inference run? It is not possible to say whether your calculations are correct from the information provided.