{"id":49602279,"url":"https://github.com/jorgemunozl/psiformer_torch","last_synced_at":"2026-05-04T09:36:36.499Z","repository":{"id":323051844,"uuid":"1022657614","full_name":"jorgemunozl/psiformer_torch","owner":"jorgemunozl","description":"self-attention neural torch ansatz for electronic structure calculations introduced in the paper  “A Self-attention Ansatz for ab initio Quantum Chemistry” from Google Deep Mind.","archived":false,"fork":false,"pushed_at":"2026-04-13T20:48:24.000Z","size":30472,"stargazers_count":12,"open_issues_count":10,"forks_count":1,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-04-13T22:35:04.288Z","etag":null,"topics":["attention-mechanism","computational-physics","computational-quantum-mechanics","google-deepmind","kfca","mlp","natural-gradient","quantum-mechanics","torch","training","transformer","wandb","wave-equation-solver"],"latest_commit_sha":null,"homepage":"https://jorgemunozl.github.io/portfolio/project_transformers","language":"TeX","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/jorgemunozl.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"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,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-07-19T14:47:33.000Z","updated_at":"2026-04-13T20:48:28.000Z","dependencies_parsed_at":"2026-01-08T21:05:47.089Z","dependency_job_id":null,"html_url":"https://github.com/jorgemunozl/psiformer_torch","commit_stats":null,"previous_names":["jorgemunozl/transformer_archi","jorgemunozl/psiformer_torch"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/jorgemunozl/psiformer_torch","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jorgemunozl%2Fpsiformer_torch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jorgemunozl%2Fpsiformer_torch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jorgemunozl%2Fpsiformer_torch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jorgemunozl%2Fpsiformer_torch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/jorgemunozl","download_url":"https://codeload.github.com/jorgemunozl/psiformer_torch/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jorgemunozl%2Fpsiformer_torch/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":32601761,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-03T22:12:39.696Z","status":"online","status_checked_at":"2026-05-04T02:00:06.625Z","response_time":58,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"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":["attention-mechanism","computational-physics","computational-quantum-mechanics","google-deepmind","kfca","mlp","natural-gradient","quantum-mechanics","torch","training","transformer","wandb","wave-equation-solver"],"created_at":"2026-05-04T09:36:36.391Z","updated_at":"2026-05-04T09:36:36.472Z","avatar_url":"https://github.com/jorgemunozl.png","language":"TeX","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n  \u003ch1\u003eSolving the Schrödinger equation with Transformers\u003c/h1\u003e\n  \u003cp\u003e\n    \u003ca href=\"https://pytorch.org/\"\u003e\u003cimg alt=\"PyTorch\" src=\"https://img.shields.io/badge/PyTorch-ready-EE4C2C?logo=pytorch\u0026logoColor=white\"\u003e\u003c/a\u003e\n    \u003ca href=\"https://www.python.org/\"\u003e\u003cimg alt=\"Python\" src=\"https://img.shields.io/badge/Python-3.10+-3776AB?logo=python\u0026logoColor=white\"\u003e\u003c/a\u003e\n    \u003ca href=\"https://wandb.ai/\"\u003e\u003cimg alt=\"Weights \u0026 Biases\" src=\"https://img.shields.io/badge/Weights_\u0026_Biases-Track_runs-FFBE00?logo=weightsandbiases\u0026logoColor=white\"\u003e\u003c/a\u003e\n    \u003ca href=\"https://deepmind.google/\"\u003e\u003cimg alt=\"DeepMind\" src=\"https://img.shields.io/badge/DeepMind-inspired-4285F4?logo=google\u0026logoColor=white\"\u003e\u003c/a\u003e\n    \u003ca href=\"#\"\u003e\u003cimg alt=\"Tests\" src=\"https://img.shields.io/badge/CI-tests_pending-lightgrey?logo=githubactions\u0026logoColor=white\"\u003e\u003c/a\u003e\n    \u003ca href=\"#\"\u003e\u003cimg alt=\"PRs Welcome\" src=\"https://img.shields.io/badge/PRs-welcome-brightgreen.svg\"\u003e\u003c/a\u003e\n    \u003ca href=\"LICENSE\"\u003e\u003cimg alt=\"License\" src=\"https://img.shields.io/badge/License-repo-default-lightgrey.svg\"\u003e\u003c/a\u003e\n  \u003c/p\u003e\n\u003c/div\u003e\n\nThis repository demonstrates an experimental approach to approximating solutions to the many-electron Schrödinger equation using transformer architectures. It contains reference code and utilities to build and run scale experiments, for more details on how this repository works and was created look [PsiFormer](src/PsiFormer_Torch%20Documentation.md) for the full documentation.  Watch the [paper](assets/research/main.pdf) for a detailed theoretical framework result,the [beamer](assets/research/beamer.pdf) for a more friendly explanation or the [poster](assets/research/poster.pdf) for a resume of this work.\n\n\u003e My sincerest thanks to Angel A. Flores and Alejandro D. Paredes for their feedback and helpful technical discussion. I would also like to thank the Faculty of Computer Science at UNI for the GPU used to make this project.\n\n\u003cdiv style=\"text-align:center;\"\u003e\u003cimg src=\"assets/research/poster.png\" width=\"400\"\u003e\u003c/div\u003e\n\n## Table of contents\n\n- [What this does](#what-this-does)\n- [Key ideas (short)](#key-ideas-short)\n- [Quick install (CPU-only)](#quick-install-cpu-only)\n- [Usage](#usage)\n- [Files and layout](#files-and-layout)\n- [Example workflow](#example-workflow)\n- [Tips and caveats](#tips-and-caveats)\n- [Reproducibility](#reproducibility)\n- [References and further reading](#references-and-further-reading)\n- [License](#license)\n\n\n## What this does\n\n- Uses transformer-style attention to model electronic wavefunction structure.\n- Provides utilities and example scripts to run  problems and evaluate model performance against simple quantum chemistry references.\n\nThis project is intended for research .It is not production-ready for large-scale quantum chemistry simulations but can be used to prototype ideas and compare modeling choices.\n\n## Key ideas (short)\n\n- Represent electronic configurations or basis expansions as sequences and\n\tlearn interactions using attention layers.\n- Use permutation-equivariant input encodings or ordered basis sequences to\n\tcapture antisymmetry constraints through learned modules and loss terms.\n- Train with physics-informed losses (energy expectation, cusp conditions,\n\tor density overlaps) and supervised or self-supervised pretraining.\n\n## Quick install (CPU-only)\n\nIf you use `uv` (fast installer), you can install a CPU-only PyTorch wheel and\nother dependencies like this:\n\n```bash\n# install uv if you don't have it\npip install -U uv\n\n# Sync dependencies\nuv sync\n```\n\nThen, install PyTorch CPU-only version:\n```bash\nuv pip install --index-url https://download.pytorch.org/whl/cpu torch\n```\n\n## Usage\n\n1. Prepare data / basis encodings for the target system.\n2. Edit `config.py` to set model size, number of heads, block (sequence)\n\t length, and training hyperparameters.\n3. Run the training/experiment script (example):\n\nReplace the example command above with your own runner or flags as needed.\n\n## Files and layout\n\n- `src/` — main source code\n\t- `main.py` — training / experiment entrypoint\n\t- `utils.py` — helper utilities, device selection, etc.\n\t- `config.py` — config and hyperparameters\n- `pyproject.toml` — Python project metadata\n- `template.tex` — report template (optional)\n\n## Example workflow\n\n1. Choose a small system (2–10 electrons) and a compact basis set.\n2. Build sequence encodings for orbitals/electron coordinates.\n3. Train the transformer to predict minimal-energy coefficients or approximate the wavefunction amplitude for sampled configurations.\n4. Evaluate energy and compare against reference (Hartree–Fock, small CI).\n\n## Tips and caveats\n\n- Enforcing antisymmetry exactly (Slater determinants) is nontrivial when\n\tusing sequence models; consider hybrid approaches (learn corrections to a\n\tSlater determinant) or include antisymmetry in the loss.\n- Transformers scale quadratically with sequence length; start with small\n\tbasis sizes and toy molecules.\n- Use physics-informed losses wherever possible to improve sample efficiency.\n\n\n## References and further reading\n\n- Vaswani et al., \"Attention Is All You Need\" (transformers).\n- A Self Attention Ansatz for Quantum Chemistry\n- Ab-Initio Solution of the Many-Electron Schrodinger Equation With Deep Neural Networks\n- QMC Torch Molecular Wave functions with Neural Components for energy and force calculations.","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjorgemunozl%2Fpsiformer_torch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fjorgemunozl%2Fpsiformer_torch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjorgemunozl%2Fpsiformer_torch/lists"}