https://github.com/google-deepmind/torax Skip to content Navigation Menu Toggle navigation Sign in * Product + Actions Automate any workflow + Packages Host and manage packages + Security Find and fix vulnerabilities + Codespaces Instant dev environments + Copilot Write better code with AI + Code review Manage code changes + Issues Plan and track work + Discussions Collaborate outside of code Explore + All features + Documentation + GitHub Skills + Blog * Solutions For + Enterprise + Teams + Startups + Education By Solution + CI/CD & Automation + DevOps + DevSecOps Resources + Learning Pathways + White papers, Ebooks, Webinars + Customer Stories + Partners * Open Source + GitHub Sponsors Fund open source developers + The ReadME Project GitHub community articles Repositories + Topics + Trending + Collections * Pricing Search or jump to... Search code, repositories, users, issues, pull requests... Search [ ] Clear Search syntax tips Provide feedback We read every piece of feedback, and take your input very seriously. [ ] [ ] Include my email address so I can be contacted Cancel Submit feedback Saved searches Use saved searches to filter your results more quickly Name [ ] Query [ ] To see all available qualifiers, see our documentation. Cancel Create saved search Sign in Sign up You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session. Dismiss alert {{ message }} google-deepmind / torax Public * Notifications * Fork 0 * Star 39 * TORAX: Tokamak transport simulation in JAX torax.readthedocs.io License View license 39 stars 0 forks Branches Tags Activity Star Notifications * Code * Issues 1 * Pull requests 3 * Discussions * Actions * Projects 0 * Security * Insights Additional navigation options * Code * Issues * Pull requests * Discussions * Actions * Projects * Security * Insights google-deepmind/torax This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository. main BranchesTags Go to file Code Folders and files Last commit Last Name Name message commit date Latest commit History 125 Commits .github/workflows .github/workflows .vscode .vscode docs docs torax torax .gitignore .gitignore .pylintrc .pylintrc .readthedocs.yaml .readthedocs.yaml CONTRIBUTING.md CONTRIBUTING.md LICENSE LICENSE README.md README.md pyproject.toml pyproject.toml run_simulation_main.py run_simulation_main.py View all files Repository files navigation * README * License Unittests What is TORAX? TORAX is a differentiable tokamak core transport simulator aimed for fast and accurate forward modelling, pulse-design, trajectory optimization, and controller design workflows. TORAX is written in Python-JAX, with the following motivations: * Open-source and extensible, aiding with flexible workflow coupling * JAX provides auto-differentiation capabilities and code compilation for fast runtimes. Differentiability allows for gradient-based nonlinear PDE solvers for fast and accurate modelling, and for sensitivity analysis of simulation results to arbitrary parameter inputs, enabling applications such as trajectory optimization and data-driven parameter identification for semi-empirical models. Auto-differentiability allows for these applications to be easily extended with the addition of new physics models, or new parameter inputs, by avoiding the need to hand-derive Jacobians * Python-JAX is a natural framework for the coupling of ML-surrogates of physics models For more comprehensive documentation, see our readthedocs page. TORAX now has the following physics feature set: * Coupled PDEs of ion and electron heat transport, electron particle transport, and current diffusion + Finite-volume-method + Multiple solver options: linear with Pereverzev-Corrigan terms, nonlinear with Newton-Raphson, nonlinear with optimization using the jaxopt library * Ohmic power, ion-electron heat exchange, fusion power, bootstrap current with the analytical Sauter model * Time dependent boundary conditions and sources * Coupling to the QLKNN10D [van de Plassche et al, Phys. Plasmas 2020] QuaLiKiz-neural-network surrogate for physics-based turbulent transport * General geometry, provided via CHEASE equilibrium files + For testing and demonstration purposes, a single CHEASE equilibrium file is available in the data/geo directory. It corresponds to an ITER hybrid scenario equilibrium based on simulations in [Citrin et al, Nucl. Fusion 2010], and was obtained from PINT. A PINT license file is available in data/ geo. Additional heating and current drive sources can be provided by prescribed formulas, or user-provided analytical models. Model implementation was verified through direct comparison of simulation outputs to the RAPTOR [Felici et al, Plasma Phys. Control. Fusion 2012] tokamak transport simulator. This is not an officially supported Google product. Feature roadmap Short term development plans include: * Time dependent geometry * More flexible initial conditions * Implementation of forward sensitivity calculations w.r.t. control inputs and parameters * Implementation of persistent compilation cache for CPU * More extensive documentation and tutorials * Visualisation improvements Longer term desired features include: * Sawtooth model (Porcelli + reconnection) * Neoclassical tearing modes (modified Rutherford equation) * Radiation sinks + Cyclotron radiation + Bremsstrahlung + Line radiation * Neoclassical transport + multi-ion transport, with a focus on heavy impurities * IMAS coupling * Stationary-state solver * Momentum transport Contributions in line with the roadmap are welcome. In particular, TORAX is envisaged as a natural framework for coupling of various ML-surrogates of physics models. These could include surrogates for turbulent transport, neoclassical transport, heat and particle sources, line radiation, pedestal physics, and core-edge integration, MHD, among others. Installation guide Requirements Install Python 3.10 or greater. Make sure that tkinter is installed: sudo apt-get install python3-tk How to install Install virtualenv (if not already installed): pip install --upgrade pip pip install virtualenv Create a code directory where you will install the virtual env and other TORAX dependencies. mkdir /path/to/torax_dir && cd "$_" Where /path/to/torax_dir should be replaced by a path of your choice. Create a TORAX virtual env: python3 -m venv toraxvenv Activate the virtual env: source toraxvenv/bin/activate Download QLKNN dependencies: git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper It is recommended to automate the environment variable export. For example, if using bash, run: echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc The above command only needs to be run once on a given system. Download and install the TORAX codebase via http: git clone https://github.com/google-deepmind/torax.git or ssh (ensure that you have the appropriate SSH key uploaded to github). git clone git@github.com:google-deepmind/torax.git Enter the TORAX directory and pip install the dependencies. cd torax; pip install -e . If you want to install with the dev dependencies (useful for running pytest and installing pyink for lint checking), then run with the [dev]: cd torax; pip install -e .[dev] Optional: Install additional GPU support for JAX if your machine has a GPU: https://jax.readthedocs.io/en/latest/installation.html# supported-platforms Running an example The following command will run TORAX using the default configuration file examples/basic_config.py. python3 run_simulation_main.py \ --config='torax.examples.basic_config' --log_progress To run more involved, ITER-inspired simulations, run: python3 run_simulation_main.py \ --config='torax.examples.iterhybrid_rampup' --log_progress and python3 run_simulation_main.py \ --config='torax.examples.iterhybrid_predictor_corrector' --log_progress Additional configuration is provided through flags which append the above run command, and environment variables: Set environment variables Path to the QuaLiKiz-neural-network parameters. Note: if installation instructions above were followed, this may already be set. $ export TORAX_QLKNN_MODEL_PATH="" Path to the geometry file directory. This prefixes the path and filename provided in the geometry_file geometry constructor argument in the run config file. If not set, TORAX_GEOMETRY_DIR defaults to the relative path torax/data/third_party/geo. $ export TORAX_GEOMETRY_DIR="" If true, error checking is enabled in internal routines. Used for debugging. Default is false since it is incompatible with the persistent compilation cache. $ export TORAX_ERRORS_ENABLED= If false, JAX does not compile internal TORAX functions. Used for debugging. Default is true. $ export TORAX_COMPILATION_ENABLED= Set flags Output simulation time, dt, and number of stepper iterations (dt backtracking with nonlinear solver) carried out at each timestep. python3 run_simulation_main.py \ --config='torax.examples.iterhybrid_predictor_corrector' \ --log_progress Live plotting of simulation state and derived quantities. python3 run_simulation_main.py \ --config='torax.examples.iterhybrid_predictor_corrector' \ --plot_progress Combination of the above. python3 run_simulation_main.py \ --config='torax.examples.iterhybrid_predictor_corrector' \ --log_progress --plot_progress Post-simulation Once complete, the time history of a simulation state and derived quantities is written to state_history.nc. The output path is written to stdout. To take advantage of the in-memory (non-persistent) cache, the process does not end upon simulation termination. It is possible to modify the runtime_params, toggle the log_progress and plot_progress flags, and rerun the simulation. Only the following modifications will then trigger a recompilation: * Grid resolution * Evolved variables (equations being solved) * Changing internal functions used, e.g. transport model, or time_step_calculator Cleaning up You can get out of the Python virtual env by deactivating it: deactivate Simulation tutorials Under construction Citing TORAX A technical report is in preparation and will soon be made available for citing. About TORAX: Tokamak transport simulation in JAX torax.readthedocs.io Resources Readme License View license Activity Custom properties Stars 39 stars Watchers 13 watching Forks 0 forks Report repository Releases 1 v0.1.0 Latest May 1, 2024 Packages 0 No packages published Contributors 5 * * * * * Languages * Python 100.0% Footer (c) 2024 GitHub, Inc. Footer navigation * Terms * Privacy * Security * Status * Docs * Contact * Manage cookies * Do not share my personal information You can't perform that action at this time.