{"id":13699112,"url":"https://github.com/google-deepmind/ferminet","last_synced_at":"2025-06-17T00:39:14.609Z","repository":{"id":39709839,"uuid":"301716934","full_name":"google-deepmind/ferminet","owner":"google-deepmind","description":"An implementation of the Fermionic Neural Network for ab-initio electronic structure calculations","archived":false,"fork":false,"pushed_at":"2025-06-02T16:31:03.000Z","size":500,"stargazers_count":768,"open_issues_count":3,"forks_count":149,"subscribers_count":32,"default_branch":"main","last_synced_at":"2025-06-03T06:14:48.066Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google-deepmind.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2020-10-06T12:21:06.000Z","updated_at":"2025-06-02T16:31:06.000Z","dependencies_parsed_at":"2023-09-07T20:34:48.062Z","dependency_job_id":"9c9a40d3-f140-4126-9700-05325c491295","html_url":"https://github.com/google-deepmind/ferminet","commit_stats":null,"previous_names":["google-deepmind/ferminet","deepmind/ferminet"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/google-deepmind/ferminet","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fferminet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fferminet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fferminet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fferminet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-deepmind","download_url":"https://codeload.github.com/google-deepmind/ferminet/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fferminet/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":260268635,"owners_count":22983601,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-08-02T19:00:58.565Z","updated_at":"2025-06-17T00:39:14.600Z","avatar_url":"https://github.com/google-deepmind.png","language":"Python","funding_links":[],"categories":["🔬 Domain-Specific Applications","Wavefunction methods (ML-WFT)","Python"],"sub_categories":["🌌 Physics \u0026 Astronomy"],"readme":"# FermiNet: Fermionic Neural Networks\n\nFermiNet is a neural network for learning highly accurate ground state\nwavefunctions of atoms and molecules using a variational Monte Carlo approach.\n\nThis repository contains an implementation of the algorithm and experiments\nfirst described in \"Ab-Initio Solution of the Many-Electron Schroedinger\nEquation with Deep Neural Networks\", David Pfau, James S. Spencer, Alex G de G\nMatthews and W.M.C. Foulkes, Phys. Rev. Research 2, 033429 (2020), along with\nsubsequent research and developments.\n\nWARNING: This is a research-level release of a JAX implementation and is under\nactive development. The original TensorFlow implementation can be found in the\n`tf` branch.\n\n## Installation\n\n`pip install -e .` will install all required dependencies. This is best done\ninside a [virtual environment](https://docs.python-guide.org/dev/virtualenvs/).\n\n```shell\nvirtualenv ~/venv/ferminet\nsource ~/venv/ferminet/bin/activate\npip install -e .\n```\n\nIf you have a GPU available (highly recommended for fast training), then you can\ninstall JAX with CUDA support, using e.g.:\n\n```shell\npip install --upgrade jax jaxlib==0.1.57+cuda110 -f\nhttps://storage.googleapis.com/jax-releases/jax_releases.html\n```\n\nNote that the jaxlib version must correspond to the existing CUDA installation\nyou wish to use. Please see the\n[JAX documentation](https://github.com/jax-ml/jax#installation) for more\ndetails.\n\nThe tests are easiest run using pytest:\n\n```shell\npip install -e '.[testing]'\npython -m pytest\n```\n\n## Usage\n\nferminet uses the `ConfigDict` from\n[ml_collections](https://github.com/google/ml_collections) to configure the\nsystem. A few example scripts are included under `ferminet/configs/`. These are\nmostly for testing so may need additional settings for a production-level\ncalculation.\n\n```shell\nferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100\n```\n\nor\n\n```shell\npython3 ferminet/main.py --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100\n```\n\nwill train FermiNet to find the ground-state wavefunction of the Li atom using a\nbatch size of 1024 MCMC configurations (\"walkers\" in variational Monte Carlo\nlanguage), and 100 iterations of pretraining (the default of 1000 is overkill\nfor such a small system). The system and hyperparameters can be controlled by\nmodifying the config file or (better, for one-off changes) using flags. See the\n[ml_collections](https://github.com/google/ml_collections)' documentation for\nfurther details on the flag syntax. Details of all available config settings are\nin `ferminet/base_config.py`.\n\nOther systems can easily be set up, by creating a new config file or `ferminet`,\nor writing a custom training script. For example, to run on the H2 molecule, you\ncan create a config file containing:\n\n```python\nfrom ferminet import base_config\nfrom ferminet.utils import system\n\n# Settings in a config files are loaded by executing the the get_config\n# function.\ndef get_config():\n  # Get default options.\n  cfg = base_config.default()\n  # Set up molecule\n  cfg.system.electrons = (1,1)\n  cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]\n\n  # Set training hyperparameters\n  cfg.batch_size = 256\n  cfg.pretrain.iterations = 100\n\n  return cfg\n```\n\nand then run it using\n\n```\nferminet --config /path/to/h2_config.py\n```\n\nor equivalently write the following script (or execute it interactively):\n\n```python\nimport sys\n\nfrom absl import logging\nfrom ferminet.utils import system\nfrom ferminet import base_config\nfrom ferminet import train\n\n# Optional, for also printing training progress to STDOUT.\n# If running a script, you can also just use the --alsologtostderr flag.\nlogging.get_absl_handler().python_handler.stream = sys.stdout\nlogging.set_verbosity(logging.INFO)\n\n# Define H2 molecule\ncfg = base_config.default()\ncfg.system.electrons = (1,1)  # (alpha electrons, beta electrons)\ncfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]\n\n# Set training parameters\ncfg.batch_size = 256\ncfg.pretrain.iterations = 100\n\ntrain.train(cfg)\n```\n\nAlternatively, you can directly pass in a PySCF ['Molecule'](http://pyscf.org).\nYou can create PySCF Molecules with the following:\n\n```python\nfrom pyscf import gto\nmol = gto.Mole()\nmol.build(\n    atom = 'H  0 0 1; H 0 0 -1',\n    basis = 'sto-3g', unit='bohr')\n```\n\nOnce you have this molecule, you can pass it directly into the configuration by\nrunning\n\n```python\nfrom ferminet import base_config\nfrom ferminet import train\n\n# Add H2 molecule\ncfg = base_config.default()\ncfg.system.pyscf_mol = mol\n\n# Set training parameters\ncfg.batch_size = 256\ncfg.pretrain.iterations = 100\n\ntrain.train(cfg)\n```\n\nNote: to train on larger atoms and molecules with large batch sizes, multi-GPU\nparallelisation is essential. This is supported via JAX's\n[pmap](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap).\nMultiple GPUs will be automatically detected and used if available.\n\n### Inference\n\nAfter training, it is useful to run calculations of the energy and other\nobservables over many time steps with the parameters fixed to accumulate\nlow-variance estimates of physical quantities. To do this, just re-run the same\ncommand used for training with the flag `--config.optim.optimizer 'none'`. Make\nsure that either the value of `cfg.log.save_path` is the same, or that the value\nof `cfg.log.restore_path` is set to the value of `cfg.log.save_path` from the\noriginal training run.\n\nIt can also be useful to accumulate statistics about observables at inference\ntime which were not included in the original training run. Spin magnitude,\ndipole moments and density matrices can be tracked by adding\n`--config.observables.s2`, `--config.observables.dipole` and\n`--config.observables.density` to the command line if they are not set to true\nin the config file.\n\n## Excited States\n\nExcited state properties of systems can be calculated using either the\n[Natural Excited States for VMC (NES-VMC) algorithm](https://arxiv.org/abs/2308.16848)\nor an [ensemble penalty method](https://arxiv.org/abs/2312.00693).\nTo enable the calculation of `k` states of a system, simply set\n`cfg.system.states=k` in the config file. By default, NES-VMC is used, but to\nenable the ensemble penalty method, add `cfg.optim.objective='vmc_overlap'` to\nthe config. NES-VMC does not have any parameters to set, but the ensemble\npenalty method has a free choice of weights on the energies and overlap penalty,\nwhich can be set in `cfg.optim.overlap`. If the weights are not set for the\nenergies in the config, they are automatically set to 1/k for state k. We have\nfound that NES-VMC is generally more accurate than the ensemble penalty method,\nbut include both for completeness. Config files for all experiments from the\npaper which introduced NES-VMC can be found in the folder `configs/excited`, and\nall experiments can be tested (on smaller networks) by running\n`tests/excited_test.py`.\n\n## Output\n\nThe results directory contains `train_stats.csv` which contains the local energy\nand MCMC acceptance probability for each iteration, and the `checkpoints`\ndirectory, which contains the checkpoints generated during training. When\ncomputing observables of excited states or the density matrix for the ground\nstate, `.npy` files are also saved to the same folder. A single NumPy array is\nsaved for every iteration of optimization into the same file. An example Colab\nnotebook analyzing these outputs is given in\n`notebooks/excited_states_analysis.ipynb`.\n\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/google-deepmind/ferminet/blob/main/ferminet/notebooks/excited_states_analysis.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"(Open in Colab!)\"/\u003e\u003c/a\u003e\n\n## Giving Credit\n\nIf you use this code in your work, please cite the associated papers. The\ninitial paper details the architecture and results on a range of systems:\n\n```\n@article{pfau2020ferminet,\n  title={Ab-Initio Solution of the Many-Electron Schr{\\\"o}dinger Equation with Deep Neural Networks},\n  author={D. Pfau and J.S. Spencer and A.G. de G. Matthews and W.M.C. Foulkes},\n  journal={Phys. Rev. Research},\n  year={2020},\n  volume={2},\n  issue = {3},\n  pages={033429},\n  doi = {10.1103/PhysRevResearch.2.033429},\n  url = {https://link.aps.org/doi/10.1103/PhysRevResearch.2.033429}\n}\n```\n\nand a NeurIPS Workshop Machine Learning and Physics paper describes the JAX\nimplementation:\n\n```\n@misc{spencer2020better,\n  title={Better, Faster Fermionic Neural Networks},\n  author={James S. Spencer and David Pfau and Aleksandar Botev and W. M.C. Foulkes},\n  year={2020},\n  eprint={2011.07125},\n  archivePrefix={arXiv},\n  primaryClass={physics.comp-ph},\n  url={https://arxiv.org/abs/2011.07125}\n}\n```\n\nThe PsiFormer architecture is detailed in an ICLR 2023 paper:\n\n```\n@misc{vonglehn2023psiformer,\n  title={A Self-Attention Ansatz for Ab-initio Quantum Chemistry},\n  author={Ingrid von Glehn and James S Spencer and David Pfau},\n  journal={ICLR},\n  year={2023},\n}\n```\n\nPeriodic boundary conditions were originally introduced in a Physical Review\nLetters article:\n\n```\n@article{cassella2023discovering,\n  title={Discovering quantum phase transitions with fermionic neural networks},\n  author={Cassella, Gino and Sutterud, Halvard and Azadi, Sam and Drummond, ND and Pfau, David and Spencer, James S and Foulkes, W Matthew C},\n  journal={Physical review letters},\n  volume={130},\n  number={3},\n  pages={036401},\n  year={2023},\n  publisher={APS}\n}\n```\n\nWasserstein QMC (thanks to Kirill Neklyudov) is described in a NeurIPS 2023\narticle:\n\n```\n@article{neklyudov2023wasserstein,\n  title={Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schr{\\\"o}dinger Equation},\n  author={Neklyudov, Kirill and Nys, Jannes and Thiede, Luca and Carrasquilla, Juan and Liu, Qiang and Welling, Max and Makhzani, Alireza},\n  journal={NeurIPS},\n  year={2023}\n}\n```\n\nNatural excited states was introduced in this article, which is also the first\npaper from our group using pseudopotentials\n\n```\n@article{pfau2024excited,\n  title={Accurate computation of quantum excited states with neural networks},\n  author={Pfau, David and Axelrod, Simon and Sutterud, Halvard and von Glehn, Ingrid and Spencer, James S},\n  journal={Science},\n  volume={385},\n  number={6711},\n  pages={eadn0137},\n  year={2024},\n  url={https://doi.org/10.1126/science.adn0137},\n}\n```\n\nThis repository can be cited using:\n\n```\n@software{ferminet_github,\n  author = {James S. Spencer, David Pfau and FermiNet Contributors},\n  title = {{FermiNet}},\n  url = {http://github.com/deepmind/ferminet},\n  year = {2020},\n}\n```\n\n## Disclaimer\n\nThis is not an official Google product.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fferminet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-deepmind%2Fferminet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fferminet/lists"}