diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 00000000..9670e542 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,20 @@ +[run] +source = src +omit = + src/data_assimilation/* + src/inputs/* + src/utils/debug_helpers.py + */tests/* + */test_* + setup.py + +[report] +exclude_lines = + pragma: no cover + def __repr__ + raise AssertionError + raise NotImplementedError + if __name__ == .__main__.: + +[html] +directory = htmlcov \ No newline at end of file diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml new file mode 100644 index 00000000..45d1ab20 --- /dev/null +++ b/.github/workflows/deploy.yml @@ -0,0 +1,94 @@ +name: deploy + +on: [push, pull_request, workflow_dispatch] + +permissions: + contents: write + +jobs: + integration-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + # Skip setup-python if running under act + - name: Set up Python (GitHub only) + if: ${{ !env.ACT }} + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + # Show system Python version when running under act + - name: Use system Python (act) + if: ${{ env.ACT }} + run: | + echo "Using system Python" + which python + python --version + + - name: Install dependencies + run: | + pip install -r dev-requirements.txt + pip install . + + - name: Run tests without coverage + run: | + # pytest --cov=src --cov-report=xml --cov-report=html --cov-report=term + pytest ./test_scripts/test_blending.py + pytest ./test_scripts/test_flow_solver.py + + # - name: Upload coverage HTML report + # if: ${{ !env.ACT }} + # uses: actions/upload-artifact@v4 + # with: + # name: coverage-html + # path: htmlcov/ + + - name: Upload test plots and logs + if: ${{ !env.ACT }} + uses: actions/upload-artifact@v4 + with: + name: test-outputs + path: | + outputs/**/*.png + logs/**/*.log + !outputs/**/*.dat + !outputs/**/*.h5 + + # docs: + # runs-on: ubuntu-latest + # steps: + # - uses: actions/checkout@v4 + + # - uses: ConorMacBride/install-package@v1 + # with: + # apt: graphviz + + # - uses: actions/setup-python@v5 + # with: + # python-version: '3.10' + + # - uses: actions/cache@v4 + # with: + # path: ~/.cache/pip + # key: ${{ runner.os }}-pip-${{ hashFiles('**/dev-requirements.txt') }} + # restore-keys: | + # ${{ runner.os }}-pip- + + # - name: Install dependencies + # run: | + # pip install -r dev-requirements.txt + # pip install . + + # - name: Sphinx build + # run: | + # sphinx-build docs/source _build + + # - name: Deploy to GitHub Pages + # uses: peaceiris/actions-gh-pages@v4 + # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} + # with: + # publish_branch: gh-pages + # github_token: ${{ secrets.GITHUB_TOKEN }} + # publish_dir: _build/ + # force_orphan: true \ No newline at end of file diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml deleted file mode 100644 index e7778d47..00000000 --- a/.github/workflows/documentation.yml +++ /dev/null @@ -1,33 +0,0 @@ -name: docs - -on: [push, pull_request, workflow_dispatch] - -permissions: - contents: write - -jobs: - docs: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - uses: ConorMacBride/install-package@v1 - with: - apt: graphviz - - uses: actions/setup-python@v3 - with: - python-version: '3.10.5' - - name: Install dependencies - run: | - pip install -r dev-requirements.txt - pip install . - - name: Sphinx build - run: | - sphinx-build docs/source _build - - name: Deploy to GitHub Pages - uses: peaceiris/actions-gh-pages@v3 - if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/develop' }} - with: - publish_branch: gh-pages - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: _build/ - force_orphan: true \ No newline at end of file diff --git a/.gitignore b/.gitignore index 83df0228..76231175 100644 --- a/.gitignore +++ b/.gitignore @@ -4,9 +4,19 @@ **dask-worker-space* *.ipynb_checkpoints* *.egg-info* +*.coverage* +profile.html +profile.json + +# Ignore everything under outputs/ +outputs/** + +# Do not ignore directories (so unignore rules can match their contents) +!outputs/ +!outputs/**/ +!outputs/**/*_stripped.h5 # Ignore local development and output subdirectories -**outputs* **build* **changelog.d* diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4660e191..f2058adb 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,3 +1,21 @@ +0.60.0 (2024-12-23) +------------------- + +Changed +^^^^^^^ + +- Restructured outermost looping in main() (23a7817db225fe2188bcf3045cfdc863435b6888) + + +0.50.6 (2024-12-10) +------------------- + +Changed +^^^^^^^ + +- Cleaned __main__ with restructuring of code (61e48c2c60e32aa8ac7699979a9836c8d17ded64) + + 0.50.5 (2024-03-24) ------------------- diff --git a/README.md b/README.md index ee8e9883..6a02af2b 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,13 @@ A simple test can be found in [`run_scripts.test_dycore`](https://github.com/ray pytest ./run_scripts/test_dycore.py -v ``` -However, the codebase is structured such that the user can easily assemble a run script to define their own experiments. Refer to the documentation for the [available APIs](https://ray-chew.github.io/pyBELLA/apis.html). +To run a simulation: +```console +pybella -ic rb -N 1 +``` +Note that only the rising bubble initial condition is implemented for now with ensemble size of 1. + +The codebase is structured such that the user can easily assemble a run script to define their own experiments. Refer to the documentation for the [available APIs](https://ray-chew.github.io/pyBELLA/apis.html). ## License @@ -86,3 +92,14 @@ Refer to the [open issues](https://github.com/ray-chew/pyBELLA/issues), in parti Any changes, improvements, or bug fixes can be submitted from your remote to upstream via a pull request. +## To be tidied up + +### Profiler +```bash +$ PYTHONPATH=src scalene -m pybella -ic test_lamb_wave -N 1 +``` + +### Coverage +```bash +$ pytest --cov test_scripts/ --cov-report=html +``` \ No newline at end of file diff --git a/dev-requirements.txt b/dev-requirements.txt index b077c9f3..c609f6a9 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,17 +1,15 @@ -dask==2022.7.0 -dask[distributed]==2022.7.0 -dill==0.3.6 -h5py==3.7.0 -matplotlib==3.5.1 -numba==0.56.4 -numpy==1.22.1 -pyemd==1.0.0 -pytest==8.1.1 -PyYAML==6.0 -scipy==1.7.3 -termcolor==2.4.0 +dill +h5py +matplotlib +numba +numpy +pytest +scipy +dask +pytest +pytest-cov # dependencies for documentation -sphinx==7.2.6 -sphinx_changelog==1.5.0 -sphinx-math-dollar==1.2.1 -furo==2024.1.29 \ No newline at end of file +sphinx +sphinx_changelog +sphinx-math-dollar +furo \ No newline at end of file diff --git a/docs/source/apis/src.dycore.discretisation.rst b/docs/source/apis/src.dycore.discretisation.rst deleted file mode 100644 index f45c9533..00000000 --- a/docs/source/apis/src.dycore.discretisation.rst +++ /dev/null @@ -1,32 +0,0 @@ -src.dycore.discretisation -========================= - -.. automodule:: src.dycore.discretisation - - - - - - - - - - - - - - - - - - - -.. rubric:: Modules - -.. autosummary:: - :toctree: - :recursive: - - src.dycore.discretisation.grid - src.dycore.discretisation.time_update - diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.rst b/docs/source/apis/src.dycore.physics.gas_dynamics.rst deleted file mode 100644 index c6e2eec4..00000000 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.rst +++ /dev/null @@ -1,36 +0,0 @@ -src.dycore.physics.gas\_dynamics -================================ - -.. automodule:: src.dycore.physics.gas_dynamics - - - - - - - - - - - - - - - - - - - -.. rubric:: Modules - -.. autosummary:: - :toctree: - :recursive: - - src.dycore.physics.gas_dynamics.cfl - src.dycore.physics.gas_dynamics.eos - src.dycore.physics.gas_dynamics.explicit - src.dycore.physics.gas_dynamics.numerical_flux - src.dycore.physics.gas_dynamics.recovery - src.dycore.physics.gas_dynamics.thermodynamic - diff --git a/docs/source/apis/src.dycore.physics.low_mach.rst b/docs/source/apis/src.dycore.physics.low_mach.rst deleted file mode 100644 index 64652733..00000000 --- a/docs/source/apis/src.dycore.physics.low_mach.rst +++ /dev/null @@ -1,33 +0,0 @@ -src.dycore.physics.low\_mach -============================ - -.. automodule:: src.dycore.physics.low_mach - - - - - - - - - - - - - - - - - - - -.. rubric:: Modules - -.. autosummary:: - :toctree: - :recursive: - - src.dycore.physics.low_mach.laplacian - src.dycore.physics.low_mach.mpv - src.dycore.physics.low_mach.second_projection - diff --git a/docs/source/apis/src.dycore.physics.rst b/docs/source/apis/src.dycore.physics.rst deleted file mode 100644 index 2fa44ee6..00000000 --- a/docs/source/apis/src.dycore.physics.rst +++ /dev/null @@ -1,33 +0,0 @@ -src.dycore.physics -================== - -.. automodule:: src.dycore.physics - - - - - - - - - - - - - - - - - - - -.. rubric:: Modules - -.. autosummary:: - :toctree: - :recursive: - - src.dycore.physics.gas_dynamics - src.dycore.physics.hydrostatics - src.dycore.physics.low_mach - diff --git a/docs/source/apis/src.dycore.discretisation.grid.rst b/docs/source/apis/src.flow_solver.discretisation.grid.rst similarity index 75% rename from docs/source/apis/src.dycore.discretisation.grid.rst rename to docs/source/apis/src.flow_solver.discretisation.grid.rst index 05d2afdc..e00ef993 100644 --- a/docs/source/apis/src.dycore.discretisation.grid.rst +++ b/docs/source/apis/src.flow_solver.discretisation.grid.rst @@ -1,7 +1,7 @@ -src.dycore.discretisation.grid +src.flow_solver.discretisation.grid ============================== -.. automodule:: src.dycore.discretisation.grid +.. automodule:: src.flow_solver.discretisation.grid diff --git a/docs/source/apis/src.flow_solver.discretisation.rst b/docs/source/apis/src.flow_solver.discretisation.rst new file mode 100644 index 00000000..54d80d3d --- /dev/null +++ b/docs/source/apis/src.flow_solver.discretisation.rst @@ -0,0 +1,32 @@ +src.flow_solver.discretisation +========================= + +.. automodule:: src.flow_solver.discretisation + + + + + + + + + + + + + + + + + + + +.. rubric:: Modules + +.. autosummary:: + :toctree: + :recursive: + + src.flow_solver.discretisation.grid + src.flow_solver.discretisation.time_update + diff --git a/docs/source/apis/src.dycore.discretisation.time_update.rst b/docs/source/apis/src.flow_solver.discretisation.time_update.rst similarity index 62% rename from docs/source/apis/src.dycore.discretisation.time_update.rst rename to docs/source/apis/src.flow_solver.discretisation.time_update.rst index fe2d9a50..f8b3c193 100644 --- a/docs/source/apis/src.dycore.discretisation.time_update.rst +++ b/docs/source/apis/src.flow_solver.discretisation.time_update.rst @@ -1,7 +1,7 @@ -src.dycore.discretisation.time\_update +src.flow_solver.discretisation.time\_update ====================================== -.. automodule:: src.dycore.discretisation.time_update +.. automodule:: src.flow_solver.discretisation.time_update diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.cfl.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.cfl.rst similarity index 63% rename from docs/source/apis/src.dycore.physics.gas_dynamics.cfl.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.cfl.rst index 1a639956..32a0164f 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.cfl.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.cfl.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.cfl +src.flow_solver.physics.gas\_dynamics.cfl ==================================== -.. automodule:: src.dycore.physics.gas_dynamics.cfl +.. automodule:: src.flow_solver.physics.gas_dynamics.cfl diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.eos.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.eos.rst similarity index 73% rename from docs/source/apis/src.dycore.physics.gas_dynamics.eos.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.eos.rst index fd93d49d..bf1829a0 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.eos.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.eos.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.eos +src.flow_solver.physics.gas\_dynamics.eos ==================================== -.. automodule:: src.dycore.physics.gas_dynamics.eos +.. automodule:: src.flow_solver.physics.gas_dynamics.eos diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.explicit.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.explicit.rst similarity index 65% rename from docs/source/apis/src.dycore.physics.gas_dynamics.explicit.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.explicit.rst index e4306f28..5bde211e 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.explicit.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.explicit.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.explicit +src.flow_solver.physics.gas\_dynamics.explicit ========================================= -.. automodule:: src.dycore.physics.gas_dynamics.explicit +.. automodule:: src.flow_solver.physics.gas_dynamics.explicit diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.numerical_flux.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.numerical_flux.rst similarity index 63% rename from docs/source/apis/src.dycore.physics.gas_dynamics.numerical_flux.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.numerical_flux.rst index 637efaaa..18faffc9 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.numerical_flux.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.numerical_flux.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.numerical\_flux +src.flow_solver.physics.gas\_dynamics.numerical\_flux ================================================ -.. automodule:: src.dycore.physics.gas_dynamics.numerical_flux +.. automodule:: src.flow_solver.physics.gas_dynamics.numerical_flux diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.recovery.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.recovery.rst similarity index 66% rename from docs/source/apis/src.dycore.physics.gas_dynamics.recovery.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.recovery.rst index 577eda1a..20f16205 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.recovery.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.recovery.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.recovery +src.flow_solver.physics.gas\_dynamics.recovery ========================================= -.. automodule:: src.dycore.physics.gas_dynamics.recovery +.. automodule:: src.flow_solver.physics.gas_dynamics.recovery diff --git a/docs/source/apis/src.flow_solver.physics.gas_dynamics.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.rst new file mode 100644 index 00000000..ed8ade27 --- /dev/null +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.rst @@ -0,0 +1,36 @@ +src.flow_solver.physics.gas\_dynamics +================================ + +.. automodule:: src.flow_solver.physics.gas_dynamics + + + + + + + + + + + + + + + + + + + +.. rubric:: Modules + +.. autosummary:: + :toctree: + :recursive: + + src.flow_solver.physics.gas_dynamics.cfl + src.flow_solver.physics.gas_dynamics.eos + src.flow_solver.physics.gas_dynamics.explicit + src.flow_solver.physics.gas_dynamics.numerical_flux + src.flow_solver.physics.gas_dynamics.recovery + src.flow_solver.physics.gas_dynamics.thermodynamic + diff --git a/docs/source/apis/src.dycore.physics.gas_dynamics.thermodynamic.rst b/docs/source/apis/src.flow_solver.physics.gas_dynamics.thermodynamic.rst similarity index 58% rename from docs/source/apis/src.dycore.physics.gas_dynamics.thermodynamic.rst rename to docs/source/apis/src.flow_solver.physics.gas_dynamics.thermodynamic.rst index 012c83ec..b43d9272 100644 --- a/docs/source/apis/src.dycore.physics.gas_dynamics.thermodynamic.rst +++ b/docs/source/apis/src.flow_solver.physics.gas_dynamics.thermodynamic.rst @@ -1,7 +1,7 @@ -src.dycore.physics.gas\_dynamics.thermodynamic +src.flow_solver.physics.gas\_dynamics.thermodynamic ============================================== -.. automodule:: src.dycore.physics.gas_dynamics.thermodynamic +.. automodule:: src.flow_solver.physics.gas_dynamics.thermodynamic diff --git a/docs/source/apis/src.dycore.physics.hydrostatics.rst b/docs/source/apis/src.flow_solver.physics.hydrostatics.rst similarity index 67% rename from docs/source/apis/src.dycore.physics.hydrostatics.rst rename to docs/source/apis/src.flow_solver.physics.hydrostatics.rst index ae2e2adb..79e833ef 100644 --- a/docs/source/apis/src.dycore.physics.hydrostatics.rst +++ b/docs/source/apis/src.flow_solver.physics.hydrostatics.rst @@ -1,7 +1,7 @@ -src.dycore.physics.hydrostatics +src.flow_solver.physics.hydrostatics =============================== -.. automodule:: src.dycore.physics.hydrostatics +.. automodule:: src.flow_solver.physics.hydrostatics diff --git a/docs/source/apis/src.dycore.physics.low_mach.laplacian.rst b/docs/source/apis/src.flow_solver.physics.low_mach.laplacian.rst similarity index 81% rename from docs/source/apis/src.dycore.physics.low_mach.laplacian.rst rename to docs/source/apis/src.flow_solver.physics.low_mach.laplacian.rst index acb21b8a..53bfa5e4 100644 --- a/docs/source/apis/src.dycore.physics.low_mach.laplacian.rst +++ b/docs/source/apis/src.flow_solver.physics.low_mach.laplacian.rst @@ -1,7 +1,7 @@ -src.dycore.physics.low\_mach.laplacian +src.flow_solver.physics.low\_mach.laplacian ====================================== -.. automodule:: src.dycore.physics.low_mach.laplacian +.. automodule:: src.flow_solver.physics.low_mach.laplacian diff --git a/docs/source/apis/src.dycore.physics.low_mach.mpv.rst b/docs/source/apis/src.flow_solver.physics.low_mach.mpv.rst similarity index 61% rename from docs/source/apis/src.dycore.physics.low_mach.mpv.rst rename to docs/source/apis/src.flow_solver.physics.low_mach.mpv.rst index 2a96391f..abfa1b22 100644 --- a/docs/source/apis/src.dycore.physics.low_mach.mpv.rst +++ b/docs/source/apis/src.flow_solver.physics.low_mach.mpv.rst @@ -1,7 +1,7 @@ -src.dycore.physics.low\_mach.mpv +src.flow_solver.physics.low\_mach.mpv ================================ -.. automodule:: src.dycore.physics.low_mach.mpv +.. automodule:: src.flow_solver.physics.low_mach.mpv diff --git a/docs/source/apis/src.flow_solver.physics.low_mach.rst b/docs/source/apis/src.flow_solver.physics.low_mach.rst new file mode 100644 index 00000000..3df4d745 --- /dev/null +++ b/docs/source/apis/src.flow_solver.physics.low_mach.rst @@ -0,0 +1,33 @@ +src.flow_solver.physics.low\_mach +============================ + +.. automodule:: src.flow_solver.physics.low_mach + + + + + + + + + + + + + + + + + + + +.. rubric:: Modules + +.. autosummary:: + :toctree: + :recursive: + + src.flow_solver.physics.low_mach.laplacian + src.flow_solver.physics.low_mach.mpv + src.flow_solver.physics.low_mach.second_projection + diff --git a/docs/source/apis/src.dycore.physics.low_mach.second_projection.rst b/docs/source/apis/src.flow_solver.physics.low_mach.second_projection.rst similarity index 81% rename from docs/source/apis/src.dycore.physics.low_mach.second_projection.rst rename to docs/source/apis/src.flow_solver.physics.low_mach.second_projection.rst index 9d9cb675..e725be47 100644 --- a/docs/source/apis/src.dycore.physics.low_mach.second_projection.rst +++ b/docs/source/apis/src.flow_solver.physics.low_mach.second_projection.rst @@ -1,7 +1,7 @@ -src.dycore.physics.low\_mach.second\_projection +src.flow_solver.physics.low\_mach.second\_projection =============================================== -.. automodule:: src.dycore.physics.low_mach.second_projection +.. automodule:: src.flow_solver.physics.low_mach.second_projection diff --git a/docs/source/apis/src.flow_solver.physics.rst b/docs/source/apis/src.flow_solver.physics.rst new file mode 100644 index 00000000..3ef67ab3 --- /dev/null +++ b/docs/source/apis/src.flow_solver.physics.rst @@ -0,0 +1,33 @@ +src.flow_solver.physics +================== + +.. automodule:: src.flow_solver.physics + + + + + + + + + + + + + + + + + + + +.. rubric:: Modules + +.. autosummary:: + :toctree: + :recursive: + + src.flow_solver.physics.gas_dynamics + src.flow_solver.physics.hydrostatics + src.flow_solver.physics.low_mach + diff --git a/docs/source/apis/src.dycore.rst b/docs/source/apis/src.flow_solver.rst similarity index 55% rename from docs/source/apis/src.dycore.rst rename to docs/source/apis/src.flow_solver.rst index e622abfd..527cd633 100644 --- a/docs/source/apis/src.dycore.rst +++ b/docs/source/apis/src.flow_solver.rst @@ -1,7 +1,7 @@ -src.dycore +src.flow_solver ========== -.. automodule:: src.dycore +.. automodule:: src.flow_solver @@ -27,6 +27,6 @@ src.dycore :toctree: :recursive: - src.dycore.discretisation - src.dycore.physics + src.flow_solver.discretisation + src.flow_solver.physics diff --git a/docs/source/apis/src.rst b/docs/source/apis/src.rst index c7df569f..233812ee 100644 --- a/docs/source/apis/src.rst +++ b/docs/source/apis/src.rst @@ -28,7 +28,7 @@ :recursive: src.data_assimilation - src.dycore + src.flow_solver src.tests src.utils src.vis diff --git a/docs/source/conf.py b/docs/source/conf.py index a3f661f0..35360dc7 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -20,9 +20,9 @@ # -- Project information ----------------------------------------------------- -project = 'pyBELLA' -copyright = '2024, Ray Chew, Tommaso Benacchio, Rupert Klein' -author = 'Ray Chew, Tommaso Benacchio, Rupert Klein' +project = "pyBELLA" +copyright = "2024, Ray Chew, Tommaso Benacchio, Rupert Klein" +author = "Ray Chew, Tommaso Benacchio, Rupert Klein" # The full version, including alpha/beta/rc tags release = "v0.50.1" @@ -33,26 +33,26 @@ # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.autosummary', - 'sphinx_changelog', - 'sphinx.ext.doctest', - 'sphinx.ext.graphviz', - 'sphinx.ext.imgconverter', - 'sphinx.ext.todo', - 'sphinx_math_dollar', - 'sphinx.ext.mathjax' + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx_changelog", + "sphinx.ext.doctest", + "sphinx.ext.graphviz", + "sphinx.ext.imgconverter", + "sphinx.ext.todo", + "sphinx_math_dollar", + "sphinx.ext.mathjax", ] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. # # This is also used if you do content translation via gettext catalogs. # Usually you set "language" from the command line for these cases. -language = 'en' +language = "en" # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. @@ -89,9 +89,9 @@ # see details: https://www.sympy.org/sphinx-math-dollar/ mathjax3_config = { - 'tex2jax': { - 'inlineMath': [ ["\\(","\\)"] ], - 'displayMath': [["\\[","\\]"] ], + "tex2jax": { + "inlineMath": [["\\(", "\\)"]], + "displayMath": [["\\[", "\\]"]], }, } @@ -101,4 +101,4 @@ # -- Options for todo extension ---------------------------------------------- # If true, `todo` and `todoList` produce output, else they produce nothing. -todo_include_todos = True \ No newline at end of file +todo_include_todos = True diff --git a/outputs/target_blending_warm_bubble/target_blending_warm_bubble_64_48_stripped.h5 b/outputs/target_blending_warm_bubble/target_blending_warm_bubble_64_48_stripped.h5 new file mode 100644 index 00000000..aedaf501 Binary files /dev/null and b/outputs/target_blending_warm_bubble/target_blending_warm_bubble_64_48_stripped.h5 differ diff --git a/outputs/target_internal_long_wave/target_internal_long_wave_301_10_stripped.h5 b/outputs/target_internal_long_wave/target_internal_long_wave_301_10_stripped.h5 new file mode 100644 index 00000000..119ed571 Binary files /dev/null and b/outputs/target_internal_long_wave/target_internal_long_wave_301_10_stripped.h5 differ diff --git a/outputs/target_lamb_wave/target_lamb_wave_151_60_stripped.h5 b/outputs/target_lamb_wave/target_lamb_wave_151_60_stripped.h5 new file mode 100644 index 00000000..7bd460b8 Binary files /dev/null and b/outputs/target_lamb_wave/target_lamb_wave_151_60_stripped.h5 differ diff --git a/outputs/target_travelling_vortex/target_travelling_vortex_64_64_stripped.h5 b/outputs/target_travelling_vortex/target_travelling_vortex_64_64_stripped.h5 new file mode 100644 index 00000000..182a4ef9 Binary files /dev/null and b/outputs/target_travelling_vortex/target_travelling_vortex_64_64_stripped.h5 differ diff --git a/outputs/target_unstable_lamb/target_unstable_lamb_301_30_stripped.h5 b/outputs/target_unstable_lamb/target_unstable_lamb_301_30_stripped.h5 new file mode 100644 index 00000000..3caef678 Binary files /dev/null and b/outputs/target_unstable_lamb/target_unstable_lamb_301_30_stripped.h5 differ diff --git a/pyproject.toml b/pyproject.toml index a29c2eaf..0c926353 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,28 +1,25 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + [project] -name = "pyBELLA" +name = "pybella" version = "0.50.1" dependencies = [ - "dask==2022.7.0", - "dill==0.3.6", - "h5py==3.7.0", - "matplotlib==3.5.1", - "numba==0.56.4", - "numpy==1.22.1", - "PyYAML==6.0", - "scipy==1.7.3", - "termcolor==2.4.0" + "dask", + "dill", + "h5py", + "matplotlib", + "numba", + "numpy", + "PyYAML", + "scipy", + "termcolor" ] - -# Packaging -[build-system] -requires = ["setuptools"] -build-backend = "setuptools.build_meta" - -[tool.setuptools] -package-dir = {"pybella" = "src"} - +[project.scripts] +pybella = "pybella.__main__:main" # Change Log [tool.towncrier] diff --git a/requirements.txt b/requirements.txt index 26d29b2b..2ff59943 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,10 @@ -dask==2022.7.0 -dask[distributed]==2022.7.0 -dill==0.3.6 -h5py==3.7.0 -matplotlib==3.5.1 -numba==0.56.4 -numpy==1.22.1 -pyemd==1.0.0 -pytest==8.1.1 -PyYAML==6.0 -scipy==1.7.3 -termcolor==2.4.0 +dill +h5py +matplotlib +numba +numpy +pytest +scipy +dask +pytest +pytest-cov diff --git a/run_scripts/run.py b/run_scripts/driver.py similarity index 91% rename from run_scripts/run.py rename to run_scripts/driver.py index c014579e..672dea3a 100644 --- a/run_scripts/run.py +++ b/run_scripts/driver.py @@ -7,7 +7,7 @@ class run_params(object): N = 1 tc = "rb" # tc = 'mark' - tc = "tv" + # tc = "tv" def __init__(self): self.N = self.N @@ -18,21 +18,18 @@ def __init__(self): self.restart = False def single_run(self): - subprocess.call( - [sys.executable, "src/__main__.py", "-ic", self.tc, "-N", "%i" % self.N] - ) + subprocess.run(["pybella", "-ic", self.tc, "-N", "1"]) def queue_run(self): if self.ud is None and self.dap is None: assert 0, "ud or params must be defined" - subprocess.call( + subprocess.run( [ - sys.executable, - "src/__main__.py", + "pybella", "-ic", self.tc, "-N", - "%i" % self.N, + f"{self.N}", "queue", "-w", self.ud, diff --git a/run_scripts/test_dycore.py b/run_scripts/test_dycore.py deleted file mode 100644 index 3460b199..00000000 --- a/run_scripts/test_dycore.py +++ /dev/null @@ -1,16 +0,0 @@ -import pytest -import sys -import subprocess - -@pytest.mark.parametrize("ic", - ["test_travelling_vortex", - "test_internal_long_wave", - "test_lamb_wave"]) -def test_single_run(ic): - run = subprocess.Popen( - [sys.executable, "src/__main__.py", "-ic", ic, "-N", "1"], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - _, stderr = run.communicate() - assert run.returncode == 0, stderr.splitlines()[-3:] \ No newline at end of file diff --git a/run_scripts/test_suite.py b/run_scripts/test_suite.py index c5b97e7c..4cd596a5 100644 --- a/run_scripts/test_suite.py +++ b/run_scripts/test_suite.py @@ -1,5 +1,5 @@ # %% -from run import run_params as rp +from .driver import run_params as rp import pybella.tests.diagnostics as td import json @@ -17,46 +17,47 @@ rp.N = 1 ud = {} -if gen_targets: - # set target user data parameters - ud["output_type"] = "target" - ud["diag"] = False +if __name__ == "__main__": + if gen_targets: + # set target user data parameters + ud["output_type"] = "target" + ud["diag"] = False + + # generate target + # define horizontal slice test case + rp.tc = "test_travelling_vortex" + rp.ud = json.dumps(ud) + rp.queue_run() + + # define vertical slice test case + rp.tc = "test_internal_long_wave" + rp.ud = json.dumps(ud) + rp.queue_run() + + # define test case + rp.tc = "test_lamb_wave" + rp.ud = json.dumps(ud) + rp.queue_run() + + if updt_targets: + diag = td.compare_sol("gen_target") + diag.update_targets() + + + ud["output_type"] = "test" + # Do diagnostics + ud["diag"] = True - # generate target - # define horizontal slice test case rp.tc = "test_travelling_vortex" rp.ud = json.dumps(ud) rp.queue_run() - # define vertical slice test case rp.tc = "test_internal_long_wave" rp.ud = json.dumps(ud) rp.queue_run() - # define test case rp.tc = "test_lamb_wave" rp.ud = json.dumps(ud) rp.queue_run() -if updt_targets: - diag = td.compare_sol("gen_target") - diag.update_targets() - - -ud["output_type"] = "test" -# Do diagnostics -ud["diag"] = True - -rp.tc = "test_travelling_vortex" -rp.ud = json.dumps(ud) -rp.queue_run() - -rp.tc = "test_internal_long_wave" -rp.ud = json.dumps(ud) -rp.queue_run() - -rp.tc = "test_lamb_wave" -rp.ud = json.dumps(ud) -rp.queue_run() - # %% diff --git a/src/__init__.py b/src/__init__.py deleted file mode 100644 index 8d1c8b69..00000000 --- a/src/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/src/__main__.py b/src/__main__.py deleted file mode 100644 index 211adfbc..00000000 --- a/src/__main__.py +++ /dev/null @@ -1,400 +0,0 @@ -import numpy as np - -# dependencies of the atmospheric flow solver -import dycore.discretisation.grid as dis_grid -import dycore.discretisation.time_update as dis_time_update -import dycore.utils.boundary as bdry -import dycore.utils.variable as var -import dycore.physics.low_mach.mpv as lm_var -import dycore.physics.hydrostatics as hydrostatic -import dycore.physics.gas_dynamics.thermodynamic as gd_thermodynamics - -# dependencies of the parallelisation by dask -from dask.distributed import Client, progress - -# dependencies of the data assimilation subpackag -from data_assimilation import etpf as da_etpf -from data_assimilation import blending as da_blending -from data_assimilation import post_processing as da_post_processing -from data_assimilation import letkf as da_letkf -from data_assimilation import params as da_params -from data_assimilation import utils as da_utils - -# input file -import utils.user_data as user_data -import utils.io as io -import utils.sim_params as params - -# some diagnostics -import copy -import time -import termcolor -import logging - -# test module -import tests.diagnostics as diag - -debug =params.debug -da_debug = params.da_debug -output_timesteps = False -if debug == True: - output_timesteps = True -label_type = "TIME" -np.set_printoptions(precision = params.print_precision) - -step = 0 -t = 0.0 - -########################################################## -# Initialisation of data containers and helper classes -########################################################## -# get arguments for initial condition and ensemble size -N, UserData, sol_init, restart, ud_rewrite, dap_rewrite, r_params = io.get_args() -if N == 1: - da_debug = False - -initial_data = vars(UserData()) -ud = user_data.UserDataInit(**initial_data) -if ud_rewrite is not None: - ud.update_ud(ud_rewrite) -if hasattr(ud, "rayleigh_bc"): - ud.rayleigh_bc(ud) -if ud.output_timesteps: - output_timesteps = True -ud.coriolis_strength = np.array(ud.coriolis_strength) - -elem, node = dis_grid.grid_init(ud) - -Sol = var.Vars(elem.sc, ud) - -flux = np.empty((3), dtype=object) -flux[0] = var.States(elem.sfx, ud) -if elem.ndim > 1: - flux[1] = var.States(elem.sfy, ud) -if elem.ndim > 2: - flux[2] = var.States(elem.sfz, ud) - -th = gd_thermodynamics.init(ud) -mpv = lm_var.MPV(elem, node, ud) -bld = da_blending.Blend(ud) - -io.init_logger(ud) - - -########################################################## -# Initialise test module -########################################################## -if ud.diag: - diag_comparison = diag.compare_sol(ud.diag_current_run) - -########################################################## -# Initialisation of data assimilation module -########################################################## - -# possible da_types: -# 1) batch_obs for the LETKF with batch observations -# 2) rloc for LETKF with grid-point localisation -# 3) etpf for the ETPF algorithm -dap = da_params.init(N, da_type="rloc") -if dap_rewrite is not None: - dap.update_dap(dap_rewrite) - -# if elem.ndim == 2: -if dap.da_type == "rloc" and N > 1: - rloc = da_letkf.prepare_rloc(ud, elem, node, dap, N) - -logging.info(termcolor.colored("Generating initial ensemble...", "yellow")) -sol_ens = np.zeros((N), dtype=object) - -# Set random seed for reproducibility -np.random.seed(params.random_seed) - -seeds = np.random.randint(10000, size=N) if N > 1 else None -if seeds is not None and restart == False: - logging.info("Seeds used in generating initial ensemble spread = ", seeds) - for n in range(N): - Sol0 = copy.deepcopy(Sol) - mpv0 = copy.deepcopy(mpv) - Sol0 = sol_init(Sol0, mpv0, elem, node, th, ud, seed=seeds[n]) - sol_ens[n] = [Sol0, copy.deepcopy(flux), mpv0, [-np.inf, step]] -elif restart == False: - sol_ens = [[sol_init(Sol, mpv, elem, node, th, ud), flux, mpv, [-np.inf, step]]] -elif restart == True: - hydrostatic.state(mpv, elem, node, th, ud) - ud.old_suffix = np.copy(ud.output_suffix) - ud.old_suffix = "_ensemble=%i%s" % (N, ud.old_suffix) - Sol0, mpv0, touts = io.sim_restart( - r_params[0], r_params[1], elem, node, ud, Sol, mpv, r_params[2] - ) - sol_ens = [[Sol0, flux, mpv0, [-np.inf, step]]] - # ud.tout = touts[1:] - ud.tout = [touts[-1]] - t = touts[0] - - if ud.bdry_type[1].value == "radiation": - ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) - -ens = da_utils.ensemble(sol_ens) - -########################################################## -# Load data assimilation observations -########################################################## - -# where are my observations? -if N > 1: - obs = dap.load_obs(dap.obs_path) - # obs_mask, no calculations where entries are True - obs_mask = da_utils.sparse_obs_selector(obs, elem, node, ud, dap) - obs_noisy, obs_covar = da_utils.obs_noiser(obs, obs_mask, dap, rloc, elem) - # obs_noisy_interp, obs_mask = sparse_obs_selector(obs_noisy, elem, node, ud, dap) - - -# add ensemble info to filename -if ud.autogen_fn: - ud.output_suffix = io.fn_gen(ud, dap, N) -# ud.output_suffix = '_ensemble=%i%s' %(N, ud.output_suffix) - -# ud.output_suffix = '%s_%s' %(ud.output_suffix, 'nr') - -########################################################## -# Start main looping -########################################################## -if __name__ == "__main__": - - ###################################################### - # Initialise writer class for I/O operations - ###################################################### - writer = io.init(ud, restart) - writer.check_jar() - writer.jar([ud, mpv, elem, node, dap]) - # sys.exit("Let's just dill the stuff and quit!") - - writer.write_attrs() - wrtr = None - if N > 1: - writer.write_da_attrs(dap) - elif output_timesteps == True: - wrtr = writer - for n in range(N): # write initial ensemble - Sol = ens.members(ens)[n][0] - mpv = ens.members(ens)[n][2] - if label_type == "STEP": - label = "ensemble_mem=%i_%.3d" % (n, step) - else: - label = "ensemble_mem=%i_%.3f" % (n, 0.0) - if not restart: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_ic") - - if da_debug: - # writer.jar([obs,obs_noisy,obs_noisy_interp,obs_mask,obs_covar]) - # obs = obs_noisy_interp - writer.jar([obs, obs_noisy, obs_mask, obs_covar]) - - # initialise dask parallelisation and timer - # client = Client(threads_per_worker=1, n_workers=1) - tic = time.time() - - ###################################################### - # Time looping over data assimilation windows - ###################################################### - tout_old = -np.inf - tout_cnt = 0 - outer_step = 0 - for tout in ud.tout: - futures = [] - - # In ensemble case, do blending for each DA window - if N > 1: - blend = bld if tout_old in dap.da_times else None - else: - blend = bld - - # initial blending? - if ud.initial_blending == True and (outer_step == 0 or outer_step == 1): - blend = bld - - ###################################################### - # Forecast step - ###################################################### - logging.info("##############################################") - logging.info(termcolor.colored("Next tout = %.3f" % tout, "yellow")) - logging.info(termcolor.colored("Starting forecast...", "green")) - mem_cnt = 0 - for mem in ens.members(ens): - # future = client.submit(time_update, *[mem[0],mem[1],mem[2], t, tout, ud, elem, node, mem[3], th, bld, None, False]) - - # handling of DA window step counter - if N > 1: - mem[3][0] = 0 if tout_old in dap.da_times else mem[3][0] - if N == 1: - mem[3][0] = mem[3][1] - logging.info(termcolor.colored("For ensemble member = %i..." % mem_cnt, "yellow")) - future = dis_time_update.do( - mem[0], - mem[1], - mem[2], - t, - tout, - ud, - elem, - node, - mem[3], - th, - blend, - wrtr, - debug, - ) - - if ud.diag: - diag_comparison.test_do( - future[0], future[2].p2_nodes, plot=ud.diag_plot_compare - ) - - futures.append(future) - mem_cnt += 1 - - # Dask commands, used only when parallelisation is - # enabled - # results = client.gather(futures) - results = np.copy(futures) - results = np.array(results) - # s_res = client.scatter(results) - - ###################################################### - # Analysis step - ###################################################### - tout = np.around(tout, 3) - if N > 1 and tout in dap.da_times: - futures = [] - - ###################################################### - # Update ensemble with forecast - ###################################################### - for n in range(N): - Sol = results[n][dap.loc_c] - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - results[n][dap.loc_c] = Sol - p2_nodes = getattr(results[n][dap.loc_n], "p2_nodes") - bdry.set_ghostnodes_p2(p2_nodes, node, ud) - setattr(results[n][dap.loc_n], "p2_nodes", p2_nodes) - - ens.set_members(results, tout) - - ###################################################### - # Write output before assimilating data - ###################################################### - logging.info(termcolor.colored("Starting output...", "yellow")) - for n in range(N): - Sol = ens.members(ens)[n][0] - mpv = ens.members(ens)[n][2] - - if label_type == "STEP": - step = outer_step - label = "ensemble_mem=%i_%.3d" % (n, step) - else: - label = "ensemble_mem=%i_%.3f" % (n, tout) - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_before_da") - - ################################################## - # LETKF with batch observations - ################################################## - if dap.da_type == "batch_obs": - logging.info("Starting analysis... for batch observations") - for attr in dap.obs_attributes: - logging.info("Assimilating %s..." % attr) - logging.info("Assimilating %s..." % attr) - # future = client.submit(da_interface, *[s_res,obs_current,dap.inflation_factor,attr,N,ud,dap.loc[attr]]) - future = da_letkf.da_interface(results, dap, obs, attr, tout, N, ud) - futures.append(future) - - # analysis = client.gather(futures) - analysis = futures - # analysis = np.array(analysis) - - logging.info("Writing analysis...") - cnt = 0 - for attr in dap.obs_attributes: - current = analysis[cnt] - for n in range(N): - setattr(results[:, dap.loc[attr], ...][n], attr, current[n]) - cnt += 1 - - ################################################## - # LETKF with grid-point localisation - ################################################## - elif dap.da_type == "rloc": - logging.info( - termcolor.colored("Starting analysis... for rloc algorithm", "green") - ) - results = da_utils.HSprojector_3t2D(results, elem, dap, N) - results = rloc.analyse(results, obs, obs_covar, obs_mask, N, tout) - results = da_utils.HSprojector_2t3D(results, elem, node, dap, N) - # if hasattr(dap, 'converter'): - # results = dap.converter(results, N, mpv, elem, node, th, ud) - - ################################################## - # ETPF - ################################################## - elif dap.da_type == "etpf": - da_utils.ensemble_inflation(results, dap.attributes, dap.inflation_factor, N) - results = da_etpf.da_interface( - results, - obs, - dap.obs_attributes, - dap.rejuvenation_factor, - dap.da_times, - tout, - N, - ) - - ################################################## - # Post-processing - ################################################## - elif dap.da_type == "pprocess": - results = da_post_processing.interface() - - else: - assert 0, "DA type not implemented: use 'rloc', 'batch_obs' or 'etpf'." - - ###################################################### - # Update ensemble with analysis - ###################################################### - for n in range(N): - Sol = results[n][dap.loc_c] - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - results[n][dap.loc_c] = Sol - p2_nodes = getattr(results[n][dap.loc_n], "p2_nodes") - bdry.set_ghostnodes_p2(p2_nodes, node, ud) - setattr(results[n][dap.loc_n], "p2_nodes", p2_nodes) - - ens.set_members(results, tout) - - ###################################################### - # Write output at tout - ###################################################### - logging.info(termcolor.colored("Starting output...", "yellow")) - for n in range(N): - Sol = ens.members(ens)[n][0] - mpv = ens.members(ens)[n][2] - - if label_type == "STEP": - step = outer_step - label = "ensemble_mem=%i_%.3d" % (n, step) - else: - label = "ensemble_mem=%i_%.3f" % (n, tout) - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_full_step") - - # synchronise_variables(mpv, Sol, elem, node, ud, th) - t = tout - tout_old = np.copy(tout) - logging.info(termcolor.colored("tout = %.3f" % tout, "yellow")) - - tout_cnt += 1 - outer_step += 1 - if outer_step > ud.stepmax: - break - - toc = time.time() - logging.info(termcolor.colored("Time taken = %.6f" % (toc - tic), "yellow")) - - writer.close_everything() diff --git a/src/data_assimilation/etpf.py b/src/data_assimilation/etpf.py deleted file mode 100644 index 0f0dfbc3..00000000 --- a/src/data_assimilation/etpf.py +++ /dev/null @@ -1,106 +0,0 @@ -import numpy as np -import logging - -def da_interface(results, obs, obs_attributes, delta, times, tout, N, loc=0): - ig = 2 - inner = (slice(ig,-ig),slice(ig,-ig)) - - - # local_ens = np.array([getattr(results[:,loc,...][n],attr) for n in range(N)]) - # local_ens = np.array([mem[inner] for mem in local_ens]) - # local_ens = analysis(local_ens,attr) - # print(results[:,0,...][0].rho) - attributes = ['rho','rhou','rhov','rhow','rhoY','rhoX'] - # attributes = ['rho', 'rhou', 'rhov'] - tmp = np.array([getattr(results[:,loc,...][n],obs_attributes[0])[inner] for n in range(N)]) - tmp = tmp[:,np.newaxis,...] - - ensemble = np.array([getattr(results[:,loc,...][n],attributes[0])[inner] for n in range(N)]) - ensemble = ensemble[:,np.newaxis,...] - - obs_current = np.array(obs[np.where(np.isclose(times,tout))[0][0]][obs_attributes[0]])[inner] - - # obs_current = bin_func(obs_current,(Nx,Ny)) - obs_current = obs_current[np.newaxis,...] - - for attr in obs_attributes[1:]: - tmp = np.hstack((tmp,np.array([getattr(results[:,loc,...][n],attr)[inner] for n in range(N)])[:,np.newaxis,...])) - tmp01 = np.array(obs[np.where(np.isclose(times,tout))[0][0]][attr])[inner] - # tmp01 = bin_func(tmp01,(Nx,Ny)) - # print(tmp01.shape) - tmp01 = tmp01[np.newaxis,...] - obs_current = np.vstack((obs_current,tmp01)) - - for attr in attributes[1:]: - ensemble = np.hstack((ensemble,np.array([getattr(results[:,loc,...][n],attr)[inner] for n in range(N)])[:,np.newaxis,...])) - - - obs_current = obs_current[np.newaxis,...] - # r = tmp - obs_current - - Hx = tmp.reshape(N,-1) - obs_current = obs_current.reshape(-1) - # ensemble = ensemble.reshape(N,-1) - # print(ensemble.shape) - - etpf = analysis(ensemble, delta) - etpf.analyse(obs_current,1.0,Hx,N) - - analysis_ens = etpf.get_ensemble_from_X() - - for n in range(N): - cnt = 0 - current = analysis_ens[n] - for attr in attributes: - data = current[cnt] - data = np.pad(data,2,mode='wrap') - - setattr(results[:,loc,...][n],attr,data) - cnt += 1 - return results - -class analysis(object): - def __init__(self,ensemble,delta,identifier=None): - self.ensemble = np.array(ensemble) - self.ensemble_shape = self.ensemble.shape - - self.X = self.state_vector(ensemble) - self.identifier = identifier - - # rejuvenation factor - self.delta = delta - - def analyse(self,obs_current,obs_covar,Hx,N): - import pyemd - logging.info("starting ETPF analysis...") - - r = (Hx - obs_current)**2 - r = np.sum(r, axis=1) - # print(r.shape) - - ww = np.exp(-r / (2. * obs_covar)) - ww /= np.sum(ww) - - # print(ww) - - Co = self.X @ self.X.T - diag = np.diag(Co) - Co = diag * np.ones((1,N)) - 2. * Co + np.ones((N,1)) * diag.T - - # print(Co) - - _, T = pyemd.emd_with_flow(ww,np.ones(N)/N, Co, -1) - T = np.array(T) - T = T*N - - self.X = np.dot(self.X.T,T).T #+ self.delta * np.random.randn(self.X.shape[0],self.X.shape[1]) - # print(self.X.shape) - - @staticmethod - def state_vector(ensemble): - return ensemble.reshape(ensemble.shape[0],-1) - - def get_ensemble_from_X(self): - return self.X.reshape(self.ensemble_shape) - - \ No newline at end of file diff --git a/src/dycore/discretisation/time_update.py b/src/dycore/discretisation/time_update.py deleted file mode 100644 index c3432c54..00000000 --- a/src/dycore/discretisation/time_update.py +++ /dev/null @@ -1,708 +0,0 @@ -# dependencies of the atmospheric flow solver -import utils.io as io - -import dycore.utils.boundary as bdry -import dycore.utils.options as opts - -import dycore.physics.gas_dynamics.numerical_flux as gd_flux -import dycore.physics.gas_dynamics.explicit as gd_explicit -import dycore.physics.gas_dynamics.eos as gd_eos -import dycore.physics.gas_dynamics.cfl as gd_cfl - -import dycore.physics.low_mach.second_projection as lm_sp -import dycore.discretisation.grid as dis_grid - -# for blending module -import data_assimilation as da - -import numpy as np -import copy -import termcolor -import logging - - -def data_init(ud): - """ - Helper function to initialise the `elem` and `node` grids, corresponding to the cell and node grids, from a given user iniital data file. - - Parameters - ---------- - ud : :class:`inputs.user_data.UserDataInit` - Data container for the initial conditions. - - Returns - ------- - elem : :class:`discretization.kgrid.ElemSpaceDiscr` - Cells grid. - node : :class:`discretization.kgrid.NodeSpaceDiscr` - Nodes grid. - - """ - inx = ud.inx - iny = ud.iny - inz = ud.inz - x0 = ud.xmin - x1 = ud.xmax - y0 = ud.ymin - y1 = ud.ymax - z0 = ud.zmin - z1 = ud.zmax - - grid = dis_grid.Grid(inx, iny, inz, x0, x1, y0, y1, z0, z1) - - elem = dis_grid.ElemSpaceDiscr(grid, ud) - node = dis_grid.NodeSpaceDiscr(grid, ud) - - return elem, node - - -def do( - Sol, - flux, - mpv, - t, - tout, - ud, - elem, - node, - steps, - th, - bld=None, - writer=None, - debug=False, -): - """ - For more details, refer to the write-up :ref:`time-stepping`. - - Does a time-step for the atmospheric solver. - - Parameters - ---------- - Sol : :class:`management.variable.Vars` - Solution data container. - flux : :class:`management.variable.States` - Data container for the fluxes. - mpv : :class:`physics.low_mach.mpv.MPV` - Variables relating to the elliptic solver. - t : float - Current time - tout : float - Next output time - ud : :class:`inputs.user_data.UserDataInit` - Data container for the initial conditions - elem : :class:`discretization.kgrid.ElemSpaceDiscr` - Cells grid. - node : :class:`discretization.kgrid.NodeSpaceDiscr` - Nodes grid. - step : int - Current step. - th : :class:`physics.gas_dynamics.thermodynamic.init` - Thermodynamic variables of the system - bld : :class:`data_assimilation.blending.Blend()` - Blending class used to initalise interface blending methods. - writer : :class:`management.io.io`, optional - `default == None`. If given, output after each time-step will be written in the hdf5 format. - debug : boolean, optional - `default == False`. If `True`, then writer will output `Sol`: - 1. before flux calculation - 2. before advection routine - 3. after advection routine - 4. after explicit solver - 5. after implicit solver - - during both the half-step for the prediction of advective flux and the full-step. - - Returns - ------- - list - A list of `[Sol,flux,mpv,[window_step,step]]` data containers at time `tout`. - """ - - window_step = steps[0] - step = steps[1] - swe_to_lake = False - if "best" not in ud.aux: - test_hydrob = True - else: - test_hydrob = False - - while (t < tout) and (step < ud.stepmax): - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - label = "%.3d" % step - - if step == 0 and writer != None: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_ic") - - dt, cfl, cfl_ac = gd_cfl.dynamic_timestep(Sol, t, tout, elem, ud, th, step) - - if "CFLfixed" in ud.aux: - if step < 2: - dt = 21.69 / ud.t_ref - - c1 = ( - step == 0 - and ud.is_nonhydrostatic == 0 - and bld is not None - and "imbal" in ud.aux - ) - c2 = ( - step == 0 - and ud.is_nonhydrostatic == 1 - and ud.initial_blending == True - and bld is not None - and "imbal" in ud.aux - ) - - if c1 or c2: - logging.info(termcolor.colored("nonhydrostatic to hydrostatic conversion...", "blue")) - ud.is_nonhydrostatic = 0 - if test_hydrob == False: - dt *= 0.5 - # elif test_hydrob == True: - # dt *= 0.25 - - ###################################################### - # Blending : Do blending before timestep - ###################################################### - swe_to_lake, Sol, mpv, t = da.blending.blending_before_timestep( - Sol, - flux, - mpv, - bld, - elem, - node, - th, - ud, - label, - writer, - step, - window_step, - t, - dt, - swe_to_lake, - debug, - ) - - ud.is_nonhydrostatic = gd_eos.is_nonhydrostatic(ud, window_step) - ud.nonhydrostasy = gd_eos.nonhydrostasy(ud, t, window_step) - - if ud.continuous_blending == True or ud.initial_blending == True: - if window_step >= 0: - logging.info("step = %i, window_step = %i" % (step, window_step)) - else: - logging.info("step = %i, window_step = %f" % (step, window_step)) - logging.info( - "is_compressible = %i, is_nonhydrostatic = %i" - % (ud.is_compressible, ud.is_nonhydrostatic) - ) - logging.info( - "compressibility = %.3f, nonhydrostasy = %.3f" - % (ud.compressibility, ud.nonhydrostasy) - ) - logging.info("-------") - - Sol0 = copy.deepcopy(Sol) - flux0 = copy.deepcopy(flux) - mpv0 = copy.deepcopy(mpv) - - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_before_flux") - - gd_flux.recompute_advective_fluxes(flux, Sol) - - if debug == True: - writer.populate(str(label) + "_before_advect", "rhoYu", flux[0].rhoY) - if debug == True: - writer.populate(str(label) + "_before_advect", "rhoYv", flux[1].rhoY) - if debug == True and elem.ndim == 3: - writer.populate(str(label) + "_before_advect", "rhoYw", flux[2].rhoY) - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_before_advect") - - # advect(Sol, flux, 0.5*dt, elem, step%2, ud, th, mpv, node, str(label)+'_half', writer) - if ud.do_advection: - gd_explicit.advect_rk( - Sol, - flux, - 0.5 * dt, - elem, - step % 2, - ud, - th, - mpv, - node, - str(label) + "_half", - writer, - ) - - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_advect") - - if debug: - writer.populate(str(label) + "_after_full_step", "p2_nodes0", mpv.p2_nodes) - - mpv.p2_nodes0[...] = mpv.p2_nodes - - lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_ebnaexp") - if ud.is_compressible == 0: - lm_sp.euler_backward_non_advective_impl_part( - Sol, - mpv, - elem, - node, - ud, - th, - t, - 0.5 * dt, - 1.0, - Sol0=Sol0, - label=str(label + "_after_ebnaimp"), - writer=writer, - ) - else: - lm_sp.euler_backward_non_advective_impl_part( - Sol, - mpv, - elem, - node, - ud, - th, - t, - 0.5 * dt, - 1.0, - label=str(label + "_after_ebnaimp"), - writer=writer, - ) - - if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: - # top rayleight damping - bdry.rayleigh_damping(Sol, mpv, ud, elem, node) - - # bottom rayleigh forcing - if hasattr(ud, "rayleigh_forcing"): - if ud.rayleigh_forcing: - - if ud.rayleigh_forcing_type == "file": - reader = io.read_input( - ud.rayleigh_forcing_fn, ud.rayleigh_forcing_path - ) - - Sol_half_new = copy.deepcopy(Sol) - mpv_half_new = copy.deepcopy(mpv) - - # misusing hydrostatic blending data containers - time_tag = "%.3d_after_full_step" % step - reader.get_data(Sol_half_new, mpv_half_new, time_tag, half=True) - - # assuming constant background state - up = Sol_half_new.rhou / Sol_half_new.rho - vp = Sol_half_new.rhov / Sol_half_new.rho - Yp = ( - Sol_half_new.rhoY / Sol_half_new.rho - - mpv.HydroState.Y0.reshape(1, -1) - ) - - pi = mpv_half_new.p2_nodes - - bdry.rayleigh_damping( - Sol, mpv, ud, elem, node, [up, vp, Yp, pi, t + 0.5 * dt] - ) - - elif ud.rayleigh_forcing_type == "func": - # boundary.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - s = 5.0e-3 + 1e-4 + 0e-5 - ud.rf_bot.eigenfunction((t + 0.5 * dt), s) - up, vp, Yp, pi = ud.rf_bot.dehatter(th) - - ud.rf_bot.eigenfunction((t + 0.5 * dt), s, grid="n") - _, _, _, pi_n = ud.rf_bot.dehatter(th, grid="n") - - bdry.rayleigh_damping( - Sol, mpv, ud, elem, node, [up, vp, Yp, pi_n, t + 0.5 * dt] - ) - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_ebnaimp") - - # if test_hydrob == True and writer is not None and step==0: - # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - - flux_half_new = copy.deepcopy(flux) - - gd_flux.recompute_advective_fluxes(flux, Sol) - - if debug == True: - writer.populate(str(label) + "_after_half_step", "rhoYu", flux[0].rhoY) - if debug == True: - writer.populate(str(label) + "_after_half_step", "rhoYv", flux[1].rhoY) - if debug == True and elem.ndim == 3: - writer.populate(str(label) + "_after_half_step", "rhoYw", flux[2].rhoY) - - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_half_step") - - Sol_half_new = copy.deepcopy(Sol) - mpv_half_new = copy.deepcopy(mpv) - rho_half = np.copy(Sol.rho) - rhou_half = np.copy(Sol.rhou) - rhov_half = np.copy(Sol.rhov) - rhow_half = np.copy(Sol.rhow) - rhoX_half = np.copy(Sol.rhoX) - rhoY_half = np.copy(Sol.rhoY) - # pwchi = np.copy(Sol.pwchi) - p2_nodes_half = np.copy(mpv.p2_nodes) - - # if test_hydrob == True and writer is not None and step==0: - # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - # print(dt*0.5) - - # takes care of the non-hydrostatic, compressible case - if ud.is_compressible == 1 and ud.is_nonhydrostatic == 1: - mpv.p2_nodes[...] = mpv.p2_nodes0 - Sol = copy.deepcopy(Sol0) - # takes care of the hydrostatic case - elif ud.is_nonhydrostatic == 0: - # if step == 1 : - # Sol.rho = (Sol.rho_half + Sol.rho) * 0.5 - # Sol.rhou = (Sol.rhou_half + Sol.rhou) * 0.5 - # Sol.rhov = (Sol.rhov_half + Sol.rhov) * 0.5 - # Sol.rhow = (Sol.rhow_half + Sol.rhow) * 0.5 - # Sol.rhoX = (Sol.rhoX_half + Sol.rhoX) * 0.5 - # Sol.rhoY = (Sol.rhoY_half + Sol.rhoY) * 0.5 - # v = (Sol.rhov_half / Sol.rho_half + Sol.rhov / Sol.rho) * 0.5 - # mpv.p2_nodes = (mpv.p2_nodes_half + mpv.p2_nodes) * 0.5 - # v = rhov_half / rho_half - # Sol = copy.deepcopy(Sol0) - # Sol.rhov = (Sol.rhov_half + rhov_half) * 0.5 - # Sol.rhov = Sol.rho * v - # mpv.p2_nodes[...] = mpv.p2_nodes0 - # None - # Sol = copy.deepcopy(Sol_tu) - # mpv = copy.deepcopy(mpv_tu) - # None - # else: - Sol = copy.deepcopy(Sol0) - mpv.p2_nodes[...] = mpv.p2_nodes0 - - # takes care of the pseudo-incompressible case - elif ud.is_compressible == 0: - Sol = copy.deepcopy(Sol0) - - # Sol.rhov0 = np.copy(Sol.rhov) - Sol.rho_half = rho_half - Sol.rhou_half = rhou_half - Sol.rhov_half = rhov_half - Sol.rhow_half = rhow_half - Sol.rhoX_half = rhoX_half - Sol.rhoY_half = rhoY_half - mpv.p2_nodes_half = p2_nodes_half - # Sol.pwchi = pwchi - - # Sol.rho_half = Sol.rho - # Sol.rhou_half = Sol.rhou - # Sol.rhov_half = Sol.rhov - # Sol.rhow_half = Sol.rhow - # Sol.rhoX_half = Sol.rhoX - # Sol.rhoY_half = Sol.rhoY - # mpv.p2_nodes_half = mpv.p2_nodes - Sol_half_old = copy.deepcopy(Sol_half_new) - flux_half_old = copy.deepcopy(flux_half_new) - mpv_half_old = copy.deepcopy(mpv_half_new) - - # mpv.p2_nodes[...] = ud.compressibility * mpv.p2_nodes0 + (1.0-ud.compressibility) * mpv.p2_nodes - - lm_sp.euler_forward_non_advective( - Sol, - mpv, - elem, - node, - 0.5 * dt, - ud, - th, - writer=writer, - label=str(label) + "_after_efna", - ) - - if debug == True: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_efna") - - if ud.do_advection: - gd_explicit.advect( - Sol, - flux, - dt, - elem, - step % 2, - ud, - th, - mpv, - node, - str(label) + "_full", - writer, - ) - # advect_rk(Sol, flux, dt, elem, step%2, ud, th, mpv, node, str(label)+'_full', writer) - - if debug == True: - writer.write_all( - Sol, mpv, elem, node, th, str(label) + "_after_full_advect" - ) - - lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) - - if debug == True: - writer.write_all( - Sol, mpv, elem, node, th, str(label) + "_after_full_ebnaexp" - ) - - lm_sp.euler_backward_non_advective_impl_part( - Sol, - mpv, - elem, - node, - ud, - th, - t, - 0.5 * dt, - 2.0, - writer=writer, - label=str(label) + "_after_full_step", - ) - - if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: - # top rayleight damping - bdry.rayleigh_damping(Sol, mpv, ud, elem, node) - - # bottom rayleigh forcing - if hasattr(ud, "rayleigh_forcing"): - if ud.rayleigh_forcing: - - if ud.rayleigh_forcing_type == "file": - - reader = io.read_input( - ud.rayleigh_forcing_fn, ud.rayleigh_forcing_path - ) - - # misusing hydrostatic blending data containers - time_tag = "%.3d_after_full_step" % step - reader.get_data(Sol_half_new, mpv_half_new, time_tag) - - # assuming constant background state - up = Sol_half_new.rhou / Sol_half_new.rho - vp = Sol_half_new.rhov / Sol_half_new.rho - Yp = ( - Sol_half_new.rhoY / Sol_half_new.rho - - mpv.HydroState.Y0.reshape(1, -1) - ) - # vp = 0.0 - # Yp = 0.0 - pi = mpv_half_new.p2_nodes - - bdry.rayleigh_damping( - Sol, mpv, ud, elem, node, [up, vp, Yp, pi, t + dt] - ) - - elif ud.rayleigh_forcing_type == "func": - - s = 5.0e-3 + 1e-4 + 0e-5 - ud.rf_bot.eigenfunction((t + dt), s) - up, vp, Yp, pi = ud.rf_bot.dehatter(th) - - ud.rf_bot.eigenfunction((t + dt), s, grid="n") - _, _, _, pi_n = ud.rf_bot.dehatter(th, grid="n") - - bdry.rayleigh_damping( - Sol, mpv, ud, elem, node, [up, vp, Yp, pi_n, t + dt] - ) - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - # if writer is not None: writer.populate(str(label)+'_after_full_step','p2_half',mpv.p2_nodes_half) - - ###################################################### - # Blending : Do blending after timestep - ###################################################### - Sol, mpv = da.blending.blending_after_timestep( - Sol, - flux, - mpv, - bld, - elem, - node, - th, - ud, - label, - writer, - step, - window_step, - t, - dt, - swe_to_lake, - debug, - ) - - if c1 or c2: - logging.info(termcolor.colored("hydrostatic to nonhydrostatic conversion...", "blue")) - - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_half_full") - writer.populate(str(label) + "_ic", "pwchi", Sol.pwchi) - - if test_hydrob == False: - Sol = copy.deepcopy(Sol_half_old) - # mpv = copy.deepcopy(mpv_half_old) - - logging.info(termcolor.colored("test_hydrob == False", "red")) - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_quarter") - - writer.populate(str(label) + "_quarter", "pwchi", Sol.pwchi) - - logging.info("quarter dt = %.8f" % (dt * 0.5)) - - ret = do( - Sol_half_old, - flux_half_old, - mpv_half_old, - dt - 0.5 * dt, - dt + 0.5 * dt, - ud, - elem, - node, - [0, 0], - th, - bld=None, - writer=None, - debug=False, - ) - - Sol_tu = copy.deepcopy(ret[0]) - # mpv_tu = copy.deepcopy(ret[2]) - Sol.rho[...] = Sol_tu.rho_half - Sol.rhou[...] = Sol_tu.rhou_half - Sol.rhov[...] = Sol_tu.rhov_half - Sol.rhow[...] = Sol_tu.rhow_half - Sol.rhoX[...] = Sol_tu.rhoX_half - Sol.rhoY[...] = Sol_tu.rhoY_half - Sol.pwchi[...] = Sol_tu.pwchi - - # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half - - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_half") - - writer.populate(str(label) + "_half", "pwchi", Sol.pwchi) - - ret = do( - Sol, - flux, - mpv, - dt, - 2.0 * dt, - ud, - elem, - node, - [0, 0], - th, - bld=None, - writer=None, - debug=False, - ) - - Sol = copy.deepcopy(ret[0]) - flux = copy.deepcopy(ret[1]) - mpv = copy.deepcopy(ret[2]) - - if test_hydrob == True: - Sol = copy.deepcopy(Sol_half_old) - # mpv = copy.deepcopy(mpv_half_old) - - logging.info(termcolor.colored("test_hydrob == False", "red")) - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_quarter") - - # writer.populate(str(label)+'_quarter', 'pwchi', Sol.pwchi) - - logging.info("quarter dt = %.8f" % (dt * 0.5)) - - ret = do( - Sol_half_old, - flux_half_old, - mpv_half_old, - dt - 0.5 * dt, - dt + 0.5 * dt, - ud, - elem, - node, - [0, 0], - th, - bld=None, - writer=None, - debug=False, - ) - - Sol_tu = copy.deepcopy(ret[0]) - # mpv_tu = copy.deepcopy(ret[2]) - Sol.rho[...] = Sol_tu.rho_half - Sol.rhou[...] = Sol_tu.rhou_half - Sol.rhov[...] = Sol_tu.rhov_half - Sol.rhow[...] = Sol_tu.rhow_half - Sol.rhoX[...] = Sol_tu.rhoX_half - Sol.rhoY[...] = Sol_tu.rhoY_half - Sol.pwchi[...] = Sol_tu.pwchi - - # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half - - # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - - # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) - - ret = do( - Sol, - flux, - mpv, - dt, - 2.0 * dt, - ud, - elem, - node, - [0, 0], - th, - bld=None, - writer=None, - debug=False, - ) - - Sol = copy.deepcopy(ret[0]) - flux = copy.deepcopy(ret[1]) - mpv = copy.deepcopy(ret[2]) - # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) - - logging.info(termcolor.colored("test_hydrob == True", "red")) - - if test_hydrob == False: - dt *= 2.0 - if c2: - ud.is_nonhydrostatic = 1 - - t += dt - - if writer != None: - writer.time = t - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_full_step") - # writer.populate(str(label)+'_after_full_step', 'pwchi', Sol.pwchi) - logging.info( - "###############################################################################################" - ) - logging.info( - "step %i done, t = %.12f, dt = %.12f, CFL = %.8f, CFL_ac = %.8f" - % (step, t, dt, cfl, cfl_ac) - ) - logging.info( - "###############################################################################################" - ) - - step += 1 - window_step += 1 - - return [Sol, flux, mpv, [window_step, step]] diff --git a/src/dycore/physics/gas_dynamics/numerical_flux.py b/src/dycore/physics/gas_dynamics/numerical_flux.py deleted file mode 100644 index 548e8ad0..00000000 --- a/src/dycore/physics/gas_dynamics/numerical_flux.py +++ /dev/null @@ -1,99 +0,0 @@ -# -*- coding: utf-8 -*- -import numpy as np -import scipy as sp - -def recompute_advective_fluxes(flux, Sol, *args, **kwargs): - """ - Recompute the advective fluxes at the cell interfaces, i.e. the faces. This function updates the `flux` container in-place. - - Parameters - ---------- - flux : :py:class:`management.variable.States` - Data container for the fluxes at the cell interfaces. - Sol : :py:class:`management.variable.States` - Data container for the Solution. - - Attention - --------- - This function is a mess and requires cleaning up. - - """ - ndim = Sol.rho.ndim - inner_idx = tuple([slice(1,-1)] * ndim) - - if ndim == 2: - kernel_u = np.array([[0.5,1.,0.5],[0.5,1.,0.5]]) - kernel_v = kernel_u.T - elif ndim == 3: - kernel_u = np.array([[[1,2,1],[2,4,2],[1,2,1]],[[1,2,1],[2,4,2],[1,2,1]]]) - kernel_v = np.swapaxes(kernel_u,1,0) - kernel_w = np.swapaxes(kernel_u,2,0) - - rhoYw = Sol.rhoY * Sol.rhow / Sol.rho - flux[2].rhoY[inner_idx] = sp.signal.fftconvolve(rhoYw, kernel_w, mode='valid') / kernel_w.sum() - else: - assert 0, "Unsupported dimension in recompute_advective_flux" - - rhoYu = kwargs.get('u',Sol.rhoY * Sol.rhou / Sol.rho) - - flux[0].rhoY[inner_idx] = np.moveaxis(sp.signal.fftconvolve(rhoYu, kernel_u, mode='valid') / kernel_u.sum(), 0, -1) - - rhoYv = kwargs.get('v',Sol.rhoY * Sol.rhov / Sol.rho) - if ndim == 2: - flux[1].rhoY[inner_idx] = sp.signal.fftconvolve(rhoYv, kernel_v, mode='valid') / kernel_v.sum() - elif ndim == 3: - flux[1].rhoY[inner_idx] = np.moveaxis(sp.signal.fftconvolve(rhoYv, kernel_v, mode='valid') / kernel_v.sum(), -1,0) - # flux[1].rhoY[...,-1] = 0. - -def hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th): - """ - HLL solver for the Riemann problem. Chooses the advected quantities from `Lefts` or `Rights` based on the direction given by `flux`. - - Parameters - ---------- - flux : :py:class:`management.variable.States` - Data container for fluxes. - Lefts : :py:class:`management.variable.States` - Container for the quantities on the left of the cell interfaces. - Rights : :py:class:`management.variable.States` - Container for the quantities on the right of the cell interfaces. - Sol : :py:class:`management.variable.Vars` - Solution data container. - lmbda : float - :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep. - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for the initial condition. - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for the thermodynamical constants. - - Returns - ------- - :py:class:`management.variable.States` - `flux` data container with the solution of the Riemann problem. - """ - # flux: index 1 to end = Left[inner_idx]: index 0 to -1 = Right[inner_idx]: index 1 to end - - ndim = Sol.rho.ndim - left_idx, right_idx, remove_cols_idx = [slice(None)] * ndim, [slice(None)] * ndim, [slice(None)] * ndim - - remove_cols_idx[-1] = slice(1,-1) - left_idx[-1] = slice(0,-1) - right_idx[-1] = slice(1,None) - - left_idx, right_idx, remove_cols_idx = tuple(left_idx), tuple(right_idx), tuple(remove_cols_idx) - - Lefts.primitives(th) - Rights.primitives(th) - - upwind = 0.5 * (1.0 + np.sign(flux.rhoY)) - upl = upwind[right_idx] - upr = (1.0 - upwind[left_idx]) - - flux.rhou[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (upl[left_idx] / Lefts.Y[left_idx] * Lefts.u[left_idx] + upr[right_idx] / Rights.Y[right_idx] * Rights.u[right_idx]) - flux.rho[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (upl[left_idx] / Lefts.Y[left_idx] * 1.0 + upr[right_idx] / Rights.Y[right_idx] * 1.0) - - flux.rhov[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (upl[left_idx] / Lefts.Y[left_idx] * Lefts.v[left_idx] + upr[right_idx] / Rights.Y[right_idx] * Rights.v[right_idx]) - flux.rhow[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (upl[left_idx] / Lefts.Y[left_idx] * Lefts.w[left_idx] + upr[right_idx] / Rights.Y[right_idx] * Rights.w[right_idx]) - flux.rhoX[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (upl[left_idx] / Lefts.Y[left_idx] * Lefts.X[left_idx] + upr[right_idx] / Rights.Y[right_idx] * Rights.X[right_idx]) - - return flux diff --git a/src/dycore/physics/hydrostatics.py b/src/dycore/physics/hydrostatics.py deleted file mode 100644 index 6bc5b538..00000000 --- a/src/dycore/physics/hydrostatics.py +++ /dev/null @@ -1,249 +0,0 @@ -import dycore.utils.boundary as bdry -import numpy as np -import numba as nb - -def column(HydroState, HydroState_n, Y, Y_n, elem, node, th, ud): - Gamma = th.gm1 / th.gamm - gamm = th.gamm - gm1 = th.gm1 - Gamma_inv = 1.0 / Gamma - gm1_inv = 1.0 / gm1 - - icy = elem.icy - igy = elem.igy - - xc_idx = slice(0,-1) - yc_idx = slice(0,-1) - - c_idx = (xc_idx,yc_idx) - - rhoY0 = 1.0 - - g = ud.gravity_strength[1] - - p0 = rhoY0**gamm - pi0 = rhoY0**gm1 - HydroState_n.rho0[xc_idx,igy] = rhoY0 / Y_n[:,igy] - HydroState_n.rhoY0[xc_idx,igy] = rhoY0 - HydroState_n.Y0[xc_idx,igy] = Y_n[:,igy] - HydroState_n.S0[xc_idx,igy] = 1.0 / Y_n[:,igy] - HydroState_n.p0[xc_idx,igy] = p0 - HydroState_n.p20[xc_idx,igy] = pi0 / ud.Msq - - dys = np.array([-elem.dy] + [-elem.dy/2] + [elem.dy/2] + list(np.ones((icy-3)) * elem.dy)) - S_p = 1.0 / Y[:,:] - S_m = np.zeros_like(S_p) - S_m[:,igy-1:igy+1] = 1.0 / Y_n[:,igy].reshape(-1,1) - S_m[:,0] = 1.0 / Y[:,igy-1] - S_m[:,igy+1:] = 1.0 / Y[:,igy:-1] - - S_integral_p = dys * 0.5 * (S_p + S_m) - S_integral_p[:,:igy] = np.cumsum(S_integral_p[:,:igy][:,::-1],axis=1)[:,::-1] - S_integral_p[:,igy:] = np.cumsum(S_integral_p[:,igy:],axis=1) - - pi_hydro = pi0 - Gamma * g * S_integral_p - p_hydro = pi_hydro**Gamma_inv - rhoY_hydro = pi_hydro**gm1_inv - - HydroState.rho0[c_idx] = rhoY_hydro * S_p - HydroState.p0[c_idx] = p_hydro - HydroState.p20[c_idx] = pi_hydro / ud.Msq - HydroState.S0[c_idx] = S_p - HydroState.S10[c_idx] = 0.0 - HydroState.Y0[c_idx] = 1.0 / S_p - HydroState.rhoY0[c_idx] = rhoY_hydro - - Sn_p = 1.0 / Y[:,:] - dys = np.ones((icy)) * elem.dy - dys[:igy] *= -1 - Sn_integral_p = dys * Sn_p - Sn_integral_p[:,:igy] = np.cumsum(Sn_integral_p[:,:igy][:,::-1],axis=1)[:,::-1] - Sn_integral_p[:,igy:] = np.cumsum(Sn_integral_p[:,igy:],axis=1) - - pi_hydro_n = pi0 - Gamma * g * Sn_integral_p - rhoY_hydro_n = pi_hydro_n**gm1_inv - - HydroState_n.rhoY0[xc_idx,:igy] = rhoY_hydro_n[:,:igy] - HydroState_n.Y0[xc_idx,:igy] = Y_n[0,:igy] - HydroState_n.S0[xc_idx,:igy] = 1.0 / Y_n[:,:igy] - HydroState_n.p0[xc_idx,:igy] = rhoY_hydro_n[:,:igy]**th.gamm - HydroState_n.p20[xc_idx,:igy] = pi_hydro_n[:,:igy] / ud.Msq - - HydroState_n.rhoY0[xc_idx,igy+1:] = rhoY_hydro_n[:,igy:] - HydroState_n.Y0[xc_idx,igy+1:] = Y_n[0,igy:] - HydroState_n.S0[xc_idx,igy+1:] = 1.0 / Y_n[:,igy:] - HydroState_n.p0[xc_idx,igy+1:] = rhoY_hydro_n[:,igy:]**th.gamm - HydroState_n.p20[xc_idx,igy+1:] = pi_hydro_n[:,igy:] / ud.Msq - - -def state(mpv, elem, node, th, ud): - g = ud.gravity_strength[1] - - if g != 0.0: - Gamma = th.Gamma - Hex = 1.0 / (th.Gamma * g) - dy = elem.dy - - pi_np = np.exp(-(node.y + 0.5 * dy) / Hex) - pi_nm = np.exp(-(node.y - 0.5 * dy) / Hex) - pi_n = np.exp(-(node.y) / Hex) - - Y_n = - Gamma * g * dy / (pi_np - pi_nm) - P_n = pi_n**th.gm1inv - p_n = pi_n**th.Gammainv - rho_n = P_n / Y_n - - mpv.HydroState_n.p20[...] = pi_n / ud.Msq - mpv.HydroState_n.p0[...] = p_n - mpv.HydroState_n.rho0[...] = rho_n - mpv.HydroState_n.rhoY0[...] = P_n - mpv.HydroState_n.Y0[...] = Y_n - mpv.HydroState_n.S0[...] = 1.0 / Y_n - - pi_cp = np.exp(-(elem.y + 0.5 * dy) / Hex) - pi_cm = np.exp(-(elem.y - 0.5 * dy) / Hex) - pi_c = np.exp(-(elem.y) / Hex) - - Y_c = - Gamma * g * dy / (pi_cp - pi_cm) - P_c = pi_c**th.gm1inv - p_c = pi_c**th.Gammainv - rho_c = P_c / Y_c - - mpv.HydroState.p20[...] = pi_c / ud.Msq - mpv.HydroState.p0[...] = p_c - mpv.HydroState.rho0[...] = rho_c - mpv.HydroState.rhoY0[...] = P_c - mpv.HydroState.Y0[...] = Y_c - mpv.HydroState.S0[...] = 1.0 / Y_c - - else: - mpv.HydroState.p20[...] = 1.0 - mpv.HydroState.p0[...] = 1.0 - mpv.HydroState.rho0[...] = 1.0 - mpv.HydroState.rhoY0[...] = 1.0 - mpv.HydroState.Y0[...] = 1.0 - mpv.HydroState.S0[...] = 1.0 - - mpv.HydroState_n.p20[...] = 1.0 - mpv.HydroState_n.p0[...] = 1.0 - mpv.HydroState_n.rho0[...] = 1.0 - mpv.HydroState_n.rhoY0[...] = 1.0 - mpv.HydroState_n.Y0[...] = 1.0 - mpv.HydroState_n.S0[...] = 1.0 - - -def initial_pressure(Sol,mpv,elem,node,ud,th): - Gammainv = th.Gammainv - igy = node.igy - igx = node.igx - icx= elem.icx - dx = node.dx - dy = node.dy - - beta = np.zeros((node.icx)) - bdpdx = np.zeros((node.icx)) - - x_idx_m = slice(0,-1) - x_idx_c = slice(1,None) - y_idx = slice(igy,-igy) - xn_idx = slice(1,-1) - height = node.y[-igy-1] - - Pc = Sol.rhoY[x_idx_c,y_idx] - Pm = Sol.rhoY[x_idx_m,y_idx] - thc = Pc / Sol.rho[x_idx_c,y_idx] - thm = Pm / Sol.rho[x_idx_m,y_idx] - beta[xn_idx] = np.sum(0.5 * (Pm * thm + Pc * thc) * dy, axis=1) - bdpdx[xn_idx] = np.sum(0.5 * (Pm * thm + Pc * thc) * (mpv.p2_cells[x_idx_c,y_idx] - mpv.p2_cells[x_idx_m,y_idx]) * dy, axis=1) - - beta *= Gammainv / height - bdpdx *= Gammainv / height / dx - - coeff = np.zeros((elem.icx)) - pibot = np.zeros((elem.icx)) - coeff[igx+1:-igx+1] = np.cumsum(coeff[igx:-igx] + dx / beta[igx+1:-igx]) - pibot[igx+1:-igx+1] = np.cumsum(pibot[igx:-igx] - dx * bdpdx[igx+1:-igx] / beta[igx+1:-igx]) - - dotPU = pibot[icx-igx] / coeff[icx-igx] - pibot[igx:-igx] -= dotPU * coeff[igx:-igx] - - x_idx = slice(igx,-igx+1) - y_idx = slice(igy,-igy+1) - - mpv.p2_cells[x_idx,y_idx] += pibot[x_idx].reshape(-1,1) - 1.0 * mpv.HydroState.p20[y_idx].reshape(1,-1) - bdry.set_ghostcells_p2(mpv.p2_cells, elem, ud) - - icxn = node.icx - icyn = node.icy - x_idx = slice(1,icxn-1) - y_idx = slice(igy,-igy+1) - height = node.y[-igy] - - Pc = Sol.rhoY[1:,y_idx] - thc = Pc / Sol.rho[1:,y_idx] - - beta = np.zeros((elem.icx,)) - bdpdx = np.zeros((elem.icx)) - - beta[1:] = np.sum(Pc*thc*dy, axis=1) - beta *= Gammainv / height - - bdpdx[1:] = np.sum(Pc * thc * (mpv.p2_nodes[1:-1,igy:-igy] - mpv.p2_nodes[:-2,igy:-igy]) * dy, axis=1) - bdpdx *= Gammainv / height / dx - - coeff = np.zeros((node.icx)) - pibot = np.zeros((node.icx)) - - coeff[igx+1:-igx+1] = np.cumsum(coeff[igx:-igx] + dx / beta[igx+1:]) - pibot[igx+1:-igx+1] = np.cumsum(pibot[igx:-igx] - dx * bdpdx[igx+1:] / beta[igx+1:]) - - dotPU = pibot[icx-igx] / coeff[icx-igx] - - pibot[igx:-igx] -= dotPU * coeff[igx:-igx] - - x_idx = slice(igx,-igx+1) - y_idx = slice(igy,-igy+1) - mpv.p2_nodes[x_idx,y_idx] += pibot[x_idx].reshape(-1,1) - 1.0 * mpv.HydroState_n.p20[y_idx].reshape(1,-1) - - mpv.dp2_nodes[:,:] = mpv.p2_nodes - - # guess initial node value (at left-most node) - mpv.p2_nodes[igx,igy:-igy] = mpv.dp2_nodes[igx,igy:-igy] - - mpv.p2_nodes[:,:] = __loop_over_array(igx,igy,icxn,icyn,mpv.p2_nodes, mpv.dp2_nodes) - - assert ((node.icx+1)%2) == 1 - delp2 = 0.5 * (mpv.p2_nodes[-igx-1,igy:-igy] - mpv.p2_nodes[igx,igy:-igy]) - delp2 = delp2.reshape(1,-1) - sgn = np.ones_like(mpv.p2_nodes[:,0][igy:-igy]).reshape(-1,1) - - sgn[1::2] *= -1 - - mpv.p2_nodes[igx:-igx,igy:-igy] += sgn * delp2 - bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) - - mpv.dp2_nodes[:,:] = 0.0 - - inner_domain = (slice(igx,-igx), slice(igy,-igy)) - pi = ud.Msq * (mpv.p2_cells[inner_domain] + 1.0 * mpv.HydroState.p20[igy:-igy]) - Y = Sol.rhoY[inner_domain] / Sol.rho[inner_domain] - rhoold = np.copy(Sol.rho[inner_domain]) - Sol.rhoY[inner_domain] = pi**th.gm1inv - Sol.rho[inner_domain] = Sol.rhoY[inner_domain] / Y - Sol.rhou[inner_domain] *= Sol.rho[inner_domain] / rhoold - Sol.rhov[inner_domain] *= Sol.rho[inner_domain] / rhoold - Sol.rhow[inner_domain] *= Sol.rho[inner_domain] / rhoold - Sol.rhoX[inner_domain] *= Sol.rho[inner_domain] / rhoold - -# need details: -# populate the rest of the nodes recursively based on the left-most node. -# recursive: use numba. -@nb.jit(nopython=True) -def __loop_over_array(igx,igy,icxn,icyn,p,dp): - for j in range(igy,icyn-igy): - for i in range(igx+1,icxn-igx): - p[i,j] = 2.0 * dp[i-1,j] - p[i-1,j] - return p - - - \ No newline at end of file diff --git a/src/dycore/physics/low_mach/mpv.py b/src/dycore/physics/low_mach/mpv.py deleted file mode 100644 index e3d3c069..00000000 --- a/src/dycore/physics/low_mach/mpv.py +++ /dev/null @@ -1,47 +0,0 @@ -import numpy as np -import dycore.utils.variable as var - -class MPV(object): - def __init__(self,elem,node,ud): - sc = elem.sc - sn = node.sc - - self.p0 = 1.0 - self.p00 = 1.0 - - self.p2_cells = np.zeros((sc)) - self.dp2_cells = np.zeros((sc)) - self.p2_nodes = np.zeros((sn)) - self.p2_nodes0 = np.zeros((sn)) - self.dp2_nodes = np.zeros((sn)) - - self.u = np.zeros((sc)) - self.v = np.zeros((sc)) - self.w = np.zeros((sc)) - - self.rhs = np.zeros((node.isc)) - self.wcenter = np.zeros((node.isc)) - self.wplus = np.zeros(([elem.ndim]+list(sc))) - - self.HydroState = var.States([sc[1]],ud) - self.HydroState_n = var.States([sn[1]],ud) - - self.squeezer() - - def squeezer(self): - for key, value in vars(self).items(): - if type(value) == np.ndarray: - setattr(self,key,value.squeeze()) - - -# def acoustic_order(ud,t,step): -# if ud.is_compressible == 0: -# return 2.0 -# elif ud.is_compressible == 1: -# return 2.0 -# elif ud.is_compressible == -1: -# current_transition_step = step - ud.no_of_pi_initial -# # return np.linspace(,2.0,ud.no_of_pi_transition+2)[1:-1][current_transition_step] -# return 2.0 -# else: -# assert 0 \ No newline at end of file diff --git a/src/pybella/__init__.py b/src/pybella/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/src/pybella/__init__.py @@ -0,0 +1 @@ + diff --git a/src/pybella/__main__.py b/src/pybella/__main__.py new file mode 100644 index 00000000..b27d2848 --- /dev/null +++ b/src/pybella/__main__.py @@ -0,0 +1,125 @@ +# some diagnostics +import time +import logging + +import numpy as np + +# dependencies of the atmospheric flow solver +from .flow_solver.discretisation import time_update as dis_time_update + +# dependencies of the interface subpackage +from .interfaces.dynamics_blending import prepare as blending_prepare +from .interfaces.postprocessing import strip_target_file as strip_target + + +# dependencies of the data assimilation subpackage +from .data_assimilation import prepare as da_prepare, analysis as da_analysis + +# package imports +from .utils import prepare, io, sim_params as params + + +########################################################## +# Start main looping +########################################################## +def main(): + sim_state = prepare.initialise() + blending_prepare.initialise(sim_state) + da_prepare.initialise(sim_state) + + if sim_state.restart: + prepare.overwrite_init_with_restart(sim_state) + + writer, step_writer = io.initialise(sim_state) + + tic = time.time() + + ###################################################### + # Time looping over data assimilation windows + ###################################################### + tout_old = -np.inf + tout_cnt = 0 + outer_step = 0 + for tout in sim_state.ud.tout: + sst = sim_state + es = sst.ensemble_state + dp = sst.da_params + + futures = [] + + # based on the initial blending parameter, define if we want to blend for this assimilation window or simulaiton run. + blend = blending_prepare.init_da_window(sim_state, tout_old, outer_step) + + ###################################################### + # Forecast step + ###################################################### + logging.info("##############################################") + logging.info("Next tout = %.3f" % tout) + logging.info("Starting forecast...") + for cnt, mem in enumerate(es): + # handling of DA window step counter + if sst.N > 1: + if tout_old in dp.dap.da_times: + mem.time.window_step = 0 + if sst.N == 1: + mem.time.window_step = mem.time.step + + debug_writer = io.create_debug_writer(params.debug, writer, mem) + + logging.info("For ensemble member = %i..." % cnt) + mem = dis_time_update.do( + sst, + mem, + tout, + blend, + step_writer, + debug_writer + ) + + if sst.ud.diag: + if sst.ud.diag_updt_targets: + strip_target.do(step_writer.OUTPUT_FILENAME + step_writer.BASE_NAME + step_writer.SUFFIX + '.h5', sst.ud, time_increment=sst.diag_comparison.time_increment) + sst.diag_comparison.update_targets() + else: + sst.diag_comparison.test_do( + mem, sst.ud + ) + + futures.append(mem) + + # Dask commands, used only when parallelisation is + # enabled + results = np.array(futures) + + da_analysis.do_for_window(tout, outer_step, results, sst, writer) + + ###################################################### + # Write output at tout + ###################################################### + logging.info("Starting output...") + for n, mem in enumerate(es): + if params.label_type == "STEP": + step = outer_step + label = "ensemble_mem=%i_%.3d" % (n, step) + else: + label = "ensemble_mem=%i_%.3f" % (n, tout) + writer.write_all(mem, str(label) + "_after_full_step") + + # synchronise_variables(mpv, Sol, elem, node, ud, th) + # sst.t = tout + tout_old = np.copy(tout) + logging.info("tout = %.3f" % tout) + + tout_cnt += 1 + outer_step += 1 + if outer_step > sst.ud.stepmax: + break + + toc = time.time() + logging.info("Time taken = %.6f" % (toc - tic)) + + writer.close_everything() + + +if __name__ == "__main__": + main() diff --git a/inputs/__init__.py b/src/pybella/data_assimilation/__init__.py similarity index 100% rename from inputs/__init__.py rename to src/pybella/data_assimilation/__init__.py diff --git a/src/pybella/data_assimilation/analysis.py b/src/pybella/data_assimilation/analysis.py new file mode 100644 index 00000000..69d2e7fd --- /dev/null +++ b/src/pybella/data_assimilation/analysis.py @@ -0,0 +1,128 @@ +import logging + +import numpy as np + +from . import ( + etpf as da_etpf, + post_processing as da_post_processing, + letkf as da_letkf, + utils as da_utils, +) + +from ..flow_solver.utils import boundary as bdry + +from ..utils import sim_params as params + + +def do_for_window(tout, outer_step, results, sst, writer): + # mp = sst.model_params + dp = sst.da_params + # ens = dp.sol_ens + + ###################################################### + # Analysis step + ###################################################### + tout = np.around(tout, 3) + if sst.N > 1 and tout in dp.dap.da_times: + futures = [] + + ###################################################### + # Update ensemble with forecast + ###################################################### + for mem in results: + elem, node, sol, _, mpv, th, _ = mem + bdry.set_explicit_boundary_data(sol, elem, sst.ud, th, mpv) + bdry.set_ghostnodes_p2(mpv.p2_nodes, node, sst.ud) + + # ens.set_members(results, tout) + sst.ensemble_state.set_members(results) + + ###################################################### + # Write output before assimilating data + ###################################################### + logging.info("Starting output...") + for mem in sst.ensemble_state: + elem, node, sol, _, mpv, th, _ = mem + if params.label_type == "STEP": + step = outer_step + label = "ensemble_mem=%i_%.3d" % (n, step) + else: + label = "ensemble_mem=%i_%.3f" % (n, tout) + writer.write_all(sol, mpv, elem, node, th, str(label) + "_before_da") + + ################################################## + # LETKF with batch observations + ################################################## + if dp.dap.da_type == "batch_obs": + logging.info("Starting analysis... for batch observations") + for attr in dp.dap.obs_attributes: + logging.info("Assimilating %s..." % attr) + logging.info("Assimilating %s..." % attr) + # future = client.submit(da_interface, *[s_res,obs_current,dap.inflation_factor,attr,N,ud,dap.loc[attr]]) + future = da_letkf.da_interface( + results, dp.dap, dp.obs, attr, tout, sst.N, sst.ud + ) + futures.append(future) + + # analysis = client.gather(futures) + analysis = futures + # analysis = np.array(analysis) + + logging.info("Writing analysis...") + cnt = 0 + for attr in dp.dap.obs_attributes: + current = analysis[cnt] + for n in range(sst.N): + setattr(results[:, dp.dap.loc[attr], ...][n], attr, current[n]) + cnt += 1 + + ################################################## + # LETKF with grid-point localisation + ################################################## + elif dp.dap.da_type == "rloc": + logging.info("Starting analysis... for rloc algorithm") + elem, node = sst.ensemble_state.get_grid() + results = da_utils.HSprojector_3t2D(results, elem, dp.dap, sst.N) + results = dp.rloc.analyse( + results, dp.obs, dp.obs_covar, dp.obs_mask, sst.N, tout + ) + results = da_utils.HSprojector_2t3D(results, elem, node, dp.dap, sst.N) + # if hasattr(dap, 'converter'): + # results = dap.converter(results, N, mpv, elem, node, th, ud) + + ################################################## + # ETPF + ################################################## + elif dp.dap.da_type == "etpf": + da_utils.ensemble_inflation( + results, dp.dap.attributes, dp.dap.inflation_factor, sst.N + ) + results = da_etpf.da_interface( + results, + dp.obs, + dp.dap.obs_attributes, + dp.dap.rejuvenation_factor, + dp.dap.da_times, + tout, + sst.N, + ) + + ################################################## + # Post-processing + ################################################## + elif dp.dap.da_type == "pprocess": + results = da_post_processing.interface() + + else: + assert 0, "DA type not implemented: use 'rloc', 'batch_obs' or 'etpf'." + + ###################################################### + # Update ensemble with analysis + ###################################################### + for mem in results: + elem, node, Sol, _, mpv, th, _ = mem + bdry.set_explicit_boundary_data(Sol, elem, sst.ud, th, mpv) + p2_nodes = mpv.p2_nodes + bdry.set_ghostnodes_p2(p2_nodes, node, sst.ud) + + sst.ensemble_state.set_members(results) diff --git a/src/pybella/data_assimilation/etpf.py b/src/pybella/data_assimilation/etpf.py new file mode 100644 index 00000000..9a058350 --- /dev/null +++ b/src/pybella/data_assimilation/etpf.py @@ -0,0 +1,127 @@ +import numpy as np +import logging + + +def da_interface(results, obs, obs_attributes, delta, times, tout, N, loc=0): + ig = 2 + inner = (slice(ig, -ig), slice(ig, -ig)) + + # local_ens = np.array([getattr(results[:,loc,...][n],attr) for n in range(N)]) + # local_ens = np.array([mem[inner] for mem in local_ens]) + # local_ens = analysis(local_ens,attr) + # print(results[:,0,...][0].rho) + attributes = ["rho", "rhou", "rhov", "rhow", "rhoY", "rhoX"] + # attributes = ['rho', 'rhou', 'rhov'] + tmp = np.array( + [getattr(results[:, loc, ...][n], obs_attributes[0])[inner] for n in range(N)] + ) + tmp = tmp[:, np.newaxis, ...] + + ensemble = np.array( + [getattr(results[:, loc, ...][n], attributes[0])[inner] for n in range(N)] + ) + ensemble = ensemble[:, np.newaxis, ...] + + obs_current = np.array( + obs[np.where(np.isclose(times, tout))[0][0]][obs_attributes[0]] + )[inner] + + # obs_current = bin_func(obs_current,(Nx,Ny)) + obs_current = obs_current[np.newaxis, ...] + + for attr in obs_attributes[1:]: + tmp = np.hstack( + ( + tmp, + np.array( + [getattr(results[:, loc, ...][n], attr)[inner] for n in range(N)] + )[:, np.newaxis, ...], + ) + ) + tmp01 = np.array(obs[np.where(np.isclose(times, tout))[0][0]][attr])[inner] + # tmp01 = bin_func(tmp01,(Nx,Ny)) + # print(tmp01.shape) + tmp01 = tmp01[np.newaxis, ...] + obs_current = np.vstack((obs_current, tmp01)) + + for attr in attributes[1:]: + ensemble = np.hstack( + ( + ensemble, + np.array( + [getattr(results[:, loc, ...][n], attr)[inner] for n in range(N)] + )[:, np.newaxis, ...], + ) + ) + + obs_current = obs_current[np.newaxis, ...] + # r = tmp - obs_current + + Hx = tmp.reshape(N, -1) + obs_current = obs_current.reshape(-1) + # ensemble = ensemble.reshape(N,-1) + # print(ensemble.shape) + + etpf = analysis(ensemble, delta) + etpf.analyse(obs_current, 1.0, Hx, N) + + analysis_ens = etpf.get_ensemble_from_X() + + for n in range(N): + cnt = 0 + current = analysis_ens[n] + for attr in attributes: + data = current[cnt] + data = np.pad(data, 2, mode="wrap") + + setattr(results[:, loc, ...][n], attr, data) + cnt += 1 + return results + + +class analysis(object): + def __init__(self, ensemble, delta, identifier=None): + self.ensemble = np.array(ensemble) + self.ensemble_shape = self.ensemble.shape + + self.X = self.state_vector(ensemble) + self.identifier = identifier + + # rejuvenation factor + self.delta = delta + + def analyse(self, obs_current, obs_covar, Hx, N): + import pyemd + + logging.info("starting ETPF analysis...") + + r = (Hx - obs_current) ** 2 + r = np.sum(r, axis=1) + # print(r.shape) + + ww = np.exp(-r / (2.0 * obs_covar)) + ww /= np.sum(ww) + + # print(ww) + + Co = self.X @ self.X.T + diag = np.diag(Co) + Co = diag * np.ones((1, N)) - 2.0 * Co + np.ones((N, 1)) * diag.T + + # print(Co) + + _, T = pyemd.emd_with_flow(ww, np.ones(N) / N, Co, -1) + T = np.array(T) + T = T * N + + self.X = np.dot( + self.X.T, T + ).T # + self.delta * np.random.randn(self.X.shape[0],self.X.shape[1]) + # print(self.X.shape) + + @staticmethod + def state_vector(ensemble): + return ensemble.reshape(ensemble.shape[0], -1) + + def get_ensemble_from_X(self): + return self.X.reshape(self.ensemble_shape) diff --git a/src/data_assimilation/letkf.py b/src/pybella/data_assimilation/letkf.py similarity index 64% rename from src/data_assimilation/letkf.py rename to src/pybella/data_assimilation/letkf.py index 7db1063e..07ef0fa4 100644 --- a/src/data_assimilation/letkf.py +++ b/src/pybella/data_assimilation/letkf.py @@ -1,3 +1,5 @@ +import logging + import numpy as np import numpy.lib.stride_tricks as st @@ -8,14 +10,14 @@ import dask import matplotlib.pyplot as plt -import logging -import dycore.utils.options as opts -import data_assimilation as da +from ..flow_solver.utils import options as opts +from . import utils debug_cnt = 0 -def da_interface(results,dap,obs,attr,tout,N,ud): + +def da_interface(results, dap, obs, attr, tout, N, ud): """ Interface for batch-observations localisation with LETKF. @@ -25,24 +27,31 @@ def da_interface(results,dap,obs,attr,tout,N,ud): inf_fac = dap.inflation_factor loc = dap.loc[attr] - inner = (slice(None,),slice(None,)) + inner = ( + slice( + None, + ), + slice( + None, + ), + ) - local_ens = np.array([getattr(results[:,loc,...][n],attr) for n in range(N)]) + local_ens = np.array([getattr(results[:, loc, ...][n], attr) for n in range(N)]) local_ens = np.array([mem[inner] for mem in local_ens]) - local_ens = analysis(local_ens,inf_fac,attr) + local_ens = analysis(local_ens, inf_fac, attr) obs_current = obs_current[inner] - + x_obs, y_obs = obs_current.shape - - forward_operator = lambda ensemble : ensemble # state space == observation space + + forward_operator = lambda ensemble: ensemble # state space == observation space # forward_operator = dap.forward_operator # localisation_matrix = dap.localisation_matrix # localisation_matrix = np.ones_like(obs_current.flatten()) local_ens.forward(forward_operator) # local_ens.localisation(localisation_matrix) - obs_covar = sp.sparse.eye((x_obs*y_obs), (x_obs*y_obs), format='csr') + obs_covar = sp.sparse.eye((x_obs * y_obs), (x_obs * y_obs), format="csr") X = local_ens.analyse(obs_current.reshape(-1), obs_covar) local_ens.ensemble = local_ens.to_array(X) @@ -55,12 +64,12 @@ def da_interface(results,dap,obs,attr,tout,N,ud): # let me ust put the forward operator here for now - will need to tidy stuff up.... -def interpolation_func(ensemble,x_obs,y_obs,ud): +def interpolation_func(ensemble, x_obs, y_obs, ud): if ud.bdry_type[0] == opts.BdryType.WALL or ud.bdry_type[1] == opts.BdryType.WALL: - assert("WALL NOT IMPLEMENTED!") + assert "WALL NOT IMPLEMENTED!" x_ens, y_ens = ensemble[0].shape - x = np.linspace(0,x_ens,x_obs) - y = np.linspace(0,y_ens,y_obs) + x = np.linspace(0, x_ens, x_obs) + y = np.linspace(0, y_ens, y_obs) # x = np.linspace(-0.5,0.5,x_ens) # y = np.linspace(-0.5,0.5,y_ens) @@ -72,28 +81,34 @@ def interpolation_func(ensemble,x_obs,y_obs,ud): # ensemble = [interpn((x,y),mem,pts, method='splinef2d').reshape(x_obs,y_obs) for mem in ensemble] - x,y = np.meshgrid(x,y) - ensemble = [ndimage.map_coordinates(mem,[y,x],mode='wrap', order=3) for mem in ensemble] + x, y = np.meshgrid(x, y) + ensemble = [ + ndimage.map_coordinates(mem, [y, x], mode="wrap", order=3) for mem in ensemble + ] return np.array(ensemble) -def bin_func(obs,ens_mem_shape): - obs = obs.reshape(ens_mem_shape[0],obs.shape[0]//ens_mem_shape[0], - ens_mem_shape[1],obs.shape[1]//ens_mem_shape[1]) - return obs.mean(axis=(1,3)) +def bin_func(obs, ens_mem_shape): + obs = obs.reshape( + ens_mem_shape[0], + obs.shape[0] // ens_mem_shape[0], + ens_mem_shape[1], + obs.shape[1] // ens_mem_shape[1], + ) + return obs.mean(axis=(1, 3)) def none_func(ensemble): - return lambda ensemble : ensemble + return lambda ensemble: ensemble class analysis(object): """ - LETKF Analysis based on (Hunt et al., 2007). + LETKF Analysis based on (Hunt et al., 2007). """ - def __init__(self,ensemble, rho, X_mean, Y_mean, identifier=None): + def __init__(self, ensemble, rho, X_mean, Y_mean, identifier=None): self.ensemble = np.array(ensemble) self.X = self.state_vector(ensemble) self.no_of_members = self.ensemble.shape[0] @@ -111,13 +126,13 @@ def __init__(self,ensemble, rho, X_mean, Y_mean, identifier=None): self.X_mean = X_mean self.Y_mean = Y_mean - def forward(self,forward_operator): + def forward(self, forward_operator): self.forward_operator = forward_operator - def localisation(self,localisation_matrix): + def localisation(self, localisation_matrix): self.localisation_matrix = localisation_matrix - def analyse(self,obs,obs_covar): + def analyse(self, obs, obs_covar): """ Analysis method. 'l' is the observation space. 'm' the state space, 'k' the ensemble size. @@ -143,48 +158,57 @@ def analyse(self,obs,obs_covar): # self.X -= self.X_mean # R in (m x k) # This is step 4 of the algorithm in Hunt et. al., 2007. - C = sp.spsolve(obs_covar, self.Y.T).T # R in (k x l) + C = sp.spsolve(obs_covar, self.Y.T).T # R in (k x l) # This applies the localisation function to the local region. if self.localisation_matrix is not None: C[...] = ((np.array(self.localisation_matrix)) @ C.T).T # The next three lines are step 5 of the algorithm. - Pa = (self.no_of_members - 1.) * np.eye(self.no_of_members) / self.rho + (C @ self.Y.T) + Pa = (self.no_of_members - 1.0) * np.eye(self.no_of_members) / self.rho + ( + C @ self.Y.T + ) Lambda, P = sp.linalg.eigh(Pa) - Pa = P @ (np.diag(1./Lambda) @ P.T) + Pa = P @ (np.diag(1.0 / Lambda) @ P.T) # This is step 6 of the algorithm. - Wa = np.sqrt(self.no_of_members - 1.) * P @ (np.diag(np.sqrt(1./Lambda)) @ P.T) + Wa = ( + np.sqrt(self.no_of_members - 1.0) + * P + @ (np.diag(np.sqrt(1.0 / Lambda)) @ P.T) + ) # The following two lines are step 7 of the algorithm. wa = Pa @ (C @ (obs - self.Y_mean)) Wa += wa # This is step 8 of the algorithm. - return (np.dot(self.X.T , Wa.T) + self.X_mean.reshape(-1,1)).T + return (np.dot(self.X.T, Wa.T) + self.X_mean.reshape(-1, 1)).T - - def get_mean(self,vec): + def get_mean(self, vec): return np.array(vec).mean(axis=0) - + # More readable method needed - seems to be most efficient though. @staticmethod def state_vector(ensemble): return np.array([mem.reshape(-1) for mem in ensemble]) - def to_array(self,X): + def to_array(self, X): return np.array([x.reshape(self.member_shape) for x in X]) def debug(self, obs_current, suffix=""): global debug_cnt N = self.ensemble.shape[0] ens_mean = self.ensemble.mean(axis=0) - _, ax = plt.subplots(ncols=N+1, figsize=(15,4)) + _, ax = plt.subplots(ncols=N + 1, figsize=(15, 4)) for n in range(N): ax[n].imshow(self.ensemble[n]) - ax[n+1].imshow(obs_current) - plt.savefig("./output_images/da_debug_%s_%i_%s" %(self.identifier,debug_cnt // 3,suffix), bbox_inches='tight') + ax[n + 1].imshow(obs_current) + plt.savefig( + "./output_images/da_debug_%s_%i_%s" + % (self.identifier, debug_cnt // 3, suffix), + bbox_inches="tight", + ) plt.close() debug_cnt += 1 @@ -192,8 +216,9 @@ def debug(self, obs_current, suffix=""): class prepare_rloc(object): """ Helper class to get grid-based localisation for LETKF. Used only when da_type=='rloc' is True. - + """ + def __init__(self, ud, elem, node, dap, N, obs_X=5, obs_Y=5): # get grid properties self.iicx = elem.iicx @@ -205,8 +230,15 @@ def __init__(self, ud, elem, node, dap, N, obs_X=5, obs_Y=5): self.N = N # define inner and outer domains - self.i2 = (slice(elem.igx,-elem.igx),slice(elem.igy,-elem.igy)) - self.i0 = (slice(None,),slice(None,)) + self.i2 = (slice(elem.igx, -elem.igx), slice(elem.igy, -elem.igy)) + self.i0 = ( + slice( + None, + ), + slice( + None, + ), + ) # get da parameters self.attr_len = len(dap.obs_attributes) @@ -230,9 +262,9 @@ def __init__(self, ud, elem, node, dap, N, obs_X=5, obs_Y=5): # make sure that subdomain is not larger than inner cellular grid assert self.obs_X < elem.iicx, "obs_X > iicx" - if elem.iicy == 1: # implying horizontal slice + if elem.iicy == 1: # implying horizontal slice assert self.obs_Y < elem.iicz, "obs_Y > iicz" - else: + else: assert self.obs_Y < elem.iicy, "obs_Y > iicy" # make sure that subdomain size is odd @@ -244,12 +276,13 @@ def __init__(self, ud, elem, node, dap, N, obs_X=5, obs_Y=5): self.pad_Y = int((self.obs_Y - 1) / 2) # get mask to handle BC. Periodic mask includes ghost cells/nodes and wall masks takes only the inner domain. - self.cmask, self.nmask = da.utils.boundary_mask(ud, elem, node, self.pad_X, self.pad_Y) + self.cmask, self.nmask = utils.boundary_mask( + ud, elem, node, self.pad_X, self.pad_Y + ) # get from da parameters the localisation matrix and inflation factor self.inf_fac = dap.inflation_factor self.loc_mat = dap.localisation_matrix - def sort_locs(self): """ @@ -261,28 +294,44 @@ def sort_locs(self): """ - cell_attributes = [key for key,value in self.loc.items() if value == self.loc_c and key in self.oa] - node_attributes = [key for key,value in self.loc.items() if value == self.loc_n and key in self.oa] - face_attributes = [key for key,value in self.loc.items() if value == self.loc_f and key in self.oa] + cell_attributes = [ + key + for key, value in self.loc.items() + if value == self.loc_c and key in self.oa + ] + node_attributes = [ + key + for key, value in self.loc.items() + if value == self.loc_n and key in self.oa + ] + face_attributes = [ + key + for key, value in self.loc.items() + if value == self.loc_f and key in self.oa + ] if len(face_attributes) > 0: assert 0, "r-localisation for values on faces not supported." return cell_attributes, node_attributes - def analyse(self,results,obs,covar,mask,N,tout): + def analyse(self, results, obs, covar, mask, N, tout): """ Wrapper to do analysis by grid-type. """ if self.cattr_len > 0: - results = self.analyse_by_grid_type(results,obs,covar[0],mask,N,tout,'cell') + results = self.analyse_by_grid_type( + results, obs, covar[0], mask, N, tout, "cell" + ) if self.nattr_len > 0: - results = self.analyse_by_grid_type(results,obs,covar[1],mask,N,tout,'node') + results = self.analyse_by_grid_type( + results, obs, covar[1], mask, N, tout, "node" + ) return results - def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): + def analyse_by_grid_type(self, results, obs, covar, mask, N, tout, grid_type): """ Do analysis by grid-type. Wrapper of the LETKF analysis. @@ -290,7 +339,7 @@ def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): """ - logging.info("Analysis grid type = %s" %grid_type) + logging.info("Analysis grid type = %s" % grid_type) # get properties from class attributes Nx, Ny, obs_attr, attr_len, loc = self.get_properties(grid_type) @@ -299,35 +348,44 @@ def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): # get stacked state space and observation # this prepares X, and Y_obs - state_p, obs_p, mask_p = self.stack(results,obs,mask,obs_attr,tout) + state_p, obs_p, mask_p = self.stack(results, obs, mask, obs_attr, tout) # get boundary mask # this does the BC handling - bc_mask,mask_n = self.get_bc_mask(mask_p, grid_type, Nx, Ny, obs_X, obs_Y,attr_len) - - # Here, the 2D arrays are split into local counterparts, say of (5x5) arrays. - X, X_bar = self.get_state(state_p,mask_p,Nx,Ny,attr_len) # X in R in ((Nx * Ny) x k x m) and X_bar in ((Nx * Ny) * m) - obs_p = self.get_obs(obs_p,mask_p,obs_X,obs_Y,Nx,Ny,attr_len) # R in ((Nx * Ny) x l) - Y, Y_bar = self.get_state_in_obs_space(results,mask_p,obs_attr,obs_X,obs_Y,Nx,Ny,attr_len) # Y in R in ((Nx * Ny) x k x l) and Y_bar in ((Nx * Ny) * l) + bc_mask, mask_n = self.get_bc_mask( + mask_p, grid_type, Nx, Ny, obs_X, obs_Y, attr_len + ) + + # Here, the 2D arrays are split into local counterparts, say of (5x5) arrays. + X, X_bar = self.get_state( + state_p, mask_p, Nx, Ny, attr_len + ) # X in R in ((Nx * Ny) x k x m) and X_bar in ((Nx * Ny) * m) + obs_p = self.get_obs( + obs_p, mask_p, obs_X, obs_Y, Nx, Ny, attr_len + ) # R in ((Nx * Ny) x l) + Y, Y_bar = self.get_state_in_obs_space( + results, mask_p, obs_attr, obs_X, obs_Y, Nx, Ny, attr_len + ) # Y in R in ((Nx * Ny) x k x l) and Y_bar in ((Nx * Ny) * l) analysis_res = np.zeros_like(X) # covariance handling if covar is None: # use identity for observation covariance - obs_covar_current = sp.sparse.eye(attr_len*obs_X*obs_Y,attr_len*obs_X*obs_Y, format='csr') + obs_covar_current = sp.sparse.eye( + attr_len * obs_X * obs_Y, attr_len * obs_X * obs_Y, format="csr" + ) else: # get covar at current time covar = covar[list(self.da_times).index(tout)] - + # expand obs covar to size of the subdomain - covar = np.expand_dims(covar,axis=-1) + covar = np.expand_dims(covar, axis=-1) covar = np.repeat(covar, obs_X, axis=-1) - covar = np.expand_dims(covar,axis=-1) + covar = np.expand_dims(covar, axis=-1) covar = np.repeat(covar, obs_Y, axis=-1) - # Note: I implemented the easiest and laziest way to make the LETKF work with large localisation regions. # However, this link: # https://github.com/dask/dask/issues/7589 @@ -335,13 +393,15 @@ def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): # Calculate chunk size such that largest chunk takes 180mb of memory: mem_size = 180 - mb=Y.nbytes/1024/1024 # Y is our largest array - chunks = int(np.ceil(mb/mem_size)) - chunk_size = int(np.ceil((Nx*Ny) / chunks)) + mb = Y.nbytes / 1024 / 1024 # Y is our largest array + chunks = int(np.ceil(mb / mem_size)) + chunk_size = int(np.ceil((Nx * Ny) / chunks)) logging.info("\n===================") - logging.info("To split DA problem into %i chunks at %s mb each" %(chunks,mem_size)) - logging.info("with analysis of %i grid points per chunk" %chunk_size) + logging.info( + "To split DA problem into %i chunks at %s mb each" % (chunks, mem_size) + ) + logging.info("with analysis of %i grid points per chunk" % chunk_size) logging.info("") # Now we chunkify our dask arrays: @@ -356,7 +416,23 @@ def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): # Now for each chunk, we get the analysis: analysis_by_chunk = [] for chunk in range(chunks): - analysis_by_chunk.append(darr.from_delayed(self.do_analysis(attr_len, covar, bc_mask[chunk], mask_n[chunk], obs_p[chunk], X[chunk], X_bar[chunk], Y[chunk], Y_bar[chunk]), shape=(attr_len,self.N,chunk_size), dtype=np.int64)) + analysis_by_chunk.append( + darr.from_delayed( + self.do_analysis( + attr_len, + covar, + bc_mask[chunk], + mask_n[chunk], + obs_p[chunk], + X[chunk], + X_bar[chunk], + Y[chunk], + Y_bar[chunk], + ), + shape=(attr_len, self.N, chunk_size), + dtype=np.int64, + ) + ) # Put the results of the chunks together analysis_res = darr.concatenate(analysis_by_chunk, axis=2) @@ -375,16 +451,15 @@ def analyse_by_grid_type(self,results,obs,covar,mask,N,tout,grid_type): for n in range(N): data = current[n] data = data.reshape(Nx, Ny) - - data = np.pad(data,2,mode='constant') - setattr(results[:,loc,...][n],attr,data) + data = np.pad(data, 2, mode="constant") - return results + setattr(results[:, loc, ...][n], attr, data) + return results # loop through all grid-points, selecting either the grid-point or its corresponding local region - # This is step 3 of the LETKF algorithm in Hunt et. al. 2007. + # This is step 3 of the LETKF algorithm in Hunt et. al. 2007. @dask.delayed def do_analysis(self, attr_len, covar, bc_mask, mask_n, obs_p, X, X_bar, Y, Y_bar): analysis_res = np.zeros_like(X) @@ -394,71 +469,74 @@ def do_analysis(self, attr_len, covar, bc_mask, mask_n, obs_p, X, X_bar, Y, Y_ba # For each of the quantities in the local observation space, Y, Y_bar, Y_obs, covar and loc_mat, remove the NaNs corresponding to grid-points without observations. # using forward operator as a projection of the state into observation space. - forward_operator = lambda ensemble : np.array([mem[~np.isnan(mem)] for mem in Y[n]]) + forward_operator = lambda ensemble: np.array( + [mem[~np.isnan(mem)] for mem in Y[n]] + ) Y_mean = Y_bar[n][~np.isnan(Y_bar[n])] # setup LETKF class with local state vector - local_ens = analysis(X[n],self.inf_fac,X_bar[n],Y_mean) + local_ens = analysis(X[n], self.inf_fac, X_bar[n], Y_mean) # setup forward operator method in LETKF class local_ens.forward(forward_operator) - + # get the localisation matrix with BC handling # BC is handled by localisation matrix. Observations on wall ghost cells have zero influence. - loc_mat = self.get_loc_mat(bc_mask,mask_n,n,attr_len) + loc_mat = self.get_loc_mat(bc_mask, mask_n, n, attr_len) local_ens.localisation_matrix = loc_mat # get masked covariance in local domain - covar_current = np.ma.array(covar,mask=mask_n[n]).filled(fill_value=np.nan) + covar_current = np.ma.array(covar, mask=mask_n[n]).filled(fill_value=np.nan) covar_current = covar_current[~np.isnan(covar_current)] - obs_covar_current = sp.sparse.diags(covar_current.flatten(), format='csr') + obs_covar_current = sp.sparse.diags(covar_current.flatten(), format="csr") # get obs according to sparsity structure obs_pn = obs_p[n][~np.isnan(obs_p[n])] # do analysis given observations and obs covar. - analysis_ens = local_ens.analyse(obs_pn,obs_covar_current) + analysis_ens = local_ens.analyse(obs_pn, obs_covar_current) # This is step 9 of the algorithm, where the grid-points are re-assembled. analysis_res[n] = analysis_ens - analysis_res = np.swapaxes(analysis_res,0,2) + analysis_res = np.swapaxes(analysis_res, 0, 2) return analysis_res - def stack(self,results,obs,mask,obs_attr,tout): + def stack(self, results, obs, mask, obs_attr, tout): """ On each grid-point, stack all the quantities to be assimilated. This stacking is done separately for cells and nodes. Quantities, e.g. {rho, rhou, ...}. """ - state = self.get_quantity(results,obs_attr[0]) - state = state[:,np.newaxis,...] + state = self.get_quantity(results, obs_attr[0]) + state = state[:, np.newaxis, ...] obs_stack = np.array(obs[list(self.da_times).index(tout)][obs_attr[0]])[self.i2] - obs_stack = obs_stack[np.newaxis,...] + obs_stack = obs_stack[np.newaxis, ...] - mask_stack = np.array(mask[list(self.da_times).index(tout)][obs_attr[0]])[self.i2] - mask_stack = mask_stack[np.newaxis,...] + mask_stack = np.array(mask[list(self.da_times).index(tout)][obs_attr[0]])[ + self.i2 + ] + mask_stack = mask_stack[np.newaxis, ...] for attr in obs_attr[1:]: - next_attr = self.get_quantity(results,attr)[:,np.newaxis,...] - state = np.hstack((state,next_attr)) + next_attr = self.get_quantity(results, attr)[:, np.newaxis, ...] + state = np.hstack((state, next_attr)) next_obs = np.array(obs[list(self.da_times).index(tout)][attr])[self.i2] - next_obs = next_obs[np.newaxis,...] + next_obs = next_obs[np.newaxis, ...] - obs_stack = np.vstack((obs_stack,next_obs)) + obs_stack = np.vstack((obs_stack, next_obs)) next_mask = np.array(mask[list(self.da_times).index(tout)][attr])[self.i2] - next_mask = next_mask[np.newaxis,...] + next_mask = next_mask[np.newaxis, ...] - mask_stack = np.vstack((mask_stack,next_mask)) + mask_stack = np.vstack((mask_stack, next_mask)) return state, obs_stack, mask_stack - - def get_state(self,state,mask,Nx,Ny,attr_len): + def get_state(self, state, mask, Nx, Ny, attr_len): """ Get state vector in grid-localised view. This is step 1 of the LETKF algorithm in Hunt et. al. (2007). Prepares global X and \bar{X}. @@ -472,21 +550,25 @@ def get_state(self,state,mask,Nx,Ny,attr_len): # Decompose X[G] and bar{X[G]} into the local regions view # X = np.array([dau.st.sliding_window_view(mem, (1,1), (1,1)).reshape(Nx*Ny,attr_len) for mem in state]) - # X_bar = np.array([dau.st.sliding_window_view(mem, (1,1), (1,1)).reshape(Nx*Ny,attr_len) for mem in X_bar]) + # X_bar = np.array([dau.st.sliding_window_view(mem, (1,1), (1,1)).reshape(Nx*Ny,attr_len) for mem in X_bar]) state = darr.from_array(state) - X = st.sliding_window_view(state, (1,1), axis=(2,3)).reshape(self.N,attr_len,Nx*Ny) + X = st.sliding_window_view(state, (1, 1), axis=(2, 3)).reshape( + self.N, attr_len, Nx * Ny + ) X_bar = darr.from_array(X_bar) - X_bar = st.sliding_window_view(X_bar, (1,1), axis=(2,3)).reshape(1,attr_len,Nx*Ny) + X_bar = st.sliding_window_view(X_bar, (1, 1), axis=(2, 3)).reshape( + 1, attr_len, Nx * Ny + ) - X = np.swapaxes(X,1,2) - X_bar = np.swapaxes(X_bar,1,2) + X = np.swapaxes(X, 1, 2) + X_bar = np.swapaxes(X_bar, 1, 2) # Make first index the local regions - X = np.swapaxes(X,0,1) + X = np.swapaxes(X, 0, 1) return X, X_bar[0] - def get_obs(self,obs,mask,obs_X,obs_Y,Nx,Ny,attr_len): + def get_obs(self, obs, mask, obs_X, obs_Y, Nx, Ny, attr_len): """ Get observations vector in grid-localised view. @@ -495,52 +577,53 @@ def get_obs(self,obs,mask,obs_X,obs_Y,Nx,Ny,attr_len): x_pad = int((obs_X - 1) / 2) y_pad = int((obs_Y - 1) / 2) pads = [[x_pad, x_pad], [y_pad, y_pad]] - obs = np.array([np.pad(obs_attr,pads,mode='wrap') for obs_attr in obs]) - mask = np.array([np.pad(mask_attr, pads, mode='wrap') for mask_attr in mask]) + obs = np.array([np.pad(obs_attr, pads, mode="wrap") for obs_attr in obs]) + mask = np.array([np.pad(mask_attr, pads, mode="wrap") for mask_attr in mask]) - obs = np.ma.array(obs,mask=mask).filled(fill_value=np.nan) + obs = np.ma.array(obs, mask=mask).filled(fill_value=np.nan) # obs = np.array([dau.st.sliding_window_view(obs_attr, (obs_X,obs_Y), (1,1)) for obs_attr in obs]) obs = darr.from_array(obs) - obs = st.sliding_window_view(obs, (obs_X,obs_Y), axis=(1,2)) + obs = st.sliding_window_view(obs, (obs_X, obs_Y), axis=(1, 2)) - obs = np.swapaxes(obs,0,2) - obs = np.swapaxes(obs,0,1) - obs = obs.reshape(Nx*Ny,attr_len,obs_X,obs_Y) + obs = np.swapaxes(obs, 0, 2) + obs = np.swapaxes(obs, 0, 1) + obs = obs.reshape(Nx * Ny, attr_len, obs_X, obs_Y) return obs - - def get_state_in_obs_space(self,state,mask,obs_attr,obs_X,obs_Y,Nx,Ny,attr_len): + def get_state_in_obs_space( + self, state, mask, obs_attr, obs_X, obs_Y, Nx, Ny, attr_len + ): """ Get state vector projected onto observation space, in grid-localised view. This is step 1 of the LETKF algorithm in Hunt et. al., 2007. """ # Prepare data structure for Y[G] calculation - # stack grid by attributes - sios = self.get_quantity(state,obs_attr[0],inner=False) + # stack grid by attributes + sios = self.get_quantity(state, obs_attr[0], inner=False) - sios = sios[:,np.newaxis,...] + sios = sios[:, np.newaxis, ...] for attr in obs_attr[1:]: - next_attr = self.get_quantity(state,attr,inner=False) + next_attr = self.get_quantity(state, attr, inner=False) # next_attr = np.ma.array(next_attr,mask=mask).filled(fill_value=np.nan) - next_attr = next_attr[:,np.newaxis,...] - sios = np.hstack((sios,next_attr)) + next_attr = next_attr[:, np.newaxis, ...] + sios = np.hstack((sios, next_attr)) # prepare mask array # grid-points where there are observations is 0 x_pad = int((obs_X - 1) / 2) y_pad = int((obs_Y - 1) / 2) pads = [[x_pad, x_pad], [y_pad, y_pad]] - mask = np.array([np.pad(mask_attr, pads, mode='wrap') for mask_attr in mask]) + mask = np.array([np.pad(mask_attr, pads, mode="wrap") for mask_attr in mask]) - # apply mask to get Y[G]. Now Y[G] is in observation space. - sios = np.ma.array([np.ma.array(mem,mask=mask) for mem in sios]) + # apply mask to get Y[G]. Now Y[G] is in observation space. + sios = np.ma.array([np.ma.array(mem, mask=mask) for mem in sios]) # Get bar{y[G]}, i.e. ensemble averages of the states in obs space. - Y_bar = sios.mean(axis=0,keepdims=True) + Y_bar = sios.mean(axis=0, keepdims=True) # Get anomaly for Y[G] sios -= Y_bar @@ -553,16 +636,20 @@ def get_state_in_obs_space(self,state,mask,obs_attr,obs_X,obs_Y,Nx,Ny,attr_len): # Y_bar = np.array([st.sliding_window_view(mem, (obs_X,obs_Y), axis=(1,2)).reshape(Nx*Ny,attr_len,obs_X,obs_Y) for mem in Y_bar]) sios = darr.from_array(sios) - sios = st.sliding_window_view(sios, (obs_X,obs_Y), axis=(2,3)).reshape(self.N,attr_len,Nx*Ny,obs_X,obs_Y) + sios = st.sliding_window_view(sios, (obs_X, obs_Y), axis=(2, 3)).reshape( + self.N, attr_len, Nx * Ny, obs_X, obs_Y + ) Y_bar = darr.from_array(Y_bar) - Y_bar = st.sliding_window_view(Y_bar, (obs_X,obs_Y), axis=(2,3)).reshape(1,attr_len,Nx*Ny,obs_X,obs_Y) + Y_bar = st.sliding_window_view(Y_bar, (obs_X, obs_Y), axis=(2, 3)).reshape( + 1, attr_len, Nx * Ny, obs_X, obs_Y + ) ############################### # from skimage.util import view_as_windows # sios_tmp = np.zeros((10,attr_len,Nx*Ny,obs_X,obs_Y)) # Y_bar_tmp = np.zeros((1,attr_len,Nx*Ny,obs_X,obs_Y)) - + # # sios_tmp = darr.from_array(sios_tmp) # sios = darr.from_array(sios) @@ -576,13 +663,11 @@ def get_state_in_obs_space(self,state,mask,obs_attr,obs_X,obs_Y,Nx,Ny,attr_len): # print(idx) # sios_tmp[idx] = st.sliding_window_view(sios[idx], (obs_X,obs_Y),axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) - # sios_tmp00 = np.zeros((5,attr_len,Nx*Ny,obs_X,obs_Y)) # for idx in range(0,self.N-hN): # print(idx+hN) # sios_tmp00[idx] = st.sliding_window_view(sios[idx+hN], (obs_X,obs_Y),axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) - #################################### @@ -591,55 +676,51 @@ def get_state_in_obs_space(self,state,mask,obs_attr,obs_X,obs_Y,Nx,Ny,attr_len): # sios_tmp = view_as_windows(sios, (self.N,attr_len,obs_X,obs_Y)).reshape(self.N,Nx*Ny,attr_len,obs_X,obs_Y) # Y_bar_tmp[0] = st.sliding_window_view(Y_bar[0], (obs_X,obs_Y),axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) - - + # sios = np.array([st.sliding_window_view(mem, (obs_X,obs_Y), axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) for mem in sios]) # Y_bar = np.array([st.sliding_window_view(mem, (obs_X,obs_Y), axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) for mem in Y_bar]) - # sios = np.array([view_as_windows(mem, (attr_len, obs_X,obs_Y)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) for mem in sios]) # Y_bar = np.array([view_as_windows(mem, (attr_len, obs_X,obs_Y)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) for mem in Y_bar]) # sios = sios_tmp # Y_bar = Y_bar_tmp - # Y_bar = np.array([dau.st.sliding_window_view(mem, (obs_X,obs_Y), (1,1)).reshape(Nx*Ny,attr_len,obs_X,obs_Y) for mem in Y_bar]) - sios = np.swapaxes(sios,1,2) - Y_bar = np.swapaxes(Y_bar,1,2) + sios = np.swapaxes(sios, 1, 2) + Y_bar = np.swapaxes(Y_bar, 1, 2) # Make local region the first axis index - sios = np.swapaxes(sios,0,1) + sios = np.swapaxes(sios, 0, 1) # sios = sios.rechunk(chunks=(self.N,100,attr_len,obs_X,obs_Y)) # print(sios.chunks) return sios, Y_bar[0] - - def get_bc_mask(self, mask, type, Nx, Ny, obs_X, obs_Y,attr_len): + def get_bc_mask(self, mask, type, Nx, Ny, obs_X, obs_Y, attr_len): """ Mask handling the boundary condition for cell or node grids in the local subdomains. """ - if type == 'cell' and self.iicy > 1: - bc_mask = np.ones((self.iicx,self.iicy)) - elif type == 'cell' and self.iicy == 1: - bc_mask = np.ones((self.iicx,self.iicz)) - elif type == 'node' and self.iicyn > 2: - bc_mask = np.ones((self.iicxn,self.iicyn)) - elif type == 'node' and self.iicyn == 2: - bc_mask = np.ones((self.iicxn,self.iiczn)) + if type == "cell" and self.iicy > 1: + bc_mask = np.ones((self.iicx, self.iicy)) + elif type == "cell" and self.iicy == 1: + bc_mask = np.ones((self.iicx, self.iicz)) + elif type == "node" and self.iicyn > 2: + bc_mask = np.ones((self.iicxn, self.iicyn)) + elif type == "node" and self.iicyn == 2: + bc_mask = np.ones((self.iicxn, self.iiczn)) # bc_mask = np.pad(bc_mask, 2, mode='constant', constant_values=(1.0)) - pads = ((self.pad_X,self.pad_X),(self.pad_Y,self.pad_Y)) + pads = ((self.pad_X, self.pad_X), (self.pad_Y, self.pad_Y)) - bc_mask = np.pad(bc_mask, pads, mode='constant', constant_values=(1.0)) + bc_mask = np.pad(bc_mask, pads, mode="constant", constant_values=(1.0)) - mask = np.array([np.pad(attr, pads, mode='wrap') for attr in mask]) + mask = np.array([np.pad(attr, pads, mode="wrap") for attr in mask]) - if type == 'cell': + if type == "cell": bc_mask *= self.cmask else: bc_mask *= self.nmask @@ -649,17 +730,20 @@ def get_bc_mask(self, mask, type, Nx, Ny, obs_X, obs_Y,attr_len): # mask_n = np.array(dau.st.sliding_window_view(mask, (obs_X,obs_Y), (1,1))).reshape(Nx*Ny,attr_len,obs_X,obs_Y) bc_mask = darr.from_array(bc_mask) - bc_mask = st.sliding_window_view(bc_mask, (obs_X,obs_Y)).reshape(Nx*Ny,obs_X,obs_Y) + bc_mask = st.sliding_window_view(bc_mask, (obs_X, obs_Y)).reshape( + Nx * Ny, obs_X, obs_Y + ) mask = darr.from_array(mask) - mask_n = st.sliding_window_view(mask, (obs_X,obs_Y), axis=(1,2)).reshape(attr_len,Nx*Ny,obs_X,obs_Y) + mask_n = st.sliding_window_view(mask, (obs_X, obs_Y), axis=(1, 2)).reshape( + attr_len, Nx * Ny, obs_X, obs_Y + ) - mask_n = np.swapaxes(mask_n,0,1) + mask_n = np.swapaxes(mask_n, 0, 1) return bc_mask, mask_n - - def get_properties(self,type): + def get_properties(self, type): """ For a given grid-type (cell / node), return the 2D grid-size (Nx,Ny), the attributes {rho, rhou...} observed on this grid (obs_attr), the number of attributes (attr_len), and the index location of its data container. @@ -672,49 +756,51 @@ def get_properties(self,type): loc : int """ - if type != 'cell' and type != 'node': + if type != "cell" and type != "node": assert 0, "rloc: grid-type not supported" - Nx = self.iicx if type == 'cell' else self.iicxn + Nx = self.iicx if type == "cell" else self.iicxn if self.iicy > 1: - Ny = self.iicy if type == 'cell' else self.iicyn + Ny = self.iicy if type == "cell" else self.iicyn else: - Ny = self.iicz if type == 'cell' else self.iiczn - obs_attr = self.ca if type == 'cell' else self.na - attr_len = self.cattr_len if type == 'cell' else self.nattr_len - loc = self.loc_c if type == 'cell' else self.loc_n + Ny = self.iicz if type == "cell" else self.iiczn + obs_attr = self.ca if type == "cell" else self.na + attr_len = self.cattr_len if type == "cell" else self.nattr_len + loc = self.loc_c if type == "cell" else self.loc_n return Nx, Ny, obs_attr, attr_len, loc - - def get_quantity(self,results,quantity,inner=True): + def get_quantity(self, results, quantity, inner=True): """ Get ensemble representation of {rho, rhou...}. """ slc = self.i2 loc = self.loc[quantity] - pads = ((self.pad_X,self.pad_X),(self.pad_Y,self.pad_Y)) + pads = ((self.pad_X, self.pad_X), (self.pad_Y, self.pad_Y)) if inner: - attribute_array = [getattr(results[:,loc,...][n],quantity)[slc] for n in range(self.N)] + attribute_array = [ + getattr(results[:, loc, ...][n], quantity)[slc] for n in range(self.N) + ] else: - attribute_array = [np.pad(getattr(results[:,loc,...][n],quantity)[slc],pads, mode='wrap') for n in range(self.N)] + attribute_array = [ + np.pad( + getattr(results[:, loc, ...][n], quantity)[slc], pads, mode="wrap" + ) + for n in range(self.N) + ] return np.array(attribute_array) - - def get_loc_mat(self,bc_mask,mask_n,n,attr_len): + def get_loc_mat(self, bc_mask, mask_n, n, attr_len): loc_mat = self.loc_mat * bc_mask[n] - loc_mat = np.expand_dims(loc_mat,axis=0) + loc_mat = np.expand_dims(loc_mat, axis=0) loc_mat = np.repeat(loc_mat, attr_len, axis=0) - loc_mat = np.ma.array(loc_mat,mask=mask_n[n]).filled(fill_value=np.nan) + loc_mat = np.ma.array(loc_mat, mask=mask_n[n]).filled(fill_value=np.nan) loc_mat = loc_mat[~np.isnan(loc_mat)] loc_mat = np.diag(loc_mat.flatten()) return loc_mat - - - diff --git a/src/data_assimilation/localisation.py b/src/pybella/data_assimilation/localisation.py similarity index 73% rename from src/data_assimilation/localisation.py rename to src/pybella/data_assimilation/localisation.py index 6069a342..8c61c6af 100644 --- a/src/data_assimilation/localisation.py +++ b/src/pybella/data_assimilation/localisation.py @@ -1,7 +1,8 @@ import numpy as np -import dycore.utils.options as opts +import flow_solver.utils.options as opts -def rlocal_5pt(elem,node,ud): + +def rlocal_5pt(elem, node, ud): igx = elem.igx igy = elem.igy @@ -19,7 +20,10 @@ def rlocal_5pt(elem,node,ud): x_wall = ud.bdry_type[1] == opts.BdryType.WALL y_wall = ud.bdry_type[0] == opts.BdryType.WALL - return lambda covar : rlocal_5pt_stencil(covar, iicxn, iicyn, x_periodic, y_periodic, x_wall, y_wall) + return lambda covar: rlocal_5pt_stencil( + covar, iicxn, iicyn, x_periodic, y_periodic, x_wall, y_wall + ) + # @jit(nopython=True, nogil=True, cache=True) def rlocal_5pt_stencil(covar, iicxn, iicyn, x_periodic, y_periodic, x_wall, y_wall): @@ -52,20 +56,20 @@ def rlocal_5pt_stencil(covar, iicxn, iicyn, x_periodic, y_periodic, x_wall, y_wa botmid_idx -= iicxn - 1 if cnt_y == 0: - topmid_idx += ((iicxn) * (iicyn - 1)) + topmid_idx += (iicxn) * (iicyn - 1) if y_periodic: - midleft_idx += ((iicxn) * (iicyn - 1)) - midmid_idx += ((iicxn) * (iicyn - 1)) - midright_idx += ((iicxn) * (iicyn - 1)) + midleft_idx += (iicxn) * (iicyn - 1) + midmid_idx += (iicxn) * (iicyn - 1) + midright_idx += (iicxn) * (iicyn - 1) if cnt_y == (iicyn - 1): - botmid_idx -= ((iicxn) * (iicyn - 1)) + botmid_idx -= (iicxn) * (iicyn - 1) if y_periodic: - midleft_idx -= ((iicxn) * (iicyn - 1)) - midmid_idx -= ((iicxn) * (iicyn - 1)) - midright_idx -= ((iicxn) * (iicyn - 1)) + midleft_idx -= (iicxn) * (iicyn - 1) + midmid_idx -= (iicxn) * (iicyn - 1) + midright_idx -= (iicxn) * (iicyn - 1) midleft = covar[midleft_idx] topmid = covar[topmid_idx] @@ -81,15 +85,17 @@ def rlocal_5pt_stencil(covar, iicxn, iicyn, x_periodic, y_periodic, x_wall, y_wa # if y_wall and (cnt_y == 0): # topmid = 0.0 - + # if y_wall and (cnt_y == (iicyn - 1)): # botmid = 0.0 - - R[idx] = (0.5 * midleft + 1.0 * midmid + 0.5 * midright) + (0.5 * topmid + 1.0 * midmid + 0.5 * botmid) + + R[idx] = (0.5 * midleft + 1.0 * midmid + 0.5 * midright) + ( + 0.5 * topmid + 1.0 * midmid + 0.5 * botmid + ) cnt_x += 1 if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - - return R \ No newline at end of file + + return R diff --git a/src/data_assimilation/params.py b/src/pybella/data_assimilation/params.py similarity index 70% rename from src/data_assimilation/params.py rename to src/pybella/data_assimilation/params.py index a9e921bb..8eb4f995 100644 --- a/src/data_assimilation/params.py +++ b/src/pybella/data_assimilation/params.py @@ -1,27 +1,27 @@ +import logging + +import h5py import numpy as np import scipy as sp -import h5py -import dycore.utils.boundary as bdry +from ..flow_solver.utils import boundary as bdry -import logging class init(object): - - def __init__(self,N,da_type='rloc'): + def __init__(self, N, da_type="rloc"): # number of ensemble members self.N = N - self._da_times = np.arange(0.0,3.25,0.25)[1:] + self._da_times = np.arange(0.0, 3.25, 0.25)[1:] # self._da_times = np.arange(5.0,10.5,0.5)/10.0 # self._da_times = [0.1] - self._da_times = np.around(self.da_times,3) - - self.obs_attributes = ['rho','rhou', 'rhov', 'rhoY', 'p2_nodes'] + self._da_times = np.around(self.da_times, 3) + + self.obs_attributes = ["rho", "rhou", "rhov", "rhoY", "p2_nodes"] # self.obs_attributes = ['rhou','rhov'] # which attributes to inflate in ensemble inflation? - self.attributes = ['rho', 'rhou', 'rhov'] + self.attributes = ["rho", "rhou", "rhov"] # self.obs_path = './output_travelling_vortex/output_travelling_vortex_ensemble=1_32_32_6.0_truthgen.h5' # self.obs_path = './output_rising_bubble/output_rising_bubble_ensemble=1_100_50_10.0_psinc_ref.h5' @@ -32,7 +32,7 @@ def __init__(self,N,da_type='rloc'): # self.obs_path = './output_swe_vortex/output_swe_vortex_ensemble=1_64_1_64_3.0_comp_1.0_pps_tra_truth.h5' # self.obs_path = './output_swe_vortex/output_swe_vortex_ensemble=1_64_1_64_3.0_neg_comp_1.0_pp_tra_truth_ip.h5' # self.obs_path = './output_travelling_vortex/output_travelling_vortex_ensemble=1_64_64_3.0_comp_1.0_pp_tra_truth_ip.h5' - self.obs_path = './output_travelling_vortex/output_travelling_vortex_ensemble=1_64_64_3.0_obs.h5' + self.obs_path = "./output_travelling_vortex/output_travelling_vortex_ensemble=1_64_64_3.0_obs.h5" # forward operator (projector from state space to observation space) self.forward_operator = np.eye(N) @@ -47,14 +47,14 @@ def __init__(self,N,da_type='rloc'): if self.sparse_obs_by_attr: assert 0, "Not yet implemented." - self.obs_frac = 0.10 # fraction of the observations to pick. + self.obs_frac = 0.10 # fraction of the observations to pick. self.gen_obs_sparse() ############################################ # Parameters for measurement noise ############################################ self.add_obs_noise = True - self.noise_type = 'VarCov' + self.noise_type = "VarCov" self._noise_percentage = 0.05 self.obs_noise = {} @@ -68,7 +68,7 @@ def __init__(self,N,da_type='rloc'): self.obs_Y = 11 # constants, linear, gaussian - self.localisation_matrix = self.get_loc_mat('gaussian') + self.localisation_matrix = self.get_loc_mat("gaussian") self.da_type = da_type @@ -78,19 +78,19 @@ def __init__(self,N,da_type='rloc'): # rejuvenation factor for ETPF self.rejuvenation_factor = 0.001 - self.loc_c = 0 # container list location of cell-based arrays - self.loc_f = 1 # ... of face-based arrays - self.loc_n = 2 # ... of node-based arrays + self.loc_c = 0 # container list location of cell-based arrays + self.loc_f = 1 # ... of face-based arrays + self.loc_n = 2 # ... of node-based arrays # in which data container are the attributes involved in the DA procedure? self.loc = { - 'rho' : 0, - 'rhou' : 0, - 'rhov' : 0, - 'rhow' : 0, - 'rhoY' : 0, - 'rhoX' : 0, - 'p2_nodes' : 2, + "rho": 0, + "rhou": 0, + "rhov": 0, + "rhow": 0, + "rhoY": 0, + "rhoX": 0, + "p2_nodes": 2, } def gen_obs_sparse(self): @@ -105,62 +105,64 @@ def gen_obs_sparse(self): if len(self.sparse_obs_seeds) > 1: self.sparse_obs_seeds = self.sparse_obs_seeds.squeeze() - def gen_obs_noise(self): for cnt, key in enumerate(self.obs_attributes): - if self.noise_type == 'FixCov': - assert self.std_dev is not None, "std_dev keyword argument must be a list equal in size to dap.obs_attributes" - assert len(self.std_dev) == len(self.obs_attributes), "std_dev keyword argument must be a list equal in size to dap.obs_attributes" + if self.noise_type == "FixCov": + assert ( + self.std_dev is not None + ), "std_dev keyword argument must be a list equal in size to dap.obs_attributes" + assert len(self.std_dev) == len( + self.obs_attributes + ), "std_dev keyword argument must be a list equal in size to dap.obs_attributes" self.obs_noise[key] = float(self.std_dev[cnt]) logging.info(self.std_dev[cnt]) cnt += 1 - else: + else: self.obs_noise[key] = self._noise_percentage da_depth = len(self.obs_attributes) np.random.seed(888) if da_depth > 1: - self.obs_noise_seeds = np.random.randint(10000,size=(da_depth)).squeeze() + self.obs_noise_seeds = np.random.randint(10000, size=(da_depth)).squeeze() else: self.obs_noise_seeds = [np.random.randint(10000)] @staticmethod def converter(results, N, mpv, elem, node, th, ud): - ''' + """ Do this after data assimilation for HS balanced vortex. - ''' + """ logging.info("Post DA conversion...") g = ud.g0 for n in range(N): - bdry.set_explicit_boundary_data(results[n][0],elem,ud,th,mpv) - results[n][0].rhoY[...] = (g / 2.0 * results[n][0].rho**2)**th.gamminv + bdry.set_explicit_boundary_data(results[n][0], elem, ud, th, mpv) + results[n][0].rhoY[...] = (g / 2.0 * results[n][0].rho ** 2) ** th.gamminv igy = elem.igy - kernel = np.ones((2,2)) + kernel = np.ones((2, 2)) kernel /= kernel.sum() - pn = sp.signal.convolve(results[n][0].rhoY[:,igy,:], kernel, mode='valid') + pn = sp.signal.convolve(results[n][0].rhoY[:, igy, :], kernel, mode="valid") - bdry.set_explicit_boundary_data(results[n][0],elem,ud,th,mpv) + bdry.set_explicit_boundary_data(results[n][0], elem, ud, th, mpv) pn = np.expand_dims(pn, 1) pn = np.repeat(pn, node.icy, axis=1) - results[n][2].p2_nodes[1:-1,:,1:-1] = pn - bdry.set_ghostnodes_p2(results[n][2].p2_nodes,node,ud) + results[n][2].p2_nodes[1:-1, :, 1:-1] = pn + bdry.set_ghostnodes_p2(results[n][2].p2_nodes, node, ud) - pn = np.expand_dims(results[n][2].p2_nodes[:,igy,:], 1) + pn = np.expand_dims(results[n][2].p2_nodes[:, igy, :], 1) results[n][2].p2_nodes[...] = np.repeat(pn[...], node.icy, axis=1) return results - - def load_obs(self,obs_path,loc=0): + def load_obs(self, obs_path, loc=0): if self.N > 1: - obs_file = h5py.File(obs_path, 'r') + obs_file = h5py.File(obs_path, "r") obs_attributes = self.obs_attributes #### when were these observations taken? @@ -172,42 +174,39 @@ def load_obs(self,obs_path,loc=0): t_cnt = 0 for t in times: #### how were these dataset called? - label = '_ensemble_mem=0_%.3f_after_full_step' %t + label = "_ensemble_mem=0_%.3f_after_full_step" % t #### axis 1 stores the attributes obs[t_cnt] = {} for attribute in obs_attributes: data = obs_file[str(attribute)][str(attribute) + str(label)][:] - if data.ndim == 3: # implying horizontal slice... - data = data[:,0,:] - dict_attr = { - attribute: data - } + if data.ndim == 3: # implying horizontal slice... + data = data[:, 0, :] + dict_attr = {attribute: data} obs[t_cnt].update(dict_attr) t_cnt += 1 obs = np.array(obs) obs_file.close() return obs - def get_loc_mat(self, type, c=1.0): - x = (np.linspace(0,1,self.obs_X) - 0.5) - y = (np.linspace(0,1,self.obs_Y) - 0.5) + x = np.linspace(0, 1, self.obs_X) - 0.5 + y = np.linspace(0, 1, self.obs_Y) - 0.5 - X,Y = np.meshgrid(x,y) - if type == 'constants': + X, Y = np.meshgrid(x, y) + if type == "constants": Z = np.ones((self.obs_X, self.obs_Y)) return Z.flatten() * c - - elif type == 'linear': - Z = (np.sqrt(X**2 + Y**2)) + + elif type == "linear": + Z = np.sqrt(X**2 + Y**2) Z = Z.max() - Z return Z.flatten() * c - elif type == 'gaussian': - pos = np.array([X.flatten(),Y.flatten()]).T - norm = sp.stats.multivariate_normal([0,0],[[0.05,0.0],[0.0,0.05]]) - Z = norm.pdf(pos).reshape(X.shape[0],X.shape[1]) + elif type == "gaussian": + pos = np.array([X.flatten(), Y.flatten()]).T + norm = sp.stats.multivariate_normal([0, 0], [[0.05, 0.0], [0.0, 0.05]]) + Z = norm.pdf(pos).reshape(X.shape[0], X.shape[1]) Z = Z - Z.min() Z /= Z.max() @@ -217,7 +216,6 @@ def update_dap(self, params): for key, value in params.items(): setattr(self, key, value) - # setter functions @property def noise_percentage(self): @@ -249,11 +247,10 @@ def da_times(self): @da_times.setter def da_times(self, lst): self._da_times = lst - self._da_times = np.around(self._da_times,3) + self._da_times = np.around(self._da_times, 3) if self.sparse_obs: self.gen_obs_sparse() - @property def loc_setter(self): return self.loc_setter @@ -263,17 +260,16 @@ def loc_setter(self, tpl): self.obs_X = tpl[0] self.obs_Y = tpl[1] # constants, linear, gaussian - self.localisation_matrix = self.get_loc_mat('gaussian') - + self.localisation_matrix = self.get_loc_mat("gaussian") @property def sd_setter(self): return self.sd_setter - + @sd_setter.setter - def sd_setter(self,lst): + def sd_setter(self, lst): self.std_dev = lst - self.noise_type = 'FixCov' + self.noise_type = "FixCov" if self.add_obs_noise: self.gen_obs_noise() diff --git a/src/data_assimilation/post_processing.py b/src/pybella/data_assimilation/post_processing.py similarity index 78% rename from src/data_assimilation/post_processing.py rename to src/pybella/data_assimilation/post_processing.py index 6525fbf1..2360ddec 100644 --- a/src/data_assimilation/post_processing.py +++ b/src/pybella/data_assimilation/post_processing.py @@ -3,4 +3,4 @@ ## Interface function def interface(): - return None \ No newline at end of file + return None diff --git a/src/pybella/data_assimilation/prepare.py b/src/pybella/data_assimilation/prepare.py new file mode 100644 index 00000000..f50ddb1e --- /dev/null +++ b/src/pybella/data_assimilation/prepare.py @@ -0,0 +1,100 @@ +import logging + +# to generate ensemble from one sol init instantiation +from copy import deepcopy + +import numpy as np + +from ..utils import sim_params as params +from ..utils import io +from ..utils import data_structures + +from . import utils as da_utils +from . import params as da_params +from . import letkf as da_letkf + + +def initialise(sst): + es = sst.ensemble_state + rp = sst.restart_params + + ########################################################## + # Initialisation of data assimilation module + ########################################################## + + # possible da_types: + # 1) batch_obs for the LETKF with batch observations + # 2) rloc for LETKF with grid-point localisation + # 3) etpf for the ETPF algorithm + dap = da_params.init(sst.N, da_type="rloc") + if rp.dap_rewrite is not None: + dap.update_dap(rp.dap_rewrite) + + # if elem.ndim == 2: + if dap.da_type == "rloc" and sst.N > 1: + rloc = da_letkf.prepare_rloc(es.ud, es.elem, es.node, dap, sst.N) + else: + rloc = None + + logging.info("Generating initial ensemble...") + sol_ens = data_structures.EnsembleState() + + # Set random seed for reproducibility + np.random.seed(params.random_seed) + + seeds = np.random.randint(10000, size=sst.N) if sst.N > 1 else None + + if sst.N > 1: + logging.info("Seeds used in generating initial ensemble spread = ", seeds) + for n in range(sst.N): + Sol0 = deepcopy(sst.Sol) + mpv0 = deepcopy(sst.mpv) + Sol0 = sst.sol_init( + Sol0, mpv0, es.elem, es.node, es.th, es.ud, seed=seeds[n] + ) + # sol_ens[n] = [Sol0, deepcopy(es.flux), mpv0, [-np.inf, es.step]] + sol_ens.update_member( + es.elem, es.node, Sol0, mpv0, deepcopy(es.flux), es.th + ) + + sst.ensembble_state = sol_ens + # elif sst.restart == False: + # sol_ens = [[sst.sol_init(mp.Sol, mp.mpv, mp.elem, mp.node, mp.th, sst.ud), mp.flux, mp.mpv, [-np.inf, sst.step]]] + # sol_ens.update_member(mp.elem, mp.node, sst.sol_init(mp.Sol, mp.mpv, mp.elem, mp.node, mp.th, sst.ud), mp.mpv, deepcopy(mp.flux), mp.th) + # for n in range(sst.N): + # sol_ens.get_member(n).time.t = -np.inf + + # ens = da_utils.ensemble(sol_ens) + + ########################################################## + # Load data assimilation observations + ########################################################## + + # where are my observations? + if sst.N > 1: + obs = dap.load_obs(dap.obs_path) + # obs_mask, no calculations where entries are True + obs_mask = da_utils.sparse_obs_selector(obs, es.elem, es.node, sst.ud, dap) + obs_noisy, obs_covar = da_utils.obs_noiser(obs, obs_mask, dap, rloc, es.elem) + else: + obs, obs_noisy, obs_mask, obs_covar = None, None, None, None + + ########################################################## + # Add ensemble info into filename + ########################################################## + if sst.ud.autogen_fn: + sst.ud.output_suffix = io.fn_gen(sst.ud, dap, sst.N) + + ####################### + # Populate DA params + ####################### + + sst.da_params = data_structures.DataAssimilationParameters( + dap=dap, + rloc=rloc, + # sol_ens=ens, + obs=obs, + obs_noisy=obs_noisy, + obs_mask=obs_mask, + obs_covar=obs_covar, + ) diff --git a/src/data_assimilation/utils.py b/src/pybella/data_assimilation/utils.py similarity index 64% rename from src/data_assimilation/utils.py rename to src/pybella/data_assimilation/utils.py index f5b04596..f81dc133 100644 --- a/src/data_assimilation/utils.py +++ b/src/pybella/data_assimilation/utils.py @@ -1,38 +1,38 @@ +import copy + import numpy as np import scipy as sp -import copy import matplotlib.pyplot as plt -import dycore.utils.boundary as bdry -import dycore.utils.options as opts +from ..flow_solver.utils import options as opts, boundary as bdry + class ensemble(object): def __init__(self, input_ensemble=None): if input_ensemble is not None: cnt = 0 for mem in input_ensemble: - setattr(self,'mem_' + str(cnt),mem) + setattr(self, "mem_" + str(cnt), mem) # self.debug_im(mem[0].rho, cnt) cnt += 1 else: None - - def initialise_members(self,ic,N): + def initialise_members(self, ic, N): for cnt in range(N): mem = [copy.deepcopy(arr) for arr in ic] # mem = sampler(mem) - setattr(self,'mem_' + str(cnt),mem) + setattr(self, "mem_" + str(cnt), mem) # def state_vector(self,ensemble): # N = self.members(ensemble).shape[0] # return self.members(ensemble).reshape(N,-1) - def set_members(self,analysis_ensemble, tout): + def set_members(self, analysis_ensemble, tout): cnt = 0 for xi in analysis_ensemble: - setattr(self,'mem_' + str(cnt),np.array(xi)) + setattr(self, "mem_" + str(cnt), np.array(xi)) # self.debug_im(xi[0].rho, cnt, tout) cnt += 1 @@ -42,107 +42,129 @@ def ensemble_spreading(self, ens, sampler, attributes, loc=0): np.random.seed(555) for attribute in attributes: for n in range(N): - mem = getattr(self,'mem_' + str(n)) - value = getattr(mem[loc],attribute) + mem = getattr(self, "mem_" + str(n)) + value = getattr(mem[loc], attribute) # self.debug_im(sampler(value), n) - setattr(mem[loc],attribute,sampler(value)) + setattr(mem[loc], attribute, sampler(value)) @staticmethod def members(ensemble): - return np.array(list(ensemble.__dict__.values()),dtype='object') + return np.array(list(ensemble.__dict__.values()), dtype="object") @staticmethod def debug_im(value, n, tout): plt.figure() - plt.imshow(value, origin='lower') - plt.savefig("./output_images/ensemble_snapshots/%.3f_%03d.png" %(tout,n), bbox_inches='tight') + plt.imshow(value, origin="lower") + plt.savefig( + "./output_images/ensemble_snapshots/%.3f_%03d.png" % (tout, n), + bbox_inches="tight", + ) plt.close() + def ensemble_inflation(results, attributes, factor, N, loc=0): for attribute in attributes: - mean = [getattr(results[n][loc],attribute) for n in range(N)] + mean = [getattr(results[n][loc], attribute) for n in range(N)] mean = np.array(mean) - mean = np.mean(mean,axis=0) + mean = np.mean(mean, axis=0) for n in range(N): - inflation = mean + factor * (getattr(results[n][loc],attribute) - mean) - setattr(results[n][loc],attribute,inflation) + inflation = mean + factor * (getattr(results[n][loc], attribute) - mean) + setattr(results[n][loc], attribute, inflation) + -def set_p2_nodes(analysis,results,N,th,node,ud,loc_c=0,loc_n=2): +def set_p2_nodes(analysis, results, N, th, node, ud, loc_c=0, loc_n=2): for n in range(N): rhoY = analysis[n] - p2_n = getattr(results[n][loc_n],'p2_nodes') + p2_n = getattr(results[n][loc_n], "p2_nodes") rhoY_n = np.zeros_like(p2_n) - kernel = np.array([[1.,1.],[1.,1.]]) - rhoY_n[1:-1,1:-1] = sp.signal.fftconvolve(rhoY,kernel,mode='valid') / kernel.sum() - bdry.set_ghostnodes_p2(rhoY_n,node,ud) + kernel = np.array([[1.0, 1.0], [1.0, 1.0]]) + rhoY_n[1:-1, 1:-1] = ( + sp.signal.fftconvolve(rhoY, kernel, mode="valid") / kernel.sum() + ) + bdry.set_ghostnodes_p2(rhoY_n, node, ud) p2_n = rhoY_n**th.gm1 - 1.0 + (p2_n - p2_n.mean()) # p2_n = rhoY_n**th.gm1 - 1.0 p2_n -= p2_n.mean() # p2_n = np.pad(p2_n,2,mode='wrap') - bdry.set_ghostnodes_p2(p2_n,node,ud) - setattr(results[n][loc_n],'p2_nodes',p2_n) + bdry.set_ghostnodes_p2(p2_n, node, ud) + setattr(results[n][loc_n], "p2_nodes", p2_n) + -def set_rhoY_cells(analysis,results,N,th,ud,loc_c=0,loc_n=2): +def set_rhoY_cells(analysis, results, N, th, ud, loc_c=0, loc_n=2): for n in range(N): p2n = analysis[n] - rhoYc0 = getattr(results[n][loc_c], 'rhoY') - kernel = np.array([[1.,1.],[1.,1.]]) - p2c = sp.signal.fftconvolve(p2n,kernel,mode='valid') / kernel.sum() + rhoYc0 = getattr(results[n][loc_c], "rhoY") + kernel = np.array([[1.0, 1.0], [1.0, 1.0]]) + p2c = sp.signal.fftconvolve(p2n, kernel, mode="valid") / kernel.sum() p2c -= p2c.mean() - rhoYc = (rhoYc0**th.gm1 + ud.Msq*p2c)**th.gm1inv - setattr(results[n][loc_c], 'rhoYc', rhoYc) - -def boundary_mask(ud,elem,node,pad_X,pad_Y): + rhoYc = (rhoYc0**th.gm1 + ud.Msq * p2c) ** th.gm1inv + setattr(results[n][loc_c], "rhoYc", rhoYc) + + +def boundary_mask(ud, elem, node, pad_X, pad_Y): """ - Returns a mask for the underlying cellular and nodal grids padded such that the size of the local subdomain has been accounted for. For ghost cells on wall boundaries, values are 0.0 and 1.0 for periodic. + Returns a mask for the underlying cellular and nodal grids padded such that the size of the local subdomain has been accounted for. For ghost cells on wall boundaries, values are 0.0 and 1.0 for periodic. """ - pads = [pad_X,pad_Y] + pads = [pad_X, pad_Y] if elem.iicy > 1: cmask = np.ones(elem.iisc).squeeze() nmask = np.ones(node.iisc).squeeze() for dim in range(elem.ndim): - ghost_padding = [[0,0]] * elem.ndim + ghost_padding = [[0, 0]] * elem.ndim # ghost_padding[dim] = [elem.igs[dim],elem.igs[dim]] - ghost_padding[dim] = [pads[dim],pads[dim]] + ghost_padding[dim] = [pads[dim], pads[dim]] if ud.bdry_type[dim] == opts.BdryType.PERIODIC: - cmask = np.pad(cmask, ghost_padding, mode='constant', constant_values=(1.0)) - nmask = np.pad(nmask, ghost_padding, mode='constant', constant_values=(1.0)) + cmask = np.pad( + cmask, ghost_padding, mode="constant", constant_values=(1.0) + ) + nmask = np.pad( + nmask, ghost_padding, mode="constant", constant_values=(1.0) + ) elif ud.bdry_type[dim] == opts.BdryType.WALL: - cmask = np.pad(cmask, ghost_padding, mode='constant', constant_values=(0.0)) - nmask = np.pad(nmask, ghost_padding, mode='constant', constant_values=(0.0)) - - elif elem.iicy == 1: # implying horizontal slices + cmask = np.pad( + cmask, ghost_padding, mode="constant", constant_values=(0.0) + ) + nmask = np.pad( + nmask, ghost_padding, mode="constant", constant_values=(0.0) + ) + + elif elem.iicy == 1: # implying horizontal slices cmask = np.ones(elem.iisc).squeeze() - nmask = np.ones(node.iisc)[:,0,:] + nmask = np.ones(node.iisc)[:, 0, :] ndim = elem.ndim - 1 for dim in range(ndim): - ghost_padding = [[0,0]] * ndim + ghost_padding = [[0, 0]] * ndim # ghost_padding[dim] = [elem.igs[dim],elem.igs[dim]] - ghost_padding[dim] = [pads[dim],pads[dim]] + ghost_padding[dim] = [pads[dim], pads[dim]] if ud.bdry_type[dim] == opts.BdryType.PERIODIC: - cmask = np.pad(cmask, ghost_padding, mode='constant', constant_values=(1.0)) - nmask = np.pad(nmask, ghost_padding, mode='constant', constant_values=(1.0)) + cmask = np.pad( + cmask, ghost_padding, mode="constant", constant_values=(1.0) + ) + nmask = np.pad( + nmask, ghost_padding, mode="constant", constant_values=(1.0) + ) elif ud.bdry_type[dim] == opts.BdryType.WALL: - cmask = np.pad(cmask, ghost_padding, mode='constant', constant_values=(0.0)) - nmask = np.pad(nmask, ghost_padding, mode='constant', constant_values=(0.0)) - - return cmask.astype('bool'), nmask.astype('bool') - - + cmask = np.pad( + cmask, ghost_padding, mode="constant", constant_values=(0.0) + ) + nmask = np.pad( + nmask, ghost_padding, mode="constant", constant_values=(0.0) + ) + return cmask.astype("bool"), nmask.astype("bool") # ref: https://gist.github.com/meowklaski/4bda7c86c6168f3557657d5fb0b5395a def sliding_window_view(arr, window_shape, steps): - """ + """ Produce a view from a sliding, striding window over `arr`. The window is only placed in 'valid' positions - no overlapping over the boundary. @@ -208,9 +230,10 @@ def sliding_window_view(arr, window_shape, steps): >>> conv_out = conv_out.transpose([3,0,1,2]) """ - + from numpy.lib.stride_tricks import as_strided - in_shape = np.array(arr.shape[-len(steps):]) # [x, (...), z] + + in_shape = np.array(arr.shape[-len(steps) :]) # [x, (...), z] window_shape = np.array(window_shape) # [Wx, (...), Wz] steps = np.array(steps) # [Sx, (...), Sz] nbytes = arr.strides[-1] # size (bytes) of an element in `arr` @@ -218,13 +241,13 @@ def sliding_window_view(arr, window_shape, steps): # number of per-byte steps to take to fill window window_strides = tuple(np.cumprod(arr.shape[:0:-1])[::-1]) + (1,) # number of per-byte steps to take to place window - step_strides = tuple(window_strides[-len(steps):] * steps) + step_strides = tuple(window_strides[-len(steps) :] * steps) # number of bytes to step to populate sliding window view strides = tuple(int(i) * nbytes for i in step_strides + window_strides) outshape = tuple((in_shape - window_shape) // steps + 1) # outshape: ([X, (...), Z], ..., [Wx, (...), Wz]) - outshape = outshape + arr.shape[:-len(steps)] + tuple(window_shape) + outshape = outshape + arr.shape[: -len(steps)] + tuple(window_shape) return as_strided(arr, shape=outshape, strides=strides, writeable=False) @@ -235,7 +258,7 @@ def HSprojector_3t2D(results, elem, dap, N): Parameters ---------- results : nd.array - An array of ensemble size k. Each ensemble member has [Sol,flux,mpv,[window_step,step]]. + An array of ensemble size k. Each ensemble member has [Sol,flux,mpv,[window_step,step]]. dap : data assimilation input class . N : int @@ -247,19 +270,28 @@ def HSprojector_3t2D(results, elem, dap, N): """ - slc = (slice(None,), 2, slice(None,)) + slc = ( + slice( + None, + ), + 2, + slice( + None, + ), + ) # implying horizontal slice... - if elem.ndim == 3 and elem.iicy == 1: + if elem.ndim == 3 and elem.iicy == 1: for key, value in dap.loc.items(): - # only reshape data arrays that are involved in DA. - if key in dap.obs_attributes: + # only reshape data arrays that are involved in DA. + if key in dap.obs_attributes: for n in range(N): - p_arr = getattr(results[n][value],key)[slc] - setattr(results[n][value],key, p_arr) + p_arr = getattr(results[n][value], key)[slc] + setattr(results[n][value], key, p_arr) return results + def HSprojector_2t3D(results, elem, node, dap, N): """ Projection method from 2D array to 3D horizontal slice. To be used after data assimilation. @@ -268,39 +300,46 @@ def HSprojector_2t3D(results, elem, node, dap, N): if elem.ndim == 3 and elem.iicy == 1: for key, value in dap.loc.items(): if key in dap.obs_attributes: - if value == 0: # cell container + if value == 0: # cell container ys = elem.icy - if value == 2: # node container + if value == 2: # node container ys = node.icy for n in range(N): - p_arr = getattr(results[n][value],key)[:,np.newaxis,:] + p_arr = getattr(results[n][value], key)[:, np.newaxis, :] p_arr = np.repeat(p_arr, ys, axis=1) - setattr(results[n][value],key, p_arr) + setattr(results[n][value], key, p_arr) return results - + def sparse_obs_selector(obs, elem, node, ud, dap): sparse_obs = dap.sparse_obs if not sparse_obs or len(dap.da_times) == 0: mask = copy.deepcopy(obs) - for tt,mask_t in enumerate(mask): + for tt, mask_t in enumerate(mask): for key, _ in mask_t.items(): mask[tt][key][...] = 0.0 return obs, mask - + else: sparse_obs_by_attr = dap.sparse_obs_by_attr seeds = dap.sparse_obs_seeds K = dap.obs_frac # define inner and outer domains in 2D - i2 = (slice(elem.igx,-elem.igx),slice(elem.igy,-elem.igy)) - i0 = (slice(None,),slice(None,)) + i2 = (slice(elem.igx, -elem.igx), slice(elem.igy, -elem.igy)) + i0 = ( + slice( + None, + ), + slice( + None, + ), + ) # get inner domain size - if elem.iicy == 1: # implying horizontal slice + if elem.iicy == 1: # implying horizontal slice Ncx, Ncy = elem.iicx, elem.iicz Nnx, Nny = node.iicx, node.iicz else: @@ -314,12 +353,12 @@ def sparse_obs_selector(obs, elem, node, ud, dap): Kc = int(np.ceil(Kc)) Kn = int(np.ceil(Kn)) - if elem.iicy == 1: # implying horizontal slice + if elem.iicy == 1: # implying horizontal slice Xc, Yc = elem.x, elem.z Xc, Yc = np.meshgrid(Xc, Yc) Xn, Yn = node.x, node.z Xn, Yn = np.meshgrid(Xn, Yn) - else: # implying vertical slice + else: # implying vertical slice Xc, Yc = elem.x, elem.y Xc, Yc = np.meshgrid(Xc, Yc) Xn, Yn = node.x, node.y @@ -329,10 +368,10 @@ def sparse_obs_selector(obs, elem, node, ud, dap): # obs_noisy_interp = deepcopy(obs) # obs is a list of dictionaries, list length da_len, dictionary length attr_len. - for tt,obs_t in enumerate(obs): + for tt, obs_t in enumerate(obs): attr_cnt = tt for key, value in obs_t.items(): - if key == 'p2_nodes': + if key == "p2_nodes": grid_x, grid_y = Xn[i0], Yn[i0] K, N = Kn, Nn Nx, Ny = Nnx, Nny @@ -347,10 +386,10 @@ def sparse_obs_selector(obs, elem, node, ud, dap): # https://stackoverflow.com/questions/19597473/binary-random-array-with-a-specific-proportion-of-ones/19597805 # append mask array with new seed - mask = np.array([0] * K + [1] * (N-K)) + mask = np.array([0] * K + [1] * (N - K)) np.random.shuffle(mask) - mask = mask.reshape(Nx,Ny) - mask = np.pad(mask,(2,2),mode='constant',constant_values=0.0) + mask = mask.reshape(Nx, Ny) + mask = np.pad(mask, (2, 2), mode="constant", constant_values=0.0) # values = np.ma.array(value[i0], mask=mask).compressed() # X = np.ma.array(grid_x, mask=mask).compressed() @@ -361,83 +400,88 @@ def sparse_obs_selector(obs, elem, node, ud, dap): # points[:,1] = Y[...].flatten() # values = griddata(points, values, (grid_x, grid_y), method='cubic') - + # if dap.obs_frac < 1.0: - # obs_noisy_interp[tt][key][...] = 0.0 - # obs_noisy_interp[tt][key][i0] = values + # obs_noisy_interp[tt][key][...] = 0.0 + # obs_noisy_interp[tt][key][i0] = values mask_arr[tt][key][...] = 1 mask_arr[tt][key][i2] = mask[i2] - + if sparse_obs_by_attr: attr_cnt += 1 for mask_at_t in mask_arr: for key, mask in mask_at_t.items(): - if key == 'p2_nodes': + if key == "p2_nodes": Nx, Ny = Nnx, Nny else: Nx, Ny = Ncx, Ncy - assert (mask.shape[0] * mask.shape[1]) - mask.sum() == np.ceil(Nx*Ny*dap.obs_frac), "Mask sparsity does not match obs_frac defined" + assert (mask.shape[0] * mask.shape[1]) - mask.sum() == np.ceil( + Nx * Ny * dap.obs_frac + ), "Mask sparsity does not match obs_frac defined" # return obs_noisy_interp, mask_arr return mask_arr -def obs_noiser(obs,mask,dap,rloc,elem): +def obs_noiser(obs, mask, dap, rloc, elem): if dap.add_obs_noise: assert isinstance(dap.obs_noise, dict), "obs_noise has to be dict" - assert len(dap.obs_noise) == len(dap.obs_attributes), "obs_noise length has to be equal to obs_attributes len" + assert len(dap.obs_noise) == len( + dap.obs_attributes + ), "obs_noise length has to be equal to obs_attributes len" - obs_covar_c = np.zeros((len(dap.da_times),rloc.cattr_len)) - obs_covar_n = np.zeros((len(dap.da_times),rloc.nattr_len)) + obs_covar_c = np.zeros((len(dap.da_times), rloc.cattr_len)) + obs_covar_n = np.zeros((len(dap.da_times), rloc.nattr_len)) obs_noisy = copy.deepcopy(obs) - std_dev = np.zeros((obs.shape[0],len(dap.obs_noise_seeds))) + std_dev = np.zeros((obs.shape[0], len(dap.obs_noise_seeds))) # define inner domain in 2D VS - i2 = (slice(elem.igx,-elem.igx),slice(elem.igy,-elem.igy)) + i2 = (slice(elem.igx, -elem.igx), slice(elem.igy, -elem.igy)) for tt, obs_t in enumerate(obs): attr_cnt = 0 for key, value in obs_t.items(): - # Get inner 2D domain of sparse obs field value = value[i2] mask_at_t = mask[tt][key][i2] - value = np.ma.array(value,mask=mask_at_t) + value = np.ma.array(value, mask=mask_at_t) seed = dap.obs_noise_seeds[attr_cnt] np.random.seed(seed) - - if dap.noise_type == 'AmpCov': - field_var = (dap.obs_noise[key] * np.abs(value.max() - value.min())) + + if dap.noise_type == "AmpCov": + field_var = dap.obs_noise[key] * np.abs(value.max() - value.min()) field_sd = field_var**0.5 - elif dap.noise_type == 'VarCov': - field_var = (dap.obs_noise[key] * ((value - value.mean())**2).mean()) + elif dap.noise_type == "VarCov": + field_var = ( + dap.obs_noise[key] * ((value - value.mean()) ** 2).mean() + ) field_sd = field_var**0.5 - elif dap.noise_type == 'FixCov': + elif dap.noise_type == "FixCov": field_sd = 1.0 * dap.obs_noise[key] # here, we take the fraction defined by obs_noise multiplied by the maximum value of the observation as the standard deviation of the measurement noise. - - std_dev[tt,attr_cnt] = field_sd + + std_dev[tt, attr_cnt] = field_sd attr_cnt += 1 # var = std_dev**2 - if dap.noise_type == 'VarCov' or dap.noise_type == 'AmpCov': - sd = std_dev.mean(axis=0,keepdims=True) - std_dev[:,...] = sd + if dap.noise_type == "VarCov" or dap.noise_type == "AmpCov": + sd = std_dev.mean(axis=0, keepdims=True) + std_dev[:, ...] = sd for tt, obs_t in enumerate(obs): ccnt, ncnt, attr_cnt = 0, 0, 0 for key, value in obs_t.items(): shp = value.shape - sd = std_dev[tt,attr_cnt] - + sd = std_dev[tt, attr_cnt] + # print(tt, key, sd**2, np.abs(value.max()-value.min())) # generate gaussian noise for observations. @@ -451,13 +495,13 @@ def obs_noiser(obs,mask,dap,rloc,elem): if var == 0: var = -np.inf if key in rloc.ca: - obs_covar_c[tt,ccnt] = var + obs_covar_c[tt, ccnt] = var ccnt += 1 else: - obs_covar_n[tt,ncnt] = var + obs_covar_n[tt, ncnt] = var ncnt += 1 attr_cnt += 1 - + obs_covar = [obs_covar_c, obs_covar_n] return obs_noisy, obs_covar diff --git a/src/data_assimilation/__init__.py b/src/pybella/flow_solver/__init__.py similarity index 100% rename from src/data_assimilation/__init__.py rename to src/pybella/flow_solver/__init__.py diff --git a/src/dycore/__init__.py b/src/pybella/flow_solver/discretisation/__init__.py similarity index 100% rename from src/dycore/__init__.py rename to src/pybella/flow_solver/discretisation/__init__.py diff --git a/src/dycore/discretisation/grid.py b/src/pybella/flow_solver/discretisation/grid.py similarity index 83% rename from src/dycore/discretisation/grid.py rename to src/pybella/flow_solver/discretisation/grid.py index 85965bbc..b95bc7a1 100644 --- a/src/dycore/discretisation/grid.py +++ b/src/pybella/flow_solver/discretisation/grid.py @@ -1,6 +1,7 @@ import numpy as np -import dycore.utils.options as opts +from ..utils import options as opts + def grid_init(ud): """ @@ -29,20 +30,22 @@ def grid_init(ud): z0 = ud.zmin z1 = ud.zmax - grid = Grid(inx,iny,inz,x0,x1,y0,y1,z0,z1) + grid = Grid(inx, iny, inz, x0, x1, y0, y1, z0, z1) elem = ElemSpaceDiscr(grid, ud) node = NodeSpaceDiscr(grid, ud) return elem, node + class Grid(object): # def __init__(self, inx,iny,inz,x0,x1,y0,y1,z0,z1,left,right,bottom,top,back,front): """ Base grid class, defines the extent of grid and the grid spacing. """ - def __init__(self, inx,iny,inz,x0,x1,y0,y1,z0,z1): + + def __init__(self, inx, iny, inz, x0, x1, y0, y1, z0, z1): """ Parameters ---------- @@ -63,8 +66,8 @@ def __init__(self, inx,iny,inz,x0,x1,y0,y1,z0,z1): z0 : float Minimum extent in the z direction z1 : float - Maximum extent in the z direction - + Maximum extent in the z direction + """ assert inx > 1 assert iny >= 1 @@ -76,15 +79,15 @@ def __init__(self, inx,iny,inz,x0,x1,y0,y1,z0,z1): self.ndim = 2 else: self.ndim = 1 - + self.inx = inx self.iny = iny self.inz = inz - self.dx = (x1 - x0) / (inx - 1.) - self.dy = (y1 - y0) / (iny - 1.) if iny > 1 else 0.0 - self.dz = (z1 - z0) / (inz - 1.) if inz > 1 else 0.0 - + self.dx = (x1 - x0) / (inx - 1.0) + self.dy = (y1 - y0) / (iny - 1.0) if iny > 1 else 0.0 + self.dz = (z1 - z0) / (inz - 1.0) if inz > 1 else 0.0 + assert self.dx > 0.0 assert self.dy >= 0.0 assert self.dz >= 0.0 @@ -107,7 +110,10 @@ def __init__(self, inx,iny,inz,x0,x1,y0,y1,z0,z1): # self.back = back # self.front = front + big = 1.0 + + class SpaceDiscr(object): """ For a given grid extent and number of grid-points, this class returns an equidistant discretised grid. @@ -119,7 +125,7 @@ class SpaceDiscr(object): stride = np.zeros((3)) dxyz = np.zeros((3)) - def __init__(self,g): + def __init__(self, g): """ Parameters ---------- @@ -149,9 +155,9 @@ def __init__(self,g): self.nc = self.icx * self.icy * self.icz self.iisc = (self.iicx, self.iicy, self.iicz) - self.isc = (self.iicx+self.igx, self.iicy+self.igy, self.iicz+self.igz) + self.isc = (self.iicx + self.igx, self.iicy + self.igy, self.iicz + self.igz) self.sc = (self.icx, self.icy, self.icz) - self.igs = [self.igx,self.igy,self.igz] + self.igs = [self.igx, self.igy, self.igz] self.ifx = self.icx + 1 self.ify = self.icy + 1 if g.iny > 1 else 0 @@ -161,16 +167,16 @@ def __init__(self,g): self.nfy = self.icx * self.ify * self.icz self.nfz = self.icx * self.icy * self.ifz - self.sfx = (self.icy , self.icz , self.ifx) - self.sfy = (self.icz , self.icx , self.ify) - self.sfz = (self.icx , self.icy , self.ifz) + self.sfx = (self.icy, self.icz, self.ifx) + self.sfy = (self.icz, self.icx, self.ify) + self.sfz = (self.icx, self.icy, self.ifz) self.nf = self.nfx + self.nfy + self.nfz self.dx = g.dx self.dy = g.dy if self.icy > 1 else big self.dz = g.dz if self.icz > 1 else big - + self.dxyz[0] = self.dx self.dxyz[1] = self.dy self.dxyz[2] = self.dz @@ -178,16 +184,15 @@ def __init__(self,g): assert self.dx > 0.0 assert self.dy > 0.0 assert self.dz > 0.0 - - i1 = np.empty(self.ndim, dtype='object') - i2 = np.empty(self.ndim, dtype='object') + i1 = np.empty(self.ndim, dtype="object") + i2 = np.empty(self.ndim, dtype="object") for dim in range(self.ndim): - i1[dim] = slice(1,-1) - i2[dim] = slice(self.igs[dim],-self.igs[dim]) + i1[dim] = slice(1, -1) + i2[dim] = slice(self.igs[dim], -self.igs[dim]) self.i1 = tuple(i1) self.i2 = tuple(i2) - + class ElemSpaceDiscr(SpaceDiscr): """ @@ -195,7 +200,7 @@ class ElemSpaceDiscr(SpaceDiscr): """ - def __init__(self,g, ud): + def __init__(self, g, ud): """ Parameters ---------- @@ -225,10 +230,10 @@ def get_p_indim(self, ud): p_isc = [] pp1_isc = [] - eindim = np.empty((ndim),dtype='object') + eindim = np.empty((ndim), dtype="object") for dim in range(ndim): is_periodic = ud.bdry_type[dim] == opts.BdryType.PERIODIC - eindim[dim] = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic-1) + eindim[dim] = slice(igs[dim] - is_periodic, -igs[dim] + is_periodic - 1) p_isc.append(self.isc[dim] + 2 * is_periodic) pp1_isc.append(self.isc[dim] + 2 * is_periodic + 2) @@ -236,14 +241,15 @@ def get_p_indim(self, ud): self.periodic_indim = tuple(eindim) self.p_isc = tuple(p_isc) self.pp1_isc = tuple(pp1_isc) - - + + class NodeSpaceDiscr(SpaceDiscr): """ Inherits the class :class:`discretization.kgrid.SpaceDiscr`. For a given grid extent and number of grid-points, this class returns an equidistant discretised node-based grid. """ - def __init__(self,g, ud): + + def __init__(self, g, ud): """ Parameters ---------- @@ -268,9 +274,9 @@ def __init__(self,g, ud): self.y = y0 + self.dy * np.arange(self.icy) self.z = z0 + self.dz * np.arange(self.icz) - self.iisc = (self.iicx , self.iicy , self.iicz) - self.isc = (self.iicx+self.igx , self.iicy+self.igy , self.iicz+self.igz) - self.sc = (self.icx , self.icy , self.icz) + self.iisc = (self.iicx, self.iicy, self.iicz) + self.isc = (self.iicx + self.igx, self.iicy + self.igy, self.iicz + self.igz) + self.sc = (self.icx, self.icy, self.icz) self.get_p_indim(ud) @@ -285,14 +291,14 @@ def get_p_indim(self, ud): p_isc = [] pp1_isc = [] - nindim = np.empty((ndim),dtype='object') + nindim = np.empty((ndim), dtype="object") for dim in range(ndim): is_periodic = ud.bdry_type[dim] == opts.BdryType.PERIODIC - nindim[dim] = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic) - + nindim[dim] = slice(igs[dim] - is_periodic, -igs[dim] + is_periodic) + p_isc.append(self.isc[dim] + 2 * is_periodic) pp1_isc.append(self.isc[dim] + 2 * is_periodic + 2) self.periodic_indim = tuple(nindim) self.p_isc = tuple(p_isc) - self.pp1_isc = tuple(pp1_isc) \ No newline at end of file + self.pp1_isc = tuple(pp1_isc) diff --git a/src/pybella/flow_solver/discretisation/time_update.py b/src/pybella/flow_solver/discretisation/time_update.py new file mode 100644 index 00000000..e0f17a85 --- /dev/null +++ b/src/pybella/flow_solver/discretisation/time_update.py @@ -0,0 +1,332 @@ +import copy +import logging + +import numpy as np + +# dependencies of the flow solver subpackage +from . import grid as dis_grid +from ..utils import boundary as bdry, options as opts +from ..physics.gas_dynamics import ( + numerical_flux as gd_flux, + eos as gd_eos, + cfl as gd_cfl, +) +from ..physics.gas_dynamics import explicit as gd_explicit +from ..physics.low_mach import second_projection as lm_sp + +# for blending module +from ...interfaces.dynamics_blending import schemes +from ...interfaces.time_stepper import prestep + +def do( + sst, + mem, + tout, + bld=None, + writer=None, + debug_writer=None, +): + """ + For more details, refer to the write-up :ref:`time-stepping`. + + Does a time-step for the atmospheric solver. + + Parameters + ---------- + Sol : :class:`management.variable.Vars` + Solution data container. + flux : :class:`management.variable.States` + Data container for the fluxes. + mpv : :class:`physics.low_mach.mpv.MPV` + Variables relating to the elliptic solver. + t : float + Current time + tout : float + Next output time + ud : :class:`inputs.user_data.UserDataInit` + Data container for the initial conditions + elem : :class:`discretization.kgrid.ElemSpaceDiscr` + Cells grid. + node : :class:`discretization.kgrid.NodeSpaceDiscr` + Nodes grid. + step : int + Current step. + th : :class:`physics.gas_dynamics.thermodynamic.init` + Thermodynamic variables of the system + bld : :class:`data_assimilation.blending.Blend()` + Blending class used to initalise interface blending methods. + writer : :class:`management.io.io`, optional + `default == None`. If given, output after each time-step will be written in the hdf5 format. + debug : boolean, optional + `default == False`. If `True`, then writer will output `Sol`: + 1. before flux calculation + 2. before advection routine + 3. after advection routine + 4. after explicit solver + 5. after implicit solver + + during both the half-step for the prediction of advective flux and the full-step. + + Returns + ------- + list + A list of `[Sol,flux,mpv,[window_step,step]]` data containers at time `tout`. + """ + ud = sst.ud + elem, node, Sol, flux, mpv, th, time = mem + + window_step = time.window_step + step = time.step + t = time.t + swe_to_lake = False + + while (t < tout) and (step < ud.stepmax): + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + + label = "%.3d" % step + + if step == 0 and writer != None: + writer.write_all(mem, str(label) + "_ic") + + dt, cfl, cfl_ac = gd_cfl.dynamic_timestep(Sol, t, tout, elem, ud, th, step) + + dt = prestep.apply_modifcations(dt, ud, step) + + ###################################################### + # Blending : Do blending before timestep + ###################################################### + swe_to_lake, Sol, mpv, t = schemes.prepare_blending( + sst, + mem, + bld, + label, + writer, + step, + window_step, + t, + dt, + swe_to_lake, + debug_writer, + ) + + ud.is_nonhydrostatic = gd_eos.is_nonhydrostatic(ud, window_step) + ud.nonhydrostasy = gd_eos.nonhydrostasy(ud, t, window_step) + + if ud.continuous_blending or ud.initial_blending: + logging.info(f"step = {step}, window_step = {window_step}") + + logging.info( + f""" + ------- + is_compressible = {ud.is_compressible}, is_nonhydrostatic = {ud.is_nonhydrostatic} + compressibility = {ud.compressibility:.3f}, nonhydrostasy = {ud.nonhydrostasy:.3f} + ------- + """ + ) + + Sol0 = copy.deepcopy(Sol) + + debug_writer.write(f"{label}_before_flux") + + gd_flux.recompute_advective_fluxes(flux, Sol) + + debug_writer.populate_flux_components(f"{label}_before_advect", flux, elem) + debug_writer.write(f"{label}_before_advect") + + if ud.do_advection: + gd_explicit.advect_rk( + Sol, + flux, + 0.5 * dt, + elem, + step % 2, + ud, + th, + mpv, + node, + str(label) + "_half", + writer, + ) + + debug_writer.write(f"{label}_after_advect") + debug_writer.populate(f"{label}_after_full_step", "p2_nodes", mpv.p2_nodes) + + mpv.p2_nodes0[...] = mpv.p2_nodes + + lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) + + debug_writer.write(f"{label}_after_ebnaexp") + + Sol0_increment = Sol0 if ud.is_compressible == 0 else None + + lm_sp.euler_backward_non_advective_impl_part( + Sol, + mpv, + elem, + node, + ud, + th, + t, + 0.5 * dt, + 1.0, + Sol0=Sol0_increment, + label=f"{label}_after_ebnaimp", + writer=writer, + ) + + if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: + # top rayleight damping + bdry.rayleigh_damping(Sol, mpv, ud, elem, node) + + bdry.apply_rayleigh_forcing( + Sol, + mpv, + ud, + elem, + node, + t, + step, + dt, + th, + bdry, + ) + + debug_writer.write(f"{label}_after_ebnaimp") + + gd_flux.recompute_advective_fluxes(flux, Sol) + + debug_writer.populate_flux_components(f"{label}_after_half_step", flux, elem) + debug_writer.write(f"{label}_after_half_step") + + Sol_half_new = copy.deepcopy(Sol) + mpv_half_new = copy.deepcopy(mpv) + mpv.p2_nodes_half = np.copy(mpv.p2_nodes) + + if ud.is_nonhydrostatic == 0 or ( + ud.is_compressible == 1 and ud.is_nonhydrostatic == 1 + ): + mpv.p2_nodes[...] = mpv.p2_nodes0 + + Sol = copy.deepcopy(Sol0) + + lm_sp.euler_forward_non_advective( + Sol, + mpv, + elem, + node, + 0.5 * dt, + ud, + th, + writer=writer, + label=str(label) + "_after_efna", + ) + + debug_writer.write(f"{label}_after_efna") + + if ud.do_advection: + gd_explicit.advect( + Sol, + flux, + dt, + elem, + step % 2, + ud, + th, + mpv, + node, + str(label) + "_full", + writer, + ) + + debug_writer.write(f"{label}_after_full_advect") + + lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) + + debug_writer.write(f"{label}_after_full_ebnaexp") + + lm_sp.euler_backward_non_advective_impl_part( + Sol, + mpv, + elem, + node, + ud, + th, + t, + 0.5 * dt, + 2.0, + writer=writer, + label=str(label) + "_after_full_step", + ) + + if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: + # top rayleight damping + bdry.rayleigh_damping(Sol, mpv, ud, elem, node) + + # bottom rayleigh forcing + bdry.apply_rayleigh_forcing( + Sol, + mpv, + ud, + elem, + node, + t, + step, + dt, + th, + bdry, + half=False, + Sol_half_new=Sol_half_new, + mpv_half_new=mpv_half_new, + ) + + ###################################################### + # Blending : Do blending after timestep + ###################################################### + Sol, mpv = schemes.blending_after_timestep( + Sol, + flux, + mpv, + bld, + elem, + node, + th, + ud, + label, + writer, + step, + window_step, + t, + dt, + swe_to_lake, + debug_writer, + ) + + mem.sol = Sol + mem.flux = flux + mem.mpv = mpv + + if writer != None: + writer.time = t + writer.write_all(mem, str(label) + "_after_full_step") + + logging.info( + "###############################################################################################" + ) + logging.info( + "step %i done, t = %.12f, dt = %.12f, CFL = %.8f, CFL_ac = %.8f" + % (step, t, dt, cfl, cfl_ac) + ) + logging.info( + "###############################################################################################" + ) + + t += dt + step += 1 + window_step += 1 + + mem.time.t = t + mem.time.step = step + mem.time.window_step = window_step + + return mem + # return [Sol, flux, mpv, [window_step, step]] diff --git a/src/dycore/discretisation/__init__.py b/src/pybella/flow_solver/physics/__init__.py similarity index 100% rename from src/dycore/discretisation/__init__.py rename to src/pybella/flow_solver/physics/__init__.py diff --git a/src/dycore/physics/__init__.py b/src/pybella/flow_solver/physics/gas_dynamics/__init__.py similarity index 100% rename from src/dycore/physics/__init__.py rename to src/pybella/flow_solver/physics/gas_dynamics/__init__.py diff --git a/src/dycore/physics/gas_dynamics/cfl.py b/src/pybella/flow_solver/physics/gas_dynamics/cfl.py similarity index 71% rename from src/dycore/physics/gas_dynamics/cfl.py rename to src/pybella/flow_solver/physics/gas_dynamics/cfl.py index 01e70ee2..1badfcf4 100644 --- a/src/dycore/physics/gas_dynamics/cfl.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/cfl.py @@ -2,6 +2,7 @@ machine_epsilon = np.finfo(float).eps + def dynamic_timestep(Sol, time, time_output, elem, ud, th, step): global machine_epsilon @@ -28,21 +29,21 @@ def dynamic_timestep(Sol, time, time_output, elem, ud, th, step): v_max = max(v.max(), v_max) w_max = max(w.max(), w_max) - upc_max = max((u+c).max(), upc_max) - vpc_max = max((v+c).max(), vpc_max) - wpc_max = max((w+c).max(), wpc_max) + upc_max = max((u + c).max(), upc_max) + vpc_max = max((v + c).max(), vpc_max) + wpc_max = max((w + c).max(), wpc_max) - if (ud.acoustic_timestep == 1): + if ud.acoustic_timestep == 1: dtx = CFL * elem.dx / upc_max dty = CFL * elem.dy / vpc_max dtz = CFL * elem.dz / wpc_max dt_cfl = min(min(dtx, dty), dtz) - dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.) * (ud.dtfixed - ud.dtfixed0)) + dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0)) # if (2.0*dt > time_output - time): # dt = 0.5 * (time_output - time) + machine_epsilon - if (dt > (time_output - time)): + if dt > (time_output - time): dt = (time_output - time) + machine_epsilon return dt @@ -57,8 +58,8 @@ def dynamic_timestep(Sol, time, time_output, elem, ud, th, step): dt_cfl = min(min(dtx, dty), dtz) # dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.) * (ud.dtfixed - ud.dtfixed0)) if step >= 0: - dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.) * (ud.dtfixed - ud.dtfixed0)) - dt *= min(float(step+1), 1.0) + dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0)) + dt *= min(float(step + 1), 1.0) # else: # dt = min(dt_cfl, ud.dtfixed0 + 1. * (ud.dtfixed - ud.dtfixed0)) @@ -66,10 +67,12 @@ def dynamic_timestep(Sol, time, time_output, elem, ud, th, step): # f.write(str(dt_cfl) + "\n") # f.close() - if (dt > (time_output - time)): - dt = (time_output - time) #+ machine_epsilon + if dt > (time_output - time): + dt = time_output - time # + machine_epsilon cfl = CFL * dt / dt_cfl - cfl_ac = max(dt * upc_max / elem.dx, dt * vpc_max / elem.dy, dt * wpc_max / elem.dz) + cfl_ac = max( + dt * upc_max / elem.dx, dt * vpc_max / elem.dy, dt * wpc_max / elem.dz + ) - return dt, cfl, cfl_ac \ No newline at end of file + return dt, cfl, cfl_ac diff --git a/src/dycore/physics/gas_dynamics/eos.py b/src/pybella/flow_solver/physics/gas_dynamics/eos.py similarity index 60% rename from src/dycore/physics/gas_dynamics/eos.py rename to src/pybella/flow_solver/physics/gas_dynamics/eos.py index 8ee24444..f7c00226 100644 --- a/src/dycore/physics/gas_dynamics/eos.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/eos.py @@ -1,8 +1,9 @@ import numpy as np -import dycore.utils.boundary as bdry +from ...utils import boundary as bdry -def nonhydrostasy(ud,t,step): + +def nonhydrostasy(ud, t, step): if step >= 0: if ud.is_nonhydrostatic == 0: return 0.0 @@ -12,11 +13,14 @@ def nonhydrostasy(ud,t,step): current_transition_step = step - ud.no_of_hy_initial # print("current_transition_step =", step - ud.no_of_hy_initial) # print(np.linspace(0.0,1.0,ud.no_of_hy_transition+2)[1:-1][current_transition_step]) - return np.linspace(0.0,1.0,ud.no_of_hy_transition+2)[1:-1][current_transition_step] + return np.linspace(0.0, 1.0, ud.no_of_hy_transition + 2)[1:-1][ + current_transition_step + ] else: return float(ud.is_nonhydrostatic) -def compressibility(ud,t,step): + +def compressibility(ud, t, step): if step >= 0: if ud.is_compressible == 0: if step < ud.no_of_pi_initial: @@ -25,14 +29,17 @@ def compressibility(ud,t,step): return 0.0 else: current_transition_step = step - ud.no_of_pi_initial - return np.linspace(0.0,1.0,ud.no_of_pi_transition+2)[1:-1][current_transition_step] + return np.linspace(0.0, 1.0, ud.no_of_pi_transition + 2)[1:-1][ + current_transition_step + ] elif ud.is_compressible == 1: return 1.0 # elif ud.is_compressible == 0: else: return ud.compressibility -def is_compressible(ud,step): + +def is_compressible(ud, step): if step >= 0: if ud.continuous_blending == True: if step < ud.no_of_pi_initial: @@ -45,38 +52,33 @@ def is_compressible(ud,step): return ud.is_compressible else: return ud.is_compressible - -def is_nonhydrostatic(ud,step): - # print("is_nonhydrostatic", ud.is_nonhydrostatic) - # print("no_of_nhy_initial:", ud.no_of_hy_initial) - # print("no_of_nhy_transition:", ud.no_of_hy_transition) - if step >= 0: - if ud.continuous_blending == True: - if step < ud.no_of_hy_initial: - return 0 - elif step < (ud.no_of_hy_initial + ud.no_of_hy_transition): - return -1 - else: - return 1 - else: - return ud.is_nonhydrostatic - else: + +def is_nonhydrostatic(ud, step): + if step < 0 or not ud.continuous_blending: return ud.is_nonhydrostatic + if step < ud.no_of_hy_initial: + return 0 + elif step < (ud.no_of_hy_initial + ud.no_of_hy_transition): + return -1 + else: + return 1 + -def rhoe(rho,u,v,w,p,ud,th): +def rhoe(rho, u, v, w, p, ud, th): Msq = ud.compressibility * ud.Msq gm1inv = th.gm1inv return p * gm1inv + 0.5 * Msq * rho * (u**2 + v**2 + w**2) + def synchronise_variables(mpv, Sol, elem, node, ud, th): scale_factor = 1.0 / ud.Msq - if (ud.is_compressible): - p2bg = mpv.HydroState.p20[0,:].reshape(1,-1) - p2bg = np.repeat(p2bg, elem.icx,axis=0) + if ud.is_compressible: + p2bg = mpv.HydroState.p20[0, :].reshape(1, -1) + p2bg = np.repeat(p2bg, elem.icx, axis=0) mpv.p2_cells = scale_factor * Sol.rhoY**th.gm1 - p2bg - bdry.set_ghostcells_p2(mpv.p2_cells, elem, ud) \ No newline at end of file + bdry.set_ghostcells_p2(mpv.p2_cells, elem, ud) diff --git a/src/dycore/physics/gas_dynamics/explicit.py b/src/pybella/flow_solver/physics/gas_dynamics/explicit.py similarity index 85% rename from src/dycore/physics/gas_dynamics/explicit.py rename to src/pybella/flow_solver/physics/gas_dynamics/explicit.py index 283985a6..2d205416 100644 --- a/src/dycore/physics/gas_dynamics/explicit.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/explicit.py @@ -1,14 +1,15 @@ -import dycore.utils.boundary as bdry -import dycore.physics.gas_dynamics.recovery as gd_recovery -import dycore.physics.gas_dynamics.numerical_flux as gd_flux +from ...utils import boundary as bdry +from . import recovery as gd_recovery +from . import numerical_flux as gd_flux -def advect(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None): + +def advect(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer=None): """ Function that runs the advection routine with Strang-splitting. This function updates the `Sol` solution container with the advected solution in-place. Parameters ---------- - Sol : :py:class:`management.variable.Vars` + Sol : :py:class:`management.variable.Vars` Solution data container. flux : :py:class:`management.variable.States` Fluxes data container @@ -38,48 +39,66 @@ def advect(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None): # Sol.rhoX -= Sol.rho * mpv.HydroState.S0 - if (odd): + if odd: for split in range(ndim): lmbda = time_step / elem.dxyz[split] Sol.flip_forward() - if elem.iisc[split] > 1: - explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv) + if elem.iisc[split] > 1: + explicit_step_and_flux( + Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv + ) else: for i_split in range(ndim): split = elem.ndim - 1 - i_split lmbda = time_step / elem.dxyz[split] if elem.iisc[split] > 1: - explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, [writer,node,label]) + explicit_step_and_flux( + Sol, + flux[split], + lmbda, + elem, + split, + stage, + ud, + th, + mpv, + [writer, node, label], + ) Sol.flip_backward() stage = 1 - if (odd): + if odd: for i_split in range(ndim): split = elem.ndim - 1 - i_split lmbda = time_step / elem.dxyz[split] if elem.iisc[split] > 1: - explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv) + explicit_step_and_flux( + Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv + ) Sol.flip_backward() else: for split in range(ndim): lmbda = time_step / elem.dxyz[split] Sol.flip_forward() if elem.iisc[split] > 1: - explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv) - + explicit_step_and_flux( + Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv + ) + # Sol.rhoX += Sol.rho * mpv.HydroState.S0 bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - -def explicit_step_and_flux(Sol, flux, lmbda, elem, split_step, stage, ud, th, mpv, writer=None, tag=None): +def explicit_step_and_flux( + Sol, flux, lmbda, elem, split_step, stage, ud, th, mpv, writer=None, tag=None +): """ For each advection substep, solve the advection problem. For more details, see :ref:`advection_routine`. This function updates the solution `Sol` container in-place if a Strang-splitting is used, or returns the `flux` data container if a Runge-Kutta method is used. Parameters ---------- - Sol : :py:class:`management.variable.Vars` + Sol : :py:class:`management.variable.Vars` Solution data container. flux : :py:class:`management.variable.States` Fluxes data container @@ -145,35 +164,35 @@ def explicit_step_and_flux(Sol, flux, lmbda, elem, split_step, stage, ud, th, mp # skipped check_flux_bcs for now; first debug other functions # check_flux_bcs(Lefts, Rights, elem, split_step, ud) - flux = gd_flux.hll_solver(flux,Lefts,Rights,Sol, lmbda, ud, th) + flux = gd_flux.hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th) ndim = elem.ndim left_idx, right_idx = [slice(None)] * ndim, [slice(None)] * ndim - right_idx[-1] = slice(1,None) - left_idx[-1] = slice(0,-1) + right_idx[-1] = slice(1, None) + left_idx[-1] = slice(0, -1) left_idx, right_idx = tuple(left_idx), tuple(right_idx) - if tag != 'rk': + if tag != "rk": Sol.rho += lmbda * (flux.rho[left_idx] - flux.rho[right_idx]) Sol.rhou += lmbda * (flux.rhou[left_idx] - flux.rhou[right_idx]) Sol.rhov += lmbda * (flux.rhov[left_idx] - flux.rhov[right_idx]) Sol.rhow += lmbda * (flux.rhow[left_idx] - flux.rhow[right_idx]) Sol.rhoX += lmbda * (flux.rhoX[left_idx] - flux.rhoX[right_idx]) Sol.rhoY += lmbda * (flux.rhoY[left_idx] - flux.rhoY[right_idx]) - + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=split_step) - if tag == 'rk': + if tag == "rk": return flux -def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None): +def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer=None): """ Function that runs the advection routine with a first-order Runge-Kutta update. This function updates the `Sol` solution container with the advected solution in-place. Parameters ---------- - Sol : :py:class:`management.variable.Vars` + Sol : :py:class:`management.variable.Vars` Solution data container. flux : :py:class:`management.variable.States` Fluxes data container @@ -195,7 +214,7 @@ def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None) Tag label for the output array writer : :py:class:`management.io.io`, optional Writer class for I/O operations, by default None - + Attention --------- This function is not usually called unless commented out in the :py:meth:`management.data.time_update` routine. @@ -210,13 +229,15 @@ def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None) for split in range(ndim): lmbda = time_step / elem.dxyz[split] Sol.flip_forward() - if elem.iisc[split] > 1: - flux[split] = explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, tag='rk') + if elem.iisc[split] > 1: + flux[split] = explicit_step_and_flux( + Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, tag="rk" + ) ndim = elem.ndim left_idx, right_idx = [slice(None)] * ndim, [slice(None)] * ndim - right_idx[-1] = slice(1,None) - left_idx[-1] = slice(0,-1) + right_idx[-1] = slice(1, None) + left_idx[-1] = slice(0, -1) left_idx, right_idx = tuple(left_idx), tuple(right_idx) for dim in range(ndim): @@ -229,9 +250,9 @@ def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None) Sol.rhoX += lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) Sol.rhoY += lmbda * (flux[dim].rhoY[left_idx] - flux[dim].rhoY[right_idx]) - if dim == 1: # vertical axis + if dim == 1: # vertical axis updt = lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) - setattr(Sol, 'pwchi', updt) + setattr(Sol, "pwchi", updt) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) @@ -239,7 +260,7 @@ def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None) # for split in range(ndim): # lmbda = dt / elem.dxyz[split] # Sol.flip_forward() - # if elem.iisc[split] > 1: + # if elem.iisc[split] > 1: # flux[split] = explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, tag='rk') # Sol = deepcopy(Sol0) @@ -253,5 +274,5 @@ def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer = None) # Sol.rhow += lmbda * (flux[dim].rhow[left_idx] - flux[dim].rhow[right_idx]) # Sol.rhoX += lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) # Sol.rhoY += lmbda * (flux[dim].rhoY[left_idx] - flux[dim].rhoY[right_idx]) - + # set_explicit_boundary_data(Sol, elem, ud, th, mpv) diff --git a/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py b/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py new file mode 100644 index 00000000..3d5098a6 --- /dev/null +++ b/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- +import numpy as np +import scipy as sp + + +def recompute_advective_fluxes(flux, Sol, *args, **kwargs): + """ + Recompute the advective fluxes at the cell interfaces, i.e. the faces. This function updates the `flux` container in-place. + + Parameters + ---------- + flux : :py:class:`management.variable.States` + Data container for the fluxes at the cell interfaces. + Sol : :py:class:`management.variable.States` + Data container for the Solution. + + Attention + --------- + This function is a mess and requires cleaning up. + + """ + ndim = Sol.rho.ndim + inner_idx = tuple([slice(1, -1)] * ndim) + + if ndim == 2: + kernel_u = np.array([[0.5, 1.0, 0.5], [0.5, 1.0, 0.5]]) + kernel_v = kernel_u.T + elif ndim == 3: + kernel_u = np.array( + [[[1, 2, 1], [2, 4, 2], [1, 2, 1]], [[1, 2, 1], [2, 4, 2], [1, 2, 1]]] + ) + kernel_v = np.swapaxes(kernel_u, 1, 0) + kernel_w = np.swapaxes(kernel_u, 2, 0) + + rhoYw = Sol.rhoY * Sol.rhow / Sol.rho + flux[2].rhoY[inner_idx] = ( + sp.signal.fftconvolve(rhoYw, kernel_w, mode="valid") / kernel_w.sum() + ) + else: + assert 0, "Unsupported dimension in recompute_advective_flux" + + rhoYu = kwargs.get("u", Sol.rhoY * Sol.rhou / Sol.rho) + + flux[0].rhoY[inner_idx] = np.moveaxis( + sp.signal.fftconvolve(rhoYu, kernel_u, mode="valid") / kernel_u.sum(), 0, -1 + ) + + rhoYv = kwargs.get("v", Sol.rhoY * Sol.rhov / Sol.rho) + if ndim == 2: + flux[1].rhoY[inner_idx] = ( + sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum() + ) + elif ndim == 3: + flux[1].rhoY[inner_idx] = np.moveaxis( + sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum(), -1, 0 + ) + # flux[1].rhoY[...,-1] = 0. + + +def hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th): + """ + HLL solver for the Riemann problem. Chooses the advected quantities from `Lefts` or `Rights` based on the direction given by `flux`. + + Parameters + ---------- + flux : :py:class:`management.variable.States` + Data container for fluxes. + Lefts : :py:class:`management.variable.States` + Container for the quantities on the left of the cell interfaces. + Rights : :py:class:`management.variable.States` + Container for the quantities on the right of the cell interfaces. + Sol : :py:class:`management.variable.Vars` + Solution data container. + lmbda : float + :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep. + ud : :py:class:`inputs.user_data.UserDataInit` + Class container for the initial condition. + th : :py:class:`physics.gas_dynamics.thermodynamic.init` + Class container for the thermodynamical constants. + + Returns + ------- + :py:class:`management.variable.States` + `flux` data container with the solution of the Riemann problem. + """ + # flux: index 1 to end = Left[inner_idx]: index 0 to -1 = Right[inner_idx]: index 1 to end + + ndim = Sol.rho.ndim + left_idx, right_idx, remove_cols_idx = ( + [slice(None)] * ndim, + [slice(None)] * ndim, + [slice(None)] * ndim, + ) + + remove_cols_idx[-1] = slice(1, -1) + left_idx[-1] = slice(0, -1) + right_idx[-1] = slice(1, None) + + left_idx, right_idx, remove_cols_idx = ( + tuple(left_idx), + tuple(right_idx), + tuple(remove_cols_idx), + ) + + Lefts.primitives(th) + Rights.primitives(th) + + upwind = 0.5 * (1.0 + np.sign(flux.rhoY)) + upl = upwind[right_idx] + upr = 1.0 - upwind[left_idx] + + flux.rhou[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( + upl[left_idx] / Lefts.Y[left_idx] * Lefts.u[left_idx] + + upr[right_idx] / Rights.Y[right_idx] * Rights.u[right_idx] + ) + flux.rho[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( + upl[left_idx] / Lefts.Y[left_idx] * 1.0 + + upr[right_idx] / Rights.Y[right_idx] * 1.0 + ) + + flux.rhov[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( + upl[left_idx] / Lefts.Y[left_idx] * Lefts.v[left_idx] + + upr[right_idx] / Rights.Y[right_idx] * Rights.v[right_idx] + ) + flux.rhow[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( + upl[left_idx] / Lefts.Y[left_idx] * Lefts.w[left_idx] + + upr[right_idx] / Rights.Y[right_idx] * Rights.w[right_idx] + ) + flux.rhoX[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( + upl[left_idx] / Lefts.Y[left_idx] * Lefts.X[left_idx] + + upr[right_idx] / Rights.Y[right_idx] * Rights.X[right_idx] + ) + + return flux diff --git a/src/dycore/physics/gas_dynamics/recovery.py b/src/pybella/flow_solver/physics/gas_dynamics/recovery.py similarity index 69% rename from src/dycore/physics/gas_dynamics/recovery.py rename to src/pybella/flow_solver/physics/gas_dynamics/recovery.py index 30ee3db3..3bfb9030 100644 --- a/src/dycore/physics/gas_dynamics/recovery.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/recovery.py @@ -1,7 +1,7 @@ import numpy as np -import dycore.utils.options as opts -import dycore.utils.variable as var +from ...utils import options as opts, variable as var + def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): """ @@ -23,7 +23,7 @@ def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): Class container for the cell-grid. split_step : int Tracks the substep in the Strang-splitting. - tag : `None` or `rk` + tag : `None` or `rk` Default is `None` which uses a second-order Strang-splitting. `rk` toggles a first-order Runge-Kutta update for the advection scheme. Returns @@ -33,56 +33,78 @@ def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): """ gamm = th.gamm - - order_two = 1 # always 1 + + order_two = 1 # always 1 Sol.primitives(th) - if tag == 'rk': + if tag == "rk": lmbda = 0.0 ndim = elem.ndim - lefts_idx, rights_idx, inner_idx = [slice(None,)] * ndim, [slice(None,)] * ndim, [slice(1,-1)] * ndim - lefts_idx[-1] = slice(0,-1) - rights_idx[-1] = slice(1,None) - lefts_idx, rights_idx, inner_idx = tuple(lefts_idx), tuple(rights_idx), tuple(inner_idx) + lefts_idx, rights_idx, inner_idx = ( + [ + slice( + None, + ) + ] + * ndim, + [ + slice( + None, + ) + ] + * ndim, + [slice(1, -1)] * ndim, + ) + lefts_idx[-1] = slice(0, -1) + rights_idx[-1] = slice(1, None) + lefts_idx, rights_idx, inner_idx = ( + tuple(lefts_idx), + tuple(rights_idx), + tuple(inner_idx), + ) # inner_idx here are where the interface fluxes are calculated with non-zero values. face_inner_idx = inner_idx u = np.zeros_like(Sol.rhoY) - u[inner_idx] = 0.5 * (flux.rhoY[face_inner_idx][lefts_idx] + flux.rhoY[face_inner_idx][rights_idx]) / Sol.rhoY[inner_idx] + u[inner_idx] = ( + 0.5 + * (flux.rhoY[face_inner_idx][lefts_idx] + flux.rhoY[face_inner_idx][rights_idx]) + / Sol.rhoY[inner_idx] + ) shape = Sol.u.shape - Diffs = var.States(shape,ud) + Diffs = var.States(shape, ud) Ampls = var.Characters(shape) Lefts = var.States(shape, ud) Rights = var.States(shape, ud) - Diffs.u[...,:-1] = Sol.u[rights_idx] - Sol.u[lefts_idx] - Diffs.v[...,:-1] = Sol.v[rights_idx] - Sol.v[lefts_idx] - Diffs.w[...,:-1] = Sol.w[rights_idx] - Sol.w[lefts_idx] - Diffs.X[...,:-1] = Sol.X[rights_idx] - Sol.X[lefts_idx] - Diffs.Y[...,:-1] = 1.0 / Sol.Y[rights_idx] - 1.0 / Sol.Y[lefts_idx] + Diffs.u[..., :-1] = Sol.u[rights_idx] - Sol.u[lefts_idx] + Diffs.v[..., :-1] = Sol.v[rights_idx] - Sol.v[lefts_idx] + Diffs.w[..., :-1] = Sol.w[rights_idx] - Sol.w[lefts_idx] + Diffs.X[..., :-1] = Sol.X[rights_idx] - Sol.X[lefts_idx] + Diffs.Y[..., :-1] = 1.0 / Sol.Y[rights_idx] - 1.0 / Sol.Y[lefts_idx] Slopes = slopes(Sol, Diffs, ud, elem) - Ampls.u[...] = 0.5 * Slopes.u * (1. - lmbda * u) - Ampls.v[...] = 0.5 * Slopes.v * (1. - lmbda * u) - Ampls.w[...] = 0.5 * Slopes.w * (1. - lmbda * u) - Ampls.X[...] = 0.5 * Slopes.X * (1. - lmbda * u) - Ampls.Y[...] = 0.5 * Slopes.Y * (1. - lmbda * u) - + Ampls.u[...] = 0.5 * Slopes.u * (1.0 - lmbda * u) + Ampls.v[...] = 0.5 * Slopes.v * (1.0 - lmbda * u) + Ampls.w[...] = 0.5 * Slopes.w * (1.0 - lmbda * u) + Ampls.X[...] = 0.5 * Slopes.X * (1.0 - lmbda * u) + Ampls.Y[...] = 0.5 * Slopes.Y * (1.0 - lmbda * u) + Lefts.u[...] = Sol.u + order_two * Ampls.u Lefts.v[...] = Sol.v + order_two * Ampls.v Lefts.w[...] = Sol.w + order_two * Ampls.w Lefts.X[...] = Sol.X + order_two * Ampls.X Lefts.Y[...] = 1.0 / (1.0 / Sol.Y + order_two * Ampls.Y) - Ampls.u[...] = -0.5 * Slopes.u * (1. + lmbda * u) - Ampls.v[...] = -0.5 * Slopes.v * (1. + lmbda * u) - Ampls.w[...] = -0.5 * Slopes.w * (1. + lmbda * u) - Ampls.X[...] = -0.5 * Slopes.X * (1. + lmbda * u) - Ampls.Y[...] = -0.5 * Slopes.Y * (1. + lmbda * u) + Ampls.u[...] = -0.5 * Slopes.u * (1.0 + lmbda * u) + Ampls.v[...] = -0.5 * Slopes.v * (1.0 + lmbda * u) + Ampls.w[...] = -0.5 * Slopes.w * (1.0 + lmbda * u) + Ampls.X[...] = -0.5 * Slopes.X * (1.0 + lmbda * u) + Ampls.Y[...] = -0.5 * Slopes.Y * (1.0 + lmbda * u) Rights.u[...] = Sol.u + order_two * Ampls.u Rights.v[...] = Sol.v + order_two * Ampls.v @@ -90,14 +112,18 @@ def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): Rights.X[...] = Sol.X + order_two * Ampls.X Rights.Y[...] = 1.0 / (1.0 / Sol.Y + order_two * Ampls.Y) - vel = [Sol.u,Sol.v,Sol.w] + vel = [Sol.u, Sol.v, Sol.w] # Lefts.rhoY[lefts_idx] = Rights.rhoY[rights_idx] = 0.5 * (Sol.rhoY[lefts_idx] + Sol.rhoY[rights_idx]) \ # - order_two * 0.5 * lmbda * (Sol.u[rights_idx] * Sol.rhoY[rights_idx] - Sol.u[lefts_idx] * Sol.rhoY[lefts_idx]) - Lefts.rhoY[lefts_idx] = Rights.rhoY[rights_idx] = 0.5 * (Sol.rhoY[lefts_idx] + Sol.rhoY[rights_idx]) \ - - order_two * 0.5 * lmbda * (vel[split_step][rights_idx] * Sol.rhoY[rights_idx] - vel[split_step][lefts_idx] * Sol.rhoY[lefts_idx]) + Lefts.rhoY[lefts_idx] = Rights.rhoY[rights_idx] = 0.5 * ( + Sol.rhoY[lefts_idx] + Sol.rhoY[rights_idx] + ) - order_two * 0.5 * lmbda * ( + vel[split_step][rights_idx] * Sol.rhoY[rights_idx] + - vel[split_step][lefts_idx] * Sol.rhoY[lefts_idx] + ) - Lefts.p0[lefts_idx] = Rights.p0[rights_idx] = Lefts.rhoY[lefts_idx]**gamm + Lefts.p0[lefts_idx] = Rights.p0[rights_idx] = Lefts.rhoY[lefts_idx] ** gamm get_conservatives(Rights, ud, th) get_conservatives(Lefts, ud, th) @@ -105,6 +131,7 @@ def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): # return Lefts, Rights, u, Diffs, Ampls, Slopes return Lefts, Rights + def slopes(Sol, Diffs, ud, elem): """ Reconstruct the piecewise linear slopes in the cells from the piecewise constants in the cell and its neighbours. @@ -124,14 +151,22 @@ def slopes(Sol, Diffs, ud, elem): ------- :py:class:`management.variable.Characters` Reconstructed piecewise linear slopes in the cell. - """ + """ limiter_type_velocity = ud.limiter_type_velocity limiter_type_scalar = ud.limiter_type_scalars ndim = elem.ndim - lefts_idx, rights_idx = [slice(None,)] * ndim, [slice(None,)] * ndim - lefts_idx[-1] = slice(0,-1) - rights_idx[-1] = slice(1,None) + lefts_idx, rights_idx = [ + slice( + None, + ) + ] * ndim, [ + slice( + None, + ) + ] * ndim + lefts_idx[-1] = slice(0, -1) + rights_idx[-1] = slice(1, None) lefts_idx, rights_idx = tuple(lefts_idx), tuple(rights_idx) # amplitudes of the state differences: @@ -152,14 +187,15 @@ def slopes(Sol, Diffs, ud, elem): Slopes = var.Characters(Diffs.u.shape) - Slopes.u[...,1:-1] = limiters(limiter_type_velocity, aul, aur) - Slopes.v[...,1:-1] = limiters(limiter_type_velocity, avl, avr) - Slopes.w[...,1:-1] = limiters(limiter_type_velocity, awl, awr) - Slopes.X[...,1:-1] = limiters(limiter_type_scalar, aXl, aXr) - Slopes.Y[...,1:-1] = limiters(limiter_type_scalar, aYl, aYr) + Slopes.u[..., 1:-1] = limiters(limiter_type_velocity, aul, aur) + Slopes.v[..., 1:-1] = limiters(limiter_type_velocity, avl, avr) + Slopes.w[..., 1:-1] = limiters(limiter_type_velocity, awl, awr) + Slopes.X[..., 1:-1] = limiters(limiter_type_scalar, aXl, aXr) + Slopes.Y[..., 1:-1] = limiters(limiter_type_scalar, aYl, aYr) return Slopes + def limiters(limiter_type, al, ar): """ Applies the limiter type specified in the initial conditions to recovery the slope. @@ -187,6 +223,7 @@ def limiters(limiter_type, al, ar): if limiter_type == opts.LimiterType.NONE: return 0.5 * (al + ar) + def get_conservatives(U, ud, th): """ Get advected (conservative) quantities at the left and right of the cell interfaces. @@ -206,6 +243,6 @@ def get_conservatives(U, ud, th): U.rhow = U.w * U.rho U.rhoY = U.Y * U.rho U.rhoX = U.X * U.rho - + sgn = np.sign(U.rhoY) - p = sgn*np.abs(U.rhoY)**th.gamminv \ No newline at end of file + p = sgn * np.abs(U.rhoY) ** th.gamminv diff --git a/src/dycore/physics/gas_dynamics/thermodynamic.py b/src/pybella/flow_solver/physics/gas_dynamics/thermodynamics.py similarity index 91% rename from src/dycore/physics/gas_dynamics/thermodynamic.py rename to src/pybella/flow_solver/physics/gas_dynamics/thermodynamics.py index f4a330d7..bdf06637 100644 --- a/src/dycore/physics/gas_dynamics/thermodynamic.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/thermodynamics.py @@ -1,8 +1,9 @@ -class init(object): +class ThermodynamicalQuantities(object): """ Data container for thermodynamical quantities. """ + def __init__(self, ud): """ Parameters @@ -10,7 +11,7 @@ def __init__(self, ud): ud : :class:`inputs.user_data.UserDataInit` Data container for the initial conditions """ - + g = ud.gamm self.gamm = g self.gamminv = 1.0 / g @@ -18,4 +19,3 @@ def __init__(self, ud): self.gm1inv = 1.0 / (g - 1.0) self.Gamma = (g - 1.0) / g self.Gammainv = g / (g - 1.0) - diff --git a/src/pybella/flow_solver/physics/hydrostatics.py b/src/pybella/flow_solver/physics/hydrostatics.py new file mode 100644 index 00000000..64c720ba --- /dev/null +++ b/src/pybella/flow_solver/physics/hydrostatics.py @@ -0,0 +1,388 @@ +import numpy as np +import numba as nb + +from ..utils import boundary as bdry + + +def column(HydroState, HydroState_n, Y, Y_n, elem, node, th, ud): + Gamma = th.gm1 / th.gamm + gamm = th.gamm + gm1 = th.gm1 + Gamma_inv = 1.0 / Gamma + gm1_inv = 1.0 / gm1 + + icy = elem.icy + igy = elem.igy + + xc_idx = slice(0, -1) + yc_idx = slice(0, -1) + + c_idx = (xc_idx, yc_idx) + + rhoY0 = 1.0 + + g = ud.gravity_strength[1] + + p0 = rhoY0**gamm + pi0 = rhoY0**gm1 + HydroState_n.rho0[xc_idx, igy] = rhoY0 / Y_n[:, igy] + HydroState_n.rhoY0[xc_idx, igy] = rhoY0 + HydroState_n.Y0[xc_idx, igy] = Y_n[:, igy] + HydroState_n.S0[xc_idx, igy] = 1.0 / Y_n[:, igy] + HydroState_n.p0[xc_idx, igy] = p0 + HydroState_n.p20[xc_idx, igy] = pi0 / ud.Msq + + dys = np.array( + [-elem.dy] + [-elem.dy / 2] + [elem.dy / 2] + list(np.ones((icy - 3)) * elem.dy) + ) + S_p = 1.0 / Y[:, :] + S_m = np.zeros_like(S_p) + S_m[:, igy - 1 : igy + 1] = 1.0 / Y_n[:, igy].reshape(-1, 1) + S_m[:, 0] = 1.0 / Y[:, igy - 1] + S_m[:, igy + 1 :] = 1.0 / Y[:, igy:-1] + + S_integral_p = dys * 0.5 * (S_p + S_m) + S_integral_p[:, :igy] = np.cumsum(S_integral_p[:, :igy][:, ::-1], axis=1)[:, ::-1] + S_integral_p[:, igy:] = np.cumsum(S_integral_p[:, igy:], axis=1) + + pi_hydro = pi0 - Gamma * g * S_integral_p + p_hydro = pi_hydro**Gamma_inv + rhoY_hydro = pi_hydro**gm1_inv + + HydroState.rho0[c_idx] = rhoY_hydro * S_p + HydroState.p0[c_idx] = p_hydro + HydroState.p20[c_idx] = pi_hydro / ud.Msq + HydroState.S0[c_idx] = S_p + HydroState.S10[c_idx] = 0.0 + HydroState.Y0[c_idx] = 1.0 / S_p + HydroState.rhoY0[c_idx] = rhoY_hydro + + Sn_p = 1.0 / Y[:, :] + dys = np.ones((icy)) * elem.dy + dys[:igy] *= -1 + Sn_integral_p = dys * Sn_p + Sn_integral_p[:, :igy] = np.cumsum(Sn_integral_p[:, :igy][:, ::-1], axis=1)[:, ::-1] + Sn_integral_p[:, igy:] = np.cumsum(Sn_integral_p[:, igy:], axis=1) + + pi_hydro_n = pi0 - Gamma * g * Sn_integral_p + rhoY_hydro_n = pi_hydro_n**gm1_inv + + HydroState_n.rhoY0[xc_idx, :igy] = rhoY_hydro_n[:, :igy] + HydroState_n.Y0[xc_idx, :igy] = Y_n[0, :igy] + HydroState_n.S0[xc_idx, :igy] = 1.0 / Y_n[:, :igy] + HydroState_n.p0[xc_idx, :igy] = rhoY_hydro_n[:, :igy] ** th.gamm + HydroState_n.p20[xc_idx, :igy] = pi_hydro_n[:, :igy] / ud.Msq + + HydroState_n.rhoY0[xc_idx, igy + 1 :] = rhoY_hydro_n[:, igy:] + HydroState_n.Y0[xc_idx, igy + 1 :] = Y_n[0, igy:] + HydroState_n.S0[xc_idx, igy + 1 :] = 1.0 / Y_n[:, igy:] + HydroState_n.p0[xc_idx, igy + 1 :] = rhoY_hydro_n[:, igy:] ** th.gamm + HydroState_n.p20[xc_idx, igy + 1 :] = pi_hydro_n[:, igy:] / ud.Msq + +def integrated_state(mpv, elem, node, th, ud): + """ + Compute hydrostatic background state for atmospheric model. + Handles arbitrary stratification profiles and proper numerical integration. + """ + # Thermodynamic constants + Gamma = th.gm1 / th.gamm + gamm = th.gamm + gm1 = th.gm1 + Gamma_inv = 1.0 / Gamma + gm1_inv = 1.0 / gm1 + + # Grid parameters + icy = elem.icy + igy = elem.igy + + # Reference state at y=0 + rhoY0 = 1.0 + g = ud.gravity_strength[1] + p0 = rhoY0**gamm + pi0 = rhoY0**gm1 + + if g != 0.0: + ########################### + # Update cell hydrostates + ########################### + + # Define midpoint quadrature along vertical (y-axis) + dys = np.hstack(( + np.ones(igy-1) * -elem.dy, + [-elem.dy/2], + [elem.dy/2], + np.ones(icy-3) * elem.dy + )) + + # Cell centers and midpoints for integration + y_ps = elem.y + y_ms = np.hstack(( + elem.y[1:igy], + node.y[igy], + node.y[igy], + elem.y[igy:-1] + )) + + # Get inverse stratification at each point + S_ps = 1.0 / ud.stratification(y_ps) + S_ms = 1.0 / ud.stratification(y_ms) + + # Trapezoidal integration over inverse stratification + S_integral_p = 0.5 * dys * (S_ms + S_ps) + + # Cumulative integration (split at boundary igy) + S_integral_p[:igy] = np.cumsum(S_integral_p[:igy][::-1])[::-1] + S_integral_p[igy:] = np.cumsum(S_integral_p[igy:]) + + # Calculate hydrostatic fields + pi_hydro = pi0 - Gamma * g * S_integral_p + p_hydro = pi_hydro**Gamma_inv + rhoY_hydro = pi_hydro**gm1_inv + + # Update cell solutions + mpv.HydroState.rhoY0[:] = rhoY_hydro + mpv.HydroState.rho0[:] = rhoY_hydro * S_ps + mpv.HydroState.p0[:] = p_hydro + mpv.HydroState.p20[:] = pi_hydro / ud.Msq + mpv.HydroState.S0[:] = S_ps + mpv.HydroState.S10[:] = 0.0 + mpv.HydroState.Y0[:] = 1.0 / S_ps + + ############################ + # Update node hydrostates + ############################ + + # Bottom reference node (y=0) + mpv.HydroState_n.Y0[igy] = ud.stratification(0.0) + mpv.HydroState_n.rhoY0[igy] = rhoY0 + mpv.HydroState_n.rho0[igy] = rhoY0 / ud.stratification(0.0) + mpv.HydroState_n.S0[igy] = 1.0 / mpv.HydroState_n.Y0[igy] + mpv.HydroState_n.p0[igy] = p0 + mpv.HydroState_n.p20[igy] = pi0 / ud.Msq + + # Ghost cells below bottom (negative heights) + Sn_integral_p = np.zeros(igy) + yn_p = node.y[:igy] - node.dy + yn_m = node.y[1:igy+1] - node.dy + + Sn_integral_p[:] = -node.dy * 1.0 / ud.stratification(0.5*(yn_p + yn_m)) + Sn_integral_p = np.cumsum(Sn_integral_p[:igy][::-1])[::-1] + + # Bulk domain above reference level + yn_p = node.y[igy+1:] + yn_m = np.zeros_like(yn_p) + yn_m[1:] = yn_p[:-1] + + Sn_p = 1.0 / ud.stratification(0.5 * (yn_p + yn_m)) + Sn_integral_p = np.hstack((Sn_integral_p, np.cumsum(elem.dy * Sn_p))) + + # Calculate nodal hydrostatic fields + pi_hydro_n = pi0 - Gamma * g * Sn_integral_p + rhoY_hydro_n = pi_hydro_n**gm1_inv + + # Update node solutions - below reference + mpv.HydroState_n.rhoY0[:igy] = rhoY_hydro_n[:igy] + mpv.HydroState_n.Y0[:igy+1] = ud.stratification(0.5 * (y_ps[:igy+1] + y_ps[:igy+1] - elem.dy)) + mpv.HydroState_n.rho0[:igy] = rhoY_hydro_n[:igy] / mpv.HydroState_n.Y0[:igy] + mpv.HydroState_n.S0[:igy] = 1.0 / mpv.HydroState_n.Y0[:igy] + mpv.HydroState_n.p0[:igy] = rhoY_hydro_n[:igy]**th.gamm + mpv.HydroState_n.p20[:igy] = pi_hydro_n[:igy] / ud.Msq + + # Update node solutions - above reference + mpv.HydroState_n.rhoY0[igy+1:] = rhoY_hydro_n[igy:] + mpv.HydroState_n.Y0[igy+1:] = ud.stratification(0.5 * (y_ps[igy:] + y_ps[igy:] + elem.dy)) + mpv.HydroState_n.rho0[igy+1:] = rhoY_hydro_n[igy:] / mpv.HydroState_n.Y0[igy+1:] + mpv.HydroState_n.S0[igy+1:] = 1.0 / mpv.HydroState_n.Y0[igy+1:] + mpv.HydroState_n.p0[igy+1:] = rhoY_hydro_n[igy:]**th.gamm + mpv.HydroState_n.p20[igy+1:] = pi_hydro_n[igy:] / ud.Msq + + else: + # No gravity case - uniform atmosphere + mpv.HydroState.p20[:] = 1.0 + mpv.HydroState.p0[:] = 1.0 + mpv.HydroState.rho0[:] = 1.0 + mpv.HydroState.rhoY0[:] = 1.0 + mpv.HydroState.Y0[:] = 1.0 + mpv.HydroState.S0[:] = 1.0 + mpv.HydroState.S10[:] = 0.0 + + mpv.HydroState_n.p20[:] = 1.0 + mpv.HydroState_n.p0[:] = 1.0 + mpv.HydroState_n.rho0[:] = 1.0 + mpv.HydroState_n.rhoY0[:] = 1.0 + mpv.HydroState_n.Y0[:] = 1.0 + mpv.HydroState_n.S0[:] = 1.0 + + +def analytical_state(mpv, elem, node, th, ud): + g = ud.gravity_strength[1] + Gamma = th.Gamma + Hex = 1.0 / (th.Gamma * g) + dy = elem.dy + + pi_np = np.exp(-(node.y + 0.5 * dy) / Hex) + pi_nm = np.exp(-(node.y - 0.5 * dy) / Hex) + pi_n = np.exp(-(node.y) / Hex) + + Y_n = - Gamma * g * dy / (pi_np - pi_nm) + P_n = pi_n**th.gm1inv + p_n = pi_n**th.Gammainv + rho_n = P_n / Y_n + + mpv.HydroState_n.p20[...] = pi_n / ud.Msq + mpv.HydroState_n.p0[...] = p_n + mpv.HydroState_n.rho0[...] = rho_n + mpv.HydroState_n.rhoY0[...] = P_n + mpv.HydroState_n.Y0[...] = Y_n + mpv.HydroState_n.S0[...] = 1.0 / Y_n + + pi_cp = np.exp(-(elem.y + 0.5 * dy) / Hex) + pi_cm = np.exp(-(elem.y - 0.5 * dy) / Hex) + pi_c = np.exp(-(elem.y) / Hex) + + Y_c = - Gamma * g * dy / (pi_cp - pi_cm) + P_c = pi_c**th.gm1inv + p_c = pi_c**th.Gammainv + rho_c = P_c / Y_c + + mpv.HydroState.p20[...] = pi_c / ud.Msq + mpv.HydroState.p0[...] = p_c + mpv.HydroState.rho0[...] = rho_c + mpv.HydroState.rhoY0[...] = P_c + mpv.HydroState.Y0[...] = Y_c + mpv.HydroState.S0[...] = 1.0 / Y_c + + + +def initial_pressure(Sol, mpv, elem, node, ud, th): + Gammainv = th.Gammainv + igy = node.igy + igx = node.igx + icx = elem.icx + dx = node.dx + dy = node.dy + + beta = np.zeros((node.icx)) + bdpdx = np.zeros((node.icx)) + + x_idx_m = slice(0, -1) + x_idx_c = slice(1, None) + y_idx = slice(igy, -igy) + xn_idx = slice(1, -1) + height = node.y[-igy - 1] + + Pc = Sol.rhoY[x_idx_c, y_idx] + Pm = Sol.rhoY[x_idx_m, y_idx] + thc = Pc / Sol.rho[x_idx_c, y_idx] + thm = Pm / Sol.rho[x_idx_m, y_idx] + beta[xn_idx] = np.sum(0.5 * (Pm * thm + Pc * thc) * dy, axis=1) + bdpdx[xn_idx] = np.sum( + 0.5 + * (Pm * thm + Pc * thc) + * (mpv.p2_cells[x_idx_c, y_idx] - mpv.p2_cells[x_idx_m, y_idx]) + * dy, + axis=1, + ) + + beta *= Gammainv / height + bdpdx *= Gammainv / height / dx + + coeff = np.zeros((elem.icx)) + pibot = np.zeros((elem.icx)) + coeff[igx + 1 : -igx + 1] = np.cumsum(coeff[igx:-igx] + dx / beta[igx + 1 : -igx]) + pibot[igx + 1 : -igx + 1] = np.cumsum( + pibot[igx:-igx] - dx * bdpdx[igx + 1 : -igx] / beta[igx + 1 : -igx] + ) + + dotPU = pibot[icx - igx] / coeff[icx - igx] + pibot[igx:-igx] -= dotPU * coeff[igx:-igx] + + x_idx = slice(igx, -igx + 1) + y_idx = slice(igy, -igy + 1) + + mpv.p2_cells[x_idx, y_idx] += pibot[x_idx].reshape( + -1, 1 + ) - 1.0 * mpv.HydroState.p20[y_idx].reshape(1, -1) + bdry.set_ghostcells_p2(mpv.p2_cells, elem, ud) + + icxn = node.icx + icyn = node.icy + x_idx = slice(1, icxn - 1) + y_idx = slice(igy, -igy + 1) + height = node.y[-igy] + + Pc = Sol.rhoY[1:, y_idx] + thc = Pc / Sol.rho[1:, y_idx] + + beta = np.zeros((elem.icx,)) + bdpdx = np.zeros((elem.icx)) + + beta[1:] = np.sum(Pc * thc * dy, axis=1) + beta *= Gammainv / height + + bdpdx[1:] = np.sum( + Pc * thc * (mpv.p2_nodes[1:-1, igy:-igy] - mpv.p2_nodes[:-2, igy:-igy]) * dy, + axis=1, + ) + bdpdx *= Gammainv / height / dx + + coeff = np.zeros((node.icx)) + pibot = np.zeros((node.icx)) + + coeff[igx + 1 : -igx + 1] = np.cumsum(coeff[igx:-igx] + dx / beta[igx + 1 :]) + pibot[igx + 1 : -igx + 1] = np.cumsum( + pibot[igx:-igx] - dx * bdpdx[igx + 1 :] / beta[igx + 1 :] + ) + + dotPU = pibot[icx - igx] / coeff[icx - igx] + + pibot[igx:-igx] -= dotPU * coeff[igx:-igx] + + x_idx = slice(igx, -igx + 1) + y_idx = slice(igy, -igy + 1) + mpv.p2_nodes[x_idx, y_idx] += pibot[x_idx].reshape( + -1, 1 + ) - 1.0 * mpv.HydroState_n.p20[y_idx].reshape(1, -1) + + mpv.dp2_nodes[:, :] = mpv.p2_nodes + + # guess initial node value (at left-most node) + mpv.p2_nodes[igx, igy:-igy] = mpv.dp2_nodes[igx, igy:-igy] + + mpv.p2_nodes[:, :] = __loop_over_array( + igx, igy, icxn, icyn, mpv.p2_nodes, mpv.dp2_nodes + ) + + assert ((node.icx + 1) % 2) == 1 + delp2 = 0.5 * (mpv.p2_nodes[-igx - 1, igy:-igy] - mpv.p2_nodes[igx, igy:-igy]) + delp2 = delp2.reshape(1, -1) + sgn = np.ones_like(mpv.p2_nodes[:, 0][igy:-igy]).reshape(-1, 1) + + sgn[1::2] *= -1 + + mpv.p2_nodes[igx:-igx, igy:-igy] += sgn * delp2 + bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) + + mpv.dp2_nodes[:, :] = 0.0 + + inner_domain = (slice(igx, -igx), slice(igy, -igy)) + pi = ud.Msq * (mpv.p2_cells[inner_domain] + 1.0 * mpv.HydroState.p20[igy:-igy]) + Y = Sol.rhoY[inner_domain] / Sol.rho[inner_domain] + rhoold = np.copy(Sol.rho[inner_domain]) + Sol.rhoY[inner_domain] = pi**th.gm1inv + Sol.rho[inner_domain] = Sol.rhoY[inner_domain] / Y + Sol.rhou[inner_domain] *= Sol.rho[inner_domain] / rhoold + Sol.rhov[inner_domain] *= Sol.rho[inner_domain] / rhoold + Sol.rhow[inner_domain] *= Sol.rho[inner_domain] / rhoold + Sol.rhoX[inner_domain] *= Sol.rho[inner_domain] / rhoold + + +# need details: +# populate the rest of the nodes recursively based on the left-most node. +# recursive: use numba. +@nb.jit(nopython=True) +def __loop_over_array(igx, igy, icxn, icyn, p, dp): + for j in range(igy, icyn - igy): + for i in range(igx + 1, icxn - igx): + p[i, j] = 2.0 * dp[i - 1, j] - p[i - 1, j] + return p diff --git a/src/dycore/physics/gas_dynamics/__init__.py b/src/pybella/flow_solver/physics/low_mach/__init__.py similarity index 100% rename from src/dycore/physics/gas_dynamics/__init__.py rename to src/pybella/flow_solver/physics/low_mach/__init__.py diff --git a/src/dycore/physics/low_mach/laplacian.py b/src/pybella/flow_solver/physics/low_mach/laplacian.py similarity index 51% rename from src/dycore/physics/low_mach/laplacian.py rename to src/pybella/flow_solver/physics/low_mach/laplacian.py index c41f5d48..cdf982e9 100644 --- a/src/dycore/physics/low_mach/laplacian.py +++ b/src/pybella/flow_solver/physics/low_mach/laplacian.py @@ -2,9 +2,10 @@ import scipy as sp import numba as nb -import dycore.utils.options as opts +from ...utils import options as opts -def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): + +def stencil_9pt(elem, node, mpv, Sol, ud, diag_inv, dt, coriolis_params): igx = elem.igx igy = elem.igy @@ -19,14 +20,22 @@ def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): dx = node.dy dy = node.dx - inner_domain = (slice(igx,-igx),slice(igy,-igy)) + inner_domain = (slice(igx, -igx), slice(igy, -igy)) i1 = node.i1 - hplusx = mpv.wplus[1][i1].reshape(-1,) - hplusy = mpv.wplus[0][i1].reshape(-1,) - hcenter = mpv.wcenter[i1].reshape(-1,) + hplusx = mpv.wplus[1][i1].reshape( + -1, + ) + hplusy = mpv.wplus[0][i1].reshape( + -1, + ) + hcenter = mpv.wcenter[i1].reshape( + -1, + ) - diag_inv = diag_inv[i1].reshape(-1,) + diag_inv = diag_inv[i1].reshape( + -1, + ) oodx = 1.0 / (dx) oody = 1.0 / (dy) @@ -34,9 +43,14 @@ def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): x_periodic = ud.bdry_type[1] == opts.BdryType.PERIODIC y_periodic = ud.bdry_type[0] == opts.BdryType.PERIODIC - x_wall = ud.bdry_type[1] == opts.BdryType.WALL or ud.bdry_type[1] == opts.BdryType.RAYLEIGH - y_wall = ud.bdry_type[0] == opts.BdryType.WALL or ud.bdry_type[0] == opts.BdryType.RAYLEIGH - + x_wall = ( + ud.bdry_type[1] == opts.BdryType.WALL + or ud.bdry_type[1] == opts.BdryType.RAYLEIGH + ) + y_wall = ( + ud.bdry_type[0] == opts.BdryType.WALL + or ud.bdry_type[0] == opts.BdryType.RAYLEIGH + ) #### Compute Coriolis parameters: # nonhydro = ud.nonhydrostasy @@ -67,7 +81,7 @@ def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): # y_idx1 = slice(igs[dim]-is_periodic+1, right_idx) # innerdim1[dim] = slice(igs[dim]-1, (-igs[dim]+1)) - + # strat = (mpv.HydroState_n.S0[y_idx1] - mpv.HydroState_n.S0[y_idx]) / dy # nindim = tuple(nindim) @@ -97,7 +111,24 @@ def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): # coriolis_params = (coeff_vv, coeff_vu, coeff_uv, coeff_uu) # coriolis_params = np.array(coriolis_params, dtype=np.float64) - return lambda p : lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, x_periodic, y_periodic, x_wall, y_wall, diag_inv, coriolis_params) + return lambda p: lap2D_gather( + p, + igx, + igy, + iicxn, + iicyn, + hplusx, + hplusy, + hcenter, + oodx, + oody, + x_periodic, + y_periodic, + x_wall, + y_wall, + diag_inv, + coriolis_params, + ) # ndim = elem.ndim # periodicity = np.zeros(ndim, dtype='int') @@ -145,8 +176,26 @@ def stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params): # return lambda p : lap2Dc(p, hplusx, hplusy, hcenter, dxy, periodicity, diag_inv, coriolis_params) + @nb.jit(nopython=True, nogil=False, cache=True) -def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, x_periodic, y_periodic, x_wall, y_wall, diag_inv, coriolis): +def lap2D_gather( + p, + igx, + igy, + iicxn, + iicyn, + hplusx, + hplusy, + hcenter, + oodx, + oody, + x_periodic, + y_periodic, + x_wall, + y_wall, + diag_inv, + coriolis, +): ngnc = (iicxn) * (iicyn) lap = np.zeros((ngnc)) cnt_x = 0 @@ -159,7 +208,7 @@ def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, for idx in range(iicxn * iicyn): ne_topleft = idx - iicxn - 1 - ne_topright = idx - iicxn + ne_topright = idx - iicxn ne_botleft = idx - 1 ne_botright = idx @@ -193,20 +242,20 @@ def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, ne_botright -= iicxn - 1 if cnt_y == 0: - topleft_idx += ((iicxn) * (iicyn - 1)) - topmid_idx += ((iicxn) * (iicyn - 1)) - topright_idx += ((iicxn) * (iicyn - 1)) + topleft_idx += (iicxn) * (iicyn - 1) + topmid_idx += (iicxn) * (iicyn - 1) + topright_idx += (iicxn) * (iicyn - 1) - ne_topleft += ((iicxn) * (iicyn - 1)) - ne_topright += ((iicxn) * (iicyn - 1)) + ne_topleft += (iicxn) * (iicyn - 1) + ne_topright += (iicxn) * (iicyn - 1) if cnt_y == (iicyn - 1): - botleft_idx -= ((iicxn) * (iicyn - 1)) - botmid_idx -= ((iicxn) * (iicyn - 1)) - botright_idx -= ((iicxn) * (iicyn - 1)) + botleft_idx -= (iicxn) * (iicyn - 1) + botmid_idx -= (iicxn) * (iicyn - 1) + botright_idx -= (iicxn) * (iicyn - 1) - ne_botleft -= ((iicxn) * (iicyn - 1)) - ne_botright -= ((iicxn) * (iicyn - 1)) + ne_botleft -= (iicxn) * (iicyn - 1) + ne_botright -= (iicxn) * (iicyn - 1) topleft = p[topleft_idx] midleft = p[midleft_idx] @@ -230,51 +279,50 @@ def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, hplusy_topright = hplusy[ne_topright] hplusy_botright = hplusy[ne_botright] - cxx_tl = cxx[ne_topleft] + cxx_tl = cxx[ne_topleft] cxx_tr = cxx[ne_topright] - cxx_bl = cxx[ne_botleft] + cxx_bl = cxx[ne_botleft] cxx_br = cxx[ne_botright] - cxy_tl = cxy[ne_topleft] - cxy_tr = cxy[ne_topright] - cxy_bl = cxy[ne_botleft] - cxy_br = cxy[ne_botright] + cxy_tl = cxy[ne_topleft] + cxy_tr = cxy[ne_topright] + cxy_bl = cxy[ne_botleft] + cxy_br = cxy[ne_botright] - cyx_tl = cyx[ne_topleft] - cyx_tr = cyx[ne_topright] - cyx_bl = cyx[ne_botleft] - cyx_br = cyx[ne_botright] + cyx_tl = cyx[ne_topleft] + cyx_tr = cyx[ne_topright] + cyx_bl = cyx[ne_botleft] + cyx_br = cyx[ne_botright] - cyy_tl = cyy[ne_topleft] - cyy_tr = cyy[ne_topright] - cyy_bl = cyy[ne_botleft] - cyy_br = cyy[ne_botright] - + cyy_tl = cyy[ne_topleft] + cyy_tr = cyy[ne_topright] + cyy_bl = cyy[ne_botleft] + cyy_br = cyy[ne_botright] if x_wall and (cnt_x == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_botleft = 0. - hplusy_botleft = 0. + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 if x_wall and (cnt_x == (iicxn - 1)): - hplusx_topright = 0. - hplusy_topright = 0. - hplusx_botright = 0. - hplusy_botright = 0. + hplusx_topright = 0.0 + hplusy_topright = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 if y_wall and (cnt_y == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_topright = 0. - hplusy_topright = 0. - + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_topright = 0.0 + hplusy_topright = 0.0 + if y_wall and (cnt_y == (iicyn - 1)): - hplusx_botleft = 0. - hplusy_botleft = 0. - hplusx_botright = 0. - hplusy_botright = 0. - + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 + # dp2dxdy1 = ((midmid - midleft) - (topmid - topleft)) * nine_pt # dp2dxdy2 = ((midright - midmid) - (topright - topmid)) * nine_pt # dp2dxdy3 = ((botmid - botleft) - (midmid - midleft)) * nine_pt @@ -290,26 +338,48 @@ def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, # + hplusy_botright * oody2 * ((botmid - midmid) + dp2dxdy4) \ # + hcenter[idx] * p[idx] - Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hplusx_topleft - Dx_tr = 0.5 * (topright - topmid + midright - midmid ) * hplusx_topright - Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hplusx_botleft - Dx_br = 0.5 * (botright - botmid + midright - midmid ) * hplusx_botright + Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hplusx_topleft + Dx_tr = 0.5 * (topright - topmid + midright - midmid) * hplusx_topright + Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hplusx_botleft + Dx_br = 0.5 * (botright - botmid + midright - midmid) * hplusx_botright - Dy_tl = 0.5 * (midmid - topmid + midleft - topleft) * hplusy_topleft - Dy_tr = 0.5 * (midright - topright + midmid - topmid ) * hplusy_topright - Dy_bl = 0.5 * (botmid - midmid + botleft - midleft) * hplusy_botleft - Dy_br = 0.5 * (botright - midright + botmid - midmid ) * hplusy_botright + Dy_tl = 0.5 * (midmid - topmid + midleft - topleft) * hplusy_topleft + Dy_tr = 0.5 * (midright - topright + midmid - topmid) * hplusy_topright + Dy_bl = 0.5 * (botmid - midmid + botleft - midleft) * hplusy_botleft + Dy_br = 0.5 * (botright - midright + botmid - midmid) * hplusy_botright fac = 1.0 - Dxx = 0.5 * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) * oodx * oodx * fac - Dyy = 0.5 * (cyy_br * Dy_br - cyy_tr * Dy_tr + cyy_bl * Dy_bl - cyy_tl * Dy_tl) * oody * oody * fac - Dyx = 0.5 * (cyx_br * Dy_br - cyx_bl * Dy_bl + cyx_tr * Dy_tr - cyx_tl * Dy_tl) * oody * oodx * fac - Dxy = 0.5 * (cxy_br * Dx_br - cxy_tr * Dx_tr + cxy_bl * Dy_bl - cxy_tl * Dx_tl) * oodx * oody * fac - + Dxx = ( + 0.5 + * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) + * oodx + * oodx + * fac + ) + Dyy = ( + 0.5 + * (cyy_br * Dy_br - cyy_tr * Dy_tr + cyy_bl * Dy_bl - cyy_tl * Dy_tl) + * oody + * oody + * fac + ) + Dyx = ( + 0.5 + * (cyx_br * Dy_br - cyx_bl * Dy_bl + cyx_tr * Dy_tr - cyx_tl * Dy_tl) + * oody + * oodx + * fac + ) + Dxy = ( + 0.5 + * (cxy_br * Dx_br - cxy_tr * Dx_tr + cxy_bl * Dy_bl - cxy_tl * Dx_tl) + * oodx + * oody + * fac + ) lap[idx] = Dxx + Dyy + Dyx + Dxy + hcenter[idx] * p[idx] - # fac = 1.0 # lap[idx] = - fac * hplusx_topleft * oodx2 * ((midmid - midleft)) \ # - fac * hplusy_topleft * oody2 * ((midmid - topmid)) \ @@ -327,12 +397,28 @@ def lap2D_gather(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx, oody, if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - + return lap @nb.jit(nopython=True, nogil=False, cache=True) -def lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_periodic, y_periodic, x_wall, y_wall, diag_inv): +def lap2D( + p, + igx, + igy, + iicxn, + iicyn, + hplusx, + hplusy, + hcenter, + oodx2, + oody2, + x_periodic, + y_periodic, + x_wall, + y_wall, + diag_inv, +): ngnc = (iicxn) * (iicyn) lap = np.zeros((ngnc)) cnt_x = 0 @@ -340,10 +426,9 @@ def lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_per nine_pt = 0.25 * (2.0) * 1.0 - for idx in range(iicxn * iicyn): ne_topleft = idx - iicxn - 1 - ne_topright = idx - iicxn + ne_topright = idx - iicxn ne_botleft = idx - 1 ne_botright = idx @@ -387,30 +472,30 @@ def lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_per ne_botright -= iicxn - 1 if cnt_y == 0: - topleft_idx += ((iicxn) * (iicyn - 1)) - topmid_idx += ((iicxn) * (iicyn - 1)) - topright_idx += ((iicxn) * (iicyn - 1)) + topleft_idx += (iicxn) * (iicyn - 1) + topmid_idx += (iicxn) * (iicyn - 1) + topright_idx += (iicxn) * (iicyn - 1) # if y_periodic: # midleft_idx += ((iicxn) * (iicyn - 1)) # midmid_idx += ((iicxn) * (iicyn - 1)) # midright_idx += ((iicxn) * (iicyn - 1)) - ne_topleft += ((iicxn) * (iicyn - 1)) - ne_topright += ((iicxn) * (iicyn - 1)) + ne_topleft += (iicxn) * (iicyn - 1) + ne_topright += (iicxn) * (iicyn - 1) if cnt_y == (iicyn - 1): - botleft_idx -= ((iicxn) * (iicyn - 1)) - botmid_idx -= ((iicxn) * (iicyn - 1)) - botright_idx -= ((iicxn) * (iicyn - 1)) + botleft_idx -= (iicxn) * (iicyn - 1) + botmid_idx -= (iicxn) * (iicyn - 1) + botright_idx -= (iicxn) * (iicyn - 1) # if y_periodic: # midleft_idx -= ((iicxn) * (iicyn - 1)) # midmid_idx -= ((iicxn) * (iicyn - 1)) # midright_idx -= ((iicxn) * (iicyn - 1)) - ne_botleft -= ((iicxn) * (iicyn - 1)) - ne_botright -= ((iicxn) * (iicyn - 1)) + ne_botleft -= (iicxn) * (iicyn - 1) + ne_botright -= (iicxn) * (iicyn - 1) topleft = p[topleft_idx] midleft = p[midleft_idx] @@ -435,43 +520,45 @@ def lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_per hplusy_botright = hplusy[ne_botright] if x_wall and (cnt_x == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_botleft = 0. - hplusy_botleft = 0. + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 if x_wall and (cnt_x == (iicxn - 1)): - hplusx_topright = 0. - hplusy_topright = 0. - hplusx_botright = 0. - hplusy_botright = 0. + hplusx_topright = 0.0 + hplusy_topright = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 if y_wall and (cnt_y == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_topright = 0. - hplusy_topright = 0. - + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_topright = 0.0 + hplusy_topright = 0.0 + if y_wall and (cnt_y == (iicyn - 1)): - hplusx_botleft = 0. - hplusy_botleft = 0. - hplusx_botright = 0. - hplusy_botright = 0. - + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 + dp2dxdy1 = ((midmid - midleft) - (topmid - topleft)) * nine_pt dp2dxdy2 = ((midright - midmid) - (topright - topmid)) * nine_pt dp2dxdy3 = ((botmid - botleft) - (midmid - midleft)) * nine_pt dp2dxdy4 = ((botright - botmid) - (midright - midmid)) * nine_pt - lap[idx] = - hplusx_topleft * oodx2 * ((midmid - midleft) - dp2dxdy1) \ - - hplusy_topleft * oody2 * ((midmid - topmid) - dp2dxdy1) \ - + hplusx_topright * oodx2 * ((midright - midmid) - dp2dxdy2) \ - - hplusy_topright * oody2 * ((midmid - topmid) + dp2dxdy2) \ - - hplusx_botleft * oodx2 * ((midmid - midleft) + dp2dxdy3) \ - + hplusy_botleft * oody2 * ((botmid - midmid) - dp2dxdy3) \ - + hplusx_botright * oodx2 * ((midright - midmid) + dp2dxdy4) \ - + hplusy_botright * oody2 * ((botmid - midmid) + dp2dxdy4) \ - + hcenter[idx] * p[idx] + lap[idx] = ( + -hplusx_topleft * oodx2 * ((midmid - midleft) - dp2dxdy1) + - hplusy_topleft * oody2 * ((midmid - topmid) - dp2dxdy1) + + hplusx_topright * oodx2 * ((midright - midmid) - dp2dxdy2) + - hplusy_topright * oody2 * ((midmid - topmid) + dp2dxdy2) + - hplusx_botleft * oodx2 * ((midmid - midleft) + dp2dxdy3) + + hplusy_botleft * oody2 * ((botmid - midmid) - dp2dxdy3) + + hplusx_botright * oodx2 * ((midright - midmid) + dp2dxdy4) + + hplusy_botright * oody2 * ((botmid - midmid) + dp2dxdy4) + + hcenter[idx] * p[idx] + ) # if cnt_x == 0 and x_wall: # lap[idx] *= 2. @@ -489,11 +576,11 @@ def lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_per if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - + return lap -def stencil_9pt_numba_test(mpv,node,coriolis,diag_inv, ud): +def stencil_9pt_numba_test(mpv, node, coriolis, diag_inv, ud): dx = node.dx dy = node.dy @@ -505,46 +592,105 @@ def stencil_9pt_numba_test(mpv,node,coriolis,diag_inv, ud): shp = node.iisc - dummy_p = np.zeros((node.isc[1],node.isc[0])) + dummy_p = np.zeros((node.isc[1], node.isc[0])) ### Need to clean this up, but the Numba stencil is used in the Helmholtz solve for radiative BC! - if hasattr(ud, 'LAMB_BDRY'): - return lambda p : lap2D_numba_test(p, dummy_p, dx, dy, coeffs, diag_inv.T, coriolis, shp) + if hasattr(ud, "ATMOSPHERIC_EXTENSION"): + return lambda p: lap2D_numba_test( + p, dummy_p, dx, dy, coeffs, diag_inv.T, coriolis, shp + ) ################### else: - x_wall = ud.bdry_type[0] == opts.BdryType.WALL or ud.bdry_type[0] == opts.BdryType.RAYLEIGH - y_wall = ud.bdry_type[1] == opts.BdryType.WALL or ud.bdry_type[1] == opts.BdryType.RAYLEIGH + x_wall = ( + ud.bdry_type[0] == opts.BdryType.WALL + or ud.bdry_type[0] == opts.BdryType.RAYLEIGH + ) + y_wall = ( + ud.bdry_type[1] == opts.BdryType.WALL + or ud.bdry_type[1] == opts.BdryType.RAYLEIGH + ) y_rayleigh = ud.bdry_type[1] == opts.BdryType.RAYLEIGH - cor_slc = (slice(1,-1), slice(1,-1)) - coeff_slc = (slice(1,-1), slice(1,-1)) - - coeffs = (hplusx[coeff_slc].T.reshape(-1,), hplusy[coeff_slc].T.reshape(-1,), hcenter[node.i1].T.reshape(-1,)) - - coriolis = (coriolis[0][cor_slc].reshape(-1,),coriolis[1][cor_slc].reshape(-1,),coriolis[2][cor_slc].reshape(-1,),coriolis[3][cor_slc].reshape(-1,)) + cor_slc = (slice(1, -1), slice(1, -1)) + coeff_slc = (slice(1, -1), slice(1, -1)) + + coeffs = ( + hplusx[coeff_slc].T.reshape( + -1, + ), + hplusy[coeff_slc].T.reshape( + -1, + ), + hcenter[node.i1].T.reshape( + -1, + ), + ) + + coriolis = ( + coriolis[0][cor_slc].reshape( + -1, + ), + coriolis[1][cor_slc].reshape( + -1, + ), + coriolis[2][cor_slc].reshape( + -1, + ), + coriolis[3][cor_slc].reshape( + -1, + ), + ) + + return lambda p: lap2D_gather_new( + p, + node.iicx, + node.iicy, + coeffs, + dx, + dy, + y_rayleigh, + x_wall, + y_wall, + diag_inv[node.i1].T.reshape( + -1, + ), + coriolis, + ) - return lambda p : lap2D_gather_new(p, node.iicx, node.iicy, coeffs, dx, dy, y_rayleigh, x_wall, y_wall, diag_inv[node.i1].T.reshape(-1,), coriolis) - @nb.jit(nopython=True, cache=False) def lap2D_numba_test(p, dp, dx, dy, coeffs, diag_inv, coriolis, shp): - p = p.reshape(shp[1],shp[0]) - dp[1:-1,1:-1] = p + p = p.reshape(shp[1], shp[0]) + dp[1:-1, 1:-1] = p # dp = p dp = periodic(dp) - dp = kernel_9pt(dp, dx, dy, coeffs[0], coeffs[1], coeffs[2], diag_inv, coriolis[0], coriolis[1], coriolis[2], coriolis[3]) + dp = kernel_9pt( + dp, + dx, + dy, + coeffs[0], + coeffs[1], + coeffs[2], + diag_inv, + coriolis[0], + coriolis[1], + coriolis[2], + coriolis[3], + ) # p = dp # dp = periodic(dp) - p = dp[1:-1,1:-1] + p = dp[1:-1, 1:-1] return p.ravel() @nb.njit(cache=True) -def lap2D_gather_new(p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall, diag_inv, coriolis): +def lap2D_gather_new( + p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall, diag_inv, coriolis +): ngnc = (iicxn) * (iicyn) lap = np.zeros((ngnc)) cnt_x = 0 @@ -615,9 +761,9 @@ def lap2D_gather_new(p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall # ne_botright -= iicxn + 1 if cnt_y == 0: - topleft_idx += ((iicxn) * (iicyn - 1)) - topmid_idx += ((iicxn) * (iicyn - 1)) - topright_idx += ((iicxn) * (iicyn - 1)) + topleft_idx += (iicxn) * (iicyn - 1) + topmid_idx += (iicxn) * (iicyn - 1) + topright_idx += (iicxn) * (iicyn - 1) # ne_topleft += ((iicxn + 1) * (iicyn)) # ne_topright += ((iicxn + 1) * (iicyn)) @@ -627,9 +773,9 @@ def lap2D_gather_new(p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall # if cnt_y == (iicyn - 1) and not y_rayleigh: if cnt_y == (iicyn - 1): - botleft_idx -= ((iicxn) * (iicyn - 1)) - botmid_idx -= ((iicxn) * (iicyn - 1)) - botright_idx -= ((iicxn) * (iicyn - 1)) + botleft_idx -= (iicxn) * (iicyn - 1) + botmid_idx -= (iicxn) * (iicyn - 1) + botright_idx -= (iicxn) * (iicyn - 1) # ne_botleft -= ((iicxn + 1) * (iicyn)) # ne_botright -= ((iicxn + 1) * (iicyn)) @@ -667,72 +813,95 @@ def lap2D_gather_new(p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall hplusy_topright = hplusy[ne_topright] hplusy_botright = hplusy[ne_botright] - cxx_tl = cxx[ne_topleft] + cxx_tl = cxx[ne_topleft] cxx_tr = cxx[ne_topright] - cxx_bl = cxx[ne_botleft] + cxx_bl = cxx[ne_botleft] cxx_br = cxx[ne_botright] - cxy_tl = cxy[ne_topleft] - cxy_tr = cxy[ne_topright] - cxy_bl = cxy[ne_botleft] - cxy_br = cxy[ne_botright] + cxy_tl = cxy[ne_topleft] + cxy_tr = cxy[ne_topright] + cxy_bl = cxy[ne_botleft] + cxy_br = cxy[ne_botright] - cyx_tl = cyx[ne_topleft] - cyx_tr = cyx[ne_topright] - cyx_bl = cyx[ne_botleft] - cyx_br = cyx[ne_botright] + cyx_tl = cyx[ne_topleft] + cyx_tr = cyx[ne_topright] + cyx_bl = cyx[ne_botleft] + cyx_br = cyx[ne_botright] - cyy_tl = cyy[ne_topleft] - cyy_tr = cyy[ne_topright] - cyy_bl = cyy[ne_botleft] - cyy_br = cyy[ne_botright] - + cyy_tl = cyy[ne_topleft] + cyy_tr = cyy[ne_topright] + cyy_bl = cyy[ne_botleft] + cyy_br = cyy[ne_botright] if x_wall and (cnt_x == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_botleft = 0. - hplusy_botleft = 0. + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 if x_wall and (cnt_x == (iicxn - 1)): - hplusx_topright = 0. - hplusy_topright = 0. - hplusx_botright = 0. - hplusy_botright = 0. + hplusx_topright = 0.0 + hplusy_topright = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 if y_wall and (cnt_y == 0): - hplusx_topleft = 0. - hplusy_topleft = 0. - hplusx_topright = 0. - hplusy_topright = 0. - + hplusx_topleft = 0.0 + hplusy_topleft = 0.0 + hplusx_topright = 0.0 + hplusy_topright = 0.0 + if y_wall and (cnt_y == (iicyn - 1)): - hplusx_botleft = 0. - hplusy_botleft = 0. - hplusx_botright = 0. - hplusy_botright = 0. + hplusx_botleft = 0.0 + hplusy_botleft = 0.0 + hplusx_botright = 0.0 + hplusy_botright = 0.0 # if y_rayleigh and (cnt_y == 0): # hplusx_topleft = 0. - # hplusy_topleft = 0. + # hplusy_topleft = 0. # hplusx_topright = 0. # hplusy_topright = 0. - Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hplusx_topleft - Dx_tr = 0.5 * (topright - topmid + midright - midmid ) * hplusx_topright - Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hplusx_botleft - Dx_br = 0.5 * (botright - botmid + midright - midmid ) * hplusx_botright + Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hplusx_topleft + Dx_tr = 0.5 * (topright - topmid + midright - midmid) * hplusx_topright + Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hplusx_botleft + Dx_br = 0.5 * (botright - botmid + midright - midmid) * hplusx_botright - Dy_tl = 0.5 * (midmid - topmid + midleft - topleft) * hplusy_topleft - Dy_tr = 0.5 * (midright - topright + midmid - topmid ) * hplusy_topright - Dy_bl = 0.5 * (botmid - midmid + botleft - midleft) * hplusy_botleft - Dy_br = 0.5 * (botright - midright + botmid - midmid ) * hplusy_botright + Dy_tl = 0.5 * (midmid - topmid + midleft - topleft) * hplusy_topleft + Dy_tr = 0.5 * (midright - topright + midmid - topmid) * hplusy_topright + Dy_bl = 0.5 * (botmid - midmid + botleft - midleft) * hplusy_botleft + Dy_br = 0.5 * (botright - midright + botmid - midmid) * hplusy_botright fac = 1.0 - Dxx = 0.5 * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) * oodx * oodx * fac - Dyy = 0.5 * (cyy_br * Dy_br - cyy_tr * Dy_tr + cyy_bl * Dy_bl - cyy_tl * Dy_tl) * oody * oody * fac - Dyx = 0.5 * (cxy_br * Dy_br - cxy_bl * Dy_bl + cxy_tr * Dy_tr - cxy_tl * Dy_tl) * oody * oodx * fac - Dxy = 0.5 * (cyx_br * Dx_br - cyx_tr * Dx_tr + cyx_bl * Dx_bl - cyx_tl * Dx_tl) * oodx * oody * fac + Dxx = ( + 0.5 + * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) + * oodx + * oodx + * fac + ) + Dyy = ( + 0.5 + * (cyy_br * Dy_br - cyy_tr * Dy_tr + cyy_bl * Dy_bl - cyy_tl * Dy_tl) + * oody + * oody + * fac + ) + Dyx = ( + 0.5 + * (cxy_br * Dy_br - cxy_bl * Dy_bl + cxy_tr * Dy_tr - cxy_tl * Dy_tl) + * oody + * oodx + * fac + ) + Dxy = ( + 0.5 + * (cyx_br * Dx_br - cyx_tr * Dx_tr + cyx_bl * Dx_bl - cyx_tl * Dx_tl) + * oodx + * oody + * fac + ) lap[idx] = Dxx + Dyy + Dyx + Dxy + hcenter[idx] * p[idx] @@ -742,7 +911,7 @@ def lap2D_gather_new(p, iicxn, iicyn, coeffs, dx, dy, y_rayleigh, x_wall, y_wall if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - + return lap @@ -751,17 +920,17 @@ def kernel_9pt(a, dx, dy, hpx, hpy, hpc, diag_inv, cxx, cyy, cxy, cyx): oodx = 1.0 / dx oody = 1.0 / dy - topleft = a[1,-1] - topmid = a[1,0] - topright = a[1,1] - - midleft = a[0,-1] - midmid = a[0,0] - midright = a[0,1] - - botleft = a[-1,-1] - botmid = a[-1,0] - botright = a[-1,1] + topleft = a[1, -1] + topmid = a[1, 0] + topright = a[1, 1] + + midleft = a[0, -1] + midmid = a[0, 0] + midright = a[0, 1] + + botleft = a[-1, -1] + botmid = a[-1, 0] + botright = a[-1, 1] # shftr = 1 # shftc = 0 @@ -804,83 +973,103 @@ def kernel_9pt(a, dx, dy, hpx, hpy, hpc, diag_inv, cxx, cyy, cxy, cyx): # cyx_tl = cyx[tlr,tlc] # cyx_tr = cyx[trr,trc] - hpx_bl = hpx[0,0] - hpx_br = hpx[0,1] - hpx_tl = hpx[1,0] - hpx_tr = hpx[1,1] - - hpy_bl = hpy[0,0] - hpy_br = hpy[0,1] - hpy_tl = hpy[1,0] - hpy_tr = hpy[1,1] - - cxx_bl = cxx[0,0] - cxx_br = cxx[0,1] - cxx_tl = cxx[1,0] - cxx_tr = cxx[1,1] - - cyy_bl = cyy[0,0] - cyy_br = cyy[0,1] - cyy_tl = cyy[1,0] - cyy_tr = cyy[1,1] - - cxy_bl = cxy[0,0] - cxy_br = cxy[0,1] - cxy_tl = cxy[1,0] - cxy_tr = cxy[1,1] - - cyx_bl = cyx[0,0] - cyx_br = cyx[0,1] - cyx_tl = cyx[1,0] - cyx_tr = cyx[1,1] - - Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hpx_tl - Dx_tr = 0.5 * (topright - topmid + midright - midmid ) * hpx_tr - Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hpx_bl - Dx_br = 0.5 * (botright - botmid + midright - midmid ) * hpx_br - - Dy_tl = 0.5 * (topmid - midmid + topleft - midleft) * hpy_tl - Dy_tr = 0.5 * (topright - midright + topmid - midmid ) * hpy_tr - Dy_bl = 0.5 * (midmid - botmid + midleft - botleft) * hpy_bl - Dy_br = 0.5 * (midright - botright + midmid - botmid ) * hpy_br + hpx_bl = hpx[0, 0] + hpx_br = hpx[0, 1] + hpx_tl = hpx[1, 0] + hpx_tr = hpx[1, 1] + + hpy_bl = hpy[0, 0] + hpy_br = hpy[0, 1] + hpy_tl = hpy[1, 0] + hpy_tr = hpy[1, 1] + + cxx_bl = cxx[0, 0] + cxx_br = cxx[0, 1] + cxx_tl = cxx[1, 0] + cxx_tr = cxx[1, 1] + + cyy_bl = cyy[0, 0] + cyy_br = cyy[0, 1] + cyy_tl = cyy[1, 0] + cyy_tr = cyy[1, 1] + + cxy_bl = cxy[0, 0] + cxy_br = cxy[0, 1] + cxy_tl = cxy[1, 0] + cxy_tr = cxy[1, 1] + + cyx_bl = cyx[0, 0] + cyx_br = cyx[0, 1] + cyx_tl = cyx[1, 0] + cyx_tr = cyx[1, 1] + + Dx_tl = 0.5 * (topmid - topleft + midmid - midleft) * hpx_tl + Dx_tr = 0.5 * (topright - topmid + midright - midmid) * hpx_tr + Dx_bl = 0.5 * (botmid - botleft + midmid - midleft) * hpx_bl + Dx_br = 0.5 * (botright - botmid + midright - midmid) * hpx_br + + Dy_tl = 0.5 * (topmid - midmid + topleft - midleft) * hpy_tl + Dy_tr = 0.5 * (topright - midright + topmid - midmid) * hpy_tr + Dy_bl = 0.5 * (midmid - botmid + midleft - botleft) * hpy_bl + Dy_br = 0.5 * (midright - botright + midmid - botmid) * hpy_br # print(hpx_tl, hpx_tr, hpx_bl, hpx_br) # print(hpy_tl, hpy_tr, hpy_bl, hpy_br) # print("") - - Dxx = 0.5 * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) * oodx * oodx - Dyy = 0.5 * (cyy_tr * Dy_tr - cyy_br * Dy_br + cyy_tl * Dy_tl - cyy_bl * Dy_bl) * oody * oody - Dyx = 0.5 * (cxy_br * Dy_br - cxy_bl * Dy_bl + cxy_tr * Dy_tr - cxy_tl * Dy_tl) * oody * oodx - Dxy = 0.5 * (cyx_tr * Dx_tr - cyx_br * Dx_br + cyx_tl * Dx_tl - cyx_bl * Dx_bl) * oodx * oody - - return ((Dxx + Dyy + Dyx + Dxy) + hpc[0,0] * a[0,0]) * diag_inv[0,0] + + Dxx = ( + 0.5 + * (cxx_tr * Dx_tr - cxx_tl * Dx_tl + cxx_br * Dx_br - cxx_bl * Dx_bl) + * oodx + * oodx + ) + Dyy = ( + 0.5 + * (cyy_tr * Dy_tr - cyy_br * Dy_br + cyy_tl * Dy_tl - cyy_bl * Dy_bl) + * oody + * oody + ) + Dyx = ( + 0.5 + * (cxy_br * Dy_br - cxy_bl * Dy_bl + cxy_tr * Dy_tr - cxy_tl * Dy_tl) + * oody + * oodx + ) + Dxy = ( + 0.5 + * (cyx_tr * Dx_tr - cyx_br * Dx_br + cyx_tl * Dx_tl - cyx_bl * Dx_bl) + * oodx + * oody + ) + + return ((Dxx + Dyy + Dyx + Dxy) + hpc[0, 0] * a[0, 0]) * diag_inv[0, 0] + @nb.njit(cache=True) def periodic(arr): # periodic padding - arr[:,0] = arr[:,-3] - arr[:,-1] = arr[:,2] + arr[:, 0] = arr[:, -3] + arr[:, -1] = arr[:, 2] # wall padding - arr[0,:] = arr[2,:] - arr[-1,:] = arr[-3,:] - - return arr + arr[0, :] = arr[2, :] + arr[-1, :] = arr[-3, :] + return arr -def stencil_27pt(elem,node,mpv,ud,diag_inv,dt): +def stencil_27pt(elem, node, mpv, ud, diag_inv, dt): oodxyz = node.dxyz - oodxyz = 1./(oodxyz**2) + oodxyz = 1.0 / (oodxyz**2) oodx2, oody2, oodz2 = oodxyz[0], oodxyz[1], oodxyz[2] - odx, odz = 1./node.dx, 1./node.dz + odx, odz = 1.0 / node.dx, 1.0 / node.dz - i0 = (slice(0,-1),slice(0,-1),slice(0,-1)) - i1 = (slice(1,-1),slice(1,-1),slice(1,-1)) - i2 = (slice(2,-2),slice(2,-2),slice(2,-2)) + i0 = (slice(0, -1), slice(0, -1), slice(0, -1)) + i1 = (slice(1, -1), slice(1, -1), slice(1, -1)) + i2 = (slice(2, -2), slice(2, -2), slice(2, -2)) ndim = elem.ndim - periodicity = np.empty(ndim, dtype='int') + periodicity = np.empty(ndim, dtype="int") for dim in range(ndim): periodicity[dim] = ud.bdry_type[dim] == opts.BdryType.PERIODIC @@ -893,154 +1082,261 @@ def stencil_27pt(elem,node,mpv,ud,diag_inv,dt): corrf = dt * ud.coriolis_strength[0] - return lambda p : lap3D(p, hplusx, hplusy, hplusz, hcenter, oodx2, oody2, oodz2, periodicity, diag_inv, corrf, odx, odz) + return lambda p: lap3D( + p, + hplusx, + hplusy, + hplusz, + hcenter, + oodx2, + oody2, + oodz2, + periodicity, + diag_inv, + corrf, + odx, + odz, + ) + @nb.jit(nopython=True, cache=True, nogil=False) -def lap3D(p0, hplusx, hplusy, hplusz, hcenter, oodx2, oody2, oodz2, periodicity, diag_inv, corrf, odx, odz): +def lap3D( + p0, + hplusx, + hplusy, + hplusz, + hcenter, + oodx2, + oody2, + oodz2, + periodicity, + diag_inv, + corrf, + odx, + odz, +): shx, shy, shz = hcenter.shape - p = p0.reshape(shz+2,shy+2,shx+2) + p = p0.reshape(shz + 2, shy + 2, shx + 2) - coeff = 1./16 + coeff = 1.0 / 16 lap = np.zeros_like(p) - + # cut out four cubes from the 3d array corresponding to the nodes... in each axial direction. - toplefts = [(slice(0,None),slice(0,-1),slice(0,-1)), - (slice(0,-1),slice(0,None),slice(0,-1)), - (slice(0,-1),slice(0,-1),slice(0,None)) - ] - toprights = [(slice(0,None),slice(0,-1),slice(1,None)), - (slice(1,None),slice(0,None),slice(0,-1)), - (slice(0,-1),slice(1,None),slice(0,None)) - ] - - botlefts = [(slice(0,None),slice(1,None),slice(0,-1)), - (slice(0,-1),slice(0,None),slice(1,None)), - (slice(1,None),slice(0,-1),slice(0,None)) - ] - botrights = [(slice(0,None),slice(1,None),slice(1,None)), - (slice(1,None),slice(0,None),slice(1,None)), - (slice(1,None),slice(1,None),slice(0,None)) - ] + toplefts = [ + (slice(0, None), slice(0, -1), slice(0, -1)), + (slice(0, -1), slice(0, None), slice(0, -1)), + (slice(0, -1), slice(0, -1), slice(0, None)), + ] + toprights = [ + (slice(0, None), slice(0, -1), slice(1, None)), + (slice(1, None), slice(0, None), slice(0, -1)), + (slice(0, -1), slice(1, None), slice(0, None)), + ] + + botlefts = [ + (slice(0, None), slice(1, None), slice(0, -1)), + (slice(0, -1), slice(0, None), slice(1, None)), + (slice(1, None), slice(0, -1), slice(0, None)), + ] + botrights = [ + (slice(0, None), slice(1, None), slice(1, None)), + (slice(1, None), slice(0, None), slice(1, None)), + (slice(1, None), slice(1, None), slice(0, None)), + ] cnt = 0 for bc in periodicity: if bc == True and cnt == 0: - tmp = p[1,:,:] - p[0,:,:] = p[-3,:,:] - p[-1,:,:] = p[2,:,:] - p[1,:,:] = p[-2,:,:] - p[-2,:,:] = tmp + tmp = p[1, :, :] + p[0, :, :] = p[-3, :, :] + p[-1, :, :] = p[2, :, :] + p[1, :, :] = p[-2, :, :] + p[-2, :, :] = tmp elif bc == False and cnt == 0: - hplusx[0,:,:] = 0.0 - hplusx[-1,:,:] = 0.0 - hplusy[0,:,:] = 0.0 - hplusy[-1,:,:] = 0.0 - hplusz[0,:,:] = 0.0 - hplusz[-1,:,:] = 0.0 + hplusx[0, :, :] = 0.0 + hplusx[-1, :, :] = 0.0 + hplusy[0, :, :] = 0.0 + hplusy[-1, :, :] = 0.0 + hplusz[0, :, :] = 0.0 + hplusz[-1, :, :] = 0.0 if bc == True and cnt == 1: - tmp = p[:,1,:] - p[:,0,:] = p[:,-3,:] - p[:,-1,:] = p[:,2,:] - p[:,1,:] = p[:,-2,:] - p[:,-2,:] = tmp + tmp = p[:, 1, :] + p[:, 0, :] = p[:, -3, :] + p[:, -1, :] = p[:, 2, :] + p[:, 1, :] = p[:, -2, :] + p[:, -2, :] = tmp elif bc == False and cnt == 1: - hplusx[:,0,:] = 0.0 - hplusx[:,-1,:] = 0.0 - hplusy[:,0,:] = 0.0 - hplusy[:,-1,:] = 0.0 - hplusz[:,0,:] = 0.0 - hplusz[:,-1,:] = 0.0 + hplusx[:, 0, :] = 0.0 + hplusx[:, -1, :] = 0.0 + hplusy[:, 0, :] = 0.0 + hplusy[:, -1, :] = 0.0 + hplusz[:, 0, :] = 0.0 + hplusz[:, -1, :] = 0.0 if bc == True and cnt == 2: - tmp = p[:,:,1] - p[:,:,0] = p[:,:,-3] - p[:,:,-1] = p[:,:,2] - p[:,:,1] = p[:,:,-2] - p[:,:,-2] = tmp - elif bc == False and cnt ==2: - hplusx[:,:,0] = 0.0 - hplusx[:,:,-1] = 0.0 - hplusy[:,:,0] = 0.0 - hplusy[:,:,-1] = 0.0 - hplusz[:,:,0] = 0.0 - hplusz[:,:,-1] = 0.0 + tmp = p[:, :, 1] + p[:, :, 0] = p[:, :, -3] + p[:, :, -1] = p[:, :, 2] + p[:, :, 1] = p[:, :, -2] + p[:, :, -2] = tmp + elif bc == False and cnt == 2: + hplusx[:, :, 0] = 0.0 + hplusx[:, :, -1] = 0.0 + hplusy[:, :, 0] = 0.0 + hplusy[:, :, -1] = 0.0 + hplusz[:, :, 0] = 0.0 + hplusz[:, :, -1] = 0.0 cnt += 1 - leftz = p[:,:,:-1] - rightz = p[:,:,1:] - - z_fluxes = (rightz - leftz) + leftz = p[:, :, :-1] + rightz = p[:, :, 1:] + + z_fluxes = rightz - leftz - lefty = p[:,:-1,:] - righty = p[:,1:,:] + lefty = p[:, :-1, :] + righty = p[:, 1:, :] - y_fluxes = (righty - lefty) + y_fluxes = righty - lefty - leftx = p[:-1,:,:] - rightx = p[1:,:,:] + leftx = p[:-1, :, :] + rightx = p[1:, :, :] - x_fluxes = (rightx - leftx) + x_fluxes = rightx - leftx - x_flx = x_fluxes[toplefts[0]] + x_fluxes[toprights[0]] + x_fluxes[botlefts[0]] + x_fluxes[botrights[0]] - y_flx = y_fluxes[toplefts[1]] + y_fluxes[toprights[1]] + y_fluxes[botlefts[1]] + y_fluxes[botrights[1]] - z_flx = z_fluxes[toplefts[2]] + z_fluxes[toprights[2]] + z_fluxes[botlefts[2]] + z_fluxes[botrights[2]] + x_flx = ( + x_fluxes[toplefts[0]] + + x_fluxes[toprights[0]] + + x_fluxes[botlefts[0]] + + x_fluxes[botrights[0]] + ) + y_flx = ( + y_fluxes[toplefts[1]] + + y_fluxes[toprights[1]] + + y_fluxes[botlefts[1]] + + y_fluxes[botrights[1]] + ) + z_flx = ( + z_fluxes[toplefts[2]] + + z_fluxes[toprights[2]] + + z_fluxes[botlefts[2]] + + z_fluxes[botrights[2]] + ) hxzp = hplusx * z_flx - hxzpm = hxzp[:-1,:,:] - hxzpm = hxzpm[toplefts[0]] + hxzpm[toprights[0]] + hxzpm[botlefts[0]] + hxzpm[botrights[0]] - hxzpp = hxzp[1:,:,:] - hxzpp = hxzpp[toplefts[0]] + hxzpp[toprights[0]] + hxzpp[botlefts[0]] + hxzpp[botrights[0]] + hxzpm = hxzp[:-1, :, :] + hxzpm = ( + hxzpm[toplefts[0]] + + hxzpm[toprights[0]] + + hxzpm[botlefts[0]] + + hxzpm[botrights[0]] + ) + hxzpp = hxzp[1:, :, :] + hxzpp = ( + hxzpp[toplefts[0]] + + hxzpp[toprights[0]] + + hxzpp[botlefts[0]] + + hxzpp[botrights[0]] + ) hzxp = hplusz * x_flx - hzxpm = hzxp[:,:,:-1] - hzxpm = hzxpm[toplefts[2]] + hzxpm[toprights[2]] + hzxpm[botlefts[2]] + hzxpm[botrights[2]] - hzxpp = hzxp[:,:,1:] - hzxpp = hzxpp[toplefts[2]] + hzxpp[toprights[2]] + hzxpp[botlefts[2]] + hzxpp[botrights[2]] + hzxpm = hzxp[:, :, :-1] + hzxpm = ( + hzxpm[toplefts[2]] + + hzxpm[toprights[2]] + + hzxpm[botlefts[2]] + + hzxpm[botrights[2]] + ) + hzxpp = hzxp[:, :, 1:] + hzxpp = ( + hzxpp[toplefts[2]] + + hzxpp[toprights[2]] + + hzxpp[botlefts[2]] + + hzxpp[botrights[2]] + ) x_flx = hplusx * x_flx - x_flxm = x_flx[:-1,:,:] - x_flxm = x_flxm[toplefts[0]] + x_flxm[toprights[0]] + x_flxm[botlefts[0]] + x_flxm[botrights[0]] - x_flxp = x_flx[1:,:,:] - x_flxp = x_flxp[toplefts[0]] + x_flxp[toprights[0]] + x_flxp[botlefts[0]] + x_flxp[botrights[0]] + x_flxm = x_flx[:-1, :, :] + x_flxm = ( + x_flxm[toplefts[0]] + + x_flxm[toprights[0]] + + x_flxm[botlefts[0]] + + x_flxm[botrights[0]] + ) + x_flxp = x_flx[1:, :, :] + x_flxp = ( + x_flxp[toplefts[0]] + + x_flxp[toprights[0]] + + x_flxp[botlefts[0]] + + x_flxp[botrights[0]] + ) y_flx = hplusy * y_flx - y_flxm = y_flx[:,:-1,:] - y_flxm = y_flxm[toplefts[1]] + y_flxm[toprights[1]] + y_flxm[botlefts[1]] + y_flxm[botrights[1]] - y_flxp = y_flx[:,1:,:] - y_flxp = y_flxp[toplefts[1]] + y_flxp[toprights[1]] + y_flxp[botlefts[1]] + y_flxp[botrights[1]] + y_flxm = y_flx[:, :-1, :] + y_flxm = ( + y_flxm[toplefts[1]] + + y_flxm[toprights[1]] + + y_flxm[botlefts[1]] + + y_flxm[botrights[1]] + ) + y_flxp = y_flx[:, 1:, :] + y_flxp = ( + y_flxp[toplefts[1]] + + y_flxp[toprights[1]] + + y_flxp[botlefts[1]] + + y_flxp[botrights[1]] + ) z_flx = hplusz * z_flx - z_flxm = z_flx[:,:,:-1] - z_flxm = z_flxm[toplefts[2]] + z_flxm[toprights[2]] + z_flxm[botlefts[2]] + z_flxm[botrights[2]] - z_flxp = z_flx[:,:,1:] - z_flxp = z_flxp[toplefts[2]] + z_flxp[toprights[2]] + z_flxp[botlefts[2]] + z_flxp[botrights[2]] - - lap[1:-1,1:-1,1:-1] = oodx2 * coeff * (-x_flxm + x_flxp) + \ - oody2 * coeff * (-y_flxm + y_flxp) + \ - oodz2 * coeff * (-z_flxm + z_flxp) + \ - +1.0 * odx * odz * coeff * corrf * (hxzpp - hxzpm) + \ - -1.0 * odx * odz * coeff * corrf * (hzxpp - hzxpm) + \ - hcenter * p[1:-1,1:-1,1:-1] + z_flxm = z_flx[:, :, :-1] + z_flxm = ( + z_flxm[toplefts[2]] + + z_flxm[toprights[2]] + + z_flxm[botlefts[2]] + + z_flxm[botrights[2]] + ) + z_flxp = z_flx[:, :, 1:] + z_flxp = ( + z_flxp[toplefts[2]] + + z_flxp[toprights[2]] + + z_flxp[botlefts[2]] + + z_flxp[botrights[2]] + ) + + lap[1:-1, 1:-1, 1:-1] = ( + oodx2 * coeff * (-x_flxm + x_flxp) + + oody2 * coeff * (-y_flxm + y_flxp) + + oodz2 * coeff * (-z_flxm + z_flxp) + + +1.0 * odx * odz * coeff * corrf * (hxzpp - hxzpm) + + -1.0 * odx * odz * coeff * corrf * (hzxpp - hzxpm) + + hcenter * p[1:-1, 1:-1, 1:-1] + ) lap = lap * diag_inv return lap - -def stencil_hs(elem,node,mpv,ud,diag_inv,dt): - oodx2 = 1./node.dx**2 - oodz2 = 1./node.dz**2 - odx, odz = 1./node.dx, 1./node.dz +def stencil_hs(elem, node, mpv, ud, diag_inv, dt): + oodx2 = 1.0 / node.dx**2 + oodz2 = 1.0 / node.dz**2 + odx, odz = 1.0 / node.dx, 1.0 / node.dz igy = elem.igs[1] - proj = (slice(None,),igy,slice(None,)) - i0 = (slice(0,-1),slice(0,-1)) - i1 = (slice(1,-1),slice(1,-1)) - i2 = (slice(2,-2),slice(2,-2)) + proj = ( + slice( + None, + ), + igy, + slice( + None, + ), + ) + i0 = (slice(0, -1), slice(0, -1)) + i1 = (slice(1, -1), slice(1, -1)) + i2 = (slice(2, -2), slice(2, -2)) ndim = elem.ndim - periodicity = np.empty(ndim, dtype='int') - for dim in range(0,ndim,2): + periodicity = np.empty(ndim, dtype="int") + for dim in range(0, ndim, 2): periodicity[dim] = ud.bdry_type[dim] == opts.BdryType.PERIODIC hplusx = mpv.wplus[0][proj][i0][i1] @@ -1050,106 +1346,111 @@ def stencil_hs(elem,node,mpv,ud,diag_inv,dt): corrf = dt * ud.coriolis_strength[1] - return lambda p : lapHS(p, hplusx, hplusz, hcenter, oodx2, oodz2, periodicity, diag_inv, corrf, odx, odz) + return lambda p: lapHS( + p, hplusx, hplusz, hcenter, oodx2, oodz2, periodicity, diag_inv, corrf, odx, odz + ) @nb.jit(nopython=True, cache=True, nogil=True) -def lapHS(p0, hplusx, hplusz, hcenter, oodx2, oodz2, periodicity, diag_inv, corrf, odx, odz): +def lapHS( + p0, hplusx, hplusz, hcenter, oodx2, oodz2, periodicity, diag_inv, corrf, odx, odz +): shx, shz = hcenter.shape - p = p0.reshape(shx+2,shz+2) + p = p0.reshape(shx + 2, shz + 2) - coeff = 1./4 + coeff = 1.0 / 4 lap = np.zeros_like(p) - + cnt = 0 for bc in periodicity: if bc == True and cnt == 0: - tmp = p[1,:] - p[0,:] = p[-3,:] - p[-1,:] = p[2,:] - p[1,:] = p[-2,:] - p[-2,:] = tmp + tmp = p[1, :] + p[0, :] = p[-3, :] + p[-1, :] = p[2, :] + p[1, :] = p[-2, :] + p[-2, :] = tmp # p[0,:] = p[-4,:] # p[-1,:] = p[3,:] # p[1,:] = p[-3,:] # p[-2,:] = p[2,:] elif bc == False and cnt == 0: - hplusx[0,:] = 0.0 - hplusx[-1,:] = 0.0 - hplusz[0,:] = 0.0 - hplusz[-1,:] = 0.0 + hplusx[0, :] = 0.0 + hplusx[-1, :] = 0.0 + hplusz[0, :] = 0.0 + hplusz[-1, :] = 0.0 if bc == True and cnt == 2: - tmp = p[:,1] - p[:,0] = p[:,-3] - p[:,-1] = p[:,2] - p[:,1] = p[:,-2] - p[:,-2] = tmp + tmp = p[:, 1] + p[:, 0] = p[:, -3] + p[:, -1] = p[:, 2] + p[:, 1] = p[:, -2] + p[:, -2] = tmp # p[:,0] = p[:,-4] # p[:,-1] = p[:,3] # p[:,1] = p[:,-3] # p[:,-2] = p[:,2] elif bc == False and cnt == 2: - hplusx[:,0] = 0.0 - hplusx[:,-1] = 0.0 - hplusz[:,0] = 0.0 - hplusz[:,-1] = 0.0 + hplusx[:, 0] = 0.0 + hplusx[:, -1] = 0.0 + hplusz[:, 0] = 0.0 + hplusz[:, -1] = 0.0 cnt += 1 - - leftz = p[:,:-1] - rightz = p[:,1:] - - z_fluxes = (rightz - leftz) - leftx = p[:-1,:] - rightx = p[1:,:] + leftz = p[:, :-1] + rightz = p[:, 1:] + + z_fluxes = rightz - leftz - x_fluxes = (rightx - leftx) + leftx = p[:-1, :] + rightx = p[1:, :] - x_flx = x_fluxes[:,:-1] + x_fluxes[:,1:] - z_flx = z_fluxes[:-1,:] + z_fluxes[1:,:] + x_fluxes = rightx - leftx + + x_flx = x_fluxes[:, :-1] + x_fluxes[:, 1:] + z_flx = z_fluxes[:-1, :] + z_fluxes[1:, :] hxzp = hplusx * z_flx - hxzpm = hxzp[:-1,:] - hxzpm = hxzpm[:,:-1] + hxzpm[:,1:] + hxzpm = hxzp[:-1, :] + hxzpm = hxzpm[:, :-1] + hxzpm[:, 1:] # hxzpm = hxzpm[:-1,:] + hxzpm[1:,:] - hxzpp = hxzp[1:,:] - hxzpp = hxzpp[:,:-1] + hxzpp[:,1:] + hxzpp = hxzp[1:, :] + hxzpp = hxzpp[:, :-1] + hxzpp[:, 1:] # hxzpp = hxzpp[:-1,:] + hxzpp[1:,:] hzxp = hplusz * x_flx - hzxpm = hzxp[:,:-1] - hzxpm = hzxpm[:-1,:] + hzxpm[1:,:] + hzxpm = hzxp[:, :-1] + hzxpm = hzxpm[:-1, :] + hzxpm[1:, :] # hzxpm = hzxpm[:,:-1] + hzxpm[:,1:] - hzxpp = hzxp[:,1:] - hzxpp = hzxpp[:-1,:] + hzxpp[1:,:] + hzxpp = hzxp[:, 1:] + hzxpp = hzxpp[:-1, :] + hzxpp[1:, :] # hzxpp = hzxpp[:,:-1] + hzxpp[:,1:] x_flx = hplusx * x_flx - x_flxm = x_flx[:-1,:] - x_flxm = x_flxm[:,:-1] + x_flxm[:,1:] - x_flxp = x_flx[1:,:] - x_flxp = x_flxp[:,:-1] + x_flxp[:,1:] + x_flxm = x_flx[:-1, :] + x_flxm = x_flxm[:, :-1] + x_flxm[:, 1:] + x_flxp = x_flx[1:, :] + x_flxp = x_flxp[:, :-1] + x_flxp[:, 1:] z_flx = hplusz * z_flx - z_flxm = z_flx[:,:-1] - z_flxm = z_flxm[:-1,:] + z_flxm[1:,:] - z_flxp = z_flx[:,1:] - z_flxp = z_flxp[:-1,:] + z_flxp[1:,:] - - lap[1:-1,1:-1] = oodx2 * coeff * (-x_flxm + x_flxp) + \ - oodz2 * coeff * (-z_flxm + z_flxp) + \ - +1.0 * odx * odz * coeff * corrf * (hxzpp - hxzpm) + \ - -1.0 * odx * odz * coeff * corrf * (hzxpp - hzxpm) + \ - hcenter * p[1:-1,1:-1] + z_flxm = z_flx[:, :-1] + z_flxm = z_flxm[:-1, :] + z_flxm[1:, :] + z_flxp = z_flx[:, 1:] + z_flxp = z_flxp[:-1, :] + z_flxp[1:, :] + + lap[1:-1, 1:-1] = ( + oodx2 * coeff * (-x_flxm + x_flxp) + + oodz2 * coeff * (-z_flxm + z_flxp) + + +1.0 * odx * odz * coeff * corrf * (hxzpp - hxzpm) + + -1.0 * odx * odz * coeff * corrf * (hzxpp - hzxpm) + + hcenter * p[1:-1, 1:-1] + ) lap = lap * diag_inv return lap - -def stencil_vs(elem,node,mpv,ud,diag_inv,dt): +def stencil_vs(elem, node, mpv, ud, diag_inv, dt): igx = elem.igx igy = elem.igy igz = elem.igz @@ -1170,10 +1471,18 @@ def stencil_vs(elem,node,mpv,ud,diag_inv,dt): xx, yy = 1, 0 - proj = (slice(None,),slice(None,),igz) - i0 = (slice(0,-1),slice(0,-1)) - i1 = (slice(1,-1),slice(1,-1)) - i2 = (slice(igx,-igx),slice(igy,-igy)) + proj = ( + slice( + None, + ), + slice( + None, + ), + igz, + ) + i0 = (slice(0, -1), slice(0, -1)) + i1 = (slice(1, -1), slice(1, -1)) + i2 = (slice(igx, -igx), slice(igy, -igy)) VS = True @@ -1184,20 +1493,44 @@ def stencil_vs(elem,node,mpv,ud,diag_inv,dt): x_wall = ud.bdry_type[xx] == opts.BdryType.WALL y_wall = ud.bdry_type[yy] == opts.BdryType.WALL - hplusx = mpv.wplus[xx][proj][i2].reshape(-1,) - hplusy = mpv.wplus[yy][proj][i2].reshape(-1,) - hcenter = mpv.wcenter[proj][i2].reshape(-1,) - diag_inv = diag_inv[proj][i2].reshape(-1,) - - return lambda p : lap2D(p, igx,igy, iicxn, iicyn, hplusx, hplusy, hcenter, oodx2, oody2, x_periodic, y_periodic, x_wall, y_wall, diag_inv) + hplusx = mpv.wplus[xx][proj][i2].reshape( + -1, + ) + hplusy = mpv.wplus[yy][proj][i2].reshape( + -1, + ) + hcenter = mpv.wcenter[proj][i2].reshape( + -1, + ) + diag_inv = diag_inv[proj][i2].reshape( + -1, + ) + + return lambda p: lap2D( + p, + igx, + igy, + iicxn, + iicyn, + hplusx, + hplusy, + hcenter, + oodx2, + oody2, + x_periodic, + y_periodic, + x_wall, + y_wall, + diag_inv, + ) if VS: - oodx2 = 1./node.dx**2 - oody2 = 1./node.dy**2 + oodx2 = 1.0 / node.dx**2 + oody2 = 1.0 / node.dy**2 ndim = elem.ndim - periodicity = np.zeros(ndim, dtype='int') - for dim in range(0,ndim,2): + periodicity = np.zeros(ndim, dtype="int") + for dim in range(0, ndim, 2): periodicity[dim] = ud.bdry_type[dim] == opts.BdryType.PERIODIC hplusx = mpv.wplus[0][proj][i0][i1] @@ -1205,82 +1538,85 @@ def stencil_vs(elem,node,mpv,ud,diag_inv,dt): hcenter = mpv.wcenter[proj][i2] diag_inv = diag_inv[proj][i1] - return lambda p : lapVS(p, hplusx, hplusy, hcenter, oodx2, oody2, periodicity, diag_inv) + return lambda p: lapVS( + p, hplusx, hplusy, hcenter, oodx2, oody2, periodicity, diag_inv + ) @nb.jit(nopython=True, cache=True, nogil=False) def lapVS(p0, hplusx, hplusy, hcenter, oodx2, oody2, periodicity, diag_inv): shx, shy = hcenter.shape - p = p0.reshape(shx+2,shy+2) + p = p0.reshape(shx + 2, shy + 2) - coeff = 1./4 + coeff = 1.0 / 4 lap = np.zeros_like(p) - + cnt = 0 for bc in periodicity: if bc == True and cnt == 0: - p[0,:] = p[-3,:] - p[-1,:] = p[2,:] + p[0, :] = p[-3, :] + p[-1, :] = p[2, :] elif bc == False and cnt == 0: - hplusx[0,:] = 0.0 - hplusx[-1,:] = 0.0 - hplusy[0,:] = 0.0 - hplusy[-1,:] = 0.0 + hplusx[0, :] = 0.0 + hplusx[-1, :] = 0.0 + hplusy[0, :] = 0.0 + hplusy[-1, :] = 0.0 # p[0,:] = p[2,:] # p[-1,:] = p[-3,:] if bc == True and cnt == 1: - p[:,0] = p[:,-3] - p[:,-1] = p[:,2] + p[:, 0] = p[:, -3] + p[:, -1] = p[:, 2] elif bc == False and cnt == 1: - hplusx[:,0] = 0.0 - hplusx[:,-1] = 0.0 - hplusy[:,0] = 0.0 - hplusy[:,-1] = 0.0 + hplusx[:, 0] = 0.0 + hplusx[:, -1] = 0.0 + hplusy[:, 0] = 0.0 + hplusy[:, -1] = 0.0 # p[:,0] = p[:,2] # p[:,-1] = p[:,-3] - + cnt += 1 - - lefty = p[:,:-1] - righty = p[:,1:] - - y_fluxes = (righty - lefty) - leftx = p[:-1,:] - rightx = p[1:,:] + lefty = p[:, :-1] + righty = p[:, 1:] + + y_fluxes = righty - lefty + + leftx = p[:-1, :] + rightx = p[1:, :] - x_fluxes = (rightx - leftx) + x_fluxes = rightx - leftx - x_flx = x_fluxes[:,:-1] + x_fluxes[:,1:] - y_flx = y_fluxes[:-1,:] + y_fluxes[1:,:] + x_flx = x_fluxes[:, :-1] + x_fluxes[:, 1:] + y_flx = y_fluxes[:-1, :] + y_fluxes[1:, :] x_flx = hplusx * x_flx - x_flxm = x_flx[:-1,:] - x_flxm = x_flxm[:,:-1] + x_flxm[:,1:] - x_flxp = x_flx[1:,:] - x_flxp = x_flxp[:,:-1] + x_flxp[:,1:] + x_flxm = x_flx[:-1, :] + x_flxm = x_flxm[:, :-1] + x_flxm[:, 1:] + x_flxp = x_flx[1:, :] + x_flxp = x_flxp[:, :-1] + x_flxp[:, 1:] y_flx = hplusy * y_flx - y_flxm = y_flx[:,:-1] - y_flxm = y_flxm[:-1,:] + y_flxm[1:,:] - y_flxp = y_flx[:,1:] - y_flxp = y_flxp[:-1,:] + y_flxp[1:,:] - - lap[1:-1,1:-1] = oodx2 * coeff * (-x_flxm + x_flxp) + \ - oody2 * coeff * (-y_flxm + y_flxp) + \ - hcenter * p[1:-1,1:-1] - + y_flxm = y_flx[:, :-1] + y_flxm = y_flxm[:-1, :] + y_flxm[1:, :] + y_flxp = y_flx[:, 1:] + y_flxp = y_flxp[:-1, :] + y_flxp[1:, :] + + lap[1:-1, 1:-1] = ( + oodx2 * coeff * (-x_flxm + x_flxp) + + oody2 * coeff * (-y_flxm + y_flxp) + + hcenter * p[1:-1, 1:-1] + ) + lap = lap * diag_inv return lap - def precon_diag_prepare(mpv, elem, node, ud, coriolis): dx, dy, dz = node.dx, node.dy, node.dz @@ -1299,24 +1635,27 @@ def precon_diag_prepare(mpv, elem, node, ud, coriolis): for dim in range(ndim): if ud.bdry_type[dim] == opts.BdryType.PERIODIC: - idx_periodic[dim] = slice(1,-1) + idx_periodic[dim] = slice(1, -1) - idx_e[dim] = slice(igs[dim] - periodicity[dim], -igs[dim] + periodicity[dim] - 1) + idx_e[dim] = slice( + igs[dim] - periodicity[dim], -igs[dim] + periodicity[dim] - 1 + ) idx_n[dim] = slice(igs[dim], -igs[dim]) idx_periodic, idx_e, idx_n = tuple(idx_periodic), tuple(idx_e), tuple(idx_n) # idx_n = (slice(igx,-igx), slice(igy,-igy)) - hplusxx = mpv.wplus[0] #* (coriolis[0].T) - hplusyy = mpv.wplus[1] #* (coriolis[1].T) - hplusxy = mpv.wplus[0] #* (coriolis[3].T) - hplusyx = mpv.wplus[1] #* (coriolis[2].T) + hplusxx = mpv.wplus[0] # * (coriolis[0].T) + hplusyy = mpv.wplus[1] # * (coriolis[1].T) + hplusxy = mpv.wplus[0] # * (coriolis[3].T) + hplusyx = mpv.wplus[1] # * (coriolis[2].T) - if ndim == 3: hplusz = mpv.wplus[2] + if ndim == 3: + hplusz = mpv.wplus[2] nine_pt = 0.25 * (2.0) * 1.0 - nine_pt = 0.5 * 0.5 + nine_pt = 0.5 * 0.5 if ndim == 2: coeff = 1.0 - nine_pt elif ndim == 3: @@ -1334,7 +1673,7 @@ def precon_diag_prepare(mpv, elem, node, ud, coriolis): diag_kernel = np.array(np.ones([2] * ndim)) # diag = np.zeros((node.sc)).squeeze() - # diag[idx_n] = -wx * signal.fftconvolve(hplusx[idx_e],diag_kernel,mode='full')[idx_periodic] + # diag[idx_n] = -wx * signal.fftconvolve(hplusx[idx_e],diag_kernel,mode='full')[idx_periodic] # diag[idx_n] -= wy * signal.fftconvolve(hplusy[idx_e],diag_kernel,mode='full')[idx_periodic] # if ndim == 3: # diag[idx_n] -= wz * signal.fftconvolve(hplusz[idx_e],diag_kernel,mode='full')[idx_periodic] @@ -1343,12 +1682,12 @@ def precon_diag_prepare(mpv, elem, node, ud, coriolis): # diag[idx_n] = 1.0 / diag[idx_n] diag = np.zeros_like(mpv.wcenter) - diag[...] = -wxx * sp.signal.fftconvolve(hplusxx,diag_kernel,mode='valid') - diag[...] -= wyy * sp.signal.fftconvolve(hplusyy,diag_kernel,mode='valid') - diag[...] -= wxy * sp.signal.fftconvolve(hplusxy,diag_kernel,mode='valid') - diag[...] -= wyx * sp.signal.fftconvolve(hplusyx,diag_kernel,mode='valid') + diag[...] = -wxx * sp.signal.fftconvolve(hplusxx, diag_kernel, mode="valid") + diag[...] -= wyy * sp.signal.fftconvolve(hplusyy, diag_kernel, mode="valid") + diag[...] -= wxy * sp.signal.fftconvolve(hplusxy, diag_kernel, mode="valid") + diag[...] -= wyx * sp.signal.fftconvolve(hplusyx, diag_kernel, mode="valid") if ndim == 3: - diag[...] -= wzz * sp.signal.fftconvolve(hplusz,diag_kernel,mode='valid') + diag[...] -= wzz * sp.signal.fftconvolve(hplusz, diag_kernel, mode="valid") diag[...] += mpv.wcenter diag[...] = 1.0 / diag @@ -1356,10 +1695,7 @@ def precon_diag_prepare(mpv, elem, node, ud, coriolis): return diag - - - -def stencil_3pt(elem,node,ud): +def stencil_3pt(elem, node, ud): # 1d stencil for exner pressure perturbation constraint in 2d igx = elem.igx igy = elem.igy @@ -1375,13 +1711,14 @@ def stencil_3pt(elem,node,ud): dx = node.dx # print(node.dx) - oodx2 = 1. / (dx**2) + oodx2 = 1.0 / (dx**2) # x_periodic = ud.bdry_type[1] == BdryType.PERIODIC x_periodic = ud.bdry_type[0] == opts.BdryType.PERIODIC x_wall = ud.bdry_type[0] == opts.BdryType.WALL - return lambda p : lap2D_exner(p,iicxn, iicyn, oodx2, x_periodic, x_wall) + return lambda p: lap2D_exner(p, iicxn, iicyn, oodx2, x_periodic, x_wall) + @nb.jit(nopython=True, nogil=True, cache=True) def lap2D_exner(p, iicxn, iicyn, oodx2, x_periodic, x_wall): @@ -1413,29 +1750,29 @@ def lap2D_exner(p, iicxn, iicyn, oodx2, x_periodic, x_wall): if cnt_y == 0: if x_periodic: - right_idx += ((iicxn) * (iicyn - 1)) - mid_idx += ((iicxn) * (iicyn - 1)) + right_idx += (iicxn) * (iicyn - 1) + mid_idx += (iicxn) * (iicyn - 1) if cnt_y == (iicyn - 1): if x_periodic: - left_idx -= ((iicxn) * (iicyn - 1)) - mid_idx -= ((iicxn) * (iicyn - 1)) + left_idx -= (iicxn) * (iicyn - 1) + mid_idx -= (iicxn) * (iicyn - 1) left = p[left_idx] mid = p[mid_idx] right = p[right_idx] - lap[idx] = oodx2 * (left - 2. * mid + right) + lap[idx] = oodx2 * (left - 2.0 * mid + right) cnt_x += 1 if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - + return lap -def stencil_5pt(elem,node,ud): +def stencil_5pt(elem, node, ud): igx = elem.igx igy = elem.igy @@ -1462,7 +1799,10 @@ def stencil_5pt(elem,node,ud): x_wall = ud.bdry_type[1] == opts.BdryType.WALL y_wall = ud.bdry_type[0] == opts.BdryType.WALL - return lambda p : lap2D_5pt(p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_wall) + return lambda p: lap2D_5pt( + p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_wall + ) + @nb.jit(nopython=True, nogil=True, cache=True) def lap2D_5pt(p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_wall): @@ -1472,7 +1812,6 @@ def lap2D_5pt(p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_w cnt_y = 0 for idx in range(iicxn * iicyn): - # get indices of the 9pt stencil midleft_idx = idx - 1 topmid_idx = idx - iicxn @@ -1496,20 +1835,20 @@ def lap2D_5pt(p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_w botmid_idx -= iicxn - 1 if cnt_y == 0: - topmid_idx += ((iicxn) * (iicyn - 1)) + topmid_idx += (iicxn) * (iicyn - 1) if y_periodic: - midleft_idx += ((iicxn) * (iicyn - 1)) - midmid_idx += ((iicxn) * (iicyn - 1)) - midright_idx += ((iicxn) * (iicyn - 1)) + midleft_idx += (iicxn) * (iicyn - 1) + midmid_idx += (iicxn) * (iicyn - 1) + midright_idx += (iicxn) * (iicyn - 1) if cnt_y == (iicyn - 1): - botmid_idx -= ((iicxn) * (iicyn - 1)) + botmid_idx -= (iicxn) * (iicyn - 1) if y_periodic: - midleft_idx -= ((iicxn) * (iicyn - 1)) - midmid_idx -= ((iicxn) * (iicyn - 1)) - midright_idx -= ((iicxn) * (iicyn - 1)) + midleft_idx -= (iicxn) * (iicyn - 1) + midmid_idx -= (iicxn) * (iicyn - 1) + midright_idx -= (iicxn) * (iicyn - 1) midleft = p[midleft_idx] topmid = p[topmid_idx] @@ -1518,22 +1857,24 @@ def lap2D_5pt(p, iicxn, iicyn, oodx2, oody2, x_periodic, y_periodic, x_wall, y_w midright = p[midright_idx] if x_wall and (cnt_x == 0): - midleft = 0. + midleft = 0.0 if x_wall and (cnt_x == (iicxn - 1)): - midright = 0. + midright = 0.0 if y_wall and (cnt_y == 0): - topmid = 0. - + topmid = 0.0 + if y_wall and (cnt_y == (iicyn - 1)): - botmid = 0. - - lap[idx] = oodx2 * (midleft - 2.0 * midmid + midright) + oody2 * (topmid - 2.0 * midmid + botmid) + botmid = 0.0 + + lap[idx] = oodx2 * (midleft - 2.0 * midmid + midright) + oody2 * ( + topmid - 2.0 * midmid + botmid + ) cnt_x += 1 if cnt_x % iicxn == 0: cnt_y += 1 cnt_x = 0 - - return lap \ No newline at end of file + + return lap diff --git a/src/pybella/flow_solver/physics/low_mach/mpv.py b/src/pybella/flow_solver/physics/low_mach/mpv.py new file mode 100644 index 00000000..4e08a7af --- /dev/null +++ b/src/pybella/flow_solver/physics/low_mach/mpv.py @@ -0,0 +1,36 @@ +import numpy as np + +from ...utils import variable as var + + +class MPV(object): + def __init__(self, elem, node, ud): + sc = elem.sc + sn = node.sc + + self.p0 = 1.0 + self.p00 = 1.0 + + self.p2_cells = np.zeros((sc)) + self.dp2_cells = np.zeros((sc)) + self.p2_nodes = np.zeros((sn)) + self.p2_nodes0 = np.zeros((sn)) + self.dp2_nodes = np.zeros((sn)) + + self.u = np.zeros((sc)) + self.v = np.zeros((sc)) + self.w = np.zeros((sc)) + + self.rhs = np.zeros((node.isc)) + self.wcenter = np.zeros((node.isc)) + self.wplus = np.zeros(([elem.ndim] + list(sc))) + + self.HydroState = var.States([sc[1]], ud) + self.HydroState_n = var.States([sn[1]], ud) + + self.squeezer() + + def squeezer(self): + for key, value in vars(self).items(): + if type(value) == np.ndarray: + setattr(self, key, value.squeeze()) diff --git a/src/dycore/physics/low_mach/second_projection.py b/src/pybella/flow_solver/physics/low_mach/second_projection.py similarity index 58% rename from src/dycore/physics/low_mach/second_projection.py rename to src/pybella/flow_solver/physics/low_mach/second_projection.py index 5a87793a..3f208c94 100644 --- a/src/dycore/physics/low_mach/second_projection.py +++ b/src/pybella/flow_solver/physics/low_mach/second_projection.py @@ -1,31 +1,36 @@ +import itertools as it +import logging + import numpy as np import scipy as sp -import itertools as it -import dycore.utils.options as opts -import dycore.utils.boundary as bdry -import dycore.physics.low_mach.laplacian as lm_lp +from ...utils import options as opts, boundary as bdry +from . import laplacian as lm_lp -import logging class solver_counter(object): """ taken from https://stackoverflow.com/questions/33512081/getting-the-number-of-iterations-of-scipys-gmres-iterative-method """ + def __init__(self, disp=True): self.niter = 0 + def __call__(self, rk=None): self.niter += 1 self.rk = rk -def euler_forward_non_advective(Sol, mpv, elem, node, dt, ud, th, writer = None, label = None, debug = False): + +def euler_forward_non_advective( + Sol, mpv, elem, node, dt, ud, th, writer=None, label=None, debug=False +): nonhydro = ud.nonhydrostasy g = ud.gravity_strength[1] Msq = ud.Msq Ginv = th.Gammainv corr_h1 = ud.coriolis_strength[0] - corr_v = ud.coriolis_strength[1] + corr_v = ud.coriolis_strength[1] corr_h2 = ud.coriolis_strength[2] u0 = ud.u_wind_speed v0 = ud.v_wind_speed @@ -36,17 +41,23 @@ def euler_forward_non_advective(Sol, mpv, elem, node, dt, ud, th, writer = None, ndim = elem.ndim S0c = mpv.HydroState.get_S0c(elem) - dSdy = mpv.HydroState_n.get_dSdy(elem,node) + dSdy = mpv.HydroState_n.get_dSdy(elem, node) - mpv.rhs[...] = divergence_nodes(mpv.rhs,elem,node,Sol,ud) - if not hasattr(ud, 'LAMB_BDRY'): scale_wall_node_values(mpv.rhs, node, ud, 2.0) + mpv.rhs[...] = divergence_nodes(mpv.rhs, elem, node, Sol, ud) + if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): + scale_wall_node_values(mpv.rhs, node, ud, 2.0) div = mpv.rhs - if debug: writer.populate(str(label), 'rhs', div) + if debug: + writer.populate(str(label), "rhs", div) - rhoY = Sol.rhoY**(th.gamm - 2.0) - dpidP_kernel = np.ones([2]*ndim) - dpidP = (th.gm1 / ud.Msq) * sp.signal.fftconvolve(rhoY, dpidP_kernel, mode='valid') / dpidP_kernel.sum() + rhoY = Sol.rhoY ** (th.gamm - 2.0) + dpidP_kernel = np.ones([2] * ndim) + dpidP = ( + (th.gm1 / ud.Msq) + * sp.signal.fftconvolve(rhoY, dpidP_kernel, mode="valid") + / dpidP_kernel.sum() + ) rhoYovG = Ginv * Sol.rhoY dbuoy = Sol.rhoY * (Sol.rhoX / Sol.rho) @@ -56,21 +67,28 @@ def euler_forward_non_advective(Sol, mpv, elem, node, dt, ud, th, writer = None, drhov = Sol.rhov - v0 * Sol.rho drhow = Sol.rhow - w0 * Sol.rho v = Sol.rhov / Sol.rho - - Sol.rhou = Sol.rhou - dt * ( rhoYovG * dpdx - corr_h2 * drhov + corr_v * drhow) - Sol.rhov = Sol.rhov - dt * ( rhoYovG * dpdy + (g/Msq) * dbuoy * nonhydro - corr_h1 * drhow + corr_h2 * drhou) * (1 - ud.is_ArakawaKonor) + Sol.rhou = Sol.rhou - dt * (rhoYovG * dpdx - corr_h2 * drhov + corr_v * drhow) + + Sol.rhov = Sol.rhov - dt * ( + rhoYovG * dpdy + + (g / Msq) * dbuoy * nonhydro + - corr_h1 * drhow + + corr_h2 * drhou + ) * (1 - ud.is_ArakawaKonor) - Sol.rhow = Sol.rhow - dt * ( rhoYovG * dpdz - corr_v * drhou + corr_h1 * drhov) * (ndim == 3) + Sol.rhow = Sol.rhow - dt * (rhoYovG * dpdz - corr_v * drhou + corr_h1 * drhov) * ( + ndim == 3 + ) Sol.rhoX = (Sol.rho * (Sol.rho / Sol.rhoY - S0c)) - dt * (v * dSdy) * Sol.rho - dp2n[node.i1] -= dt * dpidP * div#[node.i1] + dp2n[node.i1] -= dt * dpidP * div # [node.i1] weight = ud.compressibility mpv.p2_nodes += weight * dp2n - bdry.set_ghostnodes_p2(mpv.p2_nodes,node, ud) + bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) @@ -80,7 +98,7 @@ def euler_backward_non_advective_expl_part(Sol, mpv, elem, dt, ud, th): Msq = ud.Msq dbuoy = Sol.rhoY * (Sol.rhoX / Sol.rho) - Sol.rhov = (nonhydro * Sol.rhov) - dt * (g/Msq) * dbuoy + Sol.rhov = (nonhydro * Sol.rhov) - dt * (g / Msq) * dbuoy # Sol.rhov[np.where(Sol.rhov < 1e-15)] = 0.0 Sol.mod_bg_wind(ud, -1.0) @@ -94,7 +112,23 @@ def euler_backward_non_advective_expl_part(Sol, mpv, elem, dt, ud, th): total_iter = 0 total_calls = 0 -def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, alpha_diff, Sol0 = None, writer = None, label=None, debug=False): + + +def euler_backward_non_advective_impl_part( + Sol, + mpv, + elem, + node, + ud, + th, + t, + dt, + alpha_diff, + Sol0=None, + writer=None, + label=None, + debug=False, +): if not debug: writer = None nc = node.sc @@ -106,9 +140,9 @@ def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, # p2 = mpv.p2_nodes[node.p_isc] # elif elem.ndim == 3: # p2 = np.copy(mpv.p2_nodes[1:-1,1:-1,1:-1]) - + if writer != None: - writer.populate(str(label),'p2_initial',mpv.p2_nodes) + writer.populate(str(label), "p2_initial", mpv.p2_nodes) if Sol0 is not None: bdry.set_explicit_boundary_data(Sol0, elem, ud, th, mpv) @@ -116,44 +150,47 @@ def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, else: bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt) - - i0 = node.ndim * [(slice(0,-1))] + i0 = node.ndim * [(slice(0, -1))] i0 = tuple(i0) if writer != None: - writer.populate(str(label),'hcenter',mpv.wcenter) - writer.populate(str(label),'wplusx',mpv.wplus[0]) - writer.populate(str(label),'wplusy',mpv.wplus[1]) - writer.populate(str(label),'wplusz',mpv.wplus[2]) if elem.ndim == 3 else writer.populate(str(label),'wplusz',np.zeros_like(mpv.wplus[0])) - - bdry.set_ghostnodes_p2(mpv.p2_nodes,node,ud) - correction_nodes(Sol,elem,node,mpv,mpv.p2_nodes,dt,ud,th,0) + writer.populate(str(label), "hcenter", mpv.wcenter) + writer.populate(str(label), "wplusx", mpv.wplus[0]) + writer.populate(str(label), "wplusy", mpv.wplus[1]) + writer.populate( + str(label), "wplusz", mpv.wplus[2] + ) if elem.ndim == 3 else writer.populate( + str(label), "wplusz", np.zeros_like(mpv.wplus[0]) + ) + + bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) + correction_nodes(Sol, elem, node, mpv, mpv.p2_nodes, dt, ud, th, 0) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - rhs[...] = divergence_nodes(rhs,elem,node,Sol,ud) + rhs[...] = divergence_nodes(rhs, elem, node, Sol, ud) if writer != None: - writer.populate(str(label),'rhs',rhs) + writer.populate(str(label), "rhs", rhs) rhs /= dt if ud.is_compressible == 1: - rhs = rhs_from_p_old(rhs,node,mpv) - # if + rhs = rhs_from_p_old(rhs, node, mpv) + # if elif ud.is_compressible == 0: if ud.is_ArakawaKonor: rhs -= mpv.wcenter * mpv.dp2_nodes mpv.wcenter[...] = 0.0 else: - rhs_new = rhs_from_p_old(rhs,node,mpv) + rhs_new = rhs_from_p_old(rhs, node, mpv) rhs = ud.compressibility * rhs_new + (1.0 - ud.compressibility) * rhs mpv.wcenter[...] *= ud.compressibility else: mpv.wcenter *= ud.compressibility if writer != None: - writer.populate(str(label),'rhs_nodes',rhs) + writer.populate(str(label), "rhs_nodes", rhs) mpv.rhs[...] = rhs @@ -162,7 +199,9 @@ def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, # prepare initial left-hand side and the laplacian stencil if elem.ndim == 2: Vec = mpv - coriolis_params = multiply_inverse_coriolis(Vec, Sol, mpv, ud, elem, node, dt, attrs=('u', 'v', 'w'), get_coeffs = True) + coriolis_params = multiply_inverse_coriolis( + Vec, Sol, mpv, ud, elem, node, dt, attrs=("u", "v", "w"), get_coeffs=True + ) # lap = stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params) # sh = (ud.inx)*(ud.iny) @@ -172,62 +211,65 @@ def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, # diag_inv = np.ones_like(mpv.rhs) p2 = mpv.p2_nodes[node.i2].T - lap = lm_lp.stencil_9pt_numba_test(mpv,node,coriolis_params,diag_inv, ud) + lap = lm_lp.stencil_9pt_numba_test(mpv, node, coriolis_params, diag_inv, ud) sh = p2.shape[0] * p2.shape[1] # p2 = np.copy(mpv.p2_nodes[1:-1,1:-1]) # sh = p2.reshape(-1).shape[0] - - - elif elem.ndim == 3 and elem.icy - 2*elem.igs[1] <= 2: + elif elem.ndim == 3 and elem.icy - 2 * elem.igs[1] <= 2: # horizontal slice hack - p2 = np.copy(mpv.p2_nodes[1:-1,elem.igs[1],1:-1]) - lap = lm_lp.stencil_hs(elem,node,mpv,ud,diag_inv,dt) + p2 = np.copy(mpv.p2_nodes[1:-1, elem.igs[1], 1:-1]) + lap = lm_lp.stencil_hs(elem, node, mpv, ud, diag_inv, dt) sh = p2.reshape(-1).shape[0] elif elem.ndim == 3 and elem.iicy > 1 and elem.iicz == 1: # vertical slice hack if not VS: - p2 = np.copy(mpv.p2_nodes[...,elem.igz][node.igx:-node.igx,node.igy:-node.igy]) - lap = lm_lp.stencil_vs(elem,node,mpv,ud,diag_inv,dt) - sh = (node.iicx)*(node.iicy) + p2 = np.copy( + mpv.p2_nodes[..., elem.igz][node.igx : -node.igx, node.igy : -node.igy] + ) + lap = lm_lp.stencil_vs(elem, node, mpv, ud, diag_inv, dt) + sh = (node.iicx) * (node.iicy) if VS: - p2 = np.copy(mpv.p2_nodes[1:-1,1:-1,elem.igs[2]]) - lap = lm_lp.stencil_vs(elem,node,mpv,ud,diag_inv,dt) + p2 = np.copy(mpv.p2_nodes[1:-1, 1:-1, elem.igs[2]]) + lap = lm_lp.stencil_vs(elem, node, mpv, ud, diag_inv, dt) sh = p2.reshape(-1).shape[0] - elif elem.ndim == 3 and elem.icy - 2*elem.igs[1] > 2: - lap = lm_lp.stencil_27pt(elem,node,mpv,ud,diag_inv,dt) + elif elem.ndim == 3 and elem.icy - 2 * elem.igs[1] > 2: + lap = lm_lp.stencil_27pt(elem, node, mpv, ud, diag_inv, dt) sh = p2.reshape(-1).shape[0] - lap = sp.sparse.linalg.LinearOperator((sh,sh),lap) + lap = sp.sparse.linalg.LinearOperator((sh, sh), lap) # lap = LinearOperator(sh,lap) - + counter = solver_counter() # prepare right-hand side if elem.ndim == 2: # rhs_inner = rhs[node.igx:-node.igx,node.igy:-node.igy].ravel() # rhs_inner = rhs[1:-1,1:-1].ravel() - rhs_inner = rhs[1:-1,1:-1].T.ravel() + rhs_inner = rhs[1:-1, 1:-1].T.ravel() # rhs_inner = rhs.T.ravel() elif elem.ndim == 3 and elem.iicy > 1 and elem.iicz == 1: - if not VS: - rhs_inner = rhs[...,elem.igs[2]][node.igx:-node.igx,node.igy:-node.igy].ravel() + rhs_inner = rhs[..., elem.igs[2]][ + node.igx : -node.igx, node.igy : -node.igy + ].ravel() if VS: - rhs_inner = rhs[1:-1,1:-1,elem.igs[2]].ravel() + rhs_inner = rhs[1:-1, 1:-1, elem.igs[2]].ravel() - elif elem.ndim == 3 and elem.icy - 2*elem.igs[1] > 2: - rhs_inner = rhs[1:-1,1:-1,1:-1].ravel() + elif elem.ndim == 3 and elem.icy - 2 * elem.igs[1] > 2: + rhs_inner = rhs[1:-1, 1:-1, 1:-1].ravel() else: - rhs_inner = rhs[1:-1,elem.igs[1],1:-1].ravel() + rhs_inner = rhs[1:-1, elem.igs[1], 1:-1].ravel() # p2, _ = bicgstab(lap,rhs_inner,tol=ud.tol,maxiter=ud.max_iterations,callback=counter) - p2, _ = sp.sparse.linalg.bicgstab(lap,rhs_inner,tol=ud.tol,maxiter=ud.max_iterations,callback=counter) + p2, _ = sp.sparse.linalg.bicgstab( + lap, rhs_inner, atol=ud.tol, maxiter=ud.max_iterations, callback=counter + ) # p2, _ = gmres(lap,rhs_inner,tol=ud.tol,maxiter=ud.max_iterations) # p2,info = bicgstab(lap,rhs.ravel(),x0=p2.ravel(),tol=1e-16,maxiter=6000,callback=counter) # print("Convergence info = %i, no. of iterations = %i" %(info,counter.niter)) @@ -236,48 +278,50 @@ def euler_backward_non_advective_impl_part(Sol, mpv, elem, node, ud, th, t, dt, total_iter += counter.niter total_calls += 1 logging.info(counter.niter) - logging.info("Total calls to BiCGStab routine = %i, total iterations = %i" %(total_calls, total_iter)) + logging.info( + "Total calls to BiCGStab routine = %i, total iterations = %i" + % (total_calls, total_iter) + ) p2_full = np.zeros(nc).squeeze() if elem.ndim == 2: # p2_full[node.igx:-node.igx,node.igy:-node.igy] = p2.reshape(ud.inx,ud.iny) # p2_full[node.i2] = p2.reshape(rhs[node.i1].shape[0],rhs[node.i1].shape[1]) - p2_full[node.i2] = p2.reshape(rhs[node.i1].shape[1],rhs[node.i1].shape[0]).T + p2_full[node.i2] = p2.reshape(rhs[node.i1].shape[1], rhs[node.i1].shape[0]).T # p2_full[node.i1] = p2.reshape(rhs.shape[1],rhs.shape[0]).T # p2 = p2.reshape(ud.inx+2, ud.iny+2) # p2_full[1:-1,1:-1] = p2 - elif elem.ndim == 3 and elem.icy - 2*elem.igs[1] <= 2: + elif elem.ndim == 3 and elem.icy - 2 * elem.igs[1] <= 2: # horizontal slice hack - p2 = p2.reshape(ud.inx+2, ud.inz+2) - p2 = np.expand_dims(p2,1) + p2 = p2.reshape(ud.inx + 2, ud.inz + 2) + p2 = np.expand_dims(p2, 1) p2 = np.repeat(p2, node.icy, axis=1) - p2_full[1:-1,:,1:-1] = p2 + p2_full[1:-1, :, 1:-1] = p2 elif elem.ndim == 3 and elem.iicy > 1 and elem.iicz == 1: if not VS: - p2 = p2.reshape(node.iicx,node.iicy) - p2 = np.repeat(p2[...,np.newaxis], node.icz, axis=2) - p2_full[node.igx:-node.igx,node.igy:-node.igy] = p2 + p2 = p2.reshape(node.iicx, node.iicy) + p2 = np.repeat(p2[..., np.newaxis], node.icz, axis=2) + p2_full[node.igx : -node.igx, node.igy : -node.igy] = p2 if VS: - p2 = p2.reshape(ud.inx+2, ud.iny+2) - p2 = np.expand_dims(p2,2) + p2 = p2.reshape(ud.inx + 2, ud.iny + 2) + p2 = np.expand_dims(p2, 2) p2 = np.repeat(p2, node.icz, axis=2) - p2_full[1:-1,1:-1,:] = p2 - + p2_full[1:-1, 1:-1, :] = p2 - elif elem.ndim == 3 and elem.icy - 2*elem.igs[1] > 2: - p2_full[1:-1,1:-1,1:-1] = p2.reshape(ud.inx+2,ud.iny+2,ud.inz+2) + elif elem.ndim == 3 and elem.icy - 2 * elem.igs[1] > 2: + p2_full[1:-1, 1:-1, 1:-1] = p2.reshape(ud.inx + 2, ud.iny + 2, ud.inz + 2) if writer != None: - writer.populate(str(label),'p2_full',p2_full) + writer.populate(str(label), "p2_full", p2_full) - bdry.set_ghostnodes_p2(p2_full,node,ud) - correction_nodes(Sol,elem,node,mpv,p2_full,dt,ud,th,1) + bdry.set_ghostnodes_p2(p2_full, node, ud) + correction_nodes(Sol, elem, node, mpv, p2_full, dt, ud, th, 1) mpv.p2_nodes[...] += p2_full - bdry.set_ghostnodes_p2(mpv.p2_nodes,node,ud) + bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) -def correction_nodes(Sol,elem,node,mpv,p,dt,ud,th,updt_chi): +def correction_nodes(Sol, elem, node, mpv, p, dt, ud, th, updt_chi): ndim = node.ndim Gammainv = th.Gammainv @@ -285,26 +329,27 @@ def correction_nodes(Sol,elem,node,mpv,p,dt,ud,th,updt_chi): Dpx, Dpy, Dpz = grad_nodes(p, elem.ndim, node.dxyz) - thinv = (Sol.rho / Sol.rhoY) + thinv = Sol.rho / Sol.rhoY Y = Sol.rhoY / Sol.rho - coeff = (Gammainv * Sol.rhoY * Y) + coeff = Gammainv * Sol.rhoY * Y mpv.u = -dt * coeff * Dpx mpv.v = -dt * coeff * Dpy mpv.w = -dt * coeff * Dpz - multiply_inverse_coriolis(mpv, Sol, mpv, ud, elem, elem, dt, attrs=['u', 'v', 'w']) + multiply_inverse_coriolis(mpv, Sol, mpv, ud, elem, elem, dt, attrs=["u", "v", "w"]) Sol.rhou += thinv * mpv.u Sol.rhov += thinv * mpv.v Sol.rhow += thinv * mpv.w if ndim == 3 else 0.0 - Sol.rhoX += - updt_chi * dt * dSdy * Sol.rhov + Sol.rhoX += -updt_chi * dt * dSdy * Sol.rhov # set_explicit_boundary_data(Sol, elem, ud, th, mpv) assert True + def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): g = ud.gravity_strength[1] Msq = ud.Msq @@ -316,7 +361,7 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): wh1, wv, wh2 = dt * ud.coriolis_strength - ccenter = - ud.Msq * th.gm1inv / (dt**2) + ccenter = -ud.Msq * th.gm1inv / (dt**2) cexp = 2.0 - th.gamm igs = elem.igs @@ -327,24 +372,27 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): coeff = Gammainv * Sol.rhoY * Y nu = np.zeros_like(mpv.wcenter) - nu = -dt**2 * (g / Msq) * strat * Y + nu = -(dt**2) * (g / Msq) * strat * Y - setattr(mpv, 'nu_c', nu) + setattr(mpv, "nu_c", nu) nu = nu - for dim in range(ndim): ## Assuming 2D vertical slice! if dim == 1: - mpv.wplus[dim][...] = coeff #* gimp #* (wv**2 + 1.0) + mpv.wplus[dim][...] = coeff # * gimp #* (wv**2 + 1.0) else: - mpv.wplus[dim][...] = coeff #* fimp #* (wh1**2 + nu + nonhydro) + mpv.wplus[dim][...] = coeff # * fimp #* (wh1**2 + nu + nonhydro) kernel = np.ones([2] * ndim) - mpv.wcenter = (ccenter * sp.signal.fftconvolve(Sol.rhoY**cexp,kernel,mode='valid') / kernel.sum()) + mpv.wcenter = ( + ccenter + * sp.signal.fftconvolve(Sol.rhoY**cexp, kernel, mode="valid") + / kernel.sum() + ) - tmp_wplus = sp.signal.fftconvolve(mpv.wplus[0],kernel,mode='valid') / kernel.sum() + tmp_wplus = sp.signal.fftconvolve(mpv.wplus[0], kernel, mode="valid") / kernel.sum() # set_ghostcells_p2(mpv.wplus[0], elem, ud) # set_ghostcells_p2(mpv.wplus[1], elem, ud) @@ -353,7 +401,8 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): # mpv.wcenter[:,-1] = mpv.wcenter[:,-2] assert True - if not hasattr(ud, 'LAMB_BDRY'): scale_wall_node_values(mpv.wcenter, node, ud) + if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): + scale_wall_node_values(mpv.wcenter, node, ud) # def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): @@ -389,7 +438,7 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): # y_idx1 = slice(igs[dim]-is_periodic+1, right_idx) # innerdim1[dim] = slice(igs[dim]-1, (-igs[dim]+1)) - + # strat = (mpv.HydroState_n.S0[y_idx1] - mpv.HydroState_n.S0[y_idx]) / dy # nindim = tuple(nindim) @@ -430,7 +479,7 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): # scale_wall_node_values(mpv.wcenter, node, ud) -def scale_wall_node_values(rhs, node, ud, factor=.5): +def scale_wall_node_values(rhs, node, ud, factor=0.5): # if factor < 1.0: # if factor < 1.0: # factor = 1.0 @@ -451,13 +500,16 @@ def scale_wall_node_values(rhs, node, ud, factor=.5): wall_idx = np.empty((ndim), dtype=object) for dim in range(ndim): - wall_idx[dim] = slice(igs[dim],-igs[dim]) + wall_idx[dim] = slice(igs[dim], -igs[dim]) for dim in range(ndim): - is_wall = ud.bdry_type[dim] == opts.BdryType.WALL or ud.bdry_type[dim] == opts.BdryType.RAYLEIGH + is_wall = ( + ud.bdry_type[dim] == opts.BdryType.WALL + or ud.bdry_type[dim] == opts.BdryType.RAYLEIGH + ) if is_wall: - for direction in [-1,1]: - wall_idx[dim] = (igs[dim]-1) * direction + for direction in [-1, 1]: + wall_idx[dim] = (igs[dim] - 1) * direction if direction == -1: wall_idx[dim] -= 1 wall_idx_tuple = tuple(wall_idx) @@ -473,40 +525,54 @@ def scale_wall_node_values(rhs, node, ud, factor=.5): # rhs[wall_idx_tuple] *= factor - def grad_nodes_fft(p2n, elem, node): ndim = node.ndim dx, dy, dz = node.dx, node.dy, node.dz kernels = [] for dim in range(ndim): - kernel = np.ones([2]*ndim) - slc = [slice(None,)]*ndim - slc[dim] = slice(0,1) + kernel = np.ones([2] * ndim) + slc = [ + slice( + None, + ) + ] * ndim + slc[dim] = slice(0, 1) kernel[tuple(slc)] *= -1.0 kernels.append(kernel) - dpdx = -0.5**(ndim-1) * sp.signal.fftconvolve(p2n, kernels[0], mode='valid') / dx - dpdy = -0.5**(ndim-1) * sp.signal.fftconvolve(p2n, kernels[1], mode='valid') / dy if elem.iicy > 1 else 0.0 - dpdz = -0.5**(ndim-1) * sp.signal.fftconvolve(p2n, kernels[2], mode='valid') / dz if (ndim == 3) else 0.0 + dpdx = ( + -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[0], mode="valid") / dx + ) + dpdy = ( + -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[1], mode="valid") / dy + if elem.iicy > 1 + else 0.0 + ) + dpdz = ( + -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[2], mode="valid") / dz + if (ndim == 3) + else 0.0 + ) return dpdx, dpdy, dpdz + # @jit(nopython=True, nogil=False, cache=True) def grad_nodes(p, ndim, dxy): dx, dy, dz = dxy - indices = [idx for idx in it.product([slice(0,-1),slice(1,None)], repeat=ndim)] + indices = [idx for idx in it.product([slice(0, -1), slice(1, None)], repeat=ndim)] if ndim == 2: - signs_x = (-1., -1., +1., +1.) - signs_y = (-1., +1., -1., +1.) - signs_z = ( 0., 0., 0., 0.) + signs_x = (-1.0, -1.0, +1.0, +1.0) + signs_y = (-1.0, +1.0, -1.0, +1.0) + signs_z = (0.0, 0.0, 0.0, 0.0) elif ndim == 3: - signs_x = (-1., -1., -1., -1., +1., +1., +1., +1.) - signs_y = (-1., -1., +1., +1., -1., -1., +1., +1.) - signs_z = (-1., +1., -1., +1., -1., +1., -1., +1.) + signs_x = (-1.0, -1.0, -1.0, -1.0, +1.0, +1.0, +1.0, +1.0) + signs_y = (-1.0, -1.0, +1.0, +1.0, -1.0, -1.0, +1.0, +1.0) + signs_z = (-1.0, +1.0, -1.0, +1.0, -1.0, +1.0, -1.0, +1.0) - Dpx, Dpy, Dpz = 0., 0., 0. + Dpx, Dpy, Dpz = 0.0, 0.0, 0.0 cnt = 0 for index in indices: Dpx += signs_x[cnt] * p[index] @@ -514,15 +580,14 @@ def grad_nodes(p, ndim, dxy): Dpz += signs_z[cnt] * p[index] cnt += 1 - Dpx *= 0.5**(ndim-1) / dx - Dpy *= 0.5**(ndim-1) / dy - Dpz *= 0.5**(ndim-1) / dz + Dpx *= 0.5 ** (ndim - 1) / dx + Dpy *= 0.5 ** (ndim - 1) / dy + Dpz *= 0.5 ** (ndim - 1) / dz return Dpx, Dpy, Dpz - -def divergence_nodes(rhs,elem,node,Sol,ud): +def divergence_nodes(rhs, elem, node, Sol, ud): ndim = elem.ndim igs = elem.igs dxyz = node.dxyz @@ -530,61 +595,63 @@ def divergence_nodes(rhs,elem,node,Sol,ud): for dim in range(ndim): is_periodic = ud.bdry_type[dim] == opts.BdryType.PERIODIC - inner_idx[dim] = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic) + inner_idx[dim] = slice(igs[dim] - is_periodic, -igs[dim] + is_periodic) inner_idx_p1y = np.copy(inner_idx) - inner_idx_p1y[1] = slice(1,-1) - - indices = [idx for idx in it.product([slice(0,-1),slice(1,None)], repeat = ndim)] - signs = [sgn for sgn in it.product([1,-1], repeat = ndim)] + inner_idx_p1y[1] = slice(1, -1) + + indices = [idx for idx in it.product([slice(0, -1), slice(1, None)], repeat=ndim)] + signs = [sgn for sgn in it.product([1, -1], repeat=ndim)] inner_idx = tuple(inner_idx) inner_idx_p1y = tuple(inner_idx_p1y) - if not hasattr(ud, 'LAMB_BDRY'): - if ud.bdry_type[1] == opts.BdryType.WALL or ud.bdry_type[1] == opts.BdryType.RAYLEIGH: - Sol.rhou[:,:2,...] = 0.0 - Sol.rhov[:,:2,...] = 0.0 - Sol.rhow[:,:2,...] = 0.0 - - # if ud.bdry_type[1] == BdryType.WALL: - Sol.rhou[:,-2:,...] = 0.0 - Sol.rhov[:,-2:,...] = 0.0 - Sol.rhow[:,-2:,...] = 0.0 + if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): + if ( + ud.bdry_type[1] == opts.BdryType.WALL + or ud.bdry_type[1] == opts.BdryType.RAYLEIGH + ): + Sol.rhou[:, :2, ...] = 0.0 + Sol.rhov[:, :2, ...] = 0.0 + Sol.rhow[:, :2, ...] = 0.0 + # if ud.bdry_type[1] == BdryType.WALL: + Sol.rhou[:, -2:, ...] = 0.0 + Sol.rhov[:, -2:, ...] = 0.0 + Sol.rhow[:, -2:, ...] = 0.0 Y = Sol.rhoY / Sol.rho - Ux = np.diff(Sol.rhou * Y,axis=0) / elem.dx - Ux = 0.5 * (Ux[:,:-1,...] + Ux[:,1:,...]) + Ux = np.diff(Sol.rhou * Y, axis=0) / elem.dx + Ux = 0.5 * (Ux[:, :-1, ...] + Ux[:, 1:, ...]) - Vy = np.diff(Sol.rhov * Y,axis=1) / elem.dy - Vy = 0.5 * (Vy[:-1,...] + Vy[1:,...]) + Vy = np.diff(Sol.rhov * Y, axis=1) / elem.dy + Vy = 0.5 * (Vy[:-1, ...] + Vy[1:, ...]) if ndim == 3: - Ux = -0.5 * (Ux[...,:-1] + Ux[...,1:]) - Vy = 0.5 * (Vy[...,:-1] + Vy[...,1:]) + Ux = -0.5 * (Ux[..., :-1] + Ux[..., 1:]) + Vy = 0.5 * (Vy[..., :-1] + Vy[..., 1:]) - Wz = np.diff(Sol.rhow * Y,axis=2) / elem.dz - Wz = 0.5 * (Wz[:-1,...] + Wz[1:,...]) - Wz = 0.5 * (Wz[:,:-1,...] + Wz[:,1:,...]) + Wz = np.diff(Sol.rhow * Y, axis=2) / elem.dz + Wz = 0.5 * (Wz[:-1, ...] + Wz[1:, ...]) + Wz = 0.5 * (Wz[:, :-1, ...] + Wz[:, 1:, ...]) - rhs[1:-1,1:-1,1:-1] = (Ux + Vy + Wz) + rhs[1:-1, 1:-1, 1:-1] = Ux + Vy + Wz else: - rhs = (Ux + Vy) + rhs = Ux + Vy rhs_max = np.max(rhs[inner_idx]) if np.max(rhs[inner_idx]) > 0 else 0 return rhs -def rhs_from_p_old(rhs,node,mpv): +def rhs_from_p_old(rhs, node, mpv): igs = node.igs ndim = node.ndim assert ndim != 1, "Not implemented for 1D" - + # inner_idx = np.empty((ndim), dtype=object) # for dim in range(ndim): # inner_idx[dim] = slice(igs[dim],-igs[dim]) - + # inner_idx = tuple(inner_idx) # rhs_n = np.zeros_like(rhs) # rhs_hh = mpv.wcenter[inner_idx] * mpv.p2_nodes[inner_idx] @@ -592,8 +659,8 @@ def rhs_from_p_old(rhs,node,mpv): inner_idx = np.empty((ndim), dtype=object) for dim in range(ndim): - inner_idx[dim] = slice(igs[dim],-igs[dim]) - + inner_idx[dim] = slice(igs[dim], -igs[dim]) + inner_idx = tuple(inner_idx) rhs_n = np.zeros_like(rhs) rhs_hh = mpv.wcenter * mpv.p2_nodes[node.i1] @@ -601,7 +668,9 @@ def rhs_from_p_old(rhs,node,mpv): return rhs_n -def multiply_inverse_coriolis(Vec, Sol, mpv, ud, elem, node, dt, attrs=('rhou', 'rhov', 'rhow'), get_coeffs = False): +def multiply_inverse_coriolis( + Vec, Sol, mpv, ud, elem, node, dt, attrs=("rhou", "rhov", "rhow"), get_coeffs=False +): nonhydro = ud.nonhydrostasy g = ud.gravity_strength[1] Msq = ud.Msq @@ -611,26 +680,26 @@ def multiply_inverse_coriolis(Vec, Sol, mpv, ud, elem, node, dt, attrs=('rhou', strat = mpv.HydroState_n.get_dSdy(elem, node) Y = Sol.rhoY / Sol.rho - nu = -dt**2 * (g / Msq) * strat * Y + nu = -(dt**2) * (g / Msq) * strat * Y # get coefficients of the explicit terms # common denominator denom = 1.0 / (wh1**2 + wh2**2 + (nu + nonhydro) * (wv**2 + 1.0)) # U update - coeff_uu = (wh1**2 + nu + nonhydro) + coeff_uu = wh1**2 + nu + nonhydro coeff_uv = nonhydro * (wh1 * wv + wh2) - coeff_uw = (wh1 * wh2 - (nu + nonhydro) * wv) + coeff_uw = wh1 * wh2 - (nu + nonhydro) * wv # V update - coeff_vu = (wh1 * wv - wh2) + coeff_vu = wh1 * wv - wh2 coeff_vv = nonhydro * (1 + wv**2) - coeff_vw = (wh2 * wv + wh1) + coeff_vw = wh2 * wv + wh1 # W update - coeff_wu = (wh1 * wh2 + (nu + nonhydro) * wv) + coeff_wu = wh1 * wh2 + (nu + nonhydro) * wv coeff_wv = nonhydro * (wh2 * wv - wh1) - coeff_ww = (nu + nonhydro + wh2**2) + coeff_ww = nu + nonhydro + wh2**2 VecU = getattr(Vec, attrs[0]) VecV = getattr(Vec, attrs[1]) @@ -648,5 +717,10 @@ def multiply_inverse_coriolis(Vec, Sol, mpv, ud, elem, node, dt, attrs=('rhou', i1 = node.i1 if get_coeffs: # coriolis_parameters = ((coeff_uu * denom)[i1].reshape(-1,), (coeff_vv * denom)[i1].reshape(-1,), (coeff_uv * denom)[i1].reshape(-1,), (coeff_vu * denom)[i1].reshape(-1,)) - coriolis_parameters = ((coeff_uu * denom).T, (coeff_vv * denom).T, (coeff_uv * denom).T, (coeff_vu * denom).T) - return coriolis_parameters \ No newline at end of file + coriolis_parameters = ( + (coeff_uu * denom).T, + (coeff_vv * denom).T, + (coeff_uv * denom).T, + (coeff_vu * denom).T, + ) + return coriolis_parameters diff --git a/src/dycore/physics/low_mach/__init__.py b/src/pybella/flow_solver/utils/__init__.py similarity index 100% rename from src/dycore/physics/low_mach/__init__.py rename to src/pybella/flow_solver/utils/__init__.py diff --git a/src/dycore/utils/boundary.py b/src/pybella/flow_solver/utils/boundary.py similarity index 63% rename from src/dycore/utils/boundary.py rename to src/pybella/flow_solver/utils/boundary.py index 9992cddd..74209bf9 100644 --- a/src/dycore/utils/boundary.py +++ b/src/pybella/flow_solver/utils/boundary.py @@ -1,8 +1,11 @@ """ For more details on this module, refer to the write-up :ref:`boundary_handling`. """ +import copy import numpy as np -import dycore.utils.options as opts +from . import options as opts +from ...utils import io + def set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=None): """ @@ -31,23 +34,23 @@ def set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=None): if step == None: dims = np.arange(ndim) else: - dims = [ndim-1] + dims = [ndim - 1] for dim in dims: if step is not None: current_step = step else: current_step = dim - ghost_padding, idx = get_ghost_padding(ndim,dim,igs) + ghost_padding, idx = get_ghost_padding(ndim, dim, igs) if ud.gravity_strength[current_step] == 0.0: # Do this for the axes that do not have gravity. # Periodic BC. if ud.bdry_type[current_step] == opts.BdryType.PERIODIC: - set_boundary(Sol,ghost_padding,'wrap',idx,step=None) + set_boundary(Sol, ghost_padding, "wrap", idx, step=None) # Wall BC. elif ud.bdry_type[current_step] == opts.BdryType.WALL: - set_boundary(Sol,ghost_padding,'symmetric',idx,step=None) + set_boundary(Sol, ghost_padding, "symmetric", idx, step=None) elif ud.bdry_type[current_step] == opts.BdryType.RAYLEIGH: assert 0, "Rayleigh boundary not defined on x-direction." # set_boundary(Sol,((0,0),(0,2)),'constant',(slice(None,),slice(0,-2)),step=None) @@ -56,8 +59,8 @@ def set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=None): else: # get current axis that has gravity. gravity_axis = dim - - direction = -1. + + direction = -1.0 offset = 0 # get gravity strength specified in the user data file. @@ -69,34 +72,51 @@ def set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=None): # loop through each of these ghost cells. for current_idx in np.arange(side)[::-1]: if step != None: - y_axs = (ndim - 1) + y_axs = ndim - 1 else: y_axs = 1 - nlast, nsource, nimage = get_gravity_padding(ndim,current_idx,direction,offset,elem,y_axs=y_axs) + nlast, nsource, nimage = get_gravity_padding( + ndim, current_idx, direction, offset, elem, y_axs=y_axs + ) Y_last = Sol.rhoY[nlast] / Sol.rho[nlast] - rhoYv_image = -Sol.rhov[nsource] * Sol.rhoY[nsource] / Sol.rho[nsource] + rhoYv_image = ( + -Sol.rhov[nsource] * Sol.rhoY[nsource] / Sol.rho[nsource] + ) - S = 1. / ud.stratification(elem.y[nimage[y_axs]]) + S = 1.0 / ud.stratification(elem.y[nimage[y_axs]]) - if hasattr(ud, 'LAMB_BDRY'): - dpi = (mpv.HydroState.p20[nimage[y_axs]] - mpv.HydroState.p20[nlast[y_axs]]) * ud.Msq + if hasattr(ud, "ATMOSPHERIC_EXTENSION"): + dpi = ( + mpv.HydroState.p20[nimage[y_axs]] + - mpv.HydroState.p20[nlast[y_axs]] + ) * ud.Msq else: - dpi = direction * (th.Gamma*g) * 0.5 * elem.dy * (1.0 / Y_last + S) - - rhoY = ((Sol.rhoY[nlast]**th.gm1) + dpi)**th.gm1inv if ud.is_compressible == 1 else mpv.HydroState.rhoY0[nimage[y_axs]] + dpi = ( + direction + * (th.Gamma * g) + * 0.5 + * elem.dy + * (1.0 / Y_last + S) + ) + + rhoY = ( + ((Sol.rhoY[nlast] ** th.gm1) + dpi) ** th.gm1inv + if ud.is_compressible == 1 + else mpv.HydroState.rhoY0[nimage[y_axs]] + ) rho = rhoY * S Y_source = Sol.rhoY[nsource] / Sol.rho[nsource] Y_image = rhoY / rho - if hasattr(ud, 'LAMB_BDRY'): - if direction > 0: # if bottom boundary + if hasattr(ud, "ATMOSPHERIC_EXTENSION"): + if direction > 0: # if bottom boundary v = Sol.rhov[nsource] * Y_source / Sol.rho[nsource] * rho - else: # if top boundary - v = Sol.rhov[nsource] * Y_source + else: # if top boundary + v = Sol.rhov[nsource] * Y_source Th_slc = rhoY / rho / Y_last @@ -119,20 +139,20 @@ def set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=None): # v = -Sol.rhov[nsource] / Sol.rho[nsource] Sol.rho[nimage] = rho - Sol.rhou[nimage] = rho*u * Th_slc + Sol.rhou[nimage] = rho * u * Th_slc # Sol.rhov[nimage] = 0.0#rho*v - if hasattr(ud, 'LAMB_BDRY'): + if hasattr(ud, "ATMOSPHERIC_EXTENSION"): Sol.rhov[nimage] = -v / Y_image else: - Sol.rhov[nimage] = rho*v - Sol.rhow[nimage] = rho*w * Th_slc + Sol.rhov[nimage] = rho * v + Sol.rhow[nimage] = rho * w * Th_slc Sol.rhoY[nimage] = rhoY Sol.rhoX[nimage] = rho * X offset += 1 -def set_boundary(Sol,pads,btype,idx,step=None): +def set_boundary(Sol, pads, btype, idx, step=None): """ Called by the function :func:`inputs.boundary.set_explicit_boundary_data`. Pads in-place the ghost cells for a given boundary type. @@ -153,27 +173,27 @@ def set_boundary(Sol,pads,btype,idx,step=None): If we are in the advection routine with the flipped arrays according to the directional Strang-splitting, we want to pad the correct direction. `step=0`, pads the x-direction while `step=1` pads the y-direction. """ - Sol.rho[...] = np.pad(Sol.rho[idx],pads,btype) + Sol.rho[...] = np.pad(Sol.rho[idx], pads, btype) # Sol.rhou[...] = np.pad(Sol.rhou[idx],pads,btype) # Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,btype) # Sol.rhow[...] = np.pad(Sol.rhow[idx],pads,btype) - if btype == 'symmetric': + if btype == "symmetric": # Sol.rhou[...] = np.pad(Sol.rhou[idx],pads,negative_symmetric) - Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,negative_symmetric) + Sol.rhov[...] = np.pad(Sol.rhov[idx], pads, negative_symmetric) # Sol.rhow[...] = np.pad(Sol.rhow[idx],pads,negative_symmetric) - Sol.rho[...] = np.pad(Sol.rho[idx],pads,'symmetric') - Sol.rhou[...] = np.pad(Sol.rhou[idx],pads,'symmetric') - elif btype == 'constant': - Sol.rho[...] = np.pad(Sol.rho[idx],pads,'symmetric') - Sol.rhou[...] = np.pad(Sol.rhou[idx],pads,'symmetric') - Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,btype) - Sol.rhow[...] = np.pad(Sol.rhow[idx],pads,'symmetric') - btype = 'symmetric' + Sol.rho[...] = np.pad(Sol.rho[idx], pads, "symmetric") + Sol.rhou[...] = np.pad(Sol.rhou[idx], pads, "symmetric") + elif btype == "constant": + Sol.rho[...] = np.pad(Sol.rho[idx], pads, "symmetric") + Sol.rhou[...] = np.pad(Sol.rhou[idx], pads, "symmetric") + Sol.rhov[...] = np.pad(Sol.rhov[idx], pads, btype) + Sol.rhow[...] = np.pad(Sol.rhow[idx], pads, "symmetric") + btype = "symmetric" else: - Sol.rhou[...] = np.pad(Sol.rhou[idx],pads,btype) - Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,btype) - Sol.rhow[...] = np.pad(Sol.rhow[idx],pads,btype) + Sol.rhou[...] = np.pad(Sol.rhou[idx], pads, btype) + Sol.rhov[...] = np.pad(Sol.rhov[idx], pads, btype) + Sol.rhow[...] = np.pad(Sol.rhow[idx], pads, btype) # if step == 0: # Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,btype) @@ -212,11 +232,11 @@ def set_boundary(Sol,pads,btype,idx,step=None): # Sol.rhov[...] = np.pad(Sol.rhov[idx],pads,btype) # Sol.rhow[...] = np.pad(Sol.rhow[idx],pads,btype) - Sol.rhoY[...] = np.pad(Sol.rhoY[idx],pads,btype) - Sol.rhoX[...] = np.pad(Sol.rhoX[idx],pads,btype) + Sol.rhoY[...] = np.pad(Sol.rhoY[idx], pads, btype) + Sol.rhoX[...] = np.pad(Sol.rhoX[idx], pads, btype) -def negative_symmetric(vector,pad_width,iaxis,kwargs=None): +def negative_symmetric(vector, pad_width, iaxis, kwargs=None): """ Taken from the reference: @@ -238,14 +258,14 @@ def negative_symmetric(vector,pad_width,iaxis,kwargs=None): """ if pad_width[1] > 0: sign = -1 - vector[:pad_width[0]] = sign * vector[pad_width[0]:2*pad_width[0]][::-1] - vector[-pad_width[1]:] = sign * vector[-2*pad_width[1]:-pad_width[1]][::-1] + vector[: pad_width[0]] = sign * vector[pad_width[0] : 2 * pad_width[0]][::-1] + vector[-pad_width[1] :] = sign * vector[-2 * pad_width[1] : -pad_width[1]][::-1] return vector - else: # axis must have length > 0 for padding + else: # axis must have length > 0 for padding return vector -def get_gravity_padding(ndim,cur_idx,direction,offset,elem,y_axs=None): +def get_gravity_padding(ndim, cur_idx, direction, offset, elem, y_axs=None): """ Parameters ---------- @@ -264,7 +284,7 @@ def get_gravity_padding(ndim,cur_idx,direction,offset,elem,y_axs=None): """ cur_i = np.copy(cur_idx) - cur_idx += offset * ((elem.icy - 1) - 2*cur_idx) + cur_idx += offset * ((elem.icy - 1) - 2 * cur_idx) gravity_padding = [(slice(None))] * ndim if y_axs == None: # y_axs = ndim - 1 @@ -274,43 +294,46 @@ def get_gravity_padding(ndim,cur_idx,direction,offset,elem,y_axs=None): nlast[y_axs] = int(cur_idx + direction) nsource = np.copy(gravity_padding) - nsource[y_axs] = int(offset*(elem.icy) + direction * (2 * elem.igy - (1 - offset) - cur_i)) + nsource[y_axs] = int( + offset * (elem.icy) + direction * (2 * elem.igy - (1 - offset) - cur_i) + ) nimage = np.copy(gravity_padding) nimage[y_axs] = int(cur_idx) return tuple(nlast), tuple(nsource), tuple(nimage) -def set_ghostcells_p2(p,elem,ud): +def set_ghostcells_p2(p, elem, ud): igs = elem.igs - + for dim in range(elem.ndim): - ghost_padding, idx = get_ghost_padding(elem.ndim,dim,igs) + ghost_padding, idx = get_ghost_padding(elem.ndim, dim, igs) if ud.bdry_type[dim] == opts.BdryType.PERIODIC: - p[...] = np.pad(p[idx],ghost_padding,'wrap') - else: # WALL - p[...] = np.pad(p[idx],ghost_padding,'symmetric') + p[...] = np.pad(p[idx], ghost_padding, "wrap") + else: # WALL + p[...] = np.pad(p[idx], ghost_padding, "symmetric") -def set_ghostnodes_p2(p,node,ud, igs=None): - if igs is None: igs = node.igs +def set_ghostnodes_p2(p, node, ud, igs=None): + if igs is None: + igs = node.igs for dim in range(node.ndim): - ghost_padding, idx = get_ghost_padding(node.ndim,dim,igs) + ghost_padding, idx = get_ghost_padding(node.ndim, dim, igs) if ud.bdry_type[dim] == opts.BdryType.PERIODIC: p[...] = np.pad(p[idx], ghost_padding, periodic_plus_one) - else: # ud.bdry_type[dim] == opts.BdryType.WALL: - p[...] = np.pad(p[idx], ghost_padding, 'reflect') + else: # ud.bdry_type[dim] == opts.BdryType.WALL: + p[...] = np.pad(p[idx], ghost_padding, "reflect") - # if periodic_plus_one - if node.iicy == 2: # implying horizontal slices - pn = p[:,2,:] + # if periodic_plus_one + if node.iicy == 2: # implying horizontal slices + pn = p[:, 2, :] pn = np.expand_dims(pn, axis=1) p[...] = np.repeat(pn, node.icy, axis=1) -def get_ghost_padding(ndim,dim,igs): +def get_ghost_padding(ndim, dim, igs): """ For a given direction, return the number of ghost cells to pad the current direction, and the index slice of the inner array without the ghost cells. @@ -331,18 +354,18 @@ def get_ghost_padding(ndim,dim,igs): Index slice of the inner domain of the array. """ - ghost_padding = [(0,0)] * ndim - ghost_padding[dim] = (igs[dim],igs[dim]) + ghost_padding = [(0, 0)] * ndim + ghost_padding[dim] = (igs[dim], igs[dim]) padded_idx = np.empty((ndim), dtype=object) for idim in range(ndim): - padded_idx[idim] = slice(igs[idim],-igs[idim]) + padded_idx[idim] = slice(igs[idim], -igs[idim]) padded_idx[dim] = slice(None) inner_domain = [slice(None)] * ndim - inner_domain[dim] = slice(igs[dim],-igs[dim]) + inner_domain[dim] = slice(igs[dim], -igs[dim]) - return tuple(ghost_padding), tuple(inner_domain) + return tuple(ghost_padding), tuple(inner_domain) def periodic_plus_one(vector, pad_width, iaxis, kwargs=None): @@ -365,7 +388,10 @@ def periodic_plus_one(vector, pad_width, iaxis, kwargs=None): """ if all(pad_width) > 0: - vector[:pad_width[0]+1], vector[-pad_width[1]-1:] = vector[-pad_width[1]-pad_width[1]-1:-pad_width[1]] , vector[pad_width[0]:pad_width[0]+pad_width[0]+1].copy() + vector[: pad_width[0] + 1], vector[-pad_width[1] - 1 :] = ( + vector[-pad_width[1] - pad_width[1] - 1 : -pad_width[1]], + vector[pad_width[0] : pad_width[0] + pad_width[0] + 1].copy(), + ) return vector @@ -374,7 +400,7 @@ def get_tau_y(ud, elem, node, alpha): taun_y = np.zeros_like(node.y) # ud.bcy = elem.y[-ud.inbcy-2] - ud.bny = node.y[-ud.inbcy-3] + ud.bny = node.y[-ud.inbcy - 3] # ud.bny = ud.bcy # c1c = elem.y <= ud.bcy @@ -394,8 +420,22 @@ def get_tau_y(ud, elem, node, alpha): # tauc_y[np.where(c3c)] = - alpha / 2.0 * (1.0 + ((elem.y[np.where(c3c)] - ud.bcy) / (elem.y[-1] - ud.bcy) - 0.5) * np.pi) taun_y[np.where(c1n)] = 0.0 - taun_y[np.where(c2n)] = - alpha / 2.0 * (1.0 - np.cos( (node.y[np.where(c2n)] - ud.bny) / (node.y[-1] - ud.bny) * np.pi )) - taun_y[np.where(c3n)] = - alpha / 2.0 * (1.0 + ((node.y[np.where(c3n)] - ud.bny) / (node.y[-1] - ud.bny) - 0.5) * np.pi) + taun_y[np.where(c2n)] = ( + -alpha + / 2.0 + * ( + 1.0 + - np.cos((node.y[np.where(c2n)] - ud.bny) / (node.y[-1] - ud.bny) * np.pi) + ) + ) + taun_y[np.where(c3n)] = ( + -alpha + / 2.0 + * ( + 1.0 + + ((node.y[np.where(c3n)] - ud.bny) / (node.y[-1] - ud.bny) - 0.5) * np.pi + ) + ) taun_y[-2:] = -np.abs(taun_y).max() tauc_y[...] = np.interp(elem.y, node.y, taun_y) @@ -420,8 +460,30 @@ def get_bottom_tau_y(ud, elem, node, alpha, cutoff=0.5): c3n = np.logical_and(ccn > 0.5, ccn <= 1.0) taun_y[np.where(c1n)] = 0.0 - taun_y[np.where(c2n)] = - alpha / 2.0 * (1.0 - np.cos( (node.y[np.where(c2n)] - ud.forcing_bny) / (node.y[-1] - ud.forcing_bny) * np.pi )) - taun_y[np.where(c3n)] = - alpha / 2.0 * (1.0 + ((node.y[np.where(c3n)] - ud.forcing_bny) / (node.y[-1] - ud.forcing_bny) - 0.5) * np.pi) + taun_y[np.where(c2n)] = ( + -alpha + / 2.0 + * ( + 1.0 + - np.cos( + (node.y[np.where(c2n)] - ud.forcing_bny) + / (node.y[-1] - ud.forcing_bny) + * np.pi + ) + ) + ) + taun_y[np.where(c3n)] = ( + -alpha + / 2.0 + * ( + 1.0 + + ( + (node.y[np.where(c3n)] - ud.forcing_bny) / (node.y[-1] - ud.forcing_bny) + - 0.5 + ) + * np.pi + ) + ) taun_y[-3:] = -np.abs(taun_y).max() taun_y[...] = taun_y[::-1] @@ -436,28 +498,82 @@ def get_bottom_tau_y(ud, elem, node, alpha, cutoff=0.5): return tauc_y, taun_y - +def apply_rayleigh_forcing( + Sol, + mpv, + ud, + elem, + node, + t, + step, + dt, + th, + bdry, + half=True, + Sol_half_new=None, + mpv_half_new=None, +): + """Apply Rayleigh forcing boundary condition (file or function based).""" + if not (hasattr(ud, "rayleigh_forcing") and ud.rayleigh_forcing): + return + + t_offset = 0.5 * dt if half else dt + + if ud.rayleigh_forcing_type == "file": + reader = io.read_input( + ud.rayleigh_forcing_fn, ud.rayleigh_forcing_path + ) + + if Sol_half_new is None or mpv_half_new is None: + Sol_half_new = copy.deepcopy(Sol) + mpv_half_new = copy.deepcopy(mpv) + + time_tag = "%.3d_after_full_step" % step + reader.get_data(Sol_half_new, mpv_half_new, time_tag, half=half) + + up = Sol_half_new.rhou / Sol_half_new.rho + vp = Sol_half_new.rhov / Sol_half_new.rho + Yp = Sol_half_new.rhoY / Sol_half_new.rho - mpv.HydroState.Y0.reshape(1, -1) + pi = mpv_half_new.p2_nodes + + bdry.rayleigh_damping( + Sol, mpv, ud, elem, node, [up, vp, Yp, pi, t + t_offset] + ) + + elif ud.rayleigh_forcing_type == "func": + s = 5.0e-3 + 1e-4 + 0e-5 + ud.rf_bot.eigenfunction(t + t_offset, s) + up, vp, Yp, pi = ud.rf_bot.dehatter(th) + + ud.rf_bot.eigenfunction(t + t_offset, s, grid="n") + _, _, _, pi_n = ud.rf_bot.dehatter(th, grid="n") + + bdry.rayleigh_damping( + Sol, mpv, ud, elem, node, [up, vp, Yp, pi_n, t + t_offset] + ) + + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) def rayleigh_damping(Sol, mpv, ud, elem, node, forcing=None): - u = (Sol.rhou / Sol.rho)#[elem.i2] - v = (Sol.rhov / Sol.rho)#[elem.i2] - Y = (Sol.rhoY / Sol.rho)#[elem.i2] - rho = Sol.rho#[elem.i2] + u = Sol.rhou / Sol.rho # [elem.i2] + v = Sol.rhov / Sol.rho # [elem.i2] + Y = Sol.rhoY / Sol.rho # [elem.i2] + rho = Sol.rho # [elem.i2] if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: tcy, tny = ud.tcy, ud.tny else: tcy, tny = 0.0, 0.0 - if (forcing is not None): + if forcing is not None: tcy_f, tny_f = ud.forcing_tcy, ud.forcing_tny tcy, tny = 0.0, 0.0 u_f, v_f, Y_f, pi_f, t = forcing - if ud.rayleigh_forcing_type == 'file': - G = np.sqrt( 9.0 / 40.0 ) - N = np.sqrt(ud.Nsq_ref) + if ud.rayleigh_forcing_type == "file": + G = np.sqrt(9.0 / 40.0) + N = np.sqrt(ud.Nsq_ref) C = ud.Cs * ud.u_ref Gam = N * G / C Om = ud.coriolis_strength[2] / 2.0 / ud.t_ref @@ -467,7 +583,7 @@ def rayleigh_damping(Sol, mpv, ud, elem, node, forcing=None): # if np.all(ud.coriolis_strength) == 0.0: # if ud.trad_forcing: - # mfac = 1.0 + # mfac = 1.0 else: mfac = 1.0 @@ -481,10 +597,14 @@ def rayleigh_damping(Sol, mpv, ud, elem, node, forcing=None): mfac = 0.0 # assuming 2D vertical slice - not dimension agnostic - u += tcy * (u - ud.u_wind_speed) + c_f * (tcy_f * (u - ud.u_wind_speed) + np.abs(tcy_f) * mfac * u_f) - v += tcy * (v - ud.v_wind_speed) + c_f * (tcy_f * (v - ud.v_wind_speed) + np.abs(tcy_f) * mfac * v_f) - - Ybar = mpv.HydroState.Y0.reshape(1,-1) + u += tcy * (u - ud.u_wind_speed) + c_f * ( + tcy_f * (u - ud.u_wind_speed) + np.abs(tcy_f) * mfac * u_f + ) + v += tcy * (v - ud.v_wind_speed) + c_f * ( + tcy_f * (v - ud.v_wind_speed) + np.abs(tcy_f) * mfac * v_f + ) + + Ybar = mpv.HydroState.Y0.reshape(1, -1) Y += tcy * (Y - Ybar) + c_f * (tcy_f * (Y - Ybar) + np.abs(tcy_f) * mfac * Y_f) Sol.rhou[...] = rho * u @@ -492,23 +612,23 @@ def rayleigh_damping(Sol, mpv, ud, elem, node, forcing=None): Sol.rhoY[...] = rho * Y - def check_flux_bcs(Lefts, Rights, elem, split_step, ud): igx = elem.igx igy = elem.igy if split_step == 1: if ud.bdry_type[split_step] == opts.BdryType.WALL: + left_inner = (slice(None), slice(igy, igy + 1)) + left_ghost = (slice(None), slice(igy - 1, igy)) - left_inner = (slice(None),slice(igy,igy+1)) - left_ghost = (slice(None),slice(igy-1,igy)) + right_inner = (slice(None), slice(-igy - 1, -igy)) + right_ghost = (slice(None), slice(-igy, -igy + 1)) - right_inner = (slice(None),slice(-igy-1,-igy)) - right_ghost = (slice(None),slice(-igy,-igy+1)) - - rhou_wall = 0. + rhou_wall = 0.0 # Lefts.rhou[left_ghost] = Rights.rhou[left_inner] = rhou_wall - Lefts.rhoY[left_ghost] = Rights.rhoY[left_inner] = Rights.rho[left_inner] * ud.stratification(0.0) + Lefts.rhoY[left_ghost] = Rights.rhoY[left_inner] = Rights.rho[ + left_inner + ] * ud.stratification(0.0) Lefts.rho[left_ghost] = Rights.rho[left_inner] Lefts.rhov[left_ghost] = Rights.rhov[left_inner] @@ -516,25 +636,24 @@ def check_flux_bcs(Lefts, Rights, elem, split_step, ud): Lefts.rhoX[left_ghost] = Rights.rhoX[left_inner] Rights.rho[right_ghost] = Lefts.rho[right_inner] - Rights.rhou[right_ghost] = -1. * Lefts.rhou[right_inner] + Rights.rhou[right_ghost] = -1.0 * Lefts.rhou[right_inner] Rights.rhov[right_ghost] = Lefts.rhov[right_inner] Rights.rhow[right_ghost] = Lefts.rhow[right_inner] Rights.rhoY[right_ghost] = Lefts.rhoY[right_inner] else: if ud.bdry_type[split_step] == opts.BdryType.WALL: - - assert(0) # INCOMPLETE!!! - Lefts.rho[left_inner] = Rights.rho[:,igx-2] - Lefts.rhou[left_inner] = -1. * Rights.rhou[:,igx-2] - Lefts.rhov[left_inner] = Rights.rhov[:,igx-2] - Lefts.rhow[left_inner] = Rights.rhow[:,igx-2] - Lefts.rhoY[left_inner] = Rights.rhoY[:,igx-2] + assert 0 # INCOMPLETE!!! + Lefts.rho[left_inner] = Rights.rho[:, igx - 2] + Lefts.rhou[left_inner] = -1.0 * Rights.rhou[:, igx - 2] + Lefts.rhov[left_inner] = Rights.rhov[:, igx - 2] + Lefts.rhow[left_inner] = Rights.rhow[:, igx - 2] + Lefts.rhoY[left_inner] = Rights.rhoY[:, igx - 2] # print("#################### TRUE ########################") - Rights.rho[right_ghost] = Lefts.rho[:,-igx-2] - Rights.rhou[right_ghost] = -1. * Lefts.rhou[:,-igx-2] - Rights.rhov[right_ghost] = Lefts.rhov[:,-igx-2] - Rights.rhow[right_ghost] = Lefts.rhow[:,-igx-2] - Rights.rhoY[right_ghost] = Lefts.rhoY[:,-igx-2] - # print(Rights.rhoY[right_ghost]) \ No newline at end of file + Rights.rho[right_ghost] = Lefts.rho[:, -igx - 2] + Rights.rhou[right_ghost] = -1.0 * Lefts.rhou[:, -igx - 2] + Rights.rhov[right_ghost] = Lefts.rhov[:, -igx - 2] + Rights.rhow[right_ghost] = Lefts.rhow[:, -igx - 2] + Rights.rhoY[right_ghost] = Lefts.rhoY[:, -igx - 2] + # print(Rights.rhoY[right_ghost]) diff --git a/src/dycore/utils/options.py b/src/pybella/flow_solver/utils/options.py similarity index 87% rename from src/dycore/utils/options.py rename to src/pybella/flow_solver/utils/options.py index 3c56a361..621bea7a 100644 --- a/src/dycore/utils/options.py +++ b/src/pybella/flow_solver/utils/options.py @@ -1,4 +1,4 @@ -from enum import Enum # ! Version > Python 3.4 +from enum import Enum # ! Version > Python 3.4 # class RecoveryOrder(Enum): # FIRST = 0 @@ -27,6 +27,7 @@ # ZERO_ORDER_EXTRAPOL = 0 # BOTTOM_BC_DEFAULT = 1 + class LimiterType(Enum): NONE = 0 # MINMOD = 1 @@ -39,11 +40,12 @@ class LimiterType(Enum): # NO_SLOPE = 8 # NUMBER_OF_LIMITER = 9 + class BdryType(Enum): """ An enumeration class that defines the accepted boundary condition types. """ - - WALL = 'symmetric' - PERIODIC = 'wrap' - RAYLEIGH = 'radiation' \ No newline at end of file + + WALL = "symmetric" + PERIODIC = "wrap" + RAYLEIGH = "radiation" diff --git a/src/pybella/flow_solver/utils/solver_diagnostics.py b/src/pybella/flow_solver/utils/solver_diagnostics.py new file mode 100644 index 00000000..22983dcd --- /dev/null +++ b/src/pybella/flow_solver/utils/solver_diagnostics.py @@ -0,0 +1,22 @@ +import numpy as np +import scipy as sp + + +def get_p_from_pressure_related_fields(mem, ud, psinc=False): + p2n = mem.mpv.p2_nodes + rhoY = mem.sol.rhoY + dp2n = (p2n - p2n.mean()) * ud.Msq + + th = mem.th + + kernel = np.ones((2, 2)) + dp2c = sp.signal.fftconvolve(dp2n, kernel, mode="valid") / kernel.sum() + + if psinc: + P0 = (rhoY ** (th.gamm - 1.0) + dp2c) ** (1.0 / (th.gamm - 1.0)) + p = P0 ** (th.gamm) + else: + P0 = (rhoY ** (th.gamm - 1.0) - dp2c) ** (1.0 / (th.gamm - 1.0)) + p = rhoY ** (th.gamm) - P0 ** (th.gamm) + + return p diff --git a/src/dycore/utils/variable.py b/src/pybella/flow_solver/utils/variable.py similarity index 89% rename from src/dycore/utils/variable.py rename to src/pybella/flow_solver/utils/variable.py index 52e5dc93..48811fb0 100644 --- a/src/dycore/utils/variable.py +++ b/src/pybella/flow_solver/utils/variable.py @@ -1,13 +1,14 @@ import numpy as np import scipy as sp -# equivalent to States_new + class Vars(object): """ The data container for the solution state variables, i.e. `Sol`. """ - def __init__(self,size,ud): + + def __init__(self, size, ud): """ Parameters ---------- @@ -28,7 +29,7 @@ def __init__(self,size,ud): Notes ----- 2. `rhoX` has to be extended by `ud.nspec` for moist process. - + """ self.rho = np.zeros((size)) self.rhou = np.zeros((size)) @@ -45,10 +46,11 @@ def squeezer(self): """ for key, value in vars(self).items(): - setattr(self,key,value.squeeze()) + setattr(self, key, value.squeeze()) + # method written for 2D - def primitives(self,th): + def primitives(self, th): """ Calculate the primitive quantities from the state variables and extend the data container to include these quantities. @@ -65,7 +67,7 @@ def primitives(self,th): Y : ndarray(size_of_rhoY) X : ndarray(size_of_rhoX) p : ndarray(size_of_rhoY) - + """ nonzero_idx = np.nonzero(self.rho) @@ -83,44 +85,45 @@ def primitives(self,th): self.w[nonzero_idx] = self.rhow[nonzero_idx] / self.rho[nonzero_idx] self.Y[nonzero_idx] = self.rhoY[nonzero_idx] / self.rho[nonzero_idx] self.X[nonzero_idx] = self.rhoX[nonzero_idx] / self.rho[nonzero_idx] - self.p[nonzero_idx] = self.rhoY[nonzero_idx]**th.gamm + self.p[nonzero_idx] = self.rhoY[nonzero_idx] ** th.gamm def flip(self): """ Flips the solution variables arrays for the advection routine. `rhou` and `rhov` are also flipped, i.e.:: - self.rhou, self.rhov = self.rhov, self.rhou + self.rhou, self.rhov = self.rhov, self.rhou """ for key, value in vars(self).items(): - setattr(self,key,value.T) + setattr(self, key, value.T) self.rhou, self.rhov = self.rhov, self.rhou def flip_forward(self): for key, value in vars(self).items(): - setattr(self,key,np.moveaxis(value,0,-1)) + setattr(self, key, np.moveaxis(value, 0, -1)) def flip_backward(self): for key, value in vars(self).items(): - setattr(self,key,np.moveaxis(value,-1,0)) - + setattr(self, key, np.moveaxis(value, -1, 0)) def mod_bg_wind(self, ud, fac): u0 = ud.u_wind_speed v0 = ud.v_wind_speed w0 = ud.w_wind_speed - + self.rhou = self.rhou + fac * u0 * self.rho self.rhov = self.rhov + fac * v0 * self.rho self.rhow = self.rhow + fac * w0 * self.rho - + + class States(Vars): """ Data container for `Lefts` and `Rights` for the Riemann solver. Inherits the solution class :class:`management.variable.Vars`. """ - def __init__(self,size,ud): + + def __init__(self, size, ud): """ Parameters ---------- @@ -134,7 +137,7 @@ def __init__(self,size,ud): Many variables in this data container are no longer used and can be removed. """ - super().__init__(size,ud) + super().__init__(size, ud) self.u = np.zeros((size)) self.v = np.zeros((size)) self.w = np.zeros((size)) @@ -159,35 +162,40 @@ def __init__(self,size,ud): self.get_dSdy = self.get_dSdy self.get_S0c = self.get_S0c + self.init_dSdy = False + self.init_S0c = False + def get_dSdy(self, elem, node): - if hasattr(self, 'dSdy'): + if self.init_dSdy: return self.dSdy else: ndim = node.ndim dy = node.dy dSdy = self.S0 - dSdy = sp.signal.convolve(dSdy,[1.,-1.],mode='valid') / dy + dSdy = sp.signal.convolve(dSdy, [1.0, -1.0], mode="valid") / dy - for dim in range(0,ndim,2): + for dim in range(0, ndim, 2): dSdy = np.expand_dims(dSdy, dim) dSdy = np.repeat(dSdy, elem.sc[dim], axis=dim) self.dSdy = dSdy + self.init_dSdy = True return dSdy def get_S0c(self, elem): - if hasattr(self, 'S0c'): + if self.init_S0c: return self.S0c else: ndim = elem.ndim S0c = self.S0 - for dim in range(0,ndim,2): + for dim in range(0, ndim, 2): S0c = np.expand_dims(S0c, dim) S0c = np.repeat(S0c, elem.sc[dim], axis=dim) self.S0c = S0c + self.init_S0c = True return S0c @@ -196,6 +204,7 @@ class Characters(object): Data container for the slope and amplitude of the interpolation to the faces for the Riemann solver. """ + def __init__(self, size): """ Parameters diff --git a/src/tests/__init__.py b/src/pybella/inputs/__init__.py similarity index 100% rename from src/tests/__init__.py rename to src/pybella/inputs/__init__.py diff --git a/inputs/rising_bubble.py b/src/pybella/inputs/rising_bubble.py similarity index 62% rename from inputs/rising_bubble.py rename to src/pybella/inputs/rising_bubble.py index 8c94527a..30e86cec 100644 --- a/inputs/rising_bubble.py +++ b/src/pybella/inputs/rising_bubble.py @@ -1,28 +1,27 @@ import numpy as np -import dycore.physics.hydrostatics as hydrostatic -import dycore.utils.boundary as bdry +from ..flow_solver.physics import hydrostatics +from ..flow_solver.utils import boundary as bdry + class UserData(object): # Nsq_ref = grav * 1.3e-05 def __init__(self): - self.grav = 10.0 # [m/s^2] - self.t_ref = 1000.0 # [s] + self.grav = 10.0 # [m/s^2] + self.t_ref = 1000.0 # [s] ########################################## # NUMERICS ########################################## - self.CFL = 0.5 + self.CFL = 0.5 self.dtfixed0 = 100.0 self.dtfixed = 100.0 - self.inx = 160+1 - self.iny = 80+1 + self.inx = 160 + 1 + self.iny = 80 + 1 self.inz = 1 - - - self.tout = np.arange(0.0,1.01,0.01)[10:] + self.tout = np.arange(0.0, 1.01, 0.01)[10:] self.stepmax = 10000 self.output_base_name = "_rising_bubble" @@ -32,8 +31,8 @@ def __init__(self): # self.output_suffix = "_%i_%i_%.1f_psinc" %(self.inx-1,self.iny-1,self.tout[-1]) # if self.continuous_blending == True: # self.output_suffix = "_%i_%i_%.1f" %(self.inx-1,self.iny-1,self.tout[-1]) - - aux = 'debug_imbal_CFLfixed' + + aux = "debug_imbal_CFLfixed" self.aux = aux # self.output_suffix = "_%i_%i_%.1f_%s" %(self.inx-1,self.iny-1,self.tout[-1],aux) # self.output_suffix += '_w=%i-%i' %(self.blending_weight*16.0,16.0-(self.blending_weight*16.0)) @@ -43,39 +42,39 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): u0 = ud.u_wind_speed v0 = ud.v_wind_speed w0 = ud.w_wind_speed - delth = 2.0 # [K] - + delth = 2.0 # [K] + y0 = 0.2 r0 = 0.2 g = ud.gravity_strength[1] # print(ud.rho_ref) - hydrostatic.state(mpv, elem, node, th, ud) + hydrostatics.state(mpv, elem, node, th, ud) x = elem.x y = elem.y - x, y = np.meshgrid(x,y) + x, y = np.meshgrid(x, y) if seed != None: np.random.seed(seed) # y0 += (np.random.random()-.5)/2.0 # delth += 10.0*(np.random.random()-.5) - delth += 10.0*(np.random.random()) - - if 'truth' in ud.aux: + delth += 10.0 * (np.random.random()) + + if "truth" in ud.aux: np.random.seed(1234) # delth += 10.0*(np.random.random()-.5) - delth += 10.0*(np.random.random()) + delth += 10.0 * (np.random.random()) print(delth) - - r = np.sqrt((x)**2 + (y-y0)**2) / r0 - p = np.repeat(mpv.HydroState.p0.reshape(1,-1),elem.icx,axis=0) - rhoY = np.repeat(mpv.HydroState.rhoY0.reshape(1,-1),elem.icx,axis=0) + r = np.sqrt((x) ** 2 + (y - y0) ** 2) / r0 + + p = np.repeat(mpv.HydroState.p0.reshape(1, -1), elem.icx, axis=0) + rhoY = np.repeat(mpv.HydroState.rhoY0.reshape(1, -1), elem.icx, axis=0) - perturbation = (delth/300.0) * (np.cos(0.5 * np.pi * r)**2) + perturbation = (delth / 300.0) * (np.cos(0.5 * np.pi * r) ** 2) perturbation[np.where(r > 1.0)] = 0.0 rho = rhoY / (ud.stratification(y) + perturbation.T) @@ -84,16 +83,16 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): u, v, w = u0, v0, w0 - Sol.rho[x_idx,y_idx] = rho - Sol.rhou[x_idx,y_idx] = rho * u - Sol.rhov[x_idx,y_idx] = rho * v - Sol.rhow[x_idx,y_idx] = rho * w - Sol.rhoY[x_idx,y_idx] = rhoY + Sol.rho[x_idx, y_idx] = rho + Sol.rhou[x_idx, y_idx] = rho * u + Sol.rhov[x_idx, y_idx] = rho * v + Sol.rhow[x_idx, y_idx] = rho * w + Sol.rhoY[x_idx, y_idx] = rhoY p = mpv.HydroState_n.p0[0] rhoY = mpv.HydroState_n.rhoY0[0] mpv.p2_nodes[...] = (p - mpv.HydroState_n.p0[0]) / rhoY / ud.Msq - bdry.set_explicit_boundary_data(Sol,elem,ud,th,mpv) + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - return Sol \ No newline at end of file + return Sol diff --git a/src/utils/__init__.py b/src/pybella/interfaces/__init__.py similarity index 100% rename from src/utils/__init__.py rename to src/pybella/interfaces/__init__.py diff --git a/src/vis/__init__.py b/src/pybella/interfaces/dynamics_blending/__init__.py similarity index 100% rename from src/vis/__init__.py rename to src/pybella/interfaces/dynamics_blending/__init__.py diff --git a/src/pybella/interfaces/dynamics_blending/prepare.py b/src/pybella/interfaces/dynamics_blending/prepare.py new file mode 100644 index 00000000..65e4dc36 --- /dev/null +++ b/src/pybella/interfaces/dynamics_blending/prepare.py @@ -0,0 +1,23 @@ +from . import schemes + + +def initialise(sst): + bld = schemes.Blend(sst.ud) + + sst.interface_params.bld = bld + + +def init_da_window(sst, tout_old, outer_step): + # In ensemble case, do blending for each DA window + if sst.N > 1: + blend = ( + sst.interface_params.bld if tout_old in sst.da_params.dap.da_times else None + ) + else: + blend = sst.interface_params.bld + + # initial blending? + if sst.ud.initial_blending == True and (outer_step == 0 or outer_step == 1): + blend = sst.interface_params.bld + + return blend diff --git a/src/data_assimilation/blending.py b/src/pybella/interfaces/dynamics_blending/schemes.py similarity index 64% rename from src/data_assimilation/blending.py rename to src/pybella/interfaces/dynamics_blending/schemes.py index 51f540ff..bad652ce 100644 --- a/src/data_assimilation/blending.py +++ b/src/pybella/interfaces/dynamics_blending/schemes.py @@ -1,12 +1,11 @@ -import numpy as np -from scipy import signal - -import dycore.physics.gas_dynamics.eos as gd_eos - import logging -import termcolor import copy +import numpy as np +from scipy import signal + +from ...flow_solver.physics.gas_dynamics import eos as gd_eos +from ...utils import io class Blend(object): """ @@ -49,11 +48,15 @@ def convert_p2n(self, p2n): return dp2c - def update_Sol(self, Sol, elem, node, th, ud, mpv, sgn, label=None, writer=None): + def update_sol(self, mem, ud, sgn, label=None, writer=None): if writer != None: writer.populate(str(label) + "_before_blending", "dp2n", self.dp2n) if writer != None: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_before_blending") + writer.write_all(mem, str(label) + "_before_blending") + + Sol = mem.sol + mpv = mem.mpv + th = mem.th if sgn == "bef": sign = -1.0 @@ -94,9 +97,9 @@ def update_Sol(self, Sol, elem, node, th, ud, mpv, sgn, label=None, writer=None) assert 0, "ud.blending_conv undefined." if writer != None: - writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_blending") + writer.write_all(mem, str(label) + "_after_blending") - def update_p2n(self, Sol, mpv, node, th, ud): + def update_p2n(self, mpv): mpv.p2_nodes = self.dp2n @@ -105,45 +108,46 @@ def update_p2n(self, Sol, mpv, node, th, ud): ###################################################### -def do_comp_to_psinc_conv(Sol, mpv, bld, elem, node, th, ud, label, writer): - logging.info(termcolor.colored("Converting COMP to PSINC", "blue")) - dp2n = mpv.p2_nodes +def do_comp_to_psinc_conv(mem, bld, ud, label, writer): + logging.info("Converting COMP to PSINC") + dp2n = mem.mpv.p2_nodes bld.convert_p2n(dp2n) - bld.update_Sol(Sol, elem, node, th, ud, mpv, "bef", label=label, writer=writer) - bld.update_p2n(Sol, mpv, node, th, ud) + bld.update_sol(mem, ud, "bef", label=label, writer=writer) + bld.update_p2n(mem.mpv) - return Sol, mpv + return mem def do_psinc_to_comp_conv( - Sol, flux, mpv, bld, elem, node, th, ud, label, writer, step, window_step, t, dt + sst, + mem, + bld, + label, + writer, + step, + tout, ): - from dycore.discretisation import time_update + from ...flow_solver.discretisation import time_update - logging.info(termcolor.colored("Blending... step = %i" % step, "blue")) - Sol_freeze = copy.deepcopy(Sol) - mpv_freeze = copy.deepcopy(mpv) + logging.info(f"Blending... step = {step}") + Sol_freeze = copy.deepcopy(mem.sol) + mpv_freeze = copy.deepcopy(mem.mpv) - ret = time_update.time_update( - Sol, - flux, - mpv, - t, - t + dt, - ud, - elem, - node, - [0, step - 1], - th, + ret = time_update.do( + sst, + mem, + tout, bld=None, writer=None, - debug=False, + debug_writer=io.NullDebugWriter(), ) + ud = sst.ud + fac_old = ud.blending_weight fac_new = 1.0 - fac_old - dp2n_0 = fac_new * ret[2].p2_nodes_half + fac_old * mpv_freeze.p2_nodes_half - dp2n_1 = fac_new * ret[2].p2_nodes + fac_old * mpv_freeze.p2_nodes + dp2n_0 = fac_new * ret.mpv.p2_nodes_half + fac_old * mpv_freeze.p2_nodes_half + dp2n_1 = fac_new * ret.mpv.p2_nodes + fac_old * mpv_freeze.p2_nodes if ud.blending_type == "half": dp2n = dp2n_0 @@ -157,18 +161,23 @@ def do_psinc_to_comp_conv( str(label) + "_after_full_step", "p2_start", mpv_freeze.p2_nodes ) if writer != None: - writer.populate(str(label) + "_after_full_step", "p2_end", ret[2].p2_nodes) - Sol = Sol_freeze - mpv = mpv_freeze + writer.populate(str(label) + "_after_full_step", "p2_end", ret.mpv.p2_nodes) + mem.Sol = Sol_freeze + mem.mpv = mpv_freeze + + # elem, node, _, _, _, th, _ = mem if writer != None: writer.populate(str(label) + "_after_full_step", "dp2n", dp2n) - logging.info(termcolor.colored("Converting PSINC to COMP", "blue")) + logging.info("Converting PSINC to COMP") bld.convert_p2n(dp2n) - bld.update_Sol(Sol, elem, node, th, ud, mpv, "aft", label=label, writer=writer) - bld.update_p2n(Sol, mpv, node, th, ud) + bld.update_sol(mem, ud, "aft", label=label, writer=writer) + bld.update_p2n(mem.mpv) - return Sol, mpv + # mem.time.step -= 1 + # mem.time.window_step -= 1 + + return mem ###################################################### @@ -177,7 +186,7 @@ def do_psinc_to_comp_conv( def do_swe_to_lake_conv(Sol, mpv, elem, node, ud, th, writer, label, debug): - logging.info(termcolor.colored("swe to lake conversion...", "blue")) + logging.info("swe to lake conversion...") H1 = Sol.rho[ :, @@ -216,7 +225,7 @@ def do_swe_to_lake_conv(Sol, mpv, elem, node, ud, th, writer, label, debug): def do_lake_to_swe_conv( Sol, flux, mpv, elem, node, ud, th, writer, label, debug, step, window_step, t, dt ): - from dycore.discretisation import time_update + from ...flow_solver.discretisation import time_update if debug == True: writer.write_all(Sol, mpv, elem, node, th, str(label) + "_after_lake_time_step") @@ -224,7 +233,7 @@ def do_lake_to_swe_conv( Sol_freeze = copy.deepcopy(Sol) mpv_freeze = copy.deepcopy(mpv) - logging.info(termcolor.colored("doing lake-to-swe time-update...", "blue")) + logging.info("doing lake-to-swe time-update...") ret = time_update.time_update( Sol, flux, @@ -260,7 +269,7 @@ def do_lake_to_swe_conv( mpv.p2_nodes[...] = dp2n H10 = mpv.p2_nodes[:, 2:-2, :].mean(axis=1) - logging.info(termcolor.colored("lake to swe conversion...", "blue")) + logging.info("lake to swe conversion...") H10 -= H10.mean() # define 2D kernel @@ -294,8 +303,7 @@ def do_lake_to_swe_conv( def do_nonhydro_to_hydro_conv( Sol, flux, mpv, bld, elem, node, th, ud, label, writer, step, window_step, t, dt ): - - logging.info(termcolor.colored("nonhydrostatic to hydrostatic conversion...", "blue")) + logging.info("nonhydrostatic to hydrostatic conversion...") # bld.convert_p2n(mpv.p2_nodes) # bld.update_Sol(Sol,elem,node,th,ud,mpv,'bef',label=label,writer=writer) # Sol.rhov = Sol.rhov_half @@ -324,9 +332,8 @@ def do_nonhydro_to_hydro_conv( def do_hydro_to_nonhydro_conv( Sol, flux, mpv, bld, elem, node, th, ud, label, writer, step, window_step, t, dt ): - - logging.info(termcolor.colored("hydrostatic to nonhydrostatic conversion...", "blue")) - logging.info(termcolor.colored("Blending... step = %i" % step, "blue")) + logging.info("hydrostatic to nonhydrostatic conversion...") + logging.info(f"Blending... step = {step}") # Sol_tmp = deepcopy(Sol) # flux_tmp = deepcopy(flux) @@ -377,6 +384,154 @@ def do_hydro_to_nonhydro_conv( # bld.update_Sol(Sol,elem,node,th,ud,mpv,'aft',label=label,writer=writer) # bld.update_p2n(Sol,mpv,node,th,ud) # + + ############################### + # alternative version + ############################### + + # if c1 or c2: + # logging.info( + # termcolor.colored("hydrostatic to nonhydrostatic conversion...", "blue") + # ) + + # writer.write_all(mem, str(label) + "_half_full") + # writer.populate(str(label) + "_ic", "pwchi", Sol.pwchi) + + # if test_hydrob == False: + # Sol = copy.deepcopy(Sol_half_old) + # # mpv = copy.deepcopy(mpv_half_old) + + # logging.info(termcolor.colored("test_hydrob == False", "red")) + # writer.write_all(mem, str(label) + "_quarter") + + # writer.populate(str(label) + "_quarter", "pwchi", Sol.pwchi) + + # logging.info("quarter dt = %.8f" % (dt * 0.5)) + + # ret = do( + # Sol_half_old, + # flux_half_old, + # mpv_half_old, + # dt - 0.5 * dt, + # dt + 0.5 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol_tu = copy.deepcopy(ret[0]) + # # mpv_tu = copy.deepcopy(ret[2]) + # Sol.rho[...] = Sol_tu.rho_half + # Sol.rhou[...] = Sol_tu.rhou_half + # Sol.rhov[...] = Sol_tu.rhov_half + # Sol.rhow[...] = Sol_tu.rhow_half + # Sol.rhoX[...] = Sol_tu.rhoX_half + # Sol.rhoY[...] = Sol_tu.rhoY_half + # Sol.pwchi[...] = Sol_tu.pwchi + + # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half + + # writer.write_all(mem, str(label) + "_half") + + # writer.populate(str(label) + "_half", "pwchi", Sol.pwchi) + + # ret = do( + # Sol, + # flux, + # mpv, + # dt, + # 2.0 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol = copy.deepcopy(ret[0]) + # flux = copy.deepcopy(ret[1]) + # mpv = copy.deepcopy(ret[2]) + + # if test_hydrob == True: + # Sol = copy.deepcopy(Sol_half_old) + # # mpv = copy.deepcopy(mpv_half_old) + + # logging.info(termcolor.colored("test_hydrob == False", "red")) + # writer.write_all(mem, str(label) + "_quarter") + + # # writer.populate(str(label)+'_quarter', 'pwchi', Sol.pwchi) + + # logging.info("quarter dt = %.8f" % (dt * 0.5)) + + # ret = do( + # Sol_half_old, + # flux_half_old, + # mpv_half_old, + # dt - 0.5 * dt, + # dt + 0.5 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol_tu = copy.deepcopy(ret[0]) + # # mpv_tu = copy.deepcopy(ret[2]) + # Sol.rho[...] = Sol_tu.rho_half + # Sol.rhou[...] = Sol_tu.rhou_half + # Sol.rhov[...] = Sol_tu.rhov_half + # Sol.rhow[...] = Sol_tu.rhow_half + # Sol.rhoX[...] = Sol_tu.rhoX_half + # Sol.rhoY[...] = Sol_tu.rhoY_half + # Sol.pwchi[...] = Sol_tu.pwchi + + # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half + + # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') + + # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) + + # ret = do( + # Sol, + # flux, + # mpv, + # dt, + # 2.0 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol = copy.deepcopy(ret[0]) + # flux = copy.deepcopy(ret[1]) + # mpv = copy.deepcopy(ret[2]) + # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') + # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) + + # logging.info(termcolor.colored("test_hydrob == True", "red")) + + # if test_hydrob == False: + # dt *= 2.0 + # if c2: + # ud.is_nonhydrostatic = 1 + return Sol, mpv @@ -384,14 +539,9 @@ def do_hydro_to_nonhydro_conv( # Blending calls from data.py ###################################################### def blending_before_timestep( - Sol, - flux, - mpv, + sst, + mem, bld, - elem, - node, - th, - ud, label, writer, step, @@ -404,6 +554,10 @@ def blending_before_timestep( ###################################################### # Blending : Do full regime to limit regime conversion ###################################################### + # do unpacking + elem, node, Sol, flux, mpv, th, _ = mem + ud = sst.ud + # these make sure that we are the correct window step if bld is not None and window_step == 0: # these make sure that blending switches are on @@ -413,9 +567,7 @@ def blending_before_timestep( do_swe_to_lake_conv(Sol, mpv, elem, node, ud, th, writer, label, debug) swe_to_lake = True else: - Sol, mpv = do_comp_to_psinc_conv( - Sol, mpv, bld, elem, node, th, ud, label, writer - ) + mem = do_comp_to_psinc_conv(mem, bld, ud, label, writer) ###################################################### # Blending : Do full steps or transition steps? @@ -432,21 +584,15 @@ def blending_before_timestep( if c_init and bld.cb and ud.blending_conv is not None: # distinguish between Euler and SWE blending if ud.blending_conv != "swe": - Sol, mpv = do_psinc_to_comp_conv( - Sol, - flux, - mpv, + do_psinc_to_comp_conv( + sst, + mem, bld, - elem, - node, - th, ud, label, writer, step, - window_step, - t, - dt, + t + dt, ) ###################################################### @@ -459,9 +605,7 @@ def blending_before_timestep( if bld.psinc_init > 0: ud.is_compressible = 0 ud.compressibility = 0.0 - Sol, mpv = do_comp_to_psinc_conv( - Sol, mpv, bld, elem, node, th, ud, label, writer - ) + mem = do_comp_to_psinc_conv(mem, bld, ud, label, writer) elif bld.hydro_init > 0: Sol, mpv, t = do_nonhydro_to_hydro_conv( Sol, @@ -494,21 +638,14 @@ def blending_before_timestep( ): # Distinguish between SWE and Euler blendings if ud.blending_conv != "swe": - Sol, mpv = do_psinc_to_comp_conv( - Sol, - flux, - mpv, + do_psinc_to_comp_conv( + sst, + mem, bld, - elem, - node, - th, - ud, label, writer, step, - window_step, - t, - dt, + t + dt, ) ud.is_compressible = 1 ud.compressibility = 1.0 @@ -602,3 +739,62 @@ def blending_after_timestep( ud.compressibility = 1.0 return Sol, mpv + + +def prepare_blending( + sst, + mem, + bld, + label, + writer, + step, + window_step, + t, + dt, + swe_to_lake, + debug, + ): + + ud = sst.ud + if check_and_apply_initial_hydrostatic_conversion(step, ud, bld): + ud.is_nonhydrostatic = 0 + + swe_to_lake, Sol, mpv, t = blending_before_timestep( + sst, + mem, + bld, + label, + writer, + step, + window_step, + t, + dt, + swe_to_lake, + debug, + ) + + return swe_to_lake, Sol, mpv, t + + +def check_and_apply_initial_hydrostatic_conversion(step, ud, bld): + """ + Check and apply initial blending conversion if needed. + + Returns: + bool: True if conversion was applied, False otherwise + """ + if step != 0 or bld is None or "imbal" not in ud.aux: + return False + + hydrostatic_case = ud.is_nonhydrostatic == 0 + nonhydrostatic_case = ( + ud.is_nonhydrostatic == 1 and ud.initial_blending == True + ) + + if hydrostatic_case or nonhydrostatic_case: + logging.info("nonhydrostatic to hydrostatic conversion...") + ud.is_nonhydrostatic = 0 + return True + + return False + diff --git a/src/pybella/interfaces/ic_config.py b/src/pybella/interfaces/ic_config.py new file mode 100644 index 00000000..067e5af1 --- /dev/null +++ b/src/pybella/interfaces/ic_config.py @@ -0,0 +1,27 @@ +IC_MODULES = { + "bi": "inputs.baroclinic_instability_periodic", + "tv": "inputs.travelling_vortex_2D", + "tv_2d": "inputs.travelling_vortex_2D", + "tv_neg": "inputs.travelling_vortex_2D_neg", + "tv_3d": "inputs.travelling_vortex_3D", + "tv_corr": "inputs.travelling_vortex_3D_Coriolis", + "aw": "inputs.acoustic_wave_high", + "igw": "inputs.internal_long_wave", + "igw_3d": "inputs.internal_long_wave_3D", + "lbw": "inputs.lamb_waves", + "skl": "inputs.sk_lamb_wave", + "mark": "inputs.mark", + "lw_p": "inputs.lamb_wave_perturb", + "igw_bb": "inputs.igw_baldauf_brdar", + "rb": "inputs.rising_bubble", + "rbc": "inputs.rising_bubble_cold", + "swe_bal_vortex": "inputs.swe_bal_vortex", + "swe": "inputs.shallow_water_3D", + "swe_icshear": "inputs.shallow_water_3D_icshear", + "swe_dvortex": "inputs.shallow_water_3D_dvortex", + "test_travelling_vortex": "pybella.tests.test_travelling_vortex", + "test_internal_long_wave": "pybella.tests.test_internal_long_wave", + "test_lamb_wave": "pybella.tests.test_lamb_wave", + "test_blending_warm_bubble": "pybella.tests.test_blending_warm_bubble", + "test_unstable_lamb": "pybella.tests.test_unstable_lamb", +} \ No newline at end of file diff --git a/src/pybella/interfaces/parallel_intergrator/__init__.py b/src/pybella/interfaces/parallel_intergrator/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pybella/interfaces/postprocessing/strip_target_file.py b/src/pybella/interfaces/postprocessing/strip_target_file.py new file mode 100644 index 00000000..8c9b5948 --- /dev/null +++ b/src/pybella/interfaces/postprocessing/strip_target_file.py @@ -0,0 +1,44 @@ +import h5py +import re + +def do(h5path: str, ud, output_intervals: bool = False, time_increment: bool = False) -> None: + if ud.diag_updt_targets: + stepmax = ud.stepmax + + if output_intervals: + # Determine steps to keep + if stepmax <= 101: + keep_steps = list(range(10, stepmax, 10)) + elif 101 < stepmax <= 1000: + keep_steps = list(range(100, stepmax, 100)) + else: + raise ValueError(f"stepmax {stepmax} > 1000 is too large as a diagnostic target.") + else: + keep_steps = [stepmax - 1] + + # If time_increment is True, add the step before each kept step + if time_increment: + additional_steps = [step - 1 for step in keep_steps if step > 0] + keep_steps.extend(additional_steps) + keep_steps = sorted(list(set(keep_steps))) + + output_path = h5path.replace(".h5", "_stripped.h5") + + with h5py.File(h5path, "r") as src, h5py.File(output_path, "w") as dst: + # Copy all top-level groups, filtering time-tagged datasets + for group in src.keys(): + src_group = src[group] + dst_group = dst.create_group(group) + + for dset_name in src_group.keys(): + # Identify timestep in dataset name, e.g. p2_nodes_010_after_full_step + step_match = re.search(r'_(\d+)_after_full_step$', dset_name) + if step_match: + step = int(step_match.group(1)) + if step in keep_steps: + src_group.copy(dset_name, dst_group) + else: + # Copy non-timestep fields like "p2_nodes_ic" + src_group.copy(dset_name, dst_group) + + print(f"Stripped target file written to: {output_path}") \ No newline at end of file diff --git a/src/pybella/interfaces/time_stepper/__init__.py b/src/pybella/interfaces/time_stepper/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pybella/interfaces/time_stepper/prestep.py b/src/pybella/interfaces/time_stepper/prestep.py new file mode 100644 index 00000000..3206579a --- /dev/null +++ b/src/pybella/interfaces/time_stepper/prestep.py @@ -0,0 +1,17 @@ +""" +Pre-step modifications for the time stepper. +""" + +def apply_modifcations(dt, ud, step): + """Apply modifications to the timestep based on user-defined parameters.""" + + dt = _apply_cfl_override(dt, ud, step) + + return dt + + +def _apply_cfl_override(dt, ud, step): + """Apply CFL override for fixed timestep cases.""" + if "CFLfixed" in ud.aux and step < 2: + return 21.69 / ud.t_ref + return dt \ No newline at end of file diff --git a/src/pybella/tests/__init__.py b/src/pybella/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pybella/tests/diagnostics.py b/src/pybella/tests/diagnostics.py new file mode 100644 index 00000000..35f8b6c2 --- /dev/null +++ b/src/pybella/tests/diagnostics.py @@ -0,0 +1,231 @@ +import logging +import copy + +import numpy as np +import yaml + +from ..utils.data_structures import DiagnosticState +from ..flow_solver.utils.solver_diagnostics import get_p_from_pressure_related_fields + +from ..vis import utils as vis_utils, plotting_tools as vis_pt + + +class CompareSol(object): + def __init__(self, diag_state: DiagnosticState): + self.diag_state = diag_state + self.current_run = diag_state.test_name + self.plot = diag_state.plot_compare + self.tolerances = diag_state.tolerances + self.time_increment = diag_state.time_increment + self.__init(diag_state) + self.__get_tc() + + def update_targets(self): + self.arr_dump = {} + + for tc_name, tc in self.tcs.items(): + tp = self.tps[tc_name] + dump_name = tp.name.replace("target","test") + self.arr_dump[dump_name] = {} + + for attribute in tp.attributes: + arr = self.__get_ens(tc, tp, attribute, time_increment=self.time_increment, summed=False) + self.arr_dump[dump_name][attribute] = float( + arr.sum() + ) + + if self.plot: + # vis_pt.plotter accepts a list of tuples with plot and panel title. + pl = vis_pt.plotter([(arr.T, "ref"), ], ncols=1, figsize=(4, 3), sharey=False) + _ = pl.plot(method="contour", lvls=None, suptitle=attribute) + pl.img.savefig(tp.dir + attribute + ".png") + + # with open("./src/tests/test_targets.yml", "a") as outfile: + # yaml.dump(self.arr_dump, outfile, default_flow_style=False) + + def test_do(self, mem, ud): + tc = self.tcs[self.current_run] + tp = self.tps[self.current_run] + + # populate reference arrays + ref_mem = copy.deepcopy(mem) + for attribute in tp.attributes: + try: + ref_data = self.__get_ens(tc, tp, attribute, time_increment=False, summed=False) + if attribute != "p2_nodes": + setattr(ref_mem.sol, attribute, ref_data) + else: + if self.time_increment: + # in the case where we are interested in the time increment + # of the pressure-related fields, we need to load the previous + # time step from the output file. + tc_test = copy.deepcopy(tc) + tc_test.base_fn = tc_test.base_fn.replace("target", "test") + tc_test.base_fn = tc_test.base_fn.replace("_stripped.h5", ".h5") + tc_test.py_dir = tc_test.py_dir.replace("target", "test") + data = self.__get_ens(tc_test, tp, attribute, time_increment=self.time_increment, summed=False) + setattr(mem.mpv, attribute, data) + ref_data = self.__get_ens(tc, tp, attribute, time_increment=self.time_increment, summed=False) + setattr(ref_mem.mpv, attribute, ref_data) + except Exception as e: + raise AssertionError(f"test {self.current_run} has no target for comparison: {e}") + + if self.plot: + self.__plot_comparison(mem, ref_mem, ud) + + for attribute in tp.attributes: + test = self.__get_sol_for_comparison(mem, ud, attribute) + ref = self.__get_sol_for_comparison(ref_mem, ud, attribute) + + l2_error = np.linalg.norm(test - ref) + ref_norm = np.linalg.norm(ref) + + if ref_norm > 0.0: + rel_l2_error = l2_error / ref_norm + else: + rel_l2_error = np.inf if l2_error > self.tolerances[attribute] else 0.0 + max_abs_error = np.max(np.abs(test - ref)) + + try: + assert max_abs_error < self.tolerances[attribute], ( + "Relative L2 error for attribute %s of %s exceeds tolerance:\n" + "L2 error: %.6e\nRelative L2 error: %.6e\nMax abs error: %.6e\nTolerance: %.6e" + % ( + attribute, + self.current_run, + l2_error, + rel_l2_error, + max_abs_error, + self.tolerances[attribute], + ) + ) + logging.info( + f"Test passed for {attribute} | " + f"L2: {l2_error:.2e}, Rel L2: {rel_l2_error:.2e}, Max Abs: {max_abs_error:.2e}" + ) + except AssertionError as e: + logging.info(str(e)) + raise + + logging.info( + f""" + {'#' * 10} + Test passed for {self.current_run} + {'#' * 10} + """.strip() + ) + + def __init(self, ds: DiagnosticState): + tp = test_params(ds) + + self.tps = { + ds.test_name: tp, + } + # self.tps = [tv_2D] + + def __get_tc(self): + self.tcs = {} + for test_name, test_param in self.tps.items(): + fn = test_param.fn + ".h5" + tc = vis_utils.test_case( + fn, test_param.dir, test_param.Nx, test_param.Ny, "" + ) + + self.tcs[test_name] = tc + + def __read_yaml(self): + with open("./src/tests/test_targets.yml", "r") as infile: + self.target = yaml.safe_load(infile) + + def __plot_comparison(self, mem, ref_mem, ud): + tp = self.tps[self.current_run] + + for attribute in tp.attributes: + arr_plots = [] + + test_sol = self.__get_sol_for_comparison(mem, ud, attribute) + ref_sol = self.__get_sol_for_comparison(ref_mem, ud, attribute) + + arr_plots.append([ref_sol, "ref"]) + arr_plots.append([test_sol, "test"]) + arr_plots.append([ref_sol - test_sol, "diff"]) + + pl = vis_pt.plotter(arr_plots, ncols=3, figsize=(12, 3), sharey=False) + _ = pl.plot(method="contour", lvls=None, suptitle=attribute) + pl.img.savefig(tp.dir.replace("target", "test") + attribute + ".png") + + del ref_mem + + @staticmethod + def __get_sol_for_comparison(mem, ud, attribute): + Sol = mem.sol + mpv = mem.mpv + if attribute != "p2_nodes": + test_sol = np.copy(getattr(Sol, attribute).T) + if attribute != "rho": + rho = getattr(Sol, "rho").T + test_sol /= rho + + # if attribute == 'rhoY': + # test_sol -= mpv.HydroState.Y0[:,np.newaxis] + + else: + # test_sol = mpv.p2_nodes.T * ud.Msq + # test_sol -= mpv.HydroState_n.pi0[:,np.newaxis] + test_sol = get_p_from_pressure_related_fields(mem, ud).T + # pass + + return test_sol + + @staticmethod + def __get_ens(tc, params, attribute, time_increment=False, summed=True, normed=False): + if time_increment and attribute == "p2_nodes": + times = [params.times[0]-1, params.times[0]] + else: + times = params.times + l_typ = params.l_typ + + tags = tc.get_tag_dict() + tag = "ic" if times[0] == 0.0 else tags[9] + + ens = tc.get_ensemble( + times, + 1, + attribute, + "", + label_type=l_typ, + tag=tag, + inner=False, + get_fn=False, + fn=tc.base_fn, + load_ic=False, + avg=False, + ) + + if time_increment and attribute == "p2_nodes": + ens = ens[1] - ens[0] + else: + ens = ens[0] # removes time axis + ens = ens[0] # removes ensemble axis + + if summed: + return ens.sum() + elif normed: + return np.linalg.norm(ens) + else: + return ens + + +class test_params(object): + def __init__(self, ds: DiagnosticState): + self.name = ds.test_name + self.dir = ds.path + ds.file_name + "/" + self.fn = f"{ds.file_name}_{ds.Nx}_{ds.Ny}_stripped" + + self.Nx = ds.Nx + self.Ny = ds.Ny + + self.times = ds.steps + self.l_typ = "WINDOW_STEP" + + self.attributes = ["rho", "rhou", "rhov", "rhow", "rhoY", "rhoX", "p2_nodes"] diff --git a/src/pybella/tests/test_blending_warm_bubble.py b/src/pybella/tests/test_blending_warm_bubble.py new file mode 100644 index 00000000..67022857 --- /dev/null +++ b/src/pybella/tests/test_blending_warm_bubble.py @@ -0,0 +1,118 @@ +import numpy as np +from ..flow_solver.physics import hydrostatics +from ..flow_solver.utils import boundary as bdry + +from ..utils.data_structures import DiagnosticState + + +class UserData(object): + # Nsq_ref = grav * 1.3e-05 + + def __init__(self): + self.grav = 10.0 # [m/s^2] + self.t_ref = 1000.0 # [s] + + ########################################## + # NUMERICS + ########################################## + self.CFL = 0.9 + self.dtfixed0 = 100.0 + self.dtfixed = 100.0 + + self.inx = 64 + 1 + self.iny = 48 + 1 + self.inz = 1 + + self.tout = [1000.0] + self.stepmax = 31 + + self.is_compressible = 1 + + self.continuous_blending = False + self.no_of_pi_initial = 1 + self.no_of_pi_transition = 0 + self.no_of_hy_initial = 0 + self.no_of_hy_transition = 0 + + self.initial_blending = True + + self.diag = True + self.diag_updt_targets = False + + self.output_base_name = "_blending_warm_bubble" + self.output_type = "test" if not self.diag_updt_targets else "target" + self.aux = "CFLfixed" + + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) + + self.output_timesteps = True + + self.diag_state = DiagnosticState( + test_name=f"{self.output_type}_blending_warm_bubble", + file_name="target_blending_warm_bubble", + Nx=self.inx - 1, + Ny=self.iny - 1, + steps=[self.stepmax - 1], + plot_compare=True, + time_increment=True, + # The only thing that matters here is that + # p2_nodes remains small + tolerances={ + "rho": 1.0e-0, + "rhou": 1.0e-0, + "rhov": 1.0e-0, + "rhow": 1.0e-0, + "rhoY": 1.0e-0, + "rhoX": 1.0e-0, + "p2_nodes": 1.0e-4, + } + ) + + self.autogen_fn = False + + +def sol_init(Sol, mpv, elem, node, th, ud, seed=None): + u0 = ud.u_wind_speed + v0 = ud.v_wind_speed + w0 = ud.w_wind_speed + delth = 2.0 # [K] + + y0 = 0.2 + r0 = 0.2 + + hydrostatics.integrated_state(mpv, elem, node, th, ud) + + x = elem.x + y = elem.y + + x, y = np.meshgrid(x, y) + + r = np.sqrt((x) ** 2 + (y - y0) ** 2) / r0 + + p = np.repeat(mpv.HydroState.p0.reshape(1, -1), elem.icx, axis=0) + rhoY = mpv.HydroState.rhoY0[ + np.newaxis, : + ] + + perturbation = (delth / 300.0) * (np.cos(0.5 * np.pi * r) ** 2) + perturbation[np.where(r > 1.0)] = 0.0 + rho = rhoY / (ud.stratification(y) + perturbation.T) + + x_idx = slice(None) + y_idx = slice(None) + + u, v, w = u0, v0, w0 + + Sol.rho[x_idx, y_idx] = rho + Sol.rhou[x_idx, y_idx] = rho * u + Sol.rhov[x_idx, y_idx] = rho * v + Sol.rhow[x_idx, y_idx] = rho * w + Sol.rhoY[x_idx, y_idx] = rhoY + + p = mpv.HydroState_n.p0 + rhoY = mpv.HydroState_n.rhoY0 + mpv.p2_nodes[...] = (p - mpv.HydroState_n.p0) / rhoY / ud.Msq + + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + + return Sol diff --git a/src/tests/test_internal_long_wave.py b/src/pybella/tests/test_internal_long_wave.py similarity index 85% rename from src/tests/test_internal_long_wave.py rename to src/pybella/tests/test_internal_long_wave.py index c59db5db..3f79e0af 100644 --- a/src/tests/test_internal_long_wave.py +++ b/src/pybella/tests/test_internal_long_wave.py @@ -1,8 +1,9 @@ import numpy as np -import dycore.utils.options as opts -import dycore.utils.boundary as bdry -import dycore.utils.variable as var -import dycore.physics.hydrostatics as hydrostatic + +from ..flow_solver.utils import options as opts, boundary as bdry, variable as var +from ..flow_solver.physics import hydrostatics + +from ..utils.data_structures import DiagnosticState class UserData(object): @@ -67,13 +68,13 @@ def __init__(self): # self.coriolis_strength[2] = self.omega * self.t_ref # self.coriolis_strength[1] = self.omega * self.t_ref - for i in range(3): - if (self.gravity_strength[i] > np.finfo(np.float).eps) or (i == 1): - self.i_gravity[i] = 1 - self.gravity_direction = i + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + self.i_gravity = gravity_mask.astype(int) + if np.any(gravity_mask): + self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - if self.coriolis_strength[i] > np.finfo(np.float).eps: - self.i_coriolis[i] = 1 + coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps + self.i_coriolis = coriolis_mask.astype(int) self.xmin = -15.0 * self.scale_factor self.xmax = 15.0 * self.scale_factor @@ -142,11 +143,6 @@ def __init__(self): self.autogen_fn = False - self.output_base_name = "_internal_long_wave" - self.output_type = "test" - self.aux = "" - self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) - self.output_timesteps = True self.stratification = self.stratification_function @@ -154,8 +150,20 @@ def __init__(self): self.rhoe = self.rhoe_method self.diag = True - self.diag_current_run = "test_internal_long_wave" - self.diag_plot_compare = False + self.diag_updt_targets = False + + self.output_base_name = "_internal_long_wave" + self.output_type = "test" if not self.diag_updt_targets else "target" + self.aux = "" + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) + + self.diag_state = DiagnosticState( + test_name="test_internal_long_wave", + file_name="target_internal_long_wave", + Nx=self.inx - 1, + Ny=self.iny - 1, + steps=[self.stepmax - 1], + ) def stratification_function(self, y): Nsq = self.Nsq_ref * self.t_ref * self.t_ref @@ -188,7 +196,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): xc = 0.0 a = ud.scale_factor * 5.0e3 / ud.h_ref - hydrostatic.state(mpv, elem, node, th, ud) + hydrostatics.analytical_state(mpv, elem, node, th, ud) HySt = var.States(node.sc, ud) HyStn = var.States(node.sc, ud) @@ -207,7 +215,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): 1.0 + (xn - xc) ** 2 / (a**2) ) - hydrostatic.column(HySt, HyStn, Y, Yn, elem, node, th, ud) + hydrostatics.column(HySt, HyStn, Y, Yn, elem, node, th, ud) x_idx = slice(None) y_idx = slice(elem.igy, -elem.igy + 1) @@ -238,7 +246,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): mpv.p2_nodes[:, elem.igy : -elem.igy] = HyStn.p20[:, elem.igy : -elem.igy] - hydrostatic.initial_pressure(Sol, mpv, elem, node, ud, th) + hydrostatics.initial_pressure(Sol, mpv, elem, node, ud, th) ud.nonhydrostasy = 1.0 if ud.is_nonhydrostatic == 1 else 0.0 ud.compressibility = 1.0 if ud.is_compressible == 1 else 0.0 diff --git a/src/pybella/tests/test_lamb_wave.py b/src/pybella/tests/test_lamb_wave.py new file mode 100644 index 00000000..9d885688 --- /dev/null +++ b/src/pybella/tests/test_lamb_wave.py @@ -0,0 +1,390 @@ +import numpy as np + +from ..flow_solver.utils import options as opts, boundary as bdry + +from ..flow_solver.physics import hydrostatics + +from ..utils.data_structures import DiagnosticState + + +class UserData(object): + NSPEC = 1 + grav = 9.81 # [m s^{-2}] + omega = 0.0 * 1e-5 # [s^{-1}] + + R_gas = 287.4 # [J kg^{-1} K^{-1}] + R_vap = 461.0 + Q_vap = 2.53e06 + gamma = 1.4 + cp_gas = gamma * R_gas / (gamma - 1.0) + + p_ref = 1e5 + T_ref = 300.00 # [K] + rho_ref = p_ref / (R_gas * T_ref) + N_ref = grav / np.sqrt(cp_gas * T_ref) + Cs = np.sqrt(gamma * R_gas * T_ref) + + h_ref = 10.0e3 # [m] + t_ref = 100.0 # [s] + u_ref = h_ref / t_ref + + i_gravity = np.zeros((3)) + i_coriolis = np.zeros((3)) + + def __init__(self): + self.h_ref = self.h_ref + self.t_ref = self.t_ref + self.T_ref = self.T_ref + self.p_ref = self.p_ref + self.rho_ref = self.rho_ref + self.u_ref = self.u_ref + self.Nsq_ref = self.N_ref * self.N_ref + self.g_ref = self.grav + self.gamm = self.gamma + self.Rg_over_Rv = self.R_gas / self.R_vap + self.Q = self.Q_vap / (self.R_gas * self.T_ref) + self.R_gas = self.R_gas + self.Rg = self.R_gas / (self.h_ref**2 / self.t_ref**2 / self.T_ref) + self.cp_gas = self.cp_gas + + self.nspec = self.NSPEC + + self.is_nonhydrostatic = 1 + self.is_compressible = 1 + self.is_ArakawaKonor = 0 + + self.compressibility = 1.0 + self.acoustic_timestep = 0 + self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) + + self.gravity_strength = np.zeros((3)) + self.coriolis_strength = np.zeros((3)) + + self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) + self.coriolis_strength[2] = 2.0 * self.omega * self.t_ref + + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + self.i_gravity = gravity_mask.astype(int) + if np.any(gravity_mask): + self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index + + coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps + self.i_coriolis = coriolis_mask.astype(int) + + j = 4.0 + Lx = 1.0 * np.pi * self.Cs / self.N_ref * j + self.xmin = -Lx / self.h_ref + self.xmax = Lx / self.h_ref + self.ymin = -0.0 + self.ymax = 8.0 + self.zmin = -1.0 + self.zmax = 1.0 + + self.u_wind_speed = 0.0 * self.u_ref + self.v_wind_speed = 0.0 + self.w_wind_speed = 0.0 + + self.bdry_type = np.empty((3), dtype=object) + self.bdry_type[0] = opts.BdryType.PERIODIC + self.bdry_type[1] = opts.BdryType.WALL + self.bdry_type[2] = opts.BdryType.WALL + self.ATMOSPHERIC_EXTENSION = True + self.rayleigh_bdry_switch = False + + ########################################## + # NUMERICS + ########################################## + self.CFL = 0.9 + + self.inx = 151 + 1 + self.iny = 60 + 1 + self.inz = 1 + + self.dtfixed0 = 10.0 / self.t_ref + self.dtfixed = self.dtfixed0 + + self.do_advection = True + self.limiter_type_scalars = opts.LimiterType.NONE + self.limiter_type_velocity = opts.LimiterType.NONE + + self.tol = 1.0e-30 + self.max_iterations = 10000 + + # blending parameters + self.perturb_type = "pos_perturb" + self.blending_mean = "rhoY" # 1.0, rhoY + self.blending_conv = "rho" # theta, rho + self.blending_type = "half" # half, full + + self.continuous_blending = False + self.no_of_pi_initial = 1 + self.no_of_pi_transition = 0 + self.no_of_hy_initial = 0 + self.no_of_hy_transition = 0 + + self.blending_weight = 0.0 / 16 + self.initial_blending = False + self.initial_projection = True + + self.tout = [360.0] + # self.tout = np.arange(0,361,1.0) + # self.tout = np.append(self.tout, [720.0]) + self.stepmax = 301 + self.output_timesteps = True + + self.autogen_fn = False + + self.diag = True + self.diag_updt_targets = False + + self.output_base_name = "_lamb_wave" + self.output_type = "test" if not self.diag_updt_targets else "target" + self.aux = "" + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) + + self.diag_state = DiagnosticState( + test_name="test_lamb_wave", + file_name="target_lamb_wave", + Nx=self.inx - 1, + Ny=self.iny - 1, + steps=[self.stepmax - 1], + ) + + self.stratification = self.stratification_wrapper + self.init_forcing = self.forcing + + self.rayleigh_forcing = False + self.rayleigh_forcing_type = "func" # func or file + self.rayleigh_forcing_fn = None + self.rayleigh_forcing_path = None + + def stratification_wrapper(self, dy): + return lambda y: self.stratification_function(y, dy) + + def stratification_function(self, y, dy): + g = self.gravity_strength[1] + + Gamma = (self.gamm - 1.0) / self.gamm + Hex = 1.0 / (Gamma * g) + pi_m = np.exp(-(y - 0.5 * dy) / Hex) + pi_p = np.exp(-(y + 0.5 * dy) / Hex) + + Theta = -(Gamma * g * dy) / (pi_p - pi_m) + + return Theta + + @staticmethod + def rayleigh_bc_function(ud): + if ud.bdry_type[1] == opts.BdryType.RAYLEIGH or ud.rayleigh_forcing == True: + ud.inbcy = ud.iny - 1 + ud.iny0 = np.copy(ud.iny) + ud.iny = ud.iny0 + int(3 * ud.inbcy) + + # tentative workaround + ud.bcy = ud.ymax + ud.ymax += 3.0 * ud.bcy + + class forcing(object): + def __init__( + self, + k, + mu, + Cs, + F, + N, + Gamma, + ampl, + g, + rhobar, + Ybar, + rhobar_n, + Ybar_n, + X, + Y, + Xn, + Yn, + ): + self.k = k + self.mu = mu + self.Cs = Cs + self.F = F + self.N = N + self.Gamma = Gamma + + self.oorhobarsqrt = 1.0 / np.sqrt(rhobar.T) + self.Ybar = Ybar.T + + self.oorhobarsqrt_n = 1.0 / np.sqrt(rhobar_n.T) + self.Ybar_n = Ybar_n.T + + self.g = g + self.ampl = ampl + + self.X = X + self.Y = Y + self.Xn = Xn + self.Yn = Yn + + def get_T_matrix(self): + # system matrix of linearized equations + matrix = -np.array( + [ + [0, self.F, 0, 1j * self.Cs * self.k], + [-self.F, 0, -self.N, self.Cs * (self.mu + self.Gamma)], + [0, self.N, 0, 0], + [1j * self.Cs * self.k, self.Cs * (self.mu - self.Gamma), 0, 0], + ] + ) + + self.T_matrix = matrix + + def eigenfunction(self, t, s, grid="c"): + if grid == "c": + x, z = self.X, self.Y + elif grid == "n": + ( + x, + z, + ) = ( + self.Xn, + self.Yn, + ) + + # Compute eigenvalues and eigenvectors + eigval, eigvec = np.linalg.eig(self.T_matrix) + + # Find index of eigenvalue + # with greatest real part aka the instability growth rate + ind = np.argmax(np.real(eigval)) + + # construct solution according to eq. 2.27 and 2.19 + exponentials = np.exp( + 1j * self.k * x + self.mu * z + (eigval[ind]) * (t) + 1j * s * t + ) + chi_u = self.ampl * np.real(eigvec[0, ind] * exponentials) + chi_w = self.ampl * np.real(eigvec[1, ind] * exponentials) + chi_th = self.ampl * np.real(eigvec[2, ind] * exponentials) + chi_pi = self.ampl * np.real(eigvec[3, ind] * exponentials) + + self.arrs = (chi_u, chi_w, chi_th, chi_pi) + + def dehatter(self, th, grid="c"): + if grid == "n": + Ybar = self.Ybar_n + oorhobarsqrt = self.oorhobarsqrt_n + elif grid == "c": + Ybar = self.Ybar + oorhobarsqrt = self.oorhobarsqrt + + chi_u, chi_v, chi_Y, chi_pi = self.arrs + + up = oorhobarsqrt * chi_u + vp = oorhobarsqrt * chi_v + Yp = oorhobarsqrt * self.N / self.g * Ybar * chi_Y + pi_p = oorhobarsqrt * self.Cs / Ybar / th.Gammainv * chi_pi + + return up.T, vp.T, Yp.T, pi_p.T + + +def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): + if hasattr(ud, "rayleigh_bdry_switch"): + if ud.rayleigh_bdry_switch: + ud.bdry_type[1] = opts.BdryType.RAYLEIGH + + if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: + ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) + + if ud.rayleigh_forcing: + ud.forcing_tcy, ud.forcing_tny = bdry.get_bottom_tau_y( + ud, elem, node, 0.2, cutoff=0.3 + ) + + A0 = 1.0e-1 / ud.u_ref + Msq = ud.Msq + g = ud.gravity_strength[1] * ud.Rg + + x = elem.x.reshape(-1, 1) + y = elem.y.reshape(1, -1) + X, Y = np.meshgrid(x, y) + + xn = node.x.reshape(-1, 1) + yn = node.y.reshape(1, -1) + + dy = np.diff(node.y)[0] + + Xn, Yn = np.meshgrid(xn, yn) + + ################################################## + # Following Rupert's fix, reinitialise all background quantities + # as derived from one quantity. + ud.stratification = ud.stratification(dy) + + # Use hydrostatically balanced background + hydrostatics.analytical_state(mpv, elem, node, th, ud) + rhobar = mpv.HydroState.rho0.reshape(1, -1) + Ybar = mpv.HydroState.Y0.reshape(1, -1) + pibar = mpv.HydroState.p20.reshape(1, -1) * ud.Msq + + rhobar_n = mpv.HydroState_n.rho0.reshape(1, -1) + Ybar_n = mpv.HydroState_n.Y0.reshape(1, -1) + + ################################################## + # dimensionless Brunt-Väisälä frequency + N = ud.t_ref * np.sqrt(ud.Nsq_ref) + # dimensionless speed of sound + Cs = np.sqrt(th.gamm / Msq) + ud.Cs = Cs + ud.Ns = N + # dimensionless Coriolis strength + if ud.coriolis_strength[2] == 0.0: + ud.coriolis_strength[2] += 1e-15 + F = ud.coriolis_strength[2] + + G = np.sqrt(9.0 / 40.0) + Gamma = G * N / Cs + k = N / Cs + + ud.rf_bot = ud.init_forcing( + k, -Gamma, Cs, F, N, Gamma, A0, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn + ) + ud.rf_bot.get_T_matrix() + + ud.u_wind_speed = 0.0 + + ud.rf_bot.eigenfunction(0, 1) + up, vp, Yp, pi_p = ud.rf_bot.dehatter(th) + + u = ud.u_wind_speed + up + v = ud.v_wind_speed + vp + w = ud.w_wind_speed + Y = Ybar + Yp + + rho = rhobar + + Sol.rho[...] = rho + Sol.rhou[...] = rho * u + Sol.rhov[...] = rho * v + Sol.rhow[...] = rho * w + Sol.rhoY[...] = rho * Y + Sol.rhoX[...] = 0.0 + mpv.p2_cells[...] = pi_p + + ################################################### + # initialise nodal pi + ud.rf_bot.eigenfunction(0, 1, grid="n") + _, _, _, pi_n = ud.rf_bot.dehatter(th, grid="n") + + mpv.p2_nodes[...] = pi_n + + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + + if hasattr(ud, "mixed_run"): + if ud.mixed_run: + ud.coriolis_strength[2] = 2.0 * 7.292 * 1e-5 * ud.t_ref + + if hasattr(ud, "trad_forcing"): + if ud.trad_forcing: + ud.rf_bot.F = 0.0 + ud.rf_bot.get_T_matrix() + + return Sol diff --git a/src/tests/test_targets.yml b/src/pybella/tests/test_targets.yml similarity index 74% rename from src/tests/test_targets.yml rename to src/pybella/tests/test_targets.yml index f314e7db..45f4a41c 100644 --- a/src/tests/test_targets.yml +++ b/src/pybella/tests/test_targets.yml @@ -22,3 +22,11 @@ test_travelling_vortex: rhou: 4422.43359375 rhov: 4408.03857421875 rhow: 0.0 +test_blending_warm_bubble: + p2_nodes: 0.2735753059387207 + rho: 2312.83837890625 + rhoX: -0.32045888900756836 + rhoY: 2313.16015625 + rhou: -4.76837158203125e-07 + rhov: -4.810579299926758 + rhow: 0.0 diff --git a/src/tests/test_travelling_vortex.py b/src/pybella/tests/test_travelling_vortex.py similarity index 90% rename from src/tests/test_travelling_vortex.py rename to src/pybella/tests/test_travelling_vortex.py index b965fada..0487ae18 100644 --- a/src/tests/test_travelling_vortex.py +++ b/src/pybella/tests/test_travelling_vortex.py @@ -1,8 +1,9 @@ import numpy as np -import dycore.utils.options as opts -import dycore.physics.hydrostatics as hydrostatic -import dycore.utils.boundary as bdry -import dycore.physics.low_mach.second_projection as lm_sp +from ..flow_solver.utils import options as opts, boundary as bdry +from ..flow_solver.physics import hydrostatics +from ..flow_solver.physics.low_mach import second_projection as lm_sp + +from ..utils.data_structures import DiagnosticState import logging @@ -63,13 +64,13 @@ def __init__(self): self.coriolis_strength[0] = self.omega * self.t_ref self.coriolis_strength[2] = self.omega * self.t_ref - for i in range(3): - if (self.gravity_strength[i] > np.finfo(np.float).eps) or (i == 1): - self.i_gravity[i] = 1 - self.gravity_direction = 1 + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + self.i_gravity = gravity_mask.astype(int) + if np.any(gravity_mask): + self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - if self.coriolis_strength[i] > np.finfo(np.float).eps: - self.i_coriolis[i] = 1 + coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps + self.i_coriolis = coriolis_mask.astype(int) self.xmin = -0.5 self.xmax = 0.5 @@ -127,23 +128,28 @@ def __init__(self): self.initial_projection = True self.initial_impl_Euler = False - self.tout = [1.0]#np.arange(0.0,10.1,0.1)[1:] + self.tout = [1.0] # np.arange(0.0,10.1,0.1)[1:] self.stepmax = 101 - self.output_base_name = "_travelling_vortex" - - self.output_type = "test" - self.aux = "" - - self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) - self.stratification = self.stratification_function self.rhoe = self.rhoe_function self.output_timesteps = True self.diag = True - self.diag_current_run = "test_travelling_vortex" - self.diag_plot_compare = False + self.diag_updt_targets = False + + self.output_base_name = "_travelling_vortex" + self.output_type = "test" if not self.diag_updt_targets else "target" + self.aux = "" + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) + + self.diag_state = DiagnosticState( + test_name="test_travelling_vortex", + file_name="target_travelling_vortex", + Nx=self.inx - 1, + Ny=self.iny - 1, + steps=[self.stepmax - 1], + ) self.autogen_fn = False @@ -200,7 +206,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): igxn = node.igx igyn = node.igy - hydrostatic.state(mpv, elem, node, th, ud) + hydrostatics.integrated_state(mpv, elem, node, th, ud) coe = np.zeros((25)) coe[0] = 1.0 / 24.0 diff --git a/src/tests/test_lamb_wave.py b/src/pybella/tests/test_unstable_lamb.py similarity index 60% rename from src/tests/test_lamb_wave.py rename to src/pybella/tests/test_unstable_lamb.py index bf0fde7c..34fe949a 100644 --- a/src/tests/test_lamb_wave.py +++ b/src/pybella/tests/test_unstable_lamb.py @@ -1,56 +1,31 @@ -import numpy as np -import dycore.utils.options as opts -import dycore.utils.boundary as bdry -import dycore.physics.hydrostatics as hydrostatic - -class UserData(object): - NSPEC = 1 - grav = 9.81 # [m s^{-2}] - omega = 7.292 * 1e-5 # [s^{-1}] +""" +Unstable Lamb Wave integral test involving vertical Coriolis and Rayleigh BC. +""" - R_gas = 287.4 # [J kg^{-1} K^{-1}] - R_vap = 461.0 - Q_vap = 2.53e+06 - gamma = 1.4 - cp_gas = gamma * R_gas / (gamma-1.0) - - p_ref = 1e+5 - T_ref = 300.00 # [K] - rho_ref = p_ref / (R_gas * T_ref) - N_ref = grav / np.sqrt(cp_gas * T_ref) - Cs = np.sqrt(gamma * R_gas * T_ref) +import numpy as np +from ..flow_solver.physics import hydrostatics +from ..flow_solver.utils import boundary as bdry +from ..flow_solver.utils import options as opts - h_ref = 10.0e3 # [m] - t_ref = 100.0 # [s] - u_ref = h_ref / t_ref +from ..utils.data_structures import DiagnosticState - i_gravity = np.zeros((3)) - i_coriolis = np.zeros((3)) +class UserData(object): def __init__(self): - self.h_ref = self.h_ref - self.t_ref = self.t_ref - self.T_ref = self.T_ref - self.p_ref = self.p_ref - self.rho_ref = self.rho_ref - self.u_ref = self.u_ref - self.Nsq_ref = self.N_ref * self.N_ref - self.g_ref = self.grav - self.gamm = self.gamma - self.Rg_over_Rv = self.R_gas / self.R_vap - self.Q = self.Q_vap / (self.R_gas * self.T_ref) - self.R_gas = self.R_gas - self.Rg = self.R_gas / (self.h_ref**2 / self.t_ref**2 / self.T_ref) - self.cp_gas = self.cp_gas + self.grav = 9.81 # [m/s^2] + self.omega = 7.292 * 1e-5 # [s^{-1}] + self.t_ref = 100.0 # [s] - self.nspec = self.NSPEC + self.h_ref = 10.0e3 # [m] + self.u_ref = self.h_ref / self.t_ref - self.is_nonhydrostatic = 1 - self.is_compressible = 1 - self.is_ArakawaKonor = 0 + self.T_ref = 300.00 # [K] + self.R_gas = 287.4 # [J kg^{-1} K^{-1}] + self.Rg = self.R_gas / (self.h_ref**2 / self.t_ref**2 / self.T_ref) + self.gamma = 1.4 + self.cp_gas = self.gamma * self.R_gas / (self.gamma-1.0) - self.compressibility = 1.0 - self.acoustic_timestep = 0 + self.Nsq = (self.grav / np.sqrt(self.cp_gas * self.T_ref))**2 self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) self.gravity_strength = np.zeros((3)) @@ -59,116 +34,162 @@ def __init__(self): self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) self.coriolis_strength[2] = 2.0 * self.omega * self.t_ref - for i in range(3): - if (self.gravity_strength[i] > np.finfo(np.float).eps) or (i == 1): - self.i_gravity[i] = 1 - self.gravity_direction = 1 + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + self.i_gravity = gravity_mask.astype(int) + if np.any(gravity_mask): + self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - if (self.coriolis_strength[i] > np.finfo(np.float).eps): - self.i_coriolis[i] = 1 + coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps + self.i_coriolis = coriolis_mask.astype(int) - j = 4.0 - Lx = 1.0 * np.pi * self.Cs / self.N_ref * j - self.xmin = - Lx / self.h_ref - self.xmax = Lx / self.h_ref - self.ymin = - 0.0 - self.ymax = 2.0 - self.zmin = - 1.0 - self.zmax = 1.0 - - self.u_wind_speed = 0.0 * self.u_ref - self.v_wind_speed = 0.0 - self.w_wind_speed = 0.0 + ########################################## + # SPATIAL GRID + ########################################## + self.inx = 301 + 1 + self.iny = 30 + 1 + self.inz = 1 + + self._setup_domain() + + ########################################## + # BOUNDARY CONDITIONS + ########################################## self.bdry_type = np.empty((3), dtype=object) self.bdry_type[0] = opts.BdryType.PERIODIC self.bdry_type[1] = opts.BdryType.WALL self.bdry_type[2] = opts.BdryType.WALL - self.LAMB_BDRY = False + self.ATMOSPHERIC_EXTENSION = True ########################################## # NUMERICS ########################################## self.CFL = 0.9 - - self.inx = 151+1 - self.iny = 15+1 - self.inz = 1 - self.dtfixed0 = 100.0 / self.t_ref self.dtfixed = self.dtfixed0 - - self.do_advection = False - self.limiter_type_scalars = opts.LimiterType.NONE - self.limiter_type_velocity = opts.LimiterType.NONE - self.tol = 1.e-30 - self.max_iterations = 10000 + self.tout = [36.0] + self.stepmax = 11 + + self.is_compressible = 1 + self.is_nonhydrostatic = 1 + self.is_ArakawaKonor = 0 + + self.compressibility = 1.0 + self.acoustic_timestep = 0 - # blending parameters - self.perturb_type = 'pos_perturb' - self.blending_mean = 'rhoY' # 1.0, rhoY - self.blending_conv = 'rho' #theta, rho - self.blending_type = 'half' # half, full + ########################################## + # PHYSICS AND BACKGROUND WIND + ########################################## + self.u_wind_speed = 0.0 + self.v_wind_speed = 0.0 + self.w_wind_speed = 0.0 + ########################################## + # BLENDING + ########################################## self.continuous_blending = False self.no_of_pi_initial = 1 self.no_of_pi_transition = 0 self.no_of_hy_initial = 0 self.no_of_hy_transition = 0 - self.blending_weight = 0./16 + self.blending_weight = 0.0 / 16 + self.blending_mean = "rhoY" + self.blending_conv = "rho" + self.blending_type = "half" self.initial_blending = False - self.initial_projection = True + ########################################## + # STRATIFICATION + ########################################## + self.stratification = self.stratification_wrapper + self.rayleigh_bc = self.rayleigh_bc_function + + ########################################## + # DIAGNOSTICS + ########################################## + self.diag = True + self.diag_updt_targets = False + + ########################################## + # OUTPUTS + ########################################## + self.output_base_name = "_unstable_lamb" + self.output_type = "test" if not self.diag_updt_targets else "target" + self.aux = "" + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) - self.tout = [36.0] - # self.tout = np.arange(0,361,1.0) - # self.tout = np.append(self.tout, [720.0]) - self.stepmax = 31 self.output_timesteps = True + # Rayleigh forcing parameters + self.rayleigh_forcing = True + self.rayleigh_forcing_type = "func" + self.rayleigh_forcing_fn = None + self.rayleigh_forcing_path = None + + self.init_forcing = self.forcing + self.rayleigh_bdry_switch = True + + tol = 1.0e-2 + self.diag_state = DiagnosticState( + test_name="test_unstable_lamb", + file_name="target_unstable_lamb", + Nx=self.inx - 1, + Ny=self.iny - 1, + steps=[self.stepmax - 1], + plot_compare=True, + tolerances={ + "rho": tol, + "rhou": tol, + "rhov": tol, + "rhow": tol, + "rhoY": tol, + "rhoX": tol, + "p2_nodes": tol, + } + ) + self.autogen_fn = False - self.output_base_name = "_lamb_wave" - self.output_type = 'test' - self.aux = '' - self.output_suffix = "_%i_%i" %(self.inx-1,self.iny-1) + def _setup_domain(self): + """Setup domain dimensions based on wave parameters""" + # Reference values (from global constants) + R_gas = 287.4 # [J kg^{-1} K^{-1}] + gamma = 1.4 + T_ref = 300.0 # [K] + h_ref = 10.0e3 # [m] - self.diag = True - self.diag_current_run = 'test_lamb_wave' - self.diag_plot_compare = False + # Compute derived quantities + Cs = np.sqrt(gamma * R_gas * T_ref) + cp_gas = gamma * R_gas / (gamma - 1.0) + N_ref = self.grav / np.sqrt(cp_gas * T_ref) - self.stratification = self.stratification_wrapper - self.rayleigh_bc = self.rayleigh_bc_function - self.init_forcing = self.forcing + # Domain size calculation + j = 4.0 + Lx = 1.0 * np.pi * Cs / N_ref * j - self.rayleigh_forcing = True - self.rayleigh_forcing_type = 'func' # func or file - self.rayleigh_forcing_fn = 'output_mark_wave_ensemble=1_601_240_bottom_forcing_S16.h5' - self.rayleigh_forcing_path = './output_mark_wave' - + self.xmin = -Lx / h_ref + self.xmax = Lx / h_ref + self.ymin = 0.0 + self.ymax = 2.0 + self.zmin = -1.0 + self.zmax = 1.0 def stratification_wrapper(self, dy): return lambda y : self.stratification_function(y, dy) def stratification_function(self, y, dy): g = self.gravity_strength[1] - - Gamma = (self.gamm - 1.0) / self.gamm + Gamma = (self.gamma - 1.0) / self.gamma Hex = 1.0 / (Gamma * g) - pi_m = np.exp(-(y - 0.5 * dy) / Hex) + pi_p = np.exp(-(y + 0.5 * dy) / Hex) - - Theta = - (Gamma * g * dy) / (pi_p - pi_m) + pi_m = np.exp(-(y - 0.5 * dy) / Hex) + Theta = - (Gamma * g * dy) / (pi_p - pi_m) return Theta - - # Nsq = self.Nsq_ref * self.t_ref**2 - # g = self.gravity_strength[1] / self.Msq - - # return np.exp(Nsq * y / g) - + @staticmethod def rayleigh_bc_function(ud): if ud.bdry_type[1] == opts.BdryType.RAYLEIGH or ud.rayleigh_forcing == True: @@ -180,7 +201,6 @@ def rayleigh_bc_function(ud): ud.bcy = ud.ymax ud.ymax += 3.0 * ud.bcy - class forcing(object): def __init__(self, k, mu, Cs, F, N, Gamma, ampl, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn): self.k = k @@ -205,7 +225,6 @@ def __init__(self, k, mu, Cs, F, N, Gamma, ampl, g, rhobar, Ybar, rhobar_n, Ybar self.Yn = Yn def get_T_matrix(self): - # system matrix of linearized equations matrix = -np.array([[0, self.F, 0, 1j*self.Cs*self.k], [-self.F, 0, -self.N, self.Cs*(self.mu+self.Gamma)], @@ -256,7 +275,6 @@ def dehatter(self, th, grid='c'): def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): - if hasattr(ud, 'rayleigh_bdry_switch'): if ud.rayleigh_bdry_switch: ud.bdry_type[1] = opts.BdryType.RAYLEIGH @@ -265,56 +283,53 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) if ud.rayleigh_forcing: - # ud.tcy, ud.tny = get_tau_y(ud, elem, node, 0.005) - ud.forcing_tcy, ud.forcing_tny = bdry.get_bottom_tau_y(ud, elem, node, 0.2, cutoff=0.3) + A0 = 1.0e-1 / ud.u_ref Msq = ud.Msq g = ud.gravity_strength[1] * ud.Rg - x = elem.x.reshape(-1,1) - y = elem.y.reshape(1,-1) + x = elem.x.reshape(-1, 1) + y = elem.y.reshape(1, -1) X, Y = np.meshgrid(x, y) - xn = node.x.reshape(-1,1) - yn = node.y.reshape(1,-1) + xn = node.x.reshape(-1, 1) + yn = node.y.reshape(1, -1) dy = np.diff(node.y)[0] Xn, Yn = np.meshgrid(xn, yn) ################################################## - # Following Rupert's fix, reinitialise all background quantities + # Reinitialise all background quantities # as derived from one quantity. ud.stratification = ud.stratification(dy) # Use hydrostatically balanced background - hydrostatic.state(mpv, elem, node, th, ud) - rhobar = mpv.HydroState.rho0.reshape(1,-1) - Ybar = mpv.HydroState.Y0.reshape(1,-1) - pibar = mpv.HydroState.p20.reshape(1,-1) * ud.Msq + hydrostatics.analytical_state(mpv, elem, node, th, ud) + rhobar = mpv.HydroState.rho0.reshape(1, -1) + Ybar = mpv.HydroState.Y0.reshape(1, -1) + pibar = mpv.HydroState.p20.reshape(1, -1) * ud.Msq - rhobar_n = mpv.HydroState_n.rho0.reshape(1,-1) - Ybar_n = mpv.HydroState_n.Y0.reshape(1,-1) + rhobar_n = mpv.HydroState_n.rho0.reshape(1, -1) + Ybar_n = mpv.HydroState_n.Y0.reshape(1, -1) ################################################## # dimensionless Brunt-Väisälä frequency - N = ud.t_ref * np.sqrt(ud.Nsq_ref) + N = ud.t_ref * np.sqrt(ud.Nsq) # dimensionless speed of sound - Cs = np.sqrt(th.gamm / Msq) + Cs = np.sqrt(th.gamm / Msq) ud.Cs = Cs ud.Ns = N # dimensionless Coriolis strength if ud.coriolis_strength[2] == 0.0: ud.coriolis_strength[2] += 1e-15 F = ud.coriolis_strength[2] - # if F == 0.0: - # F += 1e-15 - G = np.sqrt( 9. / 40. ) + G = np.sqrt(9. / 40.) Gamma = G * N / Cs - k = N / Cs + k = N / Cs ud.rf_bot = ud.init_forcing(k, -Gamma, Cs, F, N, Gamma, A0, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn) ud.rf_bot.get_T_matrix() @@ -347,10 +362,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): mpv.p2_nodes[...] = pi_n - # if ud.bdry_type[1] == 'RAYLEIGH': - # rayleigh_damping(Sol, mpv, ud, ud.tcy, elem, th) - - bdry.set_explicit_boundary_data(Sol,elem,ud,th,mpv) + bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) if hasattr(ud, 'mixed_run'): if ud.mixed_run: @@ -361,4 +373,4 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ud.rf_bot.F = 0.0 ud.rf_bot.get_T_matrix() - return Sol \ No newline at end of file + return Sol diff --git a/src/pybella/utils/__init__.py b/src/pybella/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pybella/utils/data_structures.py b/src/pybella/utils/data_structures.py new file mode 100644 index 00000000..15b2aef5 --- /dev/null +++ b/src/pybella/utils/data_structures.py @@ -0,0 +1,147 @@ +from typing import Optional, Callable, List, Any +from dataclasses import dataclass, field, fields + +from ..flow_solver.discretisation.grid import Grid +from ..flow_solver.utils.variable import Vars +from ..flow_solver.physics.low_mach.mpv import MPV +from ..flow_solver.physics.gas_dynamics.thermodynamics import ThermodynamicalQuantities + + +@dataclass +class IntegrationTime: + step: int = 0 + t: float = 0.0 + window_step: int = 0 + + +@dataclass +class ModelState: + elem: Grid + node: Grid + sol: Vars + flux: List[Vars] + mpv: MPV + th: ThermodynamicalQuantities + time: IntegrationTime = field(init=False) + + def __post_init__(self): + self.time = IntegrationTime() + + def __iter__(self): + return iter(getattr(self, field.name) for field in fields(self)) + + +@dataclass +class InterfaceParameters: + bld: Optional[Any] = None + + +@dataclass +class DataAssimilationParameters: + # DA user-input parameters + dap: Any + # r-localisation function + rloc: Any + # solution ensemble + # sol_ens : Any + + # observation related attributes + obs: Any + obs_noisy: Any + obs_mask: Any + obs_covar: Any + + +@dataclass +class RestartParameters: + ud_rewrite: Optional[object] = None + dap_rewrite: Optional[object] = None + r_params: Optional[object] = None + + +@dataclass +class EnsembleState: + members: List[ModelState] = field(default_factory=list) + + def update_member( + self, + elem: Grid, + node: Grid, + sol: Vars, + mpv: MPV, + flux: List[Vars], + th: ThermodynamicalQuantities, + ): + new_state = ModelState(elem, node, sol, flux, mpv, th) + self.members.append(new_state) + + def set_members(self, members: List[ModelState]): + assert len(self.set_members == members) + self.members = members + + def get_member(self, index: int) -> ModelState: + return self.members[index] + + def get_all_members(self) -> List[ModelState]: + return self.members + + def get_grid(self) -> tuple[Grid, Grid]: + # Assuming identical underlying grid for all ensemble memebers + elem = self.members[0].elem + node = self.members[0].node + return elem, node + + def __getitem__(self, index): + return self.members[index] + + +@dataclass +class DiagnosticState: + """ + Initialise diagnostic state for tests + + Consider removing run-specific parameters, e.g., (Nx, Ny), in future. + """ + + # the name to look up in test_targets.yml + test_name: str + # filename of the reference (if updt_target = True or plot_compare = True) + file_name: str + # details related to loading the reference fields + Nx: int + Ny: int + steps: list + path: str = "./outputs/" + + # plot the comparison? + plot_compare: bool = True + + tolerances: dict[str, float] = field(default_factory=lambda: { + # for most tests, we want to ensure that the + # max absolute error is within singular precision + # limit. + "rhou": 1e-5, + "rhov": 1e-5, + "rhow": 1e-5, + "rhoY": 1e-5, + "rhoX": 1e-5, + "rho": 1e-5, + "p2_nodes": 1e-5 + }) + time_increment: bool = False + + +@dataclass +class SimulationState: + N: int + restart: bool + + ud: object + sol_init: Callable + + ensemble_state: EnsembleState + restart_params: RestartParameters + interface_params: Optional[InterfaceParameters] = None + da_params: Optional[DataAssimilationParameters] = None + + diag_comparison: Optional[object] = None diff --git a/src/pybella/utils/debug_helpers.py b/src/pybella/utils/debug_helpers.py new file mode 100644 index 00000000..d7fc86d6 --- /dev/null +++ b/src/pybella/utils/debug_helpers.py @@ -0,0 +1,7 @@ +import matplotlib.pyplot as plt + +def pl_sol(arr): + plt.figure() + plt.imshow(arr, origin="lower", aspect="auto") + plt.colorbar() + plt.show() \ No newline at end of file diff --git a/src/utils/io.py b/src/pybella/utils/io.py similarity index 81% rename from src/utils/io.py rename to src/pybella/utils/io.py index 0fb2ffe9..ec8038de 100644 --- a/src/utils/io.py +++ b/src/pybella/utils/io.py @@ -12,10 +12,45 @@ import argparse -import utils.sim_params as params +from . import sim_params as params +from ..interfaces.ic_config import IC_MODULES + + +def initialise(sst): + es = sst.ensemble_state + dp = sst.da_params + ###################################################### + # Initialise writer class for I/O operations + ###################################################### + writer = hdf5(sst.ud, sst.restart) + writer.check_jar() + writer.jar([sst.ud, es[0].elem, es[0].node, dp.dap]) + # sys.exit("Let's just dill the stuff and quit!") + + writer.write_attrs() + wrtr = None + if sst.N > 1: + writer.write_da_attrs(dp.dap) + elif params.output_timesteps == True: + wrtr = writer + + for n in range(sst.N): # write initial ensemble + if params.label_type == "STEP": + label = "ensemble_mem=%i_%.3d" % (n, sst.step) + else: + label = "ensemble_mem=%i_%.3f" % (n, 0.0) + if not sst.restart: + writer.write_all(es[n], str(label) + "_ic") + + if params.da_debug: + # writer.jar([obs,obs_noisy,obs_noisy_interp,obs_mask,obs_covar]) + # obs = obs_noisy_interp + writer.jar([dp.obs, dp.obs_noisy, dp.obs_mask, dp.obs_covar]) + + return writer, wrtr -class init(object): +class hdf5(object): """ HDF5 writer class. Contains methods to create HDF5 file, create data sets and populate them with output variables. @@ -125,7 +160,7 @@ def io_create_file(self, paths, restart): file.close() - def write_all(self, Sol, mpv, elem, node, th, name): + def write_all(self, model_state, name): """ At a given time, write output from `Sol` and `mpv` to the HDF5 file. @@ -145,6 +180,10 @@ def write_all(self, Sol, mpv, elem, node, th, name): The time and additional suffix label for the dataset, e.g. "_10.0_after_full_step", where 10.0 is the time and "after_full_step" denotes when the output was made. """ + + Sol = model_state.sol + mpv = model_state.mpv + logging.info("writing hdf output..." + name) # rho self.populate(name, "rho", Sol.rho) @@ -157,22 +196,6 @@ def write_all(self, Sol, mpv, elem, node, th, name): self.populate(name, "rhow", Sol.rhow) self.populate(name, "rhoX", Sol.rhoX) - # if hasattr(Sol,'rhov_half'): - # self.populate(name,'rhov_half',Sol.rhov_half) - # if hasattr(Sol,'rhou_half'): - # self.populate(name,'rhou_half',Sol.rhou_half) - # if hasattr(Sol,'rhow_half'): - # self.populate(name,'rhow_half',Sol.rhow_half) - # if hasattr(Sol,'rhoX_half'): - # self.populate(name,'rhoX_half',Sol.rhoX_half) - # if hasattr(Sol,'rhoY_half'): - # self.populate(name,'rhoY_half',Sol.rhoY_half) - # if hasattr(Sol,'rho_half'): - # self.populate(name,'rho_half',Sol.rho_half) - # if hasattr(mpv,'p2_nodes_half'): - # self.populate(name,'p2_nodes_half',mpv.p2_nodes_half) - # self.populate(name,'buoy',Sol.rhoX / Sol.rho) - self.populate(name, "p2_nodes", mpv.p2_nodes) # vorticity in (x,z) @@ -449,7 +472,6 @@ def jar(self, content=None): class read_input(object): - def __init__(self, fn, path): self.fn = fn self.path = path @@ -499,30 +521,7 @@ def get_args(): dest="ic", help=" Set initial conditions", required=True, - choices={ - "aw", - "tv", - "tv_neg", - "tv_2d", - "tv_3d", - "tv_corr", - "rb", - "rbc", - "igw", - "igw_3d", - "lbw", - "skl", - "mark", - "lw_p", - "igw_bb", - "swe", - "swe_bal_vortex", - "swe_icshear", - "swe_dvortex", - "test_travelling_vortex", - "test_internal_long_wave", - "test_lamb_wave", - }, + choices=set(IC_MODULES.keys()), # Use the keys from IC_MODULES ) subparsers = parser.add_subparsers(dest="subcommand") @@ -564,53 +563,21 @@ def get_args(): args = parser.parse_args() # collect cmd line args ic = args.ic - if ic == "bi": - from inputs.baroclinic_instability_periodic import UserData, sol_init - elif ic == "tv" or ic == "tv_2d": - from inputs.travelling_vortex_2D import UserData, sol_init - elif ic == "tv_neg": - from inputs.travelling_vortex_2D_neg import UserData, sol_init - elif ic == "tv_3d": - from inputs.travelling_vortex_3D import UserData, sol_init - elif ic == "tv_corr": - from inputs.travelling_vortex_3D_Coriolis import UserData, sol_init - elif ic == "aw": - from inputs.acoustic_wave_high import UserData, sol_init - elif ic == "igw": - from inputs.internal_long_wave import UserData, sol_init - elif ic == "igw_3d": - from inputs.internal_long_wave_3D import UserData, sol_init - elif ic == "lbw": - from inputs.lamb_waves import UserData, sol_init - elif ic == "skl": - from inputs.sk_lamb_wave import UserData, sol_init - elif ic == "mark": - from inputs.mark import UserData, sol_init - elif ic == "lw_p": - from inputs.lamb_wave_perturb import UserData, sol_init - elif ic == "igw_bb": - from inputs.igw_baldauf_brdar import UserData, sol_init - elif ic == "rb": - from inputs.rising_bubble import UserData, sol_init - elif ic == "rbc": - from inputs.rising_bubble_cold import UserData, sol_init - elif ic == "swe_bal_vortex": - from inputs.swe_bal_vortex import UserData, sol_init - elif ic == "swe": - from inputs.shallow_water_3D import UserData, sol_init - elif ic == "swe_icshear": - from inputs.shallow_water_3D_icshear import UserData, sol_init - elif ic == "swe_dvortex": - from inputs.shallow_water_3D_dvortex import UserData, sol_init - elif ic == "test_travelling_vortex": - from tests.test_travelling_vortex import UserData, sol_init - elif ic == "test_internal_long_wave": - from tests.test_internal_long_wave import UserData, sol_init - elif ic == "test_lamb_wave": - from tests.test_lamb_wave import UserData, sol_init + # Import the appropriate module + if ic in IC_MODULES: + module_name = IC_MODULES[ic] + try: + module = __import__(module_name, fromlist=['UserData', 'sol_init']) + UserData = getattr(module, 'UserData') + sol_init = getattr(module, 'sol_init') + except ImportError as e: + raise ImportError(f"Failed to import {module_name}: {e}") + else: + raise ValueError(f"Unknown initial condition: {ic}") if UserData is None or sol_init is None: assert 0, "Initial condition file is not well defined." + if args.N is None: N = 1 else: @@ -685,7 +652,6 @@ def sim_restart(path, name, elem, node, ud, Sol, mpv, restart_touts): def fn_gen(ud, dap, N): - suffix = "" suffix += "_%i" % (ud.inx - 1) suffix += "_%i" % (ud.iny - 1) @@ -723,22 +689,24 @@ def mkdir_p(path): try: os.makedirs(path, exist_ok=True) # Python>3.2 except OSError as exc: - if exc.errno == errno.EEXIST and os.path.isdir(path): - pass - else: raise + if exc.errno == errno.EEXIST and os.path.isdir(path): + pass + else: + raise ########################################################## # Initialise logger ########################################################## - + + def init_logger(ud): now = datetime.now() date = now.strftime("%d%m%y") time = now.strftime("%H%M%S") - input_filename = "%s%s" %(ud.output_type, ud.output_base_name) - logger_filename = "./logs/%s_%s_%s.log" %(input_filename, date, time) + input_filename = "%s%s" % (ud.output_type, ud.output_base_name) + logger_filename = "./logs/%s_%s_%s.log" % (input_filename, date, time) mkdir_p(os.path.dirname(logger_filename)) @@ -760,4 +728,70 @@ def init_logger(ud): # add the handler to the root logger logging.getLogger().addHandler(console) - logging.info("Input file is %s" %input_filename) \ No newline at end of file + # Suppress library specific debug outputs + logging.getLogger("matplotlib").setLevel(logging.WARNING) + logging.getLogger("matplotlib.font_manager").setLevel(logging.WARNING) + logging.getLogger("numba").setLevel(logging.WARNING) + logging.getLogger("numba.core").setLevel(logging.WARNING) + + logging.getLogger().setLevel(logging.WARNING) + + logging.info("Input file is %s" % input_filename) + +class NullDebugWriter: + """Null object that does nothing but implements the DebugWriter interface.""" + + def populate(self, label, field_name, data): + """No-op populate method.""" + pass + + def write(self, label): + """No-op write method.""" + pass + + def populate_flux_components(self, label, flux, elem): + """No-op populate flux components method.""" + pass + +class DebugWriter: + """Enhanced debug writer with populate functionality.""" + + def __init__(self, debug, writer, mem): + self.debug = debug + self.writer = writer + self.mem = mem + self._pending_populations = {} + + def populate(self, label, field_name, data): + """Populate field data for a given label.""" + if self.debug and self.writer is not None: + if label not in self._pending_populations: + self._pending_populations[label] = [] + self._pending_populations[label].append((field_name, data)) + + def write(self, label): + """Write all data including any pending populations.""" + if self.debug and self.writer is not None: + # Apply any pending populations for this label + if label in self._pending_populations: + for field_name, data in self._pending_populations[label]: + self.writer.populate(label, field_name, data) + del self._pending_populations[label] + + # Write the data + self.writer.write_all(self.mem, label) + + def populate_flux_components(self, label, flux, elem): + """Helper method to populate flux components.""" + if self.debug and self.writer is not None: + self.populate(label, "rhoYu", flux[0].rhoY) + self.populate(label, "rhoYv", flux[1].rhoY) + if elem.ndim == 3: + self.populate(label, "rhoYw", flux[2].rhoY) + +def create_debug_writer(debug, writer, mem): + """Factory function to create appropriate debug writer.""" + if debug and writer is not None: + return DebugWriter(debug, writer, mem) + else: + return NullDebugWriter() \ No newline at end of file diff --git a/src/pybella/utils/prepare.py b/src/pybella/utils/prepare.py new file mode 100644 index 00000000..5d776d59 --- /dev/null +++ b/src/pybella/utils/prepare.py @@ -0,0 +1,140 @@ +import numpy as np + +from . import user_data, io, data_structures + +from ..flow_solver.discretisation import grid as dis_grid +from ..flow_solver.utils import variable as var +from ..flow_solver.utils import boundary as bdry +from ..flow_solver.physics import hydrostatics +from ..flow_solver.physics.low_mach import mpv as lm_var +from ..flow_solver.physics.gas_dynamics import thermodynamics as gd_thermodynamics + +# test module +from ..tests import diagnostics as diag + + +def initialise(): + #### + # Initialise simulation state + #### + from . import sim_params as params + + np.set_printoptions(precision=params.print_precision) + + ########################################################## + # Initialisation of data containers and helper classes + ########################################################## + # get arguments for initial condition and ensemble size + N, UserData, sol_init, restart, ud_rewrite, dap_rewrite, r_params = io.get_args() + if N == 1: + params.da_debug = False + + initial_data = vars(UserData()) + ud = user_data.UserDataInit(**initial_data) + if ud_rewrite is not None: + ud.update_ud(ud_rewrite) + if hasattr(ud, "rayleigh_bc"): + ud.rayleigh_bc(ud) + if ud.output_timesteps: + params.output_timesteps = True + ud.coriolis_strength = np.array(ud.coriolis_strength) + + elem, node = dis_grid.grid_init(ud) + + sol = var.Vars(elem.sc, ud) + + flux = np.empty((3), dtype=object) + flux[0] = var.States(elem.sfx, ud) + if elem.ndim > 1: + flux[1] = var.States(elem.sfy, ud) + if elem.ndim > 2: + flux[2] = var.States(elem.sfz, ud) + + th = gd_thermodynamics.ThermodynamicalQuantities(ud) + mpv = lm_var.MPV(elem, node, ud) + + io.init_logger(ud) + + # handle radiative BC + if ud.bdry_type[1].value == "radiation": + ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) + + ########################################################## + # Initialise test module + ########################################################## + if ud.diag: + diag_comparison = diag.CompareSol(ud.diag_state) + else: + diag_comparison = None + + ########################################################## + # Populate data structures + ########################################################## + + # member_state = data_structures.MemberState( + # elem=elem, + # node=node, + # Sol=Sol, + # flux=flux, + # mpv=mpv, + # th=th, + # ) + + ensemble_state = data_structures.EnsembleState() + + sol = sol_init(sol, mpv, elem, node, th, ud) + + ensemble_state.update_member( + elem=elem, + node=node, + sol=sol, + flux=flux, + mpv=mpv, + th=th, + ) + + restart_params = data_structures.RestartParameters( + ud_rewrite=ud_rewrite, + dap_rewrite=dap_rewrite, + r_params=r_params, + ) + + interface_params = data_structures.InterfaceParameters() + + sim_st = data_structures.SimulationState( + N=N, + restart=restart, + ud=ud, + sol_init=sol_init, + ensemble_state=ensemble_state, + diag_comparison=diag_comparison, + restart_params=restart_params, + interface_params=interface_params, + ) + + return sim_st + + +def overwrite_init_with_restart(sst): + es = sst.ensemble_state + rp = sst.restart_params + + hydrostatics.state(es.mpv, es.elem, es.node, es.th, es.ud) + sst.ud.old_suffix = np.copy(sst.ud.output_suffix) + sst.ud.old_suffix = "_ensemble=%i%s" % (sst.N, sst.ud.old_suffix) + Sol0, mpv0, touts = io.sim_restart( + rp.r_params[0], + rp.r_params[1], + es.elem, + es.node, + es.ud, + es.Sol, + es.mpv, + rp.r_params[2], + ) + sol_ens = [[Sol0, es.flux, mpv0, [-np.inf, sst.step]]] + # ud.tout = touts[1:] + sst.ud.tout = [touts[-1]] + sst.t = touts[0] + + sst.ensemble_state = sol_ens diff --git a/src/pybella/utils/sim_params.py b/src/pybella/utils/sim_params.py new file mode 100644 index 00000000..8bef8077 --- /dev/null +++ b/src/pybella/utils/sim_params.py @@ -0,0 +1,33 @@ +# simulation parameters +debug = False +da_debug = False + +output_timesteps = False +if debug == True: + output_timesteps = True +label_type = "TIME" + +random_seed = 888 +print_precision = 18 +output_path = "./outputs" + + +# global constants +class global_constants(object): + def __init__(self): + self.nspec = 1 + self.buoy = 0 + + self.grav = 9.81 # [m s^{-2}] + self.omega = 0.0 # [s^{-1}] + + self.R_gas = 287.4 # [J kg^{-1} K^{-1}] + self.R_vap = 461.0 + self.Q_vap = 2.53e06 + self.gamm = 1.4 + + self.p_ref = 8.61 * 1e4 # [N/m^2] + self.T_ref = 300.00 # [K] + + self.h_ref = 10000.0 # [m] + self.t_ref = 100.0 # [s] diff --git a/src/utils/user_data.py b/src/pybella/utils/user_data.py similarity index 57% rename from src/utils/user_data.py rename to src/pybella/utils/user_data.py index 0455214c..4aea5e28 100644 --- a/src/utils/user_data.py +++ b/src/pybella/utils/user_data.py @@ -1,7 +1,7 @@ import numpy as np -import dycore.utils.options as opts -import utils.sim_params as params +from ..flow_solver.utils import options as opts +from . import sim_params as params class UserDataInit(object): @@ -10,134 +10,123 @@ class UserDataInit(object): Attributes ---------- - \*\*kwargs: class object + **kwargs: class object """ - def __init__(self,**kwargs): + def __init__(self, **kwargs): gconsts = params.global_constants() for key, value in vars(gconsts).items(): setattr(self, key, value) + # else: + ########################################## + # SPATIAL GRID + ########################################## + self.inx = 64 + 1 + self.iny = 64 + 1 + self.inz = 1 + + self.xmin = -1.0 + self.xmax = 1.0 + self.ymin = 0.0 + self.ymax = 1.0 + self.zmin = -1.0 + self.zmax = 1.0 + + ########################################## + # BOUNDARY CONDITIONS + ########################################## + self.bdry_type = np.empty((3), dtype=object) + self.bdry_type[0] = opts.BdryType.PERIODIC + self.bdry_type[1] = opts.BdryType.WALL + self.bdry_type[2] = opts.BdryType.WALL + + ########################################## + # TEMPORAL + ########################################## + self.CFL = 0.5 + self.dtfixed0 = 100.0 + self.dtfixed = 100.0 + + self.acoustic_timestep = 0 + + self.tout = np.arange(0.0, 1.01, 0.01)[10:] + self.stepmax = 10000 + + ########################################## + # MODEL REGIMES + ########################################## + self.is_ArakawaKonor = 0 + self.is_nonhydrostatic = 1 + self.is_compressible = 1 + + self.compressibility = 1.0 + + ########################################## + # PHYSICS AND BACKGROUND WIND + ########################################## + self.u_wind_speed = 0.0 + self.v_wind_speed = 0.0 + self.w_wind_speed = 0.0 + + self.stratification = self.stratification_function + + ########################################## + # NUMERICS + ########################################## + # Do we solve the left-hand side? + self.do_advection = True + + # Advection limiter types + self.limiter_type_scalars = opts.LimiterType.NONE + self.limiter_type_velocity = opts.LimiterType.NONE + + # Iterative solver + self.tol = 1.0e-8 + self.max_iterations = 6000 + + ########################################## + # BLENDING + ########################################## + self.blending_weight = 0.0 / 16 + self.blending_mean = "rhoY" # 1.0, rhoY + self.blending_conv = "rho" # theta, rho + self.blending_type = "half" + + self.continuous_blending = False + self.no_of_pi_initial = 1 + self.no_of_pi_transition = 0 + self.no_of_hy_initial = 0 + self.no_of_hy_transition = 0 + + self.initial_blending = False + + ########################################## + # DIAGNOSTICS + ########################################## + self.diag = False + self.diag_state = None + + ########################################## + # OUTPUTS + ########################################## + self.autogen_fn = False + self.output_timesteps = False + self.output_type = "output" + self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) + if len(kwargs) > 0: for key, value in kwargs.items(): setattr(self, key, value) - else: - ########################################## - # SPATIAL GRID - ########################################## - self.inx = 64+1 - self.iny = 64+1 - self.inz = 1 - - self.xmin = - 1.0 - self.xmax = 1.0 - self.ymin = 0.0 - self.ymax = 1.0 - self.zmin = - 1.0 - self.zmax = 1.0 - - - ########################################## - # BOUNDARY CONDITIONS - ########################################## - self.bdry_type = np.empty((3), dtype=object) - self.bdry_type[0] = opts.BdryType.PERIODIC - self.bdry_type[1] = opts.BdryType.WALL - self.bdry_type[2] = opts.BdryType.WALL - - - ########################################## - # TEMPORAL - ########################################## - self.CFL = 0.5 - self.dtfixed0 = 100.0 - self.dtfixed = 100.0 - - self.acoustic_timestep = 0 - - self.tout = np.arange(0.0,1.01,0.01)[10:] - self.stepmax = 10000 - - - ########################################## - # MODEL REGIMES - ########################################## - self.is_ArakawaKonor = 0 - self.is_nonhydrostatic = 1 - self.is_compressible = 1 - - self.compressibility = 0.0 - - - ########################################## - # PHYSICS AND BACKGROUND WIND - ########################################## - self.u_wind_speed = 0.0 - self.v_wind_speed = 0.0 - self.w_wind_speed = 0.0 - - self.stratification = self.stratification_function - - - ########################################## - # NUMERICS - ########################################## - # Do we solve the left-hand side? - self.do_advection = True - - # Advection limiter types - self.limiter_type_scalars = opts.LimiterType.NONE - self.limiter_type_velocity = opts.LimiterType.NONE - - # Iterative solver - self.tol = 1.e-8 - self.max_iterations = 6000 - - - ########################################## - # BLENDING - ########################################## - self.blending_weight = 0./16 - self.blending_mean = 'rhoY' # 1.0, rhoY - self.blending_conv = 'rho' # theta, rho - self.blending_type = 'half' - - self.continuous_blending = False - self.no_of_pi_initial = 1 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 0 - self.no_of_hy_transition = 0 - - self.initial_blending = False - - - ########################################## - # DIAGNOSTICS - ########################################## - self.diag = False - self.diag_plot_compare = False - - - ########################################## - # OUTPUTS - ########################################## - self.autogen_fn = False - self.output_timesteps = False - self.output_type = 'output' - self.output_suffix = "_%i_%i" %(self.inx-1,self.iny-1) - - def compute_u_ref(self): self.u_ref = self.h_ref / self.t_ref self.compute_Msq() - def compute_Msq(self): self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) - def compute_gravity(self): self.i_gravity = np.zeros((3)) self.gravity_strength = np.zeros((3)) @@ -149,7 +138,6 @@ def compute_gravity(self): self.i_gravity[i] = 1 self.gravity_direction = i - def compute_coriolis(self): self.i_coriolis = np.zeros((3)) self.coriolis_strength = np.zeros((3)) @@ -157,45 +145,39 @@ def compute_coriolis(self): self.coriolis_strength[0] = self.omega * self.t_ref self.coriolis_strength[2] = self.omega * self.t_ref - def compute_cp_gas(self): - self.cp_gas = self.gamm * self.R_gas / (self.gamm-1.0) + self.cp_gas = self.gamm * self.R_gas / (self.gamm - 1.0) if all(hasattr(self, attr) for attr in ["grav", "cp_gas", "T_ref"]): self.compute_N_ref() - def compute_rho_ref(self): self.rho_ref = self.p_ref / (self.R_gas * self.T_ref) - def compute_N_ref(self): self.N_ref = self.grav / np.sqrt(self.cp_gas * self.T_ref) self.Nsq_ref = self.N_ref * self.N_ref - def compute_Cs(self): self.Cs = np.sqrt(self.gamm * self.R_gas * self.T_ref) - @staticmethod def stratification_function(y): return 1.0 - def update_ud(self, obj): for key, value in obj.items(): setattr(self, key, value) - ########################################## - # SETTER FUNCTIONS - ########################################## + # ########################################## + # # SETTER FUNCTIONS + # ########################################## # gravity and Msq arguments @property def R_gas(self): return self._R_gas - + @R_gas.setter def R_gas(self, val): self._R_gas = val @@ -218,7 +200,7 @@ def R_gas(self, val): @property def T_ref(self): return self._T_ref - + @T_ref.setter def T_ref(self, val): self._T_ref = val @@ -226,15 +208,15 @@ def T_ref(self, val): if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): self.compute_gravity() - if all(hasattr(self, attr) for attr in ["u_ref, R_gas", "T_ref"]): + if all(hasattr(self, attr) for attr in ["u_ref", "R_gas", "T_ref"]): self.compute_Msq() - if all(hasattr(self, attr) for attr in ["p_ref, R_gas", "T_ref"]): + if all(hasattr(self, attr) for attr in ["p_ref", "R_gas", "T_ref"]): self.compute_rho_ref() if all(hasattr(self, attr) for attr in ["grav", "cp_gas", "T_ref"]): self.compute_N_ref() - + if all(hasattr(self, attr) for attr in ["gamm", "R_gas", "T_ref"]): self.compute_Cs() @@ -242,7 +224,7 @@ def T_ref(self, val): @property def grav(self): return self._grav - + @grav.setter def grav(self, val): self._grav = val @@ -256,14 +238,14 @@ def grav(self, val): @property def h_ref(self): return self._h_ref - + @h_ref.setter def h_ref(self, val): self._h_ref = val if all(hasattr(self, attr) for attr in ["h_ref", "t_ref"]): self.compute_u_ref() - + if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): self.compute_gravity() @@ -271,7 +253,7 @@ def h_ref(self, val): @property def t_ref(self): return self._t_ref - + @t_ref.setter def t_ref(self, val): self._t_ref = val @@ -285,7 +267,7 @@ def t_ref(self, val): @property def omega(self): return self._omega - + @omega.setter def omega(self, val): self._omega = val @@ -293,19 +275,18 @@ def omega(self, val): if all(hasattr(self, attr) for attr in ["omega", "t_ref"]): self.compute_coriolis() - # Cs and cp_gas argument @property def gamm(self): return self._gamm - + @gamm.setter def gamm(self, val): self._gamm = val if all(hasattr(self, attr) for attr in ["gamm", "R_gas"]): self.compute_cp_gas() - + if all(hasattr(self, attr) for attr in ["gamm", "R_gas", "T_ref"]): self.compute_Cs() @@ -313,10 +294,10 @@ def gamm(self, val): @property def p_ref(self): return self._p_ref - + @p_ref.setter def p_ref(self, val): self._p_ref = val if all(hasattr(self, attr) for attr in ["p_ref, R_gas", "T_ref"]): - self.compute_rho_ref() \ No newline at end of file + self.compute_rho_ref() diff --git a/src/pybella/vis/__init__.py b/src/pybella/vis/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pybella/vis/plotting_tools.py b/src/pybella/vis/plotting_tools.py new file mode 100644 index 00000000..17bb771d --- /dev/null +++ b/src/pybella/vis/plotting_tools.py @@ -0,0 +1,397 @@ +import matplotlib.pyplot as plt +import mpl_toolkits.axes_grid1 as axes_grid +import matplotlib.animation as animation +import numpy as np +import itertools + + +class plotter(object): + def __init__( + self, + arr_lst, + ncols=4, + figsize=(12, 8), + fontsize=14, + sharexlabel=False, + shareylabel=False, + sharex=False, + sharey=False, + ): + plt.rcParams.update({"font.size": fontsize}) + self.arr_lst = np.array(arr_lst, dtype="object") + N = self.arr_lst.shape[0] + + if N > ncols: + self.nrows = int(np.ceil(N / ncols)) + self.ncols = ncols + if N <= ncols: + self.nrows = 1 + self.ncols = N + self.N = N + + ridx = np.arange(self.nrows) + cidx = np.arange(self.ncols) + + self.idx = [] + if self.nrows > 1: + for pair in itertools.product(ridx, cidx): + self.idx.append(pair) + else: + self.idx = cidx + + self.visualise = self.visualise + self.fig, self.ax = plt.subplots( + ncols=self.ncols, + nrows=self.nrows, + figsize=figsize, + sharex=sharex, + sharey=sharey, + ) + + self.sharexlabel = sharexlabel + self.shareylabel = shareylabel + + self.img = plt + + def set_axes(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + def set_cax_axes(self, cax, n): + if hasattr(self, "x_locs"): + cax.set_xticks(self.x_locs) + if hasattr(self, "x_axs"): + cax.set_xticklabels(self.x_axs) + if hasattr(self, "y_locs"): + cax.set_yticks(self.y_locs) + if hasattr(self, "y_axs"): + cax.set_yticklabels(self.y_axs) + cax.tick_params(axis="x", pad=15) + if self.sharexlabel: + if int(n // self.ncols) == self.nrows - 1: + if hasattr(self, "x_label"): + cax.set_xlabel(self.x_label) + else: + if hasattr(self, "x_label"): + cax.set_xlabel(self.x_label) + if self.shareylabel: + if n % self.ncols == 0: + if hasattr(self, "y_label"): + cax.set_ylabel(self.y_label) + else: + if hasattr(self, "y_label"): + cax.set_ylabel(self.y_label) + if hasattr(self, "axhline"): + cax.axhline(self.axhline, c="k", lw=0.5) + if hasattr(self, "axvline"): + cax.axvline(self.axvline, c="k", lw=0.5) + if hasattr(self, "marker"): + for marker in self.marker: + cax.plot(marker[0], marker[1], marker="x", c=marker[2], ms=18, mew=4) + if hasattr(self, "rects"): + for rect in self.rects: + cax.add_patch(rect) + + def plot( + self, + method="imshow", + inner=False, + suptitle="", + rect=[0, 0.03, 1, 0.95], + aspect="auto", + lvls=None, + cmaps=None, + ): + if method != "imshow" and method != "contour": + assert 0, "Visualisation method not implemented!" + + if self.N > 1: + ims, caxs, baxs = [], [], [] + for n, arr in enumerate(self.arr_lst): + arr, title = arr[0], arr[1] + if inner == True: + arr = arr[2:-2, 2:-2] + cax = self.ax[self.idx[n]] + + lvl = lvls[n] if lvls is not None else None + cmap = cmaps[n] if cmaps is not None else "viridis" + + im = self.visualise(method, cax, arr, aspect, lvl, cmap) + if type(title) == str: + cax.set_title(title) + elif type(title) == np.ndarray or type(title) == list: + cax.set_title(title[0], fontsize=title[1], fontweight=title[2]) + loc = cax.get_xticklabels() + self.set_cax_axes(cax, n) + caxs.append(cax) + divider = axes_grid.make_axes_locatable(cax) + bax = divider.append_axes("right", size="5%", pad=0.05) + if method == "imshow" and lvl is not None: + # plt.colorbar(im, cax=bax, ticks=lvl)#, format='%.3f') + plt.colorbar( + im, cax=bax, ticks=lvl, extend="both" + ) # , format='%.3f') + else: + # plt.colorbar(im, cax=bax, extend='both')#, format='%.3f') + plt.colorbar(im, cax=bax, extendrect=True) + # plt.colorbar(im, cax=bax, ticks=lvl) + baxs.append(bax) + if hasattr(self, "cbar_label"): + bax.set_xlabel(self.cbar_label) + bax.xaxis.set_label_position("top") + if hasattr(self, "cbar_label_coords"): + bax.xaxis.set_label_coords( + self.cbar_label_coords[0], self.cbar_label_coords[1] + ) + ims.append(im) + + for i in range(n + 1, self.nrows * self.ncols): + self.fig.delaxes(self.ax[self.idx[i]]) + + else: + arr, title = self.arr_lst[0][0], self.arr_lst[0][1] + if inner == True: + arr = arr[2:-2, 2:-2] + # cax = self.fig.gca() + cmap = cmaps if cmaps is not None else "viridis" + cax = self.ax + im = self.visualise(method, cax, arr, aspect, lvls, cmap) + cax.set_title(title) + self.set_cax_axes(cax, 0) + caxs = [cax] + divider = axes_grid.make_axes_locatable(cax) + bax = divider.append_axes("right", size="5%", pad=0.05) + if hasattr(self, "cbar_label"): + bax.set_xlabel(self.cbar_label) + bax.xaxis.set_label_position("top") + if method == "imshow" and lvls is not None: + # plt.colorbar(im, cax=bax, ticks=lvls)#, format='%.3f') + plt.colorbar(im, cax=bax, ticks=lvls) + else: + plt.colorbar(im, cax=bax, extendrect=True) + # plt.colorbar(im, cax=bax) + if aspect != "auto" and aspect != "equal": + bax.set_aspect(float(aspect) * 10.0) + ims = [im] + baxs = [bax] + + plt.suptitle(suptitle) + plt.tight_layout(rect=rect) + # plt.subplots_adjust(hspace = .000005) + + if self.N > 1: + return ims, caxs, baxs + else: + return ims, caxs, baxs + + def save_fig(self, fn, format=".pdf"): + self.fig.tight_layout() + self.fig.savefig(fn + format, bbox_inches="tight", pad_inches=0.1) + + @staticmethod + def visualise(method, cax, arr, aspect, lvls, cmap): + if cmap is None: + norm = None + cmap = "viridis" + elif len(cmap) == 2: + norm = cmap[1] + cmap = cmap[0] + else: + norm = None + cmap = cmap + + if method == "imshow": + if lvls is None: + im = cax.imshow( + arr, aspect=aspect, origin="lower", cmap=cmap, norm=norm + ) + else: + im = cax.imshow( + arr, + aspect=aspect, + origin="lower", + interpolation="none", + cmap=cmap, + norm=norm, + ) + elif method == "contour": + if lvls is None: + im = cax.contour(arr, colors="k") + im = cax.contourf(arr, cmap=cmap, norm=norm) + cax.set_aspect(aspect) + else: + cax.set_aspect(aspect) + # im = cax.contour(arr,linewidths=0.5,levels=lvls,colors='k',) + # lvls = lvls[1:-1] + # print(lvls) + im = cax.contour(arr, linewidths=1.0, colors="k", levels=lvls) + # im = cax.contourf(arr,levels=lvls,extend='both') + # lvls = lvls[1:-1] + im = cax.contourf(arr, levels=lvls, extend="both", cmap=cmap, norm=norm) + # im = cax.contour(arr,linewidths=0.5,levels=lvls[1:-1],colors='k',vmin=lvls[0],vmax=lvls[-1]) + # im = cax.contourf(arr,levels=lvls[1:-1],extend='both',vmin=lvls[0],vmax=lvls[-1]) + cax.set_aspect(aspect) + return im + + +class animator_2D(plotter): + def __init__(self, time_series, ncols, figsize=(16, 8)): + self.time_series = time_series + self.frns = time_series.shape[0] + super().__init__(self.time_series[0], ncols, figsize) + + self.update_plot = self.update_plot + self.fig.tight_layout() + self.suptitle = None + self.method = None + + def animate(self, interval=100, **kwargs): + self.ims, self.caxs, self.baxs = self.plot(**kwargs) + anim = animation.FuncAnimation( + self.fig, + self.update_plot, + self.frns, + fargs=( + self.time_series, + self.ims, + self.caxs, + self.baxs, + self.fig, + self.suptitle, + self.method, + ), + interval=interval, + ) + return anim + + @staticmethod + def update_plot(frame_number, time_series, ims, caxs, baxs, img, title, method): + if method == "imshow": + for ii, im in enumerate(ims): + arr = time_series[frame_number][ii][0] + im.set_array(arr) + im.set_clim(arr.min(), arr.max()) + elif method == "contour": + for ii, cax in enumerate(caxs): + arr = time_series[frame_number][ii][0] + im = ims[ii] + bax = baxs[ii] + for c in cax.collections: + cax.collections.remove(c) + # for c in cax.collections: + # cax.collections.remove(c) + # for c in cax.collections: + # cax.collections.remove(c) + im = cax.contour(arr, colors="k") + im = cax.contourf(arr) + plt.colorbar(im, cax=bax) + if title is not None: + stt = title(frame_number) + img.suptitle(stt) + else: + img.suptitle(frame_number) + + +class plotter_1d(object): + def __init__(self, ncols=3, nrows=2, figsize=(12, 12), fontsize=16): + plt.rcParams.update({"font.size": fontsize}) + self.fig, self.ax = plt.subplots( + ncols=ncols, nrows=nrows, sharex=False, sharey=False, figsize=figsize + ) + self.nrows = nrows + self.ncols = ncols + + self.img = plt + + def set_suptitle(self, suptitle): + self.img.suptitle(suptitle) + + def set_x(self, x_axs): + self.x = x_axs + + def get_ax(self, i): + # if not hasattr(self,'x'): + # assert 0, "x-axis has not been set, use set_x(x_axs)." + if self.ncols == 1 and self.nrows == 1: + return self.ax + elif self.nrows == 1 or self.ncols == 1: + return self.ax[i] + else: + row = int(np.floor(i / self.ncols)) + col = int(i % self.ncols) + + return self.ax[row, col] + + def save_fig(self, fn, format=".pdf"): + self.img.savefig(fn + format, bbox_inches="tight", pad_inches=0) + + +def labels(): + labels_dict = { + "rho": r"$\rho$, density", + "rhou": r"$\rho u$, horizontal momentum", + "rhov": r"$\rho w$, vertical momentum", + "rhow": r"$\rho v$, horizontal momentum", + "buoy": r"buoyancy", + "rhoX": r"$\rho / \Theta$, mass-weighted inverse pot. temp.", + "rhoY": r"$P$, mass-weighted potential temperature", + "p2_nodes": r"$\pi^\prime$, Exner pressure perturbation", + } + return labels_dict + + +def labels_increment(): + labels_dict = labels() + labels_dict["p2_nodes"] = r"$\delta \pi$, nodal Exner pressure increment" + return labels_dict + + +def swe_labels(): + labels_dict = { + "rho": r"$h$, water depth", + "rhou": r"$h u$, horizontal momentum", + "rhov": r"$h w$, vertical momentum", + "rhow": r"$h v$, horizontal momentum", + "buoy": r"buoyancy", + "vorty": r"vorticity", + "rhoX": r"$h / \Theta$", + "rhoY": r"$h (\rho \Theta)$, water depth", + "p2_nodes": r"$h^\prime$, water depth perturbation", + } + return labels_dict + + +def lake_labels(): + labels_dict = { + "rho": r"$h^{(0)}$, leading-order" + "\n water depth", + "rhou": r"$h^{(0)} u^{(0)}$, leading-order horizontal momentum", + "rhov": r"$h^{(0)} w^{(0)}$, leading-order vertical momentum", + "rhow": r"$h^{(0)} v^{(0)}$, leading-order horizontal momentum", + "buoy": r"buoyancy", + "vorty": r"vorticity", + "rhoX": r"$h / \Theta$", + "rhoY": r"$h^{(0)} (\rho \Theta)$, leading-order" + "\n water depth", + "p2_nodes": r"$h^{(1)}$, next-to-leading" + "\n order water depth", + } + return labels_dict + + +def swe_labels_increment(): + labels_dict = swe_labels() + labels_dict["p2_nodes"] = r"$\delta h^\prime$, water depth increment" + return labels_dict + + +def short_labels(): + labels_dict = { + "rho": r"$\rho$", + "rhou": r"$\rho u$", + "rhov": r"$\rho v$", + "rhow": r"$\rho w$", + "buoy": r"buoyancy", + "rhoX": r"$\rho / \Theta$", + "rhoY": r"$\rho \Theta$", + "p2_nodes": r"$\pi$", + } + + return labels_dict diff --git a/src/pybella/vis/utils.py b/src/pybella/vis/utils.py new file mode 100644 index 00000000..b22ee596 --- /dev/null +++ b/src/pybella/vis/utils.py @@ -0,0 +1,424 @@ +import numpy as np +import h5py +import time + + +class test_case(object): + def __init__(self, base_fn, py_dir, Nx, Ny, end_time, Nz=None): + self.base_fn = base_fn + self.py_dir = py_dir + self.grid_x = Nx + self.grid_y = Ny + self.grid_z = Nz + self.end_time = end_time + if Nz is not None and Ny > 1: + self.ndim = 3 + else: + self.ndim = 2 + + self.cb_suffix = self.cb_suffix + self.get_tag_dict = self.get_tag_dict + self.py_out = self.py_out + + self.get_filename = self.get_filename + self.get_path = self.get_path + self.get_arr = self.get_arr + + self.i0 = tuple( + [ + slice( + None, + ) + ] + * self.ndim + ) + if self.grid_z is not None and Ny == 1: + self.i2 = tuple([slice(2, -2)] * (self.ndim + 1)) + else: + self.i2 = tuple([slice(2, -2)] * self.ndim) + + def cb_suffix(self, fs, ts, suffix=""): + if suffix != "": + return "%s_cont_blend_fs=%i_ts=%i" % (suffix, fs, ts) + else: + return "cont_blend_fs=%i_ts=%i" % (fs, ts) + + @staticmethod + def get_tag_dict(): + td = { + 0: "before_flux", + 1: "before_advect", + 2: "after_advect", + 3: "after_ebnaexp", + 4: "after_ebnaimp", + 5: "after_half_step", + 6: "after_efna", + 7: "after_full_advect", + 8: "after_full_ebnaexp", + 9: "after_full_step", + } + return td + + @staticmethod + def get_debug_attrs(): + dd = { + "p2_initial": "p2_initial", + "hcenter": "hcenter", + "wplusx": "wplusx", + "wplusy": "wplusy", + "wplusz": "wplusz", + "rhs": "rhs", + "rhs_nodes": "rhs_nodes", + "p2_full": "p2_full", + } + return dd + + def get_filename(self, N, suffix, format="h5"): + if self.ndim == 2: + fn = "%s_ensemble=%i_%i_%i_%.6f_%s.%s" % ( + self.base_fn, + N, + self.grid_x, + self.grid_y, + self.end_time, + suffix, + format, + ) + if self.ndim == 3 or self.grid_z is not None: + fn = "%s_ensemble=%i_%i_%i_%i_%.6f_%s.%s" % ( + self.base_fn, + N, + self.grid_x, + self.grid_y, + self.grid_z, + self.end_time, + suffix, + format, + ) + return fn + + def get_path(self, fn): + path = self.py_dir + fn + return path + + @staticmethod + def py_out(pyfile, py_dataset, time): + return pyfile[str(py_dataset)][str(py_dataset) + time][:], pyfile[ + str(py_dataset) + ][str(py_dataset) + time].attrs.get("t") + + def get_arr( + self, + path, + time, + N, + attribute, + label_type="TIME", + tag="after_full_step", + inner=False, + avg=False, + file=None, + ): + if inner == False: + inner = self.i0 + else: + inner = self.i2 + + if file is None: + file = h5py.File(path, "r") + + array = [] + + if not hasattr(self, "t_arr"): + self.t_arr = [] + t_arr = self.t_arr + for n in range(N): + if label_type == "TIME": + t_label = "_ensemble_mem=%i_%.3f_%s" % (n, time, tag) + elif label_type == "WINDOW_STEP": + if N == 1: + t_label = "_%.3d_%s" % (time, tag) + elif label_type == "STEP": + if N == 1: + t_label = "_ensemble_mem=0_%.3d_%s" % (time, tag) + else: + t_label = "_ensemble_mem=%i_%.3d_%s" % (n, time, tag) + + arr, t = self.py_out(file, attribute, time=t_label) + array.append(arr[inner]) + t_arr.append(t) + + array = np.array(array) + if avg == True: + array = array.mean(axis=0) + + if file is None: + file.close() + self.t_arr = t_arr + return np.array(array) + + def spatially_averaged_rmse(self, arrs, refs, avg=False, grid_type="c"): + diff = [] + refs = refs[:, np.newaxis, ...] + refs = np.repeat(refs, arrs.shape[1], axis=1) + # print(arrs.shape, refs.shape) + for arr, ref in zip(arrs, refs): + # arr = (arr[self.i2]) + # ref = (ref[self.i2]) + if grid_type == "n": + # print("arr before ie1 shape", arr.shape) + arr = self.get_ie1(arr) + ref = self.get_ie1(ref) + # print(arr.shape) + if avg == True: + for ens_mem in arr: + ens_mem -= ens_mem.mean() + # arr -= arr.mean() + ref -= ref.mean() + + ref_ampl = ref.max() - ref.min() + factor = ref_ampl + factor = 1.0 + + diff.append(np.sqrt(((arr - ref) ** 2).mean())) + # diff.append(np.sqrt( ((arr - ref)**2).mean() / (ref[0]**2).mean() ) ) + # diff.append(np.sqrt((((arr - ref) / ref)**2).mean()) ) + # Method A + # err = [np.linalg.norm(mem - ref[0]) / np.linalg.norm(ref[0]) for mem in arr] + # err = np.array(err).mean() + # diff.append(err) + # Method B + # diff.append(np.linalg.norm(arr-ref) / np.linalg.norm(ref)) + # Method C + # diff.append(np.linalg.norm((arr-ref).mean(axis=0)) / np.linalg.norm(ref[0])) + + return np.array(diff) + + def ensemble_spread(self, arrs, avg=False, grid_type="c"): + diff = [] + + refs = arrs.mean(axis=1) + refs = refs[:, np.newaxis, ...] + refs = np.repeat(refs, arrs.shape[1], axis=1) + # print(arrs.shape, refs.shape) + + for arr, ref in zip(arrs, refs): + if grid_type == "n": + arr = self.get_ie1(arr) + ref = self.get_ie1(ref) + if avg == True: + for ens_mem in arr: + ens_mem -= ens_mem.mean() + # arr -= arr.mean() + ref -= ref.mean() + + diff.append(np.sqrt(((arr - ref) ** 2).mean())) + return np.array(diff) + + def probe_rmse(self, arrs, refs, probe_loc, avg=False, inner=False): + diff = [] + + for arr, ref in zip(arrs, refs): + if avg == True: + arr = arr.mean(axis=0) + ref = ref.mean(axis=0) + + if arr.ndim == 3: + arr = arr[:, 0, :] + ref = ref[:, 0, :] + + if inner == True: + arr = arr[probe_loc[0], probe_loc[1]] + ref = ref[probe_loc[0], probe_loc[1]] + else: + arr = arr[probe_loc[0], probe_loc[1]] + ref = ref[probe_loc[0], probe_loc[1]] + + # diff.append(np.sqrt(((arr - ref)**2).mean())) + diff.append(np.linalg.norm(arr - ref)) + return np.array(diff) + + # the first and last node rows / columns are repeated in periodic bcs, when we take mean we want to avoid this. + @staticmethod + def get_ie1(arr): + arr = arr.squeeze() + ndim = arr.ndim + ie1 = tuple([slice(0, -1)] * ndim) + arr = arr[ie1] + return arr + + def get_mean(self, arrs, grid_type="c"): + arr_mean = [] + + for arr in arrs: + if grid_type == "n": + arr = self.get_ie1(arr) + arr_mean.append(arr.mean()) + + return arr_mean + + @staticmethod + def get_probe_loc(arrs, probe_loc): + time_series = [] + for arr in arrs: + time_series.append(arr[probe_loc[0], probe_loc[1]]) + return time_series + + def get_ensemble( + self, + times, + N, + attribute, + suffix, + cont_blend=False, + ts=0, + fs=0, + label_type="TIME", + tag="after_full_step", + avg=False, + diff=False, + inner=True, + load_ic=True, + get_fn=True, + fn="", + ): + self.t_arr = [] + if cont_blend == True: + suffix += cb_suffix(fs, ts) + + if get_fn: + fn = self.get_filename(N, suffix) + else: + fn = fn + path = self.get_path(fn) + + file = h5py.File(path, "r") + + # arr_lst = np.zeros((times.size+1),dtype=np.ndarray) + arr_lst = [] + if load_ic: + arr = self.get_arr( + path, + 0, + N, + attribute, + tag="ic", + label_type=label_type, + avg=avg, + inner=inner, + file=file, + ) + arr_lst.append(arr) + for tt, time in enumerate(times): + arr = self.get_arr( + path, + time, + N, + attribute, + tag=tag, + label_type=label_type, + avg=avg, + inner=inner, + file=file, + ) + # arr_lst[tt] = arr + arr_lst.append(arr) + + if diff == True: + arr_lst = get_diff(arr_lst) + + file.close() + if load_ic: + self.t_arr[0] = 0.0 + return np.array(arr_lst) + + def get_time_series( + self, + times, + N, + attribute, + suffix, + probe_loc, + tag="after_full_step", + cont_blend=False, + ts=0, + fs=0, + label_type="TIME", + diff=False, + slc=[None], + avg=False, + ): + if self.ndim == 3 and slc[0] == None: + assert ( + 0 + ), "3D array has no 2D slice, define argument slc=(slice(...),slice(...),slice(...))" + + probe_row = probe_loc[0] + probe_col = probe_loc[1] + + if cont_blend == True: + suffix += "_" + self.cb_suffix(fs, ts) + + fn = self.get_filename(N, suffix) + path = self.get_path(fn) + + probe = [] + for time in times: + arr = self.get_arr( + path, time, N, attribute, tag=tag, label_type=label_type, inner=True + ) + + if avg == True: + arr -= arr.mean() + + if self.ndim == 3: + arr = arr[slc].squeeze() + probe.append(arr[probe_row, probe_col]) + + probe = np.array(probe) + if diff == True: + return get_diff(probe) + elif diff == False: + return probe + + +def bin_func(obs, ens_mem_shape): + obs = obs.reshape( + ens_mem_shape[0], + obs.shape[0] // ens_mem_shape[0], + ens_mem_shape[1], + obs.shape[1] // ens_mem_shape[1], + ) + return obs.mean(axis=(1, 3)) + + +def rmse(diff): + return np.sqrt((diff**2).mean()) + + +def get_diff(probe): + probe = np.array(probe) + return probe[1:] - probe[:-1] + + +def spatially_averaged_rmse(arr, ref): + # arr = arr[2:-2,2:-2] + # ref = ref[2:-2,2:-2] + + # arr -= arr.mean() + # ref -= ref.mean() + + return np.sqrt(((arr - ref) ** 2).mean()) + + +class prt_time(object): + # simple profiler for utils and plottting_tools + def __init__(self, debug=True): + self.tic = time.time() + self.debug = debug + + def prtt(self, label=""): + curr_time = time.time() + if self.debug == True: + print(label, curr_time - self.tic) + self.tic = curr_time diff --git a/src/tests/diagnostics.py b/src/tests/diagnostics.py deleted file mode 100644 index afa6fd08..00000000 --- a/src/tests/diagnostics.py +++ /dev/null @@ -1,157 +0,0 @@ -# import numpy as np -import numpy as np -import yaml -import termcolor -import logging -import pybella.vis.plotting_tools as vis_pt -import pybella.vis.utils as vis_utils - - -class compare_sol(object): - def __init__(self, current_run): - self.current_run = current_run - self.__init() - self.__get_tc() - - def update_targets(self): - self.arr_dump = {} - - for tc_name, tc in self.tcs.items(): - tp = self.tps[tc_name] - self.arr_dump[tp.name] = {} - - for attribute in tp.attributes: - self.arr_dump[tp.name][attribute] = float( - self.__get_ens(tc, tp, attribute, summed=True) - ) - - with open("./src/tests/test_targets.yml", "w") as outfile: - yaml.dump(self.arr_dump, outfile, default_flow_style=False) - - def test_do(self, Sol, p2n, plot=False): - self.__read_yaml() - - try: - target_values = self.target[self.current_run] - except: - assert 0, "test %s has no target for comparison" % (self.current_run) - - if plot: - self.__plot_comparison(Sol, p2n) - - for key, value in target_values.items(): - ref = value - - if key != "p2_nodes": - test = getattr(Sol, key).astype("float32").sum() - else: - test = p2n.astype("float32").sum() - - ## use try and except - assert ( - np.isclose(ref, test) - ), "sum for attribute %s of %s changed with discrepancy:\n%.6f\n%.6f" % ( - key, - self.current_run, - ref, - test, - ) - - logging.info(termcolor.colored("##########", "green")) - logging.info(termcolor.colored("Test passed for %s" %self.current_run, "green")) - logging.info(termcolor.colored("##########", "green")) - - def __init(self): - path = "./outputs/" - - tv_2D = test_params( - "test_travelling_vortex", path, "target_travelling_vortex", 64, 64, [100] - ) - igw = test_params( - "test_internal_long_wave", path, "target_internal_long_wave", 301, 10, [30] - ) - lmbw = test_params("test_lamb_wave", path, "target_lamb_wave", 151, 15, [30]) - - self.tps = { - "test_travelling_vortex": tv_2D, - "test_internal_long_wave": igw, - "test_lamb_wave": lmbw, - } - # self.tps = [tv_2D] - - def __get_tc(self): - self.tcs = {} - for test_name, test_param in self.tps.items(): - fn = test_param.fn + ".h5" - tc = vis_utils.test_case(fn, test_param.dir, test_param.Nx, test_param.Ny, "") - - self.tcs[test_name] = tc - - def __read_yaml(self): - with open("./src/tests/test_targets.yml", "r") as infile: - self.target = yaml.safe_load(infile) - - def __plot_comparison(self, Sol, p2n): - tc = self.tcs[self.current_run] - tp = self.tps[self.current_run] - - for attribute in tp.attributes: - arr_plots = [] - - ref_sol = self.__get_ens(tc, tp, attribute).T - - if attribute != "p2_nodes": - test_sol = getattr(Sol, attribute).T - else: - test_sol = p2n.T - - arr_plots.append([ref_sol, "ref"]) - arr_plots.append([test_sol, "test"]) - arr_plots.append([ref_sol - test_sol, "diff"]) - - pl = vis_pt.plotter(arr_plots, ncols=3, figsize=(12, 3), sharey=False) - _ = pl.plot(method="contour", lvls=None, suptitle=attribute) - pl.img.savefig(tp.dir + attribute + ".png") - - @staticmethod - def __get_ens(tc, params, attribute, summed=False): - times = params.times - l_typ = params.l_typ - - tags = tc.get_tag_dict() - tag = "ic" if times[0] == 0.0 else tags[9] - - ens = tc.get_ensemble( - times, - 1, - attribute, - "", - label_type=l_typ, - tag=tag, - inner=False, - get_fn=False, - fn=params.fn + ".h5", - load_ic=False, - avg=True, - )[0] - - if summed: - return ens.sum() - else: - return ens - - -class test_params(object): - def __init__(self, name, path, fn, Nx, Ny, times): - - self.name = name - self.dir = path + fn + "/" - self.fn = "%s_%i_%i" % (fn, Nx, Ny) - - self.Nx = Nx - self.Ny = Ny - - self.times = times - self.l_typ = "WINDOW_STEP" - - self.attributes = ["rho", "rhou", "rhov", "rhow", "rhoY", "rhoX", "p2_nodes"] diff --git a/src/utils/sim_params.py b/src/utils/sim_params.py deleted file mode 100644 index ae22f7cf..00000000 --- a/src/utils/sim_params.py +++ /dev/null @@ -1,28 +0,0 @@ -# simulation parameters -debug = False -da_debug = False - -random_seed = 888 -print_precision = 18 -output_path = './outputs' - - -# global constants -class global_constants(object): - def __init__(self): - self.nspec = 1 - self.buoy = 0 - - self.grav = 9.81 # [m s^{-2}] - self.omega = 0.0 # [s^{-1}] - - self.R_gas = 287.4 # [J kg^{-1} K^{-1}] - self.R_vap = 461.0 - self.Q_vap = 2.53e+06 - self.gamm = 1.4 - - self.p_ref = 8.61 * 1e4 # [N/m^2] - self.T_ref = 300.00 # [K] - - self.h_ref = 10000.0 # [m] - self.t_ref = 100.0 # [s] \ No newline at end of file diff --git a/src/vis/plotting_tools.py b/src/vis/plotting_tools.py deleted file mode 100644 index baa4b328..00000000 --- a/src/vis/plotting_tools.py +++ /dev/null @@ -1,331 +0,0 @@ -import matplotlib.pyplot as plt -import mpl_toolkits.axes_grid1 as axes_grid -import matplotlib.animation as animation -import numpy as np -import itertools - - -class plotter(object): - def __init__(self,arr_lst, ncols=4, figsize=(12,8), fontsize=14, sharexlabel=False, shareylabel=False, sharex=False, sharey=False): - plt.rcParams.update({'font.size': fontsize}) - self.arr_lst = np.array(arr_lst, dtype='object') - N = self.arr_lst.shape[0] - - if N > ncols: - self.nrows = int(np.ceil(N/ncols)) - self.ncols = ncols - if N <= ncols: - self.nrows = 1 - self.ncols = N - self.N = N - - ridx = np.arange(self.nrows) - cidx = np.arange(self.ncols) - - self.idx = [] - if self.nrows > 1: - for pair in itertools.product(ridx, cidx): - self.idx.append(pair) - else: - self.idx = cidx - - self.visualise = self.visualise - self.fig, self.ax = plt.subplots(ncols=self.ncols,nrows=self.nrows,figsize=figsize,sharex=sharex,sharey=sharey) - - self.sharexlabel = sharexlabel - self.shareylabel = shareylabel - - self.img = plt - - def set_axes(self,**kwargs): - for key, value in kwargs.items(): - setattr(self,key,value) - - def set_cax_axes(self,cax,n): - if hasattr(self, 'x_locs') : cax.set_xticks(self.x_locs) - if hasattr(self, 'x_axs') : cax.set_xticklabels(self.x_axs) - if hasattr(self, 'y_locs') : cax.set_yticks(self.y_locs) - if hasattr(self, 'y_axs') : cax.set_yticklabels(self.y_axs) - cax.tick_params(axis='x', pad=15) - if self.sharexlabel: - if int(n // self.ncols) == self.nrows - 1: - if hasattr(self, 'x_label') : cax.set_xlabel(self.x_label) - else: - if hasattr(self, 'x_label') : cax.set_xlabel(self.x_label) - if self.shareylabel: - if n % self.ncols == 0: - if hasattr(self, 'y_label') : cax.set_ylabel(self.y_label) - else: - if hasattr(self, 'y_label') : cax.set_ylabel(self.y_label) - if hasattr(self, 'axhline'): cax.axhline(self.axhline,c='k',lw=0.5) - if hasattr(self, 'axvline'): cax.axvline(self.axvline,c='k',lw=0.5) - if hasattr(self, 'marker'): - for marker in self.marker: - cax.plot(marker[0],marker[1],marker='x',c=marker[2], ms=18, mew=4) - if hasattr(self, 'rects'): - for rect in self.rects: - cax.add_patch(rect) - - - def plot(self,method='imshow',inner=False,suptitle="",rect=[0, 0.03, 1, 0.95],aspect='auto',lvls=None,cmaps=None): - if method != 'imshow' and method != 'contour': - assert 0, "Visualisation method not implemented!" - - if self.N > 1: - ims, caxs, baxs = [], [], [] - for n, arr in enumerate(self.arr_lst): - arr, title = arr[0], arr[1] - if inner == True: - arr = arr[2:-2,2:-2] - cax = self.ax[self.idx[n]] - - lvl = lvls[n] if lvls is not None else None - cmap = cmaps[n] if cmaps is not None else 'viridis' - - im = self.visualise(method,cax,arr,aspect,lvl,cmap) - if type(title) == str: - cax.set_title(title) - elif type(title) == np.ndarray or type(title) == list: - cax.set_title(title[0], fontsize=title[1], fontweight=title[2]) - loc = cax.get_xticklabels() - self.set_cax_axes(cax,n) - caxs.append(cax) - divider = axes_grid.make_axes_locatable(cax) - bax = divider.append_axes("right", size="5%", pad=0.05) - if method == 'imshow' and lvl is not None: -# plt.colorbar(im, cax=bax, ticks=lvl)#, format='%.3f') - plt.colorbar(im, cax=bax, ticks=lvl, extend='both')#, format='%.3f') - else: - # plt.colorbar(im, cax=bax, extend='both')#, format='%.3f') - plt.colorbar(im, cax=bax, extendrect=True) -# plt.colorbar(im, cax=bax, ticks=lvl) - baxs.append(bax) - if hasattr(self, 'cbar_label'): - bax.set_xlabel(self.cbar_label) - bax.xaxis.set_label_position('top') - if hasattr(self, 'cbar_label_coords'): - bax.xaxis.set_label_coords(self.cbar_label_coords[0],self.cbar_label_coords[1]) - ims.append(im) - - for i in range(n+1,self.nrows*self.ncols): - self.fig.delaxes(self.ax[self.idx[i]]) - - else: - arr, title = self.arr_lst[0][0], self.arr_lst[0][1] - if inner == True: - arr = arr[2:-2,2:-2] -# cax = self.fig.gca() - cmap = cmaps if cmaps is not None else 'viridis' - cax = self.ax - im = self.visualise(method,cax,arr,aspect,lvls,cmap) - cax.set_title(title) - self.set_cax_axes(cax,0) - caxs = [cax] - divider = axes_grid.make_axes_locatable(cax) - bax = divider.append_axes("right", size="5%", pad=0.05) - if hasattr(self, 'cbar_label'): - bax.set_xlabel(self.cbar_label) - bax.xaxis.set_label_position('top') - if method == 'imshow' and lvls is not None: -# plt.colorbar(im, cax=bax, ticks=lvls)#, format='%.3f') - plt.colorbar(im, cax=bax, ticks=lvls) - else: - plt.colorbar(im, cax=bax, extendrect=True) -# plt.colorbar(im, cax=bax) - if aspect != 'auto' and aspect != 'equal': - bax.set_aspect(float(aspect)*10.0) - ims = [im] - baxs = [bax] - - plt.suptitle(suptitle) - plt.tight_layout(rect=rect) - # plt.subplots_adjust(hspace = .000005) - - if self.N > 1: - return ims, caxs, baxs - else: - return ims, caxs, baxs - - def save_fig(self, fn, format='.pdf'): - self.fig.tight_layout() - self.fig.savefig(fn + format, bbox_inches = 'tight', pad_inches = 0.1) - - - @staticmethod - def visualise(method,cax,arr,aspect,lvls,cmap): - if cmap is None: - norm = None - cmap = 'viridis' - elif len(cmap) == 2: - norm = cmap[1] - cmap = cmap[0] - else: - norm = None - cmap = cmap - - if method == 'imshow': - if lvls is None: - im = cax.imshow(arr,aspect=aspect,origin='lower',cmap=cmap, norm=norm) - else: - im = cax.imshow(arr,aspect=aspect,origin='lower',interpolation='none',cmap=cmap, norm=norm) - elif method == 'contour': - if lvls is None: - im = cax.contour(arr,colors='k') - im = cax.contourf(arr,cmap=cmap,norm=norm) - cax.set_aspect(aspect) - else: - cax.set_aspect(aspect) -# im = cax.contour(arr,linewidths=0.5,levels=lvls,colors='k',) -# lvls = lvls[1:-1] -# print(lvls) - im = cax.contour(arr,linewidths=1.0,colors='k',levels=lvls) -# im = cax.contourf(arr,levels=lvls,extend='both') -# lvls = lvls[1:-1] - im = cax.contourf(arr,levels=lvls,extend='both',cmap=cmap,norm=norm) - # im = cax.contour(arr,linewidths=0.5,levels=lvls[1:-1],colors='k',vmin=lvls[0],vmax=lvls[-1]) - # im = cax.contourf(arr,levels=lvls[1:-1],extend='both',vmin=lvls[0],vmax=lvls[-1]) - cax.set_aspect(aspect) - return im - - -class animator_2D(plotter): - def __init__(self,time_series,ncols,figsize=(16,8)): - self.time_series = time_series - self.frns = time_series.shape[0] - super().__init__(self.time_series[0], ncols, figsize) - - self.update_plot = self.update_plot - self.fig.tight_layout() - self.suptitle = None - self.method = None - - def animate(self, interval=100, **kwargs): - self.ims, self.caxs, self.baxs = self.plot(**kwargs) - anim = animation.FuncAnimation(self.fig, self.update_plot, self.frns, fargs=(self.time_series, self.ims, self.caxs, self.baxs, self.fig, self.suptitle, self.method), interval=interval) - return anim - - @staticmethod - def update_plot(frame_number, time_series, ims, caxs, baxs, img, title, method): - if method == 'imshow': - for ii,im in enumerate(ims): - arr = time_series[frame_number][ii][0] - im.set_array(arr) - im.set_clim(arr.min(),arr.max()) - elif method == 'contour': - for ii, cax in enumerate(caxs): - arr = time_series[frame_number][ii][0] - im = ims[ii] - bax = baxs[ii] - for c in cax.collections: - cax.collections.remove(c) -# for c in cax.collections: -# cax.collections.remove(c) -# for c in cax.collections: -# cax.collections.remove(c) - im = cax.contour(arr,colors='k') - im = cax.contourf(arr) - plt.colorbar(im, cax=bax) - if title is not None: - stt = title(frame_number) - img.suptitle(stt) - else: - img.suptitle(frame_number) - - -class plotter_1d(object): - def __init__(self,ncols=3,nrows=2,figsize=(12,12),fontsize=16): - plt.rcParams.update({'font.size': fontsize}) - self.fig, self.ax = plt.subplots(ncols=ncols,nrows=nrows, sharex=False, sharey=False, figsize=figsize) - self.nrows = nrows - self.ncols = ncols - - self.img = plt - - def set_suptitle(self,suptitle): - self.img.suptitle(suptitle) - - def set_x(self,x_axs): - self.x = x_axs - - def get_ax(self,i): - - #if not hasattr(self,'x'): - #assert 0, "x-axis has not been set, use set_x(x_axs)." - if self.ncols == 1 and self.nrows == 1: - return self.ax - elif self.nrows == 1 or self.ncols == 1: - return self.ax[i] - else: - row = int(np.floor(i/self.ncols)) - col = int(i%self.ncols) - - return self.ax[row,col] - - def save_fig(self, fn, format='.pdf'): - self.img.savefig(fn + format, bbox_inches = 'tight', pad_inches = 0) - - -def labels(): - labels_dict = { - 'rho' : r'$\rho$, density', - 'rhou' : r'$\rho u$, horizontal momentum', - 'rhov' : r'$\rho w$, vertical momentum', - 'rhow' : r'$\rho v$, horizontal momentum', - 'buoy' : r'buoyancy', - 'rhoX' : r'$\rho / \Theta$, mass-weighted inverse pot. temp.', - 'rhoY' : r'$P$, mass-weighted potential temperature', - 'p2_nodes' : r'$\pi^\prime$, Exner pressure perturbation' - } - return labels_dict - -def labels_increment(): - labels_dict = labels() - labels_dict['p2_nodes'] = r'$\delta \pi$, nodal Exner pressure increment' - return labels_dict - -def swe_labels(): - labels_dict = { - 'rho' : r'$h$, water depth', - 'rhou' : r'$h u$, horizontal momentum', - 'rhov' : r'$h w$, vertical momentum', - 'rhow' : r'$h v$, horizontal momentum', - 'buoy' : r'buoyancy', - 'vorty' : r'vorticity', - 'rhoX' : r'$h / \Theta$', - 'rhoY' : r'$h (\rho \Theta)$, water depth', - 'p2_nodes' : r'$h^\prime$, water depth perturbation' - } - return labels_dict - -def lake_labels(): - labels_dict = { - 'rho' : r'$h^{(0)}$, leading-order' + '\n water depth', - 'rhou' : r'$h^{(0)} u^{(0)}$, leading-order horizontal momentum', - 'rhov' : r'$h^{(0)} w^{(0)}$, leading-order vertical momentum', - 'rhow' : r'$h^{(0)} v^{(0)}$, leading-order horizontal momentum', - 'buoy' : r'buoyancy', - 'vorty' : r'vorticity', - 'rhoX' : r'$h / \Theta$', - 'rhoY' : r'$h^{(0)} (\rho \Theta)$, leading-order' + '\n water depth', - 'p2_nodes' : r'$h^{(1)}$, next-to-leading' + '\n order water depth' - } - return labels_dict - -def swe_labels_increment(): - labels_dict = swe_labels() - labels_dict['p2_nodes'] = r'$\delta h^\prime$, water depth increment' - return labels_dict - -def short_labels(): - labels_dict = { - 'rho' : r'$\rho$', - 'rhou' : r'$\rho u$', - 'rhov' : r'$\rho v$', - 'rhow' : r'$\rho w$', - 'buoy' : r'buoyancy', - 'rhoX' : r'$\rho / \Theta$', - 'rhoY' : r'$\rho \Theta$', - 'p2_nodes' : r'$\pi$' - } - - return labels_dict diff --git a/src/vis/utils.py b/src/vis/utils.py deleted file mode 100644 index e152a04e..00000000 --- a/src/vis/utils.py +++ /dev/null @@ -1,331 +0,0 @@ -import numpy as np -import h5py -import time - -class test_case(object): - def __init__(self,base_fn,py_dir,Nx,Ny,end_time,Nz=None): - self.base_fn = base_fn - self.py_dir = py_dir - self.grid_x = Nx - self.grid_y = Ny - self.grid_z = Nz - self.end_time = end_time - if Nz is not None and Ny > 1: - self.ndim = 3 - else: - self.ndim = 2 - - self.cb_suffix = self.cb_suffix - self.get_tag_dict = self.get_tag_dict - self.py_out = self.py_out - - self.get_filename = self.get_filename - self.get_path = self.get_path - self.get_arr = self.get_arr - - self.i0 = tuple([slice(None,)]*self.ndim) - if self.grid_z is not None and Ny == 1: - self.i2 = tuple([slice(2,-2)]*(self.ndim+1)) - else: - self.i2 = tuple([slice(2,-2)]*self.ndim) - - def cb_suffix(self,fs,ts,suffix=""): - if suffix != "": - return "%s_cont_blend_fs=%i_ts=%i" %(suffix,fs,ts) - else: - return "cont_blend_fs=%i_ts=%i" %(fs,ts) - - - @staticmethod - def get_tag_dict(): - td = { - 0 : 'before_flux', - 1 : 'before_advect', - 2 : 'after_advect', - 3 : 'after_ebnaexp', - 4 : 'after_ebnaimp', - 5 : 'after_half_step', - 6 : 'after_efna', - 7 : 'after_full_advect', - 8 : 'after_full_ebnaexp', - 9 : 'after_full_step', - } - return td - - @staticmethod - def get_debug_attrs(): - dd = { - 'p2_initial' : 'p2_initial', - 'hcenter' : 'hcenter', - 'wplusx' : 'wplusx', - 'wplusy' : 'wplusy', - 'wplusz' : 'wplusz', - 'rhs' : 'rhs', - 'rhs_nodes' : 'rhs_nodes', - 'p2_full' : 'p2_full' - } - return dd - - def get_filename(self,N,suffix,format='h5'): - if self.ndim == 2: - fn = "%s_ensemble=%i_%i_%i_%.6f_%s.%s" %(self.base_fn,N,self.grid_x,self.grid_y,self.end_time,suffix,format) - if self.ndim == 3 or self.grid_z is not None: - fn = "%s_ensemble=%i_%i_%i_%i_%.6f_%s.%s" %(self.base_fn,N,self.grid_x,self.grid_y,self.grid_z,self.end_time,suffix,format) - return fn - - - def get_path(self,fn): - path = self.py_dir + fn - return path - - - @staticmethod - def py_out(pyfile,py_dataset,time): - return pyfile[str(py_dataset)][str(py_dataset)+time][:], pyfile[str(py_dataset)][str(py_dataset)+time].attrs.get('t') - - - def get_arr(self, path, time, N, attribute, label_type='TIME', tag='after_full_step', inner=False, avg=False, file=None): - if inner == False: - inner = self.i0 - else: - inner = self.i2 - - if file is None: - file = h5py.File(path,'r') - - array = [] - - if not hasattr(self,'t_arr'): self.t_arr = [] - t_arr = self.t_arr - for n in range(N): - if label_type == 'TIME': - t_label = '_ensemble_mem=%i_%.3f_%s' %(n,time, tag) - elif label_type == 'WINDOW_STEP': - if N==1: - t_label = '_%.3d_%s' %(time, tag) - elif label_type == 'STEP': - if N==1: - t_label = '_ensemble_mem=0_%.3d_%s' %(time, tag) - else: - t_label = '_ensemble_mem=%i_%.3d_%s' %(n,time, tag) - - arr, t = self.py_out(file,attribute,time=t_label) - array.append(arr[inner]) - t_arr.append(t) - - array = np.array(array) - if avg == True: - array = array.mean(axis=0) - - if file is None: - file.close() - self.t_arr = t_arr - return np.array(array) - - def spatially_averaged_rmse(self, arrs,refs,avg=False, grid_type='c'): - diff = [] - refs = refs[:,np.newaxis,...] - refs = np.repeat(refs, arrs.shape[1], axis=1) -# print(arrs.shape, refs.shape) - for arr, ref in zip(arrs,refs): - #arr = (arr[self.i2]) - #ref = (ref[self.i2]) - if grid_type == 'n': -# print("arr before ie1 shape", arr.shape) - arr = self.get_ie1(arr) - ref = self.get_ie1(ref) -# print(arr.shape) - if avg==True: - for ens_mem in arr: - ens_mem -= ens_mem.mean() -# arr -= arr.mean() - ref -= ref.mean() - - ref_ampl = ref.max() - ref.min() - factor = ref_ampl - factor = 1.0 - - diff.append( np.sqrt(((arr - ref)**2).mean()) ) -# diff.append(np.sqrt( ((arr - ref)**2).mean() / (ref[0]**2).mean() ) ) -# diff.append(np.sqrt((((arr - ref) / ref)**2).mean()) ) - # Method A - # err = [np.linalg.norm(mem - ref[0]) / np.linalg.norm(ref[0]) for mem in arr] - #err = np.array(err).mean() - # diff.append(err) - # Method B - # diff.append(np.linalg.norm(arr-ref) / np.linalg.norm(ref)) - # Method C -# diff.append(np.linalg.norm((arr-ref).mean(axis=0)) / np.linalg.norm(ref[0])) - - return np.array(diff) - - def ensemble_spread(self, arrs,avg=False, grid_type='c'): - diff = [] - - refs = arrs.mean(axis=1) - refs = refs[:,np.newaxis,...] - refs = np.repeat(refs,arrs.shape[1],axis=1) -# print(arrs.shape, refs.shape) - - for arr, ref in zip(arrs,refs): - if grid_type == 'n': - arr = self.get_ie1(arr) - ref = self.get_ie1(ref) - if avg==True: - for ens_mem in arr: - ens_mem -= ens_mem.mean() -# arr -= arr.mean() - ref -= ref.mean() - - diff.append(np.sqrt(((arr - ref)**2).mean())) - return np.array(diff) - - - def probe_rmse(self, arrs, refs, probe_loc, avg=False, inner=False): - diff = [] - - for arr, ref in zip(arrs,refs): - if avg == True: - arr = arr.mean(axis=0) - ref = ref.mean(axis=0) - - if arr.ndim == 3: - arr = arr[:,0,:] - ref = ref[:,0,:] - - if inner == True: - arr = arr[probe_loc[0],probe_loc[1]] - ref = ref[probe_loc[0],probe_loc[1]] - else: - arr = arr[probe_loc[0],probe_loc[1]] - ref = ref[probe_loc[0],probe_loc[1]] - - #diff.append(np.sqrt(((arr - ref)**2).mean())) - diff.append(np.linalg.norm(arr-ref)) - return np.array(diff) - - # the first and last node rows / columns are repeated in periodic bcs, when we take mean we want to avoid this. - @staticmethod - def get_ie1(arr): - arr = arr.squeeze() - ndim = arr.ndim - ie1 = tuple([slice(0,-1)]*ndim) - arr = arr[ie1] - return arr - - def get_mean(self, arrs, grid_type='c'): - arr_mean = [] - - for arr in arrs: - if grid_type == 'n': - arr = self.get_ie1(arr) - arr_mean.append(arr.mean()) - - return arr_mean - - - @staticmethod - def get_probe_loc(arrs, probe_loc): - time_series = [] - for arr in arrs: - time_series.append(arr[probe_loc[0],probe_loc[1]]) - return time_series - - - def get_ensemble(self, times, N, attribute, suffix, cont_blend=False, ts=0, fs=0, label_type='TIME', tag='after_full_step', avg=False, diff=False, inner=True, load_ic=True, get_fn=True, fn=""): - self.t_arr = [] - if cont_blend == True: - suffix += cb_suffix(fs,ts) - - if get_fn: - fn = self.get_filename(N,suffix) - else: - fn = fn - path = self.get_path(fn) - - file = h5py.File(path,'r') - -# arr_lst = np.zeros((times.size+1),dtype=np.ndarray) - arr_lst = [] - if load_ic: - arr = self.get_arr(path, 0, N, attribute, tag='ic', label_type=label_type, avg=avg, inner=inner, file=file) - arr_lst.append(arr) - for tt, time in enumerate(times): - arr = self.get_arr(path, time, N, attribute, tag=tag, label_type=label_type, avg=avg, inner=inner, file=file) -# arr_lst[tt] = arr - arr_lst.append(arr) - - - if diff == True: - arr_lst = get_diff(arr_lst) - - file.close() - if load_ic: - self.t_arr[0] = 0.0 - return np.array(arr_lst) - - - def get_time_series(self, times, N, attribute, suffix, probe_loc, tag='after_full_step', cont_blend=False, ts=0, fs=0, label_type='TIME', diff=False, slc=[None], avg=False): - if self.ndim == 3 and slc[0] == None: - assert 0, "3D array has no 2D slice, define argument slc=(slice(...),slice(...),slice(...))" - - probe_row = probe_loc[0] - probe_col = probe_loc[1] - - if cont_blend == True: - suffix += '_' + self.cb_suffix(fs,ts) - - fn = self.get_filename(N,suffix) - path = self.get_path(fn) - - probe = [] - for time in times: - arr = self.get_arr(path, time, N, attribute, tag=tag, label_type=label_type, inner=True) - - if avg == True: - arr -= arr.mean() - - if self.ndim == 3: - arr = arr[slc].squeeze() - probe.append(arr[probe_row,probe_col]) - - probe = np.array(probe) - if diff==True: - return get_diff(probe) - elif diff==False: - return probe - - - -def bin_func(obs,ens_mem_shape): - obs = obs.reshape(ens_mem_shape[0],obs.shape[0]//ens_mem_shape[0], - ens_mem_shape[1],obs.shape[1]//ens_mem_shape[1]) - return obs.mean(axis=(1,3)) - -def rmse(diff): - return np.sqrt((diff**2).mean()) - -def get_diff(probe): - probe = np.array(probe) - return probe[1:] - probe[:-1] - -def spatially_averaged_rmse(arr,ref): - #arr = arr[2:-2,2:-2] - #ref = ref[2:-2,2:-2] - - #arr -= arr.mean() - #ref -= ref.mean() - - return np.sqrt(((arr - ref)**2).mean()) - -class prt_time(object): - # simple profiler for utils and plottting_tools - def __init__(self, debug=True): - self.tic = time.time() - self.debug = debug - - def prtt(self, label=""): - curr_time = time.time() - if self.debug==True: - print(label, curr_time - self.tic) - self.tic = curr_time diff --git a/test_scripts/test_blending.py b/test_scripts/test_blending.py new file mode 100644 index 00000000..ad20cd26 --- /dev/null +++ b/test_scripts/test_blending.py @@ -0,0 +1,14 @@ +import pytest +import subprocess + +@pytest.mark.parametrize( + "ic", + [ + "test_blending_warm_bubble", + ], +) +def test_single_run(ic): + result = subprocess.run( + ["pybella", "-ic", ic, "-N", "1"], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr.splitlines()[-3:] diff --git a/test_scripts/test_flow_solver.py b/test_scripts/test_flow_solver.py new file mode 100644 index 00000000..99053894 --- /dev/null +++ b/test_scripts/test_flow_solver.py @@ -0,0 +1,21 @@ +import pytest +import subprocess + +@pytest.mark.parametrize("ic", + ["test_travelling_vortex", + "test_internal_long_wave", + "test_lamb_wave", + "test_unstable_lamb",] + ) +def test_single_run(ic): + result = subprocess.run( + ["pybella", "-ic", ic, "-N", "1"], + capture_output=True, + text=True + ) + + assert result.returncode == 0, ( + f"Command failed with return code {result.returncode}\n" + f"STDERR:\n{result.stderr.strip()}\n" + f"STDOUT:\n{result.stdout.strip()}" + ) \ No newline at end of file