{"id":28753256,"url":"https://github.com/google-deepmind/torax","last_synced_at":"2025-06-17T00:39:21.233Z","repository":{"id":237504720,"uuid":"767563789","full_name":"google-deepmind/torax","owner":"google-deepmind","description":"TORAX: Tokamak transport simulation in JAX","archived":false,"fork":false,"pushed_at":"2025-06-10T12:14:57.000Z","size":458748,"stargazers_count":430,"open_issues_count":35,"forks_count":57,"subscribers_count":19,"default_branch":"main","last_synced_at":"2025-06-10T13:20:33.587Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"https://torax.readthedocs.io","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","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":"docs/roadmap.rst","authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-05T14:20:02.000Z","updated_at":"2025-06-10T12:15:10.000Z","dependencies_parsed_at":"2024-05-20T19:26:26.891Z","dependency_job_id":"d7901f74-9d0d-4877-b458-73d40fefa16f","html_url":"https://github.com/google-deepmind/torax","commit_stats":null,"previous_names":["google-deepmind/torax"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/google-deepmind/torax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Ftorax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Ftorax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Ftorax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Ftorax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-deepmind","download_url":"https://codeload.github.com/google-deepmind/torax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Ftorax/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":"2025-06-17T00:39:20.250Z","updated_at":"2025-06-17T00:39:21.209Z","avatar_url":"https://github.com/google-deepmind.png","language":"Python","funding_links":[],"categories":["Tools","🔬 Domain-Specific Applications","Python","📄 Research Papers"],"sub_categories":["Simulation and Modeling Frameworks","🌌 Physics \u0026 Astronomy","Digital Twins \u0026 Simulation Acceleration"],"readme":"[![Unittests](https://github.com/google-deepmind/torax/actions/workflows/pytest.yml/badge.svg)](https://github.com/google-deepmind/torax/actions/workflows/pytest.yml)\n\n## What is TORAX?\n\nTORAX is a differentiable tokamak core transport simulator aimed for fast\nand accurate forward modelling, pulse-design, trajectory optimization, and\ncontroller design workflows. TORAX is written in Python using JAX, with\nthe following motivations:\n\n- Open-source and extensible, aiding with flexible workflow coupling\n- JAX provides auto-differentiation capabilities and code compilation for fast\nruntimes. Differentiability allows for gradient-based nonlinear PDE solvers for\nfast and accurate modelling, and for sensitivity analysis of simulation results\nto arbitrary parameter inputs, enabling applications such as trajectory\noptimization and data-driven parameter identification for semi-empirical models.\nAuto-differentiability allows for these applications to be easily extended with\nthe addition of new physics models, including ML-surrogates, or new parameter\ninputs, by avoiding the need to hand-derive Jacobians\n- Python-JAX is a natural framework for the coupling of ML-surrogates of physics\nmodels\n\nFor more comprehensive documentation, see our [readthedocs page](https://torax.readthedocs.io/).\n\nTORAX, at v1.0.0, has the following physics and numerics feature set:\n\n- Coupled PDEs of ion and electron heat transport, electron particle transport,\nand current diffusion\n    - Finite-volume-method discretization\n    - Multiple solver options: linear with Pereverzev-Corrigan terms and the\n    predictor-corrector method, nonlinear with Newton-Raphson, nonlinear with\n    optimization using the jaxopt library\n    - Poloidal flux boundary conditions based on either total current or loop\n    voltage at the last-closed-flux-surface\n- Ohmic power, ion-electron heat exchange, fusion power, Bremsstrahlung,\nimpurity line radiation, an [[ICRH ML-surrogate]](https://meetings.aps.org/Meeting/DPP24/Session/NP12.106) (as-yet covering limited regimes),\n- Neoclassical bootstrap current and conductivity with the analytical Sauter\n  model\n- Coupling to the [[QLKNN_7_11]](https://github.com/google-deepmind/fusion_surrogates/)\nand QLKNN10D\n[[van de Plassche et al, Phys. Plasmas 2020]](https://doi.org/10.1063/1.5134126)\nQuaLiKiz neural network surrogates for physics-based turbulent transport\n- General geometry, provided via CHEASE, FBT, or EQDSK equilibrium files\n    - For testing and demonstration purposes, a single CHEASE equilibrium file\n    is available in the `data/third_party/geo` directory. It corresponds to an\n    ITER hybrid scenario equilibrium based on simulations in\n    [[Citrin et al, Nucl. Fusion 2010]](https://doi.org/10.1088/0029-5515/50/11/115007),\n    and was obtained from [PINT](https://gitlab.com/qualikiz-group/pyntegrated_model).\n    A PINT license file is available in `data/third_party/geo`\n    - Time dependent geometry is supported by providing a time series of geometry files\n- Simple pedestal models using a local adaptive source to set internal boundary\nconditions\n- Sawtooth triggering and profile redistribution\n\nAdditional heating and current drive sources can be provided by user-provided\nanalytical models, or user-provided prescribed data.\n\nModel implementation was verified through direct comparison of simulation\noutputs to the RAPTOR\n[[Felici et al, Plasma Phys. Control. Fusion 2012]](https://iopscience.iop.org/article/10.1088/0741-3335/54/2/025002)\ntokamak transport simulator.\n\nThis is not an officially supported Google product.\n\n## Development roadmap\n\nA [development roadmap](https://torax.readthedocs.io/en/latest/roadmap.html#development-roadmap)\nis outlined in our readthedocs pages.\n\n## Installation guide\n\n### Requirements\n\nInstall Python 3.10 or greater.\n\nMake sure that tkinter is installed:\n\n```shell\nsudo apt-get install python3-tk\n```\n\n### How to install\n\n#### Prepare a virtual environment\n\nInstall Virtualenv (if not already installed):\n\n```shell\npip install --upgrade pip\n```\n\n```shell\npip install virtualenv\n```\n\nCreate and activate a virtual environment\n\n```shell\npython3 -m venv toraxvenv\nsource toraxvenv/bin/activate\n```\n\n#### Install from PyPI\n\nThe simplest way to use TORAX is to install it via PyPI:\n\n```shell\npip install torax\n```\n\nYou can check that everything runs as it should:\n\n```shell\nrun_torax --config=examples/basic_config.py --quit\n```\n\nIf you plan to help develop TORAX, you will need to clone the repository, see\n[contributing](https://torax.readthedocs.io/en/latest/contributing.html) and\n[contribution_tips](https://torax.readthedocs.io/en/latest/contribution_tips.html#contribution-tips)\n\n### Running an example\n\nThe following command will run TORAX using the default configuration file\n`examples/basic_config.py`.\n\n```shell\nrun_torax --config='examples/basic_config.py'\n```\n\nSimulation progress is shown by a terminal progress bar indicating the current\ntime and percentage completed.\n\nTo run more involved, ITER-inspired simulations, run:\n\n```shell\nrun_torax --config='examples/iterhybrid_rampup.py'\n```\n\nand\n\n```shell\nrun_torax --config='examples/iterhybrid_predictor_corrector.py'\n```\n\nAdditional configuration is provided through flags which append the above\nrun command, and environment variables. For example, for increased output\nverbosity, you can run with the `--log_progress` flag.\n\n```shell\nrun_torax  --config='examples/iterhybrid_rampup.py' --log_progress\n```\n\n#### Set environment variables\n\n##### Error checking\n\nIf true, error checking is enabled in internal routines. Used for debugging.\nDefault is false since it is incompatible with the persistent compilation cache.\n\n```shell\n$ export TORAX_ERRORS_ENABLED=\u003cTrue/False\u003e\n```\n\n##### JAX Compilation and Cache\n\nIf false, JAX does not compile internal TORAX functions. Used for debugging.\nDefault is true.\n\n```shell\n$ export TORAX_COMPILATION_ENABLED=\u003cTrue/False\u003e\n```\n\nThe following implements the JAX persistent cache and will cause jax to store\ncompiled programs to the filesystem, reducing recompilation time in some cases:\n\n```shell\n$ export JAX_COMPILATION_CACHE_DIR=\u003cpath of your choice, such as ~/jax_cache\u003e\n$ export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1\n$ export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0.0\n```\n\n#### Set flags\n`log_progress` - output simulation time, dt, and number of solver iterations\n(dt backtracking with nonlinear solver) carried out at each timestep.\n\n```shell\nrun_torax \\\n   --config=examples/iterhybrid_predictor_corrector.py \\\n   --log_progress\n```\n\n`output_dir` - overrides the default output directory\n\n```shell\nrun_torax \\\n   --config=examples/iterhybrid_predictor_corrector.py \\\n   --output_dir=/path/to/output/dir\n```\n\n### Post-simulation\n\nOnce complete, the time history of a simulation state and derived quantities\nis written to a timestamped file of the format `state_history_%Y%m%d_%H%M%S.nc`.\nThe output directory is user configurable, with a default `/tmp/torax_results`.\n\nTo take advantage of the in-memory (non-persistent) cache, the process does not\nend upon simulation termination. It is possible to modify the runtime_params,\ntoggle the `log_progress` and `plot_progress` flags, and rerun the simulation.\nThe following modifications are examples which will trigger a recompilation:\n\n- Grid resolution\n- Evolved variables (equations being solved)\n- Changing internal functions used, e.g. transport models, time_step_calculator.\n  source models, solver model, pedestal model\n\n### Cleaning up\n\nYou can exit the Python virtual env by deactivating it:\n\n```shell\ndeactivate\n```\n\n#### (Optional) Install QLKNN-hyper\n\nThe default transport model in TORAX is QLKNN_7_11, installed as a TORAX\ndependency. An alternative to QLKNN_7_11 is to use QLKNN-hyper-10D, also known\nas QLKNN10D [K.L. van de Plassche PoP 2020](https://doi.org/10.1063/1.5134126).\nQLKNN_7_11 is based on QuaLiKiz 2.8.1 which has an improved collision\noperator compared to the QLKNN10D training set. QLKNN_7_11 training data\nincludes impurity density gradients as an input feature and has better coverage\nof the near-LCFS region compared to QLKNN-hyper-10D. However, using QLKNN10D may\nstill be useful for specific use-cases, such as in benchmarks with other\nsimulators, or comparisons with QLKNN_7_11.\n\nDownload QLKNN-hyper dependencies:\n\n```shell\ngit clone https://gitlab.com/qualikiz-group/qlknn-hyper.git\n```\n\nTo use QLKNN10D, set `model_path` in the `transport` section of your TORAX\nconfig to the path of the cloned repository.\n\n## Simulation tutorials\n\nSee tutorials section in our\n[readthedocs pages](https://torax.readthedocs.io/en/latest/tutorials.html).\n\n## Citing TORAX\n\nA TORAX paper is [available on arXiv](https://arxiv.org/abs/2406.06718). Cite\nthis paper to cite TORAX:\n\n```\n@article{torax2024arxiv,\n  title={{TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX}},\n  author={Citrin, Jonathan and Goodfellow, Ian and Raju, Akhil and Chen, Jeremy\n  and Degrave, Jonas and Donner, Craig and Felici, Federico and Hamel, Philippe\n  and Huber, Andrea and Nikulin, Dmitry and Pfau, David and Tracey, Brendan, and\n  Riedmiller, Martin and Kohli, Pushmeet},\n  journal={arXiv preprint arXiv:2406.06718},\n  year={2024}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Ftorax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-deepmind%2Ftorax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Ftorax/lists"}