Add tests and configs for all experiments from Pfau, Axelrod, Sutteru…

…d, von Glehn and Spencer (2024)

PiperOrigin-RevId: 643440665
Change-Id: Idb855b4023906d8a73baf09dc57e17ee31ee954d

README.md

@@ -186,7 +186,10 @@ penalty method has a free choice of weights on the energies and overlap penalty,
 which can be set in `cfg.optim.overlap`. If the weights are not set for the
 energies in the config, they are automatically set to 1/k for state k. We have
 found that NES-VMC is generally more accurate than the ensemble penalty method,
-but include both for completeness.
+but include both for completeness. Config files for all experiments from the
+paper which introduced NES-VMC can be found in the folder `configs/excited`, and
+all experiments can be tested (on smaller networks) by running
+`tests/excited_test.py`.
 
 ## Output
 

ferminet/base_config.py

@@ -262,13 +262,12 @@ def default() -> ml_collections.ConfigDict:
           },
           # Only used if network_type is 'psiformer'.
           'psiformer': {
-              # PsiFormer architecture: von Glehn, Spencer, Pfau,
-              # arXiv:2211.13672 (2022), accepted to ICLR 2023.
-              'num_layers': 2,
+              # PsiFormer architecture: von Glehn, Spencer, Pfau, ICLR 2023.
+              'num_layers': 4,
               'num_heads': 4,
               'heads_dim': 64,
               'mlp_hidden_dims': (256,),
-              'use_layer_norm': False,
+              'use_layer_norm': True,
           },
           # Config common to all architectures.
           'determinants': 16,  # Number of determinants.

ferminet/configs/atom.py

@@ -20,7 +20,7 @@
 import ml_collections
 
 
-def _adjust_nuclear_charge(cfg):
+def adjust_nuclear_charge(cfg):
   """Sets the molecule, nuclear charge electrons for the atom.
 
   Note: function name predates this logic but is kept for compatibility with
@@ -70,6 +70,6 @@ def get_config():
   cfg.system.spin_polarisation = ml_collections.FieldReference(
       None, field_type=int)
   with cfg.ignore_type():
-    cfg.system.set_molecule = _adjust_nuclear_charge
+    cfg.system.set_molecule = adjust_nuclear_charge
     cfg.config_module = '.atom'
   return cfg

ferminet/configs/excited/atoms.py

@@ -0,0 +1,36 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config to reproduce Fig. 1 from Pfau et al. (2024)."""
+
+from ferminet import base_config
+from ferminet.jax.configs import atom
+from ferminet.jax.configs.excited import presets
+import ml_collections
+
+
+def get_config() -> ml_collections.ConfigDict:
+  """Returns config for running generic atoms with qmc."""
+  cfg = base_config.default()
+  cfg.system.atom = ''
+  cfg.system.charge = 0
+  cfg.system.delta_charge = 0.0
+  cfg.system.states = 5
+  cfg.system.spin_polarisation = ml_collections.FieldReference(
+      None, field_type=int)
+  cfg.pretrain.iterations = 10_000
+  cfg.update_from_flattened_dict(presets.excited_states)
+  with cfg.ignore_type():
+    cfg.system.set_molecule = atom.adjust_nuclear_charge
+  return cfg

ferminet/configs/excited/benzene.py

@@ -0,0 +1,73 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config to reproduce Fig. 6 from Pfau et al. (2024)."""
+
+from ferminet import base_config
+from ferminet.jax.configs.excited import presets
+from ferminet.utils import system
+import ml_collections
+
+
+_GEOM = """C    0.00000000            2.63144965            0.00000000
+C   -2.27890225            1.31572483            0.00000000
+C   -2.27890225           -1.31572483            0.00000000
+C    0.00000000           -2.63144965            0.00000000
+C    2.27890225           -1.31572483            0.00000000
+C    2.27890225            1.31572483            0.00000000
+H   -4.04725813            2.33668557            0.00000000
+H   -4.04725813           -2.33668557            0.00000000
+H   -0.00000000           -4.67337115            0.00000000
+H    4.04725813           -2.33668557            0.00000000
+H    4.04725813            2.33668557            0.00000000
+H    0.00000000            4.67337115            0.00000000""".split('\n')
+
+
+def finalise(
+    experiment_config: ml_collections.ConfigDict) -> ml_collections.ConfigDict:
+  """Returns the experiment config with the molecule commpletely set."""
+  molecule = []
+  for atom in _GEOM:
+    element, x, y, z = atom.split()
+    coords = [float(xx) for xx in (x, y, z)]
+    molecule.append(system.Atom(symbol=element, coords=coords, units='bohr'))
+
+  if not experiment_config.system.electrons:  # Don't override if already set
+    nelectrons = int(sum(atom.charge for atom in molecule))
+    na = nelectrons // 2
+    experiment_config.system.electrons = (na, nelectrons - na)
+  experiment_config.system.molecule = molecule
+  return experiment_config
+
+
+def get_config() -> ml_collections.ConfigDict:
+  """Returns config for running generic atoms with qmc."""
+  cfg = base_config.default()
+  cfg.system.charge = 0
+  cfg.system.delta_charge = 0.0
+  cfg.system.states = 6
+  cfg.system.spin_polarisation = ml_collections.FieldReference(
+      None, field_type=int)
+  # While this value was used in the paper, it can be lowered.
+  cfg.pretrain.iterations = 100_000
+  cfg.mcmc.blocks = 4
+  # While this envelope was used in the paper, it can be replaced with the
+  # default 'isotropic' envelope without any noticeable change in the results.
+  cfg.network.envelope_type = 'bottleneck'
+  cfg.network.num_envelopes = 32
+  cfg.update_from_flattened_dict(presets.excited_states)
+  with cfg.ignore_type():
+    cfg.system.set_molecule = finalise
+
+  return cfg

ferminet/configs/excited/carbon_dimer.py

@@ -0,0 +1,51 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config to reproduce Fig. 3 from Pfau et al. (2024)."""
+
+from ferminet import base_config
+from ferminet.jax.configs.excited import presets
+from ferminet.utils import system
+import ml_collections
+
+
+def finalise(
+    experiment_config: ml_collections.ConfigDict) -> ml_collections.ConfigDict:
+  """Returns the experiment config with the molecule set."""
+  # Equilibrium bond length is 1.244 Angstrom
+  bond_length = experiment_config.system.equilibrium_multiple * 1.244 * 1.88973
+  experiment_config.system.molecule = [
+      system.Atom('C', coords=(0, 0, bond_length / 2)),
+      system.Atom('C', coords=(0, 0, -bond_length / 2))]
+  return experiment_config
+
+
+def get_config() -> ml_collections.ConfigDict:
+  """Returns config for running generic atoms with qmc."""
+  cfg = base_config.default()
+  cfg.system.charge = 0
+  cfg.system.delta_charge = 0.0
+  cfg.system.molecule_name = 'C2'
+  cfg.system.states = 8
+  cfg.system.spin_polarisation = ml_collections.FieldReference(
+      None, field_type=int)
+  cfg.system.units = 'bohr'
+  cfg.system.electrons = (6, 6)
+  cfg.pretrain.iterations = 100_000
+  cfg.optim.iterations = 100_000
+  cfg.update_from_flattened_dict(presets.psiformer)
+  cfg.update_from_flattened_dict(presets.excited_states)
+  with cfg.ignore_type():
+    cfg.system.set_molecule = finalise
+  return cfg

ferminet/configs/excited/double_excitation.py

@@ -0,0 +1,107 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config to reproduce Fig. 5 from Pfau et al. (2024)."""
+
+from ferminet import base_config
+from ferminet.jax.configs.excited import presets
+from ferminet.utils import system
+import ml_collections
+
+
+# Geometries are taken from Loos, Boggio-Pasqua, Scemama, Caffarel and
+# Jacquemin, JCTC (2019). The units listed in the supplemental material are
+# inconsistent: nitrosomethane and glyoxal are in Bohr, the rest are Angstrom.
+_SYSTEMS = {
+    'nitrosomethane': """C -1.78426612 0.00000000 -1.07224050
+N -0.00541753 0.00000000 1.08060391
+O 2.18814985 0.00000000 0.43452135
+H -0.77343975 0.00000000 -2.86415606
+H -2.97471478 1.66801808 -0.86424584
+H -2.97471478 -1.66801808 -0.86424584""".split('\n'),
+    'butadiene': """C 1.740343 0.616556 0.00000000
+C -1.740343 -0.616556 0.00000000
+C 0.397343 0.616556 0.00000000
+C -0.397343 -0.616556 0.00000000
+H 0.126346 -1.577069 0.00000000
+H -0.126346 1.577069 0.00000000
+H 2.279054 1.568725 0.00000000
+H -2.279054 -1.568725 0.00000000
+H 2.279054 -0.335614 0.00000000
+H -2.279054 0.335614 0.00000000""".split('\n'),
+    'glyoxal': """C 1.21360282 0.75840215 0.00000000
+C -1.21360282 -0.75840215 0.00000000
+O 3.25581408 -0.26453186 0.00000000
+O -3.25581408 0.26453186 0.00000000
+H 0.96135276 2.81883243 0.00000000
+H -0.96135276 -2.81883243 0.00000000""".split('\n'),
+    'tetrazine': """C  0.00000000  0.00000000   1.26054332
+C  0.00000000  0.00000000  -1.26054332
+N  0.00000000  1.19421138   0.66133002
+N  0.00000000 -1.19421138   0.66133002
+N  0.00000000  1.19421138  -0.66133002
+N  0.00000000 -1.19421138  -0.66133002
+H  0.00000000  0.00000000   2.33817427
+H  0.00000000  0.00000000  -2.33817427""".split('\n'),
+    'cyclopentadienone': """C  0.00000000  0.00000000   0.76853878
+C  0.00000000  1.19974276  -0.13448057
+C  0.00000000 -1.19974276  -0.13448057
+C  0.00000000  0.74909075  -1.39624830
+C  0.00000000 -0.74909075  -1.39624830
+O  0.00000000  0.00000000   1.98144505
+H  0.00000000  2.21416694   0.22305399
+H  0.00000000 -2.21416694   0.22305399
+H  0.00000000  1.34284493  -2.29584273
+H  0.00000000 -1.34284493  -2.29584273""".split('\n')
+}
+
+
+def finalise(
+    experiment_config: ml_collections.ConfigDict) -> ml_collections.ConfigDict:
+  """Returns the experiment config with the molecule set."""
+  system_name = experiment_config.system.molecule_name
+  geom = _SYSTEMS[system_name]
+  units = 'bohr' if system_name in ['gloyxal', 'nitrosomethane'] else 'angstrom'
+
+  molecule = []
+  for atom in geom:
+    element, x, y, z = atom.split()
+    coords = [float(xx) for xx in (x, y, z)]
+    molecule.append(system.Atom(symbol=element,
+                                coords=coords,
+                                units=units))
+
+  if not experiment_config.system.electrons:  # Don't override if already set
+    nelectrons = int(sum(atom.charge for atom in molecule))
+    na = nelectrons // 2
+    experiment_config.system.electrons = (na, nelectrons - na)
+  experiment_config.system.molecule = molecule
+  return experiment_config
+
+
+def get_config() -> ml_collections.ConfigDict:
+  """Returns config for running generic atoms with qmc."""
+  cfg = base_config.default()
+  cfg.system.charge = 0
+  cfg.system.delta_charge = 0.0
+  cfg.system.molecule_name = 'nitrosomethane'
+  cfg.system.states = 6
+  cfg.system.spin_polarisation = ml_collections.FieldReference(
+      None, field_type=int)
+  cfg.pretrain.iterations = 50_000
+  cfg.mcmc.blocks = 2
+  cfg.update_from_flattened_dict(presets.excited_states)
+  with cfg.ignore_type():
+    cfg.system.set_molecule = finalise
+  return cfg