from ferminet import base_config
from ferminet.utils import system
import ml_collections


def get_config() -> ml_collections.ConfigDict:
  """Returns config for running O3 with FermiNet."""
  cfg = base_config.default()
  # geometry in bohr.
  cfg.system.molecule = [
      system.Atom(symbol='O', coords=(0.0, 2.0859, -0.4319)),
      system.Atom(symbol='O', coords=(0.0, 0.0, 0.8638)),
      system.Atom(symbol='O', coords=(0.0, -2.0859, -0.4319)),
  ]
  cfg.system.electrons = (12, 12)
  return cfg
