├── .github ├── dependabot.yaml └── workflows │ ├── cpu-tests.yaml │ ├── pre-commit-cron-updater.yaml │ └── pre-commit.yaml ├── .gitignore ├── .pre-commit-config.yaml ├── .readthedocs.yaml ├── LICENSE ├── README.md ├── adept ├── __init__.py ├── _base_.py ├── _lpse2d │ ├── __init__.py │ ├── core │ │ ├── __init__.py │ │ ├── driver.py │ │ ├── epw.py │ │ ├── laser.py │ │ ├── trapper.py │ │ └── vector_field.py │ ├── datamodel.py │ ├── helpers.py │ └── modules │ │ ├── __init__.py │ │ ├── base.py │ │ └── driver.py ├── _tf1d │ ├── __init__.py │ ├── datamodel.py │ ├── modules.py │ ├── solvers │ │ ├── __init__.py │ │ ├── pushers.py │ │ └── vector_field.py │ ├── storage.py │ └── train_damping.py ├── _version.py ├── _vlasov1d │ ├── __init__.py │ ├── datamodel.py │ ├── gamma_func_for_sg.nc │ ├── helpers.py │ ├── modules.py │ ├── solvers │ │ ├── __init__.py │ │ ├── pushers │ │ │ ├── __init__.py │ │ │ ├── field.py │ │ │ ├── fokker_planck.py │ │ │ └── vlasov.py │ │ └── vector_field.py │ └── storage.py ├── electrostatic.py ├── lpse2d.py ├── sh2d │ ├── __init__.py │ ├── runner.py │ ├── solvers │ │ ├── __init__.py │ │ ├── field.py │ │ ├── fokker_planck.py │ │ ├── tridiagonal.py │ │ └── vlasov.py │ └── utils │ │ ├── __init__.py │ │ ├── helpers.py │ │ └── save.py ├── tf1d.py ├── utils.py ├── vfp1d │ ├── __init__.py │ ├── base.py │ ├── fokker_planck.py │ ├── helpers.py │ ├── impact.py │ ├── oshun.py │ ├── storage.py │ └── vector_field.py ├── vlasov1d.py ├── vlasov1d2v │ ├── __init__.py │ ├── gamma_func_for_sg.nc │ ├── helpers.py │ ├── integrator.py │ ├── pushers │ │ ├── __init__.py │ │ ├── field.py │ │ ├── fokker_planck.py │ │ └── vlasov.py │ └── storage.py └── vlasov2d │ ├── __init__.py │ ├── gamma_func_for_sg.nc │ ├── helpers.py │ ├── pushers │ ├── __init__.py │ ├── field.py │ ├── fokker_planck.py │ ├── time.py │ └── vlasov.py │ ├── solver │ ├── __init__.py │ └── tridiagonal.py │ └── storage.py ├── configs ├── envelope-2d │ ├── damping.yaml │ ├── epw.yaml │ ├── reflection.yaml │ └── tpd.yaml ├── es1d │ ├── damping.yaml │ ├── epw.yaml │ ├── es1d.yaml │ └── wp.yaml ├── sh2d │ └── landau_damping.yaml ├── tf-1d │ ├── damping.yaml │ ├── epw.yaml │ ├── tf1d.yaml │ └── wp.yaml ├── vfp-1d │ ├── epp-short.yaml │ └── hotspot.yaml ├── vlasov-1d │ ├── bump-on-tail.yaml │ ├── epw.yaml │ ├── nlepw-ic.yaml │ ├── srs.yaml │ ├── twostream.yaml │ └── wavepacket.yaml ├── vlasov-1d2v │ └── epw.yaml └── vlasov-2d │ └── base.yaml ├── docs ├── Makefile ├── make.bat ├── requirements.txt └── source │ ├── ADEPTModule.rst │ ├── adept-logo.png │ ├── api.rst │ ├── conf.py │ ├── dev_guide.rst │ ├── ergoExo.rst │ ├── faq.rst │ ├── index.rst │ ├── solvers.rst │ ├── solvers │ ├── datamodels │ │ ├── lpse2d.rst │ │ └── vlasov1d.rst │ ├── lpse2d.rst │ ├── vfp1d.rst │ └── vlasov1d1v.rst │ ├── tests.rst │ ├── usage.rst │ └── usage │ ├── initialization.rst │ ├── lpse2d.rst │ ├── tf1d.rst │ ├── vlasov1d.rst │ ├── vlasov1d2v.rst │ └── vlasov2d2v.rst ├── env.yaml ├── env_gpu.yaml ├── pyproject.toml ├── pytest.ini ├── requirements-cpu.txt ├── requirements.txt ├── ruff.toml ├── run.py └── tests ├── conftest.py ├── test_base ├── configs │ └── example.yaml └── test_ergoExo.py ├── test_lpse2d ├── __init__.py ├── configs │ ├── __init__.py │ ├── epw.yaml │ ├── resonance_search.yaml │ └── tpd.yaml ├── test_epw_frequency.py └── test_tpd_threshold.py ├── test_tf1d ├── __init__.py ├── configs │ ├── resonance.yaml │ ├── resonance_search.yaml │ └── vlasov_comparison.yaml ├── test_against_vlasov.py ├── test_landau_damping.py ├── test_resonance.py ├── test_resonance_search.py └── vlasov-reference │ ├── all-fields-kx.nc │ ├── all-fields.nc │ └── config.yaml ├── test_vfp1d ├── epp-short.yaml └── test_kappa_eh.py ├── test_vlasov1d ├── __init__.py ├── configs │ └── resonance.yaml ├── test_absorbing_wave.py └── test_landau_damping.py └── test_vlasov2d ├── __init__.py ├── configs └── damping.yaml └── test_landau_damping.py /.github/dependabot.yaml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | updates: 4 | - package-ecosystem: github-actions 5 | directory: "/" 6 | schedule: 7 | interval: "weekly" 8 | day: "monday" 9 | timezone: "Europe/London" 10 | open-pull-requests-limit: 10 11 | 12 | groups: 13 | actions: 14 | patterns: 15 | - "*" 16 | -------------------------------------------------------------------------------- /.github/workflows/cpu-tests.yaml: -------------------------------------------------------------------------------- 1 | # This is a basic workflow to help you get started with Actions 2 | 3 | name: cpu-tests 4 | 5 | # Controls when the action will run. 6 | on: 7 | pull_request: 8 | 9 | push: 10 | branches: 11 | - main 12 | 13 | # Allows you to run this workflow manually from the Actions tab 14 | workflow_dispatch: 15 | 16 | # Cancel in-progress tests on new commits to the same branch 17 | concurrency: 18 | group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} 19 | cancel-in-progress: true 20 | 21 | # A workflow run is made up of one or more jobs that can run sequentially or in parallel 22 | jobs: 23 | test-adept: 24 | runs-on: ubuntu-latest 25 | steps: 26 | - name: set up Git repository 27 | uses: actions/checkout@v4 28 | 29 | - name: Set up Python 3.11 30 | uses: actions/setup-python@v5 31 | with: 32 | python-version: '3.11' 33 | 34 | - name: Install dependencies 35 | run: | 36 | pip install --upgrade pip 37 | pip install .[dev] 38 | 39 | - name: Test with pytest 40 | run: | 41 | pytest tests/test_base tests/test_lpse2d tests/test_vlasov1d tests/test_vfp1d 42 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit-cron-updater.yaml: -------------------------------------------------------------------------------- 1 | name: Auto update pre-commit hooks 2 | 3 | on: 4 | workflow_dispatch: # Allows manual trigger 5 | 6 | schedule: 7 | - cron: '0 0 * * 1' # 12AM only on Mondays 8 | 9 | jobs: 10 | auto-update-hooks: 11 | runs-on: ubuntu-latest 12 | steps: 13 | - name: Checkout code 14 | uses: actions/checkout@v4 15 | 16 | - uses: actions/setup-python@v5 17 | with: 18 | python-version-file: "pyproject.toml" 19 | 20 | - name: Install pre-commit 21 | run: pip install pre-commit 22 | 23 | - name: Run pre-commit autoupdate 24 | run: pre-commit autoupdate 25 | 26 | - name: Detect if changes were made 27 | id: git-diff 28 | run: | 29 | changes=false 30 | git diff --exit-code || changes=true 31 | echo "update_done=$changes" >> $GITHUB_OUTPUT 32 | 33 | - name: Run pre-commit 34 | id: pre-commit 35 | if: steps.git-diff.outputs.update_done == 'true' 36 | run: | 37 | # Run twice so we only fail if there are non-fixable issues 38 | rc=0 39 | pre-commit run --all-files || true 40 | pre-commit run --all-files > /tmp/pre-commit.log || rc=$? 41 | 42 | # Add log as step output 43 | echo "pre-commit-log< $GITHUB_OUTPUT 44 | cat /tmp/pre-commit.log >> $GITHUB_OUTPUT 45 | echo "EOF" >> $GITHUB_OUTPUT 46 | 47 | # Add linting outcome as step output 48 | if [ $rc -eq 0 ]; then 49 | echo "pre-commit-outcome=success" >> $GITHUB_OUTPUT 50 | else 51 | echo "pre-commit-outcome=failure" >> $GITHUB_OUTPUT 52 | fi 53 | 54 | # Distinguish 3 cases: 55 | # 1. No changes were made -> do nothing (steps below are skipped) 56 | # 2. Changes were made and pre-commit ran successfully -> create PR 57 | # 3. Changes were made but pre-commit failed -> create PR with draft status 58 | 59 | - name: Create Pull Request (all good) 60 | if: steps.pre-commit.outputs.pre-commit-outcome == 'success' 61 | uses: peter-evans/create-pull-request@v7 62 | with: 63 | token: ${{ secrets.GITHUB_TOKEN }} 64 | commit-message: Update pre-commit hooks 65 | title: ✅ Update pre-commit hooks 66 | branch: _bot/update-precommit 67 | draft: false 68 | body: | 69 | Pre-commit hooks have been updated successfully without conflicts. 70 | 71 | - name: Create Pull Request (conflicts) 72 | if: steps.pre-commit.outputs.pre-commit-outcome == 'failure' 73 | uses: peter-evans/create-pull-request@v7 74 | with: 75 | token: ${{ secrets.GITHUB_TOKEN }} 76 | commit-message: Update pre-commit hooks 77 | title: ⚠️ Update pre-commit hooks [review required] 78 | branch: _bot/update-precommit 79 | draft: true 80 | body: | 81 | Pre-commit is unable to automatically update the hooks due to unresolvable conflicts. 82 | Please review the changes and merge manually. 83 | 84 | Log: 85 | ``` 86 | ${{ steps.pre-commit.outputs.pre-commit-log }} 87 | ``` 88 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit.yaml: -------------------------------------------------------------------------------- 1 | name: Code linting 2 | 3 | on: 4 | pull_request: 5 | 6 | push: 7 | branches: 8 | - main 9 | 10 | jobs: 11 | pre-commit: 12 | runs-on: ubuntu-latest 13 | 14 | steps: 15 | - uses: actions/checkout@v4 16 | - uses: actions/setup-python@v5 17 | with: 18 | python-version-file: "pyproject.toml" 19 | - uses: pre-commit/action@v3.0.1 20 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # mflow 129 | *mlruns/ 130 | 131 | # Pyre type checker 132 | .pyre/ 133 | 134 | # hatchling-generated version file 135 | _version.py 136 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v5.0.0 4 | hooks: 5 | - id: check-added-large-files 6 | additional_dependencies: [--isolated] 7 | args: ["--maxkb=2000"] 8 | # Add exceptions here, as a regex 9 | exclude: "" 10 | 11 | - id: check-json 12 | additional_dependencies: [--isolated] 13 | 14 | - id: check-toml 15 | additional_dependencies: [--isolated] 16 | 17 | - id: check-yaml 18 | additional_dependencies: [--isolated] 19 | 20 | - id: detect-private-key 21 | additional_dependencies: [--isolated] 22 | 23 | - id: end-of-file-fixer 24 | additional_dependencies: [--isolated] 25 | 26 | - id: trailing-whitespace 27 | additional_dependencies: [--isolated] 28 | 29 | 30 | - repo: https://github.com/henryiii/validate-pyproject-schema-store 31 | rev: 2025.05.12 32 | hooks: 33 | - id: validate-pyproject 34 | 35 | - repo: https://github.com/astral-sh/ruff-pre-commit 36 | # Ruff version. 37 | rev: v0.11.10 38 | hooks: 39 | # Run the linter. 40 | - id: ruff 41 | args: [--fix] 42 | types_or: [pyi, python, jupyter] 43 | # Ignore global python configuration for private registry and install hooks from public index 44 | # Add for each hook 45 | # Reference: https://github.com/pre-commit/pre-commit/issues/1454#issuecomment-1816328894 46 | additional_dependencies: [--isolated] 47 | # Run the formatter. 48 | - id: ruff-format 49 | types_or: [pyi, python, jupyter] 50 | additional_dependencies: [--isolated] 51 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | # .readthedocs.yaml 2 | # Read the Docs configuration file 3 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details 4 | 5 | # Required 6 | version: 2 7 | 8 | # Set the OS, Python version and other tools you might need 9 | build: 10 | os: ubuntu-22.04 11 | tools: 12 | python: "3.10" 13 | # You can also specify other tool versions: 14 | # nodejs: "19" 15 | # rust: "1.64" 16 | # golang: "1.19" 17 | 18 | # Build documentation in the "docs/" directory with Sphinx 19 | sphinx: 20 | configuration: docs/source/conf.py 21 | 22 | # Optionally build your docs in additional formats such as PDF and ePub 23 | # formats: 24 | # - pdf 25 | # - epub 26 | 27 | # Optional but recommended, declare the Python requirements required 28 | # to build your documentation 29 | # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html 30 | python: 31 | install: 32 | - requirements: docs/requirements.txt 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 ergodic 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ADEPT 2 | 3 | ![Docs](https://readthedocs.org/projects/adept/badge/?version=latest) 4 | ![Tests](https://github.com/ergodicio/adept/actions/workflows/cpu-tests.yaml/badge.svg) 5 | 6 | ![ADEPT](./docs/source/adept-logo.png) 7 | 8 | ADEPT is an **A**utomatic **D**ifferentation **E**nabled **P**lasma **T**ransport code. 9 | 10 | ## Installation 11 | ### User 12 | `pip install git+https://github.com/ergodicio/adept.git` 13 | 14 | ### Developer 15 | ### Conda 16 | 1. Install `conda` (we recommend `mamba`) 17 | 2. `mamba env create -f env.yaml` or `mamba env create -f env_gpu.yaml` 18 | 3. `mamba activate adept` 19 | 20 | ### pip 21 | 1. `python3 -m venv venv` 22 | 2. `source venv/bin/activate` 23 | 3. `pip3 install -r requirements.txt` 24 | 25 | ## Docs 26 | https://adept.readthedocs.io/en/latest/ 27 | 28 | ## Examples 29 | https://github.com/ergodicio/adept-notebooks 30 | 31 | There are other ways to use ADEPT, notably as part of a neural network training pipeline that leverages differentiable simulation. In reference [1], neural networks are trained to learn forcing functions that 32 | drive the system towards previously unseen behavior. In reference [2], neural networks are trained to help bridge the micro-macro physics gap in multiphysics simulations. 33 | 34 | ## Usage 35 | `python3 run.py --cfg {config_path}` without the `.yaml` extension 36 | 37 | This runs the simulation defined in the config and stores the output to an `mlflow` server. 38 | 39 | Unless you have separately deployed an `mlflow` server somewhere, it simply writes files using the mlflow specification to the current working directory. 40 | 41 | To access and visualize the results, it is easiest to use the UI from the browser by typing `mlflow ui` in the command line from the same directory. 42 | 43 | 44 | ## Contributing guide 45 | The contributing guide is in development but for now, just make an issue / pull request and we can go from there :) 46 | 47 | ## Citation 48 | If you are using this package for your research, please cite the following 49 | 50 | ``` 51 | A. Joglekar and A. Thomas, “ADEPT - automatic differentiation enabled plasma transport,” 52 | ICML - SynS & ML Workshop (https://syns-ml.github.io/2023/contributions/), 2023 53 | 54 | ``` 55 | 56 | ## References 57 | [1] A. S. Joglekar & A. G. R. Thomas. "Unsupervised discovery of nonlinear plasma physics using differentiable kinetic simulations." J. Plasma Phys. 88, 905880608 (2022). 58 | 59 | [2] A. S. Joglekar and A. G. R. Thomas, “Machine learning of hidden variables in multiscale fluid simulation,” Mach. Learn.: Sci. Technol., vol. 4, no. 3, p. 035049, Sep. 2023, doi: 10.1088/2632-2153/acf81a. 60 | -------------------------------------------------------------------------------- /adept/__init__.py: -------------------------------------------------------------------------------- 1 | from ._base_ import ADEPTModule, ergoExo # noqa: I001 2 | from . import lpse2d, vlasov1d 3 | -------------------------------------------------------------------------------- /adept/_lpse2d/__init__.py: -------------------------------------------------------------------------------- 1 | from .helpers import calc_threshold_intensity as calc_threshold_intensity 2 | from .modules import ( 3 | ArbitraryDriver, 4 | GaussianDriver, 5 | LorentzianDriver, 6 | UniformDriver, 7 | ) 8 | from .modules import ( 9 | BaseLPSE2D as BaseLPSE2D, 10 | ) 11 | -------------------------------------------------------------------------------- /adept/_lpse2d/core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_lpse2d/core/__init__.py -------------------------------------------------------------------------------- /adept/_lpse2d/core/driver.py: -------------------------------------------------------------------------------- 1 | from jax import numpy as jnp 2 | 3 | from adept._base_ import get_envelope 4 | 5 | 6 | class Driver: 7 | def __init__(self, cfg): 8 | self.xax = cfg["grid"]["x"] 9 | self.yax = cfg["grid"]["y"] 10 | 11 | def __call__(self, this_pulse: dict, current_time: jnp.float64): 12 | kk = this_pulse["k0"] 13 | ww = this_pulse["w0"] 14 | t_L = this_pulse["tc"] - this_pulse["tw"] * 0.5 15 | t_R = this_pulse["tc"] + this_pulse["tw"] * 0.5 16 | t_wL = this_pulse["tr"] 17 | t_wR = this_pulse["tr"] 18 | x_L = this_pulse["xc"] - this_pulse["xw"] * 0.5 19 | x_R = this_pulse["xc"] + this_pulse["xw"] * 0.5 20 | x_wL = this_pulse["xr"] 21 | x_wR = this_pulse["xr"] 22 | 23 | y_L = this_pulse["yc"] - this_pulse["yw"] * 0.5 24 | y_R = this_pulse["yc"] + this_pulse["yw"] * 0.5 25 | y_wL = this_pulse["yr"] 26 | y_wR = this_pulse["yr"] 27 | 28 | envelope_t = get_envelope(t_wL, t_wR, t_L, t_R, current_time) 29 | envelope_x = get_envelope(x_wL, x_wR, x_L, x_R, self.xax) 30 | envelope_y = get_envelope(y_wL, y_wR, y_L, y_R, self.yax) 31 | 32 | return ( 33 | envelope_t 34 | * envelope_x[:, None] 35 | * envelope_y[None, :] 36 | * jnp.abs(kk) 37 | * this_pulse["a0"] 38 | * jnp.exp(-1j * (kk * self.xax[:, None] - ww * current_time)) 39 | ) 40 | -------------------------------------------------------------------------------- /adept/_lpse2d/core/laser.py: -------------------------------------------------------------------------------- 1 | from jax import Array 2 | from jax import numpy as jnp 3 | 4 | 5 | class Light: 6 | def __init__(self, cfg) -> None: 7 | self.cfg = cfg 8 | self.E0_source = cfg["units"]["derived"]["E0_source"] 9 | self.c = cfg["units"]["derived"]["c"] 10 | self.w0 = cfg["units"]["derived"]["w0"] 11 | self.dE0x = jnp.zeros((cfg["grid"]["nx"], cfg["grid"]["ny"])) 12 | self.x = cfg["grid"]["x"] 13 | 14 | def laser_update(self, t: float, y: jnp.ndarray, light_wave: dict) -> tuple[jnp.ndarray, jnp.ndarray]: 15 | """ 16 | This function updates the laser field at time t 17 | 18 | :param t: time 19 | :param y: state variables 20 | :return: updated laser field 21 | """ 22 | # method = "broadcast" 23 | # if method == "broadcast": 24 | wpe = self.w0 * jnp.sqrt(y["background_density"])[..., None] 25 | k0 = self.w0 / self.c * jnp.sqrt((1 + 0j + light_wave["delta_omega"][None, None]) ** 2 - wpe**2 / self.w0**2) 26 | 27 | E0_static = ( 28 | (1 + 0j - wpe**2.0 / (self.w0 * (1 + light_wave["delta_omega"][None, None])) ** 2) ** -0.25 29 | * self.E0_source 30 | * jnp.sqrt(light_wave["intensities"][None, None]) 31 | * jnp.exp(1j * k0 * self.x[:, None, None] + 1j * light_wave["phases"][None, None]) 32 | ) 33 | dE0y = E0_static * jnp.exp(-1j * light_wave["delta_omega"][None, None] * self.w0 * t) 34 | dE0y = jnp.sum(dE0y, axis=-1) 35 | 36 | # elif method == "map": 37 | # wpe = self.w0 * jnp.sqrt(y["background_density"]) 38 | 39 | # def _fn_(light_wave_item): 40 | # delta_omega, intensity, initial_phase = light_wave_item 41 | # k0 = self.w0 / self.c * jnp.sqrt((1 + 0j + delta_omega) ** 2 - wpe**2 / self.w0**2) 42 | # coeff = (1 + 0j - wpe**2.0 / (self.w0 * (1 + delta_omega)) ** 2) ** -0.25 43 | # E0_static = ( 44 | # coeff 45 | # * self.E0_source 46 | # * jnp.sqrt(intensity) 47 | # * jnp.exp(1j * k0 * self.x[:, None] + 1j * initial_phase) 48 | # ) 49 | # dE0y = E0_static * jnp.exp(-1j * delta_omega * self.w0 * t) 50 | # return dE0y 51 | 52 | # dE0y = jmp( 53 | # _fn_, 54 | # [light_wave["delta_omega"], light_wave["intensities"], light_wave["initial_phase"]], 55 | # batch_size=256, 56 | # ) 57 | # dE0y = jnp.sum(dE0y, axis=0) 58 | 59 | # else: 60 | # raise NotImplementedError 61 | 62 | return jnp.stack([self.dE0x, dE0y], axis=-1) 63 | 64 | def calc_ey_at_one_point(self, t: float, density: Array, light_wave: dict) -> tuple[jnp.ndarray, jnp.ndarray]: 65 | """ 66 | This function is used to calculate the coherence time of the laser 67 | 68 | :param t: time 69 | :param y: state variables 70 | :return: updated laser field 71 | """ 72 | 73 | wpe = self.w0 * jnp.sqrt(density)[None, 0, 0] 74 | k0 = self.w0 / self.c * jnp.sqrt((1 + 0j + light_wave["delta_omega"]) ** 2 - wpe**2 / self.w0**2) 75 | E0_static = ( 76 | (1 + 0j - wpe**2.0 / (self.w0 * (1 + light_wave["delta_omega"])) ** 2) ** -0.25 77 | * self.E0_source 78 | * jnp.sqrt(light_wave["intensities"]) 79 | * jnp.exp(1j * k0 * self.x[0] + 1j * light_wave["phases"]) 80 | ) 81 | dE0y = E0_static * jnp.exp(-1j * light_wave["delta_omega"] * self.w0 * t) 82 | return jnp.sum(dE0y, axis=0) 83 | -------------------------------------------------------------------------------- /adept/_lpse2d/core/trapper.py: -------------------------------------------------------------------------------- 1 | import equinox as eqx 2 | import numpy as np 3 | from jax import Array 4 | from jax import numpy as jnp 5 | 6 | from adept import electrostatic 7 | 8 | 9 | class ParticleTrapper(eqx.Module): 10 | kx: np.ndarray 11 | kax_sq: Array 12 | model_kld: float 13 | wis: Array 14 | norm_kld: jnp.float64 15 | norm_nuee: jnp.float64 16 | vph: jnp.float64 17 | fft_norm: float 18 | dx: float 19 | 20 | def __init__(self, cfg, species="electron"): 21 | self.kx = cfg["grid"]["kx"] 22 | self.dx = cfg["grid"]["dx"] 23 | self.kax_sq = cfg["grid"]["kx"][:, None] ** 2 + cfg["grid"]["ky"][None, :] ** 2 24 | table_wrs, table_wis, table_klds = electrostatic.get_complex_frequency_table( 25 | 1024, cfg["terms"]["epw"]["kinetic real part"] 26 | ) 27 | all_ks = jnp.sqrt(self.kax_sq).flatten() 28 | self.model_kld = cfg["terms"]["epw"]["trapping"]["kld"] 29 | self.wis = jnp.interp(all_ks, table_klds, table_wis, left=0.0, right=0.0).reshape(self.kax_sq.shape) 30 | self.norm_kld = (self.model_kld - 0.26) / 0.14 31 | self.norm_nuee = (jnp.log10(1.0e-7) + 7.0) / -4.0 32 | 33 | this_wr = jnp.interp(self.model_kld, table_klds, table_wrs, left=1.0, right=table_wrs[-1]) 34 | self.vph = this_wr / self.model_kld 35 | self.fft_norm = cfg["grid"]["nx"] * cfg["grid"]["ny"] / 4.0 36 | # Make models 37 | # if models is not None: 38 | # self.nu_g_model = models["nu_g"] 39 | # else: 40 | # self.nu_g_model = lambda x: -32 41 | 42 | def __call__(self, t, delta, args): 43 | e = args["eh"] 44 | ek = jnp.fft.fft2(e[..., 0]) / self.fft_norm 45 | 46 | # this is where a specific k is chosen for the growth rate and where the identity of this delta object is given 47 | chosen_ek = jnp.interp(self.model_kld, self.kx, jnp.mean(jnp.abs(ek), axis=1)) 48 | norm_e = (jnp.log10(chosen_ek + 1e-10) + 10.0) / -10.0 49 | func_inputs = jnp.stack([norm_e, self.norm_kld, self.norm_nuee], axis=-1) 50 | growth_rates = 10 ** jnp.squeeze(3 * args["nu_g"](func_inputs)) 51 | 52 | return -self.vph * jnp.gradient(jnp.pad(delta, pad_width=1, mode="wrap"), axis=0)[ 53 | 1:-1, 1:-1 54 | ] / self.dx + growth_rates * jnp.abs(jnp.fft.ifft2(ek * self.fft_norm * self.wis)) / (1.0 + delta**2.0) 55 | -------------------------------------------------------------------------------- /adept/_lpse2d/core/vector_field.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from jax import Array 3 | from jax import numpy as jnp 4 | 5 | from adept._base_ import get_envelope 6 | from adept._lpse2d.core import epw, laser 7 | 8 | 9 | class SplitStep: 10 | """ 11 | This class contains the function that updates the state 12 | 13 | All the pushers are chosen and initialized here and a single time-step is defined here. 14 | 15 | :param cfg: 16 | :return: 17 | """ 18 | 19 | def __init__(self, cfg): 20 | super().__init__() 21 | self.cfg = cfg 22 | self.dt = cfg["grid"]["dt"] 23 | self.wp0 = cfg["units"]["derived"]["wp0"] 24 | self.epw = epw.SpectralPotential(cfg) 25 | self.light = laser.Light(cfg) 26 | self.complex_state_vars = ["E0", "epw"] 27 | self.boundary_envelope = cfg["grid"]["absorbing_boundaries"] 28 | self.one_over_ksq = cfg["grid"]["one_over_ksq"] 29 | self.zero_mask = cfg["grid"]["zero_mask"] 30 | self.low_pass_filter = cfg["grid"]["low_pass_filter"] 31 | 32 | self.nu_coll = cfg["units"]["derived"]["nu_coll"] 33 | 34 | def _unpack_y_(self, y: dict[str, Array]) -> dict[str, Array]: 35 | new_y = {} 36 | for k in y.keys(): 37 | if k in self.complex_state_vars: 38 | new_y[k] = y[k].view(jnp.complex128) 39 | else: 40 | new_y[k] = y[k].view(jnp.float64) 41 | return new_y 42 | 43 | def _pack_y_(self, y: dict[str, Array], new_y: dict[str, Array]) -> tuple[dict[str, Array], dict[str, Array]]: 44 | for k in y.keys(): 45 | y[k] = y[k].view(jnp.float64) 46 | new_y[k] = new_y[k].view(jnp.float64) 47 | 48 | return y, new_y 49 | 50 | def light_split_step(self, t, y, driver_args): 51 | if "E0" in driver_args: 52 | t_coeff = get_envelope( 53 | driver_args["E0"]["tr"], 54 | driver_args["E0"]["tr"], 55 | driver_args["E0"]["tc"] - driver_args["E0"]["tw"] / 2, 56 | driver_args["E0"]["tc"] + driver_args["E0"]["tw"] / 2, 57 | t, 58 | ) 59 | y["E0"] = t_coeff * self.light.laser_update(t, y, driver_args["E0"]) 60 | 61 | # if self.cfg["terms"]["light"]["update"]: 62 | # y["E0"] = y["E0"] + self.dt * jnp.real(k1_E0) 63 | 64 | # t_coeff = get_envelope(0.1, 0.1, 0.2, 100.0, t + 0.5 * self.dt) 65 | # y["E0"] = t_coeff * self.light.laser_update(t + 0.5 * self.dt, y, args["E0"]) 66 | # if self.cfg["terms"]["light"]["update"]: 67 | # y["E0"] = y["E0"] + 1j * self.dt * jnp.imag(k1_E0) 68 | 69 | return y 70 | 71 | def landau_damping(self, epw: Array, vte_sq: float): 72 | gammaLandauEpw = ( 73 | np.sqrt(np.pi / 8) 74 | * self.wp0**4 75 | * self.one_over_ksq**1.5 76 | / (vte_sq**1.5) 77 | * jnp.exp(-(self.wp0**2.0) * self.one_over_ksq / (2 * vte_sq)) 78 | ) 79 | 80 | return jnp.fft.ifft2( 81 | jnp.fft.fft2(epw) * jnp.exp(-(gammaLandauEpw + self.nu_coll) * self.dt) * self.low_pass_filter 82 | ) 83 | 84 | def __call__(self, t, y, args): 85 | # unpack y into complex128 86 | new_y = self._unpack_y_(y) 87 | 88 | if self.cfg["terms"]["epw"]["damping"]["landau"]: 89 | new_y["epw"] = self.landau_damping(epw=new_y["epw"], vte_sq=y["vte_sq"]) 90 | 91 | # split step 92 | new_y = self.light_split_step(t, new_y, args["drivers"]) 93 | 94 | if "E2" in args["drivers"]: 95 | new_y["epw"] += self.dt * self.epw.driver(args["drivers"]["E2"], t) 96 | new_y["epw"] = self.epw(t, new_y, args) 97 | 98 | # landau and collisional damping 99 | 100 | # boundary damping 101 | ex, ey = self.epw.calc_fields_from_phi(new_y["epw"]) 102 | ex = ex * self.boundary_envelope 103 | ey = ey * self.boundary_envelope 104 | new_y["epw"] = self.epw.calc_phi_from_fields(ex, ey) 105 | # new_y["epw"] = jnp.fft.ifft2(self.zero_mask * self.low_pass_filter * jnp.fft.fft2(new_y["epw"])) 106 | 107 | # pack y into float64 108 | y, new_y = self._pack_y_(y, new_y) 109 | 110 | return new_y 111 | -------------------------------------------------------------------------------- /adept/_lpse2d/datamodel.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel 2 | 3 | 4 | class NoiseModel(BaseModel): 5 | """ 6 | Noise model for the density profile 7 | """ 8 | 9 | max: float 10 | min: float 11 | type: str 12 | 13 | 14 | class DensityModel(BaseModel): 15 | """ 16 | Density profile for the simulation 17 | 18 | """ 19 | 20 | basis: str 21 | gradient_scale_length: str 22 | max: float 23 | min: float 24 | noise: NoiseModel 25 | 26 | 27 | class EnvelopeModel(BaseModel): 28 | """ 29 | Envelope model for the driver 30 | 31 | """ 32 | 33 | tw: str 34 | tr: str 35 | tc: str 36 | xr: str 37 | xw: str 38 | xc: str 39 | yr: str 40 | yw: str 41 | yc: str 42 | 43 | 44 | class E0DriverModel(BaseModel): 45 | """ 46 | E0 driver model 47 | 48 | """ 49 | 50 | amplitude_shape: str 51 | # Uncomment the following line if the file path is needed 52 | # file: Optional[str] 53 | delta_omega_max: float 54 | num_colors: int 55 | envelope: EnvelopeModel 56 | 57 | 58 | class DriversModel(BaseModel): 59 | """ 60 | Define the drivers for the simulation 61 | 62 | """ 63 | 64 | E0: E0DriverModel 65 | 66 | 67 | class GridModel(BaseModel): 68 | """ 69 | Define the grid for the simulation 70 | 71 | """ 72 | 73 | boundary_abs_coeff: float 74 | boundary_width: str 75 | low_pass_filter: float 76 | dt: str 77 | dx: str 78 | tmax: str 79 | tmin: str 80 | ymax: str 81 | ymin: str 82 | 83 | 84 | class TimeSaveModel(BaseModel): 85 | dt: str 86 | tmax: str 87 | tmin: str 88 | 89 | 90 | class XSaveModel(BaseModel): 91 | dx: str 92 | 93 | 94 | class YSaveModel(BaseModel): 95 | dy: str 96 | 97 | 98 | class SaveModel(BaseModel): 99 | t: TimeSaveModel 100 | x: XSaveModel 101 | y: YSaveModel 102 | 103 | 104 | class BoundaryModel(BaseModel): 105 | x: str 106 | y: str 107 | 108 | 109 | class DampingModel(BaseModel): 110 | collisions: bool 111 | landau: bool 112 | 113 | 114 | class SourceModel(BaseModel): 115 | noise: bool 116 | tpd: bool 117 | 118 | 119 | class EPWModel(BaseModel): 120 | boundary: BoundaryModel 121 | damping: DampingModel 122 | density_gradient: bool 123 | linear: bool 124 | source: SourceModel 125 | 126 | 127 | class TermsModel(BaseModel): 128 | epw: EPWModel 129 | zero_mask: bool 130 | 131 | 132 | class UnitsModel(BaseModel): 133 | atomic_number: int 134 | envelope_density: float 135 | ionization_state: int 136 | laser_intensity: str 137 | laser_wavelength: str 138 | reference_electron_temperature: str 139 | reference_ion_temperature: str 140 | 141 | 142 | class MLFlowModel(BaseModel): 143 | experiment: str 144 | run: str 145 | 146 | 147 | class ConfigModel(BaseModel): 148 | density: DensityModel 149 | drivers: DriversModel 150 | grid: GridModel 151 | mlflow: MLFlowModel 152 | save: SaveModel 153 | solver: str 154 | terms: TermsModel 155 | units: UnitsModel 156 | -------------------------------------------------------------------------------- /adept/_lpse2d/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import BaseLPSE2D as BaseLPSE2D 2 | from .driver import ArbitraryDriver, GaussianDriver, LorentzianDriver, UniformDriver 3 | -------------------------------------------------------------------------------- /adept/_lpse2d/modules/base.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from diffrax import ODETerm, SaveAt, SubSaveAt, diffeqsolve 3 | from equinox import filter_jit 4 | 5 | from adept import ADEPTModule 6 | from adept._base_ import Stepper 7 | from adept._lpse2d.core.vector_field import SplitStep 8 | from adept._lpse2d.helpers import ( 9 | get_density_profile, 10 | get_derived_quantities, 11 | get_save_quantities, 12 | get_solver_quantities, 13 | post_process, 14 | write_units, 15 | ) 16 | from adept._lpse2d.modules import driver 17 | 18 | 19 | class BaseLPSE2D(ADEPTModule): 20 | def __init__(self, cfg) -> None: 21 | super().__init__(cfg) 22 | 23 | def post_process(self, run_output: dict, td: str) -> dict: 24 | return post_process(run_output["solver result"], self.cfg, td) 25 | 26 | def write_units(self) -> dict: 27 | """ 28 | Write the units to a file 29 | 30 | :param cfg: 31 | :param td: 32 | :return: cfg 33 | """ 34 | return write_units(self.cfg) 35 | 36 | def get_derived_quantities(self): 37 | self.cfg = get_derived_quantities(self.cfg) 38 | 39 | def get_solver_quantities(self): 40 | self.cfg["grid"] = get_solver_quantities(self.cfg) 41 | 42 | def init_modules(self) -> dict: 43 | modules = {} 44 | if "E0" in self.cfg["drivers"]: 45 | DriverModule = driver.choose_driver(self.cfg["drivers"]["E0"]["shape"]) 46 | if "file" in self.cfg["drivers"]["E0"]: 47 | modules["laser"] = driver.load(self.cfg, DriverModule) 48 | else: 49 | modules["laser"] = DriverModule(self.cfg) 50 | 51 | return modules 52 | 53 | def init_diffeqsolve(self): 54 | self.cfg = get_save_quantities(self.cfg) 55 | self.time_quantities = { 56 | "t0": 0.0, 57 | "t1": self.cfg["grid"]["tmax"], 58 | "max_steps": self.cfg["grid"]["max_steps"], 59 | "save_t0": 0.0, 60 | "save_t1": self.cfg["grid"]["tmax"], 61 | "save_nt": self.cfg["grid"]["tmax"], 62 | } 63 | 64 | self.diffeqsolve_quants = dict( 65 | terms=ODETerm(SplitStep(self.cfg)), 66 | solver=Stepper(), 67 | saveat=dict( 68 | subs={ 69 | k: SubSaveAt(ts=subsave["t"]["ax"], fn=subsave["func"]) for k, subsave in self.cfg["save"].items() 70 | } 71 | ), 72 | ) 73 | 74 | def init_state_and_args(self) -> dict: 75 | if self.cfg["density"]["noise"]["type"] == "uniform": 76 | random_amps = np.random.uniform( 77 | self.cfg["density"]["noise"]["min"], 78 | self.cfg["density"]["noise"]["max"], 79 | (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"]), 80 | ) 81 | 82 | elif self.cfg["density"]["noise"]["type"] == "normal": 83 | loc = 0.5 * (self.cfg["density"]["noise"]["min"] + self.cfg["density"]["noise"]["max"]) 84 | scale = 1.0 85 | random_amps = np.random.normal(loc, scale, (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"])) 86 | 87 | else: 88 | raise NotImplementedError 89 | 90 | random_phases = np.random.uniform(0, 2 * np.pi, (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"])) 91 | phi_noise = 1 * np.exp(1j * random_phases) 92 | epw = 0 * phi_noise 93 | 94 | background_density = get_density_profile(self.cfg) 95 | vte_sq = np.ones((self.cfg["grid"]["nx"], self.cfg["grid"]["ny"])) * self.cfg["units"]["derived"]["vte"] ** 2 96 | E0 = np.zeros((self.cfg["grid"]["nx"], self.cfg["grid"]["ny"], 2), dtype=np.complex128) 97 | state = {"background_density": background_density, "epw": epw, "E0": E0, "vte_sq": vte_sq} 98 | 99 | self.state = {k: v.view(dtype=np.float64) for k, v in state.items()} 100 | self.args = {"drivers": {k: v["derived"] for k, v in self.cfg["drivers"].items()}} 101 | 102 | @filter_jit 103 | def __call__(self, trainable_modules: dict, args: dict | None = None) -> dict: 104 | state = self.state 105 | 106 | if args is not None: 107 | args = self.args | args 108 | else: 109 | args = self.args 110 | 111 | for name, module in trainable_modules.items(): 112 | state, args = module(state, args) 113 | 114 | solver_result = diffeqsolve( 115 | terms=self.diffeqsolve_quants["terms"], 116 | solver=self.diffeqsolve_quants["solver"], 117 | t0=self.time_quantities["t0"], 118 | t1=self.time_quantities["t1"], 119 | max_steps=None, # self.cfg["grid"]["max_steps"], 120 | dt0=self.cfg["grid"]["dt"], 121 | y0=state, 122 | args=args, 123 | saveat=SaveAt(**self.diffeqsolve_quants["saveat"]), 124 | ) 125 | 126 | return {"solver result": solver_result, "args": args} 127 | -------------------------------------------------------------------------------- /adept/_tf1d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_tf1d/__init__.py -------------------------------------------------------------------------------- /adept/_tf1d/datamodel.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel 2 | 3 | 4 | class MLFlowModel(BaseModel): 5 | experiment: str 6 | run: str 7 | 8 | 9 | class UnitsModel(BaseModel): 10 | laser_wavelength: str 11 | normalizing_temperature: str 12 | normalizing_density: str 13 | Z: int 14 | Zp: int 15 | 16 | 17 | class GridModel(BaseModel): 18 | nx: int 19 | xmin: float 20 | xmax: float 21 | tmin: float 22 | tmax: float 23 | 24 | 25 | class TimeSaveModel(BaseModel): 26 | tmin: float 27 | tmax: float 28 | nt: int 29 | 30 | 31 | class SpaceSaveModel(BaseModel): 32 | xmin: float 33 | xmax: float 34 | nx: int 35 | 36 | 37 | class KxSaveModel(BaseModel): 38 | kxmin: float 39 | kxmax: float 40 | nkx: int 41 | 42 | 43 | class SaveModel(BaseModel): 44 | t: TimeSaveModel 45 | x: SpaceSaveModel 46 | kx: KxSaveModel 47 | 48 | 49 | class TrappingModel(BaseModel): 50 | is_on: bool 51 | kld: float 52 | 53 | 54 | class IonModel(BaseModel): 55 | is_on: bool 56 | landau_damping: bool 57 | mass: float 58 | T0: float 59 | charge: float 60 | gamma: int | str 61 | trapping: TrappingModel 62 | 63 | 64 | class ElectronModel(BaseModel): 65 | is_on: bool 66 | landau_damping: bool 67 | T0: float 68 | mass: float 69 | charge: float 70 | gamma: int | str 71 | trapping: TrappingModel 72 | 73 | 74 | class PhysicsModel(BaseModel): 75 | ion: IonModel 76 | electron: ElectronModel 77 | 78 | 79 | class ExDriverModel(BaseModel): 80 | k0: float 81 | w0: float 82 | dw0: float 83 | t_c: float 84 | t_w: float 85 | t_r: float 86 | x_c: float 87 | x_w: float 88 | x_r: float 89 | a0: float 90 | 91 | 92 | class DriversModel(BaseModel): 93 | ex: dict[str, ExDriverModel] 94 | 95 | 96 | class ConfigModel(BaseModel): 97 | solver: str 98 | mlflow: MLFlowModel 99 | adjoint: bool 100 | units: UnitsModel 101 | grid: GridModel 102 | save: SaveModel 103 | physics: PhysicsModel 104 | drivers: DriversModel 105 | -------------------------------------------------------------------------------- /adept/_tf1d/solvers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_tf1d/solvers/__init__.py -------------------------------------------------------------------------------- /adept/_tf1d/solvers/vector_field.py: -------------------------------------------------------------------------------- 1 | from collections.abc import Callable 2 | 3 | import equinox as eqx 4 | import jax.numpy as jnp 5 | 6 | from adept._tf1d.solvers import pushers 7 | 8 | 9 | class VF(eqx.Module): 10 | """ 11 | This function returns the function that defines $d_state / dt$ 12 | 13 | All the pushers are chosen and initialized here and a single time-step is defined here. 14 | 15 | We use the time-integrators provided by diffrax, and therefore, only need $d_state / dt$ here 16 | 17 | :param cfg: 18 | :return: 19 | """ 20 | 21 | cfg: dict 22 | pusher_dict: dict 23 | push_driver: Callable 24 | poisson_solver: Callable 25 | 26 | def __init__(self, cfg): 27 | super().__init__() 28 | self.cfg = cfg 29 | self.pusher_dict = {"ion": {}, "electron": {}} 30 | for species_name in ["ion", "electron"]: 31 | self.pusher_dict[species_name]["push_n"] = pushers.DensityStepper(cfg["grid"]["kx"]) 32 | self.pusher_dict[species_name]["push_u"] = pushers.VelocityStepper( 33 | cfg["grid"]["kx"], cfg["grid"]["kxr"], cfg["grid"]["one_over_kxr"], cfg["physics"][species_name] 34 | ) 35 | self.pusher_dict[species_name]["push_e"] = pushers.EnergyStepper( 36 | cfg["grid"]["kx"], cfg["physics"][species_name] 37 | ) 38 | if cfg["physics"][species_name]["trapping"]["is_on"]: 39 | self.pusher_dict[species_name]["particle_trapper"] = pushers.ParticleTrapper(cfg, species_name) 40 | 41 | self.push_driver = pushers.Driver(cfg["grid"]["x"]) 42 | # if "ey" in self.cfg["drivers"]: 43 | # self.wave_solver = pushers.WaveSolver(cfg["grid"]["c"], cfg["grid"]["dx"], cfg["grid"]["dt"]) 44 | self.poisson_solver = pushers.PoissonSolver(cfg["grid"]["one_over_kx"]) 45 | 46 | def __call__(self, t: float, y: dict, args: dict): 47 | """ 48 | This function is used by the time integrators specified in diffrax 49 | 50 | :param t: 51 | :param y: 52 | :param args: 53 | :return: 54 | """ 55 | e = self.poisson_solver( 56 | self.cfg["physics"]["ion"]["charge"] * y["ion"]["n"] 57 | + self.cfg["physics"]["electron"]["charge"] * y["electron"]["n"] 58 | ) 59 | ed = 0.0 60 | 61 | for p_ind in self.cfg["drivers"]["ex"].keys(): 62 | ed += self.push_driver(args["drivers"]["ex"][p_ind], t) 63 | 64 | # if "ey" in self.cfg["drivers"]: 65 | # ad = 0.0 66 | # for p_ind in self.cfg["drivers"]["ey"].keys(): 67 | # ad += self.push_driver(args["pulse"]["ey"][p_ind], t) 68 | # a = self.wave_solver(a, aold, djy_array, charge) 69 | # total_a = y["a"] + ad 70 | # ponderomotive_force = -0.5 * jnp.gradient(jnp.square(total_a), self.cfg["grid"]["dx"])[1:-1] 71 | 72 | total_e = e + ed # + ponderomotive_force 73 | 74 | dstate_dt = {"ion": {}, "electron": {}} 75 | for species_name in ["ion", "electron"]: 76 | n = y[species_name]["n"] 77 | u = y[species_name]["u"] 78 | p = y[species_name]["p"] 79 | delta = y[species_name]["delta"] 80 | if self.cfg["physics"][species_name]["is_on"]: 81 | q_over_m = self.cfg["physics"][species_name]["charge"] / self.cfg["physics"][species_name]["mass"] 82 | p_over_m = p / self.cfg["physics"][species_name]["mass"] 83 | 84 | dstate_dt[species_name]["n"] = self.pusher_dict[species_name]["push_n"](n, u) 85 | dstate_dt[species_name]["u"] = self.pusher_dict[species_name]["push_u"]( 86 | n, u, p_over_m, q_over_m * total_e, delta 87 | ) 88 | dstate_dt[species_name]["p"] = self.pusher_dict[species_name]["push_e"](n, u, p_over_m, q_over_m * e) 89 | else: 90 | dstate_dt[species_name]["n"] = jnp.zeros(self.cfg["grid"]["nx"]) 91 | dstate_dt[species_name]["u"] = jnp.zeros(self.cfg["grid"]["nx"]) 92 | dstate_dt[species_name]["p"] = jnp.zeros(self.cfg["grid"]["nx"]) 93 | 94 | if self.cfg["physics"][species_name]["trapping"]["is_on"]: 95 | dstate_dt[species_name]["delta"] = self.pusher_dict[species_name]["particle_trapper"](e, delta, args) 96 | else: 97 | dstate_dt[species_name]["delta"] = jnp.zeros(self.cfg["grid"]["nx"]) 98 | 99 | return dstate_dt 100 | -------------------------------------------------------------------------------- /adept/_tf1d/storage.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import xarray as xr 4 | from diffrax import Solution 5 | from flatdict import FlatDict 6 | from matplotlib import pyplot as plt 7 | 8 | 9 | def save_arrays(result: Solution, td: str, cfg: dict, label: str) -> xr.Dataset: 10 | """ 11 | This function saves the arrays to an xarray netcdf file 12 | 13 | """ 14 | if label is None: 15 | label = "x" 16 | flattened_dict = dict(FlatDict(result.ys, delimiter="-")) 17 | save_ax = cfg["grid"]["x"] 18 | else: 19 | flattened_dict = dict(FlatDict(result.ys[label], delimiter="-")) 20 | save_ax = cfg["save"][label]["ax"] 21 | data_vars = { 22 | k: xr.DataArray(v, coords=(("t", cfg["save"]["t"]["ax"]), (label, save_ax))) for k, v in flattened_dict.items() 23 | } 24 | 25 | saved_arrays_xr = xr.Dataset(data_vars) 26 | saved_arrays_xr.to_netcdf(os.path.join(td, "binary", f"state_vs_{label}.nc")) 27 | return saved_arrays_xr 28 | 29 | 30 | def plot_xrs(which: str, td: str, xrs: dict): 31 | """ 32 | This function plots the xarray datasets 33 | 34 | 35 | """ 36 | os.makedirs(os.path.join(td, "plots", which)) 37 | os.makedirs(os.path.join(td, "plots", which, "ion")) 38 | os.makedirs(os.path.join(td, "plots", which, "electron")) 39 | 40 | for k, v in xrs.items(): 41 | fname = f"{'-'.join(k.split('-')[1:])}.png" 42 | fig, ax = plt.subplots(1, 1, figsize=(7, 4), tight_layout=True) 43 | v.plot(ax=ax, cmap="gist_ncar") 44 | ax.grid() 45 | fig.savefig(os.path.join(td, "plots", which, k.split("-")[0], fname), bbox_inches="tight") 46 | plt.close(fig) 47 | 48 | if which == "kx": 49 | os.makedirs(os.path.join(td, "plots", which, "ion", "hue"), exist_ok=True) 50 | os.makedirs(os.path.join(td, "plots", which, "electron", "hue"), exist_ok=True) 51 | # only plot 52 | if v.coords["kx"].size > 8: 53 | hue_skip = v.coords["kx"].size // 8 54 | else: 55 | hue_skip = 1 56 | 57 | for log in [True, False]: 58 | fig, ax = plt.subplots(1, 1, figsize=(7, 4), tight_layout=True) 59 | v[:, ::hue_skip].plot(ax=ax, hue="kx") 60 | ax.set_yscale("log" if log else "linear") 61 | ax.grid() 62 | fig.savefig( 63 | os.path.join( 64 | td, "plots", which, k.split("-")[0], "hue", f"{'-'.join(k.split('-')[1:])}-log-{log}.png" 65 | ), 66 | bbox_inches="tight", 67 | ) 68 | plt.close(fig) 69 | -------------------------------------------------------------------------------- /adept/_version.py: -------------------------------------------------------------------------------- 1 | # file generated by setuptools-scm 2 | # don't change, don't track in version control 3 | 4 | __all__ = ["__version__", "__version_tuple__", "version", "version_tuple"] 5 | 6 | TYPE_CHECKING = False 7 | if TYPE_CHECKING: 8 | from typing import Tuple 9 | from typing import Union 10 | 11 | VERSION_TUPLE = Tuple[Union[int, str], ...] 12 | else: 13 | VERSION_TUPLE = object 14 | 15 | version: str 16 | __version__: str 17 | __version_tuple__: VERSION_TUPLE 18 | version_tuple: VERSION_TUPLE 19 | 20 | __version__ = version = '0.0.8.dev0+g713e2ef.d20250523' 21 | __version_tuple__ = version_tuple = (0, 0, 8, 'dev0', 'g713e2ef.d20250523') 22 | -------------------------------------------------------------------------------- /adept/_vlasov1d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_vlasov1d/__init__.py -------------------------------------------------------------------------------- /adept/_vlasov1d/datamodel.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel 2 | 3 | 4 | class SpaceProfileModel(BaseModel): 5 | baseline: float 6 | bump_or_trough: str 7 | center: float 8 | rise: float 9 | slope: float 10 | bump_height: float 11 | width: float 12 | 13 | 14 | class SpeciesBackgroundModel(BaseModel): 15 | noise_seed: int 16 | noise_type: str 17 | noise_val: float 18 | v0: float 19 | T0: float 20 | m: float 21 | basis: str 22 | space_profile: SpaceProfileModel 23 | 24 | 25 | class DensityModel(BaseModel): 26 | quasineutrality: bool 27 | species_background: SpeciesBackgroundModel 28 | 29 | 30 | class UnitsModel(BaseModel): 31 | laser_wavelength: str 32 | normalizing_temperature: str 33 | normalizing_density: str 34 | Z: int 35 | Zp: int 36 | 37 | 38 | class GridModel(BaseModel): 39 | dt: float 40 | nv: int 41 | nx: int 42 | tmin: float 43 | tmax: float 44 | vmax: float 45 | xmax: float 46 | xmin: float 47 | 48 | 49 | class TimeSaveModel(BaseModel): 50 | tmin: float 51 | tmax: float 52 | nt: int 53 | 54 | 55 | class SaveModel(BaseModel): 56 | fields: dict[str, TimeSaveModel] 57 | electron: dict[str, TimeSaveModel] 58 | 59 | 60 | class ExDriverModel(BaseModel): 61 | a0: float 62 | k0: float 63 | t_center: float 64 | t_rise: float 65 | t_width: float 66 | w0: float 67 | dw0: float 68 | x_center: float 69 | x_rise: float 70 | x_width: float 71 | 72 | 73 | class EyDriverModel(BaseModel): 74 | pass 75 | 76 | 77 | class DriversModel(BaseModel): 78 | ex: dict[str, ExDriverModel] 79 | ey: EyDriverModel 80 | 81 | 82 | class TimeTermModel(BaseModel): 83 | baseline: float 84 | bump_or_trough: str 85 | center: float 86 | rise: float 87 | slope: float 88 | bump_height: float 89 | width: float 90 | 91 | 92 | class SpaceTermModel(BaseModel): 93 | baseline: float 94 | bump_or_trough: str 95 | center: float 96 | rise: float 97 | slope: float 98 | bump_height: float 99 | width: float 100 | 101 | 102 | class FokkerPlanckModel(BaseModel): 103 | is_on: bool 104 | type: str 105 | time: TimeTermModel 106 | space: SpaceTermModel 107 | 108 | 109 | class KrookModel(BaseModel): 110 | is_on: bool 111 | time: TimeTermModel 112 | space: SpaceTermModel 113 | 114 | 115 | class TermsModel(BaseModel): 116 | field: str 117 | edfdv: str 118 | time: str 119 | fokker_planck: FokkerPlanckModel 120 | krook: KrookModel 121 | 122 | 123 | class MLFlowModel(BaseModel): 124 | experiment: str 125 | run: str 126 | 127 | 128 | class ConfigModel(BaseModel): 129 | units: UnitsModel 130 | density: DensityModel 131 | grid: GridModel 132 | save: SaveModel 133 | solver: str 134 | mlflow: MLFlowModel 135 | drivers: DriversModel 136 | terms: TermsModel 137 | -------------------------------------------------------------------------------- /adept/_vlasov1d/gamma_func_for_sg.nc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_vlasov1d/gamma_func_for_sg.nc -------------------------------------------------------------------------------- /adept/_vlasov1d/solvers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_vlasov1d/solvers/__init__.py -------------------------------------------------------------------------------- /adept/_vlasov1d/solvers/pushers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/_vlasov1d/solvers/pushers/__init__.py -------------------------------------------------------------------------------- /adept/_vlasov1d/solvers/pushers/fokker_planck.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | 4 | import numpy as np 5 | from jax import numpy as jnp 6 | 7 | from adept.vlasov2d.solver.tridiagonal import TridiagonalSolver 8 | 9 | 10 | class Collisions: 11 | def __init__(self, cfg): 12 | self.cfg = cfg 13 | self.fp = self.__init_fp_operator__() 14 | self.krook = Krook(self.cfg) 15 | self.td_solver = TridiagonalSolver(self.cfg) 16 | 17 | def __init_fp_operator__(self): 18 | if self.cfg["terms"]["fokker_planck"]["type"].casefold() == "lenard_bernstein": 19 | return LenardBernstein(self.cfg) 20 | elif self.cfg["terms"]["fokker_planck"]["type"].casefold() == "dougherty": 21 | return Dougherty(self.cfg) 22 | else: 23 | raise NotImplementedError 24 | 25 | def __call__(self, nu_fp: jnp.ndarray, nu_K: jnp.ndarray, f: jnp.ndarray, dt: jnp.float64) -> jnp.ndarray: 26 | if self.cfg["terms"]["fokker_planck"]["is_on"]: 27 | # The three diagonals representing collision operator for all x 28 | cee_a, cee_b, cee_c = self.fp(nu=nu_fp, f_xv=f, dt=dt) 29 | # Solve over all x 30 | f = self.td_solver(cee_a, cee_b, cee_c, f) 31 | 32 | if self.cfg["terms"]["krook"]["is_on"]: 33 | f = self.krook(nu_K, f, dt) 34 | 35 | return f 36 | 37 | 38 | class Krook: 39 | def __init__(self, cfg): 40 | self.cfg = cfg 41 | f_mx = np.exp(-(self.cfg["grid"]["v"][None, :] ** 2.0) / 2.0) 42 | self.f_mx = f_mx / np.trapz(f_mx, dx=self.cfg["grid"]["dv"], axis=1)[:, None] 43 | self.dv = self.cfg["grid"]["dv"] 44 | 45 | def vx_moment(self, f_xv): 46 | return jnp.sum(f_xv, axis=1) * self.dv 47 | 48 | def __call__(self, nu_K, f_xv, dt) -> jnp.ndarray: 49 | nu_Kxdt = dt * nu_K[:, None] 50 | exp_nuKxdt = jnp.exp(-nu_Kxdt) 51 | n_prof = self.vx_moment(f_xv) 52 | 53 | return f_xv * exp_nuKxdt + n_prof[:, None] * self.f_mx * (1.0 - exp_nuKxdt) 54 | 55 | 56 | class LenardBernstein: 57 | def __init__(self, cfg): 58 | self.cfg = cfg 59 | self.v = self.cfg["grid"]["v"] 60 | self.dv = self.cfg["grid"]["dv"] 61 | self.ones = jnp.ones((self.cfg["grid"]["nx"], self.cfg["grid"]["nv"])) 62 | 63 | def vx_moment(self, f_xv): 64 | return jnp.sum(f_xv, axis=1) * self.dv 65 | 66 | def __call__( 67 | self, nu: jnp.float64, f_xv: jnp.ndarray, dt: jnp.float64 68 | ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: 69 | """ 70 | 71 | :param nu: 72 | :param f_xv: 73 | :param dt: 74 | :return: 75 | """ 76 | 77 | v0t_sq = self.vx_moment(f_xv * self.v[None, :] ** 2.0) 78 | a = nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 + jnp.roll(self.v, 1)[None, :] / 2 / self.dv) 79 | b = 1.0 + nu[:, None] * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) 80 | c = nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 - jnp.roll(self.v, -1)[None, :] / 2 / self.dv) 81 | return a, b, c 82 | 83 | 84 | class Dougherty: 85 | def __init__(self, cfg): 86 | self.cfg = cfg 87 | self.v = self.cfg["grid"]["v"] 88 | self.dv = self.cfg["grid"]["dv"] 89 | self.ones = jnp.ones((self.cfg["grid"]["nx"], self.cfg["grid"]["nv"])) 90 | 91 | def vx_moment(self, f_xv): 92 | return jnp.sum(f_xv, axis=1) * self.dv 93 | 94 | def __call__( 95 | self, nu: jnp.float64, f_xv: jnp.ndarray, dt: jnp.float64 96 | ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: 97 | """ 98 | 99 | :param nu: 100 | :param f_xv: 101 | :param dt: 102 | :return: 103 | """ 104 | 105 | vbar = self.vx_moment(f_xv * self.v[None, :]) 106 | v0t_sq = self.vx_moment(f_xv * (self.v[None, :] - vbar[:, None]) ** 2.0) 107 | 108 | a = ( 109 | nu[:, None] 110 | * dt 111 | * (-v0t_sq[:, None] / self.dv**2.0 + (jnp.roll(self.v, 1)[None, :] - vbar[:, None]) / 2.0 / self.dv) 112 | ) 113 | b = 1.0 + nu[:, None] * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) 114 | c = ( 115 | nu[:, None] 116 | * dt 117 | * (-v0t_sq[:, None] / self.dv**2.0 - (jnp.roll(self.v, -1)[None, :] - vbar[:, None]) / 2.0 / self.dv) 118 | ) 119 | return a, b, c 120 | -------------------------------------------------------------------------------- /adept/_vlasov1d/solvers/pushers/vlasov.py: -------------------------------------------------------------------------------- 1 | from collections.abc import Callable 2 | from functools import partial 3 | 4 | import equinox as eqx 5 | from interpax import interp1d, interp2d 6 | from jax import numpy as jnp 7 | from jax import vmap 8 | 9 | 10 | class VlasovExternalE(eqx.Module): 11 | x: jnp.ndarray 12 | v: jnp.ndarray 13 | f_interp: Callable 14 | dt: int 15 | dummy_x: jnp.ndarray 16 | dummy_v: jnp.ndarray 17 | interp_e: Callable 18 | 19 | def __init__(self, cfg, interp_e): 20 | self.x = cfg["grid"]["x"] 21 | self.v = cfg["grid"]["v"] 22 | self.f_interp = partial(interp2d, x=self.x, y=self.v, period=cfg["grid"]["xmax"]) 23 | self.dt = cfg["grid"]["dt"] 24 | self.dummy_x = jnp.ones_like(self.x) 25 | self.dummy_v = jnp.ones_like(self.v)[None, :] 26 | self.interp_e = interp_e 27 | 28 | def step_vdfdx(self, t, f, frac_dt): 29 | old_x = self.x[:, None] - self.v[None, :] * frac_dt * self.dt 30 | old_v = self.dummy_x[:, None] * self.v[None, :] 31 | new_f = self.f_interp(xq=old_x.flatten(), yq=old_v.flatten(), f=f) 32 | return jnp.reshape(new_f, (self.x.size, self.v.size)) 33 | 34 | def step_edfdv(self, t, f, frac_dt): 35 | interp_e = self.interp_e(self.dummy_x * t, self.x) 36 | old_v = self.v[None, :] + interp_e[:, None] * frac_dt * self.dt 37 | old_x = self.dummy_v * self.x[:, None] 38 | new_f = self.f_interp(xq=old_x.flatten(), yq=old_v.flatten(), f=f) 39 | return jnp.reshape(new_f, (self.x.size, self.v.size)) 40 | 41 | def __call__(self, t, y, args): 42 | f = y["electron"] 43 | 44 | new_f = self.step_vdfdx(t, f, 0.5) 45 | new_f = self.step_edfdv(t, new_f, 1.0) 46 | new_f = self.step_vdfdx(t, new_f, 0.5) 47 | 48 | return {"electron": new_f} 49 | 50 | 51 | class VelocityExponential: 52 | def __init__(self, cfg): 53 | self.kv_real = cfg["grid"]["kvr"] 54 | 55 | def __call__(self, f, e, dt): 56 | return jnp.real( 57 | jnp.fft.irfft(jnp.exp(-1j * self.kv_real[None, :] * dt * e[:, None]) * jnp.fft.rfft(f, axis=1), axis=1) 58 | ) 59 | 60 | 61 | class VelocityCubicSpline: 62 | def __init__(self, cfg): 63 | self.v = jnp.repeat(cfg["grid"]["v"][None, :], repeats=cfg["grid"]["nx"], axis=0) 64 | self.interp = vmap(partial(interp1d, extrap=True), in_axes=0) # {"xq": 0, "f": 0, "x": None}) 65 | 66 | def __call__(self, f, e, dt): 67 | vq = self.v - e[:, None] * dt 68 | return self.interp(xq=vq, x=self.v, f=f) 69 | 70 | 71 | class SpaceExponential: 72 | def __init__(self, cfg): 73 | self.kx_real = cfg["grid"]["kxr"] 74 | self.v = cfg["grid"]["v"] 75 | 76 | def __call__(self, f, dt): 77 | return jnp.real( 78 | jnp.fft.irfft(jnp.exp(-1j * self.kx_real[:, None] * dt * self.v[None, :]) * jnp.fft.rfft(f, axis=0), axis=0) 79 | ) 80 | -------------------------------------------------------------------------------- /adept/lpse2d.py: -------------------------------------------------------------------------------- 1 | from ._lpse2d import ( 2 | BaseLPSE2D, 3 | calc_threshold_intensity, 4 | ArbitraryDriver, 5 | UniformDriver, 6 | GaussianDriver, 7 | LorentzianDriver, 8 | ) 9 | -------------------------------------------------------------------------------- /adept/sh2d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/sh2d/__init__.py -------------------------------------------------------------------------------- /adept/sh2d/runner.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tempfile 3 | import time 4 | 5 | import diffrax 6 | import equinox as eqx 7 | import mlflow 8 | import numpy as np 9 | import pint 10 | import yaml 11 | from diffrax import ODETerm, SaveAt, Solution, diffeqsolve 12 | from utils import misc 13 | 14 | from adept.sh2d.utils import helpers, save 15 | 16 | 17 | def write_units(cfg, td): 18 | ureg = pint.UnitRegistry() 19 | _Q = ureg.Quantity 20 | 21 | lambda0 = _Q(cfg["units"]["laser_wavelength"]) 22 | w0 = (2 * np.pi / lambda0 * ureg.c).to("rad/s") 23 | t0 = (1 / w0).to("fs") 24 | n0 = (w0**2 * ureg.m_e * ureg.epsilon_0 / ureg.e**2.0).to("1/cc") 25 | nee = _Q(cfg["units"]["density for collisions"]) 26 | T0 = _Q(cfg["units"]["electron temperature"]).to("eV") 27 | v0 = np.sqrt(2.0 * T0 / (ureg.m_e)).to("m/s") 28 | debye_length = (v0 / w0).to("nm") 29 | 30 | logLambda_ee = 23.5 - np.log(nee.magnitude**0.5 * T0.magnitude**-1.25) 31 | logLambda_ee -= (1e-5 + (np.log(T0.magnitude) - 2) ** 2.0 / 16) ** 0.5 32 | nuee = _Q(2.91e-6 * nee.magnitude * logLambda_ee / T0.magnitude**1.5, "Hz") 33 | nuee_norm = nuee / w0 34 | 35 | # if (Ti * me / mi) < Te: 36 | # if Te > 10 * Z ^ 2: 37 | # logLambda_ei = 24 - np.log(ne.magnitude**0.5 / Te.magnitude) 38 | # else: 39 | # logLambda_ei = 23 - np.log(ne.magnitude**0.5 * Z * Te.magnitude**-1.5) 40 | # else: 41 | # logLambda_ei = 30 - np.log(ni.magnitude**0.5 * Z**2 / mu * Ti.magnitude**-1.5) 42 | 43 | # nuei = _Q(2.91e-6 * n0.magnitude * logLambda_ee / T0**1.5, "Hz") 44 | # nuee_norm = nuee / w0 45 | 46 | box_length = ((cfg["grid"]["xmax"] - cfg["grid"]["xmin"]) * debye_length).to("microns") 47 | sim_duration = (cfg["grid"]["tmax"] * t0).to("ps") 48 | 49 | all_quantities = { 50 | "w0": w0, 51 | "t0": t0, 52 | "n0": n0, 53 | "v0": v0, 54 | "T0": T0, 55 | "lambda_D": debye_length, 56 | "logLambda_ee": logLambda_ee, 57 | "nuee": nuee, 58 | "nuee_norm": nuee_norm, 59 | "box_length": box_length, 60 | "sim_duration": sim_duration, 61 | } 62 | all_quantities_str = {k: str(v) for k, v in all_quantities.items()} 63 | 64 | with open(os.path.join(td, "units.yaml"), "w") as fi: 65 | yaml.dump(all_quantities_str, fi) 66 | 67 | print("units: ") 68 | print(all_quantities_str) 69 | print() 70 | 71 | mlflow.log_artifacts(td) 72 | 73 | return all_quantities 74 | 75 | 76 | def run(cfg: dict) -> Solution: 77 | with tempfile.TemporaryDirectory() as td: 78 | with open(os.path.join(td, "config.yaml"), "w") as fi: 79 | yaml.dump(cfg, fi) 80 | 81 | # get derived quantities 82 | cfg["grid"] = helpers.get_derived_quantities(cfg["grid"]) 83 | misc.log_params(cfg) 84 | 85 | cfg["grid"] = helpers.get_solver_quantities(cfg["grid"]) 86 | cfg = helpers.get_save_quantities(cfg) 87 | 88 | cfg["units"]["derived"] = write_units(cfg, td) 89 | 90 | state = helpers.init_state(cfg) 91 | 92 | # run 93 | t0 = time.time() 94 | 95 | @eqx.filter_jit 96 | def _run_(): 97 | vvf = helpers.VlasovVectorField(cfg) 98 | cvf = helpers.FokkerPlanckVectorField(cfg) 99 | args = {"driver": cfg["drivers"], "b_ext": 0.0} 100 | return diffeqsolve( 101 | terms=diffrax.MultiTerm(ODETerm(vvf), ODETerm(cvf)), 102 | solver=helpers.ExplicitEStepper(), 103 | t0=cfg["grid"]["tmin"], 104 | t1=cfg["grid"]["tmax"], 105 | max_steps=cfg["grid"]["max_steps"], 106 | dt0=cfg["grid"]["dt"], 107 | y0=state, 108 | args=args, 109 | # adjoint=diffrax.DirectAdjoint(), 110 | saveat=SaveAt(ts=cfg["save"]["t"]["ax"]), # , fn=cfg["save"]["func"]["callable"]), 111 | ) 112 | 113 | print("starting run") 114 | result = _run_() 115 | print("finished run") 116 | mlflow.log_metrics({"run_time": round(time.time() - t0, 4)}) 117 | 118 | t0 = time.time() 119 | post_process(result, cfg, td) 120 | mlflow.log_metrics({"postprocess_time": round(time.time() - t0, 4)}) 121 | # log artifacts 122 | mlflow.log_artifacts(td) 123 | 124 | # fin 125 | 126 | return result 127 | 128 | 129 | def post_process(result, cfg: dict, td: str) -> None: 130 | os.makedirs(os.path.join(td, "binary")) 131 | os.makedirs(os.path.join(td, "plots")) 132 | 133 | xrs = save.save_arrays(result, td, cfg) 134 | -------------------------------------------------------------------------------- /adept/sh2d/solvers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/sh2d/solvers/__init__.py -------------------------------------------------------------------------------- /adept/sh2d/solvers/field.py: -------------------------------------------------------------------------------- 1 | import equinox as eqx 2 | import jax 3 | import numpy as np 4 | from jax import numpy as jnp 5 | 6 | 7 | class SpectralPoissonSolver(eqx.Module): 8 | ion_charge: jnp.array 9 | one_over_kx: jnp.array 10 | one_over_ky: jnp.array 11 | dv: float 12 | v: jax.Array 13 | 14 | def __init__(self, ion_charge, one_over_kx, one_over_ky, dv, v): 15 | super().__init__() 16 | self.ion_charge = jnp.array(ion_charge) 17 | self.one_over_kx = jnp.array(one_over_kx) 18 | self.one_over_ky = jnp.array(one_over_ky) 19 | self.dv = dv 20 | self.v = v 21 | 22 | def compute_charges(self, f): 23 | return jnp.sum(f * self.v**2.0, axis=2) * self.dv 24 | 25 | def __call__(self, f: jnp.ndarray): 26 | return jnp.concatenate( 27 | [ 28 | jnp.real( 29 | jnp.fft.ifft( 30 | 1j * self.one_over_kx[:, None] * jnp.fft.fft(self.ion_charge - self.compute_charges(f), axis=0), 31 | axis=0, 32 | ) 33 | )[..., None], 34 | jnp.real( 35 | jnp.fft.ifft( 36 | 1j * self.one_over_ky[None, :] * jnp.fft.fft(self.ion_charge - self.compute_charges(f), axis=1), 37 | axis=1, 38 | ) 39 | )[..., None], 40 | jnp.zeros((self.one_over_kx.size, self.one_over_ky.size, 1)), 41 | ], 42 | axis=-1, 43 | ) 44 | 45 | 46 | class AmpereSolver(eqx.Module): 47 | dv: float 48 | v: jax.Array 49 | 50 | def __init__(self, cfg): 51 | super().__init__() 52 | self.dv = cfg["grid"]["dv"] 53 | self.v = cfg["grid"]["v"] 54 | 55 | def compute_currents(self, f): 56 | return ( 57 | -8.0 58 | / 3.0 59 | * np.pi 60 | * self.dv 61 | * jnp.concatenate( 62 | [ 63 | 0.5 * jnp.sum(f[0] * self.v[None, None, :] ** 3.0, axis=2)[..., None], 64 | jnp.sum(jnp.real(f[1]) * self.v[None, None, :] ** 3.0, axis=2)[..., None], 65 | -jnp.sum(jnp.imag(f[1]) * self.v[None, None, :] ** 3.0, axis=2)[..., None], 66 | ], 67 | axis=-1, 68 | ) 69 | ) 70 | 71 | def __call__(self, t, y, args): 72 | return self.compute_currents(y["flm"][1]) 73 | -------------------------------------------------------------------------------- /adept/sh2d/solvers/tridiagonal.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | 4 | import equinox as eqx 5 | from jax import numpy as jnp 6 | from jax.lax import scan 7 | 8 | 9 | class TridiagonalSolver(eqx.Module): 10 | num_unroll: int 11 | 12 | def __init__(self, cfg): 13 | super().__init__() 14 | self.num_unroll = 16 15 | 16 | @staticmethod 17 | def compute_primes(last_primes, x): 18 | """ 19 | This function is a single iteration of the forward pass in the non-in-place Thomas 20 | tridiagonal algorithm 21 | 22 | :param last_primes: 23 | :param x: 24 | :return: 25 | """ 26 | 27 | last_cp, last_dp = last_primes 28 | a, b, c, d = x 29 | cp = c / (b - a * last_cp) 30 | dp = (d - a * last_dp) / (b - a * last_cp) 31 | new_primes = jnp.stack((cp, dp)) 32 | return new_primes, new_primes 33 | 34 | @staticmethod 35 | def backsubstitution(last_x, x): 36 | """ 37 | This function is a single iteration of the backward pass in the non-in-place Thomas 38 | tridiagonal algorithm 39 | 40 | :param last_x: 41 | :param x: 42 | :return: 43 | """ 44 | cp, dp = x 45 | new_x = dp - cp * last_x 46 | return new_x, new_x 47 | 48 | def __call__(self, a, b, c, d): 49 | """ 50 | Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d. 51 | 52 | This uses the non-in-place Thomas tridiagonal algorithm. 53 | 54 | The NumPy version, on the other hand, uses the in-place algorithm. 55 | 56 | :param a: (2D float array (nx, nv)) represents the subdiagonal of the linear operator 57 | :param b: (2D float array (nx, nv)) represents the main diagonal of the linear operator 58 | :param c: (2D float array (nx, nv)) represents the super diagonal of the linear operator 59 | :param d: (2D float array (nx, nv)) represents the right hand side of the linear operator 60 | :return: 61 | """ 62 | 63 | diags_stacked = jnp.stack([arr.transpose((1, 0)) for arr in (a, b, c, d)], axis=1) 64 | _, primes = scan(self.compute_primes, jnp.zeros((2, *a.shape[:-1])), diags_stacked, unroll=self.num_unroll) 65 | _, sol = scan(self.backsubstitution, jnp.zeros(a.shape[:-1]), primes[::-1], unroll=self.num_unroll) 66 | return sol[::-1].transpose((1, 0)) 67 | -------------------------------------------------------------------------------- /adept/sh2d/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/sh2d/utils/__init__.py -------------------------------------------------------------------------------- /adept/tf1d.py: -------------------------------------------------------------------------------- 1 | from ._tf1d.modules import BaseTwoFluid1D 2 | from ._tf1d.datamodel import ConfigModel 3 | -------------------------------------------------------------------------------- /adept/vfp1d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vfp1d/__init__.py -------------------------------------------------------------------------------- /adept/vlasov1d.py: -------------------------------------------------------------------------------- 1 | from ._vlasov1d.modules import BaseVlasov1D 2 | -------------------------------------------------------------------------------- /adept/vlasov1d2v/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov1d2v/__init__.py -------------------------------------------------------------------------------- /adept/vlasov1d2v/gamma_func_for_sg.nc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov1d2v/gamma_func_for_sg.nc -------------------------------------------------------------------------------- /adept/vlasov1d2v/pushers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov1d2v/pushers/__init__.py -------------------------------------------------------------------------------- /adept/vlasov1d2v/pushers/vlasov.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | from interpax import interp1d 4 | from jax import numpy as jnp 5 | from jax import vmap 6 | 7 | 8 | class VelocityExponential: 9 | def __init__(self, cfg): 10 | self.kv_real = cfg["grid"]["kvr"] 11 | 12 | def __call__(self, f, e, dt): 13 | return jnp.real( 14 | jnp.fft.irfft( 15 | jnp.exp(-1j * self.kv_real[None, :, None] * dt * e[:, None, None]) * jnp.fft.rfft(f, axis=1), axis=1 16 | ) 17 | ) 18 | 19 | 20 | class VelocityCubicSpline: 21 | def __init__(self, cfg): 22 | self.v = jnp.repeat(cfg["grid"]["v"][None, :], repeats=cfg["grid"]["nx"], axis=0) 23 | self.v = jnp.repeat(self.v[..., None], repeats=cfg["grid"]["nv"], axis=-1) 24 | 25 | self.vx = cfg["grid"]["v"] 26 | scan_vy = vmap(partial(interp1d, method="cubic", extrap=True), in_axes=1, out_axes=1) 27 | self.scan_all_x_and_vy = vmap(scan_vy, in_axes=0, out_axes=0) 28 | 29 | def __call__(self, f, e, dt): 30 | vq = self.v - e[:, None, None] * dt 31 | return self.scan_all_x_and_vy(vq, self.v, f) 32 | 33 | 34 | class SpaceExponential: 35 | def __init__(self, cfg): 36 | self.kx_real = cfg["grid"]["kxr"] 37 | self.v = cfg["grid"]["v"] 38 | 39 | def __call__(self, f, dt): 40 | return jnp.real( 41 | jnp.fft.irfft( 42 | jnp.exp(-1j * self.kx_real[:, None, None] * dt * self.v[None, :, None]) * jnp.fft.rfft(f, axis=0), 43 | axis=0, 44 | ) 45 | ) 46 | -------------------------------------------------------------------------------- /adept/vlasov2d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov2d/__init__.py -------------------------------------------------------------------------------- /adept/vlasov2d/gamma_func_for_sg.nc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov2d/gamma_func_for_sg.nc -------------------------------------------------------------------------------- /adept/vlasov2d/pushers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov2d/pushers/__init__.py -------------------------------------------------------------------------------- /adept/vlasov2d/pushers/field.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | 4 | 5 | from jax import numpy as jnp 6 | 7 | 8 | class FieldSolver: 9 | def __init__(self, cfg): 10 | self.ion_charge = cfg["grid"]["ion_charge"] 11 | self.one_over_ikx = cfg["grid"]["one_over_kx"][:, None] / 1j 12 | self.one_over_iky = cfg["grid"]["one_over_ky"][None, :] / 1j 13 | self.kx = cfg["grid"]["kx"][:, None] 14 | self.kx_mask = jnp.where(jnp.abs(self.kx) > 0, 1, 0)[:, None] 15 | self.ky = cfg["grid"]["ky"][None, :] 16 | self.ky_mask = jnp.where(jnp.abs(self.ky) > 0, 1, 0)[None, :] 17 | self.dvx = cfg["grid"]["dvx"] 18 | self.dvy = cfg["grid"]["dvy"] 19 | self.vx = cfg["grid"]["vx"][None, None, :, None] 20 | self.vy = cfg["grid"]["vy"][None, None, None, :] 21 | self.zero = jnp.concatenate([jnp.array([0.0]), jnp.ones(cfg["grid"]["nx"] - 1)]) 22 | 23 | def compute_charge(self, f): 24 | return self.vx_mom(self.vy_mom(f)) 25 | 26 | def vx_mom(self, f): 27 | return jnp.sum(f, axis=2) * self.dvx 28 | 29 | def vy_mom(self, f): 30 | return jnp.sum(f, axis=3) * self.dvy 31 | 32 | def compute_jx(self, f): 33 | return jnp.sum(jnp.sum(self.vx * f, axis=3), axis=2) * self.kx_mask * self.dvx * self.dvy 34 | 35 | def compute_jy(self, f): 36 | return jnp.sum(jnp.sum(self.vy * f, axis=3), axis=2) * self.ky_mask * self.dvx * self.dvy 37 | 38 | def ampere(self, exk, eyk, bzk, dt): 39 | exkp = exk # - dt * (1j * self.ky * bzk) 40 | eykp = eyk # - dt * (-1j * self.kx * bzk) 41 | return exkp, eykp 42 | 43 | def faraday(self, bzk, exk, eyk, dt): 44 | return bzk + dt * 1j * (self.ky * exk - self.kx * eyk) 45 | 46 | def hampere_e1(self, exk, fxy, dt): 47 | fk = jnp.fft.fft2(fxy, axes=(0, 1)) 48 | return ( 49 | exk 50 | + self.one_over_ikx 51 | * jnp.sum(jnp.sum(fk * (jnp.exp(-1j * self.kx[..., None, None] * dt * self.vx) - 1), axis=3), axis=2) 52 | * self.dvx 53 | * self.dvy 54 | ) 55 | 56 | def hampere_e2(self, eyk, fxy, dt): 57 | fk = jnp.fft.fft2(fxy, axes=(0, 1)) 58 | return ( 59 | eyk 60 | + self.one_over_iky 61 | * jnp.sum(jnp.sum(fk * (jnp.exp(-1j * self.ky[..., None, None] * dt * self.vy) - 1), axis=3), axis=2) 62 | * self.dvx 63 | * self.dvy 64 | ) 65 | 66 | 67 | class Driver: 68 | xax: jnp.ndarray 69 | yax: jnp.ndarray 70 | 71 | def __init__(self, xax, yax): 72 | self.xax = xax 73 | self.yax = yax 74 | 75 | def get_this_pulse(self, this_pulse: dict, current_time: jnp.float64): 76 | kk = this_pulse["k0"] 77 | ww = this_pulse["w0"] 78 | dw = this_pulse["dw0"] 79 | t_L = this_pulse["t_center"] - this_pulse["t_width"] * 0.5 80 | t_R = this_pulse["t_center"] + this_pulse["t_width"] * 0.5 81 | t_wL = this_pulse["t_rise"] 82 | t_wR = this_pulse["t_rise"] 83 | x_L = this_pulse["x_center"] - this_pulse["x_width"] * 0.5 84 | x_R = this_pulse["x_center"] + this_pulse["x_width"] * 0.5 85 | x_wL = this_pulse["x_rise"] 86 | x_wR = this_pulse["x_rise"] 87 | 88 | y_L = this_pulse["y_center"] - this_pulse["y_width"] * 0.5 89 | y_R = this_pulse["y_center"] + this_pulse["y_width"] * 0.5 90 | y_wL = this_pulse["y_rise"] 91 | y_wR = this_pulse["y_rise"] 92 | 93 | envelope_t = get_envelope(t_wL, t_wR, t_L, t_R, current_time) 94 | envelope_x = get_envelope(x_wL, x_wR, x_L, x_R, self.xax) 95 | envelope_y = get_envelope(y_wL, y_wR, y_L, y_R, self.yax) 96 | 97 | return ( 98 | envelope_t 99 | * envelope_x[:, None] 100 | * envelope_y[None, :] 101 | * jnp.abs(kk) 102 | * this_pulse["a0"] 103 | * jnp.sin(kk * self.xax[:, None] - (ww + dw) * current_time) 104 | ) 105 | 106 | def __call__(self, current_time, args): 107 | """ 108 | Applies the driver function 109 | 110 | :param current_time: 111 | :param pulses: 112 | :return: 113 | """ 114 | 115 | total_dex = jnp.zeros((self.xax.size, self.yax.size)) 116 | # total_djy = np.zeros(current_time.shape + xs[0].shape + xs[1].shape) 117 | 118 | for key, pulse in args["drivers"]["ex"].items(): 119 | total_dex += self.get_this_pulse(pulse, current_time) 120 | 121 | # for key, pulse in pulses["ey"].items(): 122 | # total_djy += get_this_pulse(pulse, current_time) 123 | 124 | total_dey = jnp.zeros((self.xax.size, self.yax.size)) # , dtype=jnp.complex128) 125 | # total_dex = jnp.fft.fft2(total_dex, axes=(0, 1)) 126 | 127 | return total_dex, total_dey 128 | 129 | 130 | def get_envelope(p_wL, p_wR, p_L, p_R, ax): 131 | return 0.5 * (jnp.tanh((ax - p_L) / p_wL) - jnp.tanh((ax - p_R) / p_wR)) 132 | -------------------------------------------------------------------------------- /adept/vlasov2d/pushers/time.py: -------------------------------------------------------------------------------- 1 | import diffrax 2 | from jax import numpy as jnp 3 | 4 | from adept.vlasov2d.pushers import field, vlasov 5 | 6 | 7 | class Stepper(diffrax.Euler): 8 | def step(self, terms, t0, t1, y0, args, solver_state, made_jump): 9 | del solver_state, made_jump 10 | y1 = terms.vf(t0, y0, args) 11 | dense_info = dict(y0=y0, y1=y1) 12 | return y1, None, dense_info, None, diffrax.RESULTS.successful 13 | 14 | 15 | class VlasovFieldBase: 16 | def __init__(self, cfg): 17 | self.dt = cfg["grid"]["dt"] 18 | self.vdfdx = self.get_vdfdx(cfg) 19 | self.velocity_pusher = self.get_edfdv(cfg) 20 | self.field_solve = field.FieldSolver(cfg=cfg) 21 | self.driver = field.Driver(cfg["grid"]["x"], cfg["grid"]["y"]) 22 | 23 | def get_vdfdx(self, cfg): 24 | if cfg["solver"]["vdfdx"] == "exponential": 25 | vdfdx = vlasov.ExponentialSpatialAdvection(cfg) 26 | else: 27 | raise NotImplementedError("v df/dx: <" + cfg["solver"]["vdfdx"] + "> has not yet been implemented in JAX") 28 | 29 | return vdfdx 30 | 31 | def get_edfdv(self, cfg): 32 | if cfg["solver"]["edfdv"] == "exponential": 33 | edfdv = vlasov.ExponentialVelocityAdvection(cfg) 34 | elif cfg["solver"]["edfdv"] == "center_difference": 35 | edfdv = vlasov.CD2VelocityAdvection(cfg) 36 | # elif cfg["solver"]["edfdv"] == "weno": 37 | # edfdv = velocity.WENO(cfg) 38 | # elif cfg["solver"]["edfdv"] == "semilagrangian": 39 | # edfdv = velocity.SemiLagrangian(cfg) 40 | else: 41 | raise NotImplementedError("e df/dv: <" + cfg["solver"]["edfdv"] + "> has not yet been implemented in JAX") 42 | 43 | return edfdv 44 | 45 | 46 | class ChargeConservingMaxwell(VlasovFieldBase): 47 | """ 48 | This class contains the function that updates the state 49 | 50 | All the pushers are chosen and initialized here and a single time-step is defined here. 51 | 52 | 1. Li, Y. et al. Solving the Vlasov–Maxwell equations using Hamiltonian splitting. 53 | Journal of Computational Physics 396, 381–399 (2019). 54 | 55 | 56 | 57 | :param cfg: 58 | :return: 59 | """ 60 | 61 | def __init__(self, cfg: dict): 62 | super().__init__(cfg) 63 | self.push = cfg["solver"]["push_f"] 64 | self.dth = 0.5 * self.dt 65 | self.kx = cfg["grid"]["kx"] 66 | self.kvx = cfg["grid"]["kvx"][None, None, :, None] 67 | self.vx = cfg["grid"]["vx"][None, None, :, None] 68 | self.ky = cfg["grid"]["ky"] 69 | self.kvy = cfg["grid"]["kvy"][None, None, None, :] 70 | self.vy = cfg["grid"]["vy"][None, None, None, :] 71 | 72 | def step_vxB_1(self, bz, f, dt): 73 | fxykvx = jnp.fft.fft(f, axis=2) 74 | new_fxykvx = fxykvx * jnp.exp(-1j * dt * self.kvx * self.vy * bz[..., None, None]) 75 | return jnp.fft.ifft(new_fxykvx, axis=2) 76 | 77 | def step_vxB_2(self, bz, f, dt): 78 | fxykvy = jnp.fft.fft(f, axis=3) 79 | new_fxykvy = fxykvy * jnp.exp(1j * dt * self.kvy * self.vx * bz[..., None, None]) 80 | return jnp.fft.ifft(new_fxykvy, axis=3) 81 | 82 | def __call__(self, t: float, y: dict, args: dict) -> dict: 83 | ex, ey, bz, f = y["ex"], y["ey"], y["bz"], y["electron"] 84 | 85 | dex, dey = self.driver(t, args) 86 | 87 | # H_E 88 | # update e df/dv 89 | fhe_xy = self.velocity_pusher.edfdv(fxy=f, ex=ex + dex, ey=ey + dey, dt=self.dt) 90 | 91 | exk, eyk, bzk = jnp.fft.fft2(ex), jnp.fft.fft2(ey), jnp.fft.fft2(bz) 92 | 93 | # update b 94 | bzkp = self.field_solve.faraday(bzk=bzk, exk=exk, eyk=eyk, dt=self.dt) 95 | bzn = jnp.real(jnp.fft.ifft2(bzkp)) 96 | 97 | # H_1f 98 | # update vxB df/dv 99 | fb1_xy = self.step_vxB_1(bzn, fhe_xy, dt=self.dt) 100 | 101 | # update v1 df/dx1 102 | f1_xy = self.vdfdx.step_x(fb1_xy, dt=self.dt) 103 | # update e1 104 | e1p = self.field_solve.hampere_e1(exk=exk, fxy=fb1_xy, dt=self.dt) 105 | 106 | # H_2f 107 | # update vxB df/dv 108 | fb2_xy = self.step_vxB_2(bzn, f1_xy, dt=self.dt) 109 | 110 | # update v2 df/dx2 111 | f2_xy = self.vdfdx.step_y(fb2_xy, dt=self.dt) 112 | # update e2 113 | e2p = self.field_solve.hampere_e2(eyk=eyk, fxy=fb2_xy, dt=self.dt) 114 | 115 | # H_B 116 | # update E -> dE/dt = curl B 117 | exkp, eykp = self.field_solve.ampere(exk=e1p, eyk=e2p, bzk=bzkp, dt=self.dt) 118 | 119 | e1n, e2n = jnp.real(jnp.fft.ifft2(exkp)), jnp.real(jnp.fft.ifft2(eykp)) 120 | 121 | return {"electron": f2_xy, "ex": e1n, "ey": e2n, "bz": bzn, "dex": dex, "dey": dey} 122 | -------------------------------------------------------------------------------- /adept/vlasov2d/solver/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/adept/vlasov2d/solver/__init__.py -------------------------------------------------------------------------------- /adept/vlasov2d/solver/tridiagonal.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | 4 | from jax import numpy as jnp 5 | from jax.lax import scan 6 | 7 | 8 | class TridiagonalSolver: 9 | def __init__(self, cfg): 10 | super().__init__() 11 | self.num_unroll = 128 # cfg["solver"]["num_unroll"] 12 | 13 | @staticmethod 14 | def compute_primes(last_primes, x): 15 | """ 16 | This function is a single iteration of the forward pass in the non-in-place Thomas 17 | tridiagonal algorithm 18 | 19 | :param last_primes: 20 | :param x: 21 | :return: 22 | """ 23 | 24 | last_cp, last_dp = last_primes 25 | a, b, c, d = x 26 | cp = c / (b - a * last_cp) 27 | dp = (d - a * last_dp) / (b - a * last_cp) 28 | new_primes = jnp.stack((cp, dp)) 29 | return new_primes, new_primes 30 | 31 | @staticmethod 32 | def backsubstitution(last_x, x): 33 | """ 34 | This function is a single iteration of the backward pass in the non-in-place Thomas 35 | tridiagonal algorithm 36 | 37 | :param last_x: 38 | :param x: 39 | :return: 40 | """ 41 | cp, dp = x 42 | new_x = dp - cp * last_x 43 | return new_x, new_x 44 | 45 | def __call__(self, a, b, c, d): 46 | """ 47 | Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d. 48 | 49 | This uses the non-in-place Thomas tridiagonal algorithm. 50 | 51 | The NumPy version, on the other hand, uses the in-place algorithm. 52 | 53 | :param a: (2D float array (nx, nv)) represents the subdiagonal of the linear operator 54 | :param b: (2D float array (nx, nv)) represents the main diagonal of the linear operator 55 | :param c: (2D float array (nx, nv)) represents the super diagonal of the linear operator 56 | :param d: (2D float array (nx, nv)) represents the right hand side of the linear operator 57 | :return: 58 | """ 59 | 60 | diags_stacked = jnp.stack([arr.transpose((1, 0)) for arr in (a, b, c, d)], axis=1) 61 | _, primes = scan(self.compute_primes, jnp.zeros((2, *a.shape[:-1])), diags_stacked, unroll=self.num_unroll) 62 | _, sol = scan(self.backsubstitution, jnp.zeros(a.shape[:-1]), primes[::-1], unroll=self.num_unroll) 63 | return sol[::-1].transpose((1, 0)) 64 | -------------------------------------------------------------------------------- /configs/envelope-2d/damping.yaml: -------------------------------------------------------------------------------- 1 | solver: envelope-2d 2 | 3 | density: 4 | offset: 0.9 5 | slope: 0.3 6 | noise: 7 | min: 0.0 8 | max: 0.0 9 | type: uniform 10 | 11 | drivers: 12 | E0: 13 | w0: 2.0 14 | t_c: 500. 15 | t_w: 600. 16 | t_r: 20. 17 | x_c: 500. 18 | x_w: 600. 19 | x_r: 20. 20 | y_c: 500. 21 | y_w: 60000000. 22 | y_r: 20. 23 | k0: 1.0 24 | a0: 0.0 25 | intensity: 4.0e14 26 | E2: 27 | w0: 0.135 28 | t_c: 40. 29 | t_w: 40. 30 | t_r: 5. 31 | x_c: 0. 32 | x_w: 800. 33 | x_r: 20. 34 | y_c: 0. 35 | y_w: 2000000. 36 | y_r: 5. 37 | k0: 0.3 38 | a0: 10.0 39 | intensity: 4.0e14 40 | 41 | save: 42 | fields: 43 | t: 44 | tmin: 10.0 45 | tmax: 495.0 46 | nt: 384 47 | 48 | models: 49 | file: /Users/archis/Dev/code/ergodic/laplax/weights.eqx 50 | nu_g: 51 | activation: tanh 52 | depth: 4 53 | final_activation: tanh 54 | in_size: 3 55 | out_size: 1 56 | width_size: 8 57 | units: 58 | laser_wavelength: 351nm 59 | normalizing_temperature: 2000eV 60 | normalizing_density: 1.5e21/cc 61 | Z: 10 62 | Zp: 10 63 | 64 | plasma: 65 | wp0: 1.0 66 | nu_ei: 0.0 67 | Z: 2 68 | nb: 1.0 69 | temperature: 2.0 #keV 70 | density: 2.3e27 #m^3 71 | 72 | grid: 73 | xmin: 000.0 74 | xmax: 20.94 75 | nx: 64 76 | ymin: -4.0 77 | ymax: 4.0 78 | ny: 8 79 | tmin: 0. 80 | tmax: 500.0 81 | dt: 1.0 82 | 83 | mlflow: 84 | experiment: epw 85 | run: compare-against-vlasov 86 | 87 | terms: 88 | epw: 89 | linear: True 90 | density_gradient: False 91 | kinetic real part: False 92 | boundary: 93 | x: absorbing 94 | y: periodic 95 | trapping: 96 | active: True 97 | kld: 0.3 98 | nuee: 0.00000001 99 | source: 100 | tpd: False 101 | -------------------------------------------------------------------------------- /configs/envelope-2d/epw.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | basis: uniform 3 | noise: 4 | max: 1.0e-09 5 | min: 1.0e-10 6 | type: uniform 7 | drivers: 8 | E2: 9 | envelope: 10 | tw: 200fs 11 | tr: 25fs 12 | tc: 150fs 13 | xw: 500um 14 | xc: 10um 15 | xr: 0.2um 16 | yr: 0.2um 17 | yc: 0um 18 | yw: 50um 19 | a0: 1000 20 | k0: -10.0 21 | w0: 20.0 # 1.5 k^2 vth^2 / wp0 22 | 23 | grid: 24 | boundary_abs_coeff: 1.0e4 25 | boundary_width: 8um 26 | low_pass_filter: 0.7 27 | dt: 0.002ps 28 | dx: 20nm 29 | xmax: 6.28um 30 | tmax: 4.0ps 31 | tmin: 0.0ns 32 | ymax: 1um 33 | ymin: -1um 34 | 35 | mlflow: 36 | experiment: test-lpse 37 | run: epw-test 38 | save: 39 | fields: 40 | t: 41 | dt: 100fs 42 | tmax: 4ps 43 | tmin: 0ps 44 | x: 45 | dx: 22nm 46 | y: 47 | dy: 22nm 48 | solver: envelope-2d 49 | terms: 50 | epw: 51 | boundary: 52 | x: periodic 53 | y: periodic 54 | damping: 55 | collisions: false 56 | landau: false 57 | density_gradient: false 58 | linear: True 59 | source: 60 | noise: false 61 | tpd: false 62 | zero_mask: false 63 | units: 64 | atomic number: 40 65 | envelope density: 0.25 66 | ionization state: 6 67 | laser intensity: 3.5e+14W/cm^2 68 | laser_wavelength: 351nm 69 | reference electron temperature: 2000.0eV 70 | reference ion temperature: 1000eV 71 | -------------------------------------------------------------------------------- /configs/envelope-2d/reflection.yaml: -------------------------------------------------------------------------------- 1 | solver: envelope-2d 2 | 3 | density: 4 | offset: 0.9 5 | slope: 0.3 6 | noise: 7 | min: 0.0 8 | max: 0.0 9 | type: uniform 10 | 11 | drivers: 12 | E0: 13 | w0: 2.0 14 | t_c: 500. 15 | t_w: 600. 16 | t_r: 20. 17 | x_c: 500. 18 | x_w: 600. 19 | x_r: 20. 20 | y_c: 500. 21 | y_w: 60000000. 22 | y_r: 20. 23 | k0: 1.0 24 | a0: 0.0 25 | intensity: 4.0e14 W/cm^2 26 | E2: 27 | w0: 0.1176 28 | t_c: 70. 29 | t_w: 100. 30 | t_r: 5. 31 | x_c: 1400. 32 | x_w: 600. 33 | x_r: 20. 34 | y_c: 0. 35 | y_w: 2000000. 36 | y_r: 5. 37 | k0: 0.28 38 | a0: 5000.0 39 | intensity: 4.0e14 40 | 41 | save: 42 | fields: 43 | t: 44 | tmin: 0.0 45 | tmax: 1000.0 46 | nt: 512 47 | 48 | plasma: 49 | wp0: 1.0 50 | # Z: 2 51 | # nb: 1.0 52 | # temperature: 2.0 #keV 53 | # density: 2.3e27 #m^3 54 | 55 | units: 56 | laser_wavelength: 351nm 57 | normalizing_temperature: 4000eV 58 | normalizing_density: 2.24e21/cc 59 | gas fill: N 60 | ionization state: 6 61 | electron temperature: 4000eV 62 | 63 | grid: 64 | xmin: 000.0 65 | xmax: 4000.0 66 | nx: 1024 67 | ymin: -4.0 68 | ymax: 4.0 69 | ny: 8 70 | tmin: 0. 71 | tmax: 1000.0 72 | dt: 2.0 73 | 74 | mlflow: 75 | experiment: epw-finite-length-compare 76 | run: lpse-kinetic 77 | 78 | #models: 79 | # file: /Users/archis/Dev/code/ergodic/laplax/weights.eqx 80 | # nu_g: 81 | # activation: tanh 82 | # depth: 4 83 | # final_activation: tanh 84 | # in_size: 3 85 | # out_size: 1 86 | # width_size: 8 87 | 88 | terms: 89 | epw: 90 | linear: True 91 | density_gradient: False 92 | kinetic real part: True 93 | boundary: 94 | x: absorbing 95 | y: periodic 96 | trapping: 97 | active: False 98 | kld: 0.28 99 | nuee: 0.0000001 100 | source: 101 | tpd: False 102 | -------------------------------------------------------------------------------- /configs/envelope-2d/tpd.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | basis: linear 3 | gradient scale length: 62.5um 4 | max: 0.28 5 | min: 0.2 6 | noise: 7 | max: 1.0e-09 8 | min: 1.0e-10 9 | type: uniform 10 | drivers: 11 | E0: 12 | shape: uniform 13 | # file: s3://public-ergodic-continuum/87254/0bb528f5a431439e9f9f295bdcd6d9e7/artifacts/used_driver.pkl 14 | delta_omega_max: 0.025 15 | num_colors: 8 16 | envelope: 17 | tw: 50ps 18 | tr: 0.25ps 19 | tc: 20.05ps 20 | xr: 0.2um 21 | xw: 1000um 22 | xc: 50um 23 | yr: 0.2um 24 | yw: 1000um 25 | yc: 50um 26 | params: 27 | {} 28 | 29 | grid: 30 | boundary_abs_coeff: 1.0e4 31 | boundary_width: 1.5um 32 | low_pass_filter: 0.66 33 | dt: 0.01ps 34 | dx: 50nm 35 | tmax: 5ps 36 | tmin: 0.0ns 37 | ymax: 3um 38 | ymin: -3um 39 | mlflow: 40 | experiment: tpd 41 | run: test 42 | save: 43 | fields: 44 | t: 45 | dt: 0.2ps 46 | tmax: 5ps 47 | tmin: 0ps 48 | x: 49 | dx: 52nm 50 | y: 51 | dy: 52nm 52 | solver: envelope-2d 53 | terms: 54 | epw: 55 | boundary: 56 | x: absorbing 57 | y: periodic 58 | damping: 59 | collisions: 1.0 60 | landau: true 61 | density_gradient: true 62 | linear: true 63 | source: 64 | noise: true 65 | tpd: true 66 | zero_mask: true 67 | 68 | units: 69 | atomic number: 40 70 | envelope density: 0.25 71 | ionization state: 6 72 | laser intensity: 1.2e+15W/cm^2 73 | laser_wavelength: 351nm 74 | reference electron temperature: 2000.0eV 75 | reference ion temperature: 1000eV 76 | -------------------------------------------------------------------------------- /configs/es1d/damping.yaml: -------------------------------------------------------------------------------- 1 | solver: es-1d 2 | 3 | mlflow: 4 | experiment: es1d-epw-test 5 | run: test 6 | 7 | models: 8 | file: models/weights.eqx 9 | nu_g: 10 | activation: tanh 11 | depth: 4 12 | final_activation: tanh 13 | in_size: 3 14 | out_size: 1 15 | width_size: 8 16 | 17 | 18 | grid: 19 | nx: 16 20 | xmin: 0.0 21 | xmax: 20.94 22 | tmin: 0.0 23 | tmax: 500.0 24 | 25 | save: 26 | func: 27 | is_on: True 28 | t: 29 | tmin: 0.5 30 | tmax: 500.0 31 | nt: 1000 32 | x: 33 | is_on: True 34 | xmin: 0.0 35 | xmax: 20.94 36 | nx: 16 37 | kx: 38 | is_on: True 39 | kxmin: 0.0 40 | kxmax: 0.3 41 | nkx: 2 42 | 43 | physics: 44 | ion: 45 | gamma: 3 46 | T0: 1.0 47 | is_on: False 48 | mass: 1836.0 49 | charge: 1.0 50 | landau_damping: False 51 | trapping: 52 | is_on: False 53 | kld: 0.3 54 | nuee: 1.0e-9 55 | electron: 56 | gamma: kinetic 57 | T0: 1.0 58 | is_on: True 59 | mass: 1.0 60 | landau_damping: True 61 | charge: -1.0 62 | trapping: 63 | is_on: True 64 | model: zk 65 | kld: 0.3 66 | nuee: 1.0e-7 67 | nn: 8|8 68 | 69 | 70 | drivers: 71 | "ex": 72 | "0": 73 | "k0": 0.3 74 | "w0": 1.16 75 | "dw0": 0.0 76 | "t_c": 40.0 77 | "t_w": 40.0 78 | "t_r": 5.0 79 | "x_c": 400.0 80 | "x_w": 1000000.0 81 | "x_r": 10.0 82 | "a0": 1.e-3 83 | -------------------------------------------------------------------------------- /configs/es1d/epw.yaml: -------------------------------------------------------------------------------- 1 | drivers: 2 | ex: 3 | '0': 4 | a0: 0.01 5 | dw0: 0.0 6 | k0: 0.32 7 | t_c: 40.0 8 | t_r: 5.0 9 | t_w: 40.0 10 | w0: 1.1433284742365162 11 | x_c: 400.0 12 | x_r: 10.0 13 | x_w: 1000000.0 14 | grid: 15 | nx: 16 16 | tmax: 500.0 17 | tmin: 0.0 18 | xmax: 19.634954084936208 19 | xmin: 0.0 20 | mlflow: 21 | experiment: es1d-epw-test 22 | run: nl-fluid-noml 23 | solver: es-1d 24 | models: 25 | file: false 26 | nu_g: 27 | activation: tanh 28 | depth: 4 29 | final_activation: tanh 30 | in_size: 3 31 | out_size: 1 32 | width_size: 8 33 | physics: 34 | electron: 35 | T0: 1.0 36 | charge: -1.0 37 | gamma: kinetic 38 | is_on: true 39 | landau_damping: true 40 | mass: 1.0 41 | trapping: 42 | is_on: false 43 | kld: 0.32 44 | nn: 8|8 45 | nuee: 0.0001 46 | ion: 47 | T0: 1.0 48 | charge: 1.0 49 | gamma: 3 50 | is_on: false 51 | landau_damping: false 52 | mass: 1836.0 53 | trapping: 54 | is_on: false 55 | kld: 0.3 56 | nuee: 1.0e-09 57 | landau_damping: true 58 | save: 59 | func: 60 | is_on: true 61 | kx: 62 | is_on: false 63 | kxmax: 0.32 64 | kxmin: 0.0 65 | nkx: 2 66 | t: 67 | nt: 1000 68 | tmax: 500.0 69 | tmin: 0.5 70 | x: 71 | is_on: true 72 | nx: 16 73 | xmax: 19.634954084936208 74 | xmin: 0.0 75 | -------------------------------------------------------------------------------- /configs/es1d/es1d.yaml: -------------------------------------------------------------------------------- 1 | solver: es-1d 2 | 3 | mlflow: 4 | experiment: wavepackets-for-fluid 5 | run: single-k=0.3-finite-amplitude-fluid-trapping 6 | 7 | grid: 8 | nx: 6144 9 | xmin: 0.0 10 | xmax: 6000 11 | tmin: 0.0 12 | tmax: 1000.0 13 | 14 | save: 15 | func: 16 | is_on: True 17 | t: 18 | tmin: 0.0 19 | tmax: 1000.0 20 | nt: 2001 21 | x: 22 | is_on: True 23 | xmin: 0.0 24 | xmax: 4000.0 25 | nx: 1024 26 | kx: 27 | is_on: True 28 | kxmin: 0.0 29 | kxmax: 0.5 30 | nkx: 21 31 | 32 | models: 33 | file: models/weights.eqx 34 | nu_g: 35 | in_size: 3 36 | out_size: 1 37 | width_size: 8 38 | depth: 4 39 | activation: tanh 40 | final_activation: tanh 41 | # nu_d: 42 | # in_size: 3 43 | # out_size: 1 44 | # width_size: 8 45 | # depth: 3 46 | # activation: tanh 47 | # final_activation: tanh 48 | 49 | 50 | physics: 51 | ion: 52 | is_on: False 53 | mass: 1836.0 54 | gamma: 3.0 55 | charge: 1.0 56 | T0: 0.1 57 | landau_damping: False 58 | trapping: 59 | is_on: False 60 | kld: 0.3 61 | nuee: 1.0e-9 62 | electron: 63 | is_on: True 64 | mass: 1.0 65 | T0: 1.0 66 | gamma: 3.0 67 | charge: -1.0 68 | landau_damping: True 69 | trapping: 70 | is_on: True 71 | kld: 0.3 72 | nuee: 1.1e-5 73 | 74 | drivers: 75 | "ex": 76 | "0": 77 | "k0": 0.3 78 | "w0": 1.1598 79 | "dw0": 0.0 80 | "t_c": 80 81 | "t_w": 100 82 | "t_r": 20 83 | "x_c": 600 84 | "x_w": 800 85 | "x_r": 80 86 | "a0": 4.e-3 87 | "1": 88 | "k0": 0.26 89 | "w0": 1.11 90 | "dw0": 0.0 91 | "t_c": 400 92 | "t_w": 100 93 | "t_r": 20 94 | "x_c": 2000 95 | "x_w": 800 96 | "x_r": 80 97 | "a0": 0.e-2 98 | -------------------------------------------------------------------------------- /configs/es1d/wp.yaml: -------------------------------------------------------------------------------- 1 | drivers: 2 | ex: 3 | '0': 4 | a0: 0.005 5 | dw0: 0.0 6 | k0: 0.3 7 | t_c: 40.0 8 | t_r: 5.0 9 | t_w: 40.0 10 | w0: 1.1598464805919155 11 | x_c: 300.0 12 | x_r: 10.0 13 | x_w: 300.0 14 | grid: 15 | nx: 4096 16 | tmax: 500.0 17 | tmin: 0.0 18 | xmax: 4000 19 | xmin: 0.0 20 | mlflow: 21 | experiment: es1d-epw-test 22 | run: wp-nl-local 23 | solver: es-1d 24 | models: 25 | file: models/weights.eqx 26 | nu_g: 27 | activation: tanh 28 | depth: 4 29 | final_activation: tanh 30 | in_size: 3 31 | out_size: 1 32 | width_size: 8 33 | physics: 34 | electron: 35 | T0: 1.0 36 | charge: -1.0 37 | gamma: kinetic 38 | is_on: true 39 | landau_damping: true 40 | mass: 1.0 41 | trapping: 42 | is_on: true 43 | kld: 0.3 44 | nn: 8|8 45 | nuee: 1.0e-05 46 | ion: 47 | T0: 1.0 48 | charge: 1.0 49 | gamma: 3 50 | is_on: false 51 | landau_damping: false 52 | mass: 1836.0 53 | trapping: 54 | is_on: true 55 | kld: 0.3 56 | nuee: 1.0e-09 57 | save: 58 | func: 59 | is_on: true 60 | kx: 61 | is_on: true 62 | kxmax: 0.9 63 | kxmin: 0.3 64 | nkx: 3 65 | t: 66 | nt: 1000 67 | tmax: 500.0 68 | tmin: 0.5 69 | x: 70 | is_on: true 71 | nx: 256 72 | xmax: 1000 73 | xmin: 0.0 74 | -------------------------------------------------------------------------------- /configs/sh2d/landau_damping.yaml: -------------------------------------------------------------------------------- 1 | solver: sh-2d 2 | 3 | mlflow: 4 | experiment: sh2d-epw-test 5 | run: test 6 | 7 | units: 8 | laser_wavelength: 351nm 9 | electron temperature: 2000eV 10 | density for collisions: 1.5e21/cc 11 | Z: 10 12 | Zp: 10 13 | 14 | 15 | grid: 16 | nx: 32 17 | xmin: 0.0 18 | xmax: 20.94 19 | tmin: 0.0 20 | tmax: 200.0 21 | nt: 1536 22 | nl: 5 23 | ny: 2 24 | ymin: -32.0 25 | ymax: 32.0 26 | nv: 1024 27 | vmax: 8.0 28 | 29 | save: 30 | t: 31 | tmin: 0.5 32 | tmax: 190.0 33 | nt: 64 34 | 35 | drivers: 36 | "ex": 37 | "0": 38 | "k0": 0.3 39 | "w0": 1.16 40 | "dw0": 0.0 41 | "t_c": 40.0 42 | "t_w": 40.0 43 | "t_r": 5.0 44 | "x_c": 400.0 45 | "x_w": 1000000.0 46 | "x_r": 10.0 47 | "y_c": 400.0 48 | "y_w": 1000000.0 49 | "y_r": 10.0 50 | "a0": 1.e-8 51 | 52 | terms: 53 | vlasov: True 54 | fokker-planck: 55 | active: False 56 | flm: tridiagonal 57 | f00: chang_cooper 58 | -------------------------------------------------------------------------------- /configs/tf-1d/damping.yaml: -------------------------------------------------------------------------------- 1 | solver: tf-1d 2 | 3 | mlflow: 4 | experiment: tf1d-epw-test 5 | run: test 6 | 7 | models: 8 | file: models/weights.eqx 9 | nu_g: 10 | activation: tanh 11 | depth: 4 12 | final_activation: tanh 13 | in_size: 3 14 | out_size: 1 15 | width_size: 8 16 | 17 | units: 18 | laser_wavelength: 351nm 19 | normalizing_temperature: 2000eV 20 | normalizing_density: 1.5e21/cc 21 | Z: 10 22 | Zp: 10 23 | 24 | grid: 25 | nx: 16 26 | xmin: 0.0 27 | xmax: 20.94 28 | tmin: 0.0 29 | tmax: 500.0 30 | 31 | save: 32 | t: 33 | tmin: 0.5 34 | tmax: 500.0 35 | nt: 1000 36 | x: 37 | xmin: 0.0 38 | xmax: 20.94 39 | nx: 16 40 | kx: 41 | kxmin: 0.0 42 | kxmax: 0.3 43 | nkx: 2 44 | 45 | physics: 46 | ion: 47 | gamma: 3 48 | T0: 1.0 49 | is_on: False 50 | mass: 1836.0 51 | charge: 1.0 52 | landau_damping: False 53 | trapping: 54 | is_on: False 55 | kld: 0.3 56 | nuee: 1.0e-9 57 | electron: 58 | gamma: kinetic 59 | T0: 1.0 60 | is_on: True 61 | mass: 1.0 62 | landau_damping: True 63 | charge: -1.0 64 | trapping: 65 | is_on: True 66 | model: zk 67 | kld: 0.3 68 | nuee: 1.0e-7 69 | nn: 8|8 70 | 71 | 72 | drivers: 73 | "ex": 74 | "0": 75 | "k0": 0.3 76 | "w0": 1.16 77 | "dw0": 0.0 78 | "t_c": 40.0 79 | "t_w": 40.0 80 | "t_r": 5.0 81 | "x_c": 400.0 82 | "x_w": 1000000.0 83 | "x_r": 10.0 84 | "a0": 1.e-3 85 | -------------------------------------------------------------------------------- /configs/tf-1d/epw.yaml: -------------------------------------------------------------------------------- 1 | drivers: 2 | ex: 3 | '0': 4 | a0: 0.01 5 | dw0: 0.0 6 | k0: 0.32 7 | t_c: 40.0 8 | t_r: 5.0 9 | t_w: 40.0 10 | w0: 1.1433284742365162 11 | x_c: 400.0 12 | x_r: 10.0 13 | x_w: 1000000.0 14 | grid: 15 | nx: 16 16 | tmax: 500.0 17 | tmin: 0.0 18 | xmax: 19.634954084936208 19 | xmin: 0.0 20 | mlflow: 21 | experiment: tf1d-epw-test 22 | run: nl-fluid-noml 23 | solver: tf-1d 24 | models: 25 | file: false 26 | nu_g: 27 | activation: tanh 28 | depth: 4 29 | final_activation: tanh 30 | in_size: 3 31 | out_size: 1 32 | width_size: 8 33 | physics: 34 | electron: 35 | T0: 1.0 36 | charge: -1.0 37 | gamma: kinetic 38 | is_on: true 39 | landau_damping: true 40 | mass: 1.0 41 | trapping: 42 | is_on: false 43 | kld: 0.32 44 | nn: 8|8 45 | nuee: 0.0001 46 | ion: 47 | T0: 1.0 48 | charge: 1.0 49 | gamma: 3 50 | is_on: false 51 | landau_damping: false 52 | mass: 1836.0 53 | trapping: 54 | is_on: false 55 | kld: 0.3 56 | nuee: 1.0e-09 57 | landau_damping: true 58 | 59 | save: 60 | t: 61 | nt: 1000 62 | tmax: 500.0 63 | tmin: 0.5 64 | x: 65 | nx: 16 66 | xmax: 19.634954084936208 67 | xmin: 0.0 68 | 69 | units: 70 | laser_wavelength: 351nm 71 | normalizing_temperature: 2000eV 72 | normalizing_density: 1.5e21/cc 73 | Z: 10 74 | Zp: 10 75 | -------------------------------------------------------------------------------- /configs/tf-1d/tf1d.yaml: -------------------------------------------------------------------------------- 1 | solver: tf-1d 2 | 3 | mlflow: 4 | experiment: wavepackets-for-fluid 5 | run: single-k=0.3-finite-amplitude-fluid-trapping 6 | 7 | grid: 8 | nx: 6144 9 | xmin: 0.0 10 | xmax: 6000 11 | tmin: 0.0 12 | tmax: 1000.0 13 | 14 | save: 15 | t: 16 | tmin: 0.0 17 | tmax: 1000.0 18 | nt: 2001 19 | x: 20 | xmin: 0.0 21 | xmax: 4000.0 22 | nx: 1024 23 | kx: 24 | kxmin: 0.0 25 | kxmax: 0.5 26 | nkx: 21 27 | 28 | models: 29 | file: models/weights.eqx 30 | nu_g: 31 | in_size: 3 32 | out_size: 1 33 | width_size: 8 34 | depth: 4 35 | activation: tanh 36 | final_activation: tanh 37 | # nu_d: 38 | # in_size: 3 39 | # out_size: 1 40 | # width_size: 8 41 | # depth: 3 42 | # activation: tanh 43 | # final_activation: tanh 44 | 45 | 46 | physics: 47 | ion: 48 | is_on: False 49 | mass: 1836.0 50 | gamma: 3.0 51 | charge: 1.0 52 | T0: 0.1 53 | landau_damping: False 54 | trapping: 55 | is_on: False 56 | kld: 0.3 57 | nuee: 1.0e-9 58 | electron: 59 | is_on: True 60 | mass: 1.0 61 | T0: 1.0 62 | gamma: 3.0 63 | charge: -1.0 64 | landau_damping: True 65 | trapping: 66 | is_on: True 67 | kld: 0.3 68 | nuee: 1.1e-5 69 | 70 | drivers: 71 | "ex": 72 | "0": 73 | "k0": 0.3 74 | "w0": 1.1598 75 | "dw0": 0.0 76 | "t_c": 80 77 | "t_w": 100 78 | "t_r": 20 79 | "x_c": 600 80 | "x_w": 800 81 | "x_r": 80 82 | "a0": 4.e-3 83 | "1": 84 | "k0": 0.26 85 | "w0": 1.11 86 | "dw0": 0.0 87 | "t_c": 400 88 | "t_w": 100 89 | "t_r": 20 90 | "x_c": 2000 91 | "x_w": 800 92 | "x_r": 80 93 | "a0": 0.e-2 94 | 95 | units: 96 | laser_wavelength: 351nm 97 | normalizing_temperature: 2000eV 98 | normalizing_density: 1.5e21/cc 99 | Z: 10 100 | Zp: 10 101 | -------------------------------------------------------------------------------- /configs/tf-1d/wp.yaml: -------------------------------------------------------------------------------- 1 | drivers: 2 | ex: 3 | '0': 4 | a0: 0.005 5 | dw0: 0.0 6 | k0: 0.3 7 | t_c: 40.0 8 | t_r: 5.0 9 | t_w: 40.0 10 | w0: 1.1598464805919155 11 | x_c: 300.0 12 | x_r: 10.0 13 | x_w: 300.0 14 | 15 | grid: 16 | nx: 4096 17 | tmax: 500.0 18 | tmin: 0.0 19 | xmax: 4000 20 | xmin: 0.0 21 | 22 | mlflow: 23 | experiment: tf1d-epw-test 24 | run: wp-nl-local 25 | 26 | solver: tf-1d 27 | 28 | models: 29 | file: models/weights.eqx 30 | nu_g: 31 | activation: tanh 32 | depth: 4 33 | final_activation: tanh 34 | in_size: 3 35 | out_size: 1 36 | width_size: 8 37 | 38 | physics: 39 | electron: 40 | T0: 1.0 41 | charge: -1.0 42 | gamma: kinetic 43 | is_on: true 44 | landau_damping: true 45 | mass: 1.0 46 | trapping: 47 | is_on: false 48 | kld: 0.3 49 | nn: 8|8 50 | nuee: 1.0e-05 51 | ion: 52 | T0: 0.1 53 | charge: 1.0 54 | gamma: 3 55 | is_on: true 56 | landau_damping: false 57 | mass: 1836.0 58 | trapping: 59 | is_on: false 60 | kld: 0.3 61 | nuee: 1.0e-09 62 | 63 | save: 64 | kx: 65 | kxmax: 0.9 66 | kxmin: 0.3 67 | nkx: 3 68 | t: 69 | nt: 1000 70 | tmax: 500.0 71 | tmin: 0.5 72 | x: 73 | nx: 256 74 | xmax: 1000 75 | xmin: 0.0 76 | 77 | units: 78 | laser_wavelength: 351nm 79 | normalizing_temperature: 2000eV 80 | normalizing_density: 1.5e21/cc 81 | Z: 10 82 | Zp: 10 83 | -------------------------------------------------------------------------------- /configs/vfp-1d/epp-short.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | reference electron temperature: 3000eV 4 | reference ion temperature: 300eV 5 | reference electron density: 2.275e21/cm^3 6 | Z: 6 7 | Ion: Au+ 8 | logLambda: nrl 9 | 10 | 11 | density: 12 | quasineutrality: true 13 | species-background: 14 | noise_seed: 420 15 | noise_type: gaussian 16 | noise_val: 0.0 17 | v0: 0.0 18 | T0: 1.0 19 | m: 3.4 20 | n: 21 | basis: uniform 22 | baseline: 0.9 23 | bump_or_trough: bump 24 | center: 0m 25 | rise: 1m 26 | bump_height: 0.0 27 | width: 1m 28 | T: 29 | basis: sine 30 | baseline: 1. 31 | amplitude: 1.0e-3 32 | wavelength: 250um 33 | species-hote: 34 | noise_seed: 420 35 | noise_type: gaussian 36 | noise_val: 0.0 37 | v0: 0.0 38 | T0: 1.0 39 | m: 2.0 40 | n: 41 | basis: uniform 42 | baseline: 0.1 43 | bump_or_trough: bump 44 | center: 0m 45 | rise: 1m 46 | bump_height: 0.0 47 | width: 1m 48 | T: 49 | basis: sine 50 | baseline: 40.0 51 | amplitude: 0. 52 | wavelength: 250um 53 | 54 | grid: 55 | dt: 1fs 56 | nv: 256 57 | nx: 32 58 | tmin: 0.ps 59 | tmax: 2ps 60 | vmax: 8.0 61 | xmax: 500um 62 | xmin: 0.0um 63 | nl: 1 64 | 65 | save: 66 | fields: 67 | t: 68 | tmin: 0.0ps 69 | tmax: 2.0ps 70 | nt: 11 71 | electron: 72 | t: 73 | tmin: 0.0ps 74 | tmax: 2.0ps 75 | nt: 6 76 | 77 | solver: vfp-1d 78 | 79 | mlflow: 80 | experiment: vfp1d 81 | run: epperlein-short-hote-ee 82 | 83 | drivers: 84 | ex: {} 85 | ey: {} 86 | 87 | terms: 88 | fokker_planck: 89 | flm: 90 | ee: true 91 | f00: 92 | type: Dougherty 93 | e_solver: oshun 94 | -------------------------------------------------------------------------------- /configs/vfp-1d/hotspot.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | reference electron temperature: 500eV 4 | reference ion temperature: 500eV 5 | reference electron density: 1.5e21/cm^3 6 | Z: 30 7 | Ion: Au+ 8 | logLambda: nrl 9 | 10 | 11 | density: 12 | quasineutrality: true 13 | species-background: 14 | noise_seed: 420 15 | noise_type: gaussian 16 | noise_val: 0.0 17 | v0: 0.0 18 | T0: 1.0 19 | m: 2.0 20 | n: 21 | basis: uniform 22 | baseline: 1.0 23 | bump_or_trough: bump 24 | center: 0.0um 25 | rise: 25.0um 26 | bump_height: 0.0 27 | width: 100000.0m 28 | T: 29 | basis: tanh 30 | baseline: 1.0 31 | bump_or_trough: bump 32 | center: 50um 33 | rise: 5um 34 | bump_height: 0.1 35 | width: 20um 36 | 37 | grid: 38 | dt: 2.5fs 39 | nv: 512 40 | nx: 512 41 | tmin: 0. 42 | tmax: 1ps 43 | vmax: 8.0 44 | xmax: 100um 45 | xmin: 0.0um 46 | nl: 1 47 | 48 | save: 49 | fields: 50 | t: 51 | tmin: 0.0ps 52 | tmax: 1ps 53 | nt: 11 54 | electron: 55 | t: 56 | tmin: 0.0ps 57 | tmax: 1ps 58 | nt: 6 59 | 60 | solver: vfp-1d 61 | 62 | mlflow: 63 | experiment: vfp2d 64 | run: hotspot 65 | 66 | drivers: 67 | ex: {} 68 | ey: {} 69 | 70 | terms: 71 | fokker_planck: 72 | flm: 73 | ee: True 74 | f00: 75 | type: Dougherty 76 | e_solver: oshun 77 | -------------------------------------------------------------------------------- /configs/vlasov-1d/bump-on-tail.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0 16 | T0: 1. 17 | m: 2.0 18 | basis: sine 19 | baseline: 1. 20 | amplitude: 1.0e-4 21 | wavenumber: 0.3 22 | species-beam: 23 | noise_seed: 420 24 | noise_type: gaussian 25 | noise_val: 0.0 26 | v0: 3.0 27 | T0: 0.2 28 | m: 2.0 29 | basis: sine 30 | baseline: 0.1 31 | amplitude: -1.0e-4 32 | wavenumber: 0.3 33 | 34 | grid: 35 | dt: 0.1 36 | nv: 4096 37 | nx: 64 38 | tmin: 0. 39 | tmax: 500.0 40 | vmax: 6.4 41 | xmax: 20.94 42 | xmin: 0.0 43 | 44 | save: 45 | fields: 46 | t: 47 | tmin: 0.0 48 | tmax: 500.0 49 | nt: 1001 50 | electron: 51 | t: 52 | tmin: 0.0 53 | tmax: 500.0 54 | nt: 51 55 | 56 | solver: vlasov-1d 57 | 58 | mlflow: 59 | experiment: bump-on-tail 60 | run: nonlinear-fp 61 | 62 | drivers: 63 | ex: {} 64 | ey: {} 65 | 66 | diagnostics: 67 | diag-vlasov-dfdt: False 68 | diag-fp-dfdt: False 69 | 70 | terms: 71 | field: poisson 72 | edfdv: exponential 73 | time: sixth 74 | fokker_planck: 75 | is_on: False 76 | type: Dougherty 77 | time: 78 | baseline: 1.0e-5 79 | bump_or_trough: bump 80 | center: 0.0 81 | rise: 25.0 82 | slope: 0.0 83 | bump_height: 0.0 84 | width: 100000.0 85 | space: 86 | baseline: 1.0 87 | bump_or_trough: bump 88 | center: 0.0 89 | rise: 25.0 90 | slope: 0.0 91 | bump_height: 0.0 92 | width: 100000.0 93 | krook: 94 | is_on: True 95 | time: 96 | baseline: 1.0e-6 97 | bump_or_trough: bump 98 | center: 0.0 99 | rise: 25.0 100 | slope: 0.0 101 | bump_height: 0.0 102 | width: 100000.0 103 | space: 104 | baseline: 1.0 105 | bump_or_trough: bump 106 | center: 0.0 107 | rise: 25.0 108 | slope: 0.0 109 | bump_height: 0.0 110 | width: 100000.0 111 | -------------------------------------------------------------------------------- /configs/vlasov-1d/epw.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | baseline: 1.0 20 | bump_or_trough: bump 21 | center: 0.0 22 | rise: 25.0 23 | bump_height: 0.0 24 | width: 100000.0 25 | 26 | grid: 27 | dt: 0.1 28 | nv: 256 29 | nx: 32 30 | tmin: 0. 31 | tmax: 100.0 32 | vmax: 6.4 33 | xmax: 20.94 34 | xmin: 0.0 35 | 36 | save: 37 | fields: 38 | t: 39 | tmin: 0.0 40 | tmax: 100.0 41 | nt: 601 42 | electron: 43 | t: 44 | tmin: 0.0 45 | tmax: 100.0 46 | nt: 11 47 | diag-vlasov-dfdt: 48 | t: 49 | tmin: 0.0 50 | tmax: 100.0 51 | nt: 91 52 | diag-fp-dfdt: 53 | t: 54 | tmin: 0.0 55 | tmax: 100.0 56 | nt: 91 57 | 58 | solver: vlasov-1d 59 | 60 | mlflow: 61 | experiment: basic-epw-for-plots 62 | run: nonlinear 63 | 64 | drivers: 65 | ex: 66 | '0': 67 | a0: 1.e-2 68 | k0: 0.3 69 | t_center: 40.0 70 | t_rise: 5.0 71 | t_width: 30.0 72 | w0: 1.1598 73 | dw0: 0. 74 | x_center: 0.0 75 | x_rise: 10.0 76 | x_width: 4000000.0 77 | ey: {} 78 | 79 | diagnostics: 80 | diag-vlasov-dfdt: True 81 | diag-fp-dfdt: True 82 | 83 | terms: 84 | field: poisson 85 | edfdv: cubic-spline 86 | time: sixth 87 | fokker_planck: 88 | is_on: True 89 | type: Dougherty 90 | time: 91 | baseline: 1.0e-5 92 | bump_or_trough: bump 93 | center: 0.0 94 | rise: 25.0 95 | slope: 0.0 96 | bump_height: 0.0 97 | width: 100000.0 98 | space: 99 | baseline: 1.0 100 | bump_or_trough: bump 101 | center: 0.0 102 | rise: 25.0 103 | slope: 0.0 104 | bump_height: 0.0 105 | width: 100000.0 106 | krook: 107 | is_on: False 108 | time: 109 | baseline: 1.0 110 | bump_or_trough: bump 111 | center: 0.0 112 | rise: 25.0 113 | slope: 0.0 114 | bump_height: 0.0 115 | width: 100000.0 116 | space: 117 | baseline: 1.0 118 | bump_or_trough: bump 119 | center: 0.0 120 | rise: 25.0 121 | slope: 0.0 122 | bump_height: 0.0 123 | width: 100000.0 124 | -------------------------------------------------------------------------------- /configs/vlasov-1d/nlepw-ic.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: false 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: sine 19 | baseline: 1.0 20 | amplitude: 0.1 21 | wavenumber: 0.3 22 | 23 | grid: 24 | dt: 0.2 25 | nv: 512 26 | nx: 64 27 | tmin: 0. 28 | tmax: 1000.0 29 | vmax: 6.4 30 | xmax: 20.94 31 | xmin: 0.0 32 | 33 | save: 34 | fields: 35 | t: 36 | tmin: 0.0 37 | tmax: 1000.0 38 | nt: 1001 39 | electron: 40 | t: 41 | tmin: 0.0 42 | tmax: 1000.0 43 | nt: 21 44 | 45 | solver: vlasov-1d 46 | 47 | mlflow: 48 | experiment: nlepw 49 | run: nonlinear-fp 50 | 51 | drivers: 52 | ex: {} 53 | ey: {} 54 | 55 | diagnostics: 56 | diag-vlasov-dfdt: False 57 | diag-fp-dfdt: False 58 | 59 | terms: 60 | field: poisson 61 | edfdv: exponential 62 | time: sixth 63 | fokker_planck: 64 | is_on: False 65 | type: Dougherty 66 | time: 67 | baseline: 1.0e-5 68 | bump_or_trough: bump 69 | center: 0.0 70 | rise: 25.0 71 | slope: 0.0 72 | bump_height: 0.0 73 | width: 100000.0 74 | space: 75 | baseline: 1.0 76 | bump_or_trough: bump 77 | center: 0.0 78 | rise: 25.0 79 | slope: 0.0 80 | bump_height: 0.0 81 | width: 100000.0 82 | krook: 83 | is_on: false 84 | time: 85 | baseline: 1.0e-6 86 | bump_or_trough: bump 87 | center: 0.0 88 | rise: 25.0 89 | slope: 0.0 90 | bump_height: 0.0 91 | width: 100000.0 92 | space: 93 | baseline: 1.0 94 | bump_or_trough: bump 95 | center: 0.0 96 | rise: 25.0 97 | slope: 0.0 98 | bump_height: 0.0 99 | width: 100000.0 100 | -------------------------------------------------------------------------------- /configs/vlasov-1d/srs.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: tanh 19 | baseline: 0.001 20 | bump_or_trough: bump 21 | center: 2000.0 22 | rise: 25.0 23 | bump_height: 0.999 24 | width: 3900.0 25 | 26 | grid: 27 | dt: 0.5 28 | nv: 4096 29 | nx: 4096 30 | tmin: 0. 31 | tmax: 8000.0 32 | vmax: 6.4 33 | xmax: 4000 34 | xmin: 0.0 35 | 36 | save: 37 | fields: 38 | t: 39 | tmin: 0.0 40 | tmax: 8000.0 41 | nt: 2001 42 | electron: 43 | t: 44 | tmin: 0.0 45 | tmax: 2000.0 46 | nt: 21 47 | 48 | solver: vlasov-1d 49 | 50 | mlflow: 51 | experiment: vlasov1d-srs 52 | run: test 53 | 54 | drivers: 55 | ex: {} 56 | ey: 57 | '0': 58 | a0: 1.0 59 | k0: 1.0 60 | t_center: 4000.0 61 | t_rise: 20.0 62 | t_width: 7900.0 63 | w0: 2.79 64 | dw0: 0. 65 | x_center: 50.0 66 | x_rise: 0.1 67 | x_width: 1.0 68 | '1': 69 | a0: 1.e-6 70 | k0: 1.0 71 | t_center: 4000.0 72 | t_rise: 20.0 73 | t_width: 7900.0 74 | w0: 1.63 75 | dw0: 0. 76 | x_center: 3950.0 77 | x_rise: 0.1 78 | x_width: 1.0 79 | 80 | diagnostics: 81 | diag-vlasov-dfdt: False 82 | diag-fp-dfdt: False 83 | 84 | terms: 85 | field: poisson 86 | edfdv: exponential 87 | time: leapfrog 88 | fokker_planck: 89 | is_on: false 90 | type: Dougherty 91 | time: 92 | baseline: 1.0e-5 93 | bump_or_trough: bump 94 | center: 0.0 95 | rise: 25.0 96 | slope: 0.0 97 | bump_height: 0.0 98 | width: 100000.0 99 | space: 100 | baseline: 1.0 101 | bump_or_trough: bump 102 | center: 0.0 103 | rise: 25.0 104 | slope: 0.0 105 | bump_height: 0.0 106 | width: 100000.0 107 | krook: 108 | is_on: True 109 | time: 110 | baseline: 1.0 111 | bump_or_trough: bump 112 | center: 2000.0 113 | rise: 25.0 114 | slope: 0.0 115 | bump_height: 0.0 116 | width: 10000000.0 117 | space: 118 | baseline: 0.0 119 | bump_or_trough: trough 120 | center: 2000.0 121 | rise: 10.0 122 | slope: 0.0 123 | bump_height: 10.0 124 | width: 3800.0 125 | -------------------------------------------------------------------------------- /configs/vlasov-1d/twostream.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-electron1: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: -1.5 16 | T0: 0.2 17 | m: 2.0 18 | basis: sine 19 | baseline: 0.5 20 | amplitude: 1.0e-4 21 | wavenumber: 0.3 22 | species-electron2: 23 | noise_seed: 420 24 | noise_type: gaussian 25 | noise_val: 0.0 26 | v0: 1.5 27 | T0: 0.2 28 | m: 2.0 29 | basis: sine 30 | baseline: 0.5 31 | amplitude: -1.0e-4 32 | wavenumber: 0.3 33 | 34 | grid: 35 | dt: 0.1 36 | nv: 4096 37 | nx: 64 38 | tmin: 0. 39 | tmax: 100.0 40 | vmax: 6.4 41 | xmax: 20.94 42 | xmin: 0.0 43 | 44 | save: 45 | fields: 46 | t: 47 | tmin: 0.0 48 | tmax: 100.0 49 | nt: 51 50 | electron: 51 | t: 52 | tmin: 0.0 53 | tmax: 100.0 54 | nt: 51 55 | 56 | solver: vlasov-1d 57 | 58 | mlflow: 59 | experiment: twostream 60 | run: nonlinear-fp 61 | 62 | drivers: 63 | ex: {} 64 | ey: {} 65 | 66 | diagnostics: 67 | diag-vlasov-dfdt: False 68 | diag-fp-dfdt: False 69 | 70 | terms: 71 | field: poisson 72 | edfdv: exponential 73 | time: sixth 74 | fokker_planck: 75 | is_on: True 76 | type: Dougherty 77 | time: 78 | baseline: 1.0e-5 79 | bump_or_trough: bump 80 | center: 0.0 81 | rise: 25.0 82 | slope: 0.0 83 | bump_height: 0.0 84 | width: 100000.0 85 | space: 86 | baseline: 1.0 87 | bump_or_trough: bump 88 | center: 0.0 89 | rise: 25.0 90 | slope: 0.0 91 | bump_height: 0.0 92 | width: 100000.0 93 | krook: 94 | is_on: True 95 | time: 96 | baseline: 1.0e-6 97 | bump_or_trough: bump 98 | center: 0.0 99 | rise: 25.0 100 | slope: 0.0 101 | bump_height: 0.0 102 | width: 100000.0 103 | space: 104 | baseline: 1.0 105 | bump_or_trough: bump 106 | center: 0.0 107 | rise: 25.0 108 | slope: 0.0 109 | bump_height: 0.0 110 | width: 100000.0 111 | -------------------------------------------------------------------------------- /configs/vlasov-1d/wavepacket.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | quasineutrality: true 3 | species-background: 4 | T0: 1.0 5 | basis: uniform 6 | m: 2.0 7 | noise_seed: 420 8 | noise_type: gaussian 9 | noise_val: 0.0 10 | space-profile: 11 | baseline: 1.0 12 | bump_or_trough: bump 13 | center: 0.0 14 | rise: 25.0 15 | slope: 0.0 16 | wall_height: 0.0 17 | width: 100000.0 18 | v0: 0.0 19 | drivers: 20 | ex: 21 | '0': 22 | a0: 0.01 23 | dw0: 0.0 24 | k0: 0.28 25 | t_center: 60.0 26 | t_rise: 5.0 27 | t_width: 80.0 28 | w0: 1.1370461486881023 29 | x_center: 600.0 30 | x_rise: 20.0 31 | x_width: 900.0 32 | ey: {} 33 | grid: 34 | c_light: 10 35 | dt: 0.1 36 | nv: 6144 37 | nx: 4096 38 | tmax: 1000.0 39 | tmin: 0.0 40 | vmax: 6.4 41 | xmax: 4000.0 42 | xmin: 0.0 43 | mlflow: 44 | experiment: wavepacket-k-nu-a-vlasov-hr-sl 45 | run: vlasov-absorbing 46 | solver: vlasov-1d 47 | save: 48 | electron: 49 | t: 50 | nt: 51 51 | tmax: 1000.0 52 | tmin: 0.0 53 | fields: 54 | t: 55 | nt: 1001 56 | tmax: 1000.0 57 | tmin: 0.0 58 | 59 | diagnostics: 60 | diag-vlasov-dfdt: False 61 | diag-fp-dfdt: False 62 | 63 | terms: 64 | field: poisson 65 | edfdv: cubic-spline 66 | time: leapfrog 67 | fokker_planck: 68 | is_on: true 69 | space: 70 | baseline: 1.0 71 | bump_height: 0.0 72 | bump_or_trough: bump 73 | center: 0.0 74 | rise: 25.0 75 | slope: 0.0 76 | width: 100000.0 77 | time: 78 | baseline: 1.0e-05 79 | bump_height: 0.0 80 | bump_or_trough: bump 81 | center: 0.0 82 | rise: 25.0 83 | slope: 0.0 84 | width: 100000.0 85 | type: Dougherty 86 | krook: 87 | is_on: true 88 | space: 89 | baseline: 0.0001 90 | bump_height: 1.0 91 | bump_or_trough: trough 92 | center: 2000.0 93 | rise: 20.0 94 | slope: 0.0 95 | width: 3800.0 96 | time: 97 | baseline: 1.0 98 | bump_height: 0.0 99 | bump_or_trough: bump 100 | center: 0.0 101 | rise: 25.0 102 | slope: 0.0 103 | width: 100000.0 104 | units: 105 | Z: 10 106 | Zp: 10 107 | laser_wavelength: 351nm 108 | normalizing_density: 1.5e21/cc 109 | normalizing_temperature: 2000eV 110 | -------------------------------------------------------------------------------- /configs/vlasov-1d2v/epw.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | baseline: 1.0 20 | bump_or_trough: bump 21 | center: 0.0 22 | rise: 25.0 23 | bump_height: 0.0 24 | width: 100000.0 25 | 26 | grid: 27 | dt: 0.1 28 | nv: 1024 29 | nx: 32 30 | tmin: 0. 31 | tmax: 1000.0 32 | vmax: 8.0 33 | xmax: 20.94 34 | xmin: 0.0 35 | 36 | save: 37 | fields: 38 | t: 39 | tmin: 0.0 40 | tmax: 1000.0 41 | nt: 1001 42 | electron-fx-vx: 43 | t: 44 | tmin: 0.0 45 | tmax: 1000.0 46 | nt: 101 47 | x: {} 48 | vx: 49 | vxmin: 0.0 50 | vxmax: 8.0 51 | nvx: 512 52 | electron-full: 53 | t: 54 | tmin: 0.0 55 | tmax: 1000.0 56 | nt: 6 57 | 58 | solver: vlasov-1d2v 59 | 60 | mlflow: 61 | experiment: 1d2v-epw 62 | run: a1em2-both-exact 63 | 64 | drivers: 65 | ex: 66 | '0': 67 | a0: 1.e-2 68 | k0: 0.3 69 | t_center: 40.0 70 | t_rise: 5.0 71 | t_width: 30.0 72 | w0: 1.1598 73 | dw0: 0. 74 | x_center: 0.0 75 | x_rise: 10.0 76 | x_width: 4000000.0 77 | ey: {} 78 | 79 | terms: 80 | field: poisson 81 | edfdv: exponential 82 | time: sixth 83 | fokker_planck: 84 | nu_ee: 85 | is_on: true 86 | time: 87 | baseline: 1.0e-4 88 | bump_or_trough: bump 89 | center: 0.0 90 | rise: 25.0 91 | slope: 0.0 92 | bump_height: 0.0 93 | width: 100000.0 94 | space: 95 | baseline: 1.0 96 | bump_or_trough: bump 97 | center: 0.0 98 | rise: 25.0 99 | slope: 0.0 100 | bump_height: 0.0 101 | width: 100000.0 102 | nu_ei: 103 | nr: 384 104 | nth: 2048 105 | solver: exact-fft 106 | is_on: true 107 | time: 108 | baseline: 1.0e-4 109 | bump_or_trough: bump 110 | center: 0.0 111 | rise: 25.0 112 | slope: 0.0 113 | bump_height: 0.0 114 | width: 100000.0 115 | space: 116 | baseline: 1.0 117 | bump_or_trough: bump 118 | center: 0.0 119 | rise: 25.0 120 | slope: 0.0 121 | bump_height: 0.0 122 | width: 100000.0 123 | -------------------------------------------------------------------------------- /configs/vlasov-2d/base.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | space-profile: 20 | baseline: 1.0 21 | bump_or_trough: bump 22 | center: 0.0 23 | rise: 25.0 24 | slope: 0.0 25 | wall_height: 0.0 26 | width: 100000.0 27 | grid: 28 | num_checkpoints: 10 29 | num_fs: 2 30 | dt: 0.1 31 | nvx: 64 32 | nvy: 64 33 | nx: 24 34 | ny: 8 35 | tmin: 0. 36 | tmax: 100.0 37 | vmax: 6.4 38 | xmax: 20.94 39 | xmin: 0.0 40 | ymin: -8.0 41 | ymax: 8.0 42 | 43 | save: 44 | fields: 45 | t: 46 | tmin: 0.5 47 | tmax: 99.0 48 | nt: 64 49 | x: 50 | xmin: 0.5 51 | xmax: 19.0 52 | nx: 16 53 | y: 54 | ymin: -2.0 55 | ymax: 2.0 56 | ny: 2 57 | electron: 58 | t: 59 | tmin: 5. 60 | tmax: 95. 61 | nt: 6 62 | 63 | krook: 64 | space-profile: 65 | baseline: 0.0 66 | bump_or_trough: bump 67 | center: 2500 68 | rise: 10.0 69 | slope: 0.0 70 | wall_height: 0.0 71 | width: 48000000.0 72 | time-profile: 73 | baseline: 0.0 74 | bump_or_trough: bump 75 | center: 0.0 76 | rise: 10.0 77 | slope: 0.0 78 | wall_height: 0.0 79 | width: 100000.0 80 | 81 | machine: s-t4 82 | 83 | # solver: vlasov-2d 84 | 85 | mlflow: 86 | experiment: ld-2d2v 87 | run: envelope-driver-small-amp 88 | 89 | drivers: 90 | ex: 91 | '0': 92 | a0: 1.e-6 93 | k0: 0.3 94 | t_center: 40.0 95 | t_rise: 5.0 96 | t_width: 30.0 97 | w0: 1.1598 98 | dw0: 0. 99 | x_center: 0.0 100 | x_rise: 10.0 101 | x_width: 4000000.0 102 | y_center: 0.0 103 | y_rise: 10.0 104 | y_width: 200.0 105 | ey: {} 106 | 107 | solver: 108 | dfdt: leapfrog 109 | edfdv: exponential 110 | fp_operator: dougherty 111 | num_unroll: 32 112 | field: ampere 113 | vdfdx: exponential 114 | push_f: True 115 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | %SPHINXBUILD% >NUL 2>NUL 14 | if errorlevel 9009 ( 15 | echo. 16 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 17 | echo.installed, then set the SPHINXBUILD environment variable to point 18 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 19 | echo.may add the Sphinx directory to PATH. 20 | echo. 21 | echo.If you don't have Sphinx installed, grab it from 22 | echo.https://www.sphinx-doc.org/ 23 | exit /b 1 24 | ) 25 | 26 | if "%1" == "" goto help 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | -r ../requirements-cpu.txt 2 | 3 | # building the docs 4 | sphinx > 3.0.0 5 | sphinx_copybutton 6 | sphinx-rtd-theme >= 1.0, < 2.0 7 | sphinx-github-style >= 1.0, <= 1.1 8 | -------------------------------------------------------------------------------- /docs/source/ADEPTModule.rst: -------------------------------------------------------------------------------- 1 | ADEPTModule 2 | ------------- 3 | 4 | .. autoclass:: adept.ADEPTModule 5 | :members: 6 | :show-inheritance: 7 | -------------------------------------------------------------------------------- /docs/source/adept-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/docs/source/adept-logo.png -------------------------------------------------------------------------------- /docs/source/api.rst: -------------------------------------------------------------------------------- 1 | API Guide 2 | ---------- 3 | 4 | There are two primary high level classes. 5 | 6 | 1. `ergoExo` houses the solver and handles the mlflow logging and experiment management 7 | 2. `ADEPTModule` is base class for the solver 8 | 9 | If you wanted to create your own differentiable program that uses the ADEPT solvers, you could do 10 | 11 | .. code-block:: python 12 | 13 | from adept import ergoExo 14 | 15 | exo = ergoExo() 16 | modules = exo.setup(cfg) 17 | 18 | and 19 | 20 | .. code-block:: python 21 | 22 | sol, ppo, run_id = exo(modules) 23 | 24 | or 25 | 26 | .. code-block:: python 27 | 28 | sol, ppo, run_id = exo.val_and_grad(modules) 29 | 30 | This is analogous to `torch.nn.Module` and `eqx.Module` the `Module` workflows in general. 31 | 32 | You can see what each of those calls does in API documentation below. 33 | 34 | .. toctree:: 35 | ergoExo 36 | ADEPTModule 37 | :maxdepth: 3 38 | :caption: High level API: 39 | -------------------------------------------------------------------------------- /docs/source/dev_guide.rst: -------------------------------------------------------------------------------- 1 | Developer Guide 2 | --------------- 3 | 4 | In case you are interested in looking past the forward simulation use case, that is, if you are interested in running a program which is not just 5 | 6 | .. code-block:: bash 7 | 8 | python3 run.py --cfg config// 9 | 10 | This runs a forward simulation with the specified input parameters. It calls functions within `utils/runner.py` for this. 11 | The most important one to understand is the ``_run_`` function. Here is a stripped down pseudo-code version 12 | 13 | .. code-block:: python 14 | 15 | def run(cfg: Dict) -> Tuple[Solution, Dict]: 16 | """ 17 | This function is the main entry point for running a simulation. It takes a configuration dictionary and returns a 18 | ``diffrax.Solution`` object and a dictionary of datasets. 19 | 20 | Args: 21 | cfg: A dictionary containing the configuration for the simulation. 22 | 23 | Returns: 24 | A tuple of a Solution object and a dictionary of ``xarray.dataset``s. 25 | 26 | """ 27 | t__ = time.time() # starts the timer 28 | 29 | helpers = get_helpers(cfg["mode"]) # gets the right helper functions depending on the desired simulation 30 | 31 | with tempfile.TemporaryDirectory() as td: # creates a temporary directory to store the simulation data 32 | with open(os.path.join(td, "config.yaml"), "w") as fi: # writes the configuration to the temporary directory 33 | yaml.dump(cfg, fi) 34 | 35 | # NB - this is not yet solver specific but should be 36 | cfg = write_units(cfg, td) # writes the units to the temporary directory 37 | 38 | # NB - this is solver specific 39 | cfg = helpers.get_derived_quantities(cfg) # gets the derived quantities from the configuration 40 | misc.log_params(cfg) # logs the parameters to mlflow 41 | 42 | # NB - this is solver specific 43 | cfg["grid"] = helpers.get_solver_quantities(cfg) # gets the solver quantities from the configuration 44 | cfg = helpers.get_save_quantities(cfg) # gets the save quantities from the configuration 45 | 46 | # create the dictionary of time quantities that is given to the time integrator and save manager 47 | tqs = { 48 | "t0": cfg["grid"]["tmin"], 49 | "t1": cfg["grid"]["tmax"], 50 | "max_steps": cfg["grid"]["max_steps"], 51 | "save_t0": cfg["grid"]["tmin"], 52 | "save_t1": cfg["grid"]["tmax"], 53 | "save_nt": cfg["grid"]["tmax"], 54 | } 55 | 56 | # in case you are using ML models 57 | models = helpers.get_models(cfg["models"]) if "models" in cfg else None 58 | 59 | # initialize the state for the solver - NB - this is solver specific 60 | state = helpers.init_state(cfg) 61 | 62 | # NB - this is solver specific 63 | # Remember that we rely on the diffrax library to provide the ODE (time, usually) integrator 64 | # So we need to create the diffrax terms, solver, and save objects 65 | diffeqsolve_quants = helpers.get_diffeqsolve_quants(cfg) 66 | 67 | # run 68 | t0 = time.time() 69 | 70 | @eqx.filter_jit 71 | def _run_(these_models, time_quantities: Dict): 72 | args = {"drivers": cfg["drivers"]} 73 | if these_models is not None: 74 | args["models"] = these_models 75 | if "terms" in cfg.keys(): 76 | args["terms"] = cfg["terms"] 77 | 78 | return diffeqsolve( 79 | terms=diffeqsolve_quants["terms"], 80 | solver=diffeqsolve_quants["solver"], 81 | t0=time_quantities["t0"], 82 | t1=time_quantities["t1"], 83 | max_steps=time_quantities["max_steps"], 84 | dt0=cfg["grid"]["dt"], 85 | y0=state, 86 | args=args, 87 | saveat=SaveAt(**diffeqsolve_quants["saveat"]), 88 | ) 89 | 90 | result = _run_(models, tqs) 91 | mlflow.log_metrics({"run_time": round(time.time() - t0, 4)}) # logs the run time to mlflow 92 | 93 | t0 = time.time() 94 | # NB - this is solver specific 95 | datasets = helpers.post_process(result, cfg, td) # post-processes the result 96 | mlflow.log_metrics({"postprocess_time": round(time.time() - t0, 4)}) # logs the post-process time to mlflow 97 | mlflow.log_artifacts(td) # logs the temporary directory to mlflow 98 | 99 | mlflow.log_metrics({"total_time": round(time.time() - t__, 4)}) # logs the total time to mlflow 100 | 101 | # fin 102 | return result, datasets 103 | 104 | 105 | Here, we are heavily relying on two open-source libraries. 106 | 107 | 1. **MLFlow** as an experiment manager to log parameters, metrics, and artifacts 108 | 109 | 2. **Diffrax** to solve the ODEs 110 | -------------------------------------------------------------------------------- /docs/source/ergoExo.rst: -------------------------------------------------------------------------------- 1 | ergoExo 2 | ------------- 3 | 4 | .. autoclass:: adept.ergoExo 5 | :members: 6 | :show-inheritance: 7 | -------------------------------------------------------------------------------- /docs/source/faq.rst: -------------------------------------------------------------------------------- 1 | FAQ 2 | ----- 3 | 4 | Q. **What is novel about ADEPT?** 5 | 6 | A. ADEPT has the following features that make it a unique tool for plasma physics simulations: 7 | 8 | - Automatic Differentiation (AD) Enabled (bc of JAX) 9 | - GPU-capable (bc of XLA) 10 | - Experiment manager enabled (bc of mlflow) 11 | - Pythonic 12 | 13 | --------------------- 14 | 15 | Q. **What does AD do for us?** 16 | 17 | A. AD enables the calculation of derivatives of entire simulations, pieces of it, or anything in between. This can be used for 18 | 19 | - sensitivity analyses 20 | - parameter estimation 21 | - parameter optimization 22 | - model training 23 | 24 | --------------------- 25 | 26 | **Q. Why use MLflow?** 27 | 28 | A. MLFlow handles all the incredibly rote metadata management that computational scientists have historically either just 29 | completely ignored, written in an excel file, used a lab notebook, etc (You may always be an exception!). 30 | 31 | You can store parameters (inputs to the simulation, box size, driver parameters, etc.), metrics (run time, total electrostatic energy, temperature at t=200 ps etc.) 32 | and artifacts (the fields, distribution functions, plots, post processed quantities, configuration etc.) in a single place. 33 | 34 | This place can either be your local machine, or better yet, a remote server that is backed by a database and an object store. 35 | 36 | --------------------- 37 | 38 | **Q. Why use xarray?** 39 | 40 | A. Xarray is a great way to handle gridded data. It is performant, has a stable API, has high level plotting features. It is fairly portable, maybe not as much as HDF5, but it is a netCDF4 file so 41 | it can't be that bad! 42 | 43 | --------------------- 44 | 45 | **Q. Why use diffrax?** 46 | 47 | A. Diffrax provides the ODE integrator capabilities. However, you can, and we often do, side-step the actual time-integrator but only use diffrax for yet again, a stable API that enables us to 48 | save data in a consistent way. Diffrax lets us pass functions to the integrator, which is a great way to store custom post-processed (e.g. interpolated) quantities. You can also handle the result 49 | in the same consistent way. Yes, we could have just designed an API. But diffrax DOES also provide the time integrator. 50 | 51 | Another thing diffrax does is it has a great loop handling system that compiles much faster than anything I have written. I don't know why that is, but it is. 52 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | .. adept documentation master file, created by 2 | sphinx-quickstart on Wed Nov 8 15:35:24 2023. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | ADEPT 7 | ========= 8 | 9 | .. image:: adept-logo.png 10 | :alt: ADEPT 11 | :align: right 12 | 13 | **ADEPT** is a set of **A** utomatic **D** ifferentation **E** nabled **P** lasma **T** ransport solvers. 14 | 15 | 16 | Examples 17 | ---------- 18 | Examples can be found in the tests folder or in the adept-notebooks repository - http://github.com/ergodicio/adept-notebooks. Example configuration files are also provided in `configs/` 19 | 20 | -------------------------------------------------- 21 | 22 | Documentation 23 | ------------------ 24 | 25 | .. toctree:: 26 | usage 27 | solvers 28 | faq 29 | api 30 | :maxdepth: 2 31 | :caption: Contents: 32 | 33 | .. note:: 34 | 35 | This project is under active development. 36 | 37 | 38 | Contributing guide 39 | ------------------------ 40 | The contributing guide is in development but for now, just make an issue / pull request and we can go from there :) 41 | 42 | Citation 43 | ------------ 44 | [1] A. S. Joglekar and A. G. R. Thomas, “Machine learning of hidden variables in multiscale fluid simulation,” Mach. Learn.: Sci. Technol., vol. 4, no. 3, p. 035049, Sep. 2023, doi: 10.1088/2632-2153/acf81a. 45 | 46 | 47 | .. Indices and tables 48 | .. ================== 49 | 50 | .. * :ref:`genindex` 51 | .. * :ref:`modindex` 52 | .. * :ref:`search` 53 | -------------------------------------------------------------------------------- /docs/source/solvers.rst: -------------------------------------------------------------------------------- 1 | Available solvers 2 | ================== 3 | 4 | The solvers that we have implemented (and documented) so far are 5 | 6 | 7 | .. toctree:: 8 | solvers/vlasov1d1v 9 | solvers/lpse2d 10 | solvers/vfp1d 11 | :maxdepth: 3 12 | :caption: Implemented solvers: 13 | -------------------------------------------------------------------------------- /docs/source/solvers/datamodels/lpse2d.rst: -------------------------------------------------------------------------------- 1 | Configuration Options 2 | ====================== 3 | ``ADEPT`` needs a ``yaml`` file with the following datamodel to run the simulation. The datamodel is defined in the following class. 4 | 5 | .. autoclass:: adept._lpse2d.datamodel.ConfigModel 6 | :members: __init__ 7 | 8 | Each of the objects used to initialize this datamodel can be treated just like dictionaries. Each dictionary needs to be compiled into a megadictionary that is passed to the solver. 9 | The ``yaml`` configs accomplish this because a ``yaml`` is also a nested dictionary. The following documents those classes 10 | 11 | High Level 12 | ----------- 13 | These are the high level configuration options for the LPSE2D solver. Each of these either contains a fundamental type such as 14 | ``bool``, ``int``, ``float``, or ``str`` or is another nested ``datamodel`` which can be treated just like a dictionary. 15 | 16 | .. autoclass:: adept._lpse2d.datamodel.UnitsModel 17 | :members: __init__ 18 | 19 | .. autoclass:: adept._lpse2d.datamodel.DensityModel 20 | :members: __init__ 21 | 22 | .. autoclass:: adept._lpse2d.datamodel.GridModel 23 | :members: __init__ 24 | 25 | .. autoclass:: adept._lpse2d.datamodel.SaveModel 26 | :members: __init__ 27 | 28 | .. autoclass:: adept._lpse2d.datamodel.MLFlowModel 29 | :members: __init__ 30 | 31 | .. autoclass:: adept._lpse2d.datamodel.DriverModel 32 | :members: __init__ 33 | 34 | .. autoclass:: adept._lpse2d.datamodel.TermsModel 35 | :members: __init__ 36 | 37 | 38 | Low Level 39 | ---------- 40 | 41 | The documentation for the nested datamodels is still TBD. To investigate them further, go to the source code. 42 | -------------------------------------------------------------------------------- /docs/source/solvers/datamodels/vlasov1d.rst: -------------------------------------------------------------------------------- 1 | Configuration Options 2 | ====================== 3 | ``ADEPT`` needs a ``yaml`` file with the following datamodel to run the simulation. The datamodel is defined in the following class. 4 | 5 | .. autoclass:: adept._vlasov1d.datamodel.ConfigModel 6 | :members: __init__ 7 | 8 | Each of the objects used to initialize this datamodel can be treated just like dictionaries. Each dictionary needs to be compiled into a megadictionary that is passed to the solver. 9 | The ``yaml`` configs accomplish this because a ``yaml`` is also a nested dictionary. The following documents those classes 10 | 11 | High Level 12 | ----------- 13 | These are the high level configuration options for the Vlasov1D solver. Each of these either contains a fundamental type such as 14 | ``bool``, ``int``, ``float``, or ``str`` or is another nested ``datamodel`` which can be treated just like a dictionary. 15 | 16 | .. autoclass:: adept._vlasov1d.datamodel.UnitsModel 17 | :members: __init__ 18 | 19 | .. autoclass:: adept._vlasov1d.datamodel.DensityModel 20 | :members: __init__ 21 | 22 | .. autoclass:: adept._vlasov1d.datamodel.GridModel 23 | :members: __init__ 24 | 25 | .. autoclass:: adept._vlasov1d.datamodel.SaveModel 26 | :members: __init__ 27 | 28 | .. autoclass:: adept._vlasov1d.datamodel.MLFlowModel 29 | :members: __init__ 30 | 31 | .. autoclass:: adept._vlasov1d.datamodel.DriverModel 32 | :members: __init__ 33 | 34 | .. autoclass:: adept._vlasov1d.datamodel.TermsModel 35 | :members: __init__ 36 | 37 | 38 | Low Level 39 | ---------- 40 | 41 | The documentation for the nested datamodels is still TBD. To investigate them further, go to the source code. 42 | -------------------------------------------------------------------------------- /docs/source/solvers/lpse2d.rst: -------------------------------------------------------------------------------- 1 | Enveloped equations in Cartesian 2D 2 | ==================================== 3 | 4 | 5 | 6 | Equations and Quantities 7 | ------------------------- 8 | These equations model the evolution and interaction of the complex envelopes of light waves and plasma waves. This is faster than modeling the plasma waves using a fluid or kinetic solver along with modeling the light waves 9 | 10 | Note on pump depletion 11 | ^^^^^^^^^^^^^^^^^^^^^^^^ 12 | One can solve these equations with or without "pump depletion". "Pump depletion" is the effect of the plasma waves on the light waves. We do not currently have this implemented, so we have light waves that behave as external drivers for the plasma waves and we only model the plasma wave response. 13 | This approach is adequate for modeling laser plasma instabilities below the absolute instability threshold. 14 | 15 | Electron Plasma Waves 16 | ^^^^^^^^^^^^^^^^^^^^^^^ 17 | 18 | .. math:: 19 | \nabla \cdot \left[ i \left(\frac{\partial}{\partial t} + \nu_e^{\circ} \right) + \frac{3 v_{te}^2}{2 \omega_{p0}} \nabla^2 + \frac{\omega_{p0}}{2}\left(1-\frac{n_b(x)}{n_0}\right) \right] \textbf{E}_h = S_{TPD} + S_h 20 | 21 | 22 | Two Plasmon Decay 23 | ^^^^^^^^^^^^^^^^^^^^ 24 | 25 | .. math:: 26 | S_{\text{TPD}} \equiv \frac{e}{8 \omega_{p0} m_e} \frac{n_b(x)}{n_0} \nabla \cdot [\nabla (\textbf{E}_0 \cdot \textbf{E}_h^*) - \textbf{E}_0 \nabla\cdot \textbf{E}_h^*] e^{-i (\omega_0 - 2 \omega_{p0})t} 27 | 28 | Laser Driver 29 | ^^^^^^^^^^^^^^^ 30 | We only have a plane wave implementation for now 31 | 32 | .. math:: 33 | E_0(t, x, y) = \sum_j^{N_c} A_j ~ \exp(-i k_0 x - i \omega_0 \Delta\omega_j ~ t + \phi_j) 34 | 35 | 36 | Configuration Options 37 | ---------------------- 38 | 39 | As with all other solvers, the configuration is passed in via a ``yaml`` file. The datamodel for the configuration is documented below 40 | 41 | .. toctree:: 42 | datamodels/lpse2d 43 | :maxdepth: 3 44 | :caption: Configuration Options: 45 | -------------------------------------------------------------------------------- /docs/source/solvers/vfp1d.rst: -------------------------------------------------------------------------------- 1 | Implicit Cartesian-Spherical Harmonics Vlasov-Fokker-Planck in 1D3V 2 | ===================================================================== 3 | 4 | This set of solvers is usually used to model electron transport over hydrodynamic / collisional time-scales while the other Vlasov solvers in ``ADEPT`` are for 5 | modeling over electron plasma wave time-scales. 6 | 7 | Equations and Quantities 8 | ------------------------- 9 | 10 | TODO 11 | -------------------------------------------------------------------------------- /docs/source/solvers/vlasov1d1v.rst: -------------------------------------------------------------------------------- 1 | Explicit Cartesian-Cartesian Vlasov 1D1V 2 | ========================================= 3 | 4 | Equations and Quantities 5 | ------------------------- 6 | We solve the following coupled set of partial differential equations 7 | 8 | .. math:: 9 | \frac{\partial f}{\partial t} + v \frac{\partial f}{\partial x} + \frac{q}{m} (E + E_D) \frac{\partial f}{\partial v} &= \nu \partial_v (v f + v_0^2 \partial_v f) 10 | 11 | \partial_x E &= 1 - \int f dv 12 | 13 | where :math:`f` is the distribution function, :math:`E` is the electric field, :math:`C(f)` is the collision operator, :math:`q` is the charge, :math:`m` is the mass, and :math:`v` is the velocity. 14 | 15 | The distribution function is :math:`f = f(t, x, v)` and the electric field is :math:`E = E(t, x)` 16 | 17 | These simulations can be initialized via perturbing the distribution function or the electric field. 18 | The electric field can be "driven" using :math:`E_D` which is a user defined function of time and space. 19 | 20 | Configuration Options 21 | ---------------------- 22 | 23 | As with all other solvers, the configuration is passed in via a ``yaml`` file. The datamodel for the configuration is defined in the documentation but most will first care about 24 | how the equations themselves are solved and what the different options are. We list those first here. 25 | 26 | Velocity advection 27 | ^^^^^^^^^^^^^^^^^^^^^^^^ 28 | 1. ``exponential`` - This solver (incorrectly) assumes periodic boundaries in the velocity direction and uses a direct exponential solve such that 29 | 30 | .. math:: 31 | f^{n+1} = f^n \times \exp(A*dt) 32 | 33 | where :math:`A` is the advection operator. This is a much faster solver than the cubic-spline solver, but is less accurate. Use this if you are confident that the distribution function will be well behaved in the tails 34 | 35 | 2. ``cubic-spline`` - This is a semi-Lagrangian solver that uses a cubic-spline interpolator to advect the distribution function in velocity space. Use this if you have trouble with the exponential solver. 36 | 37 | 38 | Spatial advection 39 | ^^^^^^^^^^^^^^^^^^^^^^^^ 40 | 1. ``exponential`` - This is the only solver that is available. We only have periodic boundaries implemented in space (for the plasma) so this is perfectly fine. It is also very fast. 41 | 42 | 43 | Field solver 44 | ^^^^^^^^^^^^^^^^^^^^^^^^ 45 | 46 | 1. ``poisson`` - This is the only field solver that is available. It uses a spectral solver to solve the Poisson equation. This is the fastest and most accurate solver available. 47 | 2. ``hampere`` - This solver uses a Hamiltonian formulation of the Vlasov-Ampere system that conserves energy exactly. This is the 2nd most reliable solver. 48 | 3. ``ampere`` - This solver uses the Ampere's law to solve for the electric field. 49 | 50 | Collisions 51 | ^^^^^^^^^^^^^^^^^^^^^^^^ 52 | 1. ``none`` - No collisions are included in the simulation 53 | 2. ``lenard-bernstein`` - This solver uses the Lenard-Bernstein collision operator to include collisions in the simulation. 54 | 3. ``daugherty`` - This solver uses the Daugherty collision operator to include collisions in the simulation. 55 | 56 | 57 | Remaining Options 58 | ^^^^^^^^^^^^^^^^^^^^^^^^ 59 | The following pages document the configuration options for the Vlasov1D1V solver 60 | 61 | .. toctree:: 62 | datamodels/vlasov1d 63 | :maxdepth: 3 64 | :caption: Configuration Options: 65 | -------------------------------------------------------------------------------- /docs/source/tests.rst: -------------------------------------------------------------------------------- 1 | Tests 2 | ============= 3 | 4 | Run tests and examples 5 | ------------------------------ 6 | First install pytest via 7 | 8 | .. code-block:: console 9 | 10 | (venv) $ pip install pytest 11 | (venv) $ pytest 12 | 13 | This will run all the tests, which will likely include relatively expensive 2D2V Vlasov simulations. 14 | If you only want to see example usage, you can choose particular tests by using the `-k` flag. 15 | 16 | The package is tested against 17 | 18 | 1D1V Vlasov implementation 19 | -------------------------------- 20 | - `test_landau_damping.py` - recover the real part and imaginary part (Landau damping) of the resoance according to the kinetic dispersion relation 21 | - `test_absorbing_wave.py` - make sure the absorbing boundary conditions for the wave solver for the vector potential work correctly 22 | 23 | 24 | 1D two-fluid implementation 25 | -------------------------------- 26 | 27 | - `test_resonance.py` - recover the Bohm-Gross dispersion relation and kinetic dispersion relation just using the forward pass 28 | 29 | - `test_resonance_search.py` - recover the Bohm-Gross dispersion relation and kinetic dispersion relation using the backward pass 30 | 31 | - `test_landau_damping.py` - recover the Landau damping rate according to the kinetic dispersion relation using a phenomenological term 32 | 33 | - `test_against_vlasov.py` - recover a driven warm plasma wave simulation that is nearly identical to a Vlasov-Boltzmann simulation 34 | -------------------------------------------------------------------------------- /docs/source/usage.rst: -------------------------------------------------------------------------------- 1 | Usage 2 | ===== 3 | 4 | Installation 5 | ------------ 6 | 7 | To use adept, first install the requirements using pip: 8 | 9 | .. code-block:: console 10 | 11 | $ python3 -m venv venv 12 | $ source venv/bin/activate 13 | (venv) $ pip install -r requirements.txt 14 | 15 | or using conda: 16 | 17 | .. code-block:: console 18 | 19 | $ mamba env create -f env.yaml 20 | $ mamba activate adept 21 | (adept) $ 22 | 23 | -------------- 24 | 25 | 26 | Run an example 27 | -------------- 28 | 29 | The most common and obvious use case for ADEPT is a simple forward simulation that can be run from the command line. For example, to run a 1D1V Vlasov simulation of a driven electron plasma wave, use the following command: 30 | 31 | .. code-block:: bash 32 | 33 | (venv) $ python3 run.py --cfg configs/vlasov-1d/epw 34 | 35 | The input parameters are provided in `configs/vlasov-1d/epw.yaml`. 36 | 37 | **Access the output** 38 | 39 | The output will be saved and made accessible via MLFlow. To access it, 40 | 41 | 1. Launch an mlflow server via running ``mlflow ui`` from the command line 42 | 2. Open a web browser and navigate to http://localhost:5000 43 | 3. Click on the experiment name to see the results 44 | -------------------------------------------------------------------------------- /docs/source/usage/initialization.rst: -------------------------------------------------------------------------------- 1 | The Tanh profile 2 | ----------------------- 3 | 4 | This profile is used for spatial and temporal envelopes everywhere in this code. 5 | 6 | This is different per simulation type but there are some consistent concepts. These are that the density can be defined 7 | as a function of space, the driver (antenna or ponderomotive) can be a function of time and space, and the collision frequency 8 | can be a function of time and space. 9 | 10 | That is, you can specify, 11 | 12 | .. math:: 13 | n(x), E(t, x), \nu(t, x) 14 | 15 | Each of these profiles is provided using a tanh function that is parameterized using 16 | 17 | ``p_{wL}``, ``p_{wR}`` - the rise and fall of the flat-top 18 | 19 | ``p_w`` - width of the flat-top 20 | 21 | ``p_c`` - center of the flat-top 22 | 23 | ``p_L = p_c - p_w / 2`` - left edge of the flat-top 24 | 25 | ``p_R = p_c + p_w / 2`` - right edge of the flat-top 26 | 27 | where p can be the time or space coordinate depending on the context 28 | 29 | Then, the overall shape is given by 30 | 31 | .. math:: 32 | f(p) = 0.5 * \tanh((ax - p_L) / p_{wL}) - \tanh((ax - p_R) / p_{wR}) 33 | 34 | where ``ax`` is the time or space axis 35 | 36 | If you plot this, it looks like a flat top centered at ``p_c`` and has a width ``p_w``, and a rise and fall of ``p_{wL}`` and ``p_{wR}`` respectively. 37 | -------------------------------------------------------------------------------- /docs/source/usage/lpse2d.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/docs/source/usage/lpse2d.rst -------------------------------------------------------------------------------- /docs/source/usage/tf1d.rst: -------------------------------------------------------------------------------- 1 | Fluid-Poisson 1D 2 | ======================= 3 | 4 | To run the code for TwoFluid-Poisson 1D, use the configs in `configs/tf-1d/`. 5 | There is a `mode` option in the config file which tells `ADEPT` which solver to use. 6 | 7 | The normalized 1D Fluid-Poisson system is given by: 8 | 9 | .. note:: 10 | 11 | Work in progress 12 | 13 | .. math:: 14 | 15 | \partial_t n_e + \partial_x (n_e u_e) = 0 16 | 17 | \partial_t u_e + u_e \partial_x u_e = -\frac{\partial_x P_e}{n_e} - \frac{E}{n_e} 18 | 19 | \partial_x^2 E = n_i - n_e 20 | 21 | The ions are static but all the functionality is in place if someone wants to get them to move! 22 | 23 | **Things you might care about** 24 | 25 | 1. Infinite length (Single mode) plasma waves (Landau damping, trapping) 26 | 27 | 2. Finite length plasma waves (everything in 1. + Wavepackets) 28 | 29 | 3. Wave dynamics on density gradients (2 + density gradients) 30 | 31 | 4. Machine-learned fluid closures 32 | 33 | 34 | ----------------------- 35 | 36 | 37 | Things you can change in the config file 38 | ---------------------------------------------- 39 | 40 | Density profile 41 | ^^^^^^^^^^^^^^^ 42 | Uniform is easy. For a non-uniform profile, you have to specify the parameters of the profile. 43 | 44 | The density profile can be parameterized as a sinusoidal perturbation or a tanh flat top. The parameters to the tanh flat-top are referred to in 45 | 46 | Ponderomotive Driver 47 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 48 | You may want to a driver to drive up a wave. The envelope for this wave is specified via a tanh profile in space and in time. The other parameters to the wave 49 | are the wavenumber, frequency, amplitude, and so on. Refer to the config file for more details 50 | -------------------------------------------------------------------------------- /docs/source/usage/vlasov1d.rst: -------------------------------------------------------------------------------- 1 | Vlasov-Poisson-Fokker-Planck 1D1V 2 | ==================================== 3 | 4 | To run the code for Vlasov-Poisson 1D1V, use the configs in `configs/vlasov1d/`. 5 | There is a `mode` option in the config file which tells `ADEPT` which solver to use. 6 | 7 | The normalized 1D1V Vlasov-Poisson-Fokker-Planck system is given by: 8 | 9 | .. math:: 10 | 11 | \frac{\partial f}{\partial t} + v \frac{\partial f}{\partial x} + 12 | E \frac{\partial f}{\partial v} = C_{ee}(f) + C_{K}(f) 13 | 14 | \partial_x^2 E = 1 - \int f dv 15 | 16 | C_{ee}(f) = \nu_{ee} \frac{\partial}{\partial v} 17 | \left( v f + v_{th}^2 \partial_v f \right) 18 | 19 | C_K(f) = \nu_K (f - f_{Mx}) 20 | 21 | The ions are static but an enterprising individual should be able to reuse the electron code for the ions 22 | 23 | **Things you might care about** 24 | 25 | 1. Infinite length (Single mode) plasma waves (Landau damping, trapping) 26 | 27 | 2. Finite length plasma waves (everything in 1. + Wavepackets) 28 | 29 | 3. Wave dynamics on density gradients (2 + density gradients) 30 | 31 | 4. Stimulated Raman Scattering (3 + light waves) 32 | 33 | 34 | ----------------------- 35 | 36 | 37 | Things you can change in the config file 38 | ---------------------------------------------- 39 | 40 | Density profile 41 | ^^^^^^^^^^^^^^^ 42 | Uniform is easy. For a non-uniform profile, you have to specify the parameters of the profile. 43 | 44 | The density profile can be parameterized as a sinusoidal perturbation or a tanh flat top. The parameters to the tanh flat-top are referred to in 45 | 46 | Ponderomotive Driver 47 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 48 | You may want to a driver to drive up a wave. The envelope for this wave is specified via a tanh profile in space and in time. The other parameters to the wave 49 | are the wavenumber, frequency, amplitude, and so on. Refer to the config file for more details 50 | 51 | Collision frequency 52 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 53 | What is ``nu_ee``? It will modify the dynamics of the problem, possibly substantially depending on the distribution function dynamics. The envelope for this can also be specified 54 | in the same way as the driver. 55 | 56 | Krook frequency 57 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 58 | 59 | This is another dissipative operator but in terms of physical correspondance, this mostly just resembles sideloss if anything. Use this as a hard thermalization operator, say for boundaries 60 | as in the SRS example. 61 | -------------------------------------------------------------------------------- /docs/source/usage/vlasov1d2v.rst: -------------------------------------------------------------------------------- 1 | Vlasov-Poisson-Fokker-Planck 1D2V 2 | =================================== 3 | 4 | To run the code for Vlasov-Poisson 1D2V, use the configs in `configs/vlasov1d/`. 5 | There is a `mode` option in the config file which tells `ADEPT` which solver to use. 6 | 7 | The normalized 1D2V Vlasov-Poisson-Fokker-Planck system is given by: 8 | 9 | .. math:: 10 | 11 | \frac{\partial f}{\partial t} + v_x \frac{\partial f}{\partial x} + 12 | E_x \frac{\partial f}{\partial v_x} = C_{ee}(f) + C_{K}(f) 13 | 14 | \partial_x^2 E = 1 - \int f dv 15 | 16 | where the collision operators are 17 | 18 | .. math:: 19 | 20 | C_{ee}(f) = \nu_{ee} \frac{\partial}{\partial \mathbf{v}} \cdot 21 | \left( \mathbf{v} f + v_{th}^2 \partial_\mathbf{v} f \right) 22 | 23 | C_{ei}(f) = \nu_{ei}(v) \left[\partial_{v_x} \left(-v_y^2 \partial_{v_x} f + v_x v_y \partial_{v_y} f\right) + \partial_{v_y} \left(v_x v_y \partial_{v_x} f - v_x^2 \partial_{v_y} f\right)\right] 24 | 25 | C_K(f) = \nu_K (f - f_{Mx}) 26 | 27 | The ions are static but an enterprising individual should be able to reuse the electron code for the ions 28 | 29 | **Things you might care about** 30 | 31 | 1. Infinite length (Single mode) plasma waves (Landau damping, trapping) 32 | 33 | 2. Finite length plasma waves (everything in 1. + Wavepackets) 34 | 35 | 3. Wave dynamics on density gradients (2 + density gradients) 36 | 37 | 38 | ----------------------- 39 | 40 | 41 | Things you can change in the config file 42 | ---------------------------------------------- 43 | 44 | Density profile 45 | ^^^^^^^^^^^^^^^ 46 | Uniform is easy. For a non-uniform profile, you have to specify the parameters of the profile. 47 | 48 | The density profile can be parameterized as a sinusoidal perturbation or a tanh flat top. The parameters to the tanh flat-top are referred to in 49 | 50 | Ponderomotive Driver 51 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 52 | You may want to a driver to drive up a wave. The envelope for this wave is specified via a tanh profile in space and in time. The other parameters to the wave 53 | are the wavenumber, frequency, amplitude, and so on. Refer to the config file for more details 54 | 55 | Collision frequency 56 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 57 | What are ``nu_ee`` and ``nu_ei``? It will modify the dynamics of the problem, possibly substantially depending on the distribution function dynamics. The envelope for this can also be specified 58 | in the same way as the driver. 59 | 60 | Krook frequency 61 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 62 | 63 | This is another dissipative operator but in terms of physical correspondance, this mostly just resembles sideloss if anything. Use this as a hard thermalization operator, say for boundaries 64 | as in the SRS example. 65 | -------------------------------------------------------------------------------- /docs/source/usage/vlasov2d2v.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/docs/source/usage/vlasov2d2v.rst -------------------------------------------------------------------------------- /env.yaml: -------------------------------------------------------------------------------- 1 | name: adept 2 | channels: 3 | - conda-forge 4 | dependencies: 5 | - python=3.12 6 | - ipython 7 | - ipykernel 8 | - pip 9 | - pip: 10 | - "jax[cpu]" 11 | - jaxopt 12 | - numpy 13 | - scipy 14 | - matplotlib 15 | - pyhdf 16 | - xlrd 17 | - pyyaml 18 | - mlflow 19 | - boto3 20 | - typing-extensions 21 | - optax 22 | - tqdm 23 | - xarray 24 | - pint 25 | - diffrax 26 | - flatdict 27 | - plasmapy 28 | - parsl 29 | - interpax 30 | - tabulate 31 | - scienceplots 32 | -------------------------------------------------------------------------------- /env_gpu.yaml: -------------------------------------------------------------------------------- 1 | name: adept-gpu 2 | channels: 3 | - conda-forge 4 | dependencies: 5 | - python=3.12 6 | - ipython 7 | - ipykernel 8 | - pip 9 | - pip: 10 | # works for regular pip packages 11 | - --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html 12 | - jaxlib==0.4.28+cuda12.cudnn89 13 | - jax==0.4.28 14 | - jaxopt 15 | - numpy 16 | - scipy 17 | - matplotlib 18 | - pyhdf 19 | - xlrd 20 | - pyyaml 21 | - mlflow 22 | - boto3 23 | - typing-extensions 24 | - optax 25 | - tqdm 26 | - xarray 27 | - pint 28 | - diffrax 29 | - flatdict 30 | - plasmapy 31 | - parsl 32 | - interpax 33 | - tabulate 34 | - scienceplots 35 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "adept" 3 | description = "Automatic-Differentiation-Enabled Plasma Transport in JAX" 4 | readme = "README.md" 5 | authors = [ 6 | {name = "Archis Joglekar/Ergodic LLC", email="archis@ergodic.io"}, 7 | ] 8 | license = "MIT" 9 | license-files = ["LICENSE*"] 10 | requires-python = ">=3.11" 11 | dependencies = [ 12 | "diffrax", 13 | "matplotlib", 14 | "scipy", 15 | "numpy", 16 | "tqdm", 17 | "xarray", 18 | "mlflow", 19 | "flatdict", 20 | "h5netcdf", 21 | "optax", 22 | "boto3", 23 | "pint", 24 | "plasmapy", 25 | "interpax", 26 | "tabulate", 27 | "pydantic", 28 | "scienceplots", 29 | "jax" 30 | ] 31 | dynamic = ["version"] 32 | 33 | [project.urls] 34 | Homepage = "https://github.com/ergodicio/adept" 35 | Documentation = "https://adept.readthedocs.io/en/latest/" 36 | 37 | [project.optional-dependencies] 38 | cpu = [] 39 | gpu = ["jax[cuda12]"] 40 | docs = [ 41 | "sphinx", 42 | "sphinx_autodoc_typehints", 43 | "sphinx_copybutton", 44 | "sphinx-rtd-theme", 45 | "sphinx-github-style", 46 | "adept" 47 | ] 48 | dev = [ 49 | "pre-commit", 50 | "pytest", 51 | "pytest-cov", 52 | "typeguard", 53 | "adept[docs]", 54 | ] 55 | 56 | [build-system] 57 | requires = ["hatchling", "hatch-vcs", "toml"] 58 | build-backend = "hatchling.build" 59 | 60 | [tool.hatch.version] 61 | source = "vcs" 62 | 63 | [tool.hatch.build.hooks.vcs] 64 | version-file = "adept/_version.py" 65 | 66 | [tool.pytest.ini_options] 67 | # Test type hints w/ typeguard - currently disabled 68 | # addopts = ["--typeguard-packages=adept"] 69 | testpaths = ["tests"] 70 | filterwarnings = [ 71 | "error", 72 | # ignored by default 73 | "ignore::DeprecationWarning", 74 | "ignore::PendingDeprecationWarning", 75 | "ignore::ImportWarning", 76 | # raised by Cython, usually harmless 77 | "ignore:numpy.dtype size changed", 78 | "ignore:numpy.ufunc size changed", 79 | # sometimes, dependencies leak resources 80 | "ignore:.*socket\\.socket.*:pytest.PytestUnraisableExceptionWarning", 81 | ] 82 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | filterwarnings = ignore::DeprecationWarning 3 | pythonpath = . 4 | testpaths = tests 5 | -------------------------------------------------------------------------------- /requirements-cpu.txt: -------------------------------------------------------------------------------- 1 | jax 2 | diffrax 3 | matplotlib 4 | scipy 5 | numpy 6 | tqdm 7 | xarray 8 | mlflow 9 | flatdict 10 | h5netcdf 11 | optax 12 | jaxopt 13 | boto3 14 | pint 15 | plasmapy 16 | tabulate 17 | interpax 18 | pydantic 19 | scienceplots 20 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | jax[cuda12] 2 | diffrax 3 | matplotlib 4 | scipy 5 | numpy 6 | tqdm 7 | xarray 8 | mlflow 9 | flatdict 10 | h5netcdf 11 | optax 12 | jaxopt 13 | boto3 14 | pint 15 | plasmapy 16 | tabulate 17 | interpax 18 | pydantic 19 | scienceplots 20 | -------------------------------------------------------------------------------- /ruff.toml: -------------------------------------------------------------------------------- 1 | # Set to the lowest supported Python version. 2 | target-version = "py311" 3 | 4 | # Set the target line length for formatting. 5 | line-length = 120 6 | 7 | # Exclude a variety of commonly ignored directories. 8 | extend-exclude = [ 9 | "_version.py", 10 | "lpse2d.py", 11 | "tf1d.py", 12 | "vlasov1d.py" 13 | ] 14 | 15 | src = ["."] 16 | 17 | [lint] 18 | # Select and/or ignore rules for linting. 19 | # Full list of available rules: https://docs.astral.sh/ruff/rules/ 20 | extend-select = [ 21 | "B", # Flake8 bugbear 22 | "E", # Pycodestyle errors 23 | "F", # Pyflakes 24 | "I", # Isort 25 | "NPY", # Numpy 26 | "RUF", # Ruff-specific rules 27 | "UP", # Pyupgrade 28 | "W", # Pycodestyle warnings 29 | ] 30 | 31 | # Below fix only unsafe as could remove comments 32 | extend-safe-fixes = ["UP008"] # Use `super()` instead of `super(__class__, self)` 33 | 34 | ignore = [ 35 | "B007", # Loop control variable (e.g. `i` or `key`) not used within loop body 36 | "B008", # Do not perform function call `np.ones` in argument defaults; instead, 37 | # perform the call within the function, or read the default from a module-level singleton variable 38 | "E402", # Module level import not at top of file 39 | "E731", # Do not assign a lambda expression, use a def 40 | "F401", # module imported in `__init__.py` but unused; consider removing or adding to `__all__` 41 | "F841", # Local variable (e.g. mlflow_run) is assigned to but never used 42 | "NPY002", # Replace legacy `np.random.uniform` call with `np.random.Generator` 43 | "NPY201", # `np.trapz` will be removed in NumPy 2.0. Use `numpy.trapezoid` on NumPy 2.0 44 | "RUF005", # Consider unpacking instad of concatenation 45 | ] 46 | allowed-confusables = [ 47 | "–", # EN DASH 48 | ] 49 | 50 | [lint.extend-per-file-ignores] 51 | # Allow use of \int string and I variable name in SNB model for now, can address later 52 | "lagradept/lrh1d/pushers/utils/transport_coeffs.py" = ["E741", "W605"] 53 | 54 | [lint.pycodestyle] 55 | max-line-length = 120 56 | max-doc-length = 140 # Feel free to shorten but this required a lot of manual work to fully address 57 | 58 | [format] 59 | # Enable reformatting of code snippets in docstrings. 60 | docstring-code-format = true 61 | -------------------------------------------------------------------------------- /run.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" 5 | 6 | from jax import config 7 | 8 | config.update("jax_enable_x64", True) 9 | # config.update("jax_disable_jit", True) 10 | 11 | import yaml 12 | 13 | from adept import ergoExo 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser(description="Automatic Differentiation Enabled Plasma Transport") 17 | parser.add_argument("--cfg", help="enter path to cfg") 18 | parser.add_argument("--run_id", help="enter run_id to continue") 19 | args = parser.parse_args() 20 | 21 | exo = ergoExo() 22 | 23 | if args.run_id is None: 24 | with open(f"{os.path.join(os.getcwd(), args.cfg)}.yaml") as fi: 25 | cfg = yaml.safe_load(fi) 26 | modules = exo.setup(cfg=cfg) 27 | sol, post_out, run_id = exo(modules) 28 | 29 | else: 30 | exo.run_job(args.run_id, nested=None) 31 | run_id = args.run_id 32 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import jax 2 | 3 | jax.config.update("jax_enable_x64", True) 4 | -------------------------------------------------------------------------------- /tests/test_base/configs/example.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | space-profile: 20 | baseline: 1.0 21 | bump_or_trough: bump 22 | center: 0.0 23 | rise: 25.0 24 | slope: 0.0 25 | bump_height: 0.0 26 | width: 100000.0 27 | 28 | grid: 29 | dt: 0.25 30 | nv: 512 31 | nx: 32 32 | tmin: 0. 33 | tmax: 480.0 34 | vmax: 6.4 35 | xmax: 20.94 36 | xmin: 0.0 37 | 38 | save: 39 | fields: 40 | t: 41 | tmin: 0.0 42 | tmax: 480.0 43 | nt: 961 44 | electron: 45 | t: 46 | tmin: 0.0 47 | tmax: 480.0 48 | nt: 9 49 | 50 | solver: vlasov-1d 51 | 52 | mlflow: 53 | experiment: vlasov1d 54 | run: test 55 | 56 | drivers: 57 | ex: 58 | '0': 59 | a0: 1.e-6 60 | k0: 0.3 61 | t_center: 40.0 62 | t_rise: 5.0 63 | t_width: 30.0 64 | w0: 1.1598 65 | dw0: 0. 66 | x_center: 0.0 67 | x_rise: 10.0 68 | x_width: 4000000.0 69 | ey: {} 70 | 71 | terms: 72 | field: poisson 73 | edfdv: cubic-spline 74 | time: leapfrog 75 | fokker_planck: 76 | is_on: True 77 | type: Dougherty 78 | time: 79 | baseline: 1.0 80 | bump_or_trough: bump 81 | center: 0.0 82 | rise: 25.0 83 | slope: 0.0 84 | bump_height: 0.0 85 | width: 100000.0 86 | space: 87 | baseline: 1.0e-6 88 | bump_or_trough: bump 89 | center: 0.0 90 | rise: 25.0 91 | slope: 0.0 92 | bump_height: 0.0 93 | width: 100000.0 94 | krook: 95 | is_on: True 96 | time: 97 | baseline: 1.0 98 | bump_or_trough: bump 99 | center: 0.0 100 | rise: 25.0 101 | slope: 0.0 102 | bump_height: 0.0 103 | width: 100000.0 104 | space: 105 | baseline: 1.0e-6 106 | bump_or_trough: bump 107 | center: 0.0 108 | rise: 25.0 109 | slope: 0.0 110 | bump_height: 0.0 111 | width: 100000.0 112 | 113 | diagnostics: 114 | diag-vlasov-dfdt: False 115 | diag-fp-dfdt: False 116 | -------------------------------------------------------------------------------- /tests/test_base/test_ergoExo.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | 3 | from adept import ergoExo 4 | 5 | 6 | def test_reuse_config_dict(): 7 | with open("tests/test_base/configs/example.yaml") as file: 8 | cfg = yaml.safe_load(file) 9 | 10 | exo = ergoExo() 11 | exo.setup(cfg) 12 | 13 | exo = ergoExo() 14 | exo.setup(cfg) 15 | -------------------------------------------------------------------------------- /tests/test_lpse2d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_lpse2d/__init__.py -------------------------------------------------------------------------------- /tests/test_lpse2d/configs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_lpse2d/configs/__init__.py -------------------------------------------------------------------------------- /tests/test_lpse2d/configs/epw.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | basis: uniform 3 | noise: 4 | max: 1.0e-09 5 | min: 1.0e-10 6 | type: uniform 7 | drivers: 8 | E2: 9 | envelope: 10 | tw: 200fs 11 | tr: 25fs 12 | tc: 150fs 13 | xw: 500um 14 | xc: 10um 15 | xr: 0.2um 16 | yr: 0.2um 17 | yc: 0um 18 | yw: 50um 19 | a0: 1000 20 | k0: -10.0 21 | w0: 20.0 # 1.5 k^2 vth^2 / wp0 22 | 23 | grid: 24 | boundary_abs_coeff: 1.0e4 25 | boundary_width: 8um 26 | low_pass_filter: 0.7 27 | dt: 0.002ps 28 | dx: 20nm 29 | xmax: 6.28um 30 | tmax: 4.0ps 31 | tmin: 0.0ns 32 | ymax: 1um 33 | ymin: -1um 34 | 35 | mlflow: 36 | experiment: test-lpse 37 | run: epw-test 38 | save: 39 | fields: 40 | t: 41 | dt: 100fs 42 | tmax: 4ps 43 | tmin: 0ps 44 | x: 45 | dx: 22nm 46 | y: 47 | dy: 22nm 48 | solver: envelope-2d 49 | terms: 50 | epw: 51 | boundary: 52 | x: periodic 53 | y: periodic 54 | damping: 55 | collisions: false 56 | landau: false 57 | density_gradient: false 58 | linear: True 59 | source: 60 | noise: false 61 | tpd: false 62 | zero_mask: false 63 | units: 64 | atomic number: 40 65 | envelope density: 0.25 66 | ionization state: 6 67 | laser intensity: 3.5e+14W/cm^2 68 | laser_wavelength: 351nm 69 | reference electron temperature: 2000.0eV 70 | reference ion temperature: 1000eV 71 | -------------------------------------------------------------------------------- /tests/test_lpse2d/configs/resonance_search.yaml: -------------------------------------------------------------------------------- 1 | solver: envelope-2d 2 | 3 | density: 4 | basis: uniform 5 | offset: 1.0 6 | slope: 0.0 7 | noise: 8 | min: 0.0 9 | max: 0.0 10 | type: uniform 11 | 12 | drivers: 13 | E0: 14 | w0: 2.0 15 | t_c: 500. 16 | t_w: 600. 17 | t_r: 20. 18 | x_c: 500. 19 | x_w: 600. 20 | x_r: 20. 21 | y_c: 500. 22 | y_w: 60000000. 23 | y_r: 20. 24 | k0: 1.0 25 | a0: 0.0 26 | intensity: 4.0e14 27 | E2: 28 | w0: 0.1536 29 | t_c: 60. 30 | t_w: 80. 31 | t_r: 5. 32 | x_c: 0. 33 | x_w: 800. 34 | x_r: 20. 35 | y_c: 0. 36 | y_w: 2000000. 37 | y_r: 5. 38 | k0: 0.32 39 | a0: 125 40 | intensity: 4.0e14 41 | 42 | save: 43 | t: 44 | tmin: 10.0 45 | tmax: 195.0 46 | nt: 384 47 | 48 | plasma: 49 | wp0: 1.0 50 | nu_ei: 0.0 51 | Z: 2 52 | nb: 1.0 53 | temperature: 2.0 #keV 54 | density: 2.3e27 #m^3 55 | 56 | units: 57 | laser_wavelength: 351nm 58 | normalizing_temperature: 2000eV 59 | normalizing_density: 1.5e21/cc 60 | Z: 10 61 | Zp: 10 62 | 63 | grid: 64 | xmin: 000.0 65 | xmax: 16.535 66 | nx: 64 67 | ymin: -4.0 68 | ymax: 4.0 69 | ny: 8 70 | tmin: 0. 71 | tmax: 200.0 72 | dt: 0.5 73 | 74 | mlflow: 75 | experiment: epw 76 | run: compare-against-vlasov 77 | 78 | terms: 79 | epw: 80 | linear: True 81 | density_gradient: False 82 | kinetic real part: True 83 | boundary: 84 | x: periodic 85 | y: periodic 86 | trapping: 87 | kld: 0.3 88 | active: False 89 | source: 90 | tpd: False 91 | -------------------------------------------------------------------------------- /tests/test_lpse2d/configs/tpd.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | basis: linear 3 | gradient scale length: 200.0um 4 | max: 0.32 5 | min: 0.23 6 | noise: 7 | max: 1.0e-09 8 | min: 1.0e-10 9 | type: uniform 10 | drivers: 11 | E0: 12 | params: {} 13 | shape: uniform 14 | # file: s3://public-ergodic-continuum/87254/0bb528f5a431439e9f9f295bdcd6d9e7/artifacts/used_driver.pkl 15 | delta_omega_max: 0.015 16 | num_colors: 1 17 | envelope: 18 | tw: 40ps 19 | tr: 0.1ps 20 | tc: 20.25ps 21 | xr: 0.2um 22 | xw: 1000um 23 | xc: 50um 24 | yr: 0.2um 25 | yw: 1000um 26 | yc: 50um 27 | 28 | grid: 29 | boundary_abs_coeff: 1.0e4 30 | boundary_width: 3um 31 | low_pass_filter: 0.6 32 | dt: 0.002ps 33 | dx: 30nm 34 | tmax: 15.0ps 35 | tmin: 0.0ns 36 | ymax: 3um 37 | ymin: -3um 38 | mlflow: 39 | experiment: tpd 40 | run: 1.5e14 41 | save: 42 | fields: 43 | t: 44 | dt: 1.0ps 45 | tmax: 15ps 46 | tmin: 0ps 47 | x: 48 | dx: 80nm 49 | y: 50 | dy: 128nm 51 | solver: envelope-2d 52 | terms: 53 | epw: 54 | boundary: 55 | x: absorbing 56 | y: periodic 57 | damping: 58 | collisions: false 59 | landau: true 60 | density_gradient: true 61 | linear: true 62 | source: 63 | noise: true 64 | tpd: true 65 | zero_mask: true 66 | units: 67 | atomic number: 40 68 | envelope density: 0.25 69 | ionization state: 6 70 | laser intensity: 2.0e+14W/cm^2 71 | laser_wavelength: 351nm 72 | reference electron temperature: 2000.0eV 73 | reference ion temperature: 1000eV 74 | -------------------------------------------------------------------------------- /tests/test_lpse2d/test_epw_frequency.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | import yaml 4 | from numpy import testing 5 | 6 | from adept import ergoExo 7 | from adept.electrostatic import get_nlfs 8 | 9 | 10 | def _real_part_(): 11 | with open("tests/test_lpse2d/configs/epw.yaml") as fi: 12 | cfg = yaml.safe_load(fi) 13 | 14 | exo = ergoExo() 15 | _ = exo.setup(cfg) 16 | 17 | with open("tests/test_lpse2d/configs/epw.yaml") as fi: 18 | cfg = yaml.safe_load(fi) 19 | # modify config 20 | rand_k0 = np.random.uniform(10, 30) 21 | rand_scalar = np.random.uniform(0.5, 2.0) 22 | lambda_w = round(float(2 * np.pi / rand_k0), 2) 23 | k0 = 2 * np.pi / lambda_w 24 | w0 = ( 25 | 1.5 26 | * k0**2.0 27 | * exo.adept_module.cfg["units"]["derived"]["vte_sq"] 28 | / exo.adept_module.cfg["units"]["derived"]["wp0"] 29 | ) 30 | cfg["drivers"]["E2"]["k0"] = float(k0) 31 | cfg["drivers"]["E2"]["w0"] = w0 * rand_scalar 32 | cfg["grid"]["xmax"] = f"{10 * float(2 * np.pi / k0)}um" 33 | cfg["grid"]["tmax"] = "1ps" 34 | cfg["save"]["fields"]["t"]["tmax"] = "1ps" 35 | cfg["save"]["fields"]["t"]["dt"] = "5fs" 36 | 37 | exo = ergoExo() 38 | modules = exo.setup(cfg) 39 | sol, ppo, mlrunid = exo(modules) 40 | 41 | # get damping rate out of ppo 42 | # ppo["fields"]["epw"] = ppo["fields"]["epw"].view(np.complex128) 43 | flds = ppo["x"] 44 | slc = np.fft.fft(np.real(flds["phi"][:, :, 0]).data, axis=1)[:, 10] 45 | dt = flds.coords["t (ps)"].data[2] - flds.coords["t (ps)"].data[1] 46 | amplitude_envelope, instantaneous_frequency_smooth = get_nlfs(slc, dt) 47 | 48 | actual = np.mean(instantaneous_frequency_smooth[100:160]) 49 | 50 | testing.assert_allclose(desired=w0, actual=actual, rtol=0.1) 51 | 52 | 53 | def _imaginary_part_(): 54 | with open("tests/test_lpse2d/configs/epw.yaml") as fi: 55 | cfg = yaml.safe_load(fi) 56 | 57 | exo = ergoExo() 58 | _ = exo.setup(cfg) 59 | 60 | with open("tests/test_lpse2d/configs/epw.yaml") as fi: 61 | cfg = yaml.safe_load(fi) 62 | # modify config 63 | rand_k0 = np.random.uniform(22, 28) 64 | lambda_w = round(float(2 * np.pi / rand_k0), 2) 65 | k0 = 2 * np.pi / lambda_w 66 | cfg["drivers"]["E2"]["k0"] = float(k0) 67 | cfg["drivers"]["E2"]["w0"] = ( 68 | 1.5 69 | * k0**2.0 70 | * exo.adept_module.cfg["units"]["derived"]["vte_sq"] 71 | / exo.adept_module.cfg["units"]["derived"]["wp0"] 72 | ) 73 | cfg["grid"]["xmax"] = f"{10 * float(2 * np.pi / k0)}um" 74 | cfg["terms"]["epw"]["damping"]["landau"] = True 75 | 76 | exo = ergoExo() 77 | modules = exo.setup(cfg) 78 | sol, ppo, mlrunid = exo(modules) 79 | 80 | desired = ( 81 | np.sqrt(np.pi / 8) 82 | * exo.cfg["units"]["derived"]["wp0"] ** 4 83 | * k0**-3.0 84 | / (exo.cfg["units"]["derived"]["vte_sq"] ** 1.5) 85 | * np.exp(-(exo.cfg["units"]["derived"]["wp0"] ** 2.0) / k0**2.0 / (2 * exo.cfg["units"]["derived"]["vte_sq"])) 86 | ) 87 | # get damping rate out of ppo 88 | flds = ppo["x"] 89 | amplitude_envelope = np.abs(np.fft.fft(np.real(flds["phi"][:, :, 0]).data, axis=1)[:, 10]) 90 | dt = flds.coords["t (ps)"].data[2] - flds.coords["t (ps)"].data[1] 91 | # amplitude_envelope, instantaneous_frequency_smooth = get_nlfs(slc, dt) 92 | 93 | actual = np.mean((np.gradient(amplitude_envelope) / amplitude_envelope)[20:40]) / dt 94 | 95 | # mlflow.log_metrics({"actual damping rate": actual, "desired damping rate": desired}) 96 | testing.assert_allclose(desired=desired, actual=-actual, rtol=0.35) 97 | 98 | 99 | @pytest.mark.parametrize("test_func", [_real_part_, _imaginary_part_]) 100 | def test_epw_frequency(test_func): 101 | test_func() 102 | 103 | 104 | if __name__ == "__main__": 105 | test_epw_frequency(_real_part_) 106 | test_epw_frequency(_imaginary_part_) 107 | -------------------------------------------------------------------------------- /tests/test_lpse2d/test_tpd_threshold.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | from jax import devices 4 | 5 | from adept.lpse2d import calc_threshold_intensity 6 | 7 | 8 | def run_once(L, Te, I0): 9 | import yaml 10 | 11 | from adept import ergoExo 12 | 13 | with open("tests/test_lpse2d/configs/tpd.yaml") as fi: 14 | cfg = yaml.safe_load(fi) 15 | 16 | cfg["units"]["laser intensity"] = f"{I0}e14 W/cm^2" 17 | cfg["density"]["gradient scale length"] = f"{L}um" 18 | cfg["units"]["reference electron temperature"] = f"{Te}keV" 19 | cfg["mlflow"]["run"] = f"{I0}W/cm^2" # I0 20 | cfg["mlflow"]["experiment"] = f"I2-threshold-L={L}um, Te={Te}keV" 21 | 22 | exo = ergoExo() 23 | modules = exo.setup(cfg) 24 | sol, ppo, mlrunid = exo(modules) 25 | es = ppo["metrics"]["log10_total_e_sq"] 26 | 27 | return es 28 | 29 | 30 | def test_threshold(): 31 | if not any(["gpu" == device.platform for device in devices()]): 32 | pytest.skip("Takes too long without a GPU") 33 | else: 34 | ess = [] 35 | c = 3e8 36 | lam0 = 351e-9 37 | w0 = 2 * np.pi * c / lam0 / 1e12 # 1/ps 38 | for _ in range(5): 39 | L = round(np.random.uniform(200, 500)) 40 | Te = round(np.random.uniform(1, 4), 2) 41 | 42 | It = calc_threshold_intensity(Te, L, w0) 43 | I_scan = np.linspace(0.905, 1.095, 19) * It 44 | I_scan = np.round(I_scan, 2) 45 | 46 | for I0 in I_scan: 47 | es = run_once(L, Te, I0) 48 | ess.append(es) 49 | 50 | # ess = np.array(ess) 51 | 52 | # desdi2 = ess[2:] - ess[1:-1] + ess[:-2] 53 | 54 | for es in ess: 55 | print(es.result()) 56 | # max_loc = np.argmax(desdi2) 57 | # actual = I_scan[1 + max_loc] 58 | # np.testing.assert_allclose(actual, desired=It, rtol=0.25) # it is 25% because of the resolution of the scan. 59 | # The test itself is not quite working but you can examine the results visually and they make sense, 60 | # so we are leaving it this way for now 61 | 62 | 63 | if __name__ == "__main__": 64 | test_threshold() 65 | -------------------------------------------------------------------------------- /tests/test_tf1d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_tf1d/__init__.py -------------------------------------------------------------------------------- /tests/test_tf1d/configs/resonance.yaml: -------------------------------------------------------------------------------- 1 | solver: tf-1d 2 | 3 | mlflow: 4 | experiment: tf1d-ions-test 5 | run: test 6 | 7 | grid: 8 | nx: 32 9 | xmin: 0.0 10 | xmax: 20.94 11 | tmin: 0.0 12 | tmax: 50 13 | 14 | models: False 15 | 16 | units: 17 | laser_wavelength: 351nm 18 | normalizing_temperature: 2000eV 19 | normalizing_density: 1.5e21/cc 20 | Z: 10 21 | Zp: 10 22 | 23 | save: 24 | t: 25 | tmin: 0.0 26 | tmax: 50.0 27 | nt: 201 28 | x: 29 | xmin: 0.0 30 | xmax: 20.94 31 | nx: 16 32 | kx: 33 | kxmin: 0.0 34 | kxmax: 0.3 35 | nkx: 2 36 | 37 | physics: 38 | ion: 39 | is_on: False 40 | landau_damping: False 41 | mass: 1836.0 42 | T0: 0.01 43 | charge: 1.0 44 | gamma: 3 45 | trapping: 46 | is_on: False 47 | kld: 0.3 48 | electron: 49 | is_on: True 50 | landau_damping: True 51 | T0: 1.0 52 | mass: 1.0 53 | charge: -1.0 54 | gamma: 3.0 55 | trapping: 56 | is_on: False 57 | kld: 0.3 58 | 59 | 60 | drivers: 61 | "ex": 62 | "0": 63 | "k0": 0.3 64 | "w0": 1.1 65 | "dw0": 0.0 66 | "t_c": 15.0 67 | "t_w": 10.0 68 | "t_r": 2.5 69 | "x_c": 400 70 | "x_w": 1000000 71 | "x_r": 10 72 | "a0": 1.e-6 73 | -------------------------------------------------------------------------------- /tests/test_tf1d/configs/resonance_search.yaml: -------------------------------------------------------------------------------- 1 | solver: tf-1d 2 | 3 | mlflow: 4 | experiment: tf1d-resonance-search 5 | run: test 6 | 7 | adjoint: Backsolve 8 | 9 | units: 10 | laser_wavelength: 351nm 11 | normalizing_temperature: 2000eV 12 | normalizing_density: 1.5e21/cc 13 | Z: 10 14 | Zp: 10 15 | 16 | grid: 17 | nx: 16 18 | xmin: 0.0 19 | xmax: 20.94 20 | tmin: 0.0 21 | tmax: 50 22 | 23 | save: 24 | t: 25 | tmin: 0.0 26 | tmax: 50.0 27 | nt: 101 28 | x: 29 | xmin: 0.0 30 | xmax: 20.94 31 | nx: 16 32 | kx: 33 | kxmin: 0.0 34 | kxmax: 0.3 35 | nkx: 2 36 | 37 | 38 | physics: 39 | ion: 40 | is_on: False 41 | landau_damping: True 42 | mass: 1836.0 43 | T0: 0.01 44 | charge: 1.0 45 | gamma: 3 46 | trapping: 47 | is_on: False 48 | kld: 0.3 49 | electron: 50 | is_on: True 51 | landau_damping: True 52 | T0: 1.0 53 | mass: 1.0 54 | charge: -1.0 55 | gamma: 3.0 56 | trapping: 57 | is_on: False 58 | kld: 0.3 59 | 60 | 61 | drivers: 62 | "ex": 63 | "0": 64 | "k0": 0.3 65 | "w0": 1.1 66 | "dw0": 0.0 67 | "t_c": 15.0 68 | "t_w": 10.0 69 | "t_r": 2.5 70 | "x_c": 400 71 | "x_w": 1000000 72 | "x_r": 10 73 | "a0": 1.e-6 74 | -------------------------------------------------------------------------------- /tests/test_tf1d/configs/vlasov_comparison.yaml: -------------------------------------------------------------------------------- 1 | solver: tf-1d 2 | 3 | mlflow: 4 | experiment: tf1d-ions-test 5 | run: test 6 | 7 | models: False 8 | 9 | grid: 10 | nx: 32 11 | xmin: 0.0 12 | xmax: 20.94 13 | tmin: 0.0 14 | tmax: 50 15 | 16 | units: 17 | laser_wavelength: 351nm 18 | normalizing_temperature: 2000eV 19 | normalizing_density: 1.5e21/cc 20 | Z: 10 21 | Zp: 10 22 | 23 | save: 24 | t: 25 | tmin: 0.0 26 | tmax: 50.0 27 | nt: 201 28 | x: 29 | xmin: 0.0 30 | xmax: 20.94 31 | nx: 32 32 | kx: 33 | kxmin: 0.0 34 | kxmax: 0.3 35 | nkx: 2 36 | 37 | physics: 38 | ion: 39 | is_on: False 40 | landau_damping: False 41 | mass: 1836.0 42 | T0: 0.01 43 | charge: 1.0 44 | gamma: 3 45 | trapping: 46 | is_on: False 47 | kld: 0.3 48 | electron: 49 | is_on: True 50 | landau_damping: True 51 | T0: 1.0 52 | mass: 1.0 53 | charge: -1.0 54 | gamma: 3.0 55 | trapping: 56 | is_on: False 57 | kld: 0.3 58 | 59 | 60 | drivers: 61 | "ex": 62 | "0": 63 | "k0": 0.3 64 | "w0": 1.1 65 | "dw0": 0.0 66 | "t_c": 15.0 67 | "t_w": 10.0 68 | "t_r": 2.5 69 | "x_c": 400 70 | "x_w": 1000000 71 | "x_r": 10 72 | "a0": 1.e-6 73 | -------------------------------------------------------------------------------- /tests/test_tf1d/test_against_vlasov.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | import numpy as np 4 | import xarray as xr 5 | import yaml 6 | 7 | from adept import electrostatic, ergoExo 8 | 9 | 10 | def _modify_defaults_(defaults): 11 | rand_k0 = 0.358 12 | 13 | wepw = np.sqrt(1.0 + 3.0 * rand_k0**2.0) 14 | root = electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0) 15 | 16 | defaults["physics"]["landau_damping"] = True 17 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 18 | defaults["drivers"]["ex"]["0"]["w0"] = float(wepw) 19 | xmax = float(2.0 * np.pi / rand_k0) 20 | # defaults["save"]["field"]["xmax_to_store"] = float(2.0 * np.pi / rand_k0) 21 | defaults["grid"]["xmax"] = xmax 22 | defaults["save"]["x"]["xmax"] = xmax 23 | defaults["save"]["kx"]["kxmax"] = rand_k0 24 | defaults["mlflow"]["experiment"] = "test-against-vlasov" 25 | 26 | return defaults, float(np.imag(root)) 27 | 28 | 29 | def test_single_resonance(): 30 | with open("tests/test_tf1d/configs/vlasov_comparison.yaml") as file: 31 | defaults = yaml.safe_load(file) 32 | 33 | # modify config 34 | mod_defaults, actual_damping_rate = _modify_defaults_(defaults) 35 | 36 | exo = ergoExo() 37 | exo.setup(mod_defaults) 38 | result, datasets, run_id = exo(None) 39 | result = result["solver result"] 40 | vds = xr.open_dataset("tests/test_tf1d/vlasov-reference/all-fields-kx.nc", engine="h5netcdf") 41 | 42 | nk1_fluid = result.ys["kx"]["electron"]["n"]["mag"][:, 1] 43 | nk1_vlasov = vds["n-(k_x)"][:, 1].data 44 | t_fluid = result.ts 45 | t_vlasov = vds.coords["t"].data 46 | fluid_slc = slice(80, 160) 47 | vlasov_slc = slice(700, 850) 48 | 49 | vlasov_damping_rate = np.mean( 50 | np.gradient(nk1_vlasov[vlasov_slc], (t_vlasov[1] - t_vlasov[0])) / nk1_vlasov[vlasov_slc] 51 | ) 52 | fluid_damping_rate = np.mean(np.gradient(nk1_fluid[fluid_slc], (t_fluid[1] - t_fluid[0])) / nk1_fluid[fluid_slc]) 53 | 54 | print(f"{vlasov_damping_rate=}, {fluid_damping_rate=}") 55 | print(f"{np.amax(nk1_vlasov)=}, {np.amax(nk1_fluid)=}") 56 | 57 | np.testing.assert_almost_equal(vlasov_damping_rate, fluid_damping_rate, decimal=2) 58 | np.testing.assert_allclose(np.amax(nk1_fluid), np.amax(nk1_vlasov), rtol=0.05) 59 | 60 | 61 | if __name__ == "__main__": 62 | test_single_resonance() 63 | -------------------------------------------------------------------------------- /tests/test_tf1d/test_landau_damping.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | import mlflow 4 | import numpy as np 5 | import yaml 6 | from jax import numpy as jnp 7 | 8 | from adept import electrostatic, ergoExo 9 | 10 | 11 | def _modify_defaults_(defaults, rng): 12 | rand_k0 = np.round(rng.uniform(0.25, 0.4), 3) 13 | 14 | wepw = np.sqrt(1.0 + 3.0 * rand_k0**2.0) 15 | root = electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0) 16 | print(rand_k0, wepw, root) 17 | 18 | defaults["physics"]["landau_damping"] = True 19 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 20 | defaults["drivers"]["ex"]["0"]["w0"] = float(wepw) 21 | xmax = float(2.0 * np.pi / rand_k0) 22 | defaults["grid"]["xmax"] = xmax 23 | defaults["save"]["x"]["xmax"] = xmax 24 | defaults["save"]["kx"]["kxmax"] = rand_k0 25 | defaults["mlflow"]["experiment"] = "test-landau-damping" 26 | 27 | return defaults, float(np.imag(root)) 28 | 29 | 30 | def test_single_resonance(): 31 | with open("tests/test_tf1d/configs/resonance.yaml") as file: 32 | defaults = yaml.safe_load(file) 33 | 34 | # modify config 35 | rng = np.random.default_rng() 36 | mod_defaults, actual_damping_rate = _modify_defaults_(defaults, rng) 37 | 38 | exo = ergoExo() 39 | exo.setup(mod_defaults) 40 | result, datasets, run_id = exo(None) 41 | result = result["solver result"] 42 | 43 | with mlflow.start_run(run_id=run_id, log_system_metrics=True) as mlflow_run: 44 | kx = ( 45 | np.fft.fftfreq( 46 | mod_defaults["save"]["x"]["nx"], 47 | d=mod_defaults["save"]["x"]["ax"][2] - mod_defaults["save"]["x"]["ax"][1], 48 | ) 49 | * 2.0 50 | * np.pi 51 | ) 52 | one_over_kx = np.zeros_like(kx) 53 | one_over_kx[1:] = 1.0 / kx[1:] 54 | efs = jnp.real(jnp.fft.ifft(1j * one_over_kx[None, :] * jnp.fft.fft(1 - result.ys["x"]["electron"]["n"][:, :]))) 55 | ek1 = (2.0 / mod_defaults["grid"]["nx"] * np.abs(np.fft.fft(efs, axis=1)[:, 1])) ** 2.0 56 | frslc = slice(-100, -50) 57 | measured_damping_rate = np.mean(np.gradient(ek1[frslc], (result.ts[1] - result.ts[0])) / ek1[frslc]) 58 | print( 59 | f"Landau Damping rate check \n" 60 | f"measured: {np.round(measured_damping_rate, 5)}, " 61 | f"actual: {np.round(2 * actual_damping_rate, 5)}, " 62 | ) 63 | mlflow.log_metrics( 64 | {"actual damping rate": float(actual_damping_rate), "measured damping rate": float(measured_damping_rate)} 65 | ) 66 | 67 | np.testing.assert_almost_equal(measured_damping_rate, 2 * actual_damping_rate, decimal=2) 68 | 69 | 70 | if __name__ == "__main__": 71 | test_single_resonance() 72 | -------------------------------------------------------------------------------- /tests/test_tf1d/test_resonance.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | import numpy as np 4 | import pytest 5 | import yaml 6 | from jax import numpy as jnp 7 | 8 | from adept import electrostatic, ergoExo 9 | 10 | 11 | def _modify_defaults_(defaults, rng, gamma): 12 | rand_k0 = np.round(rng.uniform(0.25, 0.4), 3) 13 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 14 | defaults["physics"]["electron"]["gamma"] = gamma 15 | if gamma == "kinetic": 16 | root = np.real(electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0)) 17 | defaults["mlflow"]["run"] = "kinetic" 18 | else: 19 | root = np.sqrt(1.0 + 3.0 * rand_k0**2.0) 20 | defaults["mlflow"]["run"] = "bohm-gross" 21 | 22 | xmax = float(2.0 * np.pi / rand_k0) 23 | defaults["grid"]["xmax"] = xmax 24 | defaults["save"]["x"]["xmax"] = xmax 25 | defaults["mlflow"]["experiment"] = "test-resonance" 26 | 27 | return defaults, float(root) 28 | 29 | 30 | @pytest.mark.parametrize("gamma", ["kinetic", 3.0]) 31 | def test_single_resonance(gamma): 32 | with open("tests/test_tf1d/configs/resonance.yaml") as file: 33 | defaults = yaml.safe_load(file) 34 | 35 | # modify config 36 | rng = np.random.default_rng() 37 | mod_defaults, actual_resonance = _modify_defaults_(defaults, rng, gamma) 38 | 39 | # run 40 | exo = ergoExo() 41 | exo.setup(mod_defaults) 42 | result, datasets, run_id = exo(None) 43 | result = result["solver result"] 44 | 45 | kx = ( 46 | np.fft.fftfreq( 47 | mod_defaults["save"]["x"]["nx"], d=mod_defaults["save"]["x"]["ax"][2] - mod_defaults["save"]["x"]["ax"][1] 48 | ) 49 | * 2.0 50 | * np.pi 51 | ) 52 | one_over_kx = np.zeros_like(kx) 53 | one_over_kx[1:] = 1.0 / kx[1:] 54 | efs = jnp.real( 55 | jnp.fft.ifft( 56 | 1j 57 | * one_over_kx[None, :] 58 | * jnp.fft.fft(result.ys["x"]["ion"]["n"][:, :] - result.ys["x"]["electron"]["n"][:, :]) 59 | ) 60 | ) 61 | ek1 = np.fft.fft(efs, axis=1)[:, 1] 62 | env, freq = electrostatic.get_nlfs(ek1, result.ts[1] - result.ts[0]) 63 | frslc = slice(-80, -10) 64 | print( 65 | f"Frequency check \nmeasured: {np.round(np.mean(freq[frslc]), 5)}, desired: {np.round(actual_resonance, 5)}, " 66 | ) 67 | measured_resonance = np.mean(freq[frslc]) 68 | np.testing.assert_almost_equal(measured_resonance, actual_resonance, decimal=2) 69 | 70 | 71 | if __name__ == "__main__": 72 | for gamma in ["kinetic", 3.0]: 73 | test_single_resonance(gamma) 74 | -------------------------------------------------------------------------------- /tests/test_tf1d/test_resonance_search.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | import time 4 | from itertools import product 5 | 6 | import equinox as eqx 7 | import mlflow 8 | import numpy as np 9 | import optax 10 | import pytest 11 | import yaml 12 | from jax import Array, devices 13 | from jax import numpy as jnp 14 | from jax import random as jr 15 | from tqdm import tqdm 16 | 17 | from adept import ergoExo 18 | from adept.electrostatic import get_roots_to_electrostatic_dispersion 19 | from adept.tf1d import BaseTwoFluid1D 20 | 21 | 22 | class Resonance(BaseTwoFluid1D): 23 | def __init__(self, cfg): 24 | super().__init__(cfg) 25 | 26 | def __call__(self, params: dict, args: dict) -> dict: 27 | args = self.args 28 | args["drivers"]["ex"]["0"]["w0"] = params["w0"] 29 | solver_result = super().__call__(params, args) 30 | soln = solver_result["solver result"] 31 | nk1 = jnp.abs(jnp.fft.fft(soln.ys["x"]["electron"]["n"], axis=1)[:, 1]) 32 | return -jnp.amax(nk1), solver_result 33 | 34 | def vg(self, params: dict, args: dict) -> tuple[float, Array, dict]: 35 | return eqx.filter_jit(eqx.filter_value_and_grad(self.__call__, has_aux=True))(params, args) 36 | 37 | 38 | def load_cfg(rand_k0, gamma, adjoint): 39 | with open("tests/test_tf1d/configs/resonance_search.yaml") as file: 40 | defaults = yaml.safe_load(file) 41 | 42 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 43 | defaults["physics"]["electron"]["gamma"] = gamma 44 | defaults["adjoint"] = adjoint 45 | 46 | if gamma == "kinetic": 47 | wepw = np.real(get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0)) 48 | defaults["mlflow"]["run"] = "kinetic" 49 | else: 50 | wepw = np.sqrt(1.0 + 3.0 * rand_k0**2.0) 51 | defaults["mlflow"]["run"] = "bohm-gross" 52 | 53 | xmax = float(2.0 * np.pi / rand_k0) 54 | defaults["grid"]["xmax"] = xmax 55 | defaults["save"]["x"]["xmax"] = xmax 56 | 57 | return defaults, wepw 58 | 59 | 60 | @pytest.mark.parametrize("adjoint", ["Recursive", "Backsolve"]) 61 | @pytest.mark.parametrize("gamma", ["kinetic", 3.0]) 62 | def test_resonance_search(gamma, adjoint): 63 | mlflow.set_experiment("tf1d-resonance-search") 64 | with mlflow.start_run(run_name="res-search-opt", log_system_metrics=True) as mlflow_run: 65 | # sim_k0, actual_w0 = init_w0(gamma, adjoint) 66 | rng = np.random.default_rng(420) 67 | sim_k0 = rng.uniform(0.26, 0.4) 68 | 69 | mod_defaults, actual_w0 = load_cfg(sim_k0, gamma, adjoint) 70 | 71 | rng_key = jr.PRNGKey(420) 72 | 73 | params = {"w0": 1.1 + 0.05 * jr.normal(rng_key, [1], dtype=jnp.float64)[0]} 74 | 75 | optimizer = optax.adam(0.1) 76 | opt_state = optimizer.init(params) 77 | 78 | t0 = time.time() 79 | mlflow.log_metrics({"init_time": round(time.time() - t0, 4)}) 80 | mlflow.log_metrics({"w0": float(params["w0"])}, step=0) 81 | mlflow.log_metrics({"actual_w0": actual_w0}, step=0) 82 | for i in tqdm(range(40)): 83 | exo = ergoExo(mlflow_nested=True) 84 | mod_defaults, _ = load_cfg(sim_k0, gamma, adjoint) 85 | exo.setup(mod_defaults, Resonance) 86 | val, grad, results = exo.val_and_grad(params) 87 | updates, opt_state = optimizer.update(grad, opt_state, params) 88 | params = optax.apply_updates(params, updates) 89 | 90 | mlflow.log_metrics({"w0": float(params["w0"]), "actual_w0": actual_w0, "loss": float(val)}, step=i + 1) 91 | 92 | print(f"{gamma=}, {adjoint=}") 93 | print(f"{actual_w0=}, {float(params['w0'])=}") 94 | 95 | np.testing.assert_allclose(actual_w0, float(params["w0"]), rtol=0.03) 96 | 97 | 98 | if __name__ == "__main__": 99 | for gamma, adjoint in product(["kinetic", 3.0], ["Recursive", "Backsolve"]): 100 | if not any(["gpu" == device.platform for device in devices()]): 101 | pytest.skip("Takes too long without a GPU") 102 | else: 103 | test_resonance_search(gamma, adjoint) 104 | -------------------------------------------------------------------------------- /tests/test_tf1d/vlasov-reference/all-fields-kx.nc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_tf1d/vlasov-reference/all-fields-kx.nc -------------------------------------------------------------------------------- /tests/test_tf1d/vlasov-reference/all-fields.nc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_tf1d/vlasov-reference/all-fields.nc -------------------------------------------------------------------------------- /tests/test_tf1d/vlasov-reference/config.yaml: -------------------------------------------------------------------------------- 1 | density: 2 | quasineutrality: true 3 | species-background: 4 | T0: 1.0 5 | basis: tanh 6 | m: 2.0 7 | noise_seed: 420 8 | noise_type: gaussian 9 | noise_val: 0.0 10 | space-profile: 11 | baseline: 1.0 12 | bump_or_trough: bump 13 | center: 0.0 14 | rise: 10.0 15 | slope: 0.0 16 | wall_height: 0.0 17 | width: 100000.0 18 | v0: 0.0 19 | diagnostics: 20 | dist: 21 | axis_transforms: 22 | - (k_x, v) 23 | f_transforms: 24 | - f 25 | kxmax_to_store: 4.0 26 | kxmin_to_store: 0.0 27 | nkx: 3 28 | nv: 256 29 | nx: 32 30 | vmax_to_store: 6.4 31 | vmin_to_store: -6.4 32 | xmax_to_store: 20.94 33 | xmin_to_store: 0.0 34 | field: 35 | axis_transforms: 36 | - (x) 37 | - (k_x) 38 | f_transforms: 39 | - force 40 | - n 41 | kxmax_to_store: 4.0 42 | kxmin_to_store: 0.0 43 | nkx: 5 44 | nx: 32 45 | xmax_to_store: 17.550796947429014 46 | xmin_to_store: 0.0 47 | postprocess: 48 | chie: 49 | k_max: 0.4 50 | k_min: 0.24 51 | num_k: 9 52 | num_w: 16 53 | t_max: 30 54 | t_min: 20 55 | w_max: 1.4 56 | w_min: 1.1 57 | nlfs: 58 | metric_tmax: 40 59 | metric_tmin: 30 60 | ts_tmax: 45 61 | ts_tmin: 5 62 | series: 63 | '0': 64 | name: damping-rate 65 | quant_conf: 66 | inv_or_not: false 67 | reduce_func: ave 68 | t_range: '30:45' 69 | xmax: all 70 | xmin: all 71 | quant_func: e_dr 72 | grid: 73 | c_light: 10 74 | checkpoint_steps: 500 75 | nt: 1000 76 | nv: 512 77 | nx: 32 78 | tmax: 50.0 79 | vmax: 6.4 80 | xmax: 17.550796947429014 81 | xmin: 0.0 82 | krook: 83 | space-profile: 84 | baseline: 1.0e-05 85 | bump_or_trough: bump 86 | center: 0.0 87 | rise: 10.0 88 | slope: 0.0 89 | wall_height: 0.0 90 | width: 100000.0 91 | time-profile: 92 | baseline: 0.0001 93 | bump_or_trough: bump 94 | center: 0.0 95 | rise: 10.0 96 | slope: 0.0 97 | wall_height: 0.0 98 | width: 100000.0 99 | machine: local 100 | mlflow: 101 | experiment: landau_damping 102 | run: test 103 | nu: 104 | time-profile: 105 | baseline: 0.0001 106 | bump_or_trough: bump 107 | center: 0.0 108 | rise: 10.0 109 | slope: 0.0 110 | wall_height: 0.0 111 | width: 100000.0 112 | pulses: 113 | ex: 114 | '0': 115 | a0: 1.0e-06 116 | dw0: 0.0 117 | k0: 0.358 118 | t_center: 15.0 119 | t_rise: 2.5 120 | t_width: 10.0 121 | w0: 1.2310768180888438 122 | x_center: 400.0 123 | x_rise: 10.0 124 | x_width: 1000000.0 125 | ey: {} 126 | solver: 127 | backend: jax 128 | dfdt: leapfrog 129 | edfdv: exponential 130 | field: poisson 131 | fp_operator: dougherty 132 | num_unroll: 32 133 | vdfdx: exponential 134 | -------------------------------------------------------------------------------- /tests/test_vfp1d/epp-short.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | reference electron temperature: 300eV 4 | reference ion temperature: 300eV 5 | reference electron density: 1.5e21/cm^3 6 | Z: 6 7 | Ion: Au+ 8 | logLambda: nrl 9 | 10 | 11 | density: 12 | quasineutrality: true 13 | species-background: 14 | noise_seed: 420 15 | noise_type: gaussian 16 | noise_val: 0.0 17 | v0: 0.0 18 | T0: 1.0 19 | m: 2.0 20 | n: 21 | basis: uniform 22 | baseline: 1.0 23 | bump_or_trough: bump 24 | center: 0m 25 | rise: 1m 26 | bump_height: 0.0 27 | width: 1m 28 | T: 29 | basis: sine 30 | baseline: 1.0 31 | amplitude: 1.0e-3 32 | wavelength: 250um 33 | 34 | grid: 35 | dt: 1fs 36 | nv: 256 37 | nx: 32 38 | tmin: 0.ps 39 | tmax: 2ps 40 | vmax: 8.0 41 | xmax: 500um 42 | xmin: 0.0um 43 | nl: 1 44 | 45 | save: 46 | fields: 47 | t: 48 | tmin: 0.0ps 49 | tmax: 0.5ps 50 | nt: 11 51 | electron: 52 | t: 53 | tmin: 0.0ps 54 | tmax: 0.5ps 55 | nt: 6 56 | 57 | solver: vfp-1d 58 | 59 | mlflow: 60 | experiment: vfp1d 61 | run: epperlein-short 62 | 63 | drivers: 64 | ex: {} 65 | ey: {} 66 | 67 | terms: 68 | fokker_planck: 69 | flm: 70 | ee: false 71 | f00: 72 | type: Dougherty 73 | e_solver: oshun 74 | -------------------------------------------------------------------------------- /tests/test_vfp1d/test_kappa_eh.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import mlflow 4 | import numpy as np 5 | import pytest 6 | import yaml 7 | from jax import devices 8 | 9 | from adept import ergoExo 10 | 11 | 12 | def _run_(Z, ee): 13 | # with open("configs/tf-1d/damping.yaml", "r") as fi: 14 | with open(f"{os.path.join(os.getcwd(), 'tests/test_vfp1d/epp-short')}.yaml") as fi: 15 | cfg = yaml.safe_load(fi) 16 | 17 | cfg["units"]["Z"] = Z 18 | 19 | if ee: 20 | cfg["terms"]["fokker_planck"]["flm"]["ee"] = True 21 | cfg["grid"]["nv"] = 2048 22 | 23 | exo = ergoExo() 24 | exo.setup(cfg) 25 | sol, datasets, run_id = exo(None) 26 | dataT = datasets["fields"]["fields-T keV"].data 27 | np.testing.assert_almost_equal(np.mean(dataT[-4, :]), np.mean(dataT[4, :]), decimal=5) 28 | 29 | datan = datasets["fields"]["fields-n n_c"].data 30 | np.testing.assert_almost_equal(np.mean(datan[-4, :]), np.mean(datan[4, :]), decimal=5) 31 | 32 | kappa_eh = mlflow.get_run(run_id).data.metrics["kappa_eh"] 33 | kappa = mlflow.get_run(run_id).data.metrics["kappa"] 34 | 35 | np.testing.assert_almost_equal(kappa, kappa_eh, decimal=0) 36 | 37 | return run_id 38 | 39 | 40 | @pytest.mark.parametrize("Z", list(range(1, 21, 4)) + [40, 60, 80]) 41 | @pytest.mark.parametrize("ee", [True, False]) 42 | def test_kappa_eh(Z, ee): 43 | if not any(["gpu" == device.platform for device in devices()]): 44 | pytest.skip(f"Skipping Z={Z} to save time because no GPU is available") 45 | if Z in [1, 21, 80]: 46 | _run_(Z, ee) 47 | 48 | else: 49 | _run_(Z, ee) 50 | -------------------------------------------------------------------------------- /tests/test_vlasov1d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_vlasov1d/__init__.py -------------------------------------------------------------------------------- /tests/test_vlasov1d/configs/resonance.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | space-profile: 20 | baseline: 1.0 21 | bump_or_trough: bump 22 | center: 0.0 23 | rise: 25.0 24 | slope: 0.0 25 | bump_height: 0.0 26 | width: 100000.0 27 | 28 | grid: 29 | dt: 0.25 30 | nv: 512 31 | nx: 32 32 | tmin: 0. 33 | tmax: 480.0 34 | vmax: 6.4 35 | xmax: 20.94 36 | xmin: 0.0 37 | 38 | save: 39 | fields: 40 | t: 41 | tmin: 0.0 42 | tmax: 480.0 43 | nt: 961 44 | electron: 45 | t: 46 | tmin: 0.0 47 | tmax: 480.0 48 | nt: 9 49 | 50 | solver: vlasov-1d 51 | 52 | mlflow: 53 | experiment: vlasov1d 54 | run: test 55 | 56 | drivers: 57 | ex: 58 | '0': 59 | a0: 1.e-6 60 | k0: 0.3 61 | t_center: 40.0 62 | t_rise: 5.0 63 | t_width: 30.0 64 | w0: 1.1598 65 | dw0: 0. 66 | x_center: 0.0 67 | x_rise: 10.0 68 | x_width: 4000000.0 69 | ey: {} 70 | 71 | terms: 72 | field: poisson 73 | edfdv: cubic-spline 74 | time: leapfrog 75 | diags: false 76 | fokker_planck: 77 | is_on: True 78 | type: Dougherty 79 | time: 80 | baseline: 1.0 81 | bump_or_trough: bump 82 | center: 0.0 83 | rise: 25.0 84 | slope: 0.0 85 | bump_height: 0.0 86 | width: 100000.0 87 | space: 88 | baseline: 1.0e-6 89 | bump_or_trough: bump 90 | center: 0.0 91 | rise: 25.0 92 | slope: 0.0 93 | bump_height: 0.0 94 | width: 100000.0 95 | krook: 96 | is_on: True 97 | time: 98 | baseline: 1.0 99 | bump_or_trough: bump 100 | center: 0.0 101 | rise: 25.0 102 | slope: 0.0 103 | bump_height: 0.0 104 | width: 100000.0 105 | space: 106 | baseline: 1.0e-6 107 | bump_or_trough: bump 108 | center: 0.0 109 | rise: 25.0 110 | slope: 0.0 111 | bump_height: 0.0 112 | width: 100000.0 113 | 114 | 115 | diagnostics: 116 | diag-vlasov-dfdt: False 117 | diag-fp-dfdt: False 118 | -------------------------------------------------------------------------------- /tests/test_vlasov1d/test_absorbing_wave.py: -------------------------------------------------------------------------------- 1 | import equinox as eqx 2 | import numpy as np 3 | from diffrax import ODETerm, diffeqsolve 4 | from jax import jit 5 | 6 | from adept._base_ import Stepper 7 | from adept._vlasov1d.solvers.pushers.field import Driver, WaveSolver 8 | 9 | 10 | class VectorField(eqx.Module): 11 | wave_solver: WaveSolver 12 | driver: Driver 13 | """ 14 | This class contains the function that updates the state 15 | 16 | All the pushers are chosen and initialized here and a single time-step is defined here. 17 | 18 | :param cfg: 19 | :return: 20 | """ 21 | 22 | def __init__(self, c_light, dx, dt, xax): 23 | self.wave_solver = WaveSolver(c=c_light, dx=dx, dt=dt) 24 | self.driver = Driver(xax=xax, driver_key="ey") 25 | 26 | def __call__(self, t, y, args): 27 | djy = self.driver(t, args) 28 | return self.wave_solver(y["a"], y["prev_a"], djy, 0.0) 29 | 30 | 31 | def test_absorbing_boundaries(): 32 | xmax = 1000 33 | xmin = 0 34 | nx = 1024 35 | tmax = 200 36 | c_light = 11.3 37 | 38 | dx = (xmax - xmin) / nx 39 | dt = 0.95 * dx / c_light 40 | nt = int(tmax / dt) 41 | 42 | xax = np.linspace(xmin - dx / 2.0, xmax + dx / 2.0, nx + 2) 43 | 44 | a_old, a, charge, djy = np.zeros_like(xax), np.zeros_like(xax), np.zeros_like(xax)[1:-1], np.zeros_like(xax) 45 | 46 | ey_dict = { 47 | "a0": 1.0e-4, 48 | "k0": -1.4, 49 | "t_center": 40.0, 50 | "t_rise": 5.0, 51 | "t_width": 30.0, 52 | "w0": 15.82, 53 | "dw0": 0.0, 54 | "x_center": 800.0, 55 | "x_rise": 10.0, 56 | "x_width": 50.0, 57 | } 58 | 59 | args = {"drivers": {"ey": {"0": ey_dict}}} 60 | 61 | @jit 62 | def _run_(y, args): 63 | return diffeqsolve( 64 | terms=ODETerm(VectorField(c_light, dx, dt, xax)), 65 | solver=Stepper(), 66 | t0=0.0, 67 | t1=tmax, 68 | max_steps=nt + 5, 69 | dt0=dt, 70 | y0=y, 71 | args=args, 72 | ) 73 | 74 | test_a = _run_({"a": a, "prev_a": a_old}, args) 75 | 76 | np.testing.assert_almost_equal(np.sum(np.square(test_a.ys["a"])), 0.0, decimal=8) 77 | 78 | 79 | if __name__ == "__main__": 80 | test_absorbing_boundaries() 81 | -------------------------------------------------------------------------------- /tests/test_vlasov1d/test_landau_damping.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | import itertools 4 | 5 | import numpy as np 6 | import pytest 7 | import yaml 8 | 9 | from adept import electrostatic, ergoExo 10 | 11 | 12 | def _modify_defaults_(defaults, rng, real_or_imag, time, field, edfdv): 13 | defaults["terms"]["time"] = time 14 | defaults["terms"]["field"] = field 15 | defaults["terms"]["edfdv"] = edfdv 16 | if field == "ampere": 17 | defaults["grid"]["dt"] = 0.025 18 | 19 | if real_or_imag == "real": 20 | rand_k0 = np.round(rng.uniform(0.2, 0.28), 3) 21 | else: 22 | rand_k0 = np.round(rng.uniform(0.28, 0.35), 3) 23 | 24 | root = electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0) 25 | 26 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 27 | defaults["drivers"]["ex"]["0"]["w0"] = float(np.real(root)) 28 | xmax = float(2.0 * np.pi / rand_k0) 29 | defaults["grid"]["xmax"] = xmax 30 | defaults["mlflow"]["experiment"] = "vlasov1d-test-resonance" 31 | 32 | return defaults, root 33 | 34 | 35 | @pytest.mark.parametrize( 36 | "real_or_imag, time, field, edfdv", 37 | itertools.product( 38 | ["real", "imag"], ["sixth", "leapfrog"], ["poisson", "ampere", "hampere"], ["exponential", "cubic-spline"] 39 | ), 40 | ) 41 | def test_single_resonance(real_or_imag, time, field, edfdv): 42 | if (time == "sixth") and (field == "ampere"): 43 | print("not implemented - skipping test") 44 | elif (time == "sixth") and (field == "hampere"): 45 | print("not implemented - skipping test") 46 | else: 47 | with open("tests/test_vlasov1d/configs/resonance.yaml") as file: 48 | defaults = yaml.safe_load(file) 49 | 50 | # modify config 51 | rng = np.random.default_rng() 52 | mod_defaults, root = _modify_defaults_(defaults, rng, real_or_imag, time, field, edfdv) 53 | 54 | actual_damping_rate = np.imag(root) 55 | actual_resonance = np.real(root) 56 | 57 | exo = ergoExo() 58 | exo.setup(mod_defaults) 59 | result, datasets, run_id = exo(None) 60 | result = result["solver result"] 61 | efs = result.ys["fields"]["e"] 62 | ek1 = 2.0 / mod_defaults["grid"]["nx"] * np.fft.fft(efs, axis=1)[:, 1] 63 | ek1_mag = np.abs(ek1) 64 | if real_or_imag == "imag": 65 | frslc = slice(-100, -50) 66 | measured_damping_rate = np.mean( 67 | np.gradient(ek1_mag[frslc], (result.ts["fields"][1] - result.ts["fields"][0])) / ek1_mag[frslc] 68 | ) 69 | print( 70 | f"Landau Damping rate check \n" 71 | f"measured: {np.round(measured_damping_rate, 5)}, " 72 | f"actual: {np.round(actual_damping_rate, 5)}, " 73 | ) 74 | 75 | np.testing.assert_almost_equal(measured_damping_rate, actual_damping_rate, decimal=2) 76 | else: 77 | env, freq = electrostatic.get_nlfs(ek1, result.ts["fields"][1] - result.ts["fields"][0]) 78 | frslc = slice(-480, -240) 79 | print( 80 | f"Frequency check \n" 81 | f"measured: {np.round(np.mean(freq[frslc]), 5)}, " 82 | f"desired: {np.round(actual_resonance, 5)}, " 83 | ) 84 | measured_resonance = np.mean(freq[frslc]) 85 | np.testing.assert_almost_equal(measured_resonance, actual_resonance, decimal=2) 86 | 87 | 88 | if __name__ == "__main__": 89 | test_single_resonance(real_or_imag="real") 90 | -------------------------------------------------------------------------------- /tests/test_vlasov2d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ergodicio/adept/5e95cfa77f1c618d0a57a346ee6349d0ba3b0b9d/tests/test_vlasov2d/__init__.py -------------------------------------------------------------------------------- /tests/test_vlasov2d/configs/damping.yaml: -------------------------------------------------------------------------------- 1 | units: 2 | laser_wavelength: 351nm 3 | normalizing_temperature: 2000eV 4 | normalizing_density: 1.5e21/cc 5 | Z: 10 6 | Zp: 10 7 | 8 | 9 | density: 10 | quasineutrality: true 11 | species-background: 12 | noise_seed: 420 13 | noise_type: gaussian 14 | noise_val: 0.0 15 | v0: 0.0 16 | T0: 1.0 17 | m: 2.0 18 | basis: uniform 19 | space-profile: 20 | baseline: 1.0 21 | bump_or_trough: bump 22 | center: 0.0 23 | rise: 25.0 24 | slope: 0.0 25 | wall_height: 0.0 26 | width: 100000.0 27 | 28 | grid: 29 | dt: 0.1 30 | nvx: 128 31 | nvy: 128 32 | nx: 24 33 | ny: 8 34 | tmin: 0. 35 | tmax: 150.0 36 | vmax: 6.4 37 | xmax: 20.94 38 | xmin: 0.0 39 | ymin: -4.0 40 | ymax: 4.0 41 | 42 | save: 43 | fields: 44 | t: 45 | tmin: 0.0 46 | tmax: 150.0 47 | nt: 301 48 | # x: 49 | # xmin: 0.0 50 | # xmax: 19.0 51 | # nx: 16 52 | # y: 53 | # ymin: -6.0 54 | # ymax: 6.0 55 | # ny: 2 56 | electron: 57 | t: 58 | tmin: 0.0 59 | tmax: 150.0 60 | nt: 3 61 | 62 | krook: 63 | space-profile: 64 | baseline: 0.0 65 | bump_or_trough: bump 66 | center: 2500 67 | rise: 10.0 68 | slope: 0.0 69 | wall_height: 0.0 70 | width: 48000000.0 71 | time-profile: 72 | baseline: 0.0 73 | bump_or_trough: bump 74 | center: 0.0 75 | rise: 10.0 76 | slope: 0.0 77 | wall_height: 0.0 78 | width: 100000.0 79 | 80 | machine: s-t4 81 | 82 | # solver: vlasov-2d 83 | 84 | mlflow: 85 | experiment: ld-2d2v 86 | run: envelope-driver-small-amp 87 | 88 | drivers: 89 | ex: 90 | '0': 91 | a0: 1.e-6 92 | k0: 0.3 93 | t_center: 30.0 94 | t_rise: 5.0 95 | t_width: 30.0 96 | w0: 1.1598 97 | dw0: 0. 98 | x_center: 0.0 99 | x_rise: 10.0 100 | x_width: 4000000.0 101 | y_center: 0.0 102 | y_rise: 10.0 103 | y_width: 200.0 104 | ey: {} 105 | 106 | solver: 107 | dfdt: leapfrog 108 | edfdv: exponential 109 | fp_operator: dougherty 110 | num_unroll: 32 111 | field: ampere 112 | vdfdx: exponential 113 | push_f: True 114 | -------------------------------------------------------------------------------- /tests/test_vlasov2d/test_landau_damping.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Ergodic LLC 2023 2 | # research@ergodic.io 3 | 4 | 5 | import mlflow 6 | import numpy as np 7 | import pytest 8 | import yaml 9 | from jax import devices 10 | 11 | from adept import electrostatic 12 | 13 | 14 | def _modify_defaults_(defaults, rng, real_or_imag): 15 | if real_or_imag == "real": 16 | rand_k0 = np.round(rng.uniform(0.2, 0.28), 3) 17 | else: 18 | rand_k0 = np.round(rng.uniform(0.28, 0.35), 3) 19 | 20 | root = electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, rand_k0) 21 | 22 | defaults["drivers"]["ex"]["0"]["k0"] = float(rand_k0) 23 | defaults["drivers"]["ex"]["0"]["w0"] = float(np.real(root)) 24 | xmax = float(2.0 * np.pi / rand_k0) 25 | defaults["grid"]["xmax"] = xmax 26 | defaults["mlflow"]["experiment"] = "vlasov2d-test" 27 | 28 | return defaults, root 29 | 30 | 31 | @pytest.mark.parametrize("real_or_imag", ["real", "imag"]) 32 | def test_single_resonance(real_or_imag): 33 | if not any(["gpu" == device.platform for device in devices()]): 34 | pytest.skip("Skipping test on CPU") 35 | 36 | else: 37 | with open("tests/test_vlasov2d/configs/damping.yaml") as file: 38 | defaults = yaml.safe_load(file) 39 | 40 | # modify config 41 | rng = np.random.default_rng() 42 | mod_defaults, root = _modify_defaults_(defaults, rng, real_or_imag) 43 | 44 | actual_damping_rate = np.imag(root) 45 | actual_resonance = np.real(root) 46 | # run 47 | mlflow.set_experiment(mod_defaults["mlflow"]["experiment"]) 48 | # modify config 49 | with mlflow.start_run(run_name=mod_defaults["mlflow"]["run"], log_system_metrics=True) as mlflow_run: 50 | # This will fail as run not defined 51 | result, datasets = run(mod_defaults) # noqa: F821 52 | efs = result.ys["fields"]["ex"] 53 | ek1 = 2.0 / mod_defaults["grid"]["nx"] * np.fft.fft2(efs, axes=(1, 2))[:, 1, 0] 54 | ek1_mag = np.abs(ek1) 55 | if real_or_imag == "imag": 56 | frslc = slice(-64, -4) 57 | measured_damping_rate = np.mean( 58 | np.gradient(ek1_mag[frslc], (result.ts["fields"][1] - result.ts["fields"][0])) / ek1_mag[frslc] 59 | ) 60 | print( 61 | f"Landau Damping rate check \n" 62 | f"measured: {np.round(measured_damping_rate, 5)}, " 63 | f"actual: {np.round(actual_damping_rate, 5)}, " 64 | ) 65 | mlflow.log_metrics( 66 | { 67 | "actual damping rate": float(actual_damping_rate), 68 | "measured damping rate": float(measured_damping_rate), 69 | } 70 | ) 71 | 72 | np.testing.assert_almost_equal(measured_damping_rate, actual_damping_rate, decimal=2) 73 | else: 74 | env, freq = electrostatic.get_nlfs(ek1, result.ts["fields"][1] - result.ts["fields"][0]) 75 | frslc = slice(-200, -100) 76 | print( 77 | f"Frequency check \n" 78 | f"measured: {np.round(np.mean(freq[frslc]), 5)}, " 79 | f"desired: {np.round(actual_resonance, 5)}, " 80 | ) 81 | measured_resonance = np.mean(freq[frslc]) 82 | mlflow.log_metrics( 83 | { 84 | "actual frequency": float(actual_resonance), 85 | "measured frequency": float(measured_resonance), 86 | } 87 | ) 88 | np.testing.assert_almost_equal(measured_resonance, actual_resonance, decimal=2) 89 | 90 | 91 | if __name__ == "__main__": 92 | test_single_resonance(real_or_imag="real") 93 | --------------------------------------------------------------------------------