{"id":28573414,"url":"https://github.com/pennylaneai/generative-quantum-states","last_synced_at":"2026-04-01T21:49:56.636Z","repository":{"id":103532235,"uuid":"564948486","full_name":"PennyLaneAI/generative-quantum-states","owner":"PennyLaneAI","description":"Official code for the paper \"Predicting Properties of Quantum Systems with Conditional Generative Models\"","archived":false,"fork":false,"pushed_at":"2022-12-02T08:17:20.000Z","size":12728,"stargazers_count":36,"open_issues_count":0,"forks_count":10,"subscribers_count":3,"default_branch":"main","last_synced_at":"2026-03-28T01:53:28.668Z","etag":null,"topics":["generative-model","machine-learning","quantum-computing","quantum-machine-learning","quantum-many-body","transformer"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","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/PennyLaneAI.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"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":"2022-11-11T22:20:47.000Z","updated_at":"2025-10-02T14:37:17.000Z","dependencies_parsed_at":null,"dependency_job_id":"242bada8-3453-40e7-ba74-83c854b4f371","html_url":"https://github.com/PennyLaneAI/generative-quantum-states","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/PennyLaneAI/generative-quantum-states","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PennyLaneAI%2Fgenerative-quantum-states","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PennyLaneAI%2Fgenerative-quantum-states/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PennyLaneAI%2Fgenerative-quantum-states/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PennyLaneAI%2Fgenerative-quantum-states/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/PennyLaneAI","download_url":"https://codeload.github.com/PennyLaneAI/generative-quantum-states/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PennyLaneAI%2Fgenerative-quantum-states/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":31292562,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-04-01T21:15:39.731Z","status":"ssl_error","status_checked_at":"2026-04-01T21:15:34.046Z","response_time":53,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"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":["generative-model","machine-learning","quantum-computing","quantum-machine-learning","quantum-many-body","transformer"],"created_at":"2025-06-10T21:17:49.163Z","updated_at":"2026-04-01T21:49:56.623Z","avatar_url":"https://github.com/PennyLaneAI.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Predicting Properties of Quantum Systems with Conditional Generative Models\n\nThis repo contains the code accompanying the paper \"_Predicting Properties of Quantum Systems with Conditional\nGenerative Models_\" [[arXiv](https://arxiv.org/abs/2211.16943)].\n\n![Our proposed method](resources/figures/workflow.png \"Proposed Method\")\n\nIf you find this repo useful for your research, please consider citing our paper:\n\n```bibtex\n@article{cond-generative-quantum-states,\n  title={Predicting Properties of Quantum Systems with Conditional Generative Models},\n  author={Wang, Haoxiang and Weber, Maurice and Izaac, Josh and Lin, Cedric Yen-Yu},\n  journal={arXiv preprint arXiv:2211.16943},\n  year={2022},\n}\n```\n\n## Requirements\n\n### Software\n\n+ Python (tested with python 3.8 \u0026 3.9)\n    + Please see `requirements.txt` for the required Python packages (if you are using pip or Conda, you can\n      run `pip install -r requirements.txt`)\n    + Pytorch-Geometric (PyG): Please refer to\n      its [official documentation](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) for\n      installation.\n        + PyG is used in our experiments on 2D Heisenberg models (our conditional generative model has a GNN module in\n          that case).\n    + Jax: Please refer to its [official documentation](https://github.com/google/jax#installation) for installation (a\n      CPU version is sufficient).\n        + `jax` and `neural-tangents` (a Jax-based package) are only used when comparing our method with Neural Tangent\n          Kernel (used by Hsin-Yuan Huang et al. in their [Science 2022](https://arxiv.org/abs/2106.12627) paper). You\n          can skip this installation if you do not intend to do the comparison.\n+ Julia (tested with version 1.7 \u0026 1.8)\n    + The Julia language can be downloaded from the [official website](https://julialang.org/downloads/) (If you are\n      using pip or Conda as your python package manager, you can also install Julia with\n      the [Jill](https://github.com/johnnychen94/jill.py) package).\n    + Run `julia install_pacakges.jl` in the `rydberg/` folder to install all necessary Julia packages.\n        + If you want to use GPU for simulation, you need to install the `CUDA` and `Adapt` packages in addition (\n          see `rydberg/README.md` for details).\n\n### Hardware\n\n+ To train the models, we recommend using a GPU with at least 16GB memory (otherwise, you may need to use a smaller\n  batch size).\n\n+ For inference with our pre-trained models, a CPU is sufficient (may need to decrease the batch size to accommodate the\n  memory).\n\n## Data \u0026 Pre-trained Models\n\n**Download** - All simulation data \u0026 pre-trained models can be downloaded from this [**Google Drive\nfolder**](https://drive.google.com/drive/folders/1AWXaXjMyiFeBsjXr6EV05VXzdG1SGGqo?usp=sharing).\nIt contains two sub-folders, `data/` and `logs/`. Please put them in the root of this repo after downloading.\n\n### Simulation Data\n\nIn the paper, we use the following datasets:\n\n+ Heisenberg Models (`data/2d_heisenberg_data/`)\n    + Our simulated data with [PennyLane](https://github.com/PennyLaneAI/pennylane) (generated using our\n      code `heisenberg_generate_data.py`)\n    + [Hsin-Yuan Huang](https://hsinyuan-huang.github.io/)'s simulated data, available in\n      [this GitHub repo](https://github.com/hsinyuan-huang/provable-ml-quantum).\n+ Rydberg Atom Systems (`data/rydberg/`)\n    + Our simulated data via [Bloqade.jl](https://queracomputing.github.io/Bloqade.jl/) (generated using our code\n      in `rydberg`)\n        + 1D Rydberg-atom chains of 13,15,...,33 atoms.\n        + 2D Rydberg-atom square lattices of 5x5 atoms (with various adiabatic evolution times for the ground-state\n          preparation).\n\n#### Processing Open-source Heisenberg Data\n\nFirst clone the [provable-ml-quantum](https://github.com/hsinyuan-huang/provable-ml-quantum) repo. Since our scripts\nexpect the data to be organized slightly differently, and to create train and test splits, we provide the\n`reorganize_heisenberg_data.py` script. Run it as follows:\n\n```bash\npython reorganize_heisenberg_data.py \u003cpath/to/provable-ml-quantum/heisenberg/data\u003e\n```\n\nThis will reorganize the data and save it to `data/2d_heisenberg_data/{rows}x{cols}/{split}/`.\n\n### Pre-trained Models\n\nWe trained conditional transformers over the simulation data. If you want to directly load pre-trained models when using\nour scripts, please put them in `logs/`.\n\n#### Heisenberg Models\n\nIn Sec. III-A of the paper, we conducted various numerical experiments over Heisenberg models of various sizes.\nThe trained models are saved in `logs/2d_heisenberg_checkpoints/`.\n\n#### Rydberg Atoms\n\nIn Sec. III-B of the paper, we conducted 4 machine learning experiments over Rydberg atom systems, and we provide\npre-trained models for them in different sub-folders. Notice that we trained models for a different number of\niterations, and marked them with suffixes like `...{#iterations}.pth` (e.g., `...100k.pth` implies training with 100k\niterations).\n\n+ Sec. III-B(1) - _Predicting quantum phases_\n    + `logs/rydberg_1D/`: models trained on 1D Rydberg-atom chains of 31 atoms.\n    + `logs/rydberg_2D/`: models trained on 2D Rydberg-atom chains of 25 atoms (prepared with adiabatic evolution of 3.0\n      μs)\n+ Sec. III-B(2) - _Predicting phases of larger quantum systems_\n    + `logs/rydberg_1D-size/`: models trained on 1D Rydberg-atom chains of 13,15,...,27 atoms.\n+ Sec. III-B(3) - _Predicting phases of ground states prepared with longer adiabatic evolution time_\n    + `logs/rydberg_2D-time/`: models trained on 2D Rydberg-atom square lattices of 5x5 atoms, which are prepared with\n      adiabatic evolution of 0.4, 0.6, 0.8, 1.0 μs.\n\n## Tutorials\n\nIn this repo, we also provide tutorials that are intended to communicate the ideas of our paper in a more interactive\nmanner.\n\n+ **Heisenberg Models.** The notebook `Tutorial-2D-Heisenberg.ipynb` first introduces the 2D random Heisenberg model,\n  and\n  shows how correlations and entanglement entropies of its ground states can be calculated. In the second part, we show\n  how a pretrained conditional generative model can be used to generate samples corresponding to a ground state of a\n  random\n  Heisenberg Hamiltonian. Finally, in the last part, we walk you through the code used to train a transformer model,\n  conditioned on the coupling graph, which was embedded by a graph convolutional neural network.\n\n+ **Rydberg Atom Systems.** The notebook `Tutorial-Rydberg-1D.ipynb` introduces 1D Rydberg atom systems, loads\n  simulation data of such systems with 31 atoms (generated via Julia in `rydberg/`), and plots the phase diagram of 1D\n  Rydberg-atom chains based on the simulation data. Then, we demonstrate how to train a conditional generative model\n  over\n  a subset of the simulation data (you can also directly load a pre-trained model from `logs/rydberg_1D/`). With the\n  trained model, you can generate measurements for any new Rydberg-atom system of the same kind, and predict its phases\n  (Disordered, $Z_2$-ordered, or $Z_3$-ordered). In this tutorial, we use the trained model to predict the entire phase\n  diagram of 1D Rydberg-atom chains (31 atoms), and compare it with the predicted phase diagrams by several kernel\n  methods (introduced by Hsin-Yuan Huang et al. in their [Science 2022](https://arxiv.org/abs/2106.12627) paper) - you\n  can see\n  that the prediction by our method is much more accurate that these kernel methods.\n\n## **Conditional generative models for 2D anti-ferromagnetic random Heisenberg models**\n\n### Dataset\n\n+ **Generate your own data**\n\nYou can generate your own dataset consisting of randomized Pauli measurements using the\nscript `generate_heisenberg_data.py`.\nThe script first samples random coupling matrices of the specified dimensions, calculates the corresponding ground state\nusing exact diagonalization and then samples random Pauli measurements using the\n[Classical Shadow implementation in PennyLane](https://docs.pennylane.ai/en/stable/code/qml_shadows.html).\nFor example, if you want to create a dataset for a `4x4` lattice, with 40 train Hamiltonians and 10 test Hamiltonians,\nwith 500 measurements for each Hamiltonian, you can run\n\n```shell\npython heisenberg_generate_data.py --nh_train 40 --nh_test 10 --rows 4 --cols 4 --shots 500 \n```\n\nThe resulting data will be saved under `data/2d_heisenberg/4x4` with two sub-folders `train` and `test` for the\nrespective data splits.\n\n+ **Open source data**\n\nFor systems with 20 and more qubits, we have used publicly available data from\n[this GitHub repo](https://github.com/hsinyuan-huang/provable-ml-quantum) by Hsin-Yuan Huang. These datasets were\nsimulated using DMRG. Or, you can directly download the data from\nour [Google Drive folder](https://drive.google.com/drive/folders/1XiAAX9aoYzEbpUYzyBCZjz6PTK8BH5a7?usp=sharing).\n\n### Training models\n\nTo train a conditional generative model for the 2D anti-ferromagnetic random Heisenberg model, you can use the\nscript `heisenberg_train_transformer.py`. An example command is\n\n```shell\npython heisenberg_train_transformer.py --train-size 4x4 --train-samples 500\n```\n\nwhich uses the default hyperparameters used throughout the paper and trains a model on all the Hamiltonians\nsampled in the previous step. Note that this is compute-intensive and should be run with GPU support.\nIf you want to quickly test on CPUs, you can set the flag `--hamiltonians 1` to use data from only a single Hamiltonian.\n\n### Generating samples from trained models\n\nTo generate samples from a trained conditional generative model, you can use the script\n`heisenberg_sample_transformer.py`. The flag `--results-dir` indicates the directory pointing to the run package where\nthe results of a training run have been saved. Note that this should be the root of the results directory and not the\n`checkpoints` folder, as this script requires access to flags set during training and saved to the `args.json` file.\n\n### Evaluating properties with classical shadows\n\nEvaluating properties of ground states of the Heisenberg model can be done using the script\n`heisenberg_evaluate_properties.py`. Note that this requires you to have completed the previous steps and generated\nsamples from a trained model using the script `heisenberg_sample_transformer.py`. Similar to the previous step, the flag\n`--results-root` indicates the directory pointing to the run package where the results of a training run have been\nsaved.\nThe flag `--snapshots` indicates the number of samples (i.e., snapshots) which will be used to estimate the correlation\nfunctions and the entanglement entropies. We use the classical shadow implementation in PennyLane to compute these\nproperties.\n\n## **Conditional generative models for Rydberg atom systems**\n\n### Dataset\n\nYou can generate your own dataset using our Julia simulation code in `rydberg/`. Or, you can download the data from\nthis [Google Drive folder](https://drive.google.com/drive/folders/1XiAAX9aoYzEbpUYzyBCZjz6PTK8BH5a7?usp=sharing).\n\n### Machine Learning Experiments\n\nTo reproduce our experiments on Rydberg atom systems, you can check out the following Jupyter notebooks\n(with detailed descriptions inside):\n\n+ Sec. III-B(1) - _Predicting quantum phases_\n    + 1D Lattices: `Tutorial-Rydberg-1D.ipynb`\n    + 2D Lattices: `notebooks/Experiment-Rydberg-2D.ipynb`\n        + Furthermore, we provide a demo, `notebooks/Demo-Rydberg-2D-Phases.ipynb`, to visualize phases of 2D\n          Rydberg-atom lattices (in both the real space and the Fourier space).\n+ Sec. III-B(2) - _Predicting phases of larger quantum systems_\n    + `notebooks/Experiment-Rydberg-1D-LargerSystems.ipynb`\n+ Sec. III-B(3) - _Predicting phases of ground states prepared with longer adiabatic evolution time_\n    + `notebooks/Experiment-Rydberg-2D-LongerEvolution.ipynb`\n\n## Acknowledgements\n\nWe found several very helpful codebases when building this repo, and we sincerely thank their authors:\n\n+ [Bloqade.jl](https://github.com/QuEraComputing/Bloqade.jl): A Julia package for neutral-atom simulations.\n+ [provable-ml-quantum](https://github.com/hsinyuan-huang/provable-ml-quantum): Official implementation for [*Provably\n  efficient machine learning for quantum many-body problems*](https://arxiv.org/abs/2106.12627).\n+ PennyLane Tutorials:\n    + [Classical Shadows](https://pennylane.ai/qml/demos/tutorial_classical_shadows.html): A Python re-implementation of\n      Hsin-Yuan Huang's [classical shadows code (C++)](https://github.com/hsinyuan-huang/predicting-quantum-properties).\n    + [Machine learning for quantum many-body problems](https://pennylane.ai/qml/demos/tutorial_ml_classical_shadows.html):\n      A Python implementation of [*Provably efficient machine learning for quantum many-body\n      problems*](https://arxiv.org/abs/2106.12627).\n+ [The Annotated Transformer](http://nlp.seas.harvard.edu/annotated-transformer/): An annotated version of the paper\n  [*Attention Is All You Need*](https://arxiv.org/abs/1706.03762) by Vaswani et al., in the form of a line-by-line\n  implementation.\n\n## Version History\n\nInitial release (v1.0): Nov 2022\n\n## Team\n\nCurrent maintainers:\n\n+ [Haoxiang Wang](https://haoxiang-wang.github.io) ([@Haoxiang__Wang](https://twitter.com/haoxiang__wang),\n  hwang264@illinois.edu),\n+ [Maurice Weber](https://mauriceweber.github.io) ([@mauriceweberq](https://twitter.com/mauriceweberq),\n  maurice.weber@inf.ethz.ch)\n\n## License\n\nApache 2.0 \n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpennylaneai%2Fgenerative-quantum-states","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpennylaneai%2Fgenerative-quantum-states","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpennylaneai%2Fgenerative-quantum-states/lists"}