pax_global_header00006660000000000000000000000064151763404050014517gustar00rootroot0000000000000052 comment=1b21a88815807921a926b795f45ddf3c12ea3ee9 scientificcomputing-io4dolfinx-d21fc0e/000077500000000000000000000000001517634040500203305ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/.coveragerc000066400000000000000000000005421517634040500224520ustar00rootroot00000000000000[run] parallel = true source = io4dolfinx [html] directory = htmlcov [xml] output = coverage.xml [report] exclude_also = def __repr__ if self.debug: if settings.DEBUG raise AssertionError raise NotImplementedError if 0: if __name__ == .__main__.: if TYPE_CHECKING: class .*\bProtocol\): @(abc\.)?abstractmethod scientificcomputing-io4dolfinx-d21fc0e/.github/000077500000000000000000000000001517634040500216705ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/.github/dependabot.yml000066400000000000000000000010301517634040500245120ustar00rootroot00000000000000# To get started with Dependabot version updates, you'll need to specify which # package ecosystems to update and where the package manifests are located. # Please see the documentation for all configuration options: # https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file version: 2 updates: - package-ecosystem: "github-actions" # See documentation for possible values directory: "/" # Location of package manifests schedule: interval: "weekly" scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/000077500000000000000000000000001517634040500237255ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/build_docs.yml000066400000000000000000000030431517634040500265570ustar00rootroot00000000000000name: Build documentation on: workflow_call: inputs: tag: description: "Tag of DOLFINx docker image" default: "stable" required: true type: string workflow_dispatch: inputs: tag: description: "Tag of DOLFINx docker image" default: "stable" required: true type: string pull_request: branches: - main schedule: - cron: "0 8 * * *" env: DEB_PYTHON_INSTALL_LAYOUT: deb_system DEFAULT_TAG: stable ARTIFACT_NAME: docs PUBLISH_DIR: ./_build/html PIP_NO_BINARY: h5py jobs: get_image_tag: runs-on: ubuntu-latest outputs: image: ${{ steps.docker_tag.outputs.image }} steps: - id: docker_tag run: echo "image=${{ inputs.tag || env.DEFAULT_TAG }}" >> $GITHUB_OUTPUT build-docs: needs: get_image_tag runs-on: ubuntu-latest container: ghcr.io/fenics/dolfinx/dolfinx:${{ needs.get_image_tag.outputs.image }} steps: # This action sets the current path to the root of your github repo - uses: actions/checkout@v6 - name: Update pip run: python3 -m pip install --upgrade pip setuptools - name: Install dependencies run: python3 -m pip install --no-binary=h5py -e ".[docs]" - name: Build docs run: jupyter book build . - name: Upload documentation as artifact uses: actions/upload-artifact@v7 if: always() with: name: ${{ env.ARTIFACT_NAME }} path: ${{ env.PUBLISH_DIR }} if-no-files-found: error scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/check_formatting.yml000066400000000000000000000023551517634040500277640ustar00rootroot00000000000000name: Check formatting on: workflow_call: inputs: tag: description: "Tag of DOLFINx docker image" default: "nightly" required: true type: string workflow_dispatch: inputs: tag: description: "Tag of DOLFINx docker image" default: "nightly" required: true type: string pull_request: branches: - main schedule: - cron: "0 8 * * *" env: DEB_PYTHON_INSTALL_LAYOUT: deb_system DEFAULT_TAG: nightly jobs: get_image_tag: runs-on: ubuntu-latest outputs: image: ${{ steps.docker_tag.outputs.image }} steps: - id: docker_tag run: echo "image=${{ inputs.tag || env.DEFAULT_TAG }}" >> $GITHUB_OUTPUT build: needs: get_image_tag runs-on: ubuntu-latest container: ghcr.io/fenics/dolfinx/dolfinx:${{ needs.get_image_tag.outputs.image }} steps: - uses: actions/checkout@v6 - name: Update pip run: python3 -m pip install --upgrade pip setuptools - name: Install code run: python3 -m pip install .[dev] - name: Check code formatting with ruff run: | ruff check . ruff format --check . - name: Mypy check run: python3 -m mypy scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/create_legacy_checkpoint.yml000066400000000000000000000014631517634040500314520ustar00rootroot00000000000000name: Generate adios4dolfinx legacy data on: workflow_call: inputs: artifact_name: type: string required: true description: "Name of the artifact to be created" jobs: create-adios-data: env: data_dir: "legacy_checkpoint" adios4dolfinx_version: "0.7.1" runs-on: "ubuntu-latest" container: ghcr.io/fenics/dolfinx/dolfinx:v0.7.3 steps: - uses: actions/checkout@v6 - name: Install legacy version of adios4dolfinx run: python3 -m pip install adios4dolfinx==${adios4dolfinx_version} - name: Create datasets run: python3 ./tests/create_legacy_checkpoint.py --output-dir=$data_dir - uses: actions/upload-artifact@v7 with: name: ${{ inputs.artifact_name }} path: ./${{ env.data_dir }} scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/create_legacy_data.yml000066400000000000000000000011631517634040500302310ustar00rootroot00000000000000name: Generate data from Legacy DOLFIN on: workflow_call: inputs: artifact_name: type: string required: true description: "Name of the artifact to be created" jobs: create-dolfin-data: env: data_dir: "legacy" runs-on: "ubuntu-latest" container: ghcr.io/scientificcomputing/fenics:2024-02-19 steps: - uses: actions/checkout@v6 - name: Create datasets run: python3 ./tests/create_legacy_data.py --output-dir=$data_dir - uses: actions/upload-artifact@v7 with: name: ${{ inputs.artifact_name }} path: ./${{ env.data_dir }} scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/deploy_pages.yml000066400000000000000000000033151517634040500271250ustar00rootroot00000000000000name: Github Pages on: push: branches: [main] # pull_request: # branches: [main] permissions: contents: read pages: write id-token: write # Allow one concurrent deployment concurrency: group: "pages" cancel-in-progress: true jobs: build-docs: uses: ./.github/workflows/build_docs.yml with: tag: "stable" test-code: uses: ./.github/workflows/test_dolfinx_versions.yml with: dolfinx_tag: "stable" deploy: needs: [build-docs, test-code] environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} if: github.event_name == 'push' runs-on: ubuntu-latest steps: - name: Download docs artifact # docs artifact is uploaded by build-docs job uses: actions/download-artifact@v8 with: name: docs path: "./public" - name: Download code coverage report artifact # docs artifact is uploaded by build-docs job uses: actions/download-artifact@v8 with: name: code-coverage-report-stable path: "./public/code-coverage-report-stable" - name: Download code coverage report artifact # docs artifact is uploaded by build-docs job uses: actions/download-artifact@v8 with: name: code-coverage-report-nightly path: "./public/code-coverage-report-nightly" - name: Upload artifact uses: actions/upload-pages-artifact@v5 with: path: "./public" - name: Checkout uses: actions/checkout@v6 - name: Setup Pages uses: actions/configure-pages@v6 - name: Deploy to GitHub Pages id: deployment uses: actions/deploy-pages@v5 scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/pypi.yml000066400000000000000000000012431517634040500254310ustar00rootroot00000000000000name: pypi on: push: branches: [main] tags: - "v*" jobs: dist: runs-on: ubuntu-latest steps: - uses: actions/checkout@v6 - name: Build SDist and wheel run: pipx run build - uses: actions/upload-artifact@v7 with: path: dist/* - name: Check metadata run: pipx run twine check dist/* publish: needs: [dist] runs-on: ubuntu-latest if: startsWith(github.ref, 'refs/tags') environment: pypi permissions: id-token: write steps: - uses: actions/download-artifact@v8 with: name: artifact path: dist - uses: pypa/gh-action-pypi-publish@release/v1 scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/test_dolfinx_versions.yml000066400000000000000000000013651517634040500311070ustar00rootroot00000000000000name: Test dolfinx versions on: workflow_call: inputs: dolfinx_tag: required: true type: string pull_request: jobs: # 1. Prerequisite jobs (Data creation) create-datasets: uses: ./.github/workflows/create_legacy_data.yml with: artifact_name: "legacy_mpich" create-legacy-datasets: uses: ./.github/workflows/create_legacy_checkpoint.yml with: artifact_name: "legacy_checkpoint_mpich" # 2. The Matrix Job run-tests: needs: [create-datasets, create-legacy-datasets] strategy: fail-fast: false matrix: tag: ["nightly", "stable"] # Call the reusable workflow uses: ./.github/workflows/test_workflow.yml with: dolfinx_tag: ${{ matrix.tag }}scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/test_package_openmpi.yml000066400000000000000000000046541517634040500306420ustar00rootroot00000000000000name: Test package with different mpi compilers on: push: branches: - main pull_request: branches: - main workflow_call: workflow_dispatch: schedule: - cron: "0 8 * * *" jobs: create-datasets: uses: ./.github/workflows/create_legacy_data.yml with: artifact_name: "legacy_ompi" create-legacy-datasets: uses: ./.github/workflows/create_legacy_checkpoint.yml with: artifact_name: "legacy_checkpoint_ompi" test-code: runs-on: ubuntu-latest needs: [create-datasets, create-legacy-datasets] container: ${{ matrix.container }} env: DEB_PYTHON_INSTALL_LAYOUT: deb_system PETSC_ARCH: "linux-gnu-real64-32" OMPI_ALLOW_RUN_AS_ROOT: 1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1 PRTE_MCA_rmaps_default_mapping_policy: :oversubscribe working-directory: ./src PIP_NO_BINARY: h5py strategy: matrix: adios2: ["v2.10.2", "v2.11.0"] # "default", Removed as HDF5 2.0 in dev env is broken with adios2<2.11 container: - ghcr.io/fenics/test-env:current-openmpi - ghcr.io/fenics/test-env:current-mpich steps: - uses: actions/checkout@v6 - name: Update pip run: python3 -m pip install --upgrade pip - name: Install build requirements run: python3 -m pip install -r build-requirements.txt - name: Install DOLFINx uses: jorgensd/actions/install-dolfinx@v0.4 with: adios2: ${{ matrix.adios2 }} petsc_arch: ${{ env.PETSC_ARCH }} dolfinx: main basix: main ufl: main ffcx: main working-directory: ${{ env.working-directory}} - name: Download legacy data uses: actions/download-artifact@v8 with: name: legacy_ompi path: ./legacy - name: Download legacy data uses: actions/download-artifact@v8 with: name: legacy_checkpoint_ompi path: ./legacy_checkpoint - name: Install package run: | HDF5_MPI=ON HDF5_PKGCONFIG_NAME="hdf5" python3 -m pip install h5py --no-build-isolation --no-binary=h5py python3 -m pip install .[test] - name: Run tests run: | coverage run --rcfile=.coveragerc -m mpi4py -m pytest -xvs ./tests/ - name: Run tests in parallel run: | mpirun -n 2 coverage run --rcfile=.coveragerc -m mpi4py -m pytest -xvs ./tests/ scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/test_redhat.yml000066400000000000000000000050661517634040500267650ustar00rootroot00000000000000name: Test package with redhat on: # push: # branches: # - main # pull_request: # branches: # - main workflow_call: workflow_dispatch: # schedule: # - cron: "0 8 * * *" jobs: create-datasets: uses: ./.github/workflows/create_legacy_data.yml with: artifact_name: "legacy_ompi" create-legacy-datasets: uses: ./.github/workflows/create_legacy_checkpoint.yml with: artifact_name: "legacy_checkpoint_ompi" test-code: runs-on: "ubuntu-latest" needs: [create-datasets, create-legacy-datasets] container: docker.io/fenicsproject/test-env:current-redhat env: DEB_PYTHON_INSTALL_LAYOUT: deb_system PETSC_ARCH: "" PETSC_DIR: "/usr/local/" PYTHONPATH: "/usr/local/lib/:${PYTHONPATH}" working-directory: ./src strategy: matrix: adios2: ["master"] steps: - uses: actions/checkout@v6 - name: Get pip flags based on version id: python-version shell: bash -el {0} run: | MODERN_PIP=$(python3 -c "import sys; t = sys.version_info >= (3, 11, 0); sys.stdout.write(str(t))") if [ ${MODERN_PIP} == "True" ]; then FLAGS="--break-system-packages" else FLAGS="" python3 -m pip install --upgrade pip fi echo "PYTHON_FLAGS=${FLAGS}" >> "$GITHUB_OUTPUT" - name: Update pip run: python3 -m pip install ${{ steps.python-version.outputs.PYTHON_FLAGS}} --upgrade pip setuptools - name: Install DOLFINx uses: jorgensd/actions/install-dolfinx@v0.4 with: adios2: ${{ matrix.adios2 }} petsc_arch: ${{ env.PETSC_ARCH }} petsc_dir: ${{ env.PETSC_DIR }} dolfinx: main basix: main ufl: main ffcx: main working-directory: ${{ env.working-directory}} - name: Download legacy data uses: actions/download-artifact@v8 with: name: legacy_ompi path: ./legacy - name: Download legacy data uses: actions/download-artifact@v8 with: name: legacy_checkpoint_ompi path: ./legacy_checkpoint - name: Install package run: python3 -m pip install ${{ steps.python-version.outputs.PYTHON_FLAGS}} --check-build-dependencies .[test] - name: Run tests run: | coverage run --rcfile=.coveragerc -m mpi4py -m pytest -xvs ./tests - name: Run tests in parallel run: | mpirun -n 4 coverage run --rcfile=.coveragerc -m mpi4py -m pytest -xvs ./tests scientificcomputing-io4dolfinx-d21fc0e/.github/workflows/test_workflow.yml000066400000000000000000000041041517634040500273600ustar00rootroot00000000000000name: Test on: workflow_call: inputs: dolfinx_tag: required: true type: string jobs: test-code: runs-on: ubuntu-24.04 # The container tag is now dynamic based on the input passed by the caller container: ghcr.io/fenics/dolfinx/dolfinx:${{ inputs.dolfinx_tag }} steps: - uses: actions/checkout@v6 - name: Update pip run: python3 -m pip install --upgrade pip - name: Install build requirements run: python3 -m pip install -r build-requirements.txt # Note: We assume the caller workflow has already uploaded these artifacts - name: Download legacy mpich data uses: actions/download-artifact@v8 with: name: legacy_mpich path: ./legacy - name: Download legacy checkpoint mpich data uses: actions/download-artifact@v8 with: name: legacy_checkpoint_mpich path: ./legacy_checkpoint - name: Install package run: | HDF5_MPI=ON HDF5_PKGCONFIG_NAME="hdf5" python3 -m pip install h5py --no-build-isolation --no-binary=h5py python3 -m pip install .[test] - name: Show adios2 version run: python3 -c "import adios2; print(adios2.__version__)" - name: Show h5py version run: python3 -c "import h5py; print(h5py.__version__)" - name: Show hdf5 version run: python3 -c "import h5py; print(h5py.version.hdf5_version)" - name: Run tests run: coverage run --rcfile=.coveragerc -m pytest -xvs ./tests/ - name: Run tests in parallel run: mpirun -n 4 coverage run --rcfile=.coveragerc -m mpi4py -m pytest -xvs ./tests/ - name: Combine coverage reports run: | coverage combine coverage report -m coverage html # Use the tag in the artifact name so parallel runs don't overwrite each other - name: Upload coverage report uses: actions/upload-artifact@v7 with: name: code-coverage-report-${{ inputs.dolfinx_tag }} path: htmlcov if-no-files-found: errorscientificcomputing-io4dolfinx-d21fc0e/.gitignore000066400000000000000000000035221517634040500223220ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ _build/ output/ *.h5 .coverage.* .coverage *.h5 *.xdmf *.bp/ *.pvd *.vtuscientificcomputing-io4dolfinx-d21fc0e/.pre-commit-config.yaml000066400000000000000000000006621517634040500246150ustar00rootroot00000000000000# See https://pre-commit.com for more information # See https://pre-commit.com/hooks.html for more hooks repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. rev: 'v0.15.12' hooks: # Run the linter. - id: ruff args: [ --fix ] # Run the formatter. - id: ruff-format - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.20.2 hooks: - id: mypy scientificcomputing-io4dolfinx-d21fc0e/CITATION.cff000066400000000000000000000006031517634040500222210ustar00rootroot00000000000000cff-version: "1.2.0" authors: - family-names: Dokken given-names: Jørgen Schartum orcid: "https://orcid.org/0000-0001-6489-8858" - family-names: Finsberg given-names: Henrik Nicolay Topnes orcid: "0000-0003-3766-2393" contact: - family-names: Dokken given-names: Jørgen Schartum orcid: "https://orcid.org/0000-0001-6489-8858" doi: 10.5281/zenodo.11094985 scientificcomputing-io4dolfinx-d21fc0e/CODE_OF_CONDUCT.md000066400000000000000000000053521517634040500231340ustar00rootroot00000000000000 # Code of Conduct ### Our Standards Examples of behavior that contributes to creating a positive environment include: * Using welcoming and inclusive language * Being respectful of differing viewpoints and experiences * Gracefully accepting constructive criticism * Focusing on what is best for the community * Showing empathy towards other community members Examples of unacceptable behavior by participants include: * The use of sexualized language or imagery and unwelcome sexual attention or advances * Trolling, insulting/derogatory comments, and personal or political attacks * Public or private harassment * Publishing others' private information, such as a physical or electronic address, without explicit permission * Other conduct which could reasonably be considered inappropriate in a professional setting ### Our Responsibilities Project maintainers are responsible for clarifying the standards of acceptable behavior and are expected to take appropriate and fair corrective action in response to any instances of unacceptable behavior. Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, or to ban temporarily or permanently any contributor for other behaviors that they deem inappropriate, threatening, offensive, or harmful. ### Scope This Code of Conduct applies both within project spaces and in public spaces when an individual is representing the project or its community. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. Representation of a project may be further defined and clarified by project maintainers. ### Enforcement Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team at dokken@simula.no. All complaints will be reviewed and investigated and will result in a response that is deemed necessary and appropriate to the circumstances. The project team is obligated to maintain confidentiality with regard to the reporter of an incident. Further details of specific enforcement policies may be posted separately. Project maintainers who do not follow or enforce the Code of Conduct in good faith may face temporary or permanent repercussions as determined by other members of the project's leadership. ### Attribution This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, available at [http://contributor-covenant.org/version/1/4][version] [homepage]: http://contributor-covenant.org [version]: http://contributor-covenant.org/version/1/4/ scientificcomputing-io4dolfinx-d21fc0e/CONTRIBUTING.md000066400000000000000000000073301517634040500225640ustar00rootroot00000000000000# Contributor guidelines When contributing to this repository, please first [create an issue](https://github.com/jorgensd/io4dolfinx/issues/new/choose) containing information about the missing feature or the bug that you would like to fix. Here you can discuss the change you want to make with the maintainers of the repository. Please note we have a code of conduct, please follow it in all your interactions with the project. ## New contributor guide To get an overview of the project, read the [documentation](https://jorgensd.github.io/io4dolfinx). Here are some resources to help you get started with open source contributions: - [Finding ways to contribute to open source on GitHub](https://docs.github.com/en/get-started/exploring-projects-on-github/finding-ways-to-contribute-to-open-source-on-github) - [Set up Git](https://docs.github.com/en/get-started/quickstart/set-up-git) - [GitHub flow](https://docs.github.com/en/get-started/quickstart/github-flow) - [Collaborating with pull requests](https://docs.github.com/en/github/collaborating-with-pull-requests) ## Pull Request Process ### Pull Request - When you're finished with the changes, create a pull request, also known as a PR. It is also OK to create a [draft pull request](https://github.blog/2019-02-14-introducing-draft-pull-requests/) from the very beginning. Once you are done you can click on the ["Ready for review"] button. You can also [request a review](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/requesting-a-pull-request-review) from one of the maintainers. - Don't forget to [link PR to the issue that you opened ](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue). - Enable the checkbox to [allow maintainer edits](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/allowing-changes-to-a-pull-request-branch-created-from-a-fork) so the branch can be updated for a merge. Once you submit your PR, a team member will review your proposal. We may ask questions or request for additional information. - We may ask for changes to be made before a PR can be merged, either using [suggested changes](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/incorporating-feedback-in-your-pull-request) or pull request comments. You can apply suggested changes directly through the UI. You can make any other changes in your fork, then commit them to your branch. - As you update your PR and apply changes, mark each conversation as [resolved](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/commenting-on-a-pull-request#resolving-conversations). - If you run into any merge issues, checkout this [git tutorial](https://lab.github.com/githubtraining/managing-merge-conflicts) to help you resolve merge conflicts and other issues. - Please make sure that all tests are passing, github pages renders nicely, and code coverage are are not lower than before your contribution. You see the different github action workflows by clicking the "Action" tab in the GitHub repository. Note that for a pull request to be accepted, it has to pass all the tests on CI, which includes: - `mypy`: typechecking - `ruff`: Code formatting - `pytest`: Successfull execution of all tests in the `tests` folder. ### Our Pledge In the interest of fostering an open and welcoming environment, we as contributors and maintainers pledge to making participation in our project and our community a harassment-free experience for everyone, regardless of age, body size, disability, ethnicity, gender identity and expression, level of experience, nationality, personal appearance, race, religion, or sexual identity and orientation. scientificcomputing-io4dolfinx-d21fc0e/LICENSE000066400000000000000000000021261517634040500213360ustar00rootroot00000000000000Copyright 2023 Jørgen S. Dokken, Henrik N.T. Finsberg and Simula Research Laboratory Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. scientificcomputing-io4dolfinx-d21fc0e/README.md000066400000000000000000000134721517634040500216160ustar00rootroot00000000000000# io4dolfinx - A framework for reading and writing data to various mesh formats **io4dolfinx** is an extension for [DOLFINx](https://github.com/FEniCS/dolfinx/) that provides advanced input/output capabilities. It focuses on **N-to-M checkpointing** (writing data on N processors, reading on M processors) and supports reading/writing various mesh formats using interchangeable backends. ## Installation The library is compatible with the DOLFINx nightly release, v0.10.0, and v0.9.0. ```bash python3 -m pip install io4dolfinx ``` For specific backend requirements (like ADIOS2 or H5PY), see the [Installation Guide](./docs/installation.md). ## Quick Start Here is a minimal example of saving and loading a simulation state (Checkpointing). ```python from pathlib import Path from mpi4py import MPI import dolfinx import io4dolfinx # 1. Create a mesh and function comm = MPI.COMM_WORLD mesh = dolfinx.mesh.create_unit_square(comm, 10, 10) V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) u = dolfinx.fem.Function(V) u.interpolate(lambda x: x[0] + x[1]) u.name = "my_function" # 2. Write checkpoint # The mesh must be written before the function filename = Path("checkpoint.bp") io4dolfinx.write_mesh(filename, mesh) io4dolfinx.write_function(filename, u, time=0.0) # 3. Read checkpoint # This works even if the number of MPI processes changes (N-to-M) mesh_new = io4dolfinx.read_mesh(filename, comm) V_new = dolfinx.fem.functionspace(mesh_new, ("Lagrange", 1)) u_new = dolfinx.fem.Function(V_new) io4dolfinx.read_function(filename, u_new, time=0.0, name="my_function") ``` ## Features and Backends `io4dolfinx` supports custom user backends. You can switch backends by passing `backend="name"` to IO functions. ### Checkpointing (N-to-M) Many finite element applications requires storage of functions that cannot be associated with the nodes or cells of the mesh. Therefore, we have implemented our own, native checkpointing format that supports N-to-M checkpointing (write data on N processors, read in on M) through the following backends: - [h5py](./docs/backends/h5py.rst): Requires HDF5 with MPI support to work, but can store, meshes, partitioning info, meshtags, function data and more. - [adios2](./docs/backends/adios2.rst): Requires [ADIOS 2](https://adios2.readthedocs.io/en/latest/) compiled with MPI support and Python bindings. Supports the same set of operations as the `h5py` backend. The code uses the ADIOS2/Python-wrappers and h5py module to write DOLFINx objects to file, supporting N-to-M (_recoverable_) and N-to-N (_snapshot_) checkpointing. See: [Checkpointing in DOLFINx - FEniCS 23](https://jsdokken.com/checkpointing-presentation/#/) or the examples in the [Documentation](https://jsdokken.com/io4dolfinx/) for more information. For scalability, the code uses [MPI Neighbourhood collectives](https://www.mpi-forum.org/docs/mpi-3.1/mpi31-report/node200.htm) for communication across processes. ### Mesh IO (Import/Export) Most meshing formats supports associating data with the nodes of the mesh (the mesh can be higher order) and the cells of the mesh. The node data can be read in as P-th order Lagrange functions (where P is the order of the grid), while the cell data can be read in as piecewise constant (DG-0) functions. - [VTKHDF](./docs/backends/vtkhdf.rst): The new scalable format from VTK, called [VTKHDF](https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/index.html) is supported by the `vtkhdf` backend. - [XDMF](./docs/backends/xdmf.rst) (eXtensible Model Data Format): `.xdmf`. The `xdmf` backend supports the `HDF5` encoding, to ensure performance in parallel. - [PyVista](./docs/backends/pyvista.rst) (IO backend is meshio): The [pyvista](https://pyvista.org/) backend uses {py:func}`pyvista.read` to read in meshes, point data and cell data. `pyvista` relies on [meshio](https://github.com/nschloe/meshio) for most reading operations (including the XDMF ascii format). ## Advanced Usage The repository contains detailed documented examples in the `docs` folder: * [Reading and writing mesh checkpoints](./docs/writing_mesh_checkpoint.py) * [Storing mesh partitioning data](./docs/partitioned_mesh.py) (Avoid re-partitioning when restarting) * [Writing mesh-tags](./docs/meshtags.py) * [Writing function checkpoints](./docs/writing_functions_checkpoint.py) * [Checkpoint on input mesh](./docs/original_checkpoint.py) For a full API reference and backend details, see the [Documentation](https://jsdokken.com/io4dolfinx/). ### Legacy DOLFIN Support `io4dolfinx` can read checkpoints created by the legacy version of DOLFIN (Lagrange or DG functions). * Reading meshes from DOLFIN HDF5File-format. * Reading checkpoints from DOLFIN HDF5File and XDMFFile. ## Project Background ### Relation to adios4dolfinx This library is an evolution of [adios4dolfinx](https://doi.org/10.21105/joss.06451). It includes all functionality of the original library but has been refactored to support multiple IO backends (not just ADIOS2), making it easier to interface with different meshing formats while keeping the library structure sane. ### Statement of Need As large-scale, long-running simulations on HPC clusters become more common, the need to store intermediate solutions is crucial. If a system crashes or a computational budget is exceeded, checkpoints allow the simulation to resume without restarting from scratch. `io4dolfinx` extends DOLFINx with this essential functionality. ## Contributing Please read [CONTRIBUTING.md](CONTRIBUTING.md) for details on our code of conduct, and the process for submitting pull requests to us. ## Testing `io4dolfinx` includes a comprehensive test suite that ensures functionality across different backends and compatibility with legacy data formats, see the [Testing Guide](./docs/testing.md) for details. ## LICENSE This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.scientificcomputing-io4dolfinx-d21fc0e/_config.yml000066400000000000000000000036741517634040500224710ustar00rootroot00000000000000# Book settings # Learn more at https://jupyterbook.org/customize/config.html title: ADIOS2Wrappers author: Jørgen S. Dokken logo: "docs/logo.png" copyright: "2023" only_build_toc_files: true # Force re-execution of notebooks on each build. # See https://jupyterbook.org/content/execute.html execute: execute_notebooks: cache # Information about where the book exists on the web repository: url: https://github.com/scientificcomputing/io4dolfinx # Online location of your book branch: main html: use_issues_button: true use_repository_button: true parse: myst_enable_extensions: - amsmath - dollarmath - linkify sphinx: extra_extensions: - "sphinx.ext.autodoc" - "sphinx.ext.napoleon" - "sphinx.ext.viewcode" - "sphinx.ext.intersphinx" - "sphinx_codeautolink" config: html_last_updated_fmt: "%b %d, %Y" nb_custom_formats: .py: - jupytext.reads - fmt: py autodoc_typehints: signature autodoc_typehints_format: short codeautolink_concat_default: True intersphinx_mapping: basix: ["https://docs.fenicsproject.org/basix/main/python", null] ffcx: ["https://docs.fenicsproject.org/ffcx/main", null] ufl: ["https://docs.fenicsproject.org/ufl/main", null] dolfinx: ["https://docs.fenicsproject.org/dolfinx/main/python", null] petsc4py: ["https://petsc.org/release/petsc4py", null] mpi4py: ["https://mpi4py.readthedocs.io/en/stable", null] numpy: ["https://numpy.org/doc/stable/", null] pyvista: ["https://docs.pyvista.org/", null] packaging: ["https://packaging.pypa.io/en/stable/", null] matplotlib: ["https://matplotlib.org/stable/", null] ipyparallel: ["https://ipyparallel.readthedocs.io/en/stable/", null] python: ["https://docs.python.org/3", null] h5py: ["https://docs.h5py.org/en/stable/", null] adios2: ["https://adios2.readthedocs.io/en/latest", null] exclude_patterns: [".pytest_cache/*"] scientificcomputing-io4dolfinx-d21fc0e/_toc.yml000066400000000000000000000021021517634040500217720ustar00rootroot00000000000000format: jb-book root: README parts: - caption: How to guides chapters: - file: "docs/installation" - file: "docs/quickstart" - file: "docs/adding_backend" - file: "docs/migration_guide" - file: "docs/testing.md" - file: "docs/reading_legacy_data" - caption: Introduction to IPyParallel chapters: - file: "docs/ipyparallel_intro" - caption: Writing and reading mesh data chapters: - file: "docs/writing_mesh_checkpoint" - file: "docs/partitioned_mesh" - file: "docs/time_dependent_mesh" - file: "docs/meshtags" - caption: Writing and reading functions chapters: - file: "docs/writing_functions_checkpoint" - file: "docs/snapshot_checkpoint" - file: "docs/original_checkpoint" - caption: Python API chapters: - file: "docs/api" - file: "docs/backend" - file: "docs/backends/adios2" - file: "docs/backends/h5py" - file: "docs/backends/pyvista" - file: "docs/backends/xdmf" - file: "docs/backends/vtkhdf" - caption: Contributing chapters: - file: "CONTRIBUTING"scientificcomputing-io4dolfinx-d21fc0e/build-requirements.txt000066400000000000000000000001411517634040500247050ustar00rootroot00000000000000scikit-build-core[pyproject]>=0.11.0 nanobind>=1.3.2 setuptools>=77.0.3 packaging>=24.2 pkgconfigscientificcomputing-io4dolfinx-d21fc0e/docs/000077500000000000000000000000001517634040500212605ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/docs/adding_backend.md000066400000000000000000000112461517634040500245030ustar00rootroot00000000000000# Adding a custom backend {py:mod}`io4dolfinx` is designed to be backend-agnostic, meaning you can implement custom readers and writers for different file formats by adhering to a specific protocol. ## The IOBackend Protocol Any backend must implement the {py:class}`IOBackend` protocol defined in {py:mod}`io4dolfinx.backends`. This protocol ensures that the backend provides all necessary methods for reading and writing meshes, functions, and attributes. To use a custom backend, you simply pass the python module name (as a string) to the `backend` argument of any {py:mod}`io4dolfinx` function. The library will attempt to import the module and use it as the backend. ## Required Data Structures Your backend will interact with several data classes defined in {py:mod}`io4dolfinx.structures`. You should import these to type-hint your implementation correctly: * {py:class}`MeshData`: Contains local geometry, topology, and partitioning information for writing {py:class}`meshes`. * {py:class}`ReadMeshData`: A container for returning mesh data (cells, geometry, etc.) when reading. * {py:class}`FunctionData`: Contains function values, dofmaps, and permutation info for writing functions. * {py:class}`MeshTagsData`: Contains indices, values, and metadata for mesh tags. ## Implementation Checklist Your backend module must implement the functions listed below. Note that `comm` is always an {py:class}`MPI.Intracomm` and `filename` is a {py:class}`pathlib.Path` or {py:class}`str`. ### General Configuration * {py:func}`~io4dolfinx.backends.IOBackend.get_default_backend_args` * Returns a dictionary of default arguments (e.g., engine type) for your backend. ### Attribute IO * {py:func}`~io4dolfinx.backends.IOBackend.write_attributes` * Writes a dictionary of attributes (key-value pairs) to the file under the specified `name` (group). * {py:func}`~io4dolfinx.backends.IOBackend.read_attributes` * Reads and returns attributes associated with `name`. ### Mesh IO * {py:func}`~io4dolfinx.backends.IOBackend.write_mesh` * Writes mesh geometry, topology, and optionally partitioning data. * Must handle {py:class}`FileMode.write` (new file) and {py:class}`FileMode.append`. * {py:func}`~io4dolfinx.backends.IOBackend.read_mesh_data` * Reads mesh geometry and topology at a specific `time`. * If `read_from_partition` is True, it should read pre-calculated partitioning data to avoid re-partitioning. ### MeshTags IO * {py:func}`~io4dolfinx.backends.IOBackend.write_meshtags` * Writes mesh tag indices and values. * {py:func}`~io4dolfinx.backends.IOBackend.read_meshtags_data` * Reads mesh tags identified by `name`. ### Function IO * {py:func}`~io4dolfinx.backends.IOBackend.write_function` * Writes function values, global dofmaps, and cell permutations. * {py:func}`~io4dolfinx.backends.IOBackend.read_dofmap` * Reads the dofmap (connectivity) for the function `name`. * {py:func}`~io4dolfinx.backends.IOBackend.read_dofs` * Reads the local chunk of function values for a specific `time`. * Returns the array of values and the global starting index of that chunk. * {py:func}`~io4dolfinx.backends.IOBackend.read_cell_perms` * Reads cell permutation data used to map input cells to the current mesh. * {py:func}`~io4dolfinx.backends.IOBackend.read_timestamps` * Returns all available time-steps for a given function. ### Legacy Support (Optional but defined in protocol) * {py:func}`~io4dolfinx.backends.IOBackend.read_legacy_mesh` * Reads mesh data from legacy DOLFIN HDF5/XDMF formats. * {py:func}`~io4dolfinx.backends.IOBackend.read_hdf5_array` * Reads a raw array from an HDF5-like structure (used for legacy vector reading). ### Snapshots * {py:func}`~io4dolfinx.backends.IOBackend.snapshot_checkpoint` * Handles lightweight N-to-N checkpointing where data is saved exactly as distributed in memory without global reordering. ## Example Skeleton ```python from typing import Any from pathlib import Path from mpi4py import MPI import numpy as np import dolfinx from io4dolfinx.structures import MeshData, FunctionData, MeshTagsData, ReadMeshData from io4dolfinx.backends import FileMode def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: return arguments or {} def write_mesh(filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None, mode: FileMode, time: float): # Implementation here pass # ... Implement all other methods defined in IOBackend ... ```scientificcomputing-io4dolfinx-d21fc0e/docs/api.rst000066400000000000000000000001061517634040500225600ustar00rootroot00000000000000API Reference ============= .. automodule:: io4dolfinx :members: scientificcomputing-io4dolfinx-d21fc0e/docs/backend.rst000066400000000000000000000004031517634040500233760ustar00rootroot00000000000000Backend API ============= This module shows what functions and data-classes a new backend should implement to achieve full support of all funcntionality. .. automodule:: io4dolfinx.backends :members: .. automodule:: io4dolfinx.structures :members: scientificcomputing-io4dolfinx-d21fc0e/docs/backends/000077500000000000000000000000001517634040500230325ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/docs/backends/adios2.rst000066400000000000000000000013041517634040500247430ustar00rootroot00000000000000ADIOS2-backend -------------- The library depends on the Python-interface of [DOLFINx](https://github.com/) and an MPI-build of [ADIOS2](https://adios2.readthedocs.io/en/latest/setting_up/setting_up.html#as-package). Therefore `ADIOS2` should not be install through PYPI/pip, but has to be installed through Conda, Spack or from source. .. important:: ADIOS2<2.10.2 does not work properly with :code:`numpy>=2.0.0`. Everyone is advised to use the newest version of ADIOS2. This is for instance available through :code:`conda` or the :code:`ghcr.io/fenics/dolfinx/dolfinx:nightly` Docker-image. .. automodule:: io4dolfinx.backends.adios2.backend :members: :exclude-members: read_point_datascientificcomputing-io4dolfinx-d21fc0e/docs/backends/h5py.rst000066400000000000000000000025731517634040500244600ustar00rootroot00000000000000.. _h5pybackend: H5PY-backend ------------- The H5PY-backend relies on an MPI compatible build of :code:`HDF5`, which should be linked to the same MPI implementation as DOLFINx relies on. The DOLFINx docker images (:code:`ghcr.io/fenics/dolfinx/dolfinx:stable`) comes with an already configures MPI compatible HDF5 installation, and `h5py` can in turn be installed with .. code-block:: bash HDF5_MPI="ON" HDF5_DIR="/usr/local" python3 -m pip install --no-binary=h5py h5py --no-build-isolation If you are using `apt` on Ubuntu, this can for instance be achieved with the following commands (here using Docker). Note that this code block does not install DOLFINx, it just illustrates how to get the correct `h5py`. .. code-block:: dockerfile FROM ubuntu:24.04 AS base RUN apt-get update && apt-get install -y python3-dev python3-pip python3-venv libhdf5-mpi-dev libopenmpi-dev ENV VIRTUAL_ENV=/test-env ENV PATH=/test-env/bin:$PATH RUN python3 -m venv ${VIRTUAL_ENV} ENV HDF5_MPI="ON" ENV HDF5_DIR=/usr/lib/x86_64-linux-gnu/hdf5/openmpi/ ENV C_PATH=/usr/lib/x86_64-linux-gnu/openmpi/include/:${C_PATH} RUN python3 -m pip install setuptools cython numpy pkgconfig mpi4py RUN CC=mpicc python3 -m pip install --no-binary=h5py h5py --no-build-isolation .. automodule:: io4dolfinx.backends.h5py.backend :members: :exclude-members: read_point_datascientificcomputing-io4dolfinx-d21fc0e/docs/backends/pyvista.rst000066400000000000000000000005421517634040500252640ustar00rootroot00000000000000.. _pyvistabackend: PyVista-backend ---------------- .. automodule:: io4dolfinx.backends.pyvista.backend :members: :exclude-members: read_attributes, read_timestamps, write_attributes, write_mesh, write_meshtags, read_meshtags_data, read_dofmap, read_dofs, read_cell_perms, write_function, read_legacy_mesh, snapshot_checkpoint, read_hdf5_arrayscientificcomputing-io4dolfinx-d21fc0e/docs/backends/vtkhdf.rst000066400000000000000000000007601517634040500250550ustar00rootroot00000000000000.. _vtkhdfbackend: VTKHDF-backend -------------- Relies on MPI compatible :code:`h5py`, similar to the :ref:`H5PY backend`. See that backend for detailed installation instructions. .. automodule:: io4dolfinx.backends.vtkhdf.backend :members: :exclude-members: read_attributes, read_timestamps, write_attributes, write_mesh, write_meshtags, read_meshtags_data, read_dofmap, read_dofs, read_cell_perms, write_function, read_legacy_mesh, snapshot_checkpoint, read_hdf5_arrayscientificcomputing-io4dolfinx-d21fc0e/docs/backends/xdmf.rst000066400000000000000000000007511517634040500245250ustar00rootroot00000000000000.. _xdmfbackend: XDMF-backend ------------- Relies on MPI compatible :code:`h5py`, similar to the :ref:`H5PY backend`. See that backend for detailed installation instructions. .. automodule:: io4dolfinx.backends.xdmf.backend :members: :exclude-members: read_attributes, read_timestamps, write_attributes, write_mesh, write_meshtags, read_meshtags_data, read_dofmap, read_dofs, read_cell_perms, write_function, read_legacy_mesh, snapshot_checkpoint, read_hdf5_arrayscientificcomputing-io4dolfinx-d21fc0e/docs/installation.md000066400000000000000000000077241517634040500243150ustar00rootroot00000000000000# Installation The main way to install `io4dolfinx` is through [PYPI](https://pypi.org/project/io4dolfinx/), which provides pre-built binary wheels for most platforms. This is the recommended method for most users. ```bash python3 -m pip install io4dolfinx ``` `io4dolfinx` has some optional dependencies for specific backends (like ADIOS2 or H5PY). If you want to use these backends, you can install the library with the appropriate extras: - Test dependencies (for running the test suite): ```bash python3 -m pip install "io4dolfinx[test]" ``` - Documentation dependencies (for building the docs): ```bash python3 -m pip install "io4dolfinx[docs]" ``` - For HDF5 support with MPI, you need to have an HDF5 library installed with MPI support, and the `h5py` Python package installed with MPI support. You can install `h5py` with MPI support using pip: ```bash python3 -m pip install --no-binary=h5py h5py ``` - For pyvista support, you can install the `pyvista` package: ```bash python3 -m pip install pyvista ``` or equivalently ```bash python3 -m pip install "io4dolfinx[pyvista]" ``` - For ADIOS2 support you should have ADIOS2 installed with Python bindings, see https://adios2.readthedocs.io/en/latest/setting_up/setting_up.html for more info. ## Spack The FEniCS Spack packages uses a separate [spack repo](https://github.com/FEniCS/spack-fenics) to be possible to maintain and keep up to date. We do the same for the [packages](https://github.com/scientificcomputing/spack_repos.git) maintained by Scientific Computing at Simula Research Laboratory. To install `py-io4dolfinx`, one should first install spack on your system, then use the following commands in a new spack environment: ```bash spack repo add https://github.com/FEniCS/spack-fenics.git spack repo add https://github.com/scientificcomputing/spack_repos.git spack add py-io4dolfinx@1.1 ^py-fenics-dolfinx+petsc4py ^adios2+python+hdf5 ^petsc+mumps ``` to get an installation of `io4dolfinx` with all backends installed. If you require further petsc packages you should activate them by adding them to `^petsc+....`. See [Spack PETSc](https://packages.spack.io/package.html?name=petsc) for options. ## Docker An MPI build of ADIOS2 is installed in the official DOLFINx containers, and thus there are no additional dependencies required to install `io4dolfinx` on top of DOLFINx in these images. Create a Docker container, named for instance `dolfinx-checkpoint`. Use the `nightly` tag to get the main branch of DOLFINx, or `stable` to get the latest stable release ```bash docker run -ti -v $(pwd):/root/shared -w /root/shared --name=dolfinx-checkpoint ghcr.io/fenics/dolfinx/dolfinx:nightly ``` For the latest version compatible with nightly (with the ability to run the test suite), use ```bash export HDF5_MPI=ON export HDF5_DIR=/usr/local python3 -m pip install --no-binary-h5py --no-build-isolation io4dolfinx[test]@git+https://github.com/scientificcomputing/io4dolfinx@main ``` If you are using the `stable` image, you can install `io4dolfinx` from [PYPI](https://pypi.org/project/io4dolfinx/) with ```bash python3 -m pip install io4dolfinx[test] ``` This docker container can be opened with ```bash docker container start -i dolfinx-checkpoint ``` at a later instance ## Conda ```{note} Conda supports the stable release of DOLFINx, and thus the appropriate version should be installed, see the section above for more details. ``` Following is a minimal recipe of how to install io4dolfinx, given that you have conda installed on your system. ```bash conda create -n dolfinx-checkpoint python=3.12 conda activate dolfinx-checkpoint conda install -c conda-forge io4dolfinx ``` ```{note} Remember to download the appropriate version of `io4dolfinx` from Github [io4dolfinx: Releases](https://github.com/scientificcomputing/io4dolfinx/releases) ``` To run the test suite, you should also install `ipyparallel`, `pytest` and `coverage`, which can all be installed with conda ```bash conda install -c conda-forge ipyparallel pytest coverage ``` scientificcomputing-io4dolfinx-d21fc0e/docs/ipyparallel_intro.py000066400000000000000000000020111517634040500253550ustar00rootroot00000000000000# # Introduction to IPython parallel # The following demos heavily rely on {py:mod}`IPython-parallel` to illustrate how # checkpointing works when using multiple {py:mod}`MPI` processes. # We illustrate what happens in parallel by launching three MPI processes # using IPython-parallel's MPI engine support. import logging import ipyparallel as ipp def hello_mpi(): # We define all imports inside the function as they have to be launched on the remote engines from mpi4py import MPI print(f"Hello from rank {MPI.COMM_WORLD.rank}/{MPI.COMM_WORLD.size - 1}") with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster: # We send the query to run the function `hello_mpi` on all engines query: ipp.AsyncResult = cluster[:].apply_async(hello_mpi) # We wait for all engines to finish query.wait() # We check that all engines exited successfully assert query.successful(), query.error # We print the output from each engine print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/docs/logo.png000066400000000000000000000356451517634040500227430ustar00rootroot00000000000000PNG  IHDRgAMA a cHRMz&u0`:pQ<bKGDCtIME :IDATxw$UoٝvawY"(DA xU*bΈ0+b| PAəEKes&wj}uOWWz眓ìN/v3p8pp0%;^?S@W'm K9-oM^  tہv*(7AA!>"OQ0Yz8XakL$ppЖf333  QC[̬( kKXGk_vp.000?wXff5g0"k@k˖ .V /GY40՞7Q; CoK˶S,Prީf|`Ibffw(pfe KÁ )Eo^# =5lffV$⯠o=&cTxYvY slIg:4|.Y9Xx'bt6żxme; ->&#Ш虡bfffrejPZPYE9XC7cX膘Ymqlx`l(`ZFYmql틨JI1bY9X6#B7M1̬r,@ }v333- `*DvQ̬",8S%5B7je Kྼ @sFYplffffVe˒Ѝ033`B7[C7jeX l݈~X n/o]bfff݈> <5t[̬F8X6-({V{B7je7C7 Y8XMQzw@G膘Ymq\bm5m*FЍ033` 3t# *>333$@WnS5uafff=['om4333˖ @pl,C^#5O76˶Sx1h!J.V >9.333,[oLeJ2_.6e333 =G"@}8L/مQ~?p),%,[AgnV7_, 0.z<9BA(@{e333K֫|6%Q!dh,4ho'd333K/XvӁ?Ds&/B3lN4׈p`duKˁ:4)GBߥ2X.5|/`w`O`&0 GFW].-ubI`! 76CT~,1sTh=V {)Y6e) m88(~v[ N=mw@y+Xh9z1#0n@aѣ>>2/W[oC϶>`{B{^>脔D?&j{n;PpЎNZ)D/Uȫɤrljyi97CQ#s7~':FNpp0]h:+G`xO@t4=~g3gBQH/)mbree{ߖ8; % DFP?߾ o]c ݶ%U(+JAϡ`? }=(~.(KpjU( F$2ܑD݅nU,k& ׅqN'GhzGGu_M{ku ic6.*G)%.|3 wN󁻀{P{ =`6_DCYCo}v _.q>(p [ :ΕR,#.]wzrN@y?=^}ps=- \ n_~T=^iU3 e>OD$ Vq4o \OD8 mGsnCe]k.NAYJ\zxnTdK,] wGǡD`80#zmEekkP|#r 3PrřB:N5[GW?77w+2xt2-kSl^Cu xgB5iygcoPRE*{gc7ą@kO^Ѕߓ>f2/CvhD!NԴ" Q`|$ @,\#.(Q9Yq^){Q&5җq65'T)]R|2:+ڕ̘' J$gh.&2;VFڟ6E(h^)cBw!_Bѡ^zQ+sGeA.z?'G%jti!QAقA~(&4p/ K \>ۀU6[B%2Y$5FeL?~:xl÷rVv:k.ىҩ@?7uz5pVK_כ|KP)PAso:ʄAW6x&_'t6ʻ}%A3AVAOG=/{i5|8*6w8(a9N탆 "$h}G_'K{Ưj58q{=s١ʀ՜~f@9oc -<bw UܧeQ⸝WqۗcNm!e/2 b;peY<I߳0 RoE!K1E%ޮd/B,,硬y8l9~MITҼrʡzߗj>qЙIjW4j ʬ]jw.PUތ@9eM>Ͼȟ夃P 1}PiՕ.فM>J5ޠĉk e7мY>9T|:7t\LsDPљ@$>nY fi;&Qy-bs%!9`9Py6h_2WN`2whƤ1yd{u(PQm8Z 1bh{YESlF2wXlCha QhopQQ' ɀ;:4s'5zC g}e; m *d߈tBˈe$}˜Њu3?{@3;L ݨ;x͂й^ƉHN2Q{4CCYY*ݳ; "en9|mJ;%gU^%@2  McfV|3g3zpT)I1 9$T= sO9,4=bfZʟf;F1d@b!fi[aΉh o֔F:@NOFok"K**8O&'Pb }]N^e n1i> f6D{ÓmsZܧUkC7b1U'v+(Bg%HmNKgC?e{@`hX>U6H]Q1n,44* Toff%0M%K`}ƘՐ}]OD_txCu`KwGk Dq Tƙ?rnھ)h9ղm@[4h~kb4pN]mv̕(]෠UNAs~fUN _}yP#xJ<"9 L-XR Q4(?(ZmsZP}(8x;6+˷W_ŁIp3"4#ڷ:η?'~e eR:APiF!CY4fL.<YĂ]ݼP]G5As;:ǭ@iw`s>CQF( "Eu6F e6eP 6[|x mbr֒فn}x+]&u^br)iX"zkBoI`#j3Y;\ƌ2,z+ W-Pt⽟+b'4(y! w[.EBA_w 6Dc!v!#sDtOk^ڬ|rhqa7g$T" gf+Oaotm^]}ݾ ՜[1Tf(PGhVׅ(Cs=^-M7?@}Eg@e*h:.h$* e% c4m~(Psf,#\ Eg*?Ĥ'+ nFCOşjMhYq[70ZTR;eJSU*SԂfV#Ti!t1*h11Ё2XL=k$hɹnTDs;sPCu磌Ă=[0+Psbc B dn;p9:y tS:Ҡp󁟒hmĠ S,XAZ\>jy(0+t-棁uCK[fI` &mLY9UjDuikSDa43DҳAhz4Rqg6mCZ|HkGfeߡLJP棙20q3QA~fe;a)ZO]؎΄H# 7F(0osR0ZaX~; }C=Ϣz4wy .4mD[~ bYyF`R4}u!h C^uAVMV#sQIhĊ07ҁB֣ ݐ>b) Y]zYU45/fճ,X*j${Q?@`d ?ֶgi g*mQJj4aƄ&}K&1XV:*3@VЍkPF4m6Y&BfeQ'idVJ[F09hIՔbռMh,mo!݆D2!1佄wyͬ:!J 7vo$ꬫ-`6Kx 2Y eT ,rICHHfV+ B`yVJݢZV0;h*u:;/Ny1%|~Y6+g~z0),p p6i0[YV&J}y KPof [݆zτaVu5z,ث Zm*u>1M"PyXCaC|8Z30g7<Cun:<9-]h+ͪ:s(請hP(1_ׄ.dAcCu.wyg r mTqł嵤gQpc3 abe" 2ڼ(u&7߾ I_۬uͥn1tˣ,D1T4`7 K pj*,oEsɾި,4e}Q;7f˜dgl$|md{%hWR5CMbAdmYCAuK؎f/EkS վ^~ xi:tt4.|Xn%X؁V %ߧ>zCY `VhUhn~,[NB |ނv& =;3eAS]=Gۉݾuʬ jHO+Wqsxk38o@,3Q<#B*,Ǧ/;0WMc ~>Cϭh$&D #хo5h'4G`܄#ƢllT[%#QEUq5( -c; tZ}n߉no#-a9,Cyyw㭚UҰe@ eEEٸo^3˱gN|/EkPf6`!vlyHD񥐒u=\ߚY}p)KHR0@/"$Y0&zQVx皳VJTŇY_$(&ǠXPVن`95z_Ry_=uJ73ε>g/>̬%(9w:JTn1)tM3yWx=I  jf(L Y)InB3q +~mXց2˿"ڝ X='z7p9=ZЁ:PY_[F0Y+I>)p%p.s[{(30D4OUbsl1H$F'Xk/RaeyˀO`-xNRP^ tiYrY)$;|09tjр|}Q܁fO_びT/?@s\ (333+D<e?ml弥GEKJ p633N@`o <|x22t0?^]̪J6pv( GO6m[Ш9ot+̑fffUfaxT# Cp7Wev>P2*Qn> \L6$%м~U >0eX4aGدT*$։ZGף9~U,10eI=bTvJm!P_spʄ/f2Tc/Yuhyhe>hoBS_:u\AA+Q܎Vy%Zn?bـVR|:zlMQ2 ]jZ@y#x3/J6x?68`63^sfL;xշX'.@A(hȸ4.AJ dwV(=f2~v*ɸ6tC̬"J>@'L@Qݦ!F@GQyCh%hM=AT)n@7 Ǡ`0yT`rUWATffria2ʚuxߏJ* M;TIJ5@&4br^QH~T}}{vnU4znYUA04?eDd(k/xEe*(r\[(p *̻ $6~Oe28x NffP{2DcX,: xIFb 7*|e$b>oC2akG9ݝtꁓh3š2DvaD1J*nfZFZ 2,2H'*@ߎ?ͺ[( } |xF;{z2tC̊S Z mMne,4IAWwj9@b@ا>J7@姫ǣQ|0pHCB6:؉Q,YZ=𵫥w2ǢiЁ~} ; wr\U,K:}^Fv^vEYG{޹'r5 `1ԎV hFCo!ȡ󒧎ˎ:pTkVWH9f(>2Ci^F:=@G-4͢eP`vLE8m |͏ 8H)2xp}x=*w8x Q67%AB7,eѴY6 Z#X XEFs܇ C5?$ҍKI; Ͳ=Qe]i(rOEꝻ{P!h]HSP,61Bb b PRIJ!MؒͭhM ˉ+PW.s[ Z|H_Y"4/T<Ai08"ײaghEwʥAUQ'X(&s[;vhRFA,ffiz`Z 3ss4ܲcBhSr6%~ 7+Q,~B+h"\! nȅI"u/@5˖u='w%AAʇKYn |xE0'n#^̪RžBoA2 hP֢2̊Ő co`*%BMl߭JyHl: ͬ:m'zD22=v<923TFws^1؇@AˉҋYhs} 13 u2X63K p2\t9x3X+M7B F31`9VyҏGl2;;8:[hf֛Մ5 xdt2ʆ Vo;>%+r9Xmr9x0TyQ.U743Kt,fTi>xNi"p*o4d()sNu,,'}*p~ s=Z2TMofF;Э^xN@Zz`f^Ѹfg¬[x$qX8 \r0}˱->cb4KwIh)ay63`IFD根Ssci4/eWaM$``%s~Drfh:]?FgZNKGU|%fkO]Bb?$pZY)H-jy:8)t# ˉlt45)ݴ53QkJey2/ 5z̬g 8/q|N0H(94KfJy P XNt{o5yd6ZbTjAfV֢* FCeAךJ/ރjNGp[\`9~2؆P| &w"!®Գx<613+&vx:t{b_IV%?8Mz=Ze ܆bp8;`GHB<;\ ~J M1s7xҖ)MlolRpW6[o02|XKu*^-Vr(Ӏ7nH!b4F.5Q}CA@<mWW$4'>Бd%-h{uFff}@l[! ]oM9sP Mty\7oe->9`vGIGmsZSyoyS!ȑO+P|v[9߁C{ff4(9 g\2eARž_NQx2ǀ-iHB&]-/hBWBWiPFmwע[kz3x.)GWggG5fZ2FԙU,x K[7x -aNTo䷣z"? c 0: %[ƅ~V1+I u*0'aQ66*]~ eG"JBZЕbtȇK͛/ZLeh4TnAN,m!zDq(X- MM: MGYmوx(*0'⭙ь4|]91PNjDhd0FP`eۜ.4-(${ԡ#퍲#Bo"nAY3؎eeӀxB_ twN =GoYBFxwcЇN|+)i)73C'p3S1=Y%ʎ`ʃ m hx$TzƺOݘ!G2ͪk[e#ebVfP\c*}%0߈J0 ݡ?"cIT R,_fQmyn D-N&Z=ϬJuI LݘU_KxǢG+ۀ j@^hli6-eS> p DMn 3ݚ56qz0Ƴm1;P:ƕgAxFhXqw_vK0,[b#hv;fyP}@sCޮB],F- 0Z@2=!fwqFu#-NdYA[EH(YfP5g)](K#tCb Y',vȌ}UyhcK?mbfV"(Q3/tC2b!}l'Q˲% 82j?ʗfVUօnHʵ?A4!v^$9{Т2*8\ >tRztwqewnUizM+U4'UރgͬzMP a%r h?nSF@A..iPWɴ]SVY7r[plfU1"<;F"]bC M)4;m~ ZB'eqlf(ѷC+%!Vr' ֞ |]0Efm,} F4;єJe3n[|'f_A!>D|~-`qv XF~)p6h.BPbe3n]Tˀ`C7J+(ܯQf@K| A (J/e~4YZ* |cmnfV 8\ [v$9*7X]mAw>AJ.N-Ьwij ;MV&YPM%(xw[]aU&0j:TvtQF9 ? \F-}? MhEF,%֢>(P6މ4ZHdWS[qӨ>@[5 pۜZ ,uo6E6W1ff8lA% E_Du`(M*bP;.#M{.03tʨM{1*Mݔu42Kos$eW lE;7ffh緣B֠kX3_;M%}e e @ e(Af9)+g4hG擁݁ѡ7@,CӶSM"4O|8p#>'Զľ_!l`XvV}=mJ,n)oCwN*w[؂ '1_hZwQ\\lڦYVye+\~zDG煟 hCY5ݟ$beQ,d/vD%&FSí߇lrej;PV'h? 84ZNX\~xb@e'ϡ We*ݙ :QYIIF bZt,t.i&%AږQ|5:? ufk'u]!95?8g)ľZڇ|tQ?c(Leg*vEEЇb@fc+Qb*3Y朾 x2'= ڬNIz^s1* {uxr)P<{`':VDBrt.98ݽ L NԏEQ17FʟxoO vOF_xua>@S6GP|_wgW^/mAbT>;z;QQWօ:.4{^so36[%j iVyv%P&46 i>x RBv܋>Gw=HfeEֲtt#J0#о?01ބ(pq5c Gr-1(<uv(<wF.^oW[ݸux[Q(z,@E[O,ÁY(pCw.'-(7 G&#y>zxB( 'Q_.>4+{=0 J C>?Rо>؋x،bQP|:VgyOeW$pΫ=c 0uv `+;Y̪Y/c ,vE{q9TY(XNԕa^P<8|sl`doD27=|C5+>x3ZRzz97=f瓇.#l%tEXtdate:create2022-05-24T11:00:10+00:00x%tEXtdate:modify2022-05-24T11:00:10+00:00%~YtEXtSoftwareAdobe ImageReadyqe<IENDB`scientificcomputing-io4dolfinx-d21fc0e/docs/meshtags.py000066400000000000000000000066551517634040500234610ustar00rootroot00000000000000# # Writing MeshTags data to a checkpoint file # In many scenarios, the mesh used in a checkpoint is not trivial, and subdomains and sub-entities # have been tagged with appropriate markers. # As the mesh gets redistributed when read # (see [Writing Mesh Checkpoint](./writing_mesh_checkpoint)), # we need to store any tags together with this new mesh. # As an example we will use a unit-cube, where each entity has been tagged with a unique index. # + import logging from pathlib import Path from mpi4py import MPI import dolfinx import ipyparallel as ipp import numpy as np import io4dolfinx assert MPI.COMM_WORLD.size == 1, "This example should only be run with 1 MPI process" mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, nx=3, ny=4, nz=5) # - # We start by computing the unique global index of each (owned) entity in the mesh # as well as its corresponding midpoint entity_midpoints = {} meshtags = {} for i in range(mesh.topology.dim + 1): mesh.topology.create_entities(i) e_map = mesh.topology.index_map(i) # Compute midpoints of entities entities = np.arange(e_map.size_local, dtype=np.int32) mesh.topology.create_connectivity(i, mesh.topology.dim) entity_midpoints[i] = dolfinx.mesh.compute_midpoints(mesh, i, entities) # Associate each local index with its global index values = np.arange(e_map.size_local, dtype=np.int32) + e_map.local_range[0] meshtags[i] = dolfinx.mesh.meshtags(mesh, i, entities, values) # We use {py:func}`io4dolfinx.write_mesh` and `io4dolfinx.write_meshtags` to write the # {py:class}`dolfinx.mesh.Mesh` and {py:class}`dolfinx.mesh.MeshTags` to file. # We associate each meshtag with a name filename = Path("mesh_with_meshtags.bp") io4dolfinx.write_mesh(filename, mesh) for i, tag in meshtags.items(): io4dolfinx.write_meshtags(filename, mesh, tag, meshtag_name=f"meshtags_{i}") # Next we want to read the meshtags in on a different number of processes, # and check that the midpoints of each entity is still correct. # We do this with {py:func}`io4dolfinx.read_meshtags`. def verify_meshtags(filename: Path): # We assume that entity_midpoints have been sent to the engine from mpi4py import MPI import dolfinx import numpy as np import io4dolfinx read_mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD) prefix = f"{read_mesh.comm.rank + 1}/{read_mesh.comm.size}: " for i in range(read_mesh.topology.dim + 1): # Read mesh from file meshtags = io4dolfinx.read_meshtags(filename, read_mesh, meshtag_name=f"meshtags_{i}") # Compute midpoints for all local entities on process read_mesh.topology.create_connectivity(i, read_mesh.topology.dim) midpoints = dolfinx.mesh.compute_midpoints(read_mesh, i, meshtags.indices) # Compare locally computed midpoint with reference data for global_pos, midpoint in zip(meshtags.values, midpoints): np.testing.assert_allclose( entity_midpoints[i][global_pos], midpoint, err_msg=f"{prefix}: Midpoint ({i, global_pos}) do not match", ) print(f"{prefix} Matching of all entities of dimension {i} successful") with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster: cluster[:].push({"entity_midpoints": entity_midpoints}) query = cluster[:].apply_async(verify_meshtags, filename) query.wait() assert query.successful(), query.error print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/docs/migration_guide.md000066400000000000000000000045571517634040500247630ustar00rootroot00000000000000--- jupytext: formats: md:myst text_representation: extension: .md format_name: myst kernelspec: display_name: Python 3 language: python name: python3 --- # Migration Guide: Transitions to `io4dolfinx` This guide outlines the necessary steps to transition your code from {py:mod}`io4dolfinx` new API introduced in `io4dolfinx`. ## Major Changes The library has undergone a major refactor to support multiple IO backends. **{py:mod}`io4dolfinx` now supports both [ADIOS2](https://adios2.readthedocs.io/en/latest/) and [h5py](https://docs.h5py.org/en/stable/) backends.** This allows users to choose between the high-performance, adaptable ADIOS2 framework and the standard HDF5 format via {py:class}`h5py.File`, all using the same high-level API. It also opens the door for new backends in the future. ### Key API Updates 1. **Backend Agnosticism**: You can now switch between {py:class}`adios2.ADIOS` (default) and {py:class}`h5py.File` by passing a `backend` argument. 2. **Engine Configuration**: The explicit `engine` argument (e.g., "BP4", "HDF5", see ADIOS2 [Engine Types](https://adios2.readthedocs.io/en/latest/engines/engines.html)) has been removed from function signatures. It is now passed via a dictionary `backend_args`. This is because different backends may have different engine options. For example {py:class}`adios2.ADIOS` supports "BP4" and "HDF5", while `h5py` does not use engines. ### Example Transition #### Writing a Mesh `````{tab-set} ````{tab-item} Old API ```python import io4dolfinx io4dolfinx.write_mesh("mesh.bp", mesh, engine="BP4") ``` ```` ````{tab-item} New API (ADIOS2 Backend) ```python import io4dolfinx io4dolfinx.write_mesh("mesh.bp", mesh, backend="adios2", backend_args={"engine": "BP4"}) ``` ```` ````{tab-item} New API (h5py Backend) ```python import io4dolfinx io4dolfinx.write_mesh("mesh.bp", mesh, backend="h5py") ``` ```` ````` #### Writing a Function `````{tab-set} ````{tab-item} Old API ```python import io4dolfinx io4dolfinx.write_function("solution.bp", u, time=0.0, engine="BP4") ``` ```` ````{tab-item} New API (ADIOS2 Backend) ```python import io4dolfinx io4dolfinx.write_function("solution.bp", u, time=0.0, backend="adios2", backend_args={"engine": "BP4"}) ``` ```` ````{tab-item} New API (h5py Backend) ```python import io4dolfinx io4dolfinx.write_function("solution.bp", u, time=0.0, backend="h5py") ``` ```` ````` scientificcomputing-io4dolfinx-d21fc0e/docs/original_checkpoint.py000066400000000000000000000107151517634040500256510ustar00rootroot00000000000000# # Checkpoint on input mesh # As we have discussed earlier, one can choose to store function data in a way that # is N-to-M compatible by using {py:func}`io4dolfinx.write_checkpoint`. # This stores the distributed mesh in it's current (partitioned) ordering, and does # use the original input data ordering for the cells and connectivity. # This means that you cannot use your original mesh (from `.xdmf` files) or mesh tags # together with the checkpoint. The checkpoint has to store the mesh and associated # mesh-tags. # An optional way of store an N-to-M checkpoint is to store the function data in the same # ordering as the mesh. The write operation will be more expensive, as it requires data # communication to ensure contiguous data being written to the checkpoint. # The method is exposed as {py:func}`io4dolfinx.write_function_on_input_mesh`. # Below we will demonstrate this method. # + import logging from pathlib import Path from typing import Tuple import ipyparallel as ipp def locate_facets(x, tol=1.0e-12): return abs(x[0]) < tol def create_xdmf_mesh(filename: Path): from mpi4py import MPI import dolfinx mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim - 1, locate_facets) facet_tag = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, 1) with dolfinx.io.XDMFFile(MPI.COMM_WORLD, filename.with_suffix(".xdmf"), "w") as xdmf: xdmf.write_mesh(mesh) xdmf.write_meshtags(facet_tag, mesh.geometry) print(f"{mesh.comm.rank + 1}/{mesh.comm.size} Mesh written to {filename.with_suffix('.xdmf')}") mesh_file = Path("MyMesh.xdmf") with ipp.Cluster(engines="mpi", n=4, log_level=logging.ERROR) as cluster: # Create a mesh and write to XDMFFile cluster[:].push({"locate_facets": locate_facets}) query = cluster[:].apply_async(create_xdmf_mesh, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # - # Next, we will create a function on the mesh and write it to a checkpoint. # + def f(x): return (x[0] + x[1]) * (x[0] < 0.5), x[1], x[2] - x[1] def write_function( mesh_filename: Path, function_filename: Path, element: Tuple[str, int, Tuple[int,]] ): from mpi4py import MPI import dolfinx import io4dolfinx with dolfinx.io.XDMFFile(MPI.COMM_WORLD, mesh_filename, "r") as xdmf: mesh = xdmf.read_mesh() V = dolfinx.fem.functionspace(mesh, element) u = dolfinx.fem.Function(V) u.interpolate(f) io4dolfinx.write_function_on_input_mesh( function_filename.with_suffix(".bp"), u, mode=io4dolfinx.FileMode.write, time=0.0, name="Output", ) print( f"{mesh.comm.rank + 1}/{mesh.comm.size} Function written to ", f"{function_filename.with_suffix('.bp')}", ) # - # Read in mesh and write function to file element = ("DG", 4, (3,)) function_file = Path("MyFunction.bp") with ipp.Cluster(engines="mpi", n=2, log_level=logging.ERROR) as cluster: cluster[:].push({"f": f}) query = cluster[:].apply_async(write_function, mesh_file, function_file, element) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # Finally, we will read in the mesh from file and the function from the checkpoint # and compare it with the analytical solution. def verify_checkpoint( mesh_filename: Path, function_filename: Path, element: Tuple[str, int, Tuple[int,]] ): from mpi4py import MPI import dolfinx import numpy as np import io4dolfinx with dolfinx.io.XDMFFile(MPI.COMM_WORLD, mesh_filename, "r") as xdmf: in_mesh = xdmf.read_mesh() V = dolfinx.fem.functionspace(in_mesh, element) u_in = dolfinx.fem.Function(V) io4dolfinx.read_function(function_filename.with_suffix(".bp"), u_in, time=0.0, name="Output") # Compute exact interpolation u_ex = dolfinx.fem.Function(V) u_ex.interpolate(f) np.testing.assert_allclose(u_in.x.array, u_ex.x.array) print( "Successfully read checkpoint onto mesh on rank ", f"{in_mesh.comm.rank + 1}/{in_mesh.comm.size}", ) # Verify checkpoint by comparing to exact solution with ipp.Cluster(engines="mpi", n=5, log_level=logging.ERROR) as cluster: cluster[:].push({"f": f}) query = cluster[:].apply_async(verify_checkpoint, mesh_file, function_file, element) query.wait() assert query.successful(), query.error print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/docs/partitioned_mesh.py000066400000000000000000000071671517634040500252030ustar00rootroot00000000000000# # Storing mesh partition # This data is re-ordered when reading in a mesh, as the mesh is partitioned. # This means that when storing the mesh to disk from DOLFINx, the geometry and # connectivity arrays are re-ordered. # If we want to avoid to re-partition the mesh every time you run a simulation # (on a fixed number of processes), one can store the partitioning of the mesh # in the checkpoint. This is done by setting the flag `store_partition_info=True` # when calling {py:func}`io4dolfinx.write_mesh`. # + import logging from pathlib import Path import ipyparallel as ipp def write_partitioned_mesh(filename: Path): import subprocess from mpi4py import MPI import dolfinx import io4dolfinx # Create a simple unit square mesh mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.quadrilateral, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, ) # Write mesh checkpoint io4dolfinx.write_mesh( filename, mesh, backend="adios2", backend_args={"engine": "BP4"}, store_partition_info=True ) # Inspect checkpoint on rank 0 with `bpls` if mesh.comm.rank == 0: output = subprocess.run(["bpls", "-a", "-l", filename], capture_output=True) print(output.stdout.decode("utf-8")) # - # We inspect the partitioned mesh # + mesh_file = Path("partitioned_mesh.bp") n = 3 with ipp.Cluster(engines="mpi", n=n, log_level=logging.ERROR) as cluster: query = cluster[:].apply_async(write_partitioned_mesh, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # - # # Reading a partitioned mesh # If we try to read the mesh in on a different number of processes, we will get an error. # We illustrate this below, by first trying to read the mesh using partitioning information, # which is done by setting the flag `read_from_partition=True` when calling # {py:func}`io4dolfinx.read_mesh`. def read_partitioned_mesh(filename: Path, read_from_partition: bool = True): from mpi4py import MPI import io4dolfinx prefix = f"{MPI.COMM_WORLD.rank + 1}/{MPI.COMM_WORLD.size}: " try: mesh = io4dolfinx.read_mesh( filename, comm=MPI.COMM_WORLD, backend="adios2", backend_args={"engine": "BP4"}, read_from_partition=read_from_partition, ) print(f"{prefix} Mesh: {mesh.name} read successfully with {read_from_partition=}") except ValueError as e: print(f"{prefix} Caught exception: ", e) with ipp.Cluster(engines="mpi", n=n + 1, log_level=logging.ERROR) as cluster: # Read mesh from file with different number of processes query = cluster[:].apply_async(read_partitioned_mesh, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # Read mesh from file with different number of processes (not using partitioning information). # If we instead turn of `read_from_partition`, we can read the mesh on a # different number of processes. with ipp.Cluster(engines="mpi", n=n + 1, log_level=logging.ERROR) as cluster: query = cluster[:].apply_async(read_partitioned_mesh, mesh_file, False) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # Read mesh from file with same number of processes as was written, # re-using partitioning information. with ipp.Cluster(engines="mpi", n=n, log_level=logging.ERROR) as cluster: query = cluster[:].apply_async(read_partitioned_mesh, mesh_file, True) query.wait() assert query.successful(), query.error print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/docs/quickstart.md000066400000000000000000000116571517634040500240060ustar00rootroot00000000000000# Quick Start Guide This document provides a quick start overview of the functions available in `io4dolfinx`. The library is designed to extend DOLFINx with advanced Input/Output capabilities, focusing on flexible checkpointing and support for various data formats. ## Core Checkpointing The primary purpose of `io4dolfinx` is to support **N-to-M checkpointing**. This means you can run a simulation on $N$ processes, save the state, and restart the simulation on $M$ processes. ### Meshes Before storing any functions, the mesh must be written to the checkpoint file. The mesh topology and geometry are saved in a distributed format. ```python from mpi4py import MPI import dolfinx import io4dolfinx ``` ```python comm = MPI.COMM_WORLD mesh = dolfinx.mesh.create_unit_square(comm, 10, 10) filename = "checkpoint.bp" ``` Write mesh to file ```python # io4dolfinx.write_mesh(filename, mesh) ``` Read mesh from file. The mesh is redistributed according to the current communicator size. ```python mesh_new = io4dolfinx.read_mesh(filename, comm) ``` ### Functions Functions can be stored associated with a timestamp. They effectively store the coefficients of the finite element function. ```python V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) u = dolfinx.fem.Function(V) u.name = "my_solution" ``` Write function ```python io4dolfinx.write_function(filename, u, time=0.5) ``` Read function ```{note} You must read the mesh first (or have a compatible mesh ready, see [Checkpoint on input mesh](./original_checkpoint.py) for details). ```` ```python u_new = dolfinx.fem.Function(V) io4dolfinx.read_function(filename, u_new, time=0.5, name="my_solution") ``` ## Mesh Tags and Data `io4dolfinx` supports storing auxiliary data associated with the mesh, such as subdomain markers (`MeshTags`) or raw data arrays. ### MeshTags MeshTags (markers for cells, facets, etc.) can be written to the same checkpoint file as the mesh. They are re-distributed correctly when reading back on a different number of processes. Create some dummy tags ```python subdomains = dolfinx.mesh.meshtags(mesh, mesh.topology.dim, [0], [1]) ``` Write tags ```python io4dolfinx.write_meshtags(filename, mesh, subdomains, meshtag_name="subdomains") ``` Read tags ```python tags = io4dolfinx.read_meshtags(filename, mesh, meshtag_name="subdomains") ``` ## Advanced Checkpointing Strategies Beyond standard N-to-M checkpointing, the library offers specialized strategies for specific use cases. ### Snapshot Checkpointing A **snapshot** is a lightweight checkpoint intended for use within the *same* simulation run (N-to-N). It is ideal for temporary storage (e.g., for an adjoint solver or saving state before a risky operation) where you know the process count and mesh partitioning will not change. ```python snapshot_file = "temp_snapshot.bp" io4dolfinx.snapshot_checkpoint(u, snapshot_file, io4dolfinx.FileMode.write) ``` Read back (must be on same mesh distribution) ```python io4dolfinx.snapshot_checkpoint(u, snapshot_file, io4dolfinx.FileMode.read) ``` See the [Snapshot Checkpointing Guide](./snapshot_checkpoint.py) for more details and examples. ### Original Mesh Checkpointing Sometimes you want to save a solution that corresponds exactly to the input mesh file (e.g., an `.xdmf` file you started with), rather than the current partitioned mesh. This is useful for visualization or post-processing on the original geometry. ```python io4dolfinx.write_function_on_input_mesh("solution_on_input.bp", u) ``` See the [Checkpoint on input mesh](./original_checkpoint.py) for more details and examples. ## Legacy DOLFIN Support The library provides readers for migrating data from legacy DOLFIN (FEniCS). * **`read_mesh_from_legacy_h5`**: Reads a mesh from a legacy HDF5 file. * **`read_function_from_legacy_h5`**: Reads a function from a legacy HDF5 file (supports both `HDF5File` and `XDMFFile` archives). See the [reading_legacy_data.md](./reading_legacy_data.md) guide for detailed examples. ## Metadata and Utilities Helper functions are available to query the contents of a checkpoint file. * **`read_function_names`**: Returns a list of all functions stored in a file. * **`read_timestamps`**: Returns the time steps available for a specific function. * **`read_attributes` / `write_attributes`**: Allows storing arbitrary metadata dictionaries. ## Backends `io4dolfinx` is backend-agnostic. You can choose the storage engine by passing the `backend` argument to most functions. 1. **`adios2` (Default)**: Uses the ADIOS2 library. Best for large-scale parallel IO. Supports engines like "BP4", "BP5", and "HDF5". 2. **`h5py`**: Uses the standard HDF5 library via `h5py`. Requires an MPI-enabled HDF5 build. Good for compatibility with other HDF5 tools. 3. **`vtkhdf`**: Supports reading and writing the VTKHDF format (scalable VTK). 4. **`pyvista`**: Primarily for reading unstructured grids (`.vtu`) via PyVista/meshio. 5. **`xdmf`**: Basic support for reading XDMF data. scientificcomputing-io4dolfinx-d21fc0e/docs/reading_legacy_data.md000066400000000000000000000073131517634040500255340ustar00rootroot00000000000000# Reading Legacy DOLFIN Data `io4dolfinx` provides functionality to read meshes and functions created with the legacy version of DOLFIN (often referred to as "Old FEniCS"). This allows users to migrate data from older simulations to DOLFINx. ## Supported Formats The library supports reading data stored in the **HDF5** format used by legacy DOLFIN. This includes: * Meshes stored using `dolfin.HDF5File`. * Functions stored using `dolfin.HDF5File`. * Checkpoints stored using `dolfin.XDMFFile` (which produces an `.h5` file alongside the `.xdmf` file). **Note:** The `.xml` or `.xml.gz` formats are not supported. You must convert these to HDF5 using legacy DOLFIN before reading them with `io4dolfinx`. ## Reading Meshes To read a mesh, you must know the **group name** where the mesh is stored within the HDF5 file. In legacy DOLFIN, this was often `/mesh` or the name you provided to the write function. ```python from mpi4py import MPI import io4dolfinx ``` ```python comm = MPI.COMM_WORLD filename = "legacy_mesh.h5" ``` ### Read the mesh Here 'group' is the path to the mesh inside the HDF5 file (e.g., "/mesh") ```python mesh = io4dolfinx.read_mesh_from_legacy_h5( filename=filename, comm=comm, group="/mesh", ) ``` ```python print(f"Read mesh with topology dimension: {mesh.topology.dim}") ``` ## Reading Functions Reading a function requires you to first create a compatible **FunctionSpace** and **Function** in DOLFINx. The data is then read from the file and loaded into this function. ### Simple Function Read If the function was saved directly using `HDF5File.write(u, "name")`: ```python import dolfinx import io4dolfinx from mpi4py import MPI ``` #### Read the mesh first ```python comm = MPI.COMM_WORLD filename = "legacy_data.h5" mesh = io4dolfinx.read_mesh_from_legacy_h5(filename, comm, group="/mesh") ``` #### Create the appropriate FunctionSpace You must know the element family and degree used in the legacy simulation ```python V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) ``` #### Create the function to hold the data ```python u = dolfinx.fem.Function(V) ``` #### Read the data Here 'group' corresponds to the name given when writing in legacy DOLFIN ```python io4dolfinx.read_function_from_legacy_h5(filename=filename, comm=comm, u=u, group="v") ``` ### Reading from XDMF Checkpoints If the data was saved using `XDMFFile.write_checkpoint`, the HDF5 structure is slightly different (often containing time steps). You can specify the `step` argument to read a specific snapshot. Assuming mesh and function space V are already created as above ```python u_checkpoint = dolfinx.fem.Function(V) ``` Read the first time step (step=0). Note that the filename should be the .h5 file generated by the XDMFFile ```python io4dolfinx.read_function_from_legacy_h5( filename="checkpoint.h5", comm=comm, u=u_checkpoint, group="v", # The name of the function in the checkpoint step=0, ) ``` ## Reading Vector Functions Vector functions are read in the same way, provided the `FunctionSpace` is initialized correctly with the vector shape. ```python W = dolfinx.fem.functionspace(mesh, ("Lagrange", 1, (mesh.geometry.dim,))) w = dolfinx.fem.Function(W) ``` ```python io4dolfinx.read_function_from_legacy_h5( filename="legacy_vector.h5", comm=comm, u=w, group="velocity_function" ) ``` ## Limitations * **Function Types:** Only `Lagrange` (Continuous Galerkin) and `DG` (Discontinuous Galerkin) functions are supported. * **One Checkpoint per File:** The legacy reader generally expects one checkpoint series per file structure for XDMF checkpoints. * **Backends:** The default backend is `adios2`. Ensure you have an MPI-enabled build of HDF5 if using the `h5py` backend explicitly. scientificcomputing-io4dolfinx-d21fc0e/docs/snapshot_checkpoint.py000066400000000000000000000051401517634040500257000ustar00rootroot00000000000000# # Snapshot checkpoint (non-persistent) # The checkpoint method described in [Writing function checkpoints](./writing_functions_checkpoint) # are *N-to-M*, meaning that you can write them out on N-processes and read them in on M processes. # # As discussed in that chapter, these checkpoints need to be associated with a mesh. # This is because the function is defined on a specific function space, which in turn is # defined on a specific mesh. # # However, there are certain scenarios where you simply want to store a checkpoint associated # with the current mesh, that should only be possible to use during this simulation. # An example use-case is when running an iterative solver, and wanting a fall-back mechanism that # does not require extra RAM. # In this example, we will demonstrate how to write a # {py:func}`snapshot checkpoint` to disk. # First we define a {py:class}`function` `f` that we want to represent in # the {py:class}`function space`. # + import logging from pathlib import Path import ipyparallel as ipp def f(x): import numpy as np return np.sin(x[0]) + 0.1 * x[1] # - # Next, we create a mesh and an appropriate function space and read and write from file. # Note that for both these operations, we use {py:func}`io4dolfinx.snapshot_checkpoint`, # with different read and write modes. def read_write_snapshot(filename: Path): from mpi4py import MPI import dolfinx import numpy as np import io4dolfinx mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 3, 7, 4) V = dolfinx.fem.functionspace(mesh, ("Lagrange", 5)) u = dolfinx.fem.Function(V) u.interpolate(f) u.name = "Current_solution" # Next, we store the solution to file io4dolfinx.snapshot_checkpoint(u, filename, io4dolfinx.FileMode.write) # Next, we create a new function and load the solution into it u_new = dolfinx.fem.Function(V) u_new.name = "Read_solution" io4dolfinx.snapshot_checkpoint(u_new, filename, io4dolfinx.FileMode.read) # Next, we verify that the solution is correct np.testing.assert_allclose(u_new.x.array, u.x.array, atol=np.finfo(float).eps) print(f"{MPI.COMM_WORLD.rank + 1}/{MPI.COMM_WORLD.size}: Successfully wrote and read snapshot") # + mesh_file = Path("snapshot.bp") with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster: cluster[:].push({"f": f}) query = cluster[:].apply_async( read_write_snapshot, mesh_file, ) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # - scientificcomputing-io4dolfinx-d21fc0e/docs/testing.md000066400000000000000000000074071517634040500232670ustar00rootroot00000000000000 # Testing Strategy This document outlines how `io4dolfinx` is tested, covering both local development testing and the Continuous Integration (CI) process on GitHub Actions. ## Coverage reports: You will find the updated coverage reports for the latest version on `main` tested against `stable` and `nightly` versions of dolfinx at the following links: - [Coverage report for stable](https://scientificcomputing.github.io/io4dolfinx/code-coverage-report-stable/) - [Coverage report for nightly](https://scientificcomputing.github.io/io4dolfinx/code-coverage-report-nightly/) ## Local Testing The library uses `pytest` for testing. To execute the tests locally, you first need to install the library and its dependencies. ### Installation for Testing Install the library with the optional `test` dependencies to ensure you have packages like `pytest`, `coverage`, and `ipyparallel` ```bash python3 -m pip install ".[test]" ``` ### Running Tests To execute all tests in the repository, run: ```bash python3 -m pytest . ``` ### Generating Test Data Some tests require specific datasets to verify compatibility with older software versions. #### Testing against data from legacy dolfin Some tests check the capability of reading data created with the legacy version of DOLFIN. To create this dataset, start a docker container with legacy DOLFIN, for instance: ```bash docker run -ti -v $(pwd):/root/shared -w /root/shared --rm ghcr.io/scientificcomputing/fenics:2024-02-19 ``` Then, inside this container, call: ```bash python3 ./tests/create_legacy_data.py --output-dir=legacy ``` #### Testing against data from older versions of adios4dolfinx Some tests check the capability to read data generated by `adios4dolfinx<0.7.2`. To generate data for these tests use the following commands: ```bash docker run -ti -v $(pwd):/root/shared -w /root/shared --rm ghcr.io/fenics/dolfinx/dolfinx:v0.7.3 ``` Then, inside the container, call: ```bash python3 -m pip install adios4dolfinx==0.7.1 python3 ./tests/create_legacy_checkpoint.py --output-dir=legacy_checkpoint ``` --- ## Continuous Integration (GitHub Actions) The repository relies on several GitHub Actions workflows to ensure code quality and compatibility across different environments. ### 1. Main Test Suite (`test_package.yml`) This is the primary workflow triggered on pushes to `main`, pull requests, and scheduled nightly runs. It runs on `ubuntu-24.04` using the official DOLFINx docker container . **Workflow Steps:** 1. **Linting & Formatting:** Checks code style using `ruff` and type consistency with `mypy` . 2. **Data Generation:** * Creates legacy DOLFIN data using the `create_legacy_data.yml` workflow. * Creates legacy `adios4dolfinx` checkpoints using `create_legacy_checkpoint.yml`. 3. **Test Execution:** * Installs the package with MPI-enabled `h5py`. * Runs the standard test suite with `coverage`. * Runs parallel tests using `mpirun -n 4 ... mpi4py -m pytest`. 4. **Reporting:** Combines coverage reports and uploads them as artifacts. ### 2. Compatibility Testing To ensure broad support, specific workflows test against different configurations: * **MPI & ADIOS2 Versions (`test_package_openmpi.yml`):** * Tests against both `openmpi` and `mpich` implementations using the `ghcr.io/fenics/test-env` containers. * Verifies compatibility with different ADIOS2 versions (e.g., `v2.10.2`, `v2.11.0`) . * **Operating System (`test_redhat.yml`):** * Runs the full test suite inside a RedHat-based container (`docker.io/fenicsproject/test-env:current-redhat`) to guarantee functionality on non-Debian systems . ### 3. Documentation (`build_docs.yml`) Ensures the documentation builds correctly with `jupyter-book` on every push and pull request, preventing documentation regressions.scientificcomputing-io4dolfinx-d21fc0e/docs/time_dependent_mesh.py000066400000000000000000000062461517634040500256420ustar00rootroot00000000000000# # Time-dependent mesh checkpoints # As we have seen in the previous examples, we store information about the connectivity, # the coordinates of the mesh nodes, # as well as a reference element. Note that the only thing that can change for a mesh # during a simulation are the coordinate of the mesh nodes. # In the following example, we will demonstrate how to write a time-dependent mesh # checkpoint to disk. # First, we create a simple function to compute the volume of a mesh # + import logging from pathlib import Path from mpi4py import MPI import ipyparallel as ipp import io4dolfinx def compute_volume(mesh, time_stamp): from mpi4py import MPI import dolfinx import ufl # Compute the volume of the mesh vol_form = dolfinx.fem.form(1 * ufl.dx(domain=mesh)) vol_local = dolfinx.fem.assemble_scalar(vol_form) vol_glob = mesh.comm.allreduce(vol_local, op=MPI.SUM) if mesh.comm.rank == 0: print(f"{mesh.comm.rank + 1}/{mesh.comm.size} Time: {time_stamp} Mesh Volume: {vol_glob}") def write_meshes(filename: Path): from mpi4py import MPI import dolfinx import numpy as np import io4dolfinx # Create a unit cube mesh = dolfinx.mesh.create_unit_cube( MPI.COMM_WORLD, 3, 6, 5, cell_type=dolfinx.mesh.CellType.hexahedron, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, ) # Write mesh to file, associated with time stamp 1.5 io4dolfinx.write_mesh(filename, mesh, time=1.5) compute_volume(mesh, 1.5) mesh.geometry.x[:, 0] += 0.1 * mesh.geometry.x[:, 0] mesh.geometry.x[:, 1] += 0.3 * mesh.geometry.x[:, 1] * np.sin(mesh.geometry.x[:, 2]) compute_volume(mesh, 3.3) # Write mesh to file, associated with time stamp 3.3 # Note that we set the mode to append, as we have already created the file # and we do not want to overwrite the existing data io4dolfinx.write_mesh(filename, mesh, time=3.3, mode=io4dolfinx.FileMode.append) # - # We write the sequence of meshes to file # + mesh_file = Path("timedep_mesh.bp") n = 3 with ipp.Cluster(engines="mpi", n=n, log_level=logging.ERROR) as cluster: # Write mesh to file cluster[:].push({"compute_volume": compute_volume}) query = cluster[:].apply_async(write_meshes, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # - # # Reading a time dependent mesh # The only thing we need to do to read the mesh is to send in the associated time stamp, # which we do by adding `time=time_stamp` when calling {py:func}`io4dolfinx.read_mesh`. second_mesh = io4dolfinx.read_mesh( mesh_file, comm=MPI.COMM_WORLD, backend="adios2", backend_args={"engine": "BP4"}, time=3.3 ) compute_volume(second_mesh, 3.3) first_mesh = io4dolfinx.read_mesh( mesh_file, comm=MPI.COMM_WORLD, backend="adios2", backend_args={"engine": "BP4"}, time=1.5 ) compute_volume(first_mesh, 1.5) # We observe that the volume of the mesh has changed, as we have perturbed the mesh # between the two time stamps. # We also note that we can read the meshes in on a different number of processes than # we wrote them with and in a different order (as long as the time stamps are correct). scientificcomputing-io4dolfinx-d21fc0e/docs/writing_functions_checkpoint.py000066400000000000000000000056721517634040500276260ustar00rootroot00000000000000# # Writing a function checkpoint # In the previous sections, we have gone in to quite some detail as to how # to store meshes with io4dolfinx. # This section will explain how to store {py:class}`functions`, # and how to read them back in. # We start by creating a {py:class}`mesh` # + import logging from pathlib import Path from mpi4py import MPI import dolfinx import ipyparallel as ipp import io4dolfinx assert MPI.COMM_WORLD.size == 1, "This example should only be run with 1 MPI process" mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, nx=10, ny=10, cell_type=dolfinx.cpp.mesh.CellType.quadrilateral ) # - # Next, we create a function, and interpolate a polynomial function into the function space # + el = "N1curl" degree = 3 V = dolfinx.fem.functionspace(mesh, (el, degree)) def f(x): return -(x[1] ** 2), x[0] - 2 * x[1] u = dolfinx.fem.Function(V) u.interpolate(f) # - # For the checkpointing, we start by storing the mesh to file filename = Path("function_checkpoint.bp") io4dolfinx.write_mesh(filename, mesh) # Next, we store the function to file, and associate it with a name. # Note that we can also associate a time stamp with it, as done for meshes in # [Writing time-dependent mesh checkpoint](./time_dependent_mesh). # We use {py:func}`io4dolfinx.write_function` for this. io4dolfinx.write_function(filename, u, time=0.3, name="my_curl_function") # Next, we want to read the function back in (using multiple MPI processes) # and check that the function is correct. # We use {py:func}`io4dolfinx.read_function` for this. # # ```{admonition} What mesh to use? # Note that we have read in the mesh with {py:func}`io4dolfinx.read_mesh` # before reading in the function. # We **cannot** use {py:func}`io4dolfinx.read_function` to read function data to the # original `mesh` object. # To do this, we need to use {py:func}`io4dolfinx.write_function_on_input_mesh`, see # [Writing function on input mesh checkpoint](./original_checkpoint). # for more details. # ``` def read_function(filename: Path, timestamp: float): from mpi4py import MPI import dolfinx import numpy as np import io4dolfinx in_mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD) W = dolfinx.fem.functionspace(in_mesh, (el, degree)) u_ref = dolfinx.fem.Function(W) u_ref.interpolate(f) u_in = dolfinx.fem.Function(W) io4dolfinx.read_function(filename, u_in, time=timestamp, name="my_curl_function") np.testing.assert_allclose(u_ref.x.array, u_in.x.array, atol=1e-14) print( f"{MPI.COMM_WORLD.rank + 1}/{MPI.COMM_WORLD.size}: ", f"Function read in correctly at time {timestamp}", ) with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster: cluster[:].push({"f": f, "el": el, "degree": degree}) query = cluster[:].apply_async(read_function, filename, 0.3) query.wait() assert query.successful(), query.error print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/docs/writing_mesh_checkpoint.py000066400000000000000000000117621517634040500265470ustar00rootroot00000000000000# # Writing a mesh checkpoint # # In this example, we will demonstrate how to write a mesh checkpoint to disk. # # We start by creating a simple {py:func}`unit-square mesh`. # + import logging from pathlib import Path from mpi4py import MPI import dolfinx import ipyparallel as ipp mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) # - # Note that when a mesh is created in DOLFINx, we send in a # {py:class}`MPI communicator`. # The communicator is used to partition (distribute) the mesh across the available processes. # This means that each process only have access to a sub-set of cells and nodes of the mesh. # We can inspect these with the following commands: def print_mesh_info(mesh: dolfinx.mesh.Mesh): cell_map = mesh.topology.index_map(mesh.topology.dim) node_map = mesh.geometry.index_map() print( f"Rank {mesh.comm.rank}: number of owned cells {cell_map.size_local}", f", number of ghosted cells {cell_map.num_ghosts}\n", f"Number of owned nodes {node_map.size_local}", f", number of ghosted nodes {node_map.num_ghosts}", ) print_mesh_info(mesh) # ## Create a distributed mesh # Next, we can use {py:mod}`IPython parallel` to inspect a partitioned # {py:class}`mesh`. # We create a convenience function for creating a mesh that # {py:attr}`shares cells` on the boundary # between two processes if `ghosted=True`. def create_distributed_mesh(ghosted: bool, N: int = 10): """ Create a distributed mesh with N x N cells. Share cells on process boundaries if ghosted is set to True """ from mpi4py import MPI import dolfinx ghost_mode = dolfinx.mesh.GhostMode.shared_facet if ghosted else dolfinx.mesh.GhostMode.none mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, N, N, ghost_mode=ghost_mode) print(f"{ghost_mode=}") print_mesh_info(mesh) # Next we start up a new cluster with three engines. # As we defined `print_mesh_info` locally on this process, we need to push it to all engines. with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster: # Push print_mesh_info to all engines cluster[:].push({"print_mesh_info": print_mesh_info}) # Create mesh with ghosted cells query_true = cluster[:].apply_async(create_distributed_mesh, True) query_true.wait() assert query_true.successful(), query_true.error print("".join(query_true.stdout)) # Create mesh without ghosted cells query_false = cluster[:].apply_async(create_distributed_mesh, False) query_false.wait() assert query_false.successful(), query_false.error print("".join(query_false.stdout)) # ## Writing a mesh checkpoint # The input data to a mesh is: # - A geometry: the set of points in R^D that are part of each cell # - A two-dimensional connectivity array: A list that indicates which nodes of the geometry # is part of each cell # - A {py:func}`reference element`: Used for push data back and # forth from the reference element and computing Jacobians # We now use {py:mod}`io4dolfinx` to write a mesh to file. def write_mesh(filename: Path): import subprocess from mpi4py import MPI import dolfinx import io4dolfinx # Create a simple unit square mesh mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.quadrilateral ) # Write mesh checkpoint io4dolfinx.write_mesh(filename, mesh, backend="adios2", backend_args={"engine": "BP4"}) # Inspect checkpoint on rank 0 with `bpls` if mesh.comm.rank == 0: output = subprocess.run(["bpls", "-a", "-l", str(filename.absolute())], capture_output=True) print(output.stdout.decode("utf-8")) # + mesh_file = Path("mesh.bp") with ipp.Cluster(engines="mpi", n=2, log_level=logging.ERROR) as cluster: # Write mesh to file query = cluster[:].apply_async(write_mesh, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) # - # We observe that we have stored all the data needed to re-create the mesh in the file `mesh.bp`. # We can therefore read it (to any number of processes) with {py:func}`io4dolfinx.read_mesh` def read_mesh(filename: Path): from mpi4py import MPI import dolfinx import io4dolfinx mesh = io4dolfinx.read_mesh( filename, comm=MPI.COMM_WORLD, backend="adios2", backend_args={"engine": "BP4"}, ghost_mode=dolfinx.mesh.GhostMode.none, ) print_mesh_info(mesh) # ## Reading mesh checkpoints (N-to-M) # We can now read the checkpoint on a different number of processes than we wrote it on. with ipp.Cluster(engines="mpi", n=4, log_level=logging.ERROR) as cluster: # Write mesh to file cluster[:].push({"print_mesh_info": print_mesh_info}) query = cluster[:].apply_async(read_mesh, mesh_file) query.wait() assert query.successful(), query.error print("".join(query.stdout)) scientificcomputing-io4dolfinx-d21fc0e/pyproject.toml000066400000000000000000000044371517634040500232540ustar00rootroot00000000000000[build-system] # Require setuptool version due to https://github.com/pypa/setuptools/issues/2938 requires = ["setuptools>=61.0.0", "wheel"] [project] name = "io4dolfinx" version = "1.2.0" description = "Checkpointing functionality for DOLFINx meshes/functions with ADIOS2" authors = [ { name = "Jørgen S. Dokken", email = "dokken@simula.no" }, { name = "Henrik N.T. Finsberg", email = "henriknf@simula.no" } ] license = { file = "LICENSE" } readme = "README.md" dependencies = ["fenics-dolfinx>=0.10.0", "packaging"] [project.optional-dependencies] test = [ "pytest", "coverage", "ipyparallel", "io4dolfinx[h5py]", "io4dolfinx[pyvista]", "io4dolfinx[xdmf]" ] dev = ["pdbpp", "ipython", "mypy", "ruff"] h5py = ["h5py"] pyvista = ["pyvista"] xdmf = ["h5py"] exodus = ["netcdf4"] docs = [ "jupyter-book<2.0.0", "ipyparallel", "ipywidgets", "jupytext", "ipykernel<7.0.0", # Note: Remove once https://github.com/ipython/ipykernel/issues/1450 is in a release "sphinx-codeautolink", "io4dolfinx[h5py]", "sphinx_external_toc<1.1.0", ] all = ["io4dolfinx[test,dev,docs]"] [tool.pytest.ini_options] addopts = ["--import-mode=importlib"] testpaths = ["tests"] [tool.mypy] ignore_missing_imports = true # Folders to exclude exclude = ["docs/", "build/"] # Folder to check with mypy files = ["src", "tests"] [tool.ruff] src = ["src", "tests", "docs"] line-length = 100 indent-width = 4 [tool.ruff.lint] select = [ # Pyflakes "F", # Pycodestyle "E", "W", # isort "I001", ] [tool.ruff.lint.isort] known-first-party = ["io4dolfinx"] known-third-party = [ "basix", "dolfinx", "ffcx", "ufl", "gmsh", "numpy", "pytest", ] section-order = [ "future", "standard-library", "mpi", "third-party", "first-party", "local-folder", ] [tool.ruff.lint.isort.sections] "mpi" = ["mpi4py", "petsc4py"] [tool.bumpversion] allow_dirty = false commit = true message = "Bump version: {current_version} → {new_version}" tag = true sign_tags = false tag_name = "v{new_version}" tag_message = "Bump version: {current_version} → {new_version}" current_version = "1.2.0" [[tool.bumpversion.files]] filename = "pyproject.toml" search = 'version = "{current_version}"' replace = 'version = "{new_version}"' scientificcomputing-io4dolfinx-d21fc0e/src/000077500000000000000000000000001517634040500211175ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/000077500000000000000000000000001517634040500231765ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/__init__.py000066400000000000000000000030411517634040500253050ustar00rootroot00000000000000# Copyright (C) 2023 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT """Top-level package for ADIOS2Wrappers.""" from importlib.metadata import metadata from .backends import FileMode, get_backend from .checkpointing import ( read_attributes, read_function, read_function_names, read_mesh, read_meshtags, read_timestamps, write_attributes, write_cell_data, write_function, write_mesh, write_meshtags, write_point_data, ) from .original_checkpoint import write_function_on_input_mesh, write_mesh_input_order from .readers import ( read_cell_data, read_function_from_legacy_h5, read_mesh_from_legacy_h5, read_point_data, ) from .snapshot import snapshot_checkpoint from .utils import reconstruct_mesh meta = metadata("io4dolfinx") __version__ = meta["Version"] __author__ = meta.get("Author", "") __license__ = meta["License"] __email__ = meta["Author-email"] __program_name__ = meta["Name"] __all__ = [ "FileMode", "write_meshtags", "read_meshtags", "read_cell_data", "read_mesh", "write_mesh", "read_function_from_legacy_h5", "read_mesh_from_legacy_h5", "write_function", "read_function", "snapshot_checkpoint", "write_function_on_input_mesh", "write_mesh_input_order", "write_attributes", "write_data", "read_attributes", "read_function_names", "read_point_data", "read_timestamps", "get_backend", "write_cell_data", "write_point_data", "reconstruct_mesh", ] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/000077500000000000000000000000001517634040500247505ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/__init__.py000066400000000000000000000320561517634040500270670ustar00rootroot00000000000000from enum import Enum from importlib import import_module from pathlib import Path from typing import Any, Protocol from mpi4py import MPI import dolfinx import numpy as np import numpy.typing as npt from ..structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData __all__ = ["FileMode", "IOBackend", "get_backend"] class ReadMode(Enum): serial = 10 # This means that all data is read in on root rank # Total number of data P, num processes = i + 1. # All processes reads at least `P // (i+1)` items # The first j=P%(i+1) processes reads `P // (i+1) + 1` items # ```python # def compute_partitioning(P, J): # min_num = P // J # num_per_proc = np.full(J, min_num) # rem = P % J # num_per_proc[:int(rem)] += 1 # assert(sum(num_per_proc)) == P # return num_per_proc # ``` parallel = 20 class FileMode(Enum): """Filen mode used for opening files.""" append = 10 #: Append data to file write = 20 #: Write data to file read = 30 #: Read data from file # See https://peps.python.org/pep-0544/#modules-as-implementations-of-protocols class IOBackend(Protocol): read_mode: ReadMode def get_default_backend_args(self, arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default backend arguments given a set of input arguments. Args: arguments: Input backend arguments Returns: Updated backend arguments """ def write_attributes( self, filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None, ): """Write attributes to file. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attribute group attributes: Dictionary of attributes to write backend_args: Arguments to backend """ def read_attributes( self, filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dict[str, Any]: """Read attributes from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attribute group backend_args: Arguments to backend Returns: Dictionary of attributes read from file """ def read_timestamps( self, filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read timestamps from file. Args: filename: Path to file to read from comm: MPI communicator used in storage function_name: Name of the function to read timestamps for backend_args: Arguments to backend Returns: Numpy array of timestamps read from file """ def write_mesh( self, filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None, mode: FileMode, time: float, ): """ Write a mesh to file. Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to backend_args: Arguments to backend mode: File-mode to store the mesh time: Time stamp associated with the mesh """ def write_meshtags( self, filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ def read_mesh_data( self, filename: Path | str, comm: MPI.Intracomm, time: str | float | None, read_from_partition: bool, backend_args: dict[str, Any] | None, ) -> ReadMeshData: """Read mesh data from file. Args: filename: Path to file to read from comm: MPI communicator used in storage time: Time stamp associated with the mesh to read read_from_partition: Whether to read partition information backend_args: Arguments to backend Returns: Internal data structure for the mesh data read from file """ def read_meshtags_data( self, filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend Returns: Internal data structure for the mesh tags read from file """ def read_dofmap( self, filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` """ def read_dofs( self, filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ def read_cell_perms( self, comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ def write_function( self, filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a function to file. Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ def read_legacy_mesh( self, filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ def snapshot_checkpoint( self, filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ def read_hdf5_array( self, comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ def read_point_data( self, filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ ... def read_function_names( self, filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ ... def read_cell_data( self, filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: """Read data from the cells of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: A tuple (topology, dofs) where topology contains the vertex indices of the cells, dofs the degrees of freedom within that cell. """ ... def write_data( self, filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file. Args: filename: Path to file array_data: Data to write to file. comm: The MPI communicator to open the writer with. time: The time stamp mode: Append or write backend_args: The backend arguments """ ... def get_backend(backend: str) -> IOBackend: """Get backend class from backend name. Args: backend: Name of the backend to get Returns: Backend class """ if backend == "h5py": from .h5py import backend as H5PYInterface return H5PYInterface elif backend == "adios2": from .adios2 import backend as ADIOS2Interface return ADIOS2Interface elif backend == "pyvista": from .pyvista import backend as PYVISTAInterface return PYVISTAInterface elif backend == "xdmf": from .xdmf import backend as XDMFInterface return XDMFInterface elif backend == "vtkhdf": from .vtkhdf import backend as VTKDHFInterface return VTKDHFInterface elif backend == "exodus": from .exodus import backend as EXODUSInterface return EXODUSInterface else: return import_module(backend) scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/adios2/000077500000000000000000000000001517634040500261315ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/adios2/__init__.py000066400000000000000000000001011517634040500302320ustar00rootroot00000000000000from . import backend, helpers __all__ = ["backend", "helpers"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/adios2/backend.py000066400000000000000000001115421517634040500300760ustar00rootroot00000000000000import warnings from pathlib import Path from typing import Any from mpi4py import MPI import adios2 import dolfinx import numpy as np import numpy.typing as npt from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from ...utils import check_file_exists, compute_local_range from .. import FileMode, ReadMode from .helpers import ( ADIOSFile, adios_to_numpy_dtype, check_variable_exists, read_adjacency_list, read_array, resolve_adios_scope, ) adios2 = resolve_adios_scope(adios2) read_mode = ReadMode.parallel def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default arguements (sets engine to BP4).""" args = arguments or {} if "engine" not in args.keys(): args["engine"] = "BP4" if "legacy" not in args.keys(): args["legacy"] = False # Only used for legacy HDF5 meshtags return args def convert_file_mode(mode: FileMode) -> adios2.Mode: # type: ignore[override] match mode: case FileMode.append: return adios2.Mode.Append case FileMode.write: return adios2.Mode.Write case FileMode.read: return adios2.Mode.Read case _: raise NotImplementedError(f"FileMode {mode} not implemented.") def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None = None, ): """Write attributes to file using ADIOS2. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attributes attributes: Dictionary of attributes to write to file engine: ADIOS2 engine to use """ adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Append, io_name="AttributeWriter", engine=backend_args["engine"], ) as adios_file: adios_file.file.BeginStep() for k, v in attributes.items(): adios_file.io.DefineAttribute(f"{name}_{k}", v) adios_file.file.PerformPuts() adios_file.file.EndStep() def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None, ) -> dict[str, np.ndarray]: """Read attributes from file using ADIOS2. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attributes engine: ADIOS2 engine to use Returns: The attributes """ check_file_exists(filename) adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, engine=backend_args["engine"], io_name="AttributesReader", ) as adios_file: adios_file.file.BeginStep() attributes = {} for k in adios_file.io.AvailableAttributes().keys(): if k.startswith(f"{name}_"): a = adios_file.io.InquireAttribute(k) attributes[k[len(name) + 1 :]] = a.Data() adios_file.file.EndStep() return attributes def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None = None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read time-stamps from a checkpoint file. Args: comm: MPI communicator filename: Path to file function_name: Name of the function to read time-stamps for backend_args: Arguments for backend, for instance file type. backend: What backend to use for writing. Returns: The time-stamps """ check_file_exists(filename) adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, engine=backend_args["engine"], io_name="TimestepReader", ) as adios_file: time_name = f"{function_name}_time" time_stamps = [] for _ in range(adios_file.file.Steps()): adios_file.file.BeginStep() if time_name in adios_file.io.AvailableVariables().keys(): arr = adios_file.io.InquireVariable(time_name) time_shape = arr.Shape() arr.SetSelection([[0], [time_shape[0]]]) times = np.empty( time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()], ) adios_file.file.Get(arr, times, adios2.Mode.Sync) time_stamps.append(times[0]) adios_file.file.EndStep() return np.array(time_stamps) def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None = None, mode: FileMode = FileMode.write, time: float = 0.0, ): """Write a mesh to file using ADIOS2. Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to backend_args: File mode and potentially the io-name. mode: Mode to use (write or append) time: Time stamp """ backend_args = get_default_backend_args(backend_args) if "io_name" not in backend_args.keys(): backend_args["io_name"] = "MeshWriter" mode = convert_file_mode(mode) gdim = mesh.local_geometry.shape[1] adios = adios2.ADIOS(comm) with ADIOSFile( adios=adios, filename=filename, mode=mode, comm=comm, engine=backend_args["engine"], io_name=backend_args["io_name"], ) as adios_file: adios_file.file.BeginStep() # Write geometry pointvar = adios_file.io.DefineVariable( "Points", mesh.local_geometry, shape=[mesh.num_nodes_global, gdim], start=[mesh.local_geometry_pos[0], 0], count=[mesh.local_geometry_pos[1] - mesh.local_geometry_pos[0], gdim], ) adios_file.file.Put(pointvar, mesh.local_geometry, adios2.Mode.Sync) if mode == adios2.Mode.Write: adios_file.io.DefineAttribute("CellType", mesh.cell_type) adios_file.io.DefineAttribute("Degree", np.array([mesh.degree], dtype=np.int32)) adios_file.io.DefineAttribute( "LagrangeVariant", np.array([mesh.lagrange_variant], dtype=np.int32) ) # Write topology (on;y on first write as topology is constant) num_dofs_per_cell = mesh.local_topology.shape[1] dvar = adios_file.io.DefineVariable( "Topology", mesh.local_topology, shape=[mesh.num_cells_global, num_dofs_per_cell], start=[mesh.local_topology_pos[0], 0], count=[ mesh.local_topology_pos[1] - mesh.local_topology_pos[0], num_dofs_per_cell, ], ) adios_file.file.Put(dvar, mesh.local_topology) # Add partitioning data if mesh.store_partition: assert mesh.partition_range is not None par_data = adios_file.io.DefineVariable( "PartitioningData", mesh.ownership_array, shape=[mesh.partition_global], start=[mesh.partition_range[0]], count=[ mesh.partition_range[1] - mesh.partition_range[0], ], ) adios_file.file.Put(par_data, mesh.ownership_array) assert mesh.ownership_offset is not None par_offset = adios_file.io.DefineVariable( "PartitioningOffset", mesh.ownership_offset, shape=[mesh.num_cells_global + 1], start=[mesh.local_topology_pos[0]], count=[mesh.local_topology_pos[1] - mesh.local_topology_pos[0] + 1], ) adios_file.file.Put(par_offset, mesh.ownership_offset) assert mesh.partition_processes is not None adios_file.io.DefineAttribute( "PartitionProcesses", np.array([mesh.partition_processes], dtype=np.int32) ) if mode == adios2.Mode.Append and mesh.store_partition: warnings.warn("Partitioning data is not written in append mode") # Add time step to file t_arr = np.array([time], dtype=np.float64) time_var = adios_file.io.DefineVariable( "MeshTime", t_arr, shape=[1], start=[0], count=[1 if comm.rank == 0 else 0], ) adios_file.file.Put(time_var, t_arr) adios_file.file.PerformPuts() adios_file.file.EndStep() def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None = 0.0, read_from_partition: bool = False, backend_args: dict[str, Any] | None = None, ) -> ReadMeshData: """Read an ADIOS2 mesh data for use with DOLFINx. Args: filename: Path to input file comm: The MPI communciator to distribute the mesh over engine: ADIOS engine to use for reading (BP4, BP5 or HDF5) time: Time stamp associated with mesh legacy: If checkpoint was made prior to time-dependent mesh-writer set to True read_from_partition: Read mesh with partition from file Returns: The mesh topology, geometry, UFL domain and partition function """ adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) legacy = backend_args.get("legacy", False) io_name = backend_args.get("io_name", "MeshReader") engine = backend_args["engine"] with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, engine=engine, io_name=io_name, ) as adios_file: # Get time independent mesh variables (mesh topology and cell type info) first adios_file.file.BeginStep() # Get mesh topology (distributed) if "Topology" not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"Mesh topology not found at Topology in {filename}") topology = adios_file.io.InquireVariable("Topology") shape = topology.Shape() local_range = compute_local_range(comm, shape[0]) topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]]) mesh_topology = np.empty((local_range[1] - local_range[0], shape[1]), dtype=np.int64) adios_file.file.Get(topology, mesh_topology, adios2.Mode.Deferred) # Check validity of partitioning information if read_from_partition: if "PartitionProcesses" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Partitioning information not found in {filename}") par_num_procs = adios_file.io.InquireAttribute("PartitionProcesses") num_procs = par_num_procs.Data()[0] if num_procs != comm.size: raise ValueError(f"Number of processes in file ({num_procs})!=({comm.size=})") # Get mesh cell type if "CellType" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Mesh cell type not found at CellType in {filename}") celltype = adios_file.io.InquireAttribute("CellType") cell_type = celltype.DataString()[0] # Get basix info if "LagrangeVariant" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Mesh LagrangeVariant not found in {filename}") lvar = adios_file.io.InquireAttribute("LagrangeVariant").Data()[0] if "Degree" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Mesh degree not found in {filename}") degree = adios_file.io.InquireAttribute("Degree").Data()[0] if not legacy: time_name = "MeshTime" for i in range(adios_file.file.Steps()): if i > 0: adios_file.file.BeginStep() if time_name in adios_file.io.AvailableVariables().keys(): arr = adios_file.io.InquireVariable(time_name) time_shape = arr.Shape() arr.SetSelection([[0], [time_shape[0]]]) times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()]) adios_file.file.Get(arr, times, adios2.Mode.Sync) if times[0] == time: break if i == adios_file.file.Steps() - 1: raise KeyError( f"No data associated with {time_name}={time} found in {filename}" ) adios_file.file.EndStep() if time_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"No data associated with {time_name}={time} found in {filename}") # Get mesh geometry if "Points" not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"Mesh coordinates not found at Points in {filename}") geometry = adios_file.io.InquireVariable("Points") x_shape = geometry.Shape() geometry_range = compute_local_range(comm, x_shape[0]) geometry.SetSelection( [ [geometry_range[0], 0], [geometry_range[1] - geometry_range[0], x_shape[1]], ] ) mesh_geometry = np.empty( (geometry_range[1] - geometry_range[0], x_shape[1]), dtype=adios_to_numpy_dtype[geometry.Type()], ) adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Deferred) adios_file.file.PerformGets() adios_file.file.EndStep() if read_from_partition: partition_graph = read_adjacency_list( adios, comm, filename, "PartitioningData", "PartitioningOffset", backend_args["engine"], ) else: partition_graph = None return ReadMeshData( cells=mesh_topology, cell_type=cell_type, x=mesh_geometry, degree=degree, lvar=lvar, partition_graph=partition_graph, ) def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None = None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ backend_args = {} if backend_args is None else backend_args io_name = backend_args.get("io_name", "MeshTagWriter") engine = backend_args.get("engine", "BP4") adios = adios2.ADIOS(comm) with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Append, engine=engine, io_name=io_name, ) as adios_file: adios_file.file.BeginStep() # Write meshtag topology topology_var = adios_file.io.DefineVariable( data.name + "_topology", data.indices, shape=[data.num_entities_global, data.num_dofs_per_entity], start=[data.local_start, 0], count=[len(data.indices), data.num_dofs_per_entity], ) adios_file.file.Put(topology_var, data.indices, adios2.Mode.Sync) # Write meshtag values vals = np.array(data.values) values_var = adios_file.io.DefineVariable( data.name + "_values", vals, shape=[data.num_entities_global], start=[data.local_start], count=[len(data.indices)], ) adios_file.file.Put(values_var, vals, adios2.Mode.Sync) # Write meshtag dim adios_file.io.DefineAttribute(data.name + "_dim", np.array([data.dim], dtype=np.uint8)) adios_file.file.PerformPuts() adios_file.file.EndStep() def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend Returns: Internal data structure for the mesh tags read from file """ adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) io_name = backend_args.get("io_name", "MeshTagsReader") engine = backend_args["engine"] legacy = backend_args["legacy"] with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, engine=engine, io_name=io_name, ) as adios_file: if not legacy: # Get mesh cell type dim_attr_name = f"{name}_dim" step = 0 for i in range(adios_file.file.Steps()): adios_file.file.BeginStep() if dim_attr_name in adios_file.io.AvailableAttributes().keys(): step = i break adios_file.file.EndStep() if dim_attr_name not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"{dim_attr_name} not found in {filename}") m_dim = adios_file.io.InquireAttribute(dim_attr_name) dim = int(m_dim.Data()[0]) # Get mesh tags entites topology_name = f"{name}_topology" for i in range(step, adios_file.file.Steps()): if i > step: adios_file.file.BeginStep() if topology_name in adios_file.io.AvailableVariables().keys(): break adios_file.file.EndStep() if topology_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"{topology_name} not found in {filename}") topology = adios_file.io.InquireVariable(topology_name) top_shape = topology.Shape() topology_range = compute_local_range(comm, top_shape[0]) topology.SetSelection( [ [topology_range[0], 0], [topology_range[1] - topology_range[0], top_shape[1]], ] ) mesh_entities = np.empty( (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64 ) adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred) # Get mesh tags values values_name = f"{name}_values" if values_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"{values_name} not found") values = adios_file.io.InquireVariable(values_name) val_shape = values.Shape() assert val_shape[0] == top_shape[0] values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]]) tag_values = np.empty( (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t") ) adios_file.file.Get(values, tag_values, adios2.Mode.Deferred) adios_file.file.PerformGets() adios_file.file.EndStep() else: # Get mesh cell type dim_attr_name = f"{name}_dim" assert adios_file.file.Steps() == 0 if (ct_key := f"/{name}/topology/celltype") in adios_file.io.AvailableAttributes(): cell_type = adios_file.io.InquireAttribute(ct_key) else: raise ValueError(f"Celltype not found for meshtags {name} in {filename}.") dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type.DataString()[0])) # Get mesh tags entites if (top_key := f"/{name}/topology") in adios_file.io.AvailableVariables(): topology = adios_file.io.InquireVariable(top_key) else: raise ValueError(f"Topology not found for meshtags {name} in {filename}.") top_shape = topology.Shape() topology_range = compute_local_range(comm, top_shape[0]) topology.SetSelection( [ [topology_range[0], 0], [topology_range[1] - topology_range[0], top_shape[1]], ] ) mesh_entities = np.empty( (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64 ) adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred) # Get mesh tags values if (val_key := f"/{name}/values") in adios_file.io.AvailableVariables(): values = adios_file.io.InquireVariable(val_key) else: raise ValueError(f"Values not found for meshtags {name} in {filename}.") val_shape = values.Shape() assert val_shape[0] == top_shape[0] values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]]) tag_values = np.empty( (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t") ) adios_file.file.Get(values, tag_values, adios2.Mode.Deferred) adios_file.file.PerformGets() adios_file.file.EndStep() return MeshTagsData( name=name, values=tag_values.astype(np.int32), indices=mesh_entities, dim=dim ) def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an AdjacencyList """ backend_args = {} if backend_args is None else backend_args # Handles legacy io4dolfinx files, modern files, and custom location of dofmap. legacy = backend_args.get("legacy", False) xdofmap_path: str | None dofmap_path: str | None if (dofmap_path := backend_args.get("dofmap", None)) is None: if legacy: dofmap_path = "Dofmap" else: dofmap_path = f"{name}_dofmap" if (xdofmap_path := backend_args.get("offsets", None)) is None: if legacy: xdofmap_path = "XDofmap" else: xdofmap_path = f"{name}_XDofmap" engine = backend_args.get("engine", "BP4") adios = adios2.ADIOS(comm) check_file_exists(filename) assert isinstance(xdofmap_path, str) return read_adjacency_list(adios, comm, filename, dofmap_path, xdofmap_path, engine=engine) def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None = None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ backend_args = {} if backend_args is None else backend_args legacy = backend_args.get("legacy", False) engine = backend_args.get("engine", "BP4") io_name = backend_args.get("io_name", f"{name}_FunctionReader") # Check that file contains the function to read adios = adios2.ADIOS(comm) check_file_exists(filename) if not legacy: with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, engine=engine, io_name=io_name, ) as adios_file: variables = set( sorted( map( lambda x: x.split("_time")[0], filter(lambda x: x.endswith("_time"), adios_file.io.AvailableVariables()), ) ) ) if name not in variables: raise KeyError(f"{name} not found in {filename}. Did you mean one of {variables}?") if legacy: array_path = "Values" else: array_path = f"{name}_values" time_name = f"{name}_time" return read_array(adios, filename, array_path, engine, comm, time, time_name, legacy=legacy) def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None = None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the mesh (global). Args: adios: The ADIOS instance comm: The MPI communicator used to read the data filename: Path to input file variable: Name of cell-permutation variable num_cells_global: Number of cells in the mesh (global) engine: Type of ADIOS engine to use for reading data Returns: Cell-permutations local to the process .. note:: No MPI communication is done during this call """ adios = adios2.ADIOS(comm) check_file_exists(filename) # Open ADIOS engine backend_args = {} if backend_args is None else backend_args engine = backend_args.get("engine", "BP4") cell_perms, _ = read_array( adios, filename, "CellPermutations", engine=engine, comm=comm, legacy=True ) return cell_perms.astype(np.uint32) def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None = None, ) -> tuple[np.ndarray, int]: adios = adios2.ADIOS(comm) """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ return read_array(adios, filename, group, engine="HDF5", comm=comm, legacy=True) def write_function( filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float = 0.0, mode: FileMode = FileMode.append, backend_args: dict[str, Any] | None = None, ): """ Write a function to file using ADIOS2 Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to engine: ADIOS2 engine to use mode: ADIOS2 mode to use (write or append) time: Time stamp associated with function io_name: Internal name used for the ADIOS IO object """ adios_mode = convert_file_mode(mode) backend_args = get_default_backend_args(backend_args) engine = backend_args["engine"] io_name = backend_args.get("io_name", "{name}_writer") adios = adios2.ADIOS(comm) cell_permutations_exists = False dofmap_exists = False XDofmap_exists = False if mode == adios2.Mode.Append: cell_permutations_exists = check_variable_exists( adios, filename, "CellPermutations", engine=engine ) dofmap_exists = check_variable_exists(adios, filename, f"{u.name}_dofmap", engine=engine) XDofmap_exists = check_variable_exists(adios, filename, f"{u.name}_XDofmap", engine=engine) with ADIOSFile( adios=adios, filename=filename, mode=adios_mode, engine=engine, io_name=io_name, comm=comm ) as adios_file: adios_file.file.BeginStep() if not cell_permutations_exists: # Add mesh permutations pvar = adios_file.io.DefineVariable( "CellPermutations", u.cell_permutations, shape=[u.num_cells_global], start=[u.local_cell_range[0]], count=[u.local_cell_range[1] - u.local_cell_range[0]], ) adios_file.file.Put(pvar, u.cell_permutations) if not dofmap_exists: # Add dofmap dofmap_var = adios_file.io.DefineVariable( f"{u.name}_dofmap", u.dofmap_array, shape=[u.global_dofs_in_dofmap], start=[u.dofmap_range[0]], count=[u.dofmap_range[1] - u.dofmap_range[0]], ) adios_file.file.Put(dofmap_var, u.dofmap_array) if not XDofmap_exists: # Add XDofmap xdofmap_var = adios_file.io.DefineVariable( f"{u.name}_XDofmap", u.dofmap_offsets, shape=[u.num_cells_global + 1], start=[u.local_cell_range[0]], count=[u.local_cell_range[1] - u.local_cell_range[0] + 1], ) adios_file.file.Put(xdofmap_var, u.dofmap_offsets) val_var = adios_file.io.DefineVariable( f"{u.name}_values", u.values, shape=[u.num_dofs_global], start=[u.dof_range[0]], count=[u.dof_range[1] - u.dof_range[0]], ) adios_file.file.Put(val_var, u.values) # Add time step to file t_arr = np.array([time], dtype=np.float64) time_var = adios_file.io.DefineVariable( f"{u.name}_time", t_arr, shape=[1], start=[0], count=[1 if comm.rank == 0 else 0], ) adios_file.file.Put(time_var, t_arr) adios_file.file.PerformPuts() adios_file.file.EndStep() def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ # Create ADIOS2 reader adios = adios2.ADIOS(comm) with ADIOSFile( adios=adios, filename=filename, mode=adios2.Mode.Read, io_name="Mesh reader", engine="HDF5", ) as adios_file: # Get mesh topology (distributed) if f"{group}/topology" not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"Mesh topology not found at '{group}/topology'") topology = adios_file.io.InquireVariable(f"{group}/topology") shape = topology.Shape() local_range = compute_local_range(comm, shape[0]) topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]]) mesh_topology = np.empty( (local_range[1] - local_range[0], shape[1]), dtype=topology.Type().strip("_t"), ) adios_file.file.Get(topology, mesh_topology, adios2.Mode.Sync) # Get mesh cell type cell_type = None if f"{group}/topology/celltype" in adios_file.io.AvailableAttributes().keys(): celltype = adios_file.io.InquireAttribute(f"{group}/topology/celltype") cell_type = celltype.DataString()[0] # Get mesh geometry for geometry_key in [f"{group}/geometry", f"{group}/coordinates"]: if geometry_key in adios_file.io.AvailableVariables().keys(): break else: raise KeyError( f"Mesh coordinates not found at '{group}/coordinates' or '{group}/geometry'" ) geometry = adios_file.io.InquireVariable(geometry_key) shape = geometry.Shape() local_range = compute_local_range(comm, shape[0]) geometry.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]]) mesh_geometry = np.empty( (local_range[1] - local_range[0], shape[1]), dtype=adios_to_numpy_dtype[geometry.Type()], ) adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Sync) return mesh_topology, mesh_geometry, cell_type def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ adios_mode = convert_file_mode(mode) adios = adios2.ADIOS(u.function_space.mesh.comm) backend_args = {} if backend_args is None else backend_args io_name = backend_args.get("io_name", "SnapshotCheckPoint") engine = backend_args.get("engine", "BP4") with ADIOSFile( adios=adios, filename=filename, mode=adios_mode, io_name=io_name, engine=engine, ) as adios_file: if adios_mode == adios2.Mode.Write: dofmap = u.function_space.dofmap num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs local_dofs = u.x.array[:num_dofs_local].copy() # Write to file adios_file.file.BeginStep() dofs = adios_file.io.DefineVariable("dofs", local_dofs, count=[num_dofs_local]) adios_file.file.Put(dofs, local_dofs, adios2.Mode.Sync) adios_file.file.EndStep() elif adios_mode == adios2.Mode.Read: adios_file.file.BeginStep() in_variable = adios_file.io.InquireVariable("dofs") in_variable.SetBlockSelection(u.function_space.mesh.comm.rank) adios_file.file.Get(in_variable, u.x.array, adios2.Mode.Sync) adios_file.file.EndStep() u.x.scatter_forward() else: raise NotImplementedError(f"Mode {mode} is not implemented for snapshot checkpoint") def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ raise NotImplementedError("Reading function names are not implemented with ADIOS2") def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ raise NotImplementedError("The ADIOS2 backend cannot read point data.") def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: """Read data from the cells of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: A tuple (topology, dofs) where topology contains the vertex indices of the cells, dofs the degrees of freedom within that cell. """ raise NotImplementedError("The ADIOS2 backend does not support reading cell data.") def write_data( filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file (distributed across proceses with MPI). Args: filename: Path to file array_data: Data to write to file comm: MPI communicator to open the file with. time: Time stamp mode: Append or write backend_args: The backend arguments """ raise NotImplementedError("ADIOS2 has not implemented this yet") scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/adios2/helpers.py000066400000000000000000000236541517634040500301570ustar00rootroot00000000000000""" Helpers reading/writing data with ADIOS2 """ from __future__ import annotations import shutil from contextlib import contextmanager from pathlib import Path from typing import NamedTuple from mpi4py import MPI import adios2 import dolfinx.cpp.graph import dolfinx.graph import numpy as np import numpy.typing as npt from io4dolfinx.utils import compute_local_range, valid_function_types def resolve_adios_scope(adios2): scope = adios2.bindings if hasattr(adios2, "bindings") else adios2 if not scope.is_built_with_mpi: raise ImportError("ADIOS2 must be built with MPI support") return scope adios2 = resolve_adios_scope(adios2) __all__ = [ "AdiosFile", "ADIOSFile", "check_variable_exists", "read_array", "read_adjacency_list", "adios_to_numpy_dtype", ] adios_to_numpy_dtype = { "float": np.float32, "double": np.float64, "float complex": np.complex64, "double complex": np.complex128, "uint32_t": np.uint32, } class AdiosFile(NamedTuple): io: adios2.IO file: adios2.Engine @contextmanager def ADIOSFile( adios: adios2.ADIOS, filename: Path | str, engine: str, mode: adios2.Mode, io_name: str, comm: MPI.Intracomm | None = None, ): io = adios.DeclareIO(io_name) io.SetEngine(engine) # ADIOS2 sometimes struggles with existing files/folders it should overwrite if mode == adios2.Mode.Write: filename = Path(filename) if filename.exists() and comm is not None and comm.rank == 0: if filename.is_dir(): shutil.rmtree(filename) else: filename.unlink() if comm is not None: comm.Barrier() file = io.Open(str(filename), mode) try: yield AdiosFile(io=io, file=file) finally: file.Close() adios.RemoveIO(io_name) def check_variable_exists( adios: adios2.ADIOS, filename: Path | str, variable: str, engine: str, ) -> bool: io_name = f"{variable}_reader" if not Path(filename).exists(): return False variable_found = False with ADIOSFile( adios=adios, engine=engine, filename=filename, mode=adios2.Mode.Read, io_name=io_name, ) as adios_file: # Find step that has cell permutation for _ in range(adios_file.file.Steps()): adios_file.file.BeginStep() if variable in adios_file.io.AvailableVariables().keys(): variable_found = True break adios_file.file.EndStep() # Not sure if this is needed, but just in case if variable in adios_file.io.AvailableVariables().keys(): variable_found = True return variable_found def read_adjacency_list( adios: adios2.ADIOS, comm: MPI.Intracomm, filename: Path | str, data_name: str, offsets_name: str, engine: str, ) -> dolfinx.graph.AdjacencyList: """ Read an adjacency-list from an ADIOS file with given communicator. The adjancency list is split in to a flat array (data) and its corresponding offset. Args: adios: The ADIOS instance comm: The MPI communicator used to read the data filename: Path to input file data_name: Name of variable containing the indices of the adjacencylist dofmap_offsets: Name of variable containing offsets of the adjacencylist engine: Type of ADIOS engine to use for reading data Returns: The local part of dofmap from input dofs .. note:: No MPI communication is done during this call """ # Open ADIOS engine io_name = f"{data_name=}_reader" with ADIOSFile( adios=adios, engine=engine, filename=filename, mode=adios2.Mode.Read, io_name=io_name, ) as adios_file: # First find step with dofmap offsets, to be able to read # in a full row of the dofmap for _ in range(adios_file.file.Steps()): adios_file.file.BeginStep() if offsets_name in adios_file.io.AvailableVariables().keys(): break adios_file.file.EndStep() if offsets_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"Dof offsets not found at '{offsets_name}' in {filename}") # Get global shape of dofmap-offset, and read in data with an overlap d_offsets = adios_file.io.InquireVariable(offsets_name) shape = d_offsets.Shape() num_nodes = shape[0] - 1 local_range = compute_local_range(comm, num_nodes) # As the offsets are one longer than the number of cells, we need to read in with an overlap if len(shape) == 1: d_offsets.SetSelection([[local_range[0]], [local_range[1] + 1 - local_range[0]]]) in_offsets = np.empty( local_range[1] + 1 - local_range[0], dtype=d_offsets.Type().strip("_t"), ) else: d_offsets.SetSelection( [ [local_range[0], 0], [local_range[1] + 1 - local_range[0], shape[1]], ] ) in_offsets = np.empty( (local_range[1] + 1 - local_range[0], shape[1]), dtype=d_offsets.Type().strip("_t"), ) adios_file.file.Get(d_offsets, in_offsets, adios2.Mode.Sync) in_offsets = in_offsets.squeeze() # Assuming dofmap is saved in stame step # Get the relevant part of the dofmap if data_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"Dofs not found at {data_name} in {filename}") cell_dofs = adios_file.io.InquireVariable(data_name) if len(shape) == 1: cell_dofs.SetSelection([[in_offsets[0]], [in_offsets[-1] - in_offsets[0]]]) in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t")) else: cell_dofs.SetSelection([[in_offsets[0], 0], [in_offsets[-1] - in_offsets[0], shape[1]]]) in_dofmap = np.empty( (in_offsets[-1] - in_offsets[0], shape[1]), dtype=cell_dofs.Type().strip("_t"), ) assert shape[1] == 1 in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t")) adios_file.file.Get(cell_dofs, in_dofmap, adios2.Mode.Sync) in_offsets -= in_offsets[0] adios_file.file.EndStep() # Return local dofmap return dolfinx.graph.adjacencylist(in_dofmap, in_offsets.astype(np.int32)) def read_array( adios: adios2.ADIOS, filename: Path | str, array_name: str, engine: str, comm: MPI.Intracomm, time: float = 0.0, time_name: str = "", legacy: bool = False, ) -> tuple[npt.NDArray[valid_function_types], int]: """ Read an array from file, return the global starting position of the local array Args: adios: The ADIOS instance filename: Path to file to read array from array_name: Name of array in file engine: Name of engine to use to read file comm: MPI communicator used for reading the data time_name: Name of time variable for modern checkpoints legacy: If True ignore time_name and read the first available step Returns: Local part of array and its global starting position """ with ADIOSFile( adios=adios, engine=engine, filename=filename, mode=adios2.Mode.Read, io_name="ArrayReader", ) as adios_file: # Get time-stamp from first available step if legacy: for i in range(adios_file.file.Steps()): adios_file.file.BeginStep() if array_name in adios_file.io.AvailableVariables().keys(): break adios_file.file.EndStep() if array_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"No array found at {array_name}") else: for i in range(adios_file.file.Steps()): adios_file.file.BeginStep() if time_name in adios_file.io.AvailableVariables().keys(): arr = adios_file.io.InquireVariable(time_name) time_shape = arr.Shape() arr.SetSelection([[0], [time_shape[0]]]) times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()]) adios_file.file.Get(arr, times, adios2.Mode.Sync) if times[0] == time: break if i == adios_file.file.Steps() - 1: raise KeyError( f"No data associated with {time_name}={time} found in {filename}" ) adios_file.file.EndStep() if time_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"No data associated with {time_name}={time} found in {filename}") if array_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"No array found at {time=} for {array_name}") arr = adios_file.io.InquireVariable(array_name) arr_shape = arr.Shape() # TODO: Should we always pick the first element? assert len(arr_shape) >= 1 arr_range = compute_local_range(comm, arr_shape[0]) if len(arr_shape) == 1: arr.SetSelection([[arr_range[0]], [arr_range[1] - arr_range[0]]]) vals = np.empty(arr_range[1] - arr_range[0], dtype=adios_to_numpy_dtype[arr.Type()]) else: arr.SetSelection([[arr_range[0], 0], [arr_range[1] - arr_range[0], arr_shape[1]]]) vals = np.empty( (arr_range[1] - arr_range[0], arr_shape[1]), dtype=adios_to_numpy_dtype[arr.Type()], ) assert arr_shape[1] == 1 adios_file.file.Get(arr, vals, adios2.Mode.Sync) adios_file.file.EndStep() return vals.reshape(-1), arr_range[0] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/exodus/000077500000000000000000000000001517634040500262575ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/exodus/__init__.py000066400000000000000000000000551517634040500303700ustar00rootroot00000000000000from . import backend __all__ = ["backend"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/exodus/backend.py000066400000000000000000000675251517634040500302370ustar00rootroot00000000000000""" Exodus interface to io4dolfinx. The Exodus2 format is described in: https://src.fedoraproject.org/repo/pkgs/exodusii/922137.pdf/a45d67f4a1a8762bcf66af2ec6eb35f9/922137.pdf Further documentation from CUBIT on node numbering can be found at: https://coreform.com/cubit_help/appendix/element_numbering.htm SPDX License identifier: MIT Copyright: Jørgen S. Dokken, Henrik N.T. Finsberg, Remi Delaporte-Mathurin, and Simula Research Laboratory """ from pathlib import Path from typing import Any, Literal, cast from mpi4py import MPI import basix.ufl import dolfinx import netCDF4 import numpy as np import numpy.typing as npt from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from ...utils import check_file_exists from .. import FileMode, ReadMode _interval_to_vertex_map = {0: [0, 1]} # Based on: https://src.fedoraproject.org/repo/pkgs/exodusii/922137.pdf/a45d67f4a1a8762bcf66af2ec6eb35f9/922137.pdf _tetra_facet_to_vertex_map = {0: [0, 1, 3], 1: [1, 2, 3], 2: [0, 2, 3], 3: [0, 1, 2]} # https://coreform.com/cubit_help/appendix/element_numbering.htm # Note that triangular side-sets goes from 2 to 4 (with 0 base index) _triangle_to_vertex_map = {2: [0, 1], 3: [1, 2], 4: [2, 0]} _quad_to_vertex_map = {0: [0, 1], 1: [1, 2], 2: [2, 3], 3: [3, 0]} _hex_to_vertex_map = { 0: [0, 1, 4, 5], 1: [1, 2, 5, 6], 2: [2, 3, 6, 7], 3: [0, 3, 4, 7], 4: [0, 1, 2, 3], 5: [4, 5, 6, 7], } _side_set_to_vertex_map = { "interval": _interval_to_vertex_map, "quadrilateral": _quad_to_vertex_map, "triangle": _triangle_to_vertex_map, "tetrahedron": _tetra_facet_to_vertex_map, "hexahedron": _hex_to_vertex_map, } _exodus_to_string = { "EDGE2": ("interval", 1), "TRI3": ("triangle", 1), "QUAD": ("quadrilateral", 1), "QUAD4": ("quadrilateral", 1), "TETRA": ("tetrahedron", 1), "HEX": ("hexahedron", 1), "HEX8": ("hexahedron", 1), "BEAM2": ("interval", 1), "HEX27": ("hexahedron", 2), } read_mode = ReadMode.serial def _get_cell_type(connectivity: netCDF4.Variable) -> tuple[dolfinx.mesh.CellType, int]: cell_type, degree = _exodus_to_string[connectivity.elem_type] return dolfinx.mesh.to_type(cell_type), degree def _compute_tdim(connectivity: netCDF4.Variable) -> int: d_ct, _ = _get_cell_type(connectivity) return dolfinx.mesh.cell_dim(d_ct) def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: args = arguments or {} return args def convert_file_mode(mode: FileMode) -> str: match mode: case FileMode.append: return "a" case FileMode.read: return "r" case FileMode.write: return "w" case _: raise NotImplementedError(f"File mode {mode} not implemented") def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None = None, ): """Write attributes to file using H5PY. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attributes attributes: Dictionary of attributes to write to file backend_args: Arguments to backend """ raise NotImplementedError("The Exodus backend cannot write attributes.") def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None, ) -> dict[str, Any]: """Read attributes from file using H5PY. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attributes backend_args: Arguments to backend Returns: The attributes """ raise NotImplementedError("The Exodus backend cannot read attributes.") def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ raise NotImplementedError("The EXODUS backend cannot make checkpoints.") def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The EXODUS backend cannot read HDF5 arrays") def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None = None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read time-stamps from a checkpoint file. Args: comm: MPI communicator filename: Path to file comm: MPI communicator function_name: Name of the function to read time-stamps for backend_args: Arguments for backend, for instance file type. Returns: The time-stamps """ raise NotImplementedError("The Exodus backend cannot read timestamps.") def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None = None, mode: FileMode = FileMode.write, time: float = 0.0, ): """Write a mesh to file using H5PY Args: filename: Path to file to write to. mesh: Internal data structure for the mesh data to save to file comm: MPI communicator used in storage backend_args: Arguments to backend mode: Mode to use (write or append) time: Time stamp """ raise NotImplementedError("The Exodus backend cannot write meshes.") def _read_mesh_geometry(infile: netCDF4.Dataset) -> tuple[int, npt.NDArray[np.floating]]: # use page 171 of manual to extract data num_nodes = infile.dimensions["num_nodes"].size gdim = infile.dimensions["num_dim"].size # Get coordinates of mesh coord_var = infile.variables.get("coord") if coord_var is None: coordinates = np.zeros((num_nodes, gdim), dtype=np.float64) for i, coord in enumerate(["x", "y", "z"]): coord_i = infile.variables.get(f"coord{coord}") if coord_i is not None: coordinates[: coord_i.size, i] = coord_i[:] else: coordinates = np.asarray(coord_var) return gdim, coordinates def _get_entity_blocks( infile: netCDF4.Dataset, search_type: Literal["cell", "facet"] ) -> tuple[int, list[netCDF4.Variable]]: # use page 171 of manual to extract data num_blocks = infile.dimensions["num_el_blk"].size # Get element connectivity all_connectivity_variables = [infile.variables[f"connect{i + 1}"] for i in range(num_blocks)] # Compute max topological dimension in mesh and find the correct max_tdim = _compute_tdim(max(all_connectivity_variables, key=_compute_tdim)) # Extract only the connectivity blocks that we need if search_type == "cell": search_dim = max_tdim elif search_type == "facet": search_dim = max_tdim - 1 else: raise RuntimeError(f"Unknown entity type: {search_type}") return search_dim, list( filter(lambda el: _compute_tdim(el) == search_dim, all_connectivity_variables) ) def _extract_connectivity_data( entity_blocks: list[netCDF4.Variable], ) -> tuple[list[npt.NDArray[np.int64]], tuple[dolfinx.mesh.CellType, int], list[int]]: connectivity_arrays = [] cell_types = [] entity_block_index = [] for entity_block in entity_blocks: connectivity_arrays.append(entity_block[:] - 1) cell_types.append(_get_cell_type(entity_block)) entity_block_index.append(int(entity_block.name.removeprefix("connect")) - 1) for cell in cell_types: assert cell_types[0] == cell, "Mixed cell types not supported" cell_type = cell_types[0] return connectivity_arrays, cell_type, entity_block_index def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None = 0.0, read_from_partition: bool = False, backend_args: dict[str, Any] | None = None, ) -> ReadMeshData: """Read mesh data from EXODUS based checkpoint files. Args: filename: Path to input file comm: The MPI communciator to distribute the mesh over time: Time stamp associated with mesh read_from_partition: Read mesh with partition from file backend_args: Arguments to backend Returns: The mesh topology, geometry, UFL domain and partition function """ check_file_exists(filename) with netCDF4.Dataset(filename, "r") as infile: if comm.rank == 0: gdim, coordinates = _read_mesh_geometry(infile) _tdim, entity_blocks = _get_entity_blocks(infile, "cell") if len(entity_blocks) > 0: # Extract markers directly from entity-blocks connectivity_arrays, (cell_type, degree), _entity_block_index = ( _extract_connectivity_data(entity_blocks) ) cells = np.vstack(connectivity_arrays) if isinstance(cells, np.ma.MaskedArray): cells = cells.filled() else: raise ValueError(f"No blocks found in {filename}") if degree == 1: perm = dolfinx.cpp.io.perm_vtk(cell_type, cells.shape[1]) elif cell_type == dolfinx.mesh.CellType.hexahedron and degree == 2: # Ordering from Fig 4.14 of: https://sandialabs.github.io/seacas-docs/exodusII-new.pdf dolfinx_to_exodus = np.array( [ 0, 1, 5, 4, 2, 3, 7, 6, 8, 12, 16, 10, 9, 11, 18, 17, 13, 15, 19, 14, 26, 21, 24, 22, 23, 20, 25, ] ) perm = np.argsort(dolfinx_to_exodus) else: raise NotImplementedError( "Reading Exodus2 mesh with {cell_type} of order {degree} is not supported." ) cells = cells[:, perm] cell_type, gdim, xtype, degree, num_dofs_per_cell = comm.bcast( (cell_type, gdim, np.dtype(coordinates.dtype).name, degree, cells.shape[1]), root=0 ) else: cell_type, gdim, xtype, degree, num_dofs_per_cell = comm.bcast( (None, None, None, None), root=0 ) coordinates = np.zeros((0, gdim), dtype=xtype) cells = np.zeros((0, num_dofs_per_cell), dtype=np.int64) return ReadMeshData( cells=cells, cell_type=dolfinx.mesh.to_string(cell_type), x=coordinates, lvar=int(basix.LagrangeVariant.equispaced), degree=degree, partition_graph=None, ) def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None = None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ raise NotImplementedError("The Exodus backend cannot write meshtags.") def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None, ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend. Returns: Internal data structure for the mesh tags read from file """ if comm.rank == 0: with netCDF4.Dataset(filename, "r") as infile: # Compute max topological dimension in mesh and find the correct if name == "cell" or name == "facet": search_dim, entity_blocks = _get_entity_blocks( infile, cast(Literal["cell", "facet"], name) ) else: raise RuntimeError("Expected name='cell' or 'facet' got {name}") if len(entity_blocks) > 0: # Extract markers directly from entity-blocks connectivity_arrays, (cell_type, degree), entity_block_index = ( _extract_connectivity_data(entity_blocks) ) marked_entities = np.vstack(connectivity_arrays) entity_values = np.zeros(marked_entities.shape[0], dtype=np.int64) if "eb_prop1" in infile.variables.keys(): block_values = infile.variables["eb_prop1"][:] # First check if entities are in eb_prop1 insert_offset = np.zeros(len(connectivity_arrays) + 1, dtype=np.int64) insert_offset[1:] = np.cumsum([c_arr.shape[0] for c_arr in connectivity_arrays]) for i, index in enumerate(entity_block_index): entity_values[insert_offset[i] : insert_offset[i + 1]] = block_values[index] else: num_dofs_per_cell = basix.ufl.element( "Lagrange", dolfinx.mesh.to_string(cell_type), degree ).dim assert num_dofs_per_cell == marked_entities.shape[1] marked_entities = np.zeros((0, marked_entities.shape[1]), dtype=np.int64) entity_values = np.zeros(0, dtype=np.int64) elif name == "facet" and "ss_prop1" in infile.variables.keys(): # If we haven't found the cell type as a block, we should be extracting facets # (from side-sets), then we need the parent cell tdim, entity_blocks = _get_entity_blocks(infile, "cell") cell_types = [] for entity_block in entity_blocks: cell_types.append(_get_cell_type(entity_block)) for cell in cell_types: assert cell_types[0] == cell, "Mixed cell types not supported" cell_type, degree = cell_types[0] local_facet_index = _side_set_to_vertex_map[dolfinx.mesh.to_string(cell_type)] if "num_side_sets" not in infile.dimensions: facet_type = dolfinx.cpp.mesh.cell_entity_type(cell_type, tdim - 1, 0) num_dofs_per_cell = basix.ufl.element( "Lagrange", dolfinx.mesh.to_string(facet_type), degree ).dim marked_entities = np.zeros((0, num_dofs_per_cell), dtype=np.int64) entity_values = np.zeros(0, dtype=np.int64) else: # Extract facet values local_facet_index = _side_set_to_vertex_map[dolfinx.mesh.to_string(cell_type)] num_facet_sets = infile.dimensions["num_side_sets"].size values = infile.variables["ss_prop1"] # Extract all cell blocks to get correct look-up connectivity_arrays = [] for entity_block in entity_blocks: connectivity_arrays.append(entity_block[:] - 1) connectivity_array = np.vstack(connectivity_arrays) # Loop through all side sets to extract the correct connectivity facet_indices = [] facet_values = [] for i in range(num_facet_sets): value = values[i].reshape(-1) elements = infile.variables[f"elem_ss{i + 1}"] local_facets = infile.variables[f"side_ss{i + 1}"] for element, index in zip(elements, local_facets): facet_indices.append( connectivity_array[element - 1, local_facet_index[index - 1]] ) facet_values.append(value.data.tolist()) marked_entities = np.vstack(facet_indices) entity_values = np.array(facet_values, dtype=np.int64).flatten() else: # If we found no blocks (for instance for facets, we search through the cells) tdim, entity_blocks = _get_entity_blocks(infile, "cell") search_dim = tdim - 1 if name == "facet" else tdim cell_type, degree = _get_cell_type(entity_blocks[0]) facet_type = dolfinx.cpp.mesh.cell_entity_type(cell_type, search_dim, 0) num_dofs_per_cell = basix.ufl.element( "Lagrange", dolfinx.mesh.to_string(facet_type), degree ).dim # If we cannot find any information about the blocks we send nothing marked_entities = np.zeros((0, num_dofs_per_cell), dtype=np.int64) entity_values = np.zeros(0, dtype=np.int64) # Broadcast information read by this process to other processes (dim, _, _) = comm.bcast( (search_dim, marked_entities.shape[1], np.dtype(entity_values.dtype).name), root=0, ) else: # Other process gets info from process that read the file dim, num_dofs_per_cell, vtype = comm.bcast((None, None, None), root=0) marked_entities = np.zeros((0, num_dofs_per_cell), dtype=np.int64) entity_values = np.zeros(0, dtype=vtype) return MeshTagsData(name=name, values=entity_values, indices=marked_entities, dim=dim) def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` """ raise NotImplementedError("The Exodus backend cannot read dofmap.") def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Exodus backend cannot read dofs.") def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Exodus backend cannot read cell perms.") def write_function( filename: str | Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None = None, ): """Write a function to file. Args: filename: Path to file to write to comm: MPI communicator used in storage u: Internal data structure for the function data to save to file time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ raise NotImplementedError("The Exodus backend cannot write function.") def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ raise NotImplementedError("The Exodus backend cannot read legacy mesh.") def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ num_components = 1 # Default assumption, overriden by data read in having multiple components if comm.rank == 0: with netCDF4.Dataset(filename, "r") as infile: raw_names = infile.variables["name_nod_var"][:].data node_names = netCDF4.chartostring(raw_names) if name not in node_names: raise ValueError( f"Point data with name {name} not found in file.", f"Available variables: {node_names}", ) index = np.flatnonzero(name == node_names)[0] + 1 temporal_dataset = infile.variables[f"vals_nod_var{index}"] time_steps = infile.variables["time_whole"][:].data if time is None: time_idx = time_steps[0] else: time_indices = np.flatnonzero(np.isclose(time_steps, time)) if len(time_indices) == 0: raise ValueError( f"Could not find {name}(t={time}), available time steps are {time_steps}" ) time_idx = time_indices[0] dataset = temporal_dataset[time_idx] if len(dataset.shape) == 1: dataset = dataset.reshape(-1, num_components) else: num_components = dataset.shape[1] # Broadcast num components to all other ranks num_components = comm.bcast(num_components, root=0) # Zero data on all other processes if comm.rank != 0: dataset = np.zeros((0, num_components), dtype=np.float64) return dataset, 0 def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: """Read data from the cells of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: A tuple (topology, dofs) where topology contains the vertex indices of the cells, dofs the degrees of freedom within that cell. """ num_components = 1 # Default assumption, overriden by data read in having multiple components if comm.rank == 0: with netCDF4.Dataset(filename, "r") as infile: raw_names = infile.variables["name_elem_var"][:].data num_blocks = infile.dimensions["num_el_blk"].size node_names = netCDF4.chartostring(raw_names) if name not in node_names: raise ValueError( f"Cell data with name {name} not found in file.", f"Available variables: {node_names}", ) index = np.flatnonzero(name == node_names)[0] + 1 entity_blocks = [ infile.variables[f"vals_elem_var{index}eb{i + 1}"] for i in range(num_blocks) ] time_steps = infile.variables["time_whole"][:].data if time is None: time_idx = time_steps[0] else: time_indices = np.flatnonzero(np.isclose(time_steps, time)) if len(time_indices) == 0: raise ValueError( f"Could not find {name}(t={time}), available time steps are {time_steps}" ) time_idx = time_indices[0] if len(entity_blocks) > 0: datasets = [] for entity_block in entity_blocks: datasets.append(entity_block[time_idx]) dataset = np.hstack(datasets) if len(dataset.shape) == 1: dataset = dataset.reshape(-1, num_components) else: num_components = dataset.shape[1] # Broadcast num components to all other ranks num_components = comm.bcast(num_components, root=0) # Zero data on all other processes if comm.rank != 0: dataset = np.zeros((0, num_components), dtype=np.float64) _time = float(time) if time is not None else None topology = read_mesh_data(filename, comm, _time, False, backend_args=None).cells return topology, dataset def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ with netCDF4.Dataset(filename, "r") as infile: function_names: list[str] = [] for key in ["name_elem_var", "name_nod_var"]: raw_names = infile.variables[key][:].data decoded_names = netCDF4.chartostring(raw_names) function_names.extend(decoded_names) return function_names def write_data( filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file (distributed across proceses with MPI). Args: filename: Path to file array_data: Data to write to file comm: MPI communicator to open the file with time: Time-stamp for data. mode: Append or write backend_args: The backend arguments """ raise NotImplementedError("Exodus has not implemented this yet") scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/h5py/000077500000000000000000000000001517634040500256355ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/h5py/__init__.py000066400000000000000000000000551517634040500277460ustar00rootroot00000000000000from . import backend __all__ = ["backend"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/h5py/backend.py000066400000000000000000000737151517634040500276130ustar00rootroot00000000000000""" H5py interface to io4dolfinx SPDX License identifier: MIT Copyright: Jørgen S. Dokken, Henrik N.T. Finsberg, Simula Research Laboratory """ import contextlib from pathlib import Path from typing import Any from mpi4py import MPI import dolfinx import h5py import numpy as np import numpy.typing as npt from dolfinx.graph import adjacencylist from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from ...utils import check_file_exists, compute_local_range from .. import FileMode, ReadMode read_mode = ReadMode.parallel # try: # except ModuleNotFoundError: # raise ModuleNotFoundError("This backend requires h5py to be installed.") @contextlib.contextmanager def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None): """Context manager for opening an HDF5 file with h5py. Args: h5name: The name of the HDF5 file. filemode: The file mode. force_serial: Force serial access to the file. comm: The MPI communicator """ if comm is None: comm = MPI.COMM_WORLD if h5py.h5.get_config().mpi and not force_serial: h5file = h5py.File(h5name, filemode, driver="mpio", comm=comm) else: if comm.size > 1 and not force_serial: raise ValueError( f"h5py is not installed with MPI support, while using {comm.size} processes.", "If you really want to do this, turn on the `force_serial` flag.", ) h5file = h5py.File(h5name, filemode) try: yield h5file finally: if h5file.id: h5file.close() def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: args = arguments or {"legacy": False} # If meshtags is read from legacy return args def convert_file_mode(mode: FileMode) -> str: match mode: case FileMode.append: return "a" case FileMode.read: return "r" case FileMode.write: return "w" case _: raise NotImplementedError(f"File mode {mode} not implemented") def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None = None, ): """Write attributes to file using H5PY. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attributes attributes: Dictionary of attributes to write to file engine: ADIOS2 engine to use """ filemode = "a" if Path(filename).exists() else "w" with h5pyfile(filename, filemode=filemode, comm=comm, force_serial=False) as h5file: if name not in h5file.keys(): h5file.create_group(name, track_order=True) group = h5file[name] for key, val in attributes.items(): group.attrs[key] = val def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None, ) -> dict[str, Any]: """Read attributes from file using H5PY. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attributes Returns: The attributes """ check_file_exists(filename) output_attrs = {} with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: for key, val in h5file[name].attrs.items(): output_attrs[key] = val return output_attrs def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None = None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read time-stamps from a checkpoint file. Args: comm: MPI communicator filename: Path to file function_name: Name of the function to read time-stamps for backend_args: Arguments for backend, for instance file type. Returns: The time-stamps """ check_file_exists(filename) mesh_name = "mesh" with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: mesh_directory = h5file[mesh_name] functions = mesh_directory["functions"] u = functions[function_name] timestamps = u.attrs["timestamps"] return timestamps def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None = None, mode: FileMode = FileMode.write, time: float = 0.0, ): """Write a mesh to file using H5PY Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to. mode: Mode to use (write or append) time: Time stamp """ backend_args = get_default_backend_args(backend_args) h5_mode = convert_file_mode(mode) mesh_name = "mesh" with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: if mesh_name in h5file.keys() and h5_mode == "a": mesh_directory = h5file[mesh_name] timestamps = mesh_directory.attrs["timestamps"] if np.isclose(time, timestamps).any(): raise ValueError("Mesh has already been stored at time={time_stamp}.") else: mesh_directory.attrs["timestamps"] = np.append( mesh_directory.attrs["timestamps"], time ) idx = len(mesh_directory.attrs["timestamps"]) - 1 write_topology = False else: mesh_directory = h5file.create_group(mesh_name) mesh_directory.attrs["timestamps"] = np.array([time], dtype=np.float64) idx = 0 write_topology = True geometry_group = mesh_directory.create_group(f"{idx}") # Write geometry data gdim = mesh.local_geometry.shape[1] geometry_dataset = geometry_group.create_dataset( "Points", [mesh.num_nodes_global, gdim], dtype=mesh.local_geometry.dtype, ) geometry_dataset[slice(*mesh.local_geometry_pos), :] = mesh.local_geometry # Write static partitioning data if "PartitioningData" not in mesh_directory.keys() and mesh.store_partition: assert mesh.partition_range is not None assert mesh.ownership_array is not None par_dataset = mesh_directory.create_dataset( "PartitioningData", [mesh.partition_global], dtype=mesh.ownership_array.dtype ) par_dataset[slice(*mesh.partition_range)] = mesh.ownership_array if "PartitioningOffset" not in mesh_directory.keys() and mesh.store_partition: assert mesh.ownership_offset is not None par_dataset = mesh_directory.create_dataset( "PartitioningOffset", [mesh.num_cells_global + 1], dtype=np.int64 ) par_dataset[mesh.local_topology_pos[0] : mesh.local_topology_pos[1] + 1] = ( mesh.ownership_offset ) if "PartitionProcesses" not in mesh_directory.attrs.keys() and mesh.store_partition: mesh_directory.attrs["PartitionProcesses"] = mesh.partition_processes # Write static data if write_topology: mesh_directory.attrs["CellType"] = mesh.cell_type mesh_directory.attrs["Degree"] = mesh.degree mesh_directory.attrs["LagrangeVariant"] = mesh.lagrange_variant num_dofs_per_cell = mesh.local_topology.shape[1] topology_dataset = mesh_directory.create_dataset( "Topology", [mesh.num_cells_global, num_dofs_per_cell], dtype=np.int64 ) topology_dataset[slice(*mesh.local_topology_pos), :] = mesh.local_topology def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None = 0.0, read_from_partition: bool = False, backend_args: dict[str, Any] | None = None, ) -> ReadMeshData: """Read mesh data from h5py based checkpoint files. Args: filename: Path to input file comm: The MPI communciator to distribute the mesh over time: Time stamp associated with mesh read_from_partition: Read mesh with partition from file Returns: The mesh topology, geometry, UFL domain and partition function """ backend_args = get_default_backend_args(backend_args) with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: if "mesh" not in h5file.keys(): raise KeyError("Could not find mesh in file") mesh_group = h5file["mesh"] timestamps = mesh_group.attrs["timestamps"] assert time is not None parent_group = np.flatnonzero(np.isclose(timestamps, time)) if len(parent_group) != 1: raise RuntimeError( f"Time step {time} not found in file, available steps is {timestamps}" ) time_group = f"{parent_group[0]:d}" # Get mesh topology (distributed) topology = mesh_group["Topology"] local_range = compute_local_range(comm, topology.shape[0]) mesh_topology = topology[slice(*local_range), :] cell_type = mesh_group.attrs["CellType"] lvar = mesh_group.attrs["LagrangeVariant"] degree = mesh_group.attrs["Degree"] geometry_group = mesh_group[time_group] geometry_dataset = geometry_group["Points"] x_shape = geometry_dataset.shape geometry_range = compute_local_range(comm, x_shape[0]) mesh_geometry = geometry_dataset[slice(*geometry_range), :] # Check validity of partitioning information if read_from_partition: if "PartitionProcesses" not in mesh_group.attrs.keys(): raise KeyError(f"Partitioning information not found in {filename}") par_keys = ("PartitioningData", "PartitioningOffset") for key in par_keys: if key not in mesh_group.keys(): raise KeyError(f"Partitioning information not found in {filename}") par_num_procs = mesh_group.attrs["PartitionProcesses"] if par_num_procs != comm.size: raise ValueError(f"Number of processes in file ({par_num_procs})!=({comm.size=})") # First read in offsets based on the number of cells [0, num_cells_local] par_offsets = mesh_group["PartitioningOffset"][local_range[0] : local_range[1] + 1] # Then read the data based of offsets par_data = mesh_group["PartitioningData"][par_offsets[0] : par_offsets[-1]] # Then make offsets local par_offsets[:] -= par_offsets[0] partition_graph = adjacencylist(par_data, par_offsets.astype(np.int32)) else: partition_graph = None return ReadMeshData( cells=mesh_topology, cell_type=cell_type, x=mesh_geometry, degree=degree, lvar=lvar, partition_graph=partition_graph, ) def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None = None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ backend_args = get_default_backend_args(backend_args) with h5pyfile(filename, filemode="a", comm=comm, force_serial=False) as h5file: if "mesh" not in h5file.keys(): raise KeyError("Could not find mesh in file") mesh_group = h5file["mesh"] if "tags" not in mesh_group.keys(): tags = mesh_group.create_group("tags") else: tags = mesh_group["tags"] if data.name in tags.keys(): raise KeyError(f"MeshTags with {data.name=} already exists in this file") tag = tags.create_group(data.name) # Add topology topology = tag.create_dataset( "Topology", shape=[data.num_entities_global, data.num_dofs_per_entity], dtype=np.int64 ) assert data.local_start is not None topology[data.local_start : data.local_start + len(data.indices), :] = data.indices # Add cell_type attribute tag.attrs["CellType"] = data.cell_type # Add values values = tag.create_dataset( "Values", shape=[data.num_entities_global], dtype=data.values.dtype ) values[data.local_start : data.local_start + len(data.indices)] = data.values # Add dimension tag.attrs["dim"] = data.dim def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend. If "legacy_dolfin" is supplied as argument the HDF5 file is assumed to have been made with DOLFIN Returns: Internal data structure for the mesh tags read from file """ backend_args = get_default_backend_args(backend_args) legacy = backend_args["legacy"] with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: if legacy: if name not in h5file.keys(): raise RuntimeError(f"MeshTag {name} not found in {filename}.") mesh = h5file[name] topology = mesh["topology"] cell_type = topology.attrs["celltype"] if isinstance(cell_type, np.bytes_): cell_type = cell_type.decode("utf-8") dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type)) values = mesh["values"] else: if "mesh" not in h5file.keys(): raise KeyError("No mesh found") mesh = h5file["mesh"] if "tags" not in mesh.keys(): raise KeyError("Could not find 'tags' in file, are you sure this is a checkpoint?") tags = mesh["tags"] if name not in tags.keys(): raise KeyError(f"Could not find {name} in '/mesh/tags/' in {filename}") tag = tags[name] dim = tag.attrs["dim"] topology = tag["Topology"] values = tag["Values"] num_entities_global = topology.shape[0] topology_range = compute_local_range(comm, num_entities_global) indices = topology[slice(*topology_range), :].astype(np.int64) vals = values[slice(*topology_range)].astype(np.int32) return MeshTagsData(name=name, values=vals, indices=indices, dim=dim) def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` """ backend_args = {} if backend_args is None else backend_args with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: # If dofmap is read with full path, it is passed through backend_args dofmap_key = backend_args.get("dofmap", None) if dofmap_key is None: mesh_name = "mesh" # Prepare for multiple meshes if mesh_name not in h5file.keys(): raise KeyError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] if "functions" not in mesh.keys(): raise KeyError(f"No functions stored in '{mesh_name}' in {filename}") functions = mesh["functions"] if name not in functions.keys(): raise KeyError( f"No function with name '{name}' on '{mesh_name}' stored in {filename}" ) function = functions[name] offset_key = "dofmap_offsets" dofmap_key = "dofmap" offsets = function[offset_key] dofmap = function[dofmap_key] else: offset_key = backend_args["offsets"] dofmap = h5file[dofmap_key] offsets = h5file[offset_key] num_cells = offsets.shape[0] - 1 local_range = compute_local_range(comm, num_cells) # First read in offsets based on the number of cells [0, num_cells_local] glob_offsets = offsets[local_range[0] : local_range[1] + 1].flatten().astype(np.int64) # Then read the data based of offsets dofmap_data = dofmap[glob_offsets[0] : glob_offsets[-1]].flatten() # Then make offsets local loc_offsets = (glob_offsets - glob_offsets[0]).astype(np.int32) return adjacencylist(dofmap_data, loc_offsets) def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: mesh_name = "mesh" # Prepare for multiple meshes if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] if "functions" not in mesh.keys(): raise RuntimeError(f"No functions stored in '{mesh_name}' in {filename}") functions = mesh["functions"] if name not in functions.keys(): raise RuntimeError( f"No function with name '{name}' on '{mesh_name}' stored in {filename}" ) function = functions[name] timestamps = function.attrs["timestamps"] idx = np.flatnonzero(np.isclose(timestamps, time)) if len(idx) != 1: raise RuntimeError("Could not find {name}(t={time}) on grid {mesh_name} in {filename}.") u_t = function[f"{idx[0]:d}"] data_group = u_t["array"] num_dofs_global = data_group.shape[0] local_range = compute_local_range(comm, num_dofs_global) local_array = data_group[slice(*local_range)] return local_array, local_range[0] def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: mesh_name = "mesh" # Prepare for multiple meshes if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] data_group = mesh["CellPermutations"] num_cells_global = data_group.shape[0] local_range = compute_local_range(comm, num_cells_global) local_array = data_group[slice(*local_range)] return local_array def write_function( filename: str | Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None = None, ): """Write a function to file. Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ mesh_name = "mesh" # Prepare for multiple meshes backend_args = get_default_backend_args(backend_args) h5_mode = convert_file_mode(mode) with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: cell_permutations_exist = False dofmap_exists = False dofmap_offsets_exists = False if h5_mode == "a": if mesh_name not in h5file.keys(): mesh = h5file.create_group(mesh_name) else: mesh = h5file[mesh_name] cell_permutations_exist = "CellPermutations" in mesh.keys() if "functions" not in mesh.keys(): functions = mesh.create_group("functions") else: functions = mesh["functions"] if u.name not in functions.keys(): function = functions.create_group(u.name) else: function = functions[u.name] dofmap_exists = "dofmap" in function.keys() dofmap_offsets_exists = "dofmap_offsets" in function.keys() if not cell_permutations_exist: cell_perms = mesh.create_dataset( "CellPermutations", shape=[u.num_cells_global], dtype=np.uint32 ) cell_perms[slice(*u.local_cell_range)] = u.cell_permutations if not dofmap_exists: dofmap = function.create_dataset( "dofmap", shape=[u.global_dofs_in_dofmap], dtype=np.int64 ) dofmap[slice(*u.dofmap_range)] = u.dofmap_array if not dofmap_offsets_exists: dofmap_offsets = function.create_dataset( "dofmap_offsets", shape=[u.num_cells_global + 1], dtype=np.int64 ) dofmap_offsets[u.local_cell_range[0] : u.local_cell_range[1] + 1] = u.dofmap_offsets # Write timestamp if "timestamps" in function.attrs.keys(): timestamps = function.attrs["timestamps"] if np.isclose(time, timestamps).any(): raise RuntimeError("FUnction has already been stored at time={time_stamp}.") else: function.attrs["timestamps"] = np.append(function.attrs["timestamps"], time) else: function.attrs["timestamps"] = np.array([time]) idx = len(function.attrs["timestamps"]) - 1 data_group = function.create_group(f"{idx:d}") array = data_group.create_dataset("array", shape=[u.num_dofs_global], dtype=u.values.dtype) array[slice(*u.dof_range)] = u.values def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: if group not in h5file.keys(): raise KeyError(f"Could not find {group} in {filename}.") mesh = h5file[group] if "topology" not in mesh.keys(): raise KeyError(f"Could not find '{group}/topology' in {filename}.") topology = mesh["topology"] shape = topology.shape local_range = compute_local_range(comm, shape[0]) mesh_topology = topology[slice(*local_range)].astype(np.int64) # Get mesh cell type cell_type = None if "celltype" in topology.attrs.keys(): cell_type = topology.attrs["celltype"] if isinstance(cell_type, np.bytes_): cell_type = cell_type.decode("utf-8") for geometry_key in ["geometry", "coordinates"]: if geometry_key in mesh.keys(): break else: raise KeyError( "Could not find geometry in '/mesh/geometry' or '/mesh/coordinates'" + f" in {filename}." ) geometry = mesh[geometry_key] g_shape = geometry.shape local_g_range = compute_local_range(comm, g_shape[0]) mesh_geometry = geometry[slice(*local_g_range)] return mesh_topology, mesh_geometry, cell_type def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None = None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: data = h5file[group] shape = data.shape[0] local_range = compute_local_range(comm, shape) out_data = data[slice(*local_range)].flatten() return out_data, local_range[0] def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ comm = u.function_space.mesh.comm dofmap = u.function_space.dofmap local_range = np.array(dofmap.index_map.local_range) * dofmap.index_map_bs num_dofs_local = local_range[1] - local_range[0] num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs h5mode = convert_file_mode(mode) if h5mode == "w": with h5pyfile(filename, filemode=h5mode, comm=comm, force_serial=False) as h5file: local_dofs = u.x.array[:num_dofs_local].copy() data = h5file.create_group("snapshot") dataset = data.create_dataset("dofs", shape=num_dofs_global, dtype=local_dofs.dtype) dataset[slice(*local_range)] = local_dofs elif h5mode == "r": with h5pyfile(filename, filemode=h5mode, comm=comm, force_serial=False) as h5file: data = h5file["snapshot"]["dofs"] assert data.shape[0] == num_dofs_global u.x.array[:num_dofs_local] = data[slice(*local_range)] u.x.scatter_forward() def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ check_file_exists(filename) backend_args = get_default_backend_args(backend_args) with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: mesh_name = "mesh" # Prepare for multiple meshes if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] if "functions" not in mesh.keys(): return [] else: return list(mesh["functions"].keys()) def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ raise NotImplementedError("The h5py backend cannot read point data.") def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: """Read data from the cells of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: A tuple (topology, dofs) where topology contains the vertex indices of the cells, dofs the degrees of freedom within that cell. """ raise NotImplementedError("The h5py backend does not support reading cell data.") def write_data( filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file (distributed across proceses with MPI). Args: filename: Path to file array_data: Data to write to file comm: MPI communicator to open the file with time: Time-stamp for data. mode: Append or write backend_args: The backend arguments """ raise NotImplementedError("H5PY has not implemented this yet") scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/pyvista/000077500000000000000000000000001517634040500264475ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/pyvista/__init__.py000066400000000000000000000000551517634040500305600ustar00rootroot00000000000000from . import backend __all__ = ["backend"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/pyvista/backend.py000066400000000000000000000457301517634040500304210ustar00rootroot00000000000000""" Module that uses pyvista to import grids. """ from typing import Any import numpy as np import numpy.typing as npt try: import pyvista except ImportError: raise ModuleNotFoundError("The PyVista-backend requires pyvista") from pathlib import Path from mpi4py import MPI import basix import dolfinx from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from io4dolfinx.utils import check_file_exists from .. import FileMode, ReadMode # Cell types can be found at # https://vtk.org/doc/nightly/html/vtkCellType_8h_source.html _first_order_vtk = { 1: "point", 3: "interval", 5: "triangle", 9: "quadrilateral", 10: "tetrahedron", 12: "hexahedron", } _arbitrary_lagrange_vtk = { 68: "interval", 69: "triangle", 70: "quadrilateral", 71: "tetrahedron", 72: "hexahedron", 73: "prism", 74: "pyramid", } read_mode = ReadMode.serial def _cell_degree(ct: str, num_nodes: int) -> int: if ct == "point": return 1 elif ct == "interval": return int(num_nodes - 1) elif ct == "triangle": n = (np.sqrt(1 + 8 * num_nodes) - 1) / 2 if 2 * num_nodes != n * (n + 1): raise ValueError(f"Unknown triangle layout. Number of nodes: {num_nodes}") return int(n - 1) elif ct == "tetrahedron": n = 0 while n * (n + 1) * (n + 2) < 6 * num_nodes: n += 1 if n * (n + 1) * (n + 2) != 6 * num_nodes: raise ValueError(f"Unknown tetrahedron layout. Number of nodes: {num_nodes}") return int(n - 1) elif ct == "quadrilateral": n = np.sqrt(num_nodes) if num_nodes != n * n: raise ValueError(f"Unknown quadrilateral layout. Number of nodes: {num_nodes}") return int(n - 1) elif ct == "hexahedron": n = np.cbrt(num_nodes) if num_nodes != n * n * n: raise ValueError(f"Unknown hexahedron layout. Number of nodes: {num_nodes}") return int(n - 1) elif ct == "prism": if num_nodes == 6: return 1 elif num_nodes == 15: return 2 else: raise ValueError(f"Unknown prism layout. Number of nodes: {num_nodes}") elif ct == "pyramid": if num_nodes == 5: return 1 elif num_nodes == 13: return 2 else: raise ValueError(f"Unknown pyramid layout. Number of nodes: {num_nodes}") else: raise ValueError(f"Unknown cell type {ct} with {num_nodes=}.") def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default backend arguments given a set of input arguments. Args: arguments: Input backend arguments Returns: Updated backend arguments """ args = arguments or {} return args def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None = None, read_from_partition: bool = False, backend_args: dict[str, Any] | None = None, ) -> ReadMeshData: """Read mesh data from file. Args: filename: Path to file to read from comm: MPI communicator used in storage time: Time stamp associated with the mesh to read read_from_partition: Whether to read partition information backend_args: Arguments to backend Returns: Internal data structure for the mesh data read from file """ backend_args = get_default_backend_args(backend_args) check_file_exists(filename) if read_from_partition: raise RuntimeError("Cannot read partition data with Pyvista") cells: npt.NDArray[np.int64] geom: npt.NDArray[np.float64 | np.float32] if comm.rank == 0: in_data = pyvista.read(filename) if isinstance(in_data, pyvista.UnstructuredGrid): grid = in_data elif isinstance(in_data, pyvista.core.composite.MultiBlock): # To handle multiblock like pvd if hasattr(pyvista, "_VTK_SNAKE_CASE_STATE"): pyvista._VTK_SNAKE_CASE_STATE = "allow" else: # Compatibility with 0.47 pyvista.core.vtk_snake_case._state = "allow" number_of_blocks = in_data.number_of_blocks assert number_of_blocks == 1 b0 = in_data.get_block(0) assert isinstance(b0, pyvista.UnstructuredGrid) grid = b0 else: raise RuntimeError(f"Unknown data type {type(in_data)}") geom = grid.points num_cells_global = grid.number_of_cells cells = grid.cells.reshape(num_cells_global, -1).astype(np.int64) nodes_per_cell_type = cells[:, 0] assert np.allclose(nodes_per_cell_type, nodes_per_cell_type[0]), "Single celltype support" cells = cells[:, 1:].astype(np.int64) cell_types = grid.celltypes assert len(np.unique(cell_types)) == 1 if (cell_type := cell_types[0]) in _first_order_vtk.keys(): cell_type = _first_order_vtk[cell_type] order = 1 elif cell_type in _arbitrary_lagrange_vtk.keys(): cell_type = _arbitrary_lagrange_vtk[cell_type] order = _cell_degree(cell_type, cells.shape[1]) perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(cell_type), cells.shape[1]) cells = cells[:, perm] lvar = int(basix.LagrangeVariant.equispaced) gtype = backend_args.get("dtype", geom.dtype) order, lvar, nodes_per_cell, cell_type, gtype, gdim = comm.bcast( (order, lvar, cells.shape[1], cell_type, gtype, geom.shape[1]), root=0 ) else: order, lvar, nodes_per_cell, cell_type, gtype, gdim = comm.bcast(None, root=0) geom = np.zeros((0, gdim), dtype=gtype) cells = np.zeros((0, nodes_per_cell), dtype=np.int64) return ReadMeshData( cells=cells, cell_type=cell_type, x=geom.astype(gtype), lvar=lvar, degree=order ) def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ dataset: np.ndarray if MPI.COMM_WORLD.rank == 0: in_data = pyvista.read(filename) if isinstance(in_data, pyvista.UnstructuredGrid): grid = in_data elif isinstance(in_data, pyvista.core.composite.MultiBlock): # To handle multiblock like pvd if hasattr(pyvista, "_VTK_SNAKE_CASE_STATE"): pyvista._VTK_SNAKE_CASE_STATE = "allow" else: # Compatibility with 0.47 pyvista.core.vtk_snake_case._state = "allow" number_of_blocks = in_data.number_of_blocks assert number_of_blocks == 1 b0 = in_data.get_block(0) assert isinstance(b0, pyvista.UnstructuredGrid) grid = b0 dataset = grid.point_data[name] if len(dataset.shape) == 1: num_components = 1 dataset = dataset.reshape(-1, num_components) else: num_components = dataset.shape[1] if np.issubdtype(dataset.dtype, np.integer): gtype = grid.points.dtype dataset = dataset.astype(gtype) else: gtype = dataset.dtype num_components, gtype = comm.bcast((num_components, gtype), root=0) local_range_start = 0 else: num_components, gtype = comm.bcast(None, root=0) dataset = np.zeros((0, num_components), dtype=gtype) local_range_start = 0 return dataset, int(local_range_start) def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: dataset: np.ndarray topology: np.ndarray if MPI.COMM_WORLD.rank == 0: in_data = pyvista.read(filename) if isinstance(in_data, pyvista.UnstructuredGrid): grid = in_data elif isinstance(in_data, pyvista.core.composite.MultiBlock): # To handle multiblock like pvd if hasattr(pyvista, "_VTK_SNAKE_CASE_STATE"): pyvista._VTK_SNAKE_CASE_STATE = "allow" else: # Compatibility with 0.47 pyvista.core.vtk_snake_case._state = "allow" number_of_blocks = in_data.number_of_blocks assert number_of_blocks == 1 b0 = in_data.get_block(0) assert isinstance(b0, pyvista.UnstructuredGrid) grid = b0 dataset = grid.cell_data[name] if len(dataset.shape) == 1: num_components = 1 dataset = dataset.reshape(-1, num_components) else: num_components = dataset.shape[1] if np.issubdtype(dataset.dtype, np.integer): gtype = in_data.points.dtype dataset = dataset.astype(gtype) else: gtype = dataset.dtype num_components, gtype = comm.bcast((num_components, gtype), root=0) else: num_components, gtype = comm.bcast(None, root=0) dataset = np.zeros((0, num_components), dtype=gtype) _time = float(time) if time is not None else None topology = read_mesh_data(filename, comm, _time, False, backend_args=None).cells return topology, dataset def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None, ): """Write attributes to file. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attribute group attributes: Dictionary of attributes to write backend_args: Arguments to backend """ raise NotImplementedError("The Pyvista backend cannot write attributes.") def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dict[str, Any]: """Read attributes from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attribute group backend_args: Arguments to backend Returns: Dictionary of attributes read from file """ raise NotImplementedError("The Pyvista backend cannot read attributes.") def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read timestamps from file. Args: filename: Path to file to read from comm: MPI communicator used in storage function_name: Name of the function to read timestamps for backend_args: Arguments to backend Returns: Numpy array of timestamps read from file """ raise NotImplementedError("The Pyvista backend cannot read timestamps.") def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ in_data = pyvista.read(filename) if isinstance(in_data, pyvista.UnstructuredGrid): grid = in_data elif isinstance(in_data, pyvista.core.composite.MultiBlock): # To handle multiblock like pvd if hasattr(pyvista, "_VTK_SNAKE_CASE_STATE"): pyvista._VTK_SNAKE_CASE_STATE = "allow" else: # Compatibility with 0.47 pyvista.core.vtk_snake_case._state = "allow" number_of_blocks = in_data.number_of_blocks assert number_of_blocks == 1 b0 = in_data.get_block(0) assert isinstance(b0, pyvista.UnstructuredGrid) grid = b0 point_data = list(grid.point_data.keys()) cell_data = list(grid.cell_data.keys()) return point_data + cell_data def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None, mode: FileMode, time: float, ): """ Write a mesh to file. Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to backend_args: Arguments to backend mode: File-mode to store the mesh time: Time stamp associated with the mesh """ raise NotImplementedError("The Pyvista backend cannot write meshes.") def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ raise NotImplementedError("The Pyvista backend cannot write meshtags.") def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend Returns: Internal data structure for the mesh tags read from file """ raise NotImplementedError("The Pyvista backend cannot read meshtags.") def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an AdjacencyList """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def write_function( filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None, ): """ Write a function to file. Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ raise NotImplementedError("The Pyvista backend cannot read legacy DOLFIN meshes.") def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Pyvista backend cannot read HDF5 arrays") def write_data( filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file (distributed across proceses with MPI). Args: filename: Path to file array_data: Data to write to file comm: MPI communicator to open the file with time: Time stamp mode: Append or write backend_args: The backend arguments """ raise NotImplementedError("The pyvista backend does not support writing point data") scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/vtkhdf/000077500000000000000000000000001517634040500262365ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/vtkhdf/__init__.py000066400000000000000000000000551517634040500303470ustar00rootroot00000000000000from . import backend __all__ = ["backend"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/vtkhdf/backend.py000066400000000000000000001345201517634040500302040ustar00rootroot00000000000000""" Module that can read the VTKHDF format using h5py. """ from pathlib import Path from typing import Any from mpi4py import MPI import basix import dolfinx import h5py import numpy as np import numpy.typing as npt from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from io4dolfinx.utils import check_file_exists, compute_local_range from .. import FileMode, ReadMode from ..h5py.backend import convert_file_mode, h5pyfile from ..pyvista.backend import _arbitrary_lagrange_vtk, _cell_degree, _first_order_vtk read_mode = ReadMode.parallel _vtk_hdf_version = np.array([2, 1], dtype=np.int32) def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default backend arguments given a set of input arguments. Args: arguments: Input backend arguments Returns: Updated backend arguments """ args = arguments or {"name": "mesh"} return args def find_all_unique_cell_types(comm, cell_types, num_nodes): """ Given a set of cell types and number of nodes per cell, find all unique cell types across all ranks. Args: comm: MPI communicator cell_types: Local cell types num_nodes: Number of nodes per cell Returns: A 2D array where each row corresponds to a cell type (vtk int) and the number of nodes. """ # Combine cell_types, num_nodes as tuple c_hash = np.zeros((2, len(cell_types)), dtype=np.int32) c_hash[0] = cell_types c_hash[1] = num_nodes indexes = np.unique(c_hash.T, axis=0, return_index=True)[1] local_unique_cells = c_hash.T[indexes] all_cell_types = np.vstack(comm.allgather(local_unique_cells)) indexes = np.unique(all_cell_types, axis=0, return_index=True)[1] all_unique_cell_types = all_cell_types[indexes] return all_unique_cell_types def _decode_bytes_if_needed(value: bytes | str) -> str: """Decode bytes to string if necessary (for h5py compatibility)""" if isinstance(value, bytes): return value.decode("utf-8") return value def _get_vtk_group(h5file, name: str) -> h5py.Group: """ Navigates the VTKHDF group hierarchy to find the specific first UnstructuredGrid with a specific name. Handles both MultiBlockDataSet and direct UnstructuredGrid types. """ hdf = h5file["VTKHDF"] file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) if file_type == "MultiBlockDataSet": ass = hdf["Assembly"] def visitor(path): mesh_group = path.rsplit("/", 1) mesh_name = mesh_group[0] if len(mesh_group) == 1 else mesh_group[1] if mesh_name == name: obj = ass.get(path) # Check attributes carefully if "Type" in obj.attrs.keys(): attr_type = obj.attrs["Type"] if isinstance(attr_type, bytes): attr_type = attr_type.decode("utf-8") if attr_type == "UnstructuredGrid": return path return None mesh_node = ass.visit_links(visitor) if mesh_node is None: raise RuntimeError(f"Could not find unique mesh named '{name}' in Assembly.") return ass[mesh_node] elif file_type == "UnstructuredGrid": return hdf else: raise RuntimeError(f"Not supported file type {file_type}") def _get_time_index(hdf: h5py.Group, time: float | str, filename: str | Path) -> int: """Finds the index of a specific time stamp.""" if "Steps" not in hdf.keys(): raise RuntimeError(f"No timestepping information found in {filename}.") stamps = hdf["Steps"]["Values"][:] pos = np.flatnonzero(np.isclose(stamps, time)) if len(pos) == 0: raise RuntimeError(f"Could not find mesh at t={time} in {filename}.") elif len(pos) > 1: raise RuntimeError(f"Multiple time steps for mesh at t={time} in {filename}") return pos[0] def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None, read_from_partition: bool, backend_args: dict[str, Any] | None, ) -> ReadMeshData: """Read mesh data from file. Args: filename: Path to file to read from comm: MPI communicator used in storage time: Time stamp associated with the mesh to read read_from_partition: Whether to read partition information backend_args: Arguments to backend Returns: Internal data structure for the mesh data read from file """ backend_args = get_default_backend_args(backend_args) check_file_exists(filename) if read_from_partition: raise RuntimeError("Cannot read partition data with VTKHDF") with h5pyfile(filename, "r", comm=comm) as h5file: hdf = _get_vtk_group(h5file, backend_args["name"]) if time is None: num_cells_global = hdf["Types"].size local_cell_range = compute_local_range(comm, num_cells_global) cell_types_local = hdf["Types"][slice(*local_cell_range)] num_points_global = hdf["NumberOfPoints"][0] local_point_range = compute_local_range(comm, num_points_global) points_local = hdf["Points"][slice(*local_point_range)] # Connectivity read offsets = hdf["Offsets"] local_connectivity_offset = offsets[local_cell_range[0] : local_cell_range[1] + 1] topology = hdf["Connectivity"][ local_connectivity_offset[0] : local_connectivity_offset[-1] ] offset = local_connectivity_offset - local_connectivity_offset[0] else: time_index = _get_time_index(hdf, time, filename) stamps = hdf["Steps"]["Values"][:] # Get number of points point_node = hdf["Points"] step_start = hdf["Steps"]["PointOffsets"][time_index] # NOTE: currently, it doesn't seem like we follow: # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata # As only one num_points is stored irregardless of time steps added. if hdf["NumberOfPoints"].shape[0] != len(stamps): num_pts = hdf["NumberOfPoints"][0] else: num_pts = hdf["NumberOfPoints"][time_index] lr = compute_local_range(comm, num_pts) points_local = point_node[step_start + lr[0] : step_start + lr[1]] # Get cell-types in step cell_start = hdf["Steps"]["CellOffsets"][time_index] if hdf["NumberOfCells"].shape[0] != len(stamps): num_cells = hdf["NumberOfCells"][0] else: num_cells = hdf["NumberOfCells"][time_index] local_cell_range = compute_local_range(comm, num_cells) cell_types_local = hdf["Types"][ cell_start + local_cell_range[0] : cell_start + local_cell_range[1] ] # Get connectivity in step connectivity_start = hdf["Steps"]["ConnectivityIdOffsets"][time_index] # Connectivity read offsets = hdf["Offsets"] local_connectivity_offset = offsets[ connectivity_start + local_cell_range[0] : connectivity_start + local_cell_range[1] + 1 ] topology = hdf["Connectivity"][ local_connectivity_offset[0] : local_connectivity_offset[-1] ] offset = local_connectivity_offset - local_connectivity_offset[0] # NOTE: Currently we limit ourselfs to a single celltype, as it makes life easier, # other things have to change in `MeshReadData` to support this. num_nodes_per_cell = offset[1:] - offset[:-1] unique_cells = find_all_unique_cell_types(MPI.COMM_WORLD, cell_types_local, num_nodes_per_cell) if unique_cells.shape[0] > 1: raise NotImplementedError("io4dolfinx does not support mixed celltype grids") topology = topology.reshape(-1, num_nodes_per_cell[0]) cell_type, number_of_nodes = unique_cells[0] gtype = backend_args.get("dtype", points_local.dtype) if cell_type in _first_order_vtk.keys(): ct = _first_order_vtk[cell_type] degree = 1 elif cell_type in _arbitrary_lagrange_vtk.keys(): ct = _arbitrary_lagrange_vtk[cell_type] degree = _cell_degree(ct, num_nodes=number_of_nodes) else: raise ValueError(f"Unknown VTK cell type {cell_type} in {filename}") perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(ct), number_of_nodes) topology = topology[:, perm] lvar = int(basix.LagrangeVariant.equispaced) return ReadMeshData( cells=topology, cell_type=ct, x=points_local.astype(gtype), lvar=lvar, degree=degree ) def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ backend_args = get_default_backend_args(backend_args) check_file_exists(filename) with h5pyfile(filename, "r", comm=comm) as h5file: hdf = _get_vtk_group(h5file, backend_args["name"]) point_data = hdf["PointData"] assert point_data is not None if name not in point_data.keys(): raise ValueError(f"No point data named {name} in {filename}.") func_node = point_data[name] if time is None: data_shape = func_node.shape[0] lr = compute_local_range(comm, data_shape) data = func_node[slice(*lr)] return data, lr[0] else: time_index = _get_time_index(hdf, time, filename) stamps = hdf["Steps"]["Values"][:] step_start = hdf["Steps"]["PointDataOffsets"][name][time_index] # NOTE: currently, it doesn't seem like we follow: # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata # As only one num_points is stored irregardless of time steps added. if hdf["NumberOfPoints"].shape[0] != len(stamps): num_pts = hdf["NumberOfPoints"][0] else: num_pts = hdf["NumberOfPoints"][time_index] lr = compute_local_range(comm, num_pts) return func_node[step_start + lr[0] : step_start + lr[1]], lr[0] def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: backend_args = get_default_backend_args(backend_args) check_file_exists(filename) with h5pyfile(filename, "r", comm=comm) as h5file: hdf = _get_vtk_group(h5file, backend_args["name"]) if "CellData" not in hdf.keys(): raise RuntimeError(f"No cell data found in {filename}.") cell_data = hdf["CellData"] assert cell_data is not None if name not in cell_data.keys(): raise ValueError(f"No cell data with name {name} in {filename}") cell_data_node = cell_data[name] assert cell_data_node is not None if time is None: cell_data_shape = cell_data_node.shape num_cells_global = hdf["Types"].size assert num_cells_global == cell_data_shape[0] local_cell_range = compute_local_range(comm, num_cells_global) data = cell_data_node[slice(*local_cell_range)] else: time_index = _get_time_index(hdf, time, filename) stamps = hdf["Steps"]["Values"][:] cd_start = hdf["Steps"]["CellDataOffsets"][name][time_index] # NOTE: currently, it doesn't seem like we follow: # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata # As only one num_points is stored irregardless of time steps added. if hdf["NumberOfCells"].shape[0] != len(stamps): number_of_cells = hdf["NumberOfCells"][0] else: number_of_cells = hdf["NumberOfCells"][time_index] lr = compute_local_range(comm, number_of_cells) data = cell_data_node[cd_start + lr[0] : cd_start + lr[1]] # NOTE: THis could be optimized by hand-coding some communication in # `read_cell_data` on the frontend side md = read_mesh_data( filename, comm, time=time, read_from_partition=False, backend_args=backend_args ) if len(data.shape) == 1: data = data.reshape(-1, 1) return md.cells, data def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None, ): """Write attributes to file. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attribute group attributes: Dictionary of attributes to write backend_args: Arguments to backend """ raise NotImplementedError("The Pyvista backend cannot write attributes.") def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dict[str, Any]: """Read attributes from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attribute group backend_args: Arguments to backend Returns: Dictionary of attributes read from file """ raise NotImplementedError("The Pyvista backend cannot read attributes.") def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read timestamps from file. Args: filename: Path to file to read from comm: MPI communicator used in storage function_name: Name of the function to read timestamps for backend_args: Arguments to backend Returns: Numpy array of timestamps read from file """ backend_args = get_default_backend_args(backend_args) check_file_exists(filename) # Temporal data storage as described in # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-data with h5pyfile(filename, "r", comm=comm) as h5file: hdf = h5file["VTKHDF"] if "Steps" in hdf.keys(): timestamps = hdf["Steps"]["Values"][:] # For either point data or cell data, check if function exists, # and check if offsets in time are changing between steps. if "CellData" in hdf.keys() and function_name in hdf["CellData"].keys(): offsets = hdf["Steps"]["CellDataOffsets"][function_name] step_offsets = offsets[:] step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1]) return timestamps[step_diff] elif "PointData" in hdf.keys() and function_name in hdf["PointData"].keys(): offsets = hdf["Steps"]["PointDataOffsets"][function_name] step_offsets = offsets[:] step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1]) # This only finds when the offset changes, does not capture first step return np.hstack([[timestamps[0]], timestamps[step_diff]]) else: raise RuntimeError(f"Function {function_name} is not assoicated with a time-stamp.") else: raise RuntimeError(f"{filename} does not contain time-stepping information.") def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ backend_args = get_default_backend_args(backend_args) check_file_exists(filename) with h5pyfile(filename, "r", comm=comm) as h5file: hdf = h5file["VTKHDF"] function_names = set() if "CellData" in hdf.keys(): for item in hdf["CellData"].keys(): function_names.add(item) if "PointData" in hdf.keys(): for item in hdf["PointData"].keys(): function_names.add(item) return list(function_names) def _create_dataset( root: h5py.File | h5py.Group, name: str, shape: tuple[int, ...], dtype: npt.DTypeLike, chunks: bool, maxshape: tuple[int | None, ...], mode: str, resize: bool = True, ) -> h5py.Dataset: if name not in root.keys(): mode = "w" if mode == "w": dataset = root.create_dataset( name, shape=shape, dtype=dtype, chunks=chunks, maxshape=maxshape ) elif mode == "a": dataset = root[name] old_shape = dataset.shape # Only resize for dimension if resize: if len(old_shape) == 1: new_shape = (old_shape[0] + shape[0],) else: new_shape = (old_shape[0] + shape[0], *old_shape[1:]) dataset.resize(new_shape) else: raise ValueError(f"Unknown file mode '{mode}' when creating dataset {name} in {root}") return dataset def _create_group(root: h5py.File | h5py.Group, name: str, mode: str) -> h5py.Group: if name not in root.keys(): mode = "w" if mode == "w": # Track order has to be on to make multiblock work: # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#partitioneddatasetcollection-and-multiblockdataset group = root.create_group(name, track_order=True) elif mode == "a": group = root[name] else: raise ValueError("Unknown file mode '{h5_mode}'") return group def _compute_append_slice( dataset: h5py.Dataset, input_size: int, original_slice: tuple[int, int] | np.ndarray, mode: str ) -> slice: append_offset = dataset.shape[0] - input_size if mode == "a" else 0 return slice(*(np.asarray(original_slice) + append_offset).astype(np.int64)) def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None, mode: FileMode, time: float, ): """ Write a mesh to file. Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to backend_args: Arguments to backend mode: File-mode to store the mesh time: Time stamp associated with the mesh """ h5_mode = convert_file_mode(mode) backend_args = get_default_backend_args(backend_args) name = backend_args["name"] with h5pyfile(filename, h5_mode, comm=comm) as h5file: hdf = _create_group(h5file, "/VTKHDF", h5_mode) hdf.attrs.create("Type", "MultiBlockDataSet") hdf.attrs["Version"] = _vtk_hdf_version mesh_group = _create_group(hdf, name, h5_mode) mesh_group.attrs.create("Type", "UnstructuredGrid") mesh_group.attrs["Version"] = _vtk_hdf_version assembly = _create_group(hdf, "Assembly", h5_mode) mesh_assembly = _create_group(assembly, name, h5_mode) if name not in mesh_assembly.keys(): mesh_assembly[name] = h5py.SoftLink(f"/VTKHDF/{name}") # Write time dependent points number_of_points = _create_dataset( mesh_group, "NumberOfPoints", shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, ) number_of_points[-1] = mesh.num_nodes_global # Store nodes points = _create_dataset( mesh_group, "Points", shape=(mesh.num_nodes_global, 3), dtype=mesh.local_geometry.dtype, chunks=True, maxshape=(None, 3), mode=h5_mode, ) insert_slice = _compute_append_slice( points, mesh.num_nodes_global, mesh.local_geometry_pos, h5_mode ) points[insert_slice, : mesh.local_geometry.shape[1]] = mesh.local_geometry # NOTE: VTKHDF currently does not support reading time dependent topology in # Paraview: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 # Therefore the following data is not resized time_indep_datasets = {} for key in ["NumberOfCells", "NumberOfConnectivityIds"]: time_indep_datasets[key] = _create_dataset( mesh_group, key, shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True ) num_dofs_per_cell = mesh.local_topology.shape[1] time_indep_datasets["NumberOfCells"][-1] = mesh.num_cells_global time_indep_datasets["NumberOfConnectivityIds"][-1] = ( mesh.num_cells_global * num_dofs_per_cell ) # NOTE: The following offsets are currently overwriting the existing topology data. # This is due to the VTKHDF bug. When we switch resize=True this will automatically work # for time dependent topologies # Store topology offsets (single celltype assumption) offsets = _create_dataset( mesh_group, "Offsets", shape=(mesh.num_cells_global + 1,), dtype=np.int64, chunks=True, mode=h5_mode, maxshape=(None,), resize=False, # Resize should really be True ) offset_data = np.arange(0, mesh.local_topology.size + 1, mesh.local_topology.shape[1]) offset_data += num_dofs_per_cell * mesh.local_topology_pos[0] insert_slice = _compute_append_slice( offsets, mesh.num_cells_global + 1, (mesh.local_topology_pos[0], mesh.local_topology_pos[1] + 1), mode=h5_mode, ) offsets[insert_slice] = offset_data del offset_data # Permute and store topology data dx_ct = dolfinx.mesh.to_type(mesh.cell_type) top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell)) topology_data = mesh.local_topology[:, top_perm].flatten() topology = _create_dataset( mesh_group, "Connectivity", shape=(mesh.num_cells_global * num_dofs_per_cell,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( topology, mesh.num_cells_global * num_dofs_per_cell, np.array(mesh.local_topology_pos) * num_dofs_per_cell, mode=h5_mode, ) topology[insert_slice] = topology_data del topology_data # Store celltypes cell_types = np.full( mesh.local_topology.shape[0], dolfinx.cpp.io.get_vtk_cell_type(dx_ct, dolfinx.mesh.cell_dim(dx_ct)), ) types = _create_dataset( mesh_group, "Types", shape=(mesh.num_cells_global,), dtype=np.uint8, maxshape=(None,), chunks=True, mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( types, mesh.num_cells_global, mesh.local_topology_pos, h5_mode ) types[insert_slice] = cell_types del cell_types steps = _create_group(mesh_group, "Steps", mode=h5_mode) # First fetch time-steps to see if we have stored this timestep already values = _create_dataset( steps, "Values", shape=(1,), dtype=np.float64, chunks=True, maxshape=(None,), mode=h5_mode, resize=h5_mode == "a", ) if h5_mode == "w": values[0] = time else: existing_steps = values[:-1] if len(np.flatnonzero(np.isclose(existing_steps, time))) > 0: raise RuntimeError(f"Mesh already exists at time {time} in {filename}.") values[-1] = time steps.attrs.create("NSteps", np.int64(len(values)), dtype=np.int64) # Write offset data for current time-step all_parts = {} for key in [ "NumberOfParts", "PartOffsets", "PointOffsets", "CellOffsets", "ConnectivityIdOffsets", ]: all_parts[key] = _create_dataset( steps, key, shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, ) all_parts["NumberOfParts"][-1] = 1 all_parts["PartOffsets"][-1] = 0 all_parts["PointOffsets"][-1] = points.shape[0] - mesh.num_nodes_global all_parts["CellOffsets"][-1] = types.shape[0] - mesh.num_cells_global all_parts["ConnectivityIdOffsets"][-1] = ( topology.shape[0] - mesh.num_cells_global * num_dofs_per_cell ) # Update cell-data and point-data offsets by copying over data from previous step for key in ["CellDataOffsets", "PointDataOffsets"]: group = _create_group(steps, key, mode=h5_mode) for cd in group.keys(): group[cd].resize(steps.attrs["NSteps"], axis=0) group[cd][-1] = group[cd][-2] # Update Steps in all other parts of the mesh as well for key in mesh_assembly.keys(): if key == name: continue # Copy time-dependent geometry info (NumberOfPoints) from mesh to tag sub_group = mesh_assembly[key] sub_step = sub_group["Steps"] sub_step.attrs["NSteps"] = steps.attrs["NSteps"] # Copy time dependent and partition info from mesh to tag step_copy_keys = ["Values", "PartOffsets", "NumberOfParts"] for key in step_copy_keys: if key in sub_step.keys(): sub_step[key].resize(steps[key].shape) sub_step[key][:] = steps[key][:] else: raise RuntimeError(f"{sub_step.name} should have {key}/") # Append value from previous step for meshtags as they are time-independent append_keys = ["CellOffsets", "ConnectivityIdOffsets"] for key in append_keys: if key in sub_step.keys(): sub_step[key].resize(sub_step[key].size + 1, axis=0) sub_step[key][-1] = sub_step[key][-2] else: raise RuntimeError(f"{sub_step.name} should have {key}/") # Append value from previous step for meshtags celldata for key in ["CellDataOffsets", "PointDataOffsets"]: group = _create_group(sub_step, key, mode=h5_mode) for cd in group.keys(): group[cd].resize(sub_step.attrs["NSteps"], axis=0) group[cd][-1] = group[cd][-2] def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ backend_args = get_default_backend_args(backend_args) name = backend_args["name"] h5_mode = "a" with h5pyfile(filename, h5_mode, comm=comm) as h5file: hdf = h5file["VTKHDF"] file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) # Check for type of VTKHDF file if file_type != "MultiBlockDataSet": raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}") # We place the mesh-tags in the same subgroup of assembly as the mesh parent_mesh_group = _get_vtk_group(h5file, name) block = parent_mesh_group.parent tag_path = f"/VTKHDF/{name}_{data.name}" if data.name not in block.keys(): block[data.name] = h5py.SoftLink(tag_path) mesh_group = _create_group(hdf, tag_path, mode=h5_mode) mesh_group.attrs.create("Type", "UnstructuredGrid") mesh_group.attrs["Version"] = _vtk_hdf_version cell_data = _create_group(mesh_group, "CellData", mode=h5_mode) assert data.num_entities_global is not None dataset = _create_dataset( cell_data, data.name, shape=(data.num_entities_global,), dtype=data.values.dtype, chunks=True, maxshape=(None,), mode=h5_mode, ) assert data.local_start is not None insert_slice = _compute_append_slice( dataset, data.num_entities_global, np.array([data.local_start, data.local_start + data.indices.shape[0]]), mode=h5_mode, ) dataset[insert_slice] = data.values # NOTE: The following is more or less a copy from write_mesh, # except that we pull out the point storage and use a softlink num_cells = _create_dataset( mesh_group, "NumberOfCells", shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 num_cells[-1] = data.num_entities_global # Hardlink data should also follow hardlink for numbering for key in ["Points", "NumberOfPoints"]: if key not in mesh_group.keys(): mesh_group[key] = parent_mesh_group[key] # Single celltype assumption num_dofs_per_cell = data.num_dofs_per_entity number_of_connectivities = _create_dataset( mesh_group, "NumberOfConnectivityIds", shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 assert data.num_entities_global is not None and num_dofs_per_cell is not None number_of_connectivities[-1] = data.num_entities_global * num_dofs_per_cell # Store topology offsets (single celltype assumption) offsets = _create_dataset( mesh_group, "Offsets", shape=(data.num_entities_global + 1,), dtype=np.int64, chunks=True, mode=h5_mode, maxshape=(None,), resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 offset_data = np.arange(0, data.indices.size + 1, data.indices.shape[1]) assert data.local_start is not None offset_data += num_dofs_per_cell * data.local_start insert_slice = _compute_append_slice( offsets, data.num_entities_global + 1, (data.local_start, data.local_start + data.indices.shape[0] + 1), mode=h5_mode, ) offsets[insert_slice] = offset_data del offset_data # Permute and store topology data dx_ct = dolfinx.mesh.to_type(data.cell_type) top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell)) topology_data = data.indices[:, top_perm].flatten() topology = _create_dataset( mesh_group, "Connectivity", shape=(data.num_entities_global * num_dofs_per_cell,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( topology, data.num_entities_global * num_dofs_per_cell, np.array([data.local_start, data.local_start + data.indices.shape[0]]) * num_dofs_per_cell, mode=h5_mode, ) topology[insert_slice] = topology_data del topology_data # Store celltypes cell_types = np.full( data.indices.shape[0], dolfinx.cpp.io.get_vtk_cell_type(dx_ct, dolfinx.mesh.cell_dim(dx_ct)), ) types = _create_dataset( mesh_group, "Types", shape=(data.num_entities_global,), dtype=np.uint8, maxshape=(None,), chunks=True, mode=h5_mode, resize=False, # Resize should really be True, see issue below ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( types, data.num_entities_global, (data.local_start, data.local_start + data.indices.shape[0]), h5_mode, ) types[insert_slice] = cell_types del cell_types steps = _create_group(mesh_group, "Steps", mode=h5_mode) # Copy n-step counter steps.attrs.create("NSteps", parent_mesh_group["Steps"].attrs["NSteps"]) # Hardlink data that we know is the same across meshes hardlink_keys = ["NumberOfParts", "PartOffsets", "Values", "PointOffsets"] for key in hardlink_keys: steps[key] = parent_mesh_group["Steps"][key] # Write offset data for current time-step all_parts = {} for key in ["CellOffsets", "ConnectivityIdOffsets"]: all_parts[key] = _create_dataset( steps, key, shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, ) all_parts["CellOffsets"][-1] = types.shape[0] - data.num_entities_global all_parts["ConnectivityIdOffsets"][-1] = offsets.shape[0] - (data.num_entities_global + 1) # CellData requires an offset cd_off = _create_group(steps, "CellDataOffsets", mode=h5_mode) cd_data = _create_dataset( cd_off, data.name, shape=parent_mesh_group["Steps"]["CellOffsets"].shape, dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, ) cd_data[-1] = types.shape[0] - data.num_entities_global def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend Returns: Internal data structure for the mesh tags read from file """ backend_args = get_default_backend_args(backend_args) backend_args.update({"name": name}) # Reuse reading cell-data indices, values = read_cell_data(filename, name, comm, None, backend_args=backend_args) # Read cell-type of grid to get topological dimension with h5pyfile(filename, "r", comm=comm) as h5file: hdf = _get_vtk_group(h5file, backend_args["name"]) num_cells_global = hdf["Types"].size local_cell_range = compute_local_range(comm, num_cells_global) cell_types_local = hdf["Types"][slice(*local_cell_range)] unique_cells = find_all_unique_cell_types(comm, cell_types_local, indices.shape[1]) if unique_cells.shape[0] > 1: raise NotImplementedError("io4dolfinx does not support mixed celltype grids") vtk_cell_type = unique_cells[0][0] if vtk_cell_type in _first_order_vtk.keys(): ct = _first_order_vtk[vtk_cell_type] elif vtk_cell_type in _arbitrary_lagrange_vtk.keys(): ct = _arbitrary_lagrange_vtk[vtk_cell_type] dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(ct)) return MeshTagsData(name=name, values=values.flatten(), indices=indices, dim=dim) def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an AdjacencyList """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The Pyvista backend cannot make checkpoints.") def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") def write_function( filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None, ): """ Write a function to file. Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ raise NotImplementedError("The VTKHDF backend cannot read legacy DOLFIN meshes.") def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The VTKHDF backend cannot read HDF5 arrays") def write_data( filename: Path | str, array_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write function to file by interpolating into geometry nodes. Args: filename: Path to file array_data: Data to write to file time: Time stamp mode: Append or write backend_args: The backend arguments """ h5_mode = convert_file_mode(mode) assert h5_mode == "a" backend_args = get_default_backend_args(backend_args) mesh_name = backend_args["name"] extension = array_data.type with h5pyfile(filename, h5_mode, comm=comm) as h5file: hdf = h5file["VTKHDF"] file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) # Check for type of VTKHDF file if file_type != "MultiBlockDataSet": raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}") # Find mesh block to add data to block = _get_vtk_group(h5file, mesh_name) data_group = _create_group(block, f"{extension}Data", mode=h5_mode) dataset = _create_dataset( data_group, array_data.name, shape=array_data.global_shape, dtype=array_data.values.dtype, chunks=True, maxshape=(None, array_data.values.shape[1]), mode=h5_mode, ) insert_slice = _compute_append_slice( dataset, array_data.global_shape[0], np.array( [ array_data.local_range[0], array_data.local_range[0] + array_data.values.shape[0], ] ), mode=h5_mode, ) dataset[insert_slice] = array_data.values steps = _create_group(block, "Steps", mode=h5_mode) pdo = _create_group(steps, f"{extension}DataOffsets", mode=h5_mode) # Check if time step is already in time-stepping of mesh timestamps = steps["Values"][:] assert isinstance(timestamps, np.ndarray) assert isinstance(time, float) time_exists = np.flatnonzero(np.isclose(timestamps, time)) if len(time_exists) > 1: raise ValueError("Time-step has been written multiple times") pdo_u = _create_dataset( pdo, array_data.name, shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,), mode=h5_mode, resize=False, ) if len(time_exists) == 1: idx = time_exists[0] elif len(time_exists) == 0: # No mesh written at step, update mesh offsets # Geometry has to be time-dependent num_points = block["NumberOfPoints"] num_points.resize(len(timestamps) + 1, axis=0) num_points[-1] = num_points[-2] # Even if topology cannot be time dep at the moment, number of cells # must be updated num_cells = block["NumberOfCells"] num_cells.resize(len(timestamps) + 1, axis=0) num_cells[-1] = num_cells[-2] steps.attrs.create("NSteps", block["Steps"].attrs["NSteps"] + 1) step_vals = _create_dataset( steps, "Values", shape=(1,), dtype=np.float64, chunks=True, maxshape=(None,), mode=h5_mode, resize=True, ) step_vals[-1] = time idx = step_vals.size - 1 for key in [ "PartOffsets", "NumberOfParts", "PointOffsets", "CellOffsets", "ConnectivityIdOffsets", ]: comp = steps[key] comp.resize(steps.attrs["NSteps"], axis=0) if comp.size > 1: comp[-1] = comp[-2] else: raise ValueError(f"Time step found multiple times in {filename}") # Update offsets for current variable pdo_u.resize(steps.attrs["NSteps"], axis=0) pdo_u[idx] = dataset.shape[0] - array_data.global_shape[0] # Point and cell data of all other variables has to be updated to be synced for key in [ "CellDataOffsets", "PointDataOffsets", ]: comp = steps[key] num_steps = steps.attrs["NSteps"] for dname, dset in comp.items(): # Only resize other data arrays if dname != array_data.name: # Only resize if it hasn't already been resized if dset.size != num_steps: dset.resize(num_steps, axis=0) # Only update data if we are further than first time step if dset.size > 1: dset[-1] = dset[-2] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/xdmf/000077500000000000000000000000001517634040500257065ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/xdmf/__init__.py000066400000000000000000000000551517634040500300170ustar00rootroot00000000000000from . import backend __all__ = ["backend"] scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/backends/xdmf/backend.py000066400000000000000000000434601517634040500276560ustar00rootroot00000000000000""" Module that uses DOLFINx/H5py to import XDMF files. """ from pathlib import Path from typing import Any from xml.etree import ElementTree from mpi4py import MPI import basix import dolfinx import numpy as np import numpy.typing as npt from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from io4dolfinx.utils import check_file_exists, compute_local_range from .. import FileMode, ReadMode from ..h5py.backend import h5pyfile read_mode = ReadMode.parallel def extract_function_names_and_timesteps(filename: Path | str) -> dict[str, list[str]]: tree = ElementTree.parse(filename) root = tree.getroot() mesh_nodes = root.findall(".//Grid[@CollectionType='Temporal']") function_names = [] for mesh in mesh_nodes: function_names.append(mesh.attrib["Name"]) time_stamps: dict[str, list[str]] = {name: [] for name in function_names} for name in function_names: time_steps = root.findall(f".//Grid[@Name='{name}']") for time in time_steps: step = time.find(".//Time") if step is not None: val = step.attrib["Value"] time_stamps[name].append(val) for name in function_names: float_steps = np.argsort(np.array(list(set(time_stamps[name])), dtype=np.float64)) time_stamps[name] = np.array(list(set(time_stamps[name])), dtype=str)[float_steps].tolist() return time_stamps def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default backend arguments given a set of input arguments. Args: arguments: Input backend arguments Returns: Updated backend arguments """ args = arguments or {} return args def read_mesh_data( filename: Path | str, comm: MPI.Intracomm, time: str | float | None, read_from_partition: bool, backend_args: dict[str, Any] | None, ) -> ReadMeshData: """Read mesh data from file. Args: filename: Path to file to read from comm: MPI communicator used in storage time: Time stamp associated with the mesh to read read_from_partition: Whether to read partition information backend_args: Arguments to backend Returns: Internal data structure for the mesh data read from file """ assert not read_from_partition check_file_exists(filename) with dolfinx.io.XDMFFile(comm, filename, "r") as file: cell_shape, cell_degree = file.read_cell_type() cells = file.read_topology_data() x = file.read_geometry_data() return ReadMeshData( cells=cells, cell_type=cell_shape.name, x=x, lvar=int(basix.LagrangeVariant.equispaced), degree=cell_degree, ) def read_point_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: float | str | None, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read data from the nodes of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: Data local to process (contiguous, no mpi comm) and local start range """ # Find function with name u in xml tree check_file_exists(filename) filename = Path(filename) tree = ElementTree.parse(filename) root = tree.getroot() backend_args = get_default_backend_args(backend_args) if time is not None: time_steps = root.findall(f".//Grid[@Name='{name}']") time_found = False for time_node in time_steps: step_node = time_node.find(".//Time") assert isinstance(step_node, ElementTree.Element) if np.isclose(float(step_node.attrib["Value"]), float(time)): time_found = True break func_node = time_node.find(f".//Attribute[@Name='{name}']") if not time_found: raise RuntimeError(f"Function {name} at time={time} not found in {filename}") else: func_node = root.find(f".//Attribute[@Name='{name}']") assert isinstance(func_node, ElementTree.Element) data_node = func_node.find(".//DataItem") assert isinstance(data_node, ElementTree.Element) global_shape = data_node.attrib["Dimensions"].split(" ") func_path = data_node.text assert isinstance(func_path, str) data_file, data_loc = func_path.split(":") data_path = filename.parent / data_file with h5pyfile(data_path, "r", comm=comm) as h5file: data = h5file[data_loc] for s1, s2 in zip(data.shape, global_shape, strict=True): assert int(s1) == int(s2) lr = compute_local_range(comm, data.shape[0]) local_range_start = lr[0] dataset = data[slice(*lr), :] return dataset, local_range_start def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dict[str, Any]: """Read attributes from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attribute group backend_args: Arguments to backend Returns: Dictionary of attributes read from file """ raise NotImplementedError("The XDMF backend cannot read attributes.") def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, Any] | None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """Read timestamps from file. Args: filename: Path to file to read from comm: MPI communicator used in storage function_name: Name of the function to read timestamps for backend_args: Arguments to backend Returns: Numpy array of timestamps read from file """ tree = ElementTree.parse(filename) root = tree.getroot() time_stamps = [] time_steps = root.findall(f".//Grid[@Name='{function_name}']") for time in time_steps: step = time.find(".//Time") if step is not None: val = step.attrib["Value"] time_stamps.append(val) float_steps = np.argsort(np.array(list(set(time_stamps)), dtype=np.float64)) return np.array(list(set(time_stamps)), dtype=str)[float_steps].tolist() def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, Any] | None, ): """Write attributes to file. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attribute group attributes: Dictionary of attributes to write backend_args: Arguments to backend """ raise NotImplementedError("The XDMF backend cannot write attributes.") def write_mesh( filename: Path | str, comm: MPI.Intracomm, mesh: MeshData, backend_args: dict[str, Any] | None, mode: FileMode, time: float, ): """ Write a mesh to file. Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to backend_args: Arguments to backend mode: File-mode to store the mesh time: Time stamp associated with the mesh """ raise NotImplementedError("The XDMF backend cannot write meshes.") def write_meshtags( filename: str | Path, comm: MPI.Intracomm, data: MeshTagsData, backend_args: dict[str, Any] | None, ): """Write mesh tags to file. Args: filename: Path to file to write to comm: MPI communicator used in storage data: Internal data structure for the mesh tags to save to file backend_args: Arguments to backend """ raise NotImplementedError("The XDMF backend cannot write meshtags.") def read_meshtags_data( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> MeshTagsData: """Read mesh tags from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the mesh tags to read backend_args: Arguments to backend Returns: Internal data structure for the mesh tags read from file """ raise NotImplementedError("The XDMF backend cannot read meshtags.") def read_dofmap( filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None, ) -> dolfinx.graph.AdjacencyList: """Read the dofmap of a function with a given name. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofmap for backend_args: Arguments to backend Returns: Dofmap as an AdjacencyList """ raise NotImplementedError("The XDMF backend cannot make checkpoints.") def read_dofs( filename: str | Path, comm: MPI.Intracomm, name: str, time: float, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: """Read the dofs (values) of a function with a given name from a given timestep. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the function to read the dofs for time: Time stamp associated with the function to read backend_args: Arguments to backend Returns: Contiguous sequence of degrees of freedom (with respect to input data) and the global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The XDMF backend cannot make checkpoints.") def read_cell_perms( comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None ) -> npt.NDArray[np.uint32]: """ Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. Args: comm: MPI communicator used in storage filename: Path to file to read from backend_args: Arguments to backend Returns: Contiguous sequence of permutations (with respect to input data) Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The XDMF backend cannot make checkpoints.") def write_function( filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float, mode: FileMode, backend_args: dict[str, Any] | None, ): """ Write a function to file. Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to time: Time stamp associated with function mode: File-mode to store the function backend_args: Arguments to backend """ raise NotImplementedError("The XDMF backend cannot make checkpoints.") def read_legacy_mesh( filename: Path | str, comm: MPI.Intracomm, group: str ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: """Read in the mesh topology, geometry and (optionally) cell type from a legacy DOLFIN HDF5-file. Args: filename: Path to file to read from comm: MPI communicator used in storage group: Group in HDF5 file where mesh is stored Returns: Tuple containing: - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates - Cell type as a string (e.g. "tetrahedron") or None if not found """ raise NotImplementedError("The XDMF backend cannot read legacy DOLFIN meshes.") def snapshot_checkpoint( filename: Path | str, mode: FileMode, u: dolfinx.fem.Function, backend_args: dict[str, Any] | None, ): """Create a snapshot checkpoint of a dolfinx function. Args: filename: Path to file to read from mode: File-mode to store the function u: dolfinx function to create a snapshot checkpoint for backend_args: Arguments to backend """ raise NotImplementedError("The XDMF backend cannot make checkpoints.") def read_hdf5_array( comm: MPI.Intracomm, filename: Path | str, group: str, backend_args: dict[str, Any] | None, ) -> tuple[np.ndarray, int]: """Read an array from an HDF5 file. Args: comm: MPI communicator used in storage filename: Path to file to read from group: Group in HDF5 file where array is stored backend_args: Arguments to backend Returns: Tuple containing: - Numpy array read from file - Global starting point on the process. Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ raise NotImplementedError("The XDMF backend cannot read HDF5 arrays") def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ # Find all functions in xml tree check_file_exists(filename) filename = Path(filename) tree = ElementTree.parse(filename) root = tree.getroot() backend_args = get_default_backend_args(backend_args) # Functions in checkpoint format checkpoint_funcs = root.findall(".//Attribute[@ItemType='FiniteElementFunction']") names = [func.attrib["Name"] for func in checkpoint_funcs] # Temporal funcs temporal_funcs = root.findall(".//Grid[@GridType='Collection']") for func in temporal_funcs: names.append(func.attrib["Name"]) return list(set(names)) def read_cell_data( filename: Path | str, name: str, comm: MPI.Intracomm, time: str | float | None, backend_args: dict[str, Any] | None, ) -> tuple[npt.NDArray[np.int64], np.ndarray]: """Read data from the cells of a mesh. Args: filename: Path to file name: Name of point data comm: Communicator to launch IO on. time: The time stamp backend_args: The backend arguments Returns: A tuple (topology, dofs) where topology contains the vertex indices of the cells, dofs the degrees of freedom within that cell. """ # Find function with name u in xml tree check_file_exists(filename) filename = Path(filename) tree = ElementTree.parse(filename) root = tree.getroot() backend_args = get_default_backend_args(backend_args) if time is not None: time_steps = root.findall(f".//Grid[@Name='{name}']") time_found = False for time_node in time_steps: step_node = time_node.find(".//Time") assert isinstance(step_node, ElementTree.Element) if np.isclose(float(step_node.attrib["Value"]), time): time_found = True break func_node = time_node.find(f".//Attribute[@Name='{name}']") if not time_found: raise RuntimeError(f"Function {name} at time={time} not found in {filename}") else: func_node = root.find(f".//Attribute[@Name='{name}']") assert func_node is not None if func_node.attrib["ItemType"] == "FiniteElementFunction": if (family := func_node.attrib["ElementFamily"]) != "DG" or ( degree := int(func_node.attrib["ElementDegree"]) ) != 0: raise ValueError( f"Cannot read in finite element function ({family}, {degree}) as cell data." ) # Get vector sub-element vec_el = None for node in func_node.iter(): comp = node.text assert comp is not None if "vector" in comp: vec_el = node break assert vec_el is not None dof_dimensions = np.array(vec_el.attrib["Dimensions"].split(" "), dtype=np.int32) vtxt = vec_el.text assert vtxt is not None vector_file, vector_h5path = vtxt.split(":") grid_node = root.find(f".//Attribute[@Name='{name}']/..") assert grid_node is not None topology = grid_node.find("./Topology") assert topology is not None ttext = topology[0].text assert ttext is not None topology_file, topology_h5path = ttext.split(":") with h5pyfile(filename.parent / vector_file, "r", comm=comm) as h5_mesh: data_loc = h5_mesh[vector_h5path] data_shape = data_loc.shape[0] assert int(np.prod(data_shape)) == int(np.prod(dof_dimensions)) local_range = compute_local_range(comm, data_shape) vec_dofs = data_loc[slice(*local_range)] with h5pyfile(filename.parent / topology_file, "r", comm=comm) as h5_top: data_loc = h5_top[topology_h5path] top_data_shape = data_loc.shape[0] assert dof_dimensions[0] == top_data_shape local_range = compute_local_range(comm, data_shape) top_dofs = data_loc[slice(*local_range)].astype(np.int64) return top_dofs, vec_dofs else: raise NotImplementedError("Not implemented yet.") def write_data( filename: Path | str, point_data: ArrayData, comm: MPI.Intracomm, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, ): """Write a 2D-array to file (distributed across proceses with MPI). Args: filename: Path to file point_data: Data to write to file time: Time stamp mode: Append or write backend_args: The backend arguments """ raise NotImplementedError("XDMF has not implemented this yet") scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/checkpointing.py000066400000000000000000000557531517634040500264140ustar00rootroot00000000000000# Copyright (C) 2023-2026 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from __future__ import annotations import typing from pathlib import Path from typing import Any from mpi4py import MPI import basix import dolfinx import numpy as np import numpy.typing as npt import ufl from packaging.version import Version from . import compat from .backends import FileMode, ReadMode, get_backend from .comm_helpers import ( send_and_recv_cell_perm, send_dofmap_and_recv_values, send_dofs_and_recv_values, ) from .readers import create_geometry_function_space from .structures import ArrayData, FunctionData, MeshTagsData from .utils import ( check_file_exists, compute_dofmap_pos, compute_local_range, index_owner, unroll_dofmap, unroll_insert_position, ) from .writers import prepare_meshdata_for_storage from .writers import write_function as _internal_function_writer from .writers import write_mesh as _internal_mesh_writer __all__ = [ "read_mesh", "write_function", "read_function", "write_mesh", "read_meshtags", "write_meshtags", "read_attributes", "write_attributes", ] def write_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, attributes: dict[str, np.ndarray], backend_args: dict[str, typing.Any] | None = None, backend: str = "adios2", ): """Write attributes to file. Args: filename: Path to file to write to comm: MPI communicator used in storage name: Name of the attributes attributes: Dictionary of attributes to write to file backend_args: Arguments for backend, for instance file type. backend: What backend to use for writing. """ backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) backend_cls.write_attributes(filename, comm, name, attributes, backend_args) def read_attributes( filename: Path | str, comm: MPI.Intracomm, name: str, backend_args: dict[str, typing.Any] | None = None, backend: str = "adios2", ) -> dict[str, typing.Any]: """Read attributes from file. Args: filename: Path to file to read from comm: MPI communicator used in storage name: Name of the attributes backend_args: Arguments for backend, for instance file type. backend: What backend to use for writing. Returns: The attributes """ backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) return backend_cls.read_attributes(filename, comm, name, backend_args) def read_timestamps( filename: Path | str, comm: MPI.Intracomm, function_name: str, backend_args: dict[str, typing.Any] | None = None, backend: str = "adios2", ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """ Read time-stamps from a checkpoint file. Args: comm: MPI communicator filename: Path to file function_name: Name of the function to read time-stamps for backend_args: Arguments for backend, for instance file type. backend: What backend to use for writing. Returns: The time-stamps """ check_file_exists(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) return backend_cls.read_timestamps(filename, comm, function_name, backend_args) def write_meshtags( filename: Path | str, mesh: dolfinx.mesh.Mesh, meshtags: dolfinx.mesh.MeshTags, meshtag_name: typing.Optional[str] = None, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Write meshtags associated with input mesh to file. .. note:: For this checkpoint to work, the mesh must be written to file using :func:`write_mesh` before calling this function. Args: filename: Path to save meshtags (with file-extension) mesh: The mesh associated with the meshtags meshtags: The meshtags to write to file meshtag_name: Name of the meshtag. If None, the meshtag name is used. backend_args: Option to IO backend. backend: IO backend """ # Extract data from meshtags (convert to global geometry node indices for each entity) tag_entities = meshtags.indices dim = meshtags.dim num_tag_entities_local = mesh.topology.index_map(dim).size_local local_tag_entities = tag_entities[tag_entities < num_tag_entities_local] local_values = meshtags.values[: len(local_tag_entities)] num_saved_tag_entities = len(local_tag_entities) local_start = mesh.comm.exscan(num_saved_tag_entities, op=MPI.SUM) local_start = local_start if mesh.comm.rank != 0 else 0 global_num_tag_entities = mesh.comm.allreduce(num_saved_tag_entities, op=MPI.SUM) dof_layout = compat.cmap(mesh).create_dof_layout() if hasattr(dof_layout, "num_entity_closure_dofs"): num_dofs_per_entity = dof_layout.num_entity_closure_dofs(dim) else: num_dofs_per_entity = len(dof_layout.entity_closure_dofs(dim, 0)) entities_to_geometry = dolfinx.cpp.mesh.entities_to_geometry( mesh._cpp_object, dim, local_tag_entities, False ) indices = ( mesh.geometry.index_map() .local_to_global(entities_to_geometry.reshape(-1)) .reshape(entities_to_geometry.shape) ) name = meshtag_name or meshtags.name tag_ct = dolfinx.cpp.mesh.cell_entity_type(mesh.topology.cell_type, dim, 0).name tag_data = MeshTagsData( values=local_values, num_entities_global=global_num_tag_entities, num_dofs_per_entity=num_dofs_per_entity, indices=indices, name=name, local_start=local_start, dim=meshtags.dim, cell_type=tag_ct, ) # Get backend and default arguments backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) return backend_cls.write_meshtags(filename, mesh.comm, tag_data, backend_args=backend_args) def read_meshtags( filename: Path | str, mesh: dolfinx.mesh.Mesh, meshtag_name: str, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ) -> dolfinx.mesh.MeshTags: """ Read meshtags from file and return a :class:`dolfinx.mesh.MeshTags` object. Args: filename: Path to meshtags file (with file-extension) mesh: The mesh associated with the meshtags meshtag_name: The name of the meshtag to read engine: Adios2 Engine Returns: The meshtags """ check_file_exists(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) data = backend_cls.read_meshtags_data(filename, mesh.comm, meshtag_name, backend_args) local_entities, local_values = dolfinx.io.distribute_entity_data( mesh, int(data.dim), data.indices, data.values ) mesh.topology.create_connectivity(data.dim, 0) mesh.topology.create_connectivity(data.dim, mesh.topology.dim) adj = dolfinx.graph.adjacencylist(local_entities) local_values = np.array(local_values, dtype=np.int32) mt = dolfinx.mesh.meshtags_from_entities(mesh, int(data.dim), adj, local_values) mt.name = meshtag_name return mt def read_function( filename: Path | str, u: dolfinx.fem.Function, time: float = 0.0, name: str | None = None, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Read checkpoint from file and fill it into `u`. Args: filename: Path to checkpoint u: Function to fill time: Time-stamp associated with checkpoint name: If not provided, `u.name` is used to search through the input file for the function """ check_file_exists(filename) mesh = u.function_space.mesh comm = mesh.comm if name is None: name = u.name # ----------------------Step 1--------------------------------- # Compute index of input cells and get cell permutation num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local input_cells = mesh.topology.original_cell_index[:num_owned_cells] mesh.topology.create_entity_permutations() cell_perm = mesh.topology.get_cell_permutation_info()[:num_owned_cells] # Compute mesh->input communicator # 1.1 Compute mesh->input communicator backend_cls = get_backend(backend) owners: npt.NDArray[np.int32] if backend_cls.read_mode == ReadMode.serial: owners = np.zeros(input_cells, dtype=np.int32) elif backend_cls.read_mode == ReadMode.parallel: num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global owners = index_owner(mesh.comm, input_cells, num_cells_global) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") # -------------------Step 2------------------------------------ # Send and receive global cell index and cell perm inc_cells, inc_perms = send_and_recv_cell_perm(input_cells, cell_perm, owners, mesh.comm) # -------------------Step 3----------------------------------- # Read dofmap from file and compute dof owners check_file_exists(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) input_dofmap = backend_cls.read_dofmap(filename, comm, name, backend_args) # Compute owner of dofs in dofmap dof_owner: npt.NDArray[np.int32] if backend_cls.read_mode == ReadMode.serial: dof_owner = np.zeros(len(input_dofmap.array), dtype=np.int32) elif backend_cls.read_mode == ReadMode.parallel: num_dofs_global = ( u.function_space.dofmap.index_map.size_global * u.function_space.dofmap.index_map_bs ) dof_owner = index_owner(comm, input_dofmap.array.astype(np.int64), num_dofs_global) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") # --------------------Step 4----------------------------------- # Read array from file and communicate them to input dofmap process input_array, starting_pos = backend_cls.read_dofs(filename, comm, name, time, backend_args) recv_array = send_dofs_and_recv_values( input_dofmap.array.astype(np.int64), dof_owner, comm, input_array, starting_pos ) # -------------------Step 5-------------------------------------- # Invert permutation of input data based on input perm # Then apply current permutation to the local data element = u.function_space.element if element.needs_dof_transformations: bs = u.function_space.dofmap.bs # Read input cell permutations on dofmap process local_input_range = compute_local_range(comm, num_cells_global) input_local_cell_index = inc_cells - local_input_range[0] input_perms = backend_cls.read_cell_perms(comm, filename, backend_args) # Start by sorting data array by cell permutation num_dofs_per_cell = input_dofmap.offsets[1:] - input_dofmap.offsets[:-1] assert np.allclose(num_dofs_per_cell, num_dofs_per_cell[0]) # Sort dofmap by input local cell index input_perms_sorted = input_perms[input_local_cell_index] unrolled_dofmap_position = unroll_insert_position( input_local_cell_index, num_dofs_per_cell[0] ) dofmap_sorted_by_input = recv_array[unrolled_dofmap_position] # First invert input data to reference element then transform to current mesh element.Tt_apply(dofmap_sorted_by_input, input_perms_sorted, bs) element.Tt_inv_apply(dofmap_sorted_by_input, inc_perms, bs) # Compute invert permutation inverted_perm = np.empty_like(unrolled_dofmap_position) inverted_perm[unrolled_dofmap_position] = np.arange( len(unrolled_dofmap_position), dtype=inverted_perm.dtype ) recv_array = dofmap_sorted_by_input[inverted_perm] # ------------------Step 6---------------------------------------- # For each dof owned by a process, find the local position in the dofmap. V = u.function_space local_cells, dof_pos = compute_dofmap_pos(V) input_cells = V.mesh.topology.original_cell_index[local_cells] num_cells_global = V.mesh.topology.index_map(V.mesh.topology.dim).size_global if backend_cls.read_mode == ReadMode.serial: owners = np.zeros(len(input_cells), dtype=np.int32) elif backend_cls.read_mode == ReadMode.parallel: owners = index_owner(V.mesh.comm, input_cells, num_cells_global) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") unique_owners, owner_count = np.unique(owners, return_counts=True) # FIXME: In C++ use NBX to find neighbourhood sub_comm = V.mesh.comm.Create_dist_graph( [V.mesh.comm.rank], [len(unique_owners)], unique_owners, reorder=False ) source, dest, _ = sub_comm.Get_dist_neighbors() sub_comm.Free() owned_values = send_dofmap_and_recv_values( comm, np.asarray(source, dtype=np.int32), np.asarray(dest, dtype=np.int32), owners, owner_count.astype(np.int32), input_cells, dof_pos, num_cells_global, recv_array, input_dofmap.offsets, ) u.x.array[: len(owned_values)] = owned_values u.x.scatter_forward() def read_mesh( filename: Path | str, comm: MPI.Intracomm, ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet, time: float | str | None = 0.0, read_from_partition: bool = False, backend_args: dict[str, Any] | None = None, backend: str = "adios2", max_facet_to_cell_links: int = 2, ) -> dolfinx.mesh.Mesh: """ Read an ADIOS2 mesh into DOLFINx. Args: filename: Path to input file comm: The MPI communciator to distribute the mesh over ghost_mode: Ghost mode to use for mesh. If `read_from_partition` is set to `True` this option is ignored. time: Time stamp associated with mesh read_from_partition: Read mesh with partition from file backend_args: List of arguments to reader backend max_facet_to_cell_links: Maximum number of cells a facet can be connected to. Returns: The distributed mesh """ # Read in data in a distributed fashin check_file_exists(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) # Let each backend handle what should be default behavior when reading mesh # with or without time stamp. dist_in_data = backend_cls.read_mesh_data( filename, comm, time=time, read_from_partition=read_from_partition, backend_args=backend_args, ) # Create DOLFINx mesh element = basix.ufl.element( basix.ElementFamily.P, dist_in_data.cell_type, dist_in_data.degree, basix.LagrangeVariant(int(dist_in_data.lvar)), shape=(dist_in_data.x.shape[1],), dtype=dist_in_data.x.dtype, ) domain = ufl.Mesh(element) if (partition_graph := dist_in_data.partition_graph) is not None: def partitioner(comm: MPI.Intracomm, n, m, topo): assert len(topo[0]) % (len(partition_graph.offsets) - 1) == 0 if Version(dolfinx.__version__) > Version("0.9.0"): return partition_graph._cpp_object else: return partition_graph else: try: partitioner = dolfinx.cpp.mesh.create_cell_partitioner( ghost_mode, max_facet_to_cell_links=max_facet_to_cell_links ) except TypeError: partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode) # Should change to the commented code below when we require python # minimum version to be >=3.12 see https://github.com/python/cpython/pull/116198 # import inspect # sig = inspect.signature(dolfinx.mesh.create_cell_partitioner) # part_kwargs = {} # if "max_facet_to_cell_links" in list(sig.parameters.keys()): # part_kwargs["max_facet_to_cell_links"] = max_facet_to_cell_links # partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode, **part_kwargs) return dolfinx.mesh.create_mesh( comm, cells=dist_in_data.cells, x=dist_in_data.x, e=domain, partitioner=partitioner, ) def write_mesh( filename: Path, mesh: dolfinx.mesh.Mesh, mode: FileMode = FileMode.write, time: float = 0.0, store_partition_info: bool = False, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Write a mesh to file. Args: filename: Path to save mesh (without file-extension) mesh: The mesh to write to file store_partition_info: Store mesh partitioning (including ghosting) to file """ mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info) _internal_mesh_writer( filename, mesh.comm, mesh_data=mesh_data, time=time, backend_args=backend_args, backend=backend, mode=mode, ) def write_function( filename: Path | str, u: dolfinx.fem.Function, time: float = 0.0, mode: FileMode = FileMode.append, name: str | None = None, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Write function checkpoint to file. Args: u: Function to write to file time: Time-stamp for simulation filename: Path to write to mode: Write or append. name: Name of function to write. If None, the name of the function is used. backend_args: Arguments to the IO backend. backend: The backend to use """ dofmap = u.function_space.dofmap values = u.x.array mesh = u.function_space.mesh comm = mesh.comm mesh.topology.create_entity_permutations() cell_perm = mesh.topology.get_cell_permutation_info() num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local local_cell_range = mesh.topology.index_map(mesh.topology.dim).local_range num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global # Convert local dofmap into global_dofmap dmap = dofmap.list num_dofs_per_cell = dmap.shape[1] dofmap_bs = dofmap.bs num_dofs_local_dmap = num_cells_local * num_dofs_per_cell * dofmap_bs index_map_bs = dofmap.index_map_bs # Unroll dofmap for block size unrolled_dofmap = unroll_dofmap(dofmap.list[:num_cells_local, :], dofmap_bs) dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1) dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1) # Convert imap index to global index imap_global = dofmap.index_map.local_to_global(dmap_loc) dofmap_global = imap_global * index_map_bs + dmap_rem dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap) # Compute dofmap offsets local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64) local_dofmap_offsets[:] *= num_dofs_per_cell * dofmap_bs local_dofmap_offsets += dofmap_imap.local_range[0] num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs local_dof_range = np.asarray(dofmap.index_map.local_range) * dofmap.index_map_bs num_dofs_local = local_dof_range[1] - local_dof_range[0] # Create internal data structure for function data to write to file function_data = FunctionData( cell_permutations=cell_perm[:num_cells_local].copy(), local_cell_range=local_cell_range, num_cells_global=num_cells_global, dofmap_array=dofmap_global, dofmap_offsets=local_dofmap_offsets, dofmap_range=dofmap_imap.local_range, global_dofs_in_dofmap=dofmap_imap.size_global, values=values[:num_dofs_local].copy(), dof_range=local_dof_range, num_dofs_global=num_dofs_global, name=name or u.name, ) # Write to file fname = Path(filename) _internal_function_writer( fname, comm, function_data, time, backend_args=backend_args, backend=backend, mode=mode ) def read_function_names( filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None = None, backend: str = "h5py", ) -> list[str]: """Read all function names from a file. Args: filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend Returns: A list of function names. """ backend_cls = get_backend(backend) return backend_cls.read_function_names(filename, comm, backend_args=backend_args) def write_point_data( filename: Path | str, u: dolfinx.fem.Function, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, backend: str = "vtkhdf", ): """Write function to file by interpolating into geometry nodes. Args: filename: Path to file u: The function to store time: Time stamp mode: Append or write backend_args: The backend arguments backend: Which backend to use. """ V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape))) v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype) v_out.interpolate(u) comm = v_out.function_space.mesh.comm data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs) local_range = V.dofmap.index_map.local_range num_dofs_local = V.dofmap.index_map.size_local data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local] ad = ArrayData( name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Point" ) backend_cls = get_backend(backend) return backend_cls.write_data( filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args ) def write_cell_data( filename: Path | str, u: dolfinx.fem.Function, time: str | float | None, mode: FileMode, backend_args: dict[str, Any] | None, backend: str = "vtkhdf", ): """Write function to file by interpolating into cell midpoints. Args: filename: Path to file point_data: Data to write to file time: Time stamp mode: Append or write backend_args: The backend arguments """ V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape)) v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype) v_out.interpolate(u) comm = v_out.function_space.mesh.comm data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs) local_range = V.dofmap.index_map.local_range num_dofs_local = V.dofmap.index_map.size_local data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local] backend_cls = get_backend(backend) ad = ArrayData( name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Cell" ) backend_cls = get_backend(backend) return backend_cls.write_data( filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args ) scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/comm_helpers.py000066400000000000000000000237161517634040500262360ustar00rootroot00000000000000from __future__ import annotations from mpi4py import MPI import numpy as np import numpy.typing as npt from .utils import compute_insert_position, compute_local_range, valid_function_types __all__ = [ "send_dofmap_and_recv_values", "send_and_recv_cell_perm", "send_dofs_and_recv_values", "numpy_to_mpi", ] """ Helpers for sending and receiving values for checkpointing """ numpy_to_mpi = { np.float64: MPI.DOUBLE, np.float32: MPI.FLOAT, np.complex64: MPI.COMPLEX, np.complex128: MPI.DOUBLE_COMPLEX, np.int64: MPI.INT64_T, np.int32: MPI.INT32_T, } def send_dofmap_and_recv_values( comm: MPI.Intracomm, source_ranks: npt.NDArray[np.int32], dest_ranks: npt.NDArray[np.int32], output_owners: npt.NDArray[np.int32], dest_size: npt.NDArray[np.int32], input_cells: npt.NDArray[np.int64], dofmap_pos: npt.NDArray[np.int32], num_cells_global: np.int64, values: npt.NDArray[valid_function_types], dofmap_offsets: npt.NDArray[np.int32], ) -> npt.NDArray[valid_function_types]: """ Given a set of positions in input dofmap, give the global input index of this dofmap entry in input file. Args: comm: The MPI communicator to create the Neighbourhood-communicator from source_ranks: Ranks that will send dofmap indices to current process dest_ranks: Ranks that will receive dofmap indices from current process output_owners: The owners of each dofmap entry on this process. The unique set of these entries should be the same as the dest_ranks. dest_size: The number of entries sent to each owner input_cells: A cell associated with the degree of freedom sent (global index). dofmap_pos: The local position in the dofmap. I.e. `dof = dofmap.links(input_cells)[dofmap_pos]` num_cells_global: Number of global cells values: Values currently held by this process. These are ordered (num_cells_local, num_dofs_per_cell), flattened row-major. dofmap_offsets: Local dofmap offsets to access the correct `values`. Returns: Values corresponding to the dofs owned by this process. """ insert_position = compute_insert_position(output_owners, dest_ranks, dest_size) # Pack the cells and dofmap position for all dofs this process is distributing out_cells = np.zeros(len(output_owners), dtype=np.int64) out_cells[insert_position] = input_cells out_pos = np.zeros(len(output_owners), dtype=np.int32) out_pos[insert_position] = dofmap_pos # Compute map from the data index sent to each process and the local # number on the current process proc_to_dof = np.zeros_like(input_cells, dtype=np.int32) proc_to_dof[insert_position] = np.arange(len(input_cells), dtype=np.int32) del insert_position # Send sizes to create data structures for receiving from NeighAlltoAllv recv_size = np.zeros(len(source_ranks), dtype=np.int32) mesh_to_data_comm = comm.Create_dist_graph_adjacent( source_ranks.tolist(), dest_ranks.tolist(), reorder=False ) mesh_to_data_comm.Neighbor_alltoall(dest_size, recv_size) # Prepare data-structures for receiving total_incoming = sum(recv_size) inc_cells = np.zeros(total_incoming, dtype=np.int64) inc_pos = np.zeros(total_incoming, dtype=np.intc) # Compute incoming offset inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc) inc_offsets[1:] = np.cumsum(recv_size) # Send data s_msg = [out_cells, dest_size, MPI.INT64_T] r_msg = [inc_cells, recv_size, MPI.INT64_T] mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg) s_msg = [out_pos, dest_size, MPI.INT32_T] r_msg = [inc_pos, recv_size, MPI.INT32_T] mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg) mesh_to_data_comm.Free() local_input_range = compute_local_range(comm, num_cells_global) values_to_distribute = np.zeros_like(inc_pos, dtype=values.dtype) # Map values based on input cells and dofmap local_cells = inc_cells - local_input_range[0] values_to_distribute = values[dofmap_offsets[local_cells] + inc_pos] # Send input dofs back to owning process data_to_mesh_comm = comm.Create_dist_graph_adjacent( dest_ranks.tolist(), source_ranks.tolist(), reorder=False ) incoming_global_dofs = np.zeros(sum(dest_size), dtype=values.dtype) s_msg = [values_to_distribute, recv_size, numpy_to_mpi[values.dtype.type]] r_msg = [incoming_global_dofs, dest_size, numpy_to_mpi[values.dtype.type]] data_to_mesh_comm.Neighbor_alltoallv(s_msg, r_msg) # Sort incoming global dofs as they were inputted assert len(incoming_global_dofs) == len(input_cells) sorted_global_dofs = np.zeros_like(incoming_global_dofs, dtype=values.dtype) sorted_global_dofs[proc_to_dof] = incoming_global_dofs data_to_mesh_comm.Free() return sorted_global_dofs def send_and_recv_cell_perm( cells: npt.NDArray[np.int64], perms: npt.NDArray[np.uint32], cell_owners: npt.NDArray[np.int32], comm: MPI.Intracomm, ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.uint32]]: """ Send global cell index and permutation to corresponding entry in `dest_ranks`. Args: cells: The global input index of the cell perms: The corresponding cell permutation of the cell cell_owners: The rank to send the i-th entry of cells and perms to comm: Rank of comm to generate neighbourhood communicator from """ dest_ranks, _dest_size = np.unique(cell_owners, return_counts=True) dest_size = _dest_size.astype(np.int32) del _dest_size mesh_to_data = comm.Create_dist_graph( [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False ) source, dest, _ = mesh_to_data.Get_dist_neighbors() assert np.allclose(dest, dest_ranks) insert_position = compute_insert_position(cell_owners, dest_ranks.astype(np.int32), dest_size) # Pack cells and permutations for sending out_cells = np.zeros_like(cells, dtype=np.int64) out_perm = np.zeros_like(perms, dtype=np.uint32) out_cells[insert_position] = cells out_perm[insert_position] = perms del insert_position # Send sizes to create data structures for receiving from NeighAlltoAllv recv_size = np.zeros_like(source, dtype=np.int32) mesh_to_data.Neighbor_alltoall(dest_size, recv_size) # Prepare data-structures for receiving total_incoming = sum(recv_size) inc_cells = np.zeros(total_incoming, dtype=np.int64) inc_perm = np.zeros(total_incoming, dtype=np.uint32) # Compute incoming offset inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc) inc_offsets[1:] = np.cumsum(recv_size) # Send data s_msg = [out_cells, dest_size, MPI.INT64_T] r_msg = [inc_cells, recv_size, MPI.INT64_T] mesh_to_data.Neighbor_alltoallv(s_msg, r_msg) s_msg = [out_perm, dest_size, MPI.UINT32_T] r_msg = [inc_perm, recv_size, MPI.UINT32_T] mesh_to_data.Neighbor_alltoallv(s_msg, r_msg) mesh_to_data.Free() return inc_cells, inc_perm def send_dofs_and_recv_values( input_dofmap: npt.NDArray[np.int64], dofmap_owners: npt.NDArray[np.int32], comm: MPI.Intracomm, input_array: npt.NDArray[valid_function_types], array_start: int, ): """ Send a set of dofs (global index) to the process holding the DOF values to retrieve them. Args: input_dofmap: List of dofs (global index) that this process wants values for dofmap_owners: The process currently holding the values this process want to get. comm: MPI communicator input_array: Values for dofs array_start: The global starting index of `input_array`. """ dest_ranks, _dest_size = np.unique(dofmap_owners, return_counts=True) dest_size = _dest_size.astype(np.int32) del _dest_size dofmap_to_values = comm.Create_dist_graph( [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False ) source, dest, _ = dofmap_to_values.Get_dist_neighbors() assert np.allclose(dest_ranks, dest) # Compute amount of data to send to each process insert_position = compute_insert_position(dofmap_owners, dest_ranks, dest_size) # Pack dofs for sending out_dofs = np.zeros(len(dofmap_owners), dtype=np.int64) out_dofs[insert_position] = input_dofmap # Compute map from the data index sent to each process and the local number on # the current process proc_to_local = np.zeros_like(input_dofmap, dtype=np.int32) proc_to_local[insert_position] = np.arange(len(input_dofmap), dtype=np.int32) del insert_position # Send sizes to create data structures for receiving from NeighAlltoAllv recv_size = np.zeros_like(source, dtype=np.int32) recv_size.resize(max(len(recv_size), 1)) # Minimal resize to work with ompi dest_size.resize(max(len(dest_size), 1)) # Mininal resize to work with ompi dofmap_to_values.Neighbor_alltoall(dest_size, recv_size) dest_size.resize(len(dest)) recv_size.resize(len(source)) # Send input dofs to processes holding input array inc_dofs = np.zeros(sum(recv_size), dtype=np.int64) s_msg = [out_dofs, dest_size, MPI.INT64_T] r_msg = [inc_dofs, recv_size, MPI.INT64_T] dofmap_to_values.Neighbor_alltoallv(s_msg, r_msg) dofmap_to_values.Free() # Send back appropriate input values if len(input_array) > 0: sending_values = input_array[inc_dofs - array_start] else: sending_values = np.zeros(0, dtype=input_array.dtype) values_to_dofmap = comm.Create_dist_graph_adjacent(dest, source, reorder=False) inc_values = np.zeros_like(out_dofs, dtype=input_array.dtype) s_msg_rev = [sending_values, recv_size, numpy_to_mpi[input_array.dtype.type]] r_msg_rev = [inc_values, dest_size, numpy_to_mpi[input_array.dtype.type]] values_to_dofmap.Neighbor_alltoallv(s_msg_rev, r_msg_rev) values_to_dofmap.Free() # Sort inputs according to local dof number (input process) values = np.empty_like(inc_values, dtype=input_array.dtype) values[proc_to_local] = inc_values return values scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/compat.py000066400000000000000000000003751517634040500250400ustar00rootroot00000000000000import dolfinx.mesh def cmap(mesh) -> dolfinx.fem.element.CoordinateElement: # Due to https://github.com/FEniCS/dolfinx/pull/4169 if callable(mesh.geometry.cmap): return mesh.geometry.cmap() else: return mesh.geometry.cmap scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/original_checkpoint.py000066400000000000000000000424231517634040500275700ustar00rootroot00000000000000# Copyright (C) 2024 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from __future__ import annotations import typing from pathlib import Path from mpi4py import MPI import dolfinx import numpy as np from . import compat from .backends import FileMode, get_backend from .comm_helpers import numpy_to_mpi from .structures import FunctionData, MeshData from .utils import ( compute_insert_position, compute_local_range, index_owner, unroll_dofmap, unroll_insert_position, ) __all__ = ["write_function_on_input_mesh", "write_mesh_input_order"] def create_original_mesh_data(mesh: dolfinx.mesh.Mesh) -> MeshData: """ Store data locally on output process """ # 1. Send cell indices owned by current process to the process which owned its input # Get the input cell index for cells owned by this process num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local original_cell_index = mesh.topology.original_cell_index[:num_owned_cells] # Compute owner of cells on this process based on the original cell index num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global output_cell_owner = index_owner(mesh.comm, original_cell_index, num_cells_global) local_cell_range = compute_local_range(mesh.comm, num_cells_global) # Compute outgoing edges from current process to outputting process # Computes the number of cells sent to each process at the same time cell_destinations, _send_cells_per_proc = np.unique(output_cell_owner, return_counts=True) send_cells_per_proc = _send_cells_per_proc.astype(np.int32) del _send_cells_per_proc cell_to_output_comm = mesh.comm.Create_dist_graph( [mesh.comm.rank], [len(cell_destinations)], cell_destinations.tolist(), reorder=False, ) cell_sources, cell_dests, _ = cell_to_output_comm.Get_dist_neighbors() assert np.allclose(cell_dests, cell_destinations) # Compute number of recieving cells recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32) if len(send_cells_per_proc) == 0: send_cells_per_proc = np.zeros(1, dtype=np.int32) if len(recv_cells_per_proc) == 0: recv_cells_per_proc = np.zeros(1, dtype=np.int32) send_cells_per_proc = send_cells_per_proc.astype(np.int32) cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc) assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0] # Pack and send cell indices (used for mapping topology dofmap later) cell_insert_position = compute_insert_position( output_cell_owner, cell_destinations, send_cells_per_proc ) send_cells = np.empty_like(cell_insert_position, dtype=np.int64) send_cells[cell_insert_position] = original_cell_index recv_cells = np.empty(recv_cells_per_proc.sum(), dtype=np.int64) send_cells_msg = [send_cells, send_cells_per_proc, MPI.INT64_T] recv_cells_msg = [recv_cells, recv_cells_per_proc, MPI.INT64_T] cell_to_output_comm.Neighbor_alltoallv(send_cells_msg, recv_cells_msg) del send_cells_msg, recv_cells_msg, send_cells # Map received cells to the local index local_cell_index = recv_cells - local_cell_range[0] # 2. Create dofmap based on original geometry indices and re-order in the same order as original # cell indices on output process # Get original node index for all nodes (including ghosts) and convert dofmap to these indices original_node_index = mesh.geometry.input_global_indices _, num_nodes_per_cell = mesh.geometry.dofmap.shape local_geometry_dofmap = mesh.geometry.dofmap[:num_owned_cells, :] global_geometry_dofmap = original_node_index[local_geometry_dofmap.reshape(-1)] # Unroll insert position for geometry dofmap dofmap_insert_position = unroll_insert_position(cell_insert_position, num_nodes_per_cell) # Create and commmnicate connecitivity in original geometry indices send_geometry_dofmap = np.empty_like(dofmap_insert_position, dtype=np.int64) send_geometry_dofmap[dofmap_insert_position] = global_geometry_dofmap del global_geometry_dofmap send_sizes_dofmap = send_cells_per_proc * num_nodes_per_cell recv_sizes_dofmap = recv_cells_per_proc * num_nodes_per_cell recv_geometry_dofmap = np.empty(recv_sizes_dofmap.sum(), dtype=np.int64) send_geometry_dofmap_msg = [send_geometry_dofmap, send_sizes_dofmap, MPI.INT64_T] recv_geometry_dofmap_msg = [recv_geometry_dofmap, recv_sizes_dofmap, MPI.INT64_T] cell_to_output_comm.Neighbor_alltoallv(send_geometry_dofmap_msg, recv_geometry_dofmap_msg) del send_geometry_dofmap_msg, recv_geometry_dofmap_msg # Reshape dofmap and sort by original cell index recv_dofmap = recv_geometry_dofmap.reshape(-1, num_nodes_per_cell) sorted_recv_dofmap = np.empty_like(recv_dofmap) sorted_recv_dofmap[local_cell_index] = recv_dofmap # 3. Move geometry coordinates to input process # Compute outgoing edges from current process and create neighbourhood communicator # Also create number of outgoing cells at the same time num_owned_nodes = mesh.geometry.index_map().size_local num_nodes_global = mesh.geometry.index_map().size_global output_node_owner = index_owner( mesh.comm, original_node_index[:num_owned_nodes], num_nodes_global ) node_destinations, _send_nodes_per_proc = np.unique(output_node_owner, return_counts=True) send_nodes_per_proc = _send_nodes_per_proc.astype(np.int32) del _send_nodes_per_proc geometry_to_owner_comm = mesh.comm.Create_dist_graph( [mesh.comm.rank], [len(node_destinations)], node_destinations.tolist(), reorder=False, ) node_sources, node_dests, _ = geometry_to_owner_comm.Get_dist_neighbors() assert np.allclose(node_dests, node_destinations) # Compute send node insert positions send_nodes_position = compute_insert_position( output_node_owner, node_destinations, send_nodes_per_proc ) unrolled_nodes_positiion = unroll_insert_position(send_nodes_position, 3) send_coordinates = np.empty_like(unrolled_nodes_positiion, dtype=mesh.geometry.x.dtype) send_coordinates[unrolled_nodes_positiion] = mesh.geometry.x[:num_owned_nodes, :].reshape(-1) # Send and recieve geometry sizes send_coordinate_sizes = (send_nodes_per_proc * 3).astype(np.int32) recv_coordinate_sizes = np.zeros_like(node_sources, dtype=np.int32) geometry_to_owner_comm.Neighbor_alltoall(send_coordinate_sizes, recv_coordinate_sizes) # Send node coordinates recv_coordinates = np.empty(recv_coordinate_sizes.sum(), dtype=mesh.geometry.x.dtype) mpi_type = numpy_to_mpi[recv_coordinates.dtype.type] send_coord_msg = [send_coordinates, send_coordinate_sizes, mpi_type] recv_coord_msg = [recv_coordinates, recv_coordinate_sizes, mpi_type] geometry_to_owner_comm.Neighbor_alltoallv(send_coord_msg, recv_coord_msg) del send_coord_msg, recv_coord_msg # Send node ordering for reordering the coordinates on output process send_nodes = np.empty(num_owned_nodes, dtype=np.int64) send_nodes[send_nodes_position] = original_node_index[:num_owned_nodes] recv_indices = np.empty(recv_coordinate_sizes.sum() // 3, dtype=np.int64) send_nodes_msg = [send_nodes, send_nodes_per_proc, MPI.INT64_T] recv_nodes_msg = [recv_indices, recv_coordinate_sizes // 3, MPI.INT64_T] geometry_to_owner_comm.Neighbor_alltoallv(send_nodes_msg, recv_nodes_msg) # Compute local ording of received nodes local_node_range = compute_local_range(mesh.comm, num_nodes_global) recv_indices -= local_node_range[0] # Sort geometry based on input index and strip to gdim gdim = mesh.geometry.dim recv_nodes = recv_coordinates.reshape(-1, 3) _geometry = np.empty(recv_nodes.shape, dtype=mesh.geometry.x.dtype) _geometry[recv_indices, :] = recv_nodes geometry = _geometry[:, :gdim].copy() del _geometry, recv_nodes assert local_node_range[1] - local_node_range[0] == geometry.shape[0] cmap = compat.cmap(mesh) cell_to_output_comm.Free() geometry_to_owner_comm.Free() # NOTE: Could in theory store partitioning information, but would not work nicely # as one would need to read this data rather than the xdmffile. # NOTE: Local geometry type hint skip is only required on DOLFINX<0.10 where # proper `dolfinx.mesh.Geometry` wrapper doesn't exist return MeshData( local_geometry=geometry, # type: ignore[arg-type] local_geometry_pos=local_node_range, num_nodes_global=num_nodes_global, local_topology=sorted_recv_dofmap, local_topology_pos=local_cell_range, num_cells_global=num_cells_global, cell_type=mesh.topology.cell_name(), degree=cmap.degree, lagrange_variant=cmap.variant, store_partition=False, partition_processes=None, ownership_array=None, ownership_offset=None, partition_range=None, partition_global=None, ) def create_function_data_on_original_mesh( u: dolfinx.fem.Function, name: typing.Optional[str] = None ) -> FunctionData: """ Create data object to save with ADIOS2 """ mesh = u.function_space.mesh # Compute what cells owned by current process should be sent to what output process # FIXME: Cache this num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local original_cell_index = mesh.topology.original_cell_index[:num_owned_cells] # Compute owner of cells on this process based on the original cell index num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global output_cell_owner = index_owner(mesh.comm, original_cell_index, num_cells_global) local_cell_range = compute_local_range(mesh.comm, num_cells_global) # Compute outgoing edges from current process to outputting process # Computes the number of cells sent to each process at the same time cell_destinations, _send_cells_per_proc = np.unique(output_cell_owner, return_counts=True) send_cells_per_proc = _send_cells_per_proc.astype(np.int32) del _send_cells_per_proc cell_to_output_comm = mesh.comm.Create_dist_graph( [mesh.comm.rank], [len(cell_destinations)], cell_destinations.tolist(), reorder=False, ) cell_sources, cell_dests, _ = cell_to_output_comm.Get_dist_neighbors() assert np.allclose(cell_dests, cell_destinations) # Compute number of recieving cells recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32) send_cells_per_proc = send_cells_per_proc.astype(np.int32) cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc) assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0] # Pack and send cell indices (used for mapping topology dofmap later) cell_insert_position = compute_insert_position( output_cell_owner, cell_destinations, send_cells_per_proc ) send_cells = np.empty_like(cell_insert_position, dtype=np.int64) send_cells[cell_insert_position] = original_cell_index recv_cells = np.empty(recv_cells_per_proc.sum(), dtype=np.int64) send_cells_msg = [send_cells, send_cells_per_proc, MPI.INT64_T] recv_cells_msg = [recv_cells, recv_cells_per_proc, MPI.INT64_T] cell_to_output_comm.Neighbor_alltoallv(send_cells_msg, recv_cells_msg) del send_cells_msg, recv_cells_msg # Map received cells to the local index local_cell_index = recv_cells - local_cell_range[0] # Pack and send cell permutation info mesh.topology.create_entity_permutations() cell_permutation_info = mesh.topology.get_cell_permutation_info()[:num_owned_cells] send_perm = np.empty_like(send_cells, dtype=np.uint32) send_perm[cell_insert_position] = cell_permutation_info recv_perm = np.empty_like(recv_cells, dtype=np.uint32) send_perm_msg = [send_perm, send_cells_per_proc, MPI.UINT32_T] recv_perm_msg = [recv_perm, recv_cells_per_proc, MPI.UINT32_T] cell_to_output_comm.Neighbor_alltoallv(send_perm_msg, recv_perm_msg) cell_permutation_info = np.empty_like(recv_perm) cell_permutation_info[local_cell_index] = recv_perm # 2. Extract function data (array is the same, keeping global indices from DOLFINx) # Dofmap is moved by the original cell index similar to the mesh geometry dofmap dofmap = u.function_space.dofmap dmap = dofmap.list num_dofs_per_cell = dmap.shape[1] dofmap_bs = dofmap.bs index_map_bs = dofmap.index_map_bs # Unroll dofmap for block size unrolled_dofmap = unroll_dofmap(dofmap.list[:num_owned_cells, :], dofmap_bs) dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1) dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1) # Convert imap index to global index imap_global = dofmap.index_map.local_to_global(dmap_loc) dofmap_global = (imap_global * index_map_bs + dmap_rem).reshape(unrolled_dofmap.shape) num_dofs_per_cell = dofmap_global.shape[1] dofmap_insert_position = unroll_insert_position(cell_insert_position, num_dofs_per_cell) # Create and send array for global dofmap send_function_dofmap = np.empty(len(dofmap_insert_position), dtype=np.int64) send_function_dofmap[dofmap_insert_position] = dofmap_global.reshape(-1) send_sizes_dofmap = send_cells_per_proc * num_dofs_per_cell recv_size_dofmap = recv_cells_per_proc * num_dofs_per_cell recv_function_dofmap = np.empty(recv_size_dofmap.sum(), dtype=np.int64) cell_to_output_comm.Neighbor_alltoallv( [send_function_dofmap, send_sizes_dofmap, MPI.INT64_T], [recv_function_dofmap, recv_size_dofmap, MPI.INT64_T], ) shaped_dofmap = recv_function_dofmap.reshape( local_cell_range[1] - local_cell_range[0], num_dofs_per_cell ).copy() _final_dofmap = np.empty_like(shaped_dofmap) _final_dofmap[local_cell_index] = shaped_dofmap final_dofmap = _final_dofmap.reshape(-1) # Get offsets of dofmap num_cells_local = local_cell_range[1] - local_cell_range[0] num_dofs_local_dmap = num_cells_local * num_dofs_per_cell dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap) local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64) local_dofmap_offsets[:] *= num_dofs_per_cell local_dofmap_offsets[:] += dofmap_imap.local_range[0] num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs local_range = np.asarray(dofmap.index_map.local_range, dtype=np.int64) * dofmap.index_map_bs func_name = name if name is not None else u.name cell_to_output_comm.Free() return FunctionData( cell_permutations=cell_permutation_info, local_cell_range=local_cell_range, num_cells_global=num_cells_global, dofmap_array=final_dofmap, dofmap_offsets=local_dofmap_offsets, values=u.x.array[:num_dofs_local].copy(), dof_range=local_range, num_dofs_global=num_dofs_global, dofmap_range=dofmap_imap.local_range, global_dofs_in_dofmap=dofmap_imap.size_global, name=func_name, ) def write_function_on_input_mesh( filename: Path | str, u: dolfinx.fem.Function, time: float = 0.0, name: typing.Optional[str] = None, mode: FileMode = FileMode.append, backend_args: dict[str, typing.Any] | None = None, backend: str = "adios2", ): """ Write function checkpoint (to be read with the input mesh). Note: Requires backend to implement {py:class}`io4dolfinx.backends.write_function`. Args: filename: The filename to write to u: The function to checkpoint time: Time-stamp associated with function at current write step mode: The mode to use (write or append) name: Name of function. If None, the name of the function is used. backend_args: Arguments to backend backend: Choice of backend module """ mesh = u.function_space.mesh function_data = create_function_data_on_original_mesh(u, name) fname = Path(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) backend_cls.write_function( fname, mesh.comm, function_data, time=time, mode=mode, backend_args=backend_args, ) def write_mesh_input_order( filename: Path | str, mesh: dolfinx.mesh.Mesh, time: float = 0.0, mode: FileMode = FileMode.write, backend: str = "adios2", backend_args: dict[str, typing.Any] | None = None, ): """ Write mesh to checkpoint file in original input ordering. Note: Requires backend to implement {py:class}`io4dolfinx.backends.write_mesh`. Args: filename: The filename to write to mesh: Mesh to checkpoint time: Time-stamp associated with function at current write step mode: The mode to use (write or append) name: Name of function. If None, the name of the function is used. backend_args: Arguments to backend backend: Choice of backend module """ mesh_data = create_original_mesh_data(mesh) fname = Path(filename) backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) backend_cls.write_mesh( fname, mesh.comm, mesh_data, backend_args=backend_args, mode=mode, time=time, ) scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/py.typed000066400000000000000000000000001517634040500246630ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/readers.py000066400000000000000000000423241517634040500252020ustar00rootroot00000000000000# Copyright (C) 2023 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from __future__ import annotations import pathlib import typing from pathlib import Path from typing import Any from mpi4py import MPI import basix import dolfinx import numpy as np import numpy.typing as npt import ufl from .backends import ReadMode, get_backend from .comm_helpers import send_dofs_and_recv_values from .utils import ( check_file_exists, compute_dofmap_pos, compute_insert_position, compute_local_range, index_owner, ) __all__ = ["read_mesh_from_legacy_h5", "read_function_from_legacy_h5", "read_point_data"] def map_dofmap(dofmap: dolfinx.graph.AdjacencyList, bs: int) -> npt.NDArray[np.int64]: """ Map xxxyyyzzz to xyzxyz """ in_dofmap = dofmap.array in_offsets = dofmap.offsets mapped_dofmap = np.empty_like(in_dofmap) for i in range(len(in_offsets) - 1): pos_begin, pos_end = ( in_offsets[i] - in_offsets[0], in_offsets[i + 1] - in_offsets[0], ) dofs_i = in_dofmap[pos_begin:pos_end] assert (pos_end - pos_begin) % bs == 0 num_dofs_local = int((pos_end - pos_begin) // bs) for k in range(bs): for j in range(num_dofs_local): mapped_dofmap[int(pos_begin + j * bs + k)] = dofs_i[int(num_dofs_local * k + j)] return mapped_dofmap.astype(np.int64) def send_cells_and_receive_dofmap_index( filename: pathlib.Path, comm: MPI.Intracomm, source_ranks: npt.NDArray[np.int32], dest_ranks: npt.NDArray[np.int32], dest_size: npt.NDArray[np.int32], output_owners: npt.NDArray[np.int32], input_cells: npt.NDArray[np.int64], dofmap_pos: npt.NDArray[np.int32], num_cells_global: np.int64, dofmap_path: str, xdofmap_path: str, bs: int, backend: str, ) -> npt.NDArray[np.int64]: """ Given a set of positions in input dofmap, give the global input index of this dofmap entry in input file. """ check_file_exists(filename) recv_size = np.zeros(len(source_ranks), dtype=np.int32) mesh_to_data_comm = comm.Create_dist_graph_adjacent( source_ranks.tolist(), dest_ranks.tolist(), reorder=False ) # Send sizes to create data structures for receiving from NeighAlltoAllv mesh_to_data_comm.Neighbor_alltoall(dest_size, recv_size) # Sort output for sending and fill send data out_cells = np.zeros(len(output_owners), dtype=np.int64) out_pos = np.zeros(len(output_owners), dtype=np.int32) proc_to_dof = np.zeros_like(input_cells, dtype=np.int32) insertion_array = compute_insert_position(output_owners, dest_ranks, dest_size) out_cells[insertion_array] = input_cells out_pos[insertion_array] = dofmap_pos proc_to_dof[insertion_array] = np.arange(len(input_cells), dtype=np.int32) del insertion_array # Prepare data-structures for receiving total_incoming = sum(recv_size) inc_cells = np.zeros(total_incoming, dtype=np.int64) inc_pos = np.zeros(total_incoming, dtype=np.intc) # Send data s_msg = [out_cells, dest_size, MPI.INT64_T] r_msg = [inc_cells, recv_size, MPI.INT64_T] mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg) s_msg = [out_pos, dest_size, MPI.INT32_T] r_msg = [inc_pos, recv_size, MPI.INT32_T] mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg) mesh_to_data_comm.Free() backend_cls = get_backend(backend) # Read dofmap from file backend_args = {"dofmap": dofmap_path, "offsets": xdofmap_path} if backend == "adios2": backend_args.update({"engine": "HDF5"}) input_dofs = backend_cls.read_dofmap(filename, comm, name="", backend_args=backend_args) # Map to xyz mapped_dofmap = map_dofmap(input_dofs, bs).astype(np.int64) # Extract dofmap data local_cell_range = compute_local_range(comm, num_cells_global) input_cell_positions = inc_cells - local_cell_range[0] in_offsets = input_dofs.offsets read_pos = (in_offsets[input_cell_positions] + inc_pos - in_offsets[0]).astype(np.int32) input_dofs = mapped_dofmap[read_pos] del input_cell_positions, read_pos # Send input dofs back to owning process data_to_mesh_comm = comm.Create_dist_graph_adjacent( dest_ranks.tolist(), source_ranks.tolist(), reorder=False ) incoming_global_dofs = np.zeros(sum(dest_size), dtype=np.int64) s_msg = [input_dofs, recv_size, MPI.INT64_T] r_msg = [incoming_global_dofs, dest_size, MPI.INT64_T] data_to_mesh_comm.Neighbor_alltoallv(s_msg, r_msg) # Sort incoming global dofs as they were inputted sorted_global_dofs = np.zeros_like(incoming_global_dofs, dtype=np.int64) assert len(incoming_global_dofs) == len(input_cells) sorted_global_dofs[proc_to_dof] = incoming_global_dofs data_to_mesh_comm.Free() return sorted_global_dofs def read_mesh_from_legacy_h5( filename: pathlib.Path, comm: MPI.Intracomm, group: str, cell_type: str = "tetrahedron", backend: str = "adios2", max_facet_to_cell_links: int = 2, ) -> dolfinx.mesh.Mesh: """ Read mesh from `h5`-file generated by legacy DOLFIN `HDF5File.write` or `XDMF.write_checkpoint`. Args: comm: MPI communicator to distribute mesh over filename: Path to `h5` or `xdmf` file group: Name of mesh in `h5`-file cell_type: What type of cell type, by default tetrahedron. backend: The IO backend to use when reading the mesh (must support legacy mesh reading, e.g., "adios2"). max_facet_to_cell_links: Maximum number of cells a facet can be connected to. """ # Make sure we use the HDF5File and check that the file is present check_file_exists(filename) backend_cls = get_backend(backend) mesh_topology, mesh_geometry, ct = backend_cls.read_legacy_mesh(filename, comm, group) if ct is not None: cell_type = ct # Create DOLFINx mesh element = basix.ufl.element( basix.ElementFamily.P, cell_type, 1, basix.LagrangeVariant.equispaced, shape=(mesh_geometry.shape[1],), ) domain = ufl.Mesh(element) try: return dolfinx.mesh.create_mesh( comm=MPI.COMM_WORLD, cells=mesh_topology, x=mesh_geometry, e=domain, partitioner=None, max_facet_to_cell_links=max_facet_to_cell_links, ) except TypeError: return dolfinx.mesh.create_mesh( comm=MPI.COMM_WORLD, cells=mesh_topology, x=mesh_geometry, e=domain, partitioner=None, ) # Should change to the commented code below when we require python # minimum version to be >=3.12 see https://github.com/python/cpython/pull/116198 # import inspect # sig = inspect.signature(dolfinx.mesh.create_mesh) # kwargs: dict[str, int] = {} # if "max_facet_to_cell_links" in list(sig.parameters.keys()): # kwargs["max_facet_to_cell_links"] = max_facet_to_cell_links # return dolfinx.mesh.create_mesh( # comm=MPI.COMM_WORLD, # cells=mesh_topology, # x=mesh_geometry, # e=domain, # partitioner=None, # **kwargs, # ) def read_function_from_legacy_h5( filename: pathlib.Path, comm: MPI.Intracomm, u: dolfinx.fem.Function, group: str = "mesh", step: typing.Optional[int] = None, vector_group: str | None = None, backend: str = "adios2", ): """ Read function from a `h5`-file generated by legacy DOLFIN `HDF5File.write` or `XDMF.write_checkpoint`. Args: comm : MPI communicator to distribute mesh over filename : Path to `h5` or `xdmf` file u : The function used to stored the read values group : Group within the `h5` file where the function is stored, by default "mesh" step : The time step used when saving the checkpoint. If not provided it will assume that the function is saved as a regular function (i.e with `HDF5File.write`) backend: The IO backend """ # Make sure we use the HDF5File and check that the file is present filename = pathlib.Path(filename) if filename.suffix == ".xdmf": filename = filename.with_suffix(".h5") if not filename.is_file(): raise FileNotFoundError(f"File {filename} does not exist") V = u.function_space mesh = u.function_space.mesh if u.function_space.element.needs_dof_transformations: raise RuntimeError( "Function-spaces requiring dof permutations are not compatible with legacy data" ) # ----------------------Step 1--------------------------------- # Compute index of input cells, and position in input dofmap local_cells, dof_pos = compute_dofmap_pos(u.function_space) input_cells = mesh.topology.original_cell_index[local_cells] # Compute mesh->input communicator # 1.1 Compute mesh->input communicator num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global backend_cls = get_backend(backend) owners: npt.NDArray[np.int32] if backend_cls.read_mode == ReadMode.serial: owners = np.zeros(len(input_cells), dtype=np.int32) elif backend_cls.read_mode == ReadMode.parallel: owners = index_owner(V.mesh.comm, input_cells, num_cells_global) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") unique_owners, owner_count = np.unique(owners, return_counts=True) # FIXME: In C++ use NBX to find neighbourhood _tmp_comm = mesh.comm.Create_dist_graph( [mesh.comm.rank], [len(unique_owners)], unique_owners, reorder=False ) source, dest, _ = _tmp_comm.Get_dist_neighbors() _tmp_comm.Free() # Strip out any / group = group.strip("/") if step is not None: group = f"{group}/{group}_{step}" vector_group = vector_group or "vector" else: vector_group = vector_group or "vector_0" # ----------------------Step 2-------------------------------- # Get global dofmap indices from input process bs = V.dofmap.bs num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global dofmap_indices = send_cells_and_receive_dofmap_index( filename, comm, np.asarray(source, dtype=np.int32), np.asarray(dest, dtype=np.int32), owner_count.astype(np.int32), owners, input_cells, dof_pos, num_cells_global, f"/{group}/cell_dofs", f"/{group}/x_cell_dofs", bs, backend=backend, ) # ----------------------Step 3--------------------------------- dof_owner: npt.NDArray[np.int32] if backend_cls.read_mode == ReadMode.serial: dof_owner = np.zeros(len(dofmap_indices), dtype=np.int32) elif backend_cls.read_mode == ReadMode.parallel: # Compute owner of global dof on distributed input data num_dof_global = V.dofmap.index_map_bs * V.dofmap.index_map.size_global dof_owner = index_owner(comm=mesh.comm, indices=dofmap_indices, N=num_dof_global) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") # Create MPI neigh comm to owner. # NOTE: USE NBX in C++ # Read input data local_array, starting_pos = backend_cls.read_hdf5_array( comm, filename, f"/{group}/{vector_group}", backend_args=None ) # Send global dof indices to correct input process, and receive value of given dof local_values = send_dofs_and_recv_values( dofmap_indices, dof_owner, comm, local_array, starting_pos ) # ----------------------Step 4--------------------------------- # Populate local part of array and scatter forward u.x.array[: len(local_values)] = local_values u.x.scatter_forward() def create_geometry_function_space(mesh: dolfinx.mesh.Mesh, N: int) -> dolfinx.fem.FunctionSpace: """Reconstruct a vector space with the N components using the geometry dofmap to ensure a 1-1 mapping between mesh nodes and DOFs.""" geom_imap = mesh.geometry.index_map() geom_dofmap = mesh.geometry.dofmap ufl_domain = mesh.ufl_domain() assert ufl_domain is not None sub_el = ufl_domain.ufl_coordinate_element().sub_elements[0] adj_list = dolfinx.cpp.graph.AdjacencyList_int32(geom_dofmap) value_shape: tuple[int, ...] if N == 1: ufl_el = sub_el value_shape = () else: ufl_el = basix.ufl.blocked_element(sub_el, shape=(N,)) value_shape = (N,) if ufl_el.dtype == np.float32: _fe_constructor = dolfinx.cpp.fem.FiniteElement_float32 _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float32 elif ufl_el.dtype == np.float64: _fe_constructor = dolfinx.cpp.fem.FiniteElement_float64 _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float64 else: raise RuntimeError(f"Unsupported type {ufl_el.dtype}") try: cpp_el = _fe_constructor(ufl_el.basix_element._e, block_shape=value_shape, symmetric=False) except TypeError: cpp_el = _fe_constructor(ufl_el.basix_element._e, block_size=N, symmetric=False) dof_layout = dolfinx.cpp.fem.create_element_dof_layout(cpp_el, []) cpp_dofmap = dolfinx.cpp.fem.DofMap(dof_layout, geom_imap, N, adj_list, N) # Create function space try: cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap) except TypeError: cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap, value_shape=value_shape) return dolfinx.fem.FunctionSpace(mesh, ufl_el, cpp_space) def read_point_data( filename: Path | str, name: str, mesh: dolfinx.mesh.Mesh, time: float | None = None, backend_args: dict[str, Any] | None = None, backend: str = "xdmf", ) -> dolfinx.fem.Function: """Read data from the nodes of a mesh. Note: Backend has to implement {py:class}`io4dolfinx.backends.read_cell_data`. Args: filename: Path to file name: Name of point data mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`. time: Time-step to read from. Returns: A function in the space equivalent to the mesh coordinate element (up to shape). """ backend_cls = get_backend(backend) dataset, local_range_start = backend_cls.read_point_data( filename=filename, name=name, comm=mesh.comm, time=time, backend_args=backend_args ) num_components = dataset.shape[1] # Create appropriate function space (based on coordinate map) V = create_geometry_function_space(mesh, num_components) uh = dolfinx.fem.Function(V, name=name, dtype=dataset.dtype) # Assume that mesh is first order for now x_dofmap = mesh.geometry.dofmap igi = np.array(mesh.geometry.input_global_indices, dtype=np.int64) # This is dependent on how the data is read in. If distributed equally this is correct global_geom_input = igi[x_dofmap] if backend_cls.read_mode == ReadMode.parallel: num_nodes_global = mesh.geometry.index_map().size_global global_geom_owner = index_owner(mesh.comm, global_geom_input.reshape(-1), num_nodes_global) elif backend_cls.read_mode == ReadMode.serial: # This is correct if everything is read in on rank 0 global_geom_owner = np.zeros(len(global_geom_input.flatten()), dtype=np.int32) else: raise NotImplementedError(f"{backend_cls.read_mode} not implemented") for i in range(num_components): arr_i = send_dofs_and_recv_values( global_geom_input.reshape(-1), global_geom_owner, mesh.comm, dataset[:, i], local_range_start, ) dof_pos = x_dofmap.reshape(-1) * num_components + i uh.x.array[dof_pos] = arr_i uh.x.scatter_forward() return uh def read_cell_data( filename: Path | str, name: str, mesh: dolfinx.mesh.Mesh, time: float | None = None, backend_args: dict[str, Any] | None = None, backend: str = "xdmf", ) -> dolfinx.fem.Function: """Read data from the nodes of a mesh. Note: Backend has to implement {py:class}`io4dolfinx.backends.read_cell_data`. Args: filename: Path to file name: Name of point data mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`. time: Time-step to read from. Returns: A function in a DG-0 space on the mesh. The cells not found in input is set to zero. """ backend_cls = get_backend(backend) topology, dofs = backend_cls.read_cell_data( filename=filename, name=name, comm=mesh.comm, time=time, backend_args=backend_args ) num_components = dofs.shape[1] shape: tuple[int, ...] if num_components == 1: shape = () else: shape = (num_components,) V = dolfinx.fem.functionspace(mesh, ("DG", 0, shape)) u = dolfinx.fem.Function(V, dtype=dofs.dtype) data_array = u.x.array.reshape(-1, num_components) for i in range(num_components): local_entities, local_values = dolfinx.io.distribute_entity_data( mesh, mesh.topology.dim, topology, dofs[:, i].copy() ) adj = dolfinx.graph.adjacencylist(local_entities) order = np.arange(len(local_values), dtype=np.int32) mt = dolfinx.mesh.meshtags_from_entities(mesh, mesh.topology.dim, adj, order) data_array[mt.indices, i] = local_values[mt.values] u.x.scatter_forward() return u scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/snapshot.py000066400000000000000000000017651517634040500254200ustar00rootroot00000000000000# Copyright (C) 2024 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from pathlib import Path from typing import Any import dolfinx from .backends import FileMode, get_backend __all__ = [ "snapshot_checkpoint", ] def snapshot_checkpoint( uh: dolfinx.fem.Function, file: Path, mode: FileMode, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """Read or write a snapshot checkpoint This checkpoint is only meant to be used on the same mesh during the same simulation. :param uh: The function to write data from or read to :param file: The file to write to or read from :param mode: Either read or write """ backend_cls = get_backend(backend) default_args = backend_cls.get_default_backend_args(backend_args) if mode not in [FileMode.write, FileMode.read]: raise ValueError(f"Got invalid mode {mode}") backend_cls.snapshot_checkpoint(file, mode, uh, default_args) scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/structures.py000066400000000000000000000100661517634040500257760ustar00rootroot00000000000000# Copyright (C) 2024 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from __future__ import annotations from dataclasses import dataclass import numpy as np import numpy.typing as npt from dolfinx.graph import AdjacencyList """Internal library classes for storing mesh and function data""" __all__ = ["MeshData", "FunctionData", "ReadMeshData", "MeshTagsData"] @dataclass class MeshData: """ Container for distributed mesh data that will be stored to disk """ #: Two-dimensional array of node coordinates local_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64] local_geometry_pos: tuple[int, int] #: Insert range on current process for geometry nodes num_nodes_global: int #: Number of nodes in global geometry array local_topology: npt.NDArray[np.int64] #: Two-dimensional connectivity array for mesh topology #: Insert range on current process for topology local_topology_pos: tuple[int, int] num_cells_global: int #: NUmber of cells in global topology cell_type: str #: The cell type degree: int #: Degree of underlying Lagrange element lagrange_variant: int #: Lagrange-variant of DOFs # Partitioning_information store_partition: bool #: Indicator if one should store mesh partitioning partition_processes: int | None = None #: Number of processes in partition ownership_array: npt.NDArray[np.int32] | None = None #: Ownership array for cells ownership_offset: npt.NDArray[np.int32] | None = None #: Ownership offset for cells partition_range: tuple[int, int] | None = ( None #: Local insert position for partitioning information ) partition_global: int | None = None #: Global size of partitioning array @dataclass class FunctionData: """ Container for distributed function data that will be written to file """ cell_permutations: npt.NDArray[np.uint32] #: Cell permutations for dofmap local_cell_range: tuple[int, int] #: Range of cells on current process num_cells_global: int #: Number of cells in global topology dofmap_array: npt.NDArray[np.int64] #: Local function dofmap (using global indices) dofmap_offsets: npt.NDArray[np.int64] #: Global dofmap offsets dofmap_range: tuple[int, int] #: Range of dofmap on current process global_dofs_in_dofmap: int #: Number of entries in global dofmap values: npt.NDArray[np.floating] #: Local function values dof_range: tuple[int, int] #: Range of local function values num_dofs_global: int #: Number of global function values name: str #: Name of function @dataclass class ReadMeshData: """Container containing data that will be read into DOLFINx""" #: Two-dimensional array containing global cell->node connectivity cells: npt.NDArray[np.int64] cell_type: str #: The cell type of the mesh x: npt.NDArray[np.floating] #: The mesh coordinates lvar: int #: The Lagrange variant degree: int #: The degree of the underlying Lagrange element #: Partitioning information per cell on the process partition_graph: AdjacencyList | None = None @dataclass class MeshTagsData: name: str #: Name of tag values: npt.NDArray # Array of values indices: npt.NDArray[np.int64] # Global indices of the entities dim: int # Topological dimension of the entities # Optional entries (used for writing to disk) num_entities_global: int | None = None #: Global number of entities that will be written out num_dofs_per_entity: int | None = None #: Number of dofs per entity #: Starting index in output array `(0<=local_start Version("0.9.0"): return V.element.signature else: return V.element.signature() def compute_insert_position( data_owner: npt.NDArray[np.int32], destination_ranks: npt.NDArray[np.int32], out_size: npt.NDArray[np.int32], ) -> npt.NDArray[np.int32]: """ Giving a list of ranks, compute the local insert position for each rank in a list sorted by destination ranks. This function is used for packing data from a given process to its destination processes. Example: .. highlight:: python .. code-block:: python data_owner = [0, 1, 1, 0, 2, 3] destination_ranks = [2,0,3,1] out_size = [1, 2, 1, 2] insert_position = compute_insert_position(data_owner, destination_ranks, out_size) Insert position is then ``[1, 4, 5, 2, 0, 3]`` """ process_pos_indicator = data_owner.reshape(-1, 1) == destination_ranks # Compute offsets for insertion based on input size send_offsets = np.zeros(len(out_size) + 1, dtype=np.intc) send_offsets[1:] = np.cumsum(out_size) assert send_offsets[-1] == len(data_owner) # Compute local insert index on each process proc_row, proc_col = np.nonzero(process_pos_indicator) cum_pos = np.cumsum(process_pos_indicator, axis=0) insert_position = cum_pos[proc_row, proc_col] - 1 # Add process offset for each local index insert_position += send_offsets[proc_col] return insert_position def unroll_insert_position( insert_position: npt.NDArray[np.int32], block_size: int ) -> npt.NDArray[np.int32]: """ Unroll insert position by a block size Example: .. highlight:: python .. code-block:: python insert_position = [1, 4, 5, 2, 0, 3] unrolled_ip = unroll_insert_position(insert_position, 3) where ``unrolled_ip = [3, 4 ,5, 12, 13, 14, 15, 16, 17, 6, 7, 8, 0, 1, 2, 9, 10, 11]`` """ unrolled_ip = np.repeat(insert_position, block_size) * block_size unrolled_ip += np.tile(np.arange(block_size), len(insert_position)) return unrolled_ip def compute_local_range(comm: MPI.Intracomm, N: np.int64): """ Divide a set of `N` objects into `M` partitions, where `M` is the size of the MPI communicator `comm`. NOTE: If N is not divisible by the number of ranks, the first `r` processes gets an extra value Returns the local range of values """ rank = comm.rank size = comm.size n = N // size r = N % size # First r processes has one extra value if rank < r: return [rank * (n + 1), (rank + 1) * (n + 1)] else: return [rank * n + r, (rank + 1) * n + r] def index_owner( comm: MPI.Intracomm, indices: npt.NDArray[np.int64], N: np.int64 ) -> npt.NDArray[np.int32]: """ Find which rank (local to comm) which owns an `index`, given that data of size `N` has been split equally among the ranks. NOTE: If `N` is not divisible by the number of ranks, the first `r` processes gets an extra value. """ size = comm.size assert (indices < N).all() n = N // size r = N % size owner = np.empty_like(indices, dtype=np.int32) inc_remainder = indices < (n + 1) * r owner[inc_remainder] = indices[inc_remainder] // (n + 1) owner[~inc_remainder] = r + (indices[~inc_remainder] - r * (n + 1)) // n return owner def unroll_dofmap(dofs: npt.NDArray[np.int32], bs: int) -> npt.NDArray[np.int32]: """ Given a two-dimensional dofmap of size `(num_cells, num_dofs_per_cell)` Expand the dofmap by its block size such that the resulting array is of size `(num_cells, bs*num_dofs_per_cell)` """ num_cells, num_dofs_per_cell = dofs.shape unrolled_dofmap = np.repeat(dofs, bs).reshape(num_cells, num_dofs_per_cell * bs) * bs unrolled_dofmap += np.tile(np.arange(bs), num_dofs_per_cell) return unrolled_dofmap def compute_dofmap_pos( V: dolfinx.fem.FunctionSpace, ) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: """ Compute a map from each owned dof in the dofmap to a single cell owned by the process, and the relative position of the dof. :param V: The function space :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the number of owned dofs (unrolled for block size) """ dofs = V.dofmap.list mesh = V.mesh num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local dofmap_bs = V.dofmap.bs num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs local_cell = np.empty( num_owned_dofs, dtype=np.int32 ) # Local cell index for each dof owned by process dof_pos = np.empty(num_owned_dofs, dtype=np.int32) # Position in dofmap for said dof unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs) markers = unrolled_dofmap < num_owned_dofs local_indices = np.broadcast_to(np.arange(markers.shape[1]), markers.shape) cell_indicator = np.broadcast_to( np.arange(num_owned_cells, dtype=np.int32).reshape(-1, 1), (num_owned_cells, markers.shape[1]), ) indicator = unrolled_dofmap[markers].reshape(-1) local_cell[indicator] = cell_indicator[markers].reshape(-1) dof_pos[indicator] = local_indices[markers].reshape(-1) return local_cell, dof_pos def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh: """ Make a copy of a mesh and potentially change the element of the coordinate element. Note: The topology is shared with the original mesh but the geometry is reconstructed. Args: mesh: Mesh to reconstruct coordinate_element_degree: Degree to use for coordinate element Returns: The new mesh """ # Extract cell properties ud = mesh.ufl_domain() assert ud is not None c_el = ud.ufl_coordinate_element() family = c_el.family_name lvar = c_el.lagrange_variant ct = c_el.cell_type # Create new UFL element new_c_el = basix.ufl.element( family, ct, coordinate_element_degree, shape=(mesh.geometry.dim,), lagrange_variant=lvar, dtype=mesh.geometry.x.dtype, ) # Extract new node coordinates V_tmp = dolfinx.fem.functionspace(mesh, new_c_el) gdim = mesh.geometry.dim x = V_tmp.tabulate_dof_coordinates()[:, :gdim] # Create new geoemtry geom_imap = V_tmp.dofmap.index_map geom_dofmap = V_tmp.dofmap.list num_nodes_local = geom_imap.size_local + geom_imap.num_ghosts original_input_indices = geom_imap.local_to_global(np.arange(num_nodes_local, dtype=np.int32)) coordinate_element = dolfinx.fem.coordinate_element( mesh.topology.cell_type, coordinate_element_degree, lvar, dtype=mesh.geometry.x.dtype ) # Could use create_geometry here when things are fixed geom = dolfinx.mesh.Geometry( type(mesh.geometry._cpp_object)( geom_imap, geom_dofmap, coordinate_element._cpp_object, x, original_input_indices ) ) # Create new mesh new_top = mesh.topology cpp_mesh = type(mesh._cpp_object)(mesh.comm, new_top._cpp_object, geom._cpp_object) return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(new_c_el)) scientificcomputing-io4dolfinx-d21fc0e/src/io4dolfinx/writers.py000066400000000000000000000132471517634040500252560ustar00rootroot00000000000000# Copyright (C) 2024-2026 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT from pathlib import Path from typing import Any from mpi4py import MPI import dolfinx import numpy as np from packaging.version import Version from . import compat from .backends import FileMode, get_backend from .structures import FunctionData, MeshData def prepare_meshdata_for_storage(mesh: dolfinx.mesh.Mesh, store_partition_info: bool) -> MeshData: """ Helper function for extracting the required data from a distributed {py:class}`dolfinx.mesh.Mesh`. Args: mesh: The mesh store_partition_info: If one should store the partitioning info Returns: Data-container with the info that should be stored. """ num_xdofs_local = mesh.geometry.index_map().size_local num_xdofs_global = mesh.geometry.index_map().size_global geometry_range = mesh.geometry.index_map().local_range gdim = mesh.geometry.dim # Convert local connectivity to globa l connectivity g_imap = mesh.geometry.index_map() g_dmap = mesh.geometry.dofmap num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global cell_range = mesh.topology.index_map(mesh.topology.dim).local_range cmap = compat.cmap(mesh) geom_layout = cmap.create_dof_layout() if hasattr(geom_layout, "num_entity_closure_dofs"): num_dofs_per_cell = geom_layout.num_entity_closure_dofs(mesh.topology.dim) else: num_dofs_per_cell = len(geom_layout.entity_closure_dofs(mesh.topology.dim, 0)) dofs_out = np.zeros((num_cells_local, num_dofs_per_cell), dtype=np.int64) assert g_dmap.shape[1] == num_dofs_per_cell dofs_out[:, :] = np.asarray( g_imap.local_to_global(g_dmap[:num_cells_local, :].reshape(-1)) ).reshape(dofs_out.shape) if store_partition_info: partition_processes = mesh.comm.size # Get partitioning if Version(dolfinx.__version__) > Version("0.9.0"): consensus_tag = 1202 cell_map = mesh.topology.index_map(mesh.topology.dim).index_to_dest_ranks(consensus_tag) else: cell_map = mesh.topology.index_map(mesh.topology.dim).index_to_dest_ranks() num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local cell_offsets = cell_map.offsets[: num_cells_local + 1] if cell_offsets[-1] == 0: cell_array = np.empty(0, dtype=np.int32) else: cell_array = cell_map.array[: cell_offsets[-1]] # Compute adjacency with current process as first entry ownership_array = np.full(num_cells_local + cell_offsets[-1], -1, dtype=np.int32) ownership_offset = cell_offsets + np.arange(len(cell_offsets), dtype=np.int32) ownership_array[ownership_offset[:-1]] = mesh.comm.rank insert_position = np.flatnonzero(ownership_array == -1) ownership_array[insert_position] = cell_array partition_map = dolfinx.common.IndexMap(mesh.comm, ownership_array.size) ownership_offset += partition_map.local_range[0] partition_range = partition_map.local_range partition_global = partition_map.size_global else: partition_processes = None ownership_array = None ownership_offset = None partition_range = None partition_global = None return MeshData( local_geometry=mesh.geometry.x[:num_xdofs_local, :gdim].copy(), local_geometry_pos=geometry_range, num_nodes_global=num_xdofs_global, local_topology=dofs_out, local_topology_pos=cell_range, num_cells_global=num_cells_global, cell_type=mesh.topology.cell_name(), degree=cmap.degree, lagrange_variant=cmap.variant, store_partition=store_partition_info, partition_processes=partition_processes, ownership_array=ownership_array, ownership_offset=ownership_offset, partition_range=partition_range, partition_global=partition_global, ) def write_mesh( filename: Path, comm: MPI.Intracomm, mesh_data: MeshData, time: float = 0.0, mode: FileMode = FileMode.write, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Write a mesh to file using ADIOS2 Args: comm: MPI communicator used in storage mesh: Internal data structure for the mesh data to save to file filename: Path to file to write to engine: ADIOS2 engine to use mode: ADIOS2 mode to use (write or append) io_name: Internal name used for the ADIOS IO object """ backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) backend_cls.write_mesh(filename, comm, mesh_data, backend_args, mode, time) def write_function( filename: Path, comm: MPI.Intracomm, u: FunctionData, time: float = 0.0, mode: FileMode = FileMode.append, backend_args: dict[str, Any] | None = None, backend: str = "adios2", ): """ Write a function to file using ADIOS2 Args: comm: MPI communicator used in storage u: Internal data structure for the function data to save to file filename: Path to file to write to engine: ADIOS2 engine to use mode: ADIOS2 mode to use (write or append) time: Time stamp associated with function io_name: Internal name used for the ADIOS IO object """ backend_cls = get_backend(backend) backend_args = backend_cls.get_default_backend_args(backend_args) backend_cls.write_function(filename, comm, u=u, time=time, backend_args=backend_args, mode=mode) scientificcomputing-io4dolfinx-d21fc0e/tests/000077500000000000000000000000001517634040500214725ustar00rootroot00000000000000scientificcomputing-io4dolfinx-d21fc0e/tests/conftest.py000066400000000000000000000166371517634040500237060ustar00rootroot00000000000000import typing from collections import ChainMap from mpi4py import MPI import dolfinx import ipyparallel as ipp import numpy as np import numpy.typing import numpy.typing as npt import pytest import io4dolfinx def find_backends(): backends = [] try: import adios2 if adios2.is_built_with_mpi: backends.append("adios2") except ModuleNotFoundError: pass try: import h5py if h5py.get_config().mpi: backends.append("h5py") except ModuleNotFoundError: pass return backends @pytest.fixture(params=find_backends(), scope="function") def backend(request): yield request.param @pytest.fixture(scope="module") def cluster(): cluster = ipp.Cluster(engines="mpi", n=2) rc = cluster.start_and_connect_sync() rc.wait_for_engines(n=2) yield rc cluster.stop_cluster_sync() @pytest.fixture(scope="function") def write_function(tmp_path): def _write_function( mesh, el, f, dtype, backend: typing.Literal["adios2", "h5py"], name="uh", append: bool = False, ) -> str: V = dolfinx.fem.functionspace(mesh, el) uh = dolfinx.fem.Function(V, dtype=dtype) uh.interpolate(f) uh.name = name el_hash = ( io4dolfinx.utils.element_signature(V) .replace(" ", "") .replace(",", "") .replace("(", "") .replace(")", "") .replace("[", "") .replace("]", "") ) # Consistent tmp dir across processes f_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) file_hash = f"{el_hash}_{np.dtype(dtype).name}" if backend == "adios2": suffix = ".bp" else: suffix = ".h5" filename = (f_path / f"mesh_{file_hash}").with_suffix(suffix) if mesh.comm.size != 1: if not append: io4dolfinx.write_mesh(filename, mesh, backend=backend) io4dolfinx.write_function(filename, uh, time=0.0, backend=backend) else: if MPI.COMM_WORLD.rank == 0: if not append: io4dolfinx.write_mesh(filename, mesh, backend=backend) io4dolfinx.write_function(filename, uh, time=0.0, backend=backend) return filename return _write_function @pytest.fixture(scope="function") def read_function(): def _read_function( comm, el, f, path, dtype, backend: typing.Literal["adios2", "h5py"], name="uh" ): mesh = io4dolfinx.read_mesh( path, comm, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, backend=backend, ) V = dolfinx.fem.functionspace(mesh, el) v = dolfinx.fem.Function(V, dtype=dtype) v.name = name io4dolfinx.read_function(path, v, backend=backend) v_ex = dolfinx.fem.Function(V, dtype=dtype) v_ex.interpolate(f) res = np.finfo(dtype).resolution np.testing.assert_allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res) return _read_function @pytest.fixture(scope="function") def get_dtype(): def _get_dtype(in_dtype: np.dtype, is_complex: bool): dtype: numpy.typing.DTypeLike if in_dtype == np.float32: if is_complex: dtype = np.complex64 else: dtype = np.float32 elif in_dtype == np.float64: if is_complex: dtype = np.complex128 else: dtype = np.float64 else: raise ValueError("Unsuported dtype") return dtype return _get_dtype @pytest.fixture(scope="function") def write_function_time_dep(tmp_path): def _write_function_time_dep( mesh, el, f0, f1, t0, t1, dtype, backend: typing.Literal["adios2", "h5py"] ) -> str: V = dolfinx.fem.functionspace(mesh, el) uh = dolfinx.fem.Function(V, dtype=dtype) uh.interpolate(f0) el_hash = ( io4dolfinx.utils.element_signature(V) .replace(" ", "") .replace(",", "") .replace("(", "") .replace(")", "") .replace("[", "") .replace("]", "") ) file_hash = f"{el_hash}_{np.dtype(dtype).name}" # Consistent tmp dir across processes f_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) if backend == "adios2": suffix = ".bp" else: suffix = ".h5" filename = (f_path / f"mesh_{file_hash}").with_suffix(suffix) if mesh.comm.size != 1: io4dolfinx.write_mesh(filename, mesh, backend=backend) io4dolfinx.write_function(filename, uh, time=t0, backend=backend) uh.interpolate(f1) io4dolfinx.write_function(filename, uh, time=t1, backend=backend) else: if MPI.COMM_WORLD.rank == 0: io4dolfinx.write_mesh(filename, mesh, backend=backend) io4dolfinx.write_function(filename, uh, time=t0, backend=backend) uh.interpolate(f1) io4dolfinx.write_function(filename, uh, time=t1, backend=backend) return filename return _write_function_time_dep @pytest.fixture(scope="function") def read_function_time_dep(): def _read_function_time_dep( comm, el, f0, f1, t0, t1, path, dtype, backend: typing.Literal["adios2", "h5py"] ): mesh = io4dolfinx.read_mesh( path, comm, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, backend=backend, ) V = dolfinx.fem.functionspace(mesh, el) v = dolfinx.fem.Function(V, dtype=dtype) io4dolfinx.read_function(path, v, time=t1, backend=backend) v_ex = dolfinx.fem.Function(V, dtype=dtype) v_ex.interpolate(f1) res = np.finfo(dtype).resolution assert np.allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res) io4dolfinx.read_function(path, v, time=t0, backend=backend) v_ex = dolfinx.fem.Function(V, dtype=dtype) v_ex.interpolate(f0) res = np.finfo(dtype).resolution assert np.allclose(v.x.array, v_ex.x.array, atol=10 * res, rtol=10 * res) return _read_function_time_dep def _generate_reference_map( mesh: dolfinx.mesh.Mesh, meshtag: dolfinx.mesh.MeshTags, comm: MPI.Intracomm, root: int, ) -> typing.Optional[dict[str, tuple[int, npt.NDArray]]]: """ Helper function to generate map from meshtag value to its corresponding index and midpoint. Args: mesh: The mesh meshtag: The associated meshtag comm: MPI communicator to gather the map from all processes with root (int): Rank to store data on Returns: Root rank returns the map, all other ranks return None """ mesh.topology.create_connectivity(meshtag.dim, mesh.topology.dim) midpoints = dolfinx.mesh.compute_midpoints(mesh, meshtag.dim, meshtag.indices) e_map = mesh.topology.index_map(meshtag.dim) value_to_midpoint = {} for index, value in zip(meshtag.indices, meshtag.values): value_to_midpoint[value] = ( int(e_map.local_range[0] + index), midpoints[index], ) global_map = comm.gather(value_to_midpoint, root=root) if comm.rank == root: return dict(ChainMap(*global_map)) # type: ignore return None @pytest.fixture def generate_reference_map(): return _generate_reference_map scientificcomputing-io4dolfinx-d21fc0e/tests/create_legacy_checkpoint.py000066400000000000000000000036041517634040500270450ustar00rootroot00000000000000# Copyright (C) 2024 Jørgen Schartum Dokken # # This file is part of io4dolfinx. # # SPDX-License-Identifier: MIT """ Functions to create checkpoints with adios4dolfinx v0.7.x """ import argparse import pathlib from importlib.metadata import version from mpi4py import MPI import adios4dolfinx import dolfinx import numpy as np a4d_version = version("adios4dolfinx") assert a4d_version < "0.7.2", ( f"Creating a legacy checkpoint requires adios4dolfinx < 0.7.2, you have {a4d_version}." ) def f(x): values = np.zeros((2, x.shape[1]), dtype=np.float64) values[0] = x[0] values[1] = -x[1] return values def write_checkpoint(filename, mesh, el, f): V = dolfinx.fem.FunctionSpace(mesh, el) uh = dolfinx.fem.Function(V, dtype=np.float64) uh.interpolate(f) adios4dolfinx.write_mesh(V.mesh, filename) adios4dolfinx.write_function(uh, filename) def verify_checkpoint(filename, el, f): mesh = adios4dolfinx.read_mesh( MPI.COMM_WORLD, filename, "BP4", dolfinx.mesh.GhostMode.shared_facet ) V = dolfinx.fem.FunctionSpace(mesh, el) uh = dolfinx.fem.Function(V, dtype=np.float64) adios4dolfinx.read_function(uh, filename) u_ex = dolfinx.fem.Function(V, dtype=np.float64) u_ex.interpolate(f) np.testing.assert_allclose(u_ex.x.array, uh.x.array, atol=1e-15) if __name__ == "__main__": parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--output-dir", type=str, default="legacy_checkpoint", dest="dir") inputs = parser.parse_args() path = pathlib.Path(inputs.dir) path.mkdir(exist_ok=True, parents=True) filename = path / "adios4dolfinx_checkpoint.bp" mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) el = ("N1curl", 3) write_checkpoint(filename, mesh, el, f) MPI.COMM_WORLD.Barrier() verify_checkpoint(filename, el, f) scientificcomputing-io4dolfinx-d21fc0e/tests/create_legacy_data.py000066400000000000000000000143011517634040500256230ustar00rootroot00000000000000# Copyright (C) 2023 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT """ Functions to create checkpoints with Legacy dolfin """ import argparse import pathlib import dolfin import numpy as np import ufl_legacy as ufl def create_reference_data( h5_file: pathlib.Path, xdmf_file: pathlib.Path, mesh_name: str, function_name: str, family: str, degree: int, function_name_vec: str, ) -> dolfin.Function: mesh = dolfin.UnitCubeMesh(1, 1, 1) V = dolfin.FunctionSpace(mesh, family, degree) W = dolfin.VectorFunctionSpace(mesh, family, degree) x = dolfin.SpatialCoordinate(mesh) f0 = ufl.conditional(ufl.gt(x[0], 0.5), x[1], 2 * x[0]) v0 = dolfin.project(f0, V) w0 = dolfin.interpolate(dolfin.Expression(("x[0]", "3*x[2]", "7*x[1]"), degree=1), W) v1 = dolfin.interpolate(dolfin.Expression("x[0]", degree=1), V) w1 = dolfin.interpolate(dolfin.Expression(("x[0]", "0", "x[1]"), degree=1), W) tt = dolfin.Function(dolfin.FunctionSpace(mesh, "DG", 0)) tt.interpolate(dolfin.Expression(("x[0]+x[1]"), degree=0)) ff = dolfin.MeshFunction("size_t", mesh, mesh.topology().dim() - 1, 1) class Boundary(dolfin.SubDomain): def inside(self, x, on_boundary): return on_boundary Boundary().mark(ff, 2) with dolfin.HDF5File(mesh.mpi_comm(), str(h5_file), "w") as hdf: hdf.write(mesh, mesh_name) hdf.write(v0, function_name) hdf.write(w0, function_name_vec) hdf.write(tt, function_name + "DG") hdf.write(ff, "facet_tags") with dolfin.XDMFFile(mesh.mpi_comm(), str(xdmf_file)) as xdmf: xdmf.write(mesh) xdmf.write_checkpoint(v0, function_name, 0, dolfin.XDMFFile.Encoding.HDF5, append=True) xdmf.write_checkpoint(w0, function_name_vec, 0, dolfin.XDMFFile.Encoding.HDF5, append=True) xdmf.write_checkpoint(v1, function_name, 1, dolfin.XDMFFile.Encoding.HDF5, append=True) xdmf.write_checkpoint(w1, function_name_vec, 1, dolfin.XDMFFile.Encoding.HDF5, append=True) # Legacy DG-0 checkpoint has to be scalar xdmf.write_checkpoint( tt, function_name + "DG_checkpoint", 0.5, dolfin.XDMFFile.Encoding.HDF5, append=True ) return v0, w0, v1, w1 def verify_hdf5( v_ref: dolfin.Function, w_ref: dolfin.Function, h5_file: pathlib.Path, mesh_name: str, function_name: str, family: str, degree: int, function_name_vec: str, ): mesh = dolfin.Mesh() with dolfin.HDF5File(mesh.mpi_comm(), str(h5_file), "r") as hdf: hdf.read(mesh, mesh_name, False) V = dolfin.FunctionSpace(mesh, family, degree) v = dolfin.Function(V) hdf.read(v, function_name) W = dolfin.VectorFunctionSpace(mesh, family, degree) w = dolfin.Function(W) hdf.read(w, function_name_vec) assert np.allclose(v.vector().get_local(), v_ref.vector().get_local()) assert np.allclose(w.vector().get_local(), w_ref.vector().get_local()) def verify_xdmf( v0_ref: dolfin.Function, w0_ref: dolfin.Function, v1_ref: dolfin.Function, w1_ref: dolfin.Function, xdmf_file: pathlib.Path, function_name: str, family: str, degree: int, function_name_vec: str, ): mesh = dolfin.Mesh() with dolfin.XDMFFile(mesh.mpi_comm(), str(xdmf_file)) as xdmf: xdmf.read(mesh) V = dolfin.FunctionSpace(mesh, family, degree) v0 = dolfin.Function(V) xdmf.read_checkpoint(v0, function_name, 0) v1 = dolfin.Function(V) xdmf.read_checkpoint(v1, function_name, 1) W = dolfin.VectorFunctionSpace(mesh, family, degree) w0 = dolfin.Function(W) xdmf.read_checkpoint(w0, function_name_vec, 0) w1 = dolfin.Function(W) xdmf.read_checkpoint(w1, function_name_vec, 1) assert np.allclose(v0.vector().get_local(), v0_ref.vector().get_local()) assert np.allclose(w0.vector().get_local(), w0_ref.vector().get_local()) assert np.allclose(v1.vector().get_local(), v1_ref.vector().get_local()) assert np.allclose(w1.vector().get_local(), w1_ref.vector().get_local()) def create_reference_P1_data(filename: pathlib.Path, mesh_name: str, function_name: str): mesh = dolfin.RectangleMesh(dolfin.Point(0.1, 0.2), dolfin.Point(2, 3), 12, 13) mesh.rename(mesh_name, mesh_name) W = dolfin.VectorFunctionSpace(mesh, "Lagrange", 1) w = dolfin.Function(W, name=function_name) w.interpolate(dolfin.Expression(("x[0]", "x[1]-x[0]"), degree=1)) dolfin.VTKFile(str(filename.with_suffix(".vtu")), "ascii").write(mesh) dolfin.File(str(filename.with_name(filename.name + "_func").with_suffix(".pvd"))) << w if __name__ == "__main__": parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--family", type=str, default="DG") parser.add_argument("--degree", type=int, default=2) parser.add_argument("--output-dir", type=str, default="legacy", dest="dir") parser.add_argument("--mesh-name", type=str, default="mesh", dest="name") parser.add_argument("--function-name", type=str, default="v", dest="f_name") parser.add_argument("--function-name-vec", type=str, default="w", dest="f_name_vec") inputs = parser.parse_args() path = pathlib.Path(inputs.dir) path.mkdir(exist_ok=True, parents=True) h5_filename = path / f"{inputs.name}.h5" xdmf_filename = path / f"{inputs.name}_checkpoint.xdmf" p1_filename = path / f"{inputs.name}_P1" v0_ref, w0_ref, v1_ref, w1_ref = create_reference_data( h5_filename, xdmf_filename, inputs.name, inputs.f_name, inputs.family, inputs.degree, inputs.f_name_vec, ) verify_hdf5( v0_ref, w0_ref, h5_filename, inputs.name, inputs.f_name, inputs.family, inputs.degree, inputs.f_name_vec, ) verify_xdmf( v0_ref, w0_ref, v1_ref, w1_ref, xdmf_filename, inputs.f_name, inputs.family, inputs.degree, inputs.f_name_vec, ) P1_ref = create_reference_P1_data( filename=p1_filename, mesh_name=inputs.name, function_name=inputs.f_name_vec ) scientificcomputing-io4dolfinx-d21fc0e/tests/test_attributes.py000066400000000000000000000033421517634040500252730ustar00rootroot00000000000000from mpi4py import MPI import numpy as np import pytest from packaging.version import parse as _v import io4dolfinx @pytest.mark.parametrize("comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_attributes(comm, backend, tmp_path): if backend == "adios2": import adios2 if _v(np.__version__) >= _v("2.0.0") and _v(adios2.__version__) < _v("2.10.2"): pytest.skip(reason="Cannot use numpy>=2.0.0 and adios2<2.10.2") attributes1 = { "a": np.array([1, 2, 3], dtype=np.uint8), "b": np.array([4, 5], dtype=np.uint8), } attributes2 = { "c": np.array([6], dtype=np.uint8), "d": np.array([7, 8, 9, 10], dtype=np.float64), } fname = comm.bcast(tmp_path, root=0) fname = fname / "attributes" suffix = ".bp" if backend == "adios2" else ".h5" file = fname.with_suffix(suffix) # print(comm.size) io4dolfinx.write_attributes( comm=comm, filename=file, name="group1", attributes=attributes1, backend=backend ) io4dolfinx.write_attributes( comm=comm, filename=file, name="group2", attributes=attributes2, backend=backend ) loaded_attributes1 = io4dolfinx.read_attributes( comm=comm, filename=file, name="group1", backend=backend ) loaded_attributes2 = io4dolfinx.read_attributes( comm=comm, filename=file, name="group2", backend=backend ) for k, v in loaded_attributes1.items(): assert np.allclose(v, attributes1[k]) for k, v in attributes1.items(): assert np.allclose(v, loaded_attributes1[k]) for k, v in loaded_attributes2.items(): assert np.allclose(v, attributes2[k]) for k, v in attributes2.items(): assert np.allclose(v, loaded_attributes2[k]) scientificcomputing-io4dolfinx-d21fc0e/tests/test_checkpointing.py000066400000000000000000000225001517634040500257270ustar00rootroot00000000000000import itertools from pathlib import Path from mpi4py import MPI import basix import basix.ufl import dolfinx import numpy as np import pytest import io4dolfinx dtypes = [np.float64, np.float32] # Mesh geometry dtypes write_comm = [MPI.COMM_SELF, MPI.COMM_WORLD] # Communicators for creating mesh two_dimensional_cell_types = [ dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral, ] three_dimensional_cell_types = [ dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron, ] two_dim_combinations = itertools.product(dtypes, two_dimensional_cell_types, write_comm) three_dim_combinations = itertools.product(dtypes, three_dimensional_cell_types, write_comm) @pytest.fixture(params=two_dim_combinations, scope="module") def mesh_2D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_square(write_comm, 10, 10, cell_type=cell_type, dtype=dtype) return mesh @pytest.fixture(params=three_dim_combinations, scope="module") def mesh_3D(request): dtype, cell_type, write_comm = request.param M = 5 mesh = dolfinx.mesh.create_unit_cube(write_comm, M, M, M, cell_type=cell_type, dtype=dtype) return mesh @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_P_2D( read_comm, family, degree, is_complex, mesh_2D, get_dtype, write_function, read_function, backend, ): mesh = mesh_2D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = x[0] if is_complex: values[0] += 1j * x[1] values[1] -= 3j * x[1] return values hash = write_function(mesh, el, f, f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_P_3D( read_comm, family, degree, is_complex, mesh_3D, get_dtype, write_function, read_function, backend, ): mesh = mesh_3D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.pi + x[0] values[1] = x[1] + 2 * x[0] values[2] = np.cos(x[2]) if is_complex: values[0] -= 2j * x[2] values[2] += 1j * x[1] return values hash = write_function(mesh, el, f, f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_P_2D_time( read_comm, family, degree, is_complex, mesh_2D, get_dtype, write_function_time_dep, read_function_time_dep, backend, ): mesh = mesh_2D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) def f0(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = x[0] if is_complex: values[0] += x[1] * 1j values[1] -= 3j * x[1] return values def f1(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = 2 * np.full(x.shape[1], np.pi) + x[0] values[1] = -x[0] + 2 * x[1] if is_complex: values[0] += x[1] * 1j values[1] += 3j * x[1] return values t0 = 0.8 t1 = 0.6 hash = write_function_time_dep(mesh, el, f0, f1, t0, t1, f_dtype, backend) MPI.COMM_WORLD.Barrier() read_function_time_dep(read_comm, el, f0, f1, t0, t1, hash, f_dtype, backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_P_3D_time( read_comm, family, degree, is_complex, mesh_3D, get_dtype, write_function_time_dep, read_function_time_dep, backend, ): mesh = mesh_3D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.pi + x[0] values[1] = x[1] + 2 * x[0] values[2] = np.cos(x[2]) if is_complex: values[0] += 2j * x[2] values[2] += 5j * x[1] return values def g(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = x[0] values[1] = 2 * x[0] values[2] = x[0] if is_complex: values[0] += np.pi * 2j * x[2] values[1] += 1j * x[2] values[2] += 1j * np.cos(x[1]) return values t0 = 0.1 t1 = 1.3 hash = write_function_time_dep(mesh, el, g, f, t0, t1, f_dtype, backend) MPI.COMM_WORLD.Barrier() read_function_time_dep(read_comm, el, g, f, t0, t1, hash, f_dtype, backend) @pytest.mark.parametrize( "func, args", [ (io4dolfinx.read_attributes, ("nonexisting_file", MPI.COMM_WORLD, "")), (io4dolfinx.read_timestamps, ("nonexisting_file", MPI.COMM_WORLD, "")), (io4dolfinx.read_meshtags, ("nonexisting_file", MPI.COMM_WORLD, None, "")), (io4dolfinx.read_function, ("nonexisting_file", None)), (io4dolfinx.read_mesh, ("nonexisting_file", MPI.COMM_WORLD)), ], ) def test_read_nonexisting_file_raises_FileNotFoundError(func, args, backend): if backend == "adios": suffix = ".bp" else: suffix = ".h5" file = Path(args[0]).with_suffix(suffix) with pytest.raises(FileNotFoundError): func(file, *args[1:], backend=backend) def test_read_function_with_invalid_name_raises_KeyError(tmp_path, backend): if backend == "adios": suffix = ".bp" else: suffix = ".h5" comm = MPI.COMM_WORLD f_path = comm.bcast(tmp_path, root=0) filename = (f_path / "func").with_suffix(suffix) mesh = dolfinx.mesh.create_unit_square(comm, 10, 10, cell_type=dolfinx.mesh.CellType.triangle) V = dolfinx.fem.functionspace(mesh, ("P", 1)) u = dolfinx.fem.Function(V) io4dolfinx.write_function(filename, u, time=0, name="some_name", backend=backend) io4dolfinx.write_function(filename, u, time=0, name="some_other_name", backend=backend) # variables = set(sorted(["some_name", "some_other_name"])) with pytest.raises(KeyError): io4dolfinx.read_function(filename, u, time=0, name="nonexisting_name", backend=backend) # assert e.value.args[0] == ( # f"nonexisting_name not found in {filename}. Did you mean one of {variables}?" # ) def test_read_timestamps(get_dtype, mesh_2D, tmp_path, backend): if backend == "adios": suffix = ".bp" else: suffix = ".h5" mesh = mesh_2D dtype = get_dtype(mesh.geometry.x.dtype, False) el = basix.ufl.element( "Lagrange", mesh.basix_cell(), 1, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V, dtype=dtype, name="u") v = dolfinx.fem.Function(V, dtype=dtype, name="v") f_path = mesh.comm.bcast(tmp_path, root=0) filename = (f_path / "read_time_stamps").with_suffix(suffix) t_u = [0.1, 1.4] t_v = [0.45, 1.2] io4dolfinx.write_mesh(filename, mesh, backend=backend) io4dolfinx.write_function(filename, u, time=t_u[0], backend=backend) io4dolfinx.write_function(filename, v, time=t_v[0], backend=backend) io4dolfinx.write_function(filename, u, time=t_u[1], backend=backend) io4dolfinx.write_function(filename, v, time=t_v[1], backend=backend) timestamps_u = io4dolfinx.read_timestamps( comm=mesh.comm, filename=filename, function_name="u", backend=backend ) timestamps_v = io4dolfinx.read_timestamps( comm=mesh.comm, filename=filename, function_name="v", backend=backend ) assert np.allclose(timestamps_u, t_u) assert np.allclose(timestamps_v, t_v) scientificcomputing-io4dolfinx-d21fc0e/tests/test_checkpointing_vector.py000066400000000000000000000162331517634040500273170ustar00rootroot00000000000000import itertools from mpi4py import MPI import basix import basix.ufl import dolfinx import numpy as np import pytest dtypes = [np.float64, np.float32] # Mesh geometry dtypes write_comm = [MPI.COMM_SELF, MPI.COMM_WORLD] # Communicators for creating mesh simplex_two_dim = itertools.product(dtypes, [dolfinx.mesh.CellType.triangle], write_comm) simplex_three_dim = itertools.product(dtypes, [dolfinx.mesh.CellType.tetrahedron], write_comm) non_simplex_two_dim = itertools.product(dtypes, [dolfinx.mesh.CellType.quadrilateral], write_comm) non_simplex_three_dim = itertools.product(dtypes, [dolfinx.mesh.CellType.hexahedron], write_comm) @pytest.fixture(params=simplex_two_dim, scope="module") def simplex_mesh_2D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_square(write_comm, 10, 10, cell_type=cell_type, dtype=dtype) return mesh @pytest.fixture(params=simplex_three_dim, scope="module") def simplex_mesh_3D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_cube(write_comm, 5, 5, 5, cell_type=cell_type, dtype=dtype) return mesh @pytest.fixture(params=non_simplex_two_dim, scope="module") def non_simplex_mesh_2D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_square(write_comm, 10, 10, cell_type=cell_type, dtype=dtype) return mesh @pytest.fixture(params=non_simplex_three_dim, scope="module") def non_simplex_mesh_3D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_cube(write_comm, 5, 5, 5, cell_type=cell_type, dtype=dtype) return mesh @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["N1curl", "RT"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_2D( read_comm, family, degree, is_complex, simplex_mesh_2D, get_dtype, write_function, read_function, backend, ): mesh = simplex_mesh_2D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = x[1] if is_complex: values[0] += 2j * x[1] values[1] += 2j * x[0] return values fname = write_function(mesh, el, f, f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, fname, f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["N1curl", "RT"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_3D( read_comm, family, degree, is_complex, simplex_mesh_3D, get_dtype, write_function, read_function, backend, ): mesh = simplex_mesh_3D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) values[1] = x[1] + 2 * x[0] values[2] = np.cos(x[2]) if is_complex: values[0] += 2j * x[2] values[1] += 2j * np.cos(x[2]) return values fname = write_function(mesh, el, f, dtype=f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, fname, dtype=f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["RTCF"]) @pytest.mark.parametrize("degree", [1, 2, 3]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_2D_quad( read_comm, family, degree, is_complex, non_simplex_mesh_2D, get_dtype, write_function, read_function, backend, ): mesh = non_simplex_mesh_2D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) values[1] = x[1] + 2 * x[0] if is_complex: values[0] += 2j * x[2] values[1] += 2j * np.cos(x[2]) return values hash = write_function(mesh, el, f, f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["NCF"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_hex( read_comm, family, degree, is_complex, non_simplex_mesh_3D, get_dtype, write_function, read_function, backend, ): mesh = non_simplex_mesh_3D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = np.cos(x[2]) values[2] = x[0] if is_complex: values[0] += 2j * x[2] values[2] -= 1j * x[1] return values hash = write_function(mesh, el, f, dtype=f_dtype, backend=backend) MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash, dtype=f_dtype, backend=backend) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["RTCF"]) @pytest.mark.parametrize("degree", [1, 2, 3]) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_read_write_multiple( read_comm, family, degree, is_complex, non_simplex_mesh_2D, get_dtype, write_function, read_function, backend, ): mesh = non_simplex_mesh_2D f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) values[1] = x[1] + 2 * x[0] if is_complex: values[0] -= 2j * x[2] values[1] += 2j * np.cos(x[2]) return values def g(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = 2 * x[1] values[1] = 3 * x[0] if is_complex: values[0] += 3j * x[0] values[1] += 2j * x[0] * x[1] return values hash_f = write_function(mesh, el, f, dtype=f_dtype, name="f", append=False, backend=backend) hash_g = write_function(mesh, el, g, dtype=f_dtype, name="g", append=True, backend=backend) assert hash_f == hash_g MPI.COMM_WORLD.Barrier() read_function(read_comm, el, f, hash_f, dtype=f_dtype, name="f", backend=backend) read_function(read_comm, el, g, hash_g, dtype=f_dtype, name="g", backend=backend) scientificcomputing-io4dolfinx-d21fc0e/tests/test_exodus.py000066400000000000000000000070431517634040500244160ustar00rootroot00000000000000import urllib.request from enum import Enum from mpi4py import MPI import numpy as np import pytest import io4dolfinx netcdf4 = pytest.importorskip("netCDF4") class DownloadStatus(Enum): success = 1 failed = -1 no_connection = -2 def download_file_if_not_exists( url, filename, comm: MPI.Intracomm = MPI.COMM_WORLD, rank: int = 0 ) -> DownloadStatus: status = DownloadStatus.failed if comm.rank == rank: if not filename.exists(): try: urllib.request.urlretrieve(url, filename) status = DownloadStatus.success except urllib.error.URLError as e: if str(e) == "": status = DownloadStatus.no_connection else: status = DownloadStatus.failed else: status = DownloadStatus.success status = comm.bcast(status, root=rank) comm.Barrier() return status def test_read_mesh_and_cell_data(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = tmp_path / "openmc_master_out_openmc0.e" url = "https://github.com/neams-th-coe/cardinal/blob/devel/test/tests/neutronics/feedback/single_level/gold/openmc_master_out_openmc0.e?raw=true" status = download_file_if_not_exists(url, filename) if status == DownloadStatus.no_connection: pytest.skip("No internet connection") mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD, backend="exodus") io4dolfinx.read_meshtags(filename, mesh, meshtag_name="cell", backend="exodus") io4dolfinx.read_meshtags(filename, mesh, meshtag_name="facet", backend="exodus") io4dolfinx.read_cell_data( filename, name="cell_temperature", mesh=mesh, backend="exodus", time=1.0 ) def test_read_mesh_point_data(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = tmp_path / "openmc_master_out_openmc0.e" url = "https://github.com/idaholab/moose/blob/next/test/tests/kernels/2d_diffusion/gold/matdiffusion_out.e?raw=true" status = download_file_if_not_exists(url, filename) if status == DownloadStatus.no_connection: pytest.skip("No internet connection") mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD, backend="exodus") u = io4dolfinx.read_point_data(filename, name="u", mesh=mesh, backend="exodus", time=1.0) assert mesh.topology.index_map(mesh.topology.dim).size_global == 4 assert mesh.geometry.index_map().size_global == 9 assert np.isclose(mesh.comm.allreduce(np.max(u.x.array), op=MPI.MAX), 1.1140844375981802) def test_read_second_order_mesh(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = tmp_path / "box-test_out_nek0.e" url = "https://github.com/neams-th-coe/cardinal/blob/devel/test/tests/deformation/simple-cube/gold/box-test_out_nek0.e?raw=true" status = download_file_if_not_exists(url, filename) if status == DownloadStatus.no_connection: pytest.skip("No internet connection") mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD, backend="exodus", time=5) assert mesh.topology.index_map(mesh.topology.dim).size_global == 64 assert mesh.geometry.index_map().size_global == 1728 u_x = io4dolfinx.read_point_data(filename, name="disp_x", mesh=mesh, backend="exodus", time=5.0) ref_min = -0.7699306351843485 ref_max = 0.7699306351843485 assert np.isclose(mesh.comm.allreduce(np.min(u_x.x.array), op=MPI.MIN), ref_min) assert np.isclose(mesh.comm.allreduce(np.max(u_x.x.array), op=MPI.MAX), ref_max) scientificcomputing-io4dolfinx-d21fc0e/tests/test_geometry_dofmap_reconstruction.py000066400000000000000000000022751517634040500314330ustar00rootroot00000000000000from mpi4py import MPI import basix import dolfinx import numpy as np import pytest from io4dolfinx.readers import create_geometry_function_space @pytest.mark.parametrize( "cell_type", [dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral] ) @pytest.mark.parametrize("value_shape", [(), (1,), (4,)]) @pytest.mark.parametrize("N", [3, 20]) @pytest.mark.parametrize("M", [8, 9]) @pytest.mark.parametrize( "ghost_mode", [dolfinx.mesh.GhostMode.shared_facet, dolfinx.mesh.GhostMode.none] ) def test_dofmap_construction(cell_type, value_shape, N, M, ghost_mode): mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, N, M, cell_type=cell_type, ghost_mode=ghost_mode ) el = basix.ufl.element("Lagrange", mesh.basix_cell(), 1) if value_shape == (): b_el = el else: b_el = basix.ufl.blocked_element(el, value_shape) MPI.COMM_WORLD.barrier() V_ref = dolfinx.fem.functionspace(mesh, b_el) MPI.COMM_WORLD.barrier() bs = int(np.prod(value_shape)) V = create_geometry_function_space(mesh, bs) assert V.dofmap.bs == V_ref.dofmap.bs assert V.dofmap.bs == bs np.testing.assert_allclose(V.dofmap.list, V_ref.dofmap.list) scientificcomputing-io4dolfinx-d21fc0e/tests/test_legacy_readers.py000066400000000000000000000227161517634040500260640ustar00rootroot00000000000000# Copyright (C) 2023-2026 Jørgen Schartum Dokken # # This file is part of io4dolfinx # # SPDX-License-Identifier: MIT import inspect import pathlib from mpi4py import MPI import dolfinx import numpy as np import pytest import ufl from io4dolfinx import ( read_cell_data, read_function, read_function_from_legacy_h5, read_function_names, read_mesh, read_mesh_from_legacy_h5, read_meshtags, read_point_data, ) def test_legacy_mesh(backend): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh.h5").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") mesh = read_mesh_from_legacy_h5(filename=path, comm=comm, group="/mesh", backend=backend) assert mesh.topology.dim == 3 volume = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=mesh))), op=MPI.SUM, ) surface = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.ds(domain=mesh))), op=MPI.SUM, ) assert np.isclose(volume, 1) assert np.isclose(surface, 6) mesh.topology.create_entities(mesh.topology.dim - 1) num_facets = mesh.topology.index_map(mesh.topology.dim - 1).size_global assert num_facets == 18 def test_read_legacy_mesh_from_checkpoint(backend): comm = MPI.COMM_WORLD filename = (pathlib.Path("legacy") / "mesh_checkpoint.h5").absolute() if not filename.exists(): pytest.skip(f"{filename} does not exist") mesh = read_mesh_from_legacy_h5( filename=filename, comm=comm, group="/Mesh/mesh", backend=backend ) assert mesh.topology.dim == 3 volume = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=mesh))), op=MPI.SUM, ) surface = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.ds(domain=mesh))), op=MPI.SUM, ) assert np.isclose(volume, 1) assert np.isclose(surface, 6) mesh.topology.create_entities(mesh.topology.dim - 1) num_facets = mesh.topology.index_map(mesh.topology.dim - 1).size_global assert num_facets == 18 def test_legacy_function(backend): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh.h5").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") mesh = read_mesh_from_legacy_h5(path, comm, "/mesh", backend=backend) ff = read_meshtags( path, mesh, "facet_tags", backend_args={"legacy": True, "engine": "HDF5"}, backend=backend ) fdim = mesh.topology.dim - 1 assert ff.dim == fdim boundary_facets = ff.find(2) interior_facets = ff.find(1) mesh.topology.create_connectivity(fdim, fdim + 1) num_facets_local = ( mesh.topology.index_map(fdim).size_local + mesh.topology.index_map(fdim).num_ghosts ) true_exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) assert len(boundary_facets) == len(true_exterior_facets) assert np.isin(boundary_facets, true_exterior_facets).all() true_interior = np.ones(num_facets_local, dtype=np.int8) true_interior[true_exterior_facets] = 0 interior_marker = np.flatnonzero(true_interior) assert len(interior_marker) == len(interior_facets) assert np.isin(interior_facets, interior_marker).all() V = dolfinx.fem.functionspace(mesh, ("DG", 2)) u = ufl.TrialFunction(V) v = ufl.TestFunction(V) a = ufl.inner(u, v) * ufl.dx x = ufl.SpatialCoordinate(mesh) f = ufl.conditional(ufl.gt(x[0], 0.5), x[1], 2 * x[0]) L = ufl.inner(f, v) * ufl.dx if not dolfinx.has_petsc4py: pytest.skip("dolfinx not configured with PETSc4py") from dolfinx.fem.petsc import LinearProblem uh = dolfinx.fem.Function(V) if "petsc_options_prefix" in inspect.signature(LinearProblem.__init__).parameters.keys(): extra_options = {"petsc_options_prefix": "legacy_test"} else: extra_options = {} problem = LinearProblem( a, L, bcs=[], u=uh, petsc_options={"ksp_type": "preonly", "pc_type": "lu"}, **extra_options ) problem.solve() u_in = dolfinx.fem.Function(V) read_function_from_legacy_h5(path, mesh.comm, u_in, group="v", backend=backend) np.testing.assert_allclose(uh.x.array, u_in.x.array, atol=1e-14) W = dolfinx.fem.functionspace(mesh, ("DG", 2, (mesh.geometry.dim,))) wh = dolfinx.fem.Function(W) wh.interpolate(lambda x: (x[0], 3 * x[2], 7 * x[1])) w_in = dolfinx.fem.Function(W) read_function_from_legacy_h5(path, mesh.comm, w_in, group="w", backend=backend) np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) def test_read_legacy_function_from_checkpoint(backend): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh_checkpoint.h5").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") from dolfinx.fem.petsc import LinearProblem mesh = read_mesh_from_legacy_h5(path, comm, "/Mesh/mesh", backend=backend) V = dolfinx.fem.functionspace(mesh, ("DG", 2)) u = ufl.TrialFunction(V) v = ufl.TestFunction(V) a = ufl.inner(u, v) * ufl.dx x = ufl.SpatialCoordinate(mesh) f = ufl.conditional(ufl.gt(x[0], 0.5), x[1], 2 * x[0]) L = ufl.inner(f, v) * ufl.dx if not dolfinx.has_petsc4py: pytest.skip("dolfinx not configured with PETSc4py") uh = dolfinx.fem.Function(V) if "petsc_options_prefix" in inspect.signature(LinearProblem.__init__).parameters.keys(): extra_options = {"petsc_options_prefix": "legacy_checkpoint_test"} else: extra_options = {} problem = LinearProblem( a, L, bcs=[], u=uh, petsc_options={"ksp_type": "preonly", "pc_type": "lu"}, **extra_options ) problem.solve() u_in = dolfinx.fem.Function(V) read_function_from_legacy_h5(path, mesh.comm, u_in, group="v", step=0, backend=backend) assert np.allclose(uh.x.array, u_in.x.array) # Check second step uh.interpolate(lambda x: x[0]) read_function_from_legacy_h5(path, mesh.comm, u_in, group="v", step=1, backend=backend) assert np.allclose(uh.x.array, u_in.x.array) W = dolfinx.fem.functionspace(mesh, ("DG", 2, (mesh.geometry.dim,))) wh = dolfinx.fem.Function(W) wh.interpolate(lambda x: (x[0], 3 * x[2], 7 * x[1])) w_in = dolfinx.fem.Function(W) read_function_from_legacy_h5(path, mesh.comm, w_in, group="w", step=0, backend=backend) np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) wh.interpolate(lambda x: np.vstack((x[0], 0 * x[0], x[1]))) read_function_from_legacy_h5(path, mesh.comm, w_in, group="w", step=1, backend=backend) np.testing.assert_allclose(wh.x.array, w_in.x.array, atol=1e-14) if backend == "h5py": qh = read_cell_data( path.with_suffix(".xdmf"), "vDG_checkpoint", mesh, 0.5, backend_args={}, backend="xdmf" ) cell_map = mesh.topology.index_map(mesh.topology.dim) num_cells_local = cell_map.size_local + cell_map.num_ghosts midpoints = dolfinx.mesh.compute_midpoints( mesh, mesh.topology.dim, np.arange(num_cells_local, dtype=np.int32) ) ref_val = midpoints[:, 0] + midpoints[:, 1] np.testing.assert_allclose(qh.x.array, ref_val) def test_io4dolfinx_legacy(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy_checkpoint") / "io4dolfinx_checkpoint.bp").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") el = ("N1curl", 3) backend_args = {"engine": "BP4", "legacy": True} mesh = read_mesh( path, comm, dolfinx.mesh.GhostMode.shared_facet, backend_args=backend_args, backend="adios2" ) def f(x): values = np.zeros((2, x.shape[1]), dtype=np.float64) values[0] = x[0] values[1] = -x[1] return values V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V) read_function(path, u, backend_args=backend_args, backend="adios2") u_ex = dolfinx.fem.Function(V) u_ex.interpolate(f) np.testing.assert_allclose(u.x.array, u_ex.x.array, atol=1e-14) def test_legacy_vtu_mesh(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh_P1000000.vtu").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") pytest.importorskip("pyvista") mesh = read_mesh(path, comm, backend="pyvista") num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global assert num_cells_global == 12 * 13 * 2 area = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=mesh))), op=MPI.SUM ) assert np.isclose(area, 1.9 * 2.8) def test_legacy_pvd(): comm = MPI.COMM_WORLD path = (pathlib.Path("legacy") / "mesh_P1_func.pvd").absolute() if not path.exists(): pytest.skip(f"{path} does not exist") pytest.importorskip("pyvista") mesh = read_mesh(path, comm, backend="pyvista") num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global assert num_cells_global == 12 * 13 * 2 area = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=mesh))), op=MPI.SUM ) assert np.isclose(area, 1.9 * 2.8) u = read_point_data(path, "w", mesh, backend="pyvista") names = read_function_names(path, mesh.comm, {}, backend="pyvista") assert len(names) == 1 assert names[0] == "w" u_ref = dolfinx.fem.Function(u.function_space) u_ref.interpolate(lambda x: (x[0], x[1] - x[0], np.zeros_like(x[0]))) np.testing.assert_allclose(u.x.array, u_ref.x.array) scientificcomputing-io4dolfinx-d21fc0e/tests/test_mesh.py000066400000000000000000000036211517634040500240410ustar00rootroot00000000000000from mpi4py import MPI import dolfinx import numpy as np import pytest import ufl from io4dolfinx import reconstruct_mesh @pytest.mark.parametrize("dtype", [np.float32, np.float64]) @pytest.mark.parametrize("degree", [2, 3]) @pytest.mark.parametrize("R", [0.1, 1, 10]) def test_curve_mesh(degree, dtype, R): N = 8 mesh = dolfinx.mesh.create_rectangle( MPI.COMM_WORLD, [[-1, -1], [1, 1]], [N, N], diagonal=dolfinx.mesh.DiagonalType.crossed, dtype=dtype, ) org_area = dolfinx.fem.form(1 * ufl.dx(domain=mesh), dtype=dtype) curved_mesh = reconstruct_mesh(mesh, degree) def transform(x): u = R * x[0] * np.sqrt(1 - (x[1] ** 2 / (2))) v = R * x[1] * np.sqrt(1 - (x[0] ** 2 / (2))) return np.asarray([u, v]) curved_mesh.geometry.x[:, : curved_mesh.geometry.dim] = transform(curved_mesh.geometry.x.T).T area = dolfinx.fem.form(1 * ufl.dx(domain=curved_mesh), dtype=dtype) circumference = dolfinx.fem.form(1 * ufl.ds(domain=curved_mesh), dtype=dtype) computed_area = curved_mesh.comm.allreduce(dolfinx.fem.assemble_scalar(area), op=MPI.SUM) computed_circumference = curved_mesh.comm.allreduce( dolfinx.fem.assemble_scalar(circumference), op=MPI.SUM ) tol = 10 * np.finfo(dtype).eps assert np.isclose(computed_area, np.pi * R**2, atol=tol) assert np.isclose(computed_circumference, 2 * np.pi * R, atol=tol) linear_mesh = reconstruct_mesh(curved_mesh, 1) linear_area = dolfinx.fem.form(1 * ufl.dx(domain=linear_mesh), dtype=dtype) recovered_area = linear_mesh.comm.allreduce( dolfinx.fem.assemble_scalar(linear_area), op=MPI.SUM ) # Curve original mesh mesh.geometry.x[:, : mesh.geometry.dim] = transform(mesh.geometry.x.T).T ref_area = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(org_area), op=MPI.SUM) assert np.isclose(recovered_area, ref_area, atol=tol) scientificcomputing-io4dolfinx-d21fc0e/tests/test_mesh_writer.py000066400000000000000000000226211517634040500254360ustar00rootroot00000000000000import importlib from mpi4py import MPI import dolfinx import numpy as np import pytest import ufl from io4dolfinx import FileMode, read_mesh, write_mesh def get_hdf5_version(): """Get the HDF5 library version found on the system. Note: This function first tries to get the version via h5py. If h5py is not installed, it attempts to load the HDF5 shared library directly using ctypes to query the version. """ import ctypes from ctypes import byref, c_uint from packaging.version import parse as _v try: import h5py return _v(h5py.version.hdf5_version) except ImportError: try: # Try to load default HDF5 library names # If adios2 is already imported, the symbols might be globally available try: libhdf5 = ctypes.CDLL("libhdf5.so") except OSError: try: libhdf5 = ctypes.CDLL("libhdf5.dylib") # MacOS except OSError: print("\n[Method 2] Could not load libhdf5 shared library directly.") libhdf5 = None if libhdf5: major = c_uint() minor = c_uint() rel = c_uint() # Call the C function libhdf5.H5get_libversion(byref(major), byref(minor), byref(rel)) return _v(f"{major.value}.{minor.value}.{rel.value}") except Exception: pass except Exception: pass raise RuntimeError("Failed to get HDF5 version") @pytest.mark.parametrize( "backend, encoder, suffix", [ ("adios2", "BP4", ".bp"), ("adios2", "HDF5", ".h5"), ("adios2", "BP5", ".bp"), ("h5py", "HDF5", ".h5"), ], ) @pytest.mark.parametrize( "ghost_mode", [dolfinx.mesh.GhostMode.shared_facet, dolfinx.mesh.GhostMode.none] ) @pytest.mark.parametrize("store_partition", [True, False]) def test_mesh_read_writer(backend, encoder, suffix, ghost_mode, tmp_path, store_partition): N = 7 if backend == "adios2" and encoder == "HDF5" and get_hdf5_version().major >= 2: pytest.skip("HDF5 version >= 2 is not supported due to ADIOS2 limitations.") try: importlib.import_module(backend) except ModuleNotFoundError: pytest.skip(f"{backend} not installed") # Consistent tmp dir across processes fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / f"{backend}_mesh_{encoder}_{store_partition}" xdmf_file = fname / f"xdmf_mesh_{encoder}_{ghost_mode}_{store_partition}" mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N, ghost_mode=ghost_mode) backend_args = None if backend == "adios2": backend_args = {"engine": encoder} write_mesh( file.with_suffix(suffix), mesh, store_partition_info=store_partition, backend_args=backend_args, backend=backend, ) mesh.comm.Barrier() with dolfinx.io.XDMFFile(mesh.comm, xdmf_file.with_suffix(".xdmf"), "w") as xdmf: xdmf.write_mesh(mesh) mesh.comm.Barrier() in_mesh = read_mesh( file.with_suffix(suffix), MPI.COMM_WORLD, ghost_mode=ghost_mode, read_from_partition=store_partition, backend_args=backend_args, backend=backend, ) in_mesh.comm.Barrier() if store_partition: def compute_distance_matrix(points_A, points_B, tol=1e-12): points_A_e = np.expand_dims(points_A, 1) points_B_e = np.expand_dims(points_B, 0) distances = np.sum(np.square(points_A_e - points_B_e), axis=2) return distances < tol cell_map = mesh.topology.index_map(mesh.topology.dim) new_cell_map = in_mesh.topology.index_map(in_mesh.topology.dim) assert cell_map.size_local == new_cell_map.size_local assert cell_map.num_ghosts == new_cell_map.num_ghosts mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim) midpoints = dolfinx.mesh.compute_midpoints( mesh, mesh.topology.dim, np.arange(cell_map.size_local + cell_map.num_ghosts, dtype=np.int32), ) in_mesh.topology.create_connectivity(in_mesh.topology.dim, in_mesh.topology.dim) new_midpoints = dolfinx.mesh.compute_midpoints( in_mesh, in_mesh.topology.dim, np.arange(new_cell_map.size_local + new_cell_map.num_ghosts, dtype=np.int32), ) # Check that all points in owned by initial mesh is owned by the new mesh # (might be locally reordered) owned_distances = compute_distance_matrix( midpoints[: cell_map.size_local], new_midpoints[: new_cell_map.size_local] ) np.testing.assert_allclose(np.sum(owned_distances, axis=1), 1) # Check that all points that are ghosted in original mesh is ghosted on the # same process in the new mesh ghost_distances = compute_distance_matrix( midpoints[cell_map.size_local :], new_midpoints[new_cell_map.size_local :] ) np.testing.assert_allclose(np.sum(ghost_distances, axis=1), 1) mesh.comm.Barrier() with dolfinx.io.XDMFFile(mesh.comm, xdmf_file.with_suffix(".xdmf"), "r") as xdmf: mesh_xdmf = xdmf.read_mesh(ghost_mode=ghost_mode) for i in range(mesh.topology.dim + 1): mesh.topology.create_entities(i) mesh_xdmf.topology.create_entities(i) in_mesh.topology.create_entities(i) assert ( mesh_xdmf.topology.index_map(i).size_global == in_mesh.topology.index_map(i).size_global ) # Check that integration over different entities are consistent measures = ( [ufl.ds, ufl.dx] if ghost_mode is dolfinx.mesh.GhostMode.none else [ufl.ds, ufl.dS, ufl.dx] ) for measure in measures: form = dolfinx.fem.form(1 * measure(domain=in_mesh)) c_adios = dolfinx.fem.assemble_scalar(form) c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh))) c_xdmf = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_xdmf))) assert np.isclose( in_mesh.comm.allreduce(c_adios, MPI.SUM), mesh.comm.allreduce(c_xdmf, MPI.SUM), ) assert np.isclose( in_mesh.comm.allreduce(c_adios, MPI.SUM), mesh.comm.allreduce(c_ref, MPI.SUM), ) @pytest.mark.parametrize( "backend, encoder, suffix", [("adios2", "BP4", ".bp"), ("adios2", "BP5", ".bp"), ("h5py", None, ".h5")], ) @pytest.mark.parametrize( "ghost_mode", [dolfinx.mesh.GhostMode.shared_facet, dolfinx.mesh.GhostMode.none] ) @pytest.mark.parametrize("store_partition", [True, False]) def test_timedep_mesh(encoder, backend, suffix, ghost_mode, tmp_path, store_partition): try: importlib.import_module(backend) except ModuleNotFoundError: pytest.skip(f"{backend} not installed") # Currently unsupported, unclear why ("HDF5", ".h5"), N = 13 # Consistent tmp dir across processes fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / f"adios_time_dep_mesh_{encoder}" mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N, ghost_mode=ghost_mode) def u(x): return np.asarray([x[0] + 0.1 * np.sin(x[1]), 0.2 * np.cos(x[1]), x[2]]) backend_args = None if backend == "adios2": backend_args = {"engine": encoder} write_mesh( file.with_suffix(suffix), mesh, mode=FileMode.write, time=0.0, store_partition_info=store_partition, backend_args=backend_args, backend=backend, ) delta_x = u(mesh.geometry.x.T).T mesh.geometry.x[:] += delta_x write_mesh( file.with_suffix(suffix), mesh, mode=FileMode.append, time=3.0, backend_args=backend_args, backend=backend, ) mesh.geometry.x[:] -= delta_x mesh_first = read_mesh( file.with_suffix(suffix), MPI.COMM_WORLD, ghost_mode, time=0.0, read_from_partition=store_partition, backend_args=backend_args, backend=backend, ) mesh_first.comm.Barrier() # Check that integration over different entities are consistent measures = [ufl.ds, ufl.dx] if ghost_mode == dolfinx.mesh.GhostMode.shared_facet: measures.append(ufl.dx) for measure in measures: # try: c_adios = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_first))) c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh))) assert np.isclose( mesh_first.comm.allreduce(c_adios, MPI.SUM), mesh.comm.allreduce(c_ref, MPI.SUM), ) mesh.geometry.x[:] += delta_x mesh_second = read_mesh( file.with_suffix(suffix), MPI.COMM_WORLD, ghost_mode, time=3.0, read_from_partition=store_partition, backend_args=backend_args, backend=backend, ) mesh_second.comm.Barrier() measures = [ufl.ds, ufl.dx] if ghost_mode == dolfinx.mesh.GhostMode.shared_facet: measures.append(ufl.dx) for measure in measures: c_adios = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_second))) c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh))) assert np.isclose( mesh_second.comm.allreduce(c_adios, MPI.SUM), mesh.comm.allreduce(c_ref, MPI.SUM), ) scientificcomputing-io4dolfinx-d21fc0e/tests/test_meshtags.py000066400000000000000000000252721517634040500247260ustar00rootroot00000000000000from __future__ import annotations import itertools from mpi4py import MPI import dolfinx import numpy as np import pytest import io4dolfinx root = 0 dtypes: list["str"] = ["float64", "float32"] # Mesh geometry dtypes write_comm: list[MPI.Intracomm] = [ MPI.COMM_SELF, MPI.COMM_WORLD, ] # Communicators for creating mesh read_modes: list[dolfinx.mesh.GhostMode] = [ dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet, ] # Cell types of different dimensions two_dimensional_cell_types: list[dolfinx.mesh.CellType] = [ dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral, ] three_dimensional_cell_types: list[dolfinx.mesh.CellType] = [ dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron, ] one_dim_combinations = itertools.product(dtypes, write_comm) two_dim_combinations = itertools.product(dtypes, two_dimensional_cell_types, write_comm) three_dim_combinations = itertools.product(dtypes, three_dimensional_cell_types, write_comm) @pytest.fixture(params=one_dim_combinations, scope="module") def mesh_1D(request): dtype, write_comm = request.param mesh = dolfinx.mesh.create_unit_interval(write_comm, 8, dtype=np.dtype(dtype)) return mesh @pytest.fixture(params=two_dim_combinations, scope="module") def mesh_2D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_square( write_comm, 10, 7, cell_type=cell_type, dtype=np.dtype(dtype) ) return mesh @pytest.fixture(params=three_dim_combinations, scope="module") def mesh_3D(request): dtype, cell_type, write_comm = request.param mesh = dolfinx.mesh.create_unit_cube( write_comm, 5, 7, 3, cell_type=cell_type, dtype=np.dtype(dtype) ) return mesh @pytest.mark.parametrize("read_mode", read_modes) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_checkpointing_meshtags_1D( mesh_1D, read_comm, read_mode, tmp_path, backend, generate_reference_map ): mesh = mesh_1D # Write unique mesh file for each combination of MPI communicator and dtype hash = f"{mesh.comm.size}_{mesh.geometry.x.dtype}" fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) suffix = ".bp" if backend == "adios" else ".h5" filename = (fname / f"{backend}_meshtags_1D_{hash}").with_suffix(suffix) # If mesh communicator is more than a self communicator or serial write on all processes. # If serial or self communicator, only write on root rank if mesh.comm.size != 1: io4dolfinx.write_mesh(filename, mesh, backend=backend) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_mesh(filename, mesh, backend=backend) # Create meshtags labeling each entity (of each co-dimension) with a # unique number (their initial global index). org_maps = [] for dim in range(mesh.topology.dim + 1): mesh.topology.create_connectivity(dim, mesh.topology.dim) e_map = mesh.topology.index_map(dim) num_entities_local = e_map.size_local entities = np.arange(num_entities_local, dtype=np.int32) ft = dolfinx.mesh.meshtags(mesh, dim, entities, e_map.local_range[0] + entities) ft.name = f"entity_{dim}" # If parallel write on all processes, else write on root rank if mesh.comm.size != 1: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) # Create map from mesh tag value to its corresponding index and midpoint org_map = generate_reference_map(mesh, ft, mesh.comm, root) org_maps.append(org_map) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) # Create map from mesh tag value to its corresponding index and midpoint org_map = generate_reference_map(mesh, ft, MPI.COMM_SELF, root) org_maps.append(org_map) del ft del mesh MPI.COMM_WORLD.Barrier() # Read mesh on testing communicator new_mesh = io4dolfinx.read_mesh(filename, read_comm, ghost_mode=read_mode, backend=backend) for dim in range(new_mesh.topology.dim + 1): # Read meshtags on all processes if testing communicator has multiple ranks # else read on root 0 if read_comm.size != 1: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) # Generate meshtags map from mesh tag value to its corresponding index and midpoint # and gather on root process read_map = generate_reference_map(new_mesh, new_ft, new_mesh.comm, root) else: if MPI.COMM_WORLD.rank == root: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) read_map = generate_reference_map(new_mesh, new_ft, read_comm, root) # On root process, check that midpoints are the same for each value in the meshtag if MPI.COMM_WORLD.rank == root: org_map = org_maps[dim] assert len(org_map) == len(read_map) for value, (_, midpoint) in org_map.items(): _, read_midpoint = read_map[value] np.testing.assert_allclose(read_midpoint, midpoint) @pytest.mark.parametrize("read_mode", read_modes) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_checkpointing_meshtags_2D( mesh_2D, read_comm, read_mode, tmp_path, backend, generate_reference_map ): mesh = mesh_2D hash = f"{mesh.comm.size}_{mesh.topology.cell_name()}_{mesh.geometry.x.dtype}" fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = fname / f"meshtags_1D_{hash}.bp" if mesh.comm.size != 1: io4dolfinx.write_mesh(filename, mesh, backend=backend) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_mesh(filename, mesh, backend=backend) org_maps = [] for dim in range(mesh.topology.dim + 1): mesh.topology.create_connectivity(dim, mesh.topology.dim) e_map = mesh.topology.index_map(dim) num_entities_local = e_map.size_local entities = np.arange(num_entities_local, dtype=np.int32) ft = dolfinx.mesh.meshtags(mesh, dim, entities, e_map.local_range[0] + entities) ft.name = f"entity_{dim}" if mesh.comm.size != 1: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) org_map = generate_reference_map(mesh, ft, mesh.comm, root) org_maps.append(org_map) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) org_map = generate_reference_map(mesh, ft, MPI.COMM_SELF, root) org_maps.append(org_map) del ft del mesh MPI.COMM_WORLD.Barrier() new_mesh = io4dolfinx.read_mesh(filename, read_comm, ghost_mode=read_mode, backend=backend) for dim in range(new_mesh.topology.dim + 1): if read_comm.size != 1: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) read_map = generate_reference_map(new_mesh, new_ft, new_mesh.comm, root) else: if MPI.COMM_WORLD.rank == root: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) read_map = generate_reference_map(new_mesh, new_ft, read_comm, root) if MPI.COMM_WORLD.rank == root: org_map = org_maps[dim] assert len(org_map) == len(read_map) for value, (_, midpoint) in org_map.items(): _, read_midpoint = read_map[value] np.testing.assert_allclose(read_midpoint, midpoint) @pytest.mark.parametrize("read_mode", read_modes) @pytest.mark.parametrize("read_comm", [MPI.COMM_SELF, MPI.COMM_WORLD]) def test_checkpointing_meshtags_3D( mesh_3D, read_comm, read_mode, tmp_path, backend, generate_reference_map ): mesh = mesh_3D hash = f"{mesh.comm.size}_{mesh.topology.cell_name()}_{mesh.geometry.x.dtype}" fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = fname / f"meshtags_1D_{hash}.bp" if mesh.comm.size != 1: io4dolfinx.write_mesh(filename, mesh, backend=backend) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_mesh(filename, mesh, backend=backend) org_maps = [] for dim in range(mesh.topology.dim + 1): mesh.topology.create_connectivity(dim, mesh.topology.dim) e_map = mesh.topology.index_map(dim) num_entities_local = e_map.size_local entities = np.arange(num_entities_local, dtype=np.int32) ft = dolfinx.mesh.meshtags(mesh, dim, entities, e_map.local_range[0] + entities) ft.name = f"entity_{dim}" if mesh.comm.size != 1: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) org_map = generate_reference_map(mesh, ft, mesh.comm, root) org_maps.append(org_map) else: if MPI.COMM_WORLD.rank == root: io4dolfinx.write_meshtags(filename, mesh, ft, backend=backend) org_map = generate_reference_map(mesh, ft, MPI.COMM_SELF, root) org_maps.append(org_map) del ft del mesh MPI.COMM_WORLD.Barrier() new_mesh = io4dolfinx.read_mesh(filename, read_comm, ghost_mode=read_mode, backend=backend) for dim in range(new_mesh.topology.dim + 1): if read_comm.size != 1: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) read_map = generate_reference_map(new_mesh, new_ft, new_mesh.comm, root) else: if MPI.COMM_WORLD.rank == root: new_ft = io4dolfinx.read_meshtags( filename, new_mesh, meshtag_name=f"entity_{dim}", backend=backend, ) read_map = generate_reference_map(new_mesh, new_ft, MPI.COMM_SELF, root) if MPI.COMM_WORLD.rank == root: org_map = org_maps[dim] assert len(org_map) == len(read_map) for value, (_, midpoint) in org_map.items(): _, read_midpoint = read_map[value] np.testing.assert_allclose(read_midpoint, midpoint) scientificcomputing-io4dolfinx-d21fc0e/tests/test_numpy_vectorization.py000066400000000000000000000132321517634040500272340ustar00rootroot00000000000000import itertools from typing import Tuple from mpi4py import MPI import basix.ufl import dolfinx import numpy as np import numpy.typing as npt import pytest from io4dolfinx.utils import compute_dofmap_pos, unroll_dofmap write_comm = [MPI.COMM_SELF, MPI.COMM_WORLD] # Communicators for creating mesh ghost_mode = [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet] two_dimensional_cell_types = [ dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral, ] three_dimensional_cell_types = [dolfinx.mesh.CellType.hexahedron] two_dim_combinations = itertools.product(two_dimensional_cell_types, write_comm, ghost_mode) three_dim_combinations = itertools.product(three_dimensional_cell_types, write_comm, ghost_mode) @pytest.fixture(params=two_dim_combinations, scope="module") def mesh_2D(request): cell_type, write_comm, ghost_mode = request.param mesh = dolfinx.mesh.create_unit_square( write_comm, 10, 10, cell_type=cell_type, ghost_mode=ghost_mode ) return mesh @pytest.fixture(params=three_dim_combinations, scope="module") def mesh_3D(request): cell_type, write_comm, ghost_mode = request.param M = 5 mesh = dolfinx.mesh.create_unit_cube( write_comm, M, M, M, cell_type=cell_type, ghost_mode=ghost_mode ) return mesh def compute_positions( dofs: npt.NDArray[np.int32], dofmap_bs: int, num_owned_dofs: int, num_owned_cells: int, ) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: """ Support function for test. Given a dofmap, compute the local cell and position in the dofmap for each owned dof. The last cell (wrt) local index will be the one in the output map """ dof_to_cell_map = np.zeros(num_owned_dofs, dtype=np.int32) dof_to_pos_map = np.zeros(num_owned_dofs, dtype=np.int32) for c in range(num_owned_cells): for i, dof in enumerate(dofs[c]): for b in range(dofmap_bs): local_dof = dof * dofmap_bs + b if local_dof < num_owned_dofs: dof_to_cell_map[local_dof] = c dof_to_pos_map[local_dof] = i * dofmap_bs + b return dof_to_cell_map, dof_to_pos_map @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) def test_unroll_P(family, degree, mesh_2D): V = dolfinx.fem.functionspace(mesh_2D, (family, degree)) dofmap = V.dofmap unrolled_map = unroll_dofmap(dofmap.list, dofmap.bs) normal_unroll = np.zeros( (dofmap.list.shape[0], dofmap.list.shape[1] * dofmap.bs), dtype=np.int32 ) for i, dofs in enumerate(dofmap.list): for j, dof in enumerate(dofs): for k in range(dofmap.bs): normal_unroll[i, j * dofmap.bs + k] = dof * dofmap.bs + k np.testing.assert_allclose(unrolled_map, normal_unroll) @pytest.mark.parametrize("family", ["RTCF"]) @pytest.mark.parametrize("degree", [1, 2, 3]) def test_unroll_RTCF(family, degree, mesh_3D): el = basix.ufl.element(family, mesh_3D.basix_cell(), degree) V = dolfinx.fem.functionspace(mesh_3D, el) dofmap = V.dofmap unrolled_map = unroll_dofmap(dofmap.list, dofmap.bs) normal_unroll = np.zeros( (dofmap.list.shape[0], dofmap.list.shape[1] * dofmap.bs), dtype=np.int32 ) for i, dofs in enumerate(dofmap.list): for j, dof in enumerate(dofs): for k in range(dofmap.bs): normal_unroll[i, j * dofmap.bs + k] = dof * dofmap.bs + k np.testing.assert_allclose(unrolled_map, normal_unroll) @pytest.mark.parametrize("family", ["RTCF"]) @pytest.mark.parametrize("degree", [1, 2, 3]) def test_compute_dofmap_pos_RTCF(family, degree, mesh_3D): el = basix.ufl.element(family, mesh_3D.basix_cell(), degree) V = dolfinx.fem.functionspace(mesh_3D, el) local_cells, local_pos = compute_dofmap_pos(V) num_cells_local = mesh_3D.topology.index_map(mesh_3D.topology.dim).size_local num_dofs_local = V.dofmap.index_map.size_local * V.dofmap.index_map_bs reference_cells, reference_pos = compute_positions( V.dofmap.list, V.dofmap.bs, num_dofs_local, num_cells_local ) np.testing.assert_allclose(reference_cells, local_cells) np.testing.assert_allclose(reference_pos, local_pos) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) def test_compute_dofmap_pos_P(family, degree, mesh_2D): el = basix.ufl.element(family, mesh_2D.basix_cell(), degree) V = dolfinx.fem.functionspace(mesh_2D, el) local_cells, local_pos = compute_dofmap_pos(V) num_cells_local = mesh_2D.topology.index_map(mesh_2D.topology.dim).size_local num_dofs_local = V.dofmap.index_map.size_local * V.dofmap.index_map_bs reference_cells, reference_pos = compute_positions( V.dofmap.list, V.dofmap.bs, num_dofs_local, num_cells_local ) np.testing.assert_allclose(reference_cells, local_cells) np.testing.assert_allclose(reference_pos, local_pos) def test_compute_send_sizes(): np.random.seed(42) N = 0 M = 10 num_data = 100 # Set of ranks to recieve data dest_ranks = np.arange(N, M, dtype=np.int32) # Random data owners data_owners = np.random.randint(N, M, num_data).astype(np.int32) # Compute the number of data to send to each rank with loops out_size = np.zeros(len(dest_ranks), dtype=np.int32) for owner in data_owners: for j, rank in enumerate(dest_ranks): if owner == rank: out_size[j] += 1 break process_pos_indicator = data_owners.reshape(-1, 1) == dest_ranks vectorized_out_size = np.count_nonzero(process_pos_indicator, axis=0) np.testing.assert_allclose(vectorized_out_size, out_size) scientificcomputing-io4dolfinx-d21fc0e/tests/test_original_checkpoint.py000066400000000000000000000436401517634040500271250ustar00rootroot00000000000000from __future__ import annotations import itertools import os import typing from collections.abc import Callable from pathlib import Path from mpi4py import MPI import basix import basix.ufl import dolfinx import numpy as np import pytest import io4dolfinx dtypes = [np.float64, np.float32] # Mesh geometry dtypes two_dimensional_cell_types = [ dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral, ] three_dimensional_cell_types = [ dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron, ] two_dim_combinations = itertools.product(dtypes, two_dimensional_cell_types) three_dim_combinations = itertools.product(dtypes, three_dimensional_cell_types) @pytest.fixture(scope="module") def create_simplex_mesh_2D(tmp_path_factory): mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.triangle, dtype=np.float64, ) fname = tmp_path_factory.mktemp("output") / "original_mesh_2D_simplex.xdmf" fname = MPI.COMM_WORLD.bcast(fname, root=0) with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname @pytest.fixture(scope="module") def create_simplex_mesh_3D(tmp_path_factory): mesh = dolfinx.mesh.create_unit_cube( MPI.COMM_WORLD, 5, 5, 5, cell_type=dolfinx.mesh.CellType.tetrahedron, dtype=np.float64, ) fname = tmp_path_factory.mktemp("output") / "original_mesh_3D_simplex.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname @pytest.fixture(scope="module") def create_non_simplex_mesh_2D(tmp_path_factory): mesh = dolfinx.mesh.create_unit_square( MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.quadrilateral, dtype=np.float64, ) fname = tmp_path_factory.mktemp("output") / "original_mesh_2D_non_simplex.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname @pytest.fixture(scope="module") def create_non_simplex_mesh_3D(tmp_path_factory): mesh = dolfinx.mesh.create_unit_cube( MPI.COMM_WORLD, 5, 5, 5, cell_type=dolfinx.mesh.CellType.hexahedron, dtype=np.float64, ) fname = tmp_path_factory.mktemp("output") / "original_mesh_3D_non_simplex.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname @pytest.fixture(params=two_dim_combinations, scope="module") def create_2D_mesh(request, tmpdir_factory): dtype, cell_type = request.param mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 7, cell_type=cell_type, dtype=dtype) fname = Path(tmpdir_factory.mktemp("output")) / f"original_mesh_2D_{dtype}_{cell_type}.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname @pytest.fixture(params=three_dim_combinations, scope="module") def create_3D_mesh(request, tmpdir_factory): dtype, cell_type = request.param mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 5, 7, 3, cell_type=cell_type, dtype=dtype) fname = Path(tmpdir_factory.mktemp("output")) / f"original_mesh_3D_{dtype}_{cell_type}.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "w") as xdmf: xdmf.write_mesh(mesh) return fname def write_function_original( write_mesh: bool, mesh: dolfinx.mesh.Mesh, el: basix.ufl._ElementBase, f: Callable[[np.ndarray], np.ndarray], dtype: np.dtype, name: str, path: Path, backend: typing.Literal["adios2", "h5py"], ) -> Path: """Convenience function for writing function to file on the original input mesh""" V = dolfinx.fem.functionspace(mesh, el) uh = dolfinx.fem.Function(V, dtype=dtype) uh.interpolate(f) uh.name = name el_hash = ( io4dolfinx.utils.element_signature(V) .replace(" ", "") .replace(",", "") .replace("(", "") .replace(")", "") .replace("[", "") .replace("]", "") ) file_hash = f"{el_hash}_{np.dtype(dtype).name}" if backend == "adios2": suffix = ".bp" elif backend == "h5py": suffix = ".h5" else: raise NotImplementedError(f"Unknown backend {backend}") filename = (path / f"mesh_{file_hash}").with_suffix(suffix) if write_mesh: io4dolfinx.write_mesh_input_order(filename, mesh, backend=backend) io4dolfinx.write_function_on_input_mesh(filename, uh, time=0.0, backend=backend) return filename def read_function_original( mesh_fname: Path, u_fname: Path, u_name: str, family: str, degree: int, f: Callable[[np.ndarray], np.ndarray], u_dtype: np.dtype, backend: typing.Literal["adios2", "h5py"], ): """ Convenience function for reading mesh with IPython-parallel and compare to exact solution """ from mpi4py import MPI import dolfinx import io4dolfinx # assert MPI.COMM_WORLD.size > 1 if backend == "adios2": backend_args = {"engine": "BP4"} else: backend_args = None if mesh_fname.suffix == ".xdmf": with dolfinx.io.XDMFFile(MPI.COMM_WORLD, mesh_fname, "r") as xdmf: mesh = xdmf.read_mesh() else: mesh = io4dolfinx.read_mesh( mesh_fname, MPI.COMM_WORLD, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, backend_args=backend_args, backend=backend, ) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V, name=u_name, dtype=u_dtype) io4dolfinx.read_function(u_fname, u, time=0.0, backend_args=backend_args, backend=backend) MPI.COMM_WORLD.Barrier() u_ex = dolfinx.fem.Function(V, name="exact", dtype=u_dtype) u_ex.interpolate(f) u_ex.x.scatter_forward() atol = 10 * np.finfo(u_dtype).resolution np.testing.assert_allclose(u.x.array, u_ex.x.array, atol=atol) # type: ignore def write_function_vector( write_mesh: bool, fname: Path, family: str, degree: int, f: Callable[[np.ndarray], np.ndarray], dtype: np.dtype, name: str, dir: Path, backend: typing.Literal["adios2", "h5py"], ) -> Path: """Convenience function for writing function to file on the original input mesh""" from mpi4py import MPI import basix.ufl import dolfinx import io4dolfinx assert MPI.COMM_WORLD.size > 1 with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "r") as xdmf: mesh = xdmf.read_mesh() el = basix.ufl.element(family, mesh.basix_cell(), degree, dtype=mesh.geometry.x.dtype) V = dolfinx.fem.functionspace(mesh, el) uh = dolfinx.fem.Function(V, dtype=dtype) uh.interpolate(f) uh.name = name el_hash = ( io4dolfinx.utils.element_signature(V) .replace(" ", "") .replace(",", "") .replace("(", "") .replace(")", "") .replace("[", "") .replace("]", "") ) file_hash = f"{el_hash}_{np.dtype(dtype).name}" filename = dir / f"mesh_{file_hash}.bp" if write_mesh: io4dolfinx.write_mesh_input_order(filename, mesh, backend=backend) io4dolfinx.write_function_on_input_mesh(filename, uh, time=0.0, backend=backend) return filename def read_function_vector( mesh_fname: Path, u_fname: Path, u_name: str, family: str, degree: int, f: Callable[[np.ndarray], np.ndarray], u_dtype: np.dtype, backend: typing.Literal["adios2", "h5py"], ): """ Convenience function for reading mesh with IPython-parallel and compare to exact solution """ if mesh_fname.suffix == ".xdmf": with dolfinx.io.XDMFFile(MPI.COMM_WORLD, mesh_fname, "r") as xdmf: mesh = xdmf.read_mesh() elif mesh_fname.suffix == ".bp": mesh = io4dolfinx.read_mesh( mesh_fname, MPI.COMM_WORLD, ghost_mode=dolfinx.mesh.GhostMode.shared_facet, backend=backend, ) el = basix.ufl.element(family, mesh.basix_cell(), degree) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V, name=u_name, dtype=u_dtype) io4dolfinx.read_function(u_fname, u, time=0.0, backend=backend) MPI.COMM_WORLD.Barrier() u_ex = dolfinx.fem.Function(V, name="exact", dtype=u_dtype) u_ex.interpolate(f) u_ex.x.scatter_forward() atol = 10 * np.finfo(u_dtype).resolution np.testing.assert_allclose(u.x.array, u_ex.x.array, atol=atol) # type: ignore @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("write_mesh", [True, False]) def test_read_write_P_2D( write_mesh, family, degree, is_complex, create_2D_mesh, cluster, get_dtype, tmp_path, backend ): fname = create_2D_mesh with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "r") as xdmf: mesh = xdmf.read_mesh() f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), dtype=mesh.geometry.x.dtype, ) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = x[0] if is_complex: values[0] -= 3j * x[1] values[1] += 2j * x[0] return values hash = write_function_original( write_mesh, mesh, el, f, f_dtype, "u_original", tmp_path, backend ) if write_mesh: mesh_fname = hash else: mesh_fname = fname query = cluster[:].apply_async( read_function_original, mesh_fname, hash, "u_original", family, degree, f, f_dtype, backend ) query.wait() assert query.successful(), query.error @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) @pytest.mark.parametrize("write_mesh", [True, False]) def test_read_write_P_3D( write_mesh, family, degree, is_complex, create_3D_mesh, cluster, get_dtype, tmp_path, backend ): fname = create_3D_mesh with dolfinx.io.XDMFFile(MPI.COMM_WORLD, fname, "r") as xdmf: mesh = xdmf.read_mesh() f_dtype = get_dtype(mesh.geometry.x.dtype, is_complex) el = basix.ufl.element( family, mesh.basix_cell(), degree, basix.LagrangeVariant.gll_warped, shape=(mesh.geometry.dim,), ) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.pi + x[0] values[1] = x[1] + 2 * x[0] values[2] = np.cos(x[2]) if is_complex: values[0] -= np.pi * x[1] values[1] += 3j * x[2] values[2] += 2j return values hash = write_function_original( write_mesh, mesh, el, f, f_dtype, "u_original", tmp_path, backend=backend ) MPI.COMM_WORLD.Barrier() if write_mesh: mesh_fname = hash else: mesh_fname = fname query = cluster[:].apply_async( read_function_original, mesh_fname, hash, "u_original", family, degree, f, f_dtype, backend=backend, ) query.wait() assert query.successful(), query.error @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("write_mesh", [True, False]) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["N1curl", "RT"]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_2D_vector_simplex( write_mesh, family, degree, is_complex, create_simplex_mesh_2D, cluster, get_dtype, tmp_path, backend, ): fname = create_simplex_mesh_2D f_dtype = get_dtype(np.float64, is_complex) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = x[1] if is_complex: values[0] -= np.sin(x[1]) * 2j values[1] += 3j return values query = cluster[:].apply_async( write_function_vector, write_mesh, fname, family, degree, f, f_dtype, "u_original", tmp_path, backend, ) query.wait() assert query.successful(), query.error paths = query.result() file_path = paths[0] assert all([file_path == path for path in paths]) if write_mesh: mesh_fname = file_path else: mesh_fname = fname read_function_vector(mesh_fname, file_path, "u_original", family, degree, f, f_dtype, backend) @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("write_mesh", [True, False]) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["N1curl", "RT"]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_3D_vector_simplex( write_mesh, family, degree, is_complex, create_simplex_mesh_3D, cluster, get_dtype, tmp_path, backend, ): fname = create_simplex_mesh_3D f_dtype = get_dtype(np.float64, is_complex) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) values[1] = x[1] + 2 * x[0] values[2] = np.cos(x[2]) if is_complex: values[0] += 2j * x[2] values[1] += 2j * np.cos(x[2]) return values query = cluster[:].apply_async( write_function_vector, write_mesh, fname, family, degree, f, f_dtype, "u_original", tmp_path, backend, ) query.wait() assert query.successful(), query.error paths = query.result() file_path = paths[0] assert all([file_path == path for path in paths]) if write_mesh: mesh_fname = file_path else: mesh_fname = fname read_function_vector(mesh_fname, file_path, "u_original", family, degree, f, f_dtype, backend) @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("write_mesh", [True, False]) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["RTCF"]) @pytest.mark.parametrize("degree", [1, 2, 3]) def test_read_write_2D_vector_non_simplex( write_mesh, family, degree, is_complex, create_non_simplex_mesh_2D, cluster, get_dtype, tmp_path, backend, ): fname = create_non_simplex_mesh_2D f_dtype = get_dtype(np.float64, is_complex) def f(x): values = np.empty((2, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) values[1] = x[1] + 2 * x[0] if is_complex: values[0] += 2j * x[1] values[1] -= np.sin(x[0]) * 9j return values query = cluster[:].apply_async( write_function_vector, write_mesh, fname, family, degree, f, f_dtype, "u_original", tmp_path, backend, ) query.wait() assert query.successful(), query.error paths = query.result() file_path = paths[0] assert all([file_path == path for path in paths]) if write_mesh: mesh_fname = file_path else: mesh_fname = fname read_function_vector(mesh_fname, file_path, "u_original", family, degree, f, f_dtype, backend) @pytest.mark.skipif( os.cpu_count() == 1, reason="Test requires that the system has more than one process" ) @pytest.mark.skipif(MPI.COMM_WORLD.size > 1, reason="Test uses ipythonparallel for MPI") @pytest.mark.parametrize("write_mesh", [True, False]) @pytest.mark.parametrize("is_complex", [True, False]) @pytest.mark.parametrize("family", ["NCF"]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_3D_vector_non_simplex( write_mesh, family, degree, is_complex, create_non_simplex_mesh_3D, cluster, get_dtype, tmp_path, backend, ): fname = create_non_simplex_mesh_3D f_dtype = get_dtype(np.float64, is_complex) def f(x): values = np.empty((3, x.shape[1]), dtype=f_dtype) values[0] = np.full(x.shape[1], np.pi) + x[0] values[1] = np.cos(x[2]) values[2] = x[0] if is_complex: values[2] += x[0] * x[1] * 3j return values query = cluster[:].apply_async( write_function_vector, write_mesh, fname, family, degree, f, f_dtype, "u_original", tmp_path, backend, ) query.wait() assert query.successful(), query.error paths = query.result() file_path = paths[0] assert all([file_path == path for path in paths]) if write_mesh: mesh_fname = file_path else: mesh_fname = fname read_function_vector(mesh_fname, file_path, "u_original", family, degree, f, f_dtype, backend) scientificcomputing-io4dolfinx-d21fc0e/tests/test_pyvista.py000066400000000000000000000033261517634040500246060ustar00rootroot00000000000000from mpi4py import MPI import dolfinx import numpy as np import pytest import ufl import io4dolfinx pyvista = pytest.importorskip("pyvista") def test_read_mesh_and_cell_data(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) filename = tmp_path / "grid.vtu" grid = pyvista.examples.load_hexbeam() if MPI.COMM_WORLD.rank == 0: grid.save(filename) MPI.COMM_WORLD.barrier() mesh = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD, backend="pyvista") vol = dolfinx.fem.form(1 * ufl.dx(domain=mesh)) surf = dolfinx.fem.form(1 * ufl.ds(domain=mesh)) vol_ref = 5 * 1 * 1 surf_ref = 5 * 4 + 2 vol_glob = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(vol), op=MPI.SUM) surf_glob = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(surf), op=MPI.SUM) assert np.isclose(vol_glob, vol_ref) assert np.isclose(surf_glob, surf_ref) names = io4dolfinx.read_function_names(filename, MPI.COMM_WORLD, backend="pyvista") for name in grid.cell_data.keys(): assert name in names for name in grid.point_data.keys(): assert name in names for name in names: if name in grid.cell_data.keys(): cd = io4dolfinx.read_cell_data(filename, name, mesh, backend="pyvista") oci = mesh.topology.original_cell_index np.testing.assert_allclose(cd.x.array[:], grid.cell_data[name][oci]) elif name in grid.point_data.keys(): pd = io4dolfinx.read_point_data(filename, name, mesh, backend="pyvista") igi = mesh.geometry.input_global_indices np.testing.assert_allclose(pd.x.array[:], grid.point_data[name][igi]) else: raise RuntimeError(f"Could not find {name} in grid") scientificcomputing-io4dolfinx-d21fc0e/tests/test_snapshot_checkpoint.py000066400000000000000000000102371517634040500271540ustar00rootroot00000000000000from pathlib import Path from mpi4py import MPI import basix.ufl import dolfinx import numpy as np import pytest from io4dolfinx import FileMode, snapshot_checkpoint def suffix(backend: str) -> str: if backend == "adios2": return ".bp" elif backend == "h5py": return ".h5" else: raise NotImplementedError(f"Unsupported backend {backend}") triangle = dolfinx.mesh.CellType.triangle quad = dolfinx.mesh.CellType.quadrilateral tetra = dolfinx.mesh.CellType.tetrahedron hex = dolfinx.mesh.CellType.hexahedron @pytest.mark.parametrize( "cell_type, family", [(triangle, "N1curl"), (triangle, "RT"), (quad, "RTCF")] ) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_2D(family, degree, cell_type, tmp_path, backend): mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, cell_type=cell_type) el = basix.ufl.element(family, mesh.basix_cell(), degree) def f(x): return (np.full(x.shape[1], np.pi) + x[0], x[1]) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V) u.interpolate(f) fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / Path("snapshot_2D_vs").with_suffix(suffix(backend)) snapshot_checkpoint(u, file, FileMode.write, backend=backend) v = dolfinx.fem.Function(V) snapshot_checkpoint(v, file, FileMode.read, backend=backend) assert np.allclose(u.x.array, v.x.array) @pytest.mark.parametrize("cell_type, family", [(tetra, "N1curl"), (tetra, "RT"), (hex, "NCF")]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_3D(family, degree, cell_type, tmp_path, backend): mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 3, 3, 3, cell_type=cell_type) el = basix.ufl.element(family, mesh.basix_cell(), degree) def f(x): return (np.full(x.shape[1], np.pi) + x[0], x[1], x[1] * x[2]) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V) u.interpolate(f) fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / Path("snapshot_3D_vs").with_suffix(suffix(backend)) snapshot_checkpoint(u, file, FileMode.write, backend=backend) v = dolfinx.fem.Function(V) snapshot_checkpoint(v, file, FileMode.read, backend=backend) assert np.allclose(u.x.array, v.x.array) @pytest.mark.parametrize( "cell_type", [dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral] ) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_P_2D(family, degree, cell_type, tmp_path, backend): mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 5, cell_type=cell_type) el = basix.ufl.element(family, mesh.basix_cell(), degree, shape=(mesh.geometry.dim,)) def f(x): return (np.full(x.shape[1], np.pi) + x[0], x[1]) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V) u.interpolate(f) fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / Path("snapshot_2D_p").with_suffix(suffix(backend)) snapshot_checkpoint(u, file, FileMode.write, backend=backend) v = dolfinx.fem.Function(V) snapshot_checkpoint(v, file, FileMode.read, backend=backend) assert np.allclose(u.x.array, v.x.array) @pytest.mark.parametrize( "cell_type", [dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron] ) @pytest.mark.parametrize("family", ["Lagrange", "DG"]) @pytest.mark.parametrize("degree", [1, 4]) def test_read_write_P_3D(family, degree, cell_type, tmp_path, backend): mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 5, 5, 5, cell_type=cell_type) el = basix.ufl.element(family, mesh.basix_cell(), degree, shape=(mesh.geometry.dim,)) def f(x): return (np.full(x.shape[1], np.pi) + x[0], x[1] + 2 * x[0], np.cos(x[2])) V = dolfinx.fem.functionspace(mesh, el) u = dolfinx.fem.Function(V) u.interpolate(f) fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) file = fname / Path("snapshot_3D_p").with_suffix(suffix(backend)) snapshot_checkpoint(u, file, FileMode.write, backend=backend) v = dolfinx.fem.Function(V) snapshot_checkpoint(v, file, FileMode.read, backend=backend) assert np.allclose(u.x.array, v.x.array) scientificcomputing-io4dolfinx-d21fc0e/tests/test_version.py000066400000000000000000000001261517634040500245670ustar00rootroot00000000000000import io4dolfinx def test_version(): assert io4dolfinx.__version__ is not None scientificcomputing-io4dolfinx-d21fc0e/tests/test_vtkhdf.py000066400000000000000000000431311517634040500243730ustar00rootroot00000000000000from mpi4py import MPI import numpy as np import pytest import ufl from dolfinx.fem import Function, assemble_scalar, form, functionspace from dolfinx.io.vtkhdf import write_cell_data, write_mesh, write_point_data from dolfinx.mesh import CellType, compute_midpoints, create_unit_cube, create_unit_square, meshtags import io4dolfinx def f(x, t): return x[0] - 2 * x[1] + x[2] * t def g(x, t): return x[0], 2 * x[1], -x[2] * t @pytest.mark.parametrize( "cell_type", [CellType.tetrahedron, CellType.hexahedron, CellType.tetrahedron] ) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_read_write_timedep_mesh(dtype, tmp_path, cell_type): comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() # Write temporal data mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype, cell_type=cell_type) ref_vol = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.dx(domain=mesh), dtype=dtype)), op=MPI.SUM ) ref_surf = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.ds(domain=mesh), dtype=dtype)), op=MPI.SUM ) # Write temporal data filename = tmp_path / f"timedep_mesh_{cell_type.name}_{np.dtype(dtype).name}.vtkhdf" io4dolfinx.write_mesh(filename, mesh, time=0.3, backend="vtkhdf") mesh.geometry.x[:, 0] += 0.05 * mesh.geometry.x[:, 0] mesh.geometry.x[:, 1] *= 1.1 + np.sin(mesh.geometry.x[:, 0]) ref_pert_vol = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.dx(domain=mesh), dtype=dtype)), op=MPI.SUM ) ref_pert_surf = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.ds(domain=mesh), dtype=dtype)), op=MPI.SUM ) io4dolfinx.write_mesh( filename, mesh, time=0.5, backend="vtkhdf", mode=io4dolfinx.FileMode.append, ) in_mesh = io4dolfinx.read_mesh(filename, comm, time=0.5, backend="vtkhdf") pert_vol = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.dx(domain=in_mesh), dtype=dtype)), op=MPI.SUM ) pert_surf = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.ds(domain=in_mesh), dtype=dtype)), op=MPI.SUM ) assert np.isclose(pert_vol, ref_pert_vol) assert np.isclose(pert_surf, ref_pert_surf) in_mesh = io4dolfinx.read_mesh(filename, comm, time=0.3, backend="vtkhdf") vol = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.dx(domain=in_mesh), dtype=dtype)), op=MPI.SUM ) surf = mesh.comm.allreduce( assemble_scalar(form(1 * ufl.ds(domain=in_mesh), dtype=dtype)), op=MPI.SUM ) assert np.isclose(vol, ref_vol) assert np.isclose(surf, ref_surf) @pytest.mark.parametrize("cell_type", [CellType.hexahedron, CellType.tetrahedron]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_write_point_data(dtype, tmp_path, cell_type): comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() # Write temporal data mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype, cell_type=cell_type) filename = tmp_path / f"point_data_{cell_type.name}_{np.dtype(dtype).name}.vtkhdf" write_mesh(str(filename), mesh) t = np.linspace(0.1, 1.2, 25) num_nodes_local = mesh.geometry.index_map().size_local for ti in t: point_data = f(mesh.geometry.x.T[:, :num_nodes_local], ti) write_point_data(str(filename), mesh, point_data, float(ti)) comm.barrier() grid = io4dolfinx.read_mesh(filename=filename, comm=comm, time=None, backend="vtkhdf") # Since we shuffle time we need to shuffle in the same way on each process np.random.shuffle(t) t = comm.bcast(t, root=0) for tj in t: u = io4dolfinx.read_point_data( filename=filename, name="u", mesh=grid, time=tj, backend="vtkhdf" ) v_ref = Function(u.function_space, dtype=u.x.array.dtype) atol = 15 * np.finfo(u.x.array.dtype).eps v_ref.interpolate(lambda x: f(x, tj)) np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol) # Test blocked data as well (with shuffled input timestep) blocked_file = filename.with_stem(filename.stem + "_blocked") write_mesh(str(blocked_file), mesh) for tj in t: point_data = np.asarray(g(mesh.geometry.x.T[:, :num_nodes_local], tj)).T.flatten() write_point_data(str(blocked_file), mesh, point_data, float(tj)) comm.barrier() np.random.shuffle(t) t = comm.bcast(t, root=0) for tk in t: u = io4dolfinx.read_point_data( filename=blocked_file, name="u", mesh=grid, time=tk, backend="vtkhdf" ) v_ref = Function(u.function_space, dtype=u.x.array.dtype) atol = 15 * np.finfo(u.x.array.dtype).eps v_ref.interpolate(lambda x: g(x, tk)) np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol) @pytest.mark.parametrize("cell_type", [CellType.hexahedron, CellType.tetrahedron]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_write_cell_data(dtype, tmp_path, cell_type): comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() # Write temporal data mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype, cell_type=cell_type) filename = tmp_path / f"cell_data_{cell_type.name}_{np.dtype(dtype).name}.vtkhdf" write_mesh(str(filename), mesh) t = np.linspace(0.1, 1.2, 25) num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local midpoints = compute_midpoints( mesh, mesh.topology.dim, np.arange(num_cells_local, dtype=np.int32) ) for ti in t: cell_data = f(midpoints.T, ti) write_cell_data(str(filename), mesh, cell_data, float(ti)) comm.barrier() grid = io4dolfinx.read_mesh(filename=filename, comm=comm, time=None, backend="vtkhdf") # Since we shuffle time we need to shuffle in the same way on each process np.random.shuffle(t) t = comm.bcast(t, root=0) for tj in t: u = io4dolfinx.read_cell_data( filename=filename, name="u", mesh=grid, time=tj, backend="vtkhdf" ) v_ref = Function(u.function_space, dtype=u.x.array.dtype) atol = 15 * np.finfo(u.x.array.dtype).eps v_ref.interpolate(lambda x: f(x, tj)) np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol) # Test blocked data as well (with shuffled input timestep) blocked_file = filename.with_stem(filename.stem + "_blocked") write_mesh(str(blocked_file), mesh) for tj in t: cell_data = np.asarray(g(midpoints.T, tj)).T.flatten() write_cell_data(str(blocked_file), mesh, cell_data, float(tj)) comm.barrier() np.random.shuffle(t) t = comm.bcast(t, root=0) for tk in t: u = io4dolfinx.read_cell_data( filename=blocked_file, name="u", mesh=grid, time=tk, backend="vtkhdf" ) v_ref = Function(u.function_space, dtype=u.x.array.dtype) atol = 15 * np.finfo(u.x.array.dtype).eps v_ref.interpolate(lambda x: g(x, tk)) np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_write_meshtags(dtype, tmp_path, generate_reference_map): comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() filename = tmp_path / f"meshtags_{np.dtype(dtype).name}.vtkhdf" mesh = create_unit_cube(comm, 3, 3, 3, dtype=dtype, cell_type=CellType.hexahedron) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.write, time=1.0, backend_args={"name": "hex"}, backend="vtkhdf", ) dim = mesh.topology.dim mesh.topology.create_connectivity(dim, mesh.topology.dim) cmap = mesh.topology.index_map(dim) cells = np.arange(cmap.size_local, dtype=np.int32) ct = meshtags(mesh, dim, cells, cells + cmap.local_range[0]) io4dolfinx.write_meshtags( filename, mesh, ct, "CellTags", backend_args={"name": "hex"}, backend="vtkhdf" ) root = 0 org_map = generate_reference_map(mesh, ct, comm, root) # Move mesh mesh.geometry.x[:, 0] *= 1 + 0.2 * np.sin(mesh.geometry.x[:, 1]) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.append, time=2.5, backend_args={"name": "hex"}, backend="vtkhdf", ) # Add stationary meshtags (after time loop) mesh = create_unit_cube(comm, 7, 3, 5, dtype=dtype, cell_type=CellType.tetrahedron) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.append, time=1.0, backend_args={"name": "tet"}, backend="vtkhdf", ) dim = mesh.topology.dim org_maps = {} mesh.geometry.x[:, 0] *= 2.0 + mesh.geometry.x[:, 1] io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.append, time=2.5, backend_args={"name": "tet"}, backend="vtkhdf", ) for dim in range(mesh.topology.dim + 1): mesh.topology.create_connectivity(dim, mesh.topology.dim) entities = np.arange(mesh.topology.index_map(dim).size_local, dtype=np.int32) et = meshtags(mesh, dim, entities, entities + mesh.topology.index_map(dim).local_range[0]) io4dolfinx.write_meshtags( filename, mesh, et, f"{dim}tags", backend_args={"name": "tet"}, backend="vtkhdf" ) org_maps[dim] = generate_reference_map(mesh, et, comm, root) tol = 15 * np.finfo(dtype).eps # Read in hex grid from second time step hex_mesh = io4dolfinx.read_mesh( filename, comm, time=1.0, backend_args={"name": "hex"}, backend="vtkhdf" ) hex_tag = io4dolfinx.read_meshtags( filename, hex_mesh, "CellTags", backend_args={"name": "hex"}, backend="vtkhdf" ) read_map = generate_reference_map(hex_mesh, hex_tag, comm, root) # On root process, check that midpoints are the same for each value in the meshtag if MPI.COMM_WORLD.rank == root: assert len(org_map) == len(read_map) for value, (_, midpoint) in org_map.items(): _, read_midpoint = read_map[value] np.testing.assert_allclose(read_midpoint, midpoint, atol=tol) # Read tet grid from second time step tet_mesh = io4dolfinx.read_mesh( filename, comm, time=2.5, backend_args={"name": "tet"}, backend="vtkhdf" ) for dim in range(mesh.topology.dim + 1): tet_tag = io4dolfinx.read_meshtags( filename, tet_mesh, f"{dim}tags", backend_args={"name": "tet"}, backend="vtkhdf" ) read_map = generate_reference_map(tet_mesh, tet_tag, comm, root) if MPI.COMM_WORLD.rank == root: assert len(org_maps[dim]) == len(read_map) for value, (_, midpoint) in org_maps[dim].items(): _, read_midpoint = read_map[value] np.testing.assert_allclose(read_midpoint, midpoint, atol=tol) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_read_write_pointdata(dtype, tmp_path): tol = 15 * np.finfo(dtype).eps comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() filename = tmp_path / "point_data.vtkhdf" mesh = create_unit_cube(comm, 3, 3, 3, dtype=dtype, cell_type=CellType.hexahedron) def f(x, t): return (x[0] + np.sin(x[1]) + np.cos(x[0] * t), x[2] + x[1] - t) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.write, time=1.0, backend_args={"name": "hex"}, backend="vtkhdf", ) f_name = "point_data" V = functionspace(mesh, ("Lagrange", 2, (2,))) u = Function(V, dtype=dtype, name=f_name) u.interpolate(lambda x: f(x, 1.0)) io4dolfinx.write_point_data( filename, u, mode=io4dolfinx.FileMode.append, time=1.0, backend_args={"name": "hex"}, backend="vtkhdf", ) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.append, time=2.0, backend_args={"name": "hex"}, backend="vtkhdf", ) u.interpolate(lambda x: f(x, 2.0)) io4dolfinx.write_point_data( filename, u, mode=io4dolfinx.FileMode.append, time=2.0, backend_args={"name": "hex"}, backend="vtkhdf", ) # Read in hex grid from second time step hex_mesh = io4dolfinx.read_mesh( filename, comm, time=2.0, backend_args={"name": "hex"}, backend="vtkhdf" ) u_end = io4dolfinx.read_point_data( filename, mesh=hex_mesh, name=f_name, time=2.0, backend_args={"name": "hex"}, backend="vtkhdf", ) u_ref = Function(u_end.function_space, dtype=dtype) u_ref.interpolate(lambda x: f(x, 2.0)) np.testing.assert_allclose(u_end.x.array, u_ref.x.array, atol=tol) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_read_write_celldata(dtype, tmp_path): tol = 15 * np.finfo(dtype).eps comm = MPI.COMM_WORLD tmp_path = comm.bcast(tmp_path, root=0) comm.barrier() filename = tmp_path / "cell_data.vtkhdf" mesh = create_unit_cube(comm, 3, 3, 3, dtype=dtype, cell_type=CellType.tetrahedron) def f(x, t): return (x[0] + np.sin(x[1]) + np.cos(x[0] * t), x[1] + t, x[0] - t) t_0 = 2.2 t_1 = 3.0 backend_args = {"name": "Grid"} io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.write, time=t_0, backend_args=backend_args, backend="vtkhdf", ) f_name = "Data" V = functionspace(mesh, ("DG", 0, (3,))) u = Function(V, dtype=dtype, name=f_name) u.interpolate(lambda x: f(x, t_0)) io4dolfinx.write_cell_data( filename, u, mode=io4dolfinx.FileMode.append, time=t_0, backend_args=backend_args, backend="vtkhdf", ) io4dolfinx.write_mesh( filename, mesh, mode=io4dolfinx.FileMode.append, time=t_1, backend_args=backend_args, backend="vtkhdf", ) u.interpolate(lambda x: f(x, t_1)) io4dolfinx.write_cell_data( filename, u, mode=io4dolfinx.FileMode.append, time=t_1, backend_args=backend_args, backend="vtkhdf", ) for t in [t_1, t_0]: grid = io4dolfinx.read_mesh( filename, comm, time=t, backend_args=backend_args, backend="vtkhdf" ) u_end = io4dolfinx.read_cell_data( filename, mesh=grid, name=f_name, time=t, backend_args=backend_args, backend="vtkhdf", ) u_ref = Function(u_end.function_space, dtype=dtype) u_ref.interpolate(lambda x: f(x, t)) np.testing.assert_allclose(u_end.x.array, u_ref.x.array, atol=tol) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_read_write_mix_data(dtype, tmp_path): tol = 15 * np.finfo(dtype).eps mesh = create_unit_square(MPI.COMM_WORLD, 5, 7, dtype=dtype) def f(x, t): return x[0] + t * x[1] def g(x, t): return x[1] - x[0] * t**2 ts = [0.1, 0.3, 0.4, 0.5] V = functionspace(mesh, ("Lagrange", 1)) u = Function(V, name="points", dtype=dtype) z = Function(V, name="some_other_array", dtype=dtype) Q = functionspace(mesh, ("DG", 0)) q = Function(Q, name="cells", dtype=dtype) tmp_path = mesh.comm.bcast(tmp_path, root=0) filename = tmp_path / "mixed_data.vtkhdf" backend_args = {"name": "MyGrid"} for i, t in enumerate(ts): if np.isclose(t, 0.1): mode = io4dolfinx.FileMode.write else: mode = io4dolfinx.FileMode.append mesh.geometry.x[:] *= 1 + 0.1 * t io4dolfinx.write_mesh( filename, mesh, mode=mode, time=t, backend="vtkhdf", backend_args=backend_args ) u.interpolate(lambda x: f(x, t)) q.interpolate(lambda x: g(x, t)) io4dolfinx.write_point_data( filename, u, time=t, mode=io4dolfinx.FileMode.append, backend_args=backend_args, backend="vtkhdf", ) io4dolfinx.write_point_data( filename, z, time=t, mode=io4dolfinx.FileMode.append, backend_args=backend_args, backend="vtkhdf", ) mesh_in = io4dolfinx.read_mesh( filename, MPI.COMM_WORLD, time=t, backend_args=backend_args, backend="vtkhdf" ) u_in = io4dolfinx.read_point_data( filename, name=u.name, mesh=mesh_in, time=t, backend_args=backend_args, backend="vtkhdf" ) u_ref = Function(u_in.function_space, dtype=dtype) u_ref.interpolate(lambda x: f(x, t)) np.testing.assert_allclose(u_ref.x.array, u_in.x.array, atol=tol) c_step = i if not np.isclose(t, 0.3): io4dolfinx.write_cell_data( filename, q, time=t, backend_args=backend_args, mode=io4dolfinx.FileMode.append, backend="vtkhdf", ) else: # Read in mesh from previous step as geometry adapts, # while reading data from current step. c_step = i - 1 mesh_in = io4dolfinx.read_mesh( filename, MPI.COMM_WORLD, time=ts[c_step], backend_args=backend_args, backend="vtkhdf", ) q_in = io4dolfinx.read_cell_data( filename, name=q.name, mesh=mesh_in, time=t, backend_args=backend_args, backend="vtkhdf" ) q_ref = Function(q_in.function_space, dtype=dtype) q_ref.interpolate(lambda x: g(x, ts[c_step])) np.testing.assert_allclose(q_ref.x.array, q_in.x.array, atol=tol) scientificcomputing-io4dolfinx-d21fc0e/tests/test_xdmf.py000066400000000000000000000044131517634040500240430ustar00rootroot00000000000000from mpi4py import MPI import dolfinx import numpy as np import ufl import io4dolfinx def test_xdmf_mesh(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) tmp_file = tmp_path / "mesh.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, tmp_file, "w") as xdmf: xdmf.write_mesh(mesh) MPI.COMM_WORLD.barrier() in_grid = io4dolfinx.read_mesh(tmp_file, MPI.COMM_WORLD, backend="xdmf") assert mesh.topology.dim == in_grid.topology.dim assert mesh.geometry.dim == in_grid.geometry.dim for i in range(mesh.topology.dim): mesh.topology.create_entities(i) o_map = mesh.topology.index_map(i) in_grid.topology.create_entities(i) map = in_grid.topology.index_map(i) assert o_map.size_global == map.size_global assert mesh.geometry.index_map().size_global == in_grid.geometry.index_map().size_global org_area = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=mesh))), op=MPI.SUM ) area = mesh.comm.allreduce( dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * ufl.dx(domain=in_grid))), op=MPI.SUM ) assert np.isclose(area, org_area) def test_xdmf_function(tmp_path): tmp_path = MPI.COMM_WORLD.bcast(tmp_path, root=0) mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 10) def f(x): return (x[0], x[1], -2 * x[1], 3 * x[0]) V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1, (4,))) u = dolfinx.fem.Function(V, name="u") u.interpolate(f) tmp_file = tmp_path / "function.xdmf" with dolfinx.io.XDMFFile(MPI.COMM_WORLD, tmp_file, "w") as xdmf: xdmf.write_mesh(mesh) xdmf.write_function(u) MPI.COMM_WORLD.barrier() in_grid = io4dolfinx.read_mesh(tmp_file, MPI.COMM_WORLD, backend="xdmf") names = io4dolfinx.read_function_names( tmp_file, MPI.COMM_WORLD, backend_args={}, backend="xdmf" ) assert len(names) == 1 assert names[0] == "u" u = io4dolfinx.read_point_data(tmp_file, "u", in_grid, backend="xdmf") u_ref = dolfinx.fem.Function(u.function_space) u_ref.interpolate(f) eps = np.finfo(mesh.geometry.x.dtype).eps np.testing.assert_allclose(u.x.array, u_ref.x.array, atol=eps)