diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 00000000..3b3c9cab --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,56 @@ +## Contributing Guidelines + + +- [ ] I have read and understood the [CONTRIBUTING.md](https://github.com/Autodesk/XLB/blob/main/CONTRIBUTING.md) guidelines + + +## Description + + + + +## Type of change + + +- [ ] Bug fix (non-breaking change which fixes an issue) +- [ ] New feature (non-breaking change which adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) +- [ ] Documentation update + +## How Has This Been Tested? + + +- [ ] All pytest tests pass + + + + +## Linting and Code Formatting + +Make sure the code follows the project's linting and formatting standards. This project uses **Ruff** for linting. + +To run Ruff, execute the following command from the root of the repository: + +```bash +ruff check . +``` + + + +- [ ] Ruff passes diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml index 07cea5aa..cb0a2d99 100644 --- a/.github/workflows/cla.yml +++ b/.github/workflows/cla.yml @@ -9,14 +9,22 @@ jobs: CLA-Assistant: runs-on: ubuntu-latest steps: + - uses: actions/create-github-app-token@v2 + id: app-token + with: + app-id: ${{ vars.CLA_BOT_APP_ID }} + private-key: ${{ secrets.CLA_BOT_PRIVATE_KEY }} + owner: Autodesk + repositories: | + CLA-Signatures - name: "CLA Assistant" if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target' # Beta Release - uses: contributor-assistant/github-action@v2.3.0 + uses: contributor-assistant/github-action@v2.6.1 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # the below token should have repo scope and must be manually added by you in the repository's secret - PERSONAL_ACCESS_TOKEN : ${{ secrets.CLA_BOT_SECRET }} + PERSONAL_ACCESS_TOKEN: ${{ steps.app-token.outputs.token }} with: remote-repository-name: 'CLA-Signatures' remote-organization-name: 'Autodesk' diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..1b44c5a1 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,27 @@ +name: Lint + +on: + pull_request: + branches: + - major-refactoring # Remember to add main branch later + +jobs: + lint: + runs-on: ubuntu-latest + + steps: + - name: Check out code + uses: actions/checkout@v2 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install ruff + + - name: Run Ruff + run: ruff check . diff --git a/.github/workflows/mkdocs.yml b/.github/workflows/mkdocs.yml index 3719e6e5..413426db 100644 --- a/.github/workflows/mkdocs.yml +++ b/.github/workflows/mkdocs.yml @@ -2,7 +2,6 @@ name: mkdocs-deployment on: push: branches: - - master - main # paths: # - 'docs/**' diff --git a/.gitignore b/.gitignore index d7744518..a47b0458 100644 --- a/.gitignore +++ b/.gitignore @@ -39,7 +39,7 @@ Thumbs.db __pycache__/ *.py[cod] *$py.class - +**pyc # C extensions *.so @@ -131,6 +131,8 @@ venv/ ENV/ env.bak/ venv.bak/ +.xlb-env/ +.xlb_install_test_venvs/ # Spyder project settings .spyderproject @@ -149,4 +151,25 @@ dmypy.json # Checkpoints (default dir for XLB) -checkpoints/* \ No newline at end of file +checkpoints/* + +# Ignore Python packaging build directories +dist/ +build/ +*.egg-info/ +*.dot + +# Ignore h5 and xmf formats +*.h5 +*.xmf + +# Ignore CSV files +*.csv + +# USD files +*.usd +*.usda +*.usdc +*.usd.gz +*.usd.zip +*.usd.bz2 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..6a2bd2f4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,6 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.5.6 + hooks: + - id: ruff + args: [--fix] diff --git a/AUTHORS b/AUTHORS index bcff69f3..4fc34969 100644 --- a/AUTHORS +++ b/AUTHORS @@ -4,4 +4,6 @@ # For a comprehensive view of all contributors, please refer to the revision history in the source control. Mehdi Ataei (Autodesk Inc) -Hesam Saleipour (Autodesk Inc) \ No newline at end of file +Hesam Saleipour (Autodesk Inc) +Oliver Hennigh (NVIDIA) +Massimiliano Meneghin (Autodesk Inc) \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..48f9cc7b --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,30 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] +- _No changes yet_ + + +## [0.2.1] - 2024-12-05 + +### Fixed +- mkdocs is now configured correctly for the new project structure +- JAX installation is now handled correctly for different configurations (CPU, CUDA, TPU) +- Fixed a couple of bugs in 2D regularied_bc and kbc (Warp) that emerged after merging 2d and 3d kernels + +### Added + +- Added abstraction layer for boundary condition efficient encoding/decoding of auxiliary data +- Added the capability to add profiles to boundary conditions +- Added prepare_fields method to the Stepper class to allow for more automatic preparation of fields + +## [0.2.0] - 2024-10-22 + +### Added +- XLB is now installable via pip +- Complete rewrite of the codebase for better modularity and extensibility based on "Operators" design pattern +- Added NVIDIA's Warp backend for state-of-the-art performance \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b31110ee..1a438143 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -38,3 +38,150 @@ If you would like to contribute your code to XLB, you should: When you submit your code, please include relevant tests as part of the pull request, and ensure that your comments and coding style align with the rest of the project. You can refer to the existing code for examples of the testing and style practices that the project follows. + +Important: Ensure that your commits are atomic and self-contained. Each PR should only make a single, cohesive change. You should also squash your commits into a single commit as described below before submitting your PR. + +## Detailed Contribution Guidelines + +### 1. Setup Your Local Environment + +- **Clone Your Fork:** + If you haven't yet cloned your copy of the repository, you can do so with the following command: + + ```bash + git clone https://github.com/Autodesk/XLB + cd XLB + ``` + +- **Add Upstream Remote:** Set up the upstream remote to track the original repository. + + ```bash + git remote add upstream https://github.com/Autodesk/XLB + ``` + + You can check your remotes to ensure everything is set up correctly: + + ```bash + git remote -v + ``` + +### 2. Syncing Your Main Branch with Upstream + +- **Fetch Updates from Upstream:** + To keep your local repository up to date with the upstream `main` branch: + + ```bash + git fetch upstream + ``` + +- **Sync Your Main Branch:** + Checkout to your local `main` and merge the upstream changes to ensure it's always up to date: + + ```bash + git checkout main + git merge upstream/main + ``` + +- **Push to Your Fork (Optional):** + It is a good practice to also keep the fork on GitHub in sync: + + ```bash + git push origin main + ``` + +### 3. Create a Feature Branch for Your Contribution + +- **Create and Checkout a New Branch:** + Always work on a new branch for each feature or issue to keep things organized: + ```bash + git checkout -b + ``` + Choose a descriptive branch name that makes it clear what your contribution is. + +### 4. Make Your Changes + +- **Make Changes and Commit:** + Make all the changes you need, then stage and commit them: + + ```bash + git add . + git commit -m "Description of the changes made" + ``` + +- **Amend or Squash Commits (Optional):** + If you need to update the commit message or add more changes before pushing, you can amend your commit: + + ```bash + git add . + git commit --amend + ``` + + This will let you update the commit message or include additional changes in a single commit. + + To consolidate your existing commits into a single PR, first, you need to reset your branch to the point where you want the single commit to start from (likely the last commit on origin/main before your work began). + + ```bash + git reset --soft + ``` + + Replace `` with the hash of the commit that should be the base of your PR (it is likely the last commit on origin/main before your changes, use that ID). + + Now, all your work will be staged as if it’s a single set of changes. You can commit this with a new message that represents the combined work. + +```bash +git commit -m "Combined changes" +``` + +Finally, you can push your changes to your fork with the following command: + +```bash +git push --force origin +``` + +### 5. Pushing Your Branch and Creating a Pull Request + +- **Push Your Branch to Your Fork:** + + ```bash + git push origin + ``` + +- **Create a Pull Request (PR):** + Go to the repository on GitHub, and you should see an option to create a Pull Request from your recently pushed branch. Follow the steps to create the PR. + +### 6. Handling Feedback and Updating PR + +- **Make Changes Based on Feedback:** + If changes are requested in the PR, make those changes in your local branch and amend the commit if needed: + ```bash + git add . + git commit --amend + git push --force origin + ``` + The `--force` flag is necessary because you amended an existing commit, and you need to update the remote branch accordingly. + +### 7. Finalizing and Merging + +- **Squash Commits on Maintainer Side:** + When the PR is ready to be merged, the maintainer *will* squash multiple commits into a single one, or you can amend and force push until only a single commit is present. + +- **Sync Your Fork Main Branch Again:** + Once your PR is merged, make sure to sync your local and forked `main` branch again: + + ```bash + git checkout main + git fetch upstream + git merge upstream/main + git push origin main + ``` + +### 8. Start a New Contribution + +- **Create a New Branch:** + For each new contribution, repeat the branching step: + ```bash + git checkout main + git checkout -b + ``` +--- +This workflow ensures every contribution is separate and cleanly managed. diff --git a/README.md b/README.md index dcb67eeb..cbf56a0a 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,68 @@ [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![GitHub star chart](https://img.shields.io/github/stars/Autodesk/XLB?style=social)](https://star-history.com/#Autodesk/XLB)

- +

# XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning -XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It's built on top of the [JAX](https://github.com/google/jax) library and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. +XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax), [NVIDIA Warp](https://github.com/NVIDIA/warp), and [Neon](https://github.com/Autodesk/Neon) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the Warp backend, XLB offers state-of-the-art single-GPU performance, and with the new Neon backend it extends to multi-GPU (single-resolution). More importantly, the Neon backend provides grid refinement capabilities for multi-resolution simulations. + +## Getting Started +To get started with XLB, you can install it using pip. There are different installation options depending on your hardware and needs: + +### Basic Installation (CPU-only) +```bash +pip install xlb +``` + +### Installation with Warp support (single-GPU) +For the NVIDIA Warp backend (single-GPU, state-of-the-art performance): +```bash +pip install "xlb[warp]" +``` + +### Installation with CUDA support (for NVIDIA GPUs) +This installation is for the JAX backend with CUDA support: +```bash +pip install "xlb[cuda]" +``` + +### Installation with TPU support +This installation is for the JAX backend with TPU support: +```bash +pip install "xlb[tpu]" +``` + +### Installation with Neon support +Neon backend enables multi-GPU dense and single-GPU multi-resolution representations. +Install XLB with Neon support using: + +```bash +git clone https://github.com/Autodesk/XLB.git +cd XLB +pip install -r requirements.txt +pip install '.[neon]' +``` + +**Requirements:** The Neon wheel supports **Python 3.11** to **Python 3.14** on **Linux x86_64** and **Linux ARM**. + +**Note:** Neon uses a custom fork of warp. + +### Notes: +- For Mac users: Use the basic CPU installation command as JAX's GPU support is not available on MacOS +- Use `xlb[warp]` for the Warp backend (single-GPU) or `xlb[neon]` for the Neon backend (multi-GPU / multi-resolution). Do not install both in the same environment. +- The installation options for CUDA and TPU only affect the JAX backend + +To install the latest development version from source: + +```bash +pip install git+https://github.com/Autodesk/XLB.git +``` + +The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md). + +For examples to get you started please refer to the [examples](https://github.com/Autodesk/XLB/tree/main/examples) folder. ## Accompanying Paper @@ -28,10 +84,36 @@ If you use XLB in your research, please cite the following paper: } ``` +If you use the grid refinement capabilities in your work, please also cite: + +``` +@inproceedings{mahmoud2024optimized, + title={Optimized {GPU} implementation of grid refinement in lattice {Boltzmann} method}, + author={Mahmoud, Ahmed H and Salehipour, Hesam and Meneghin, Massimiliano}, + booktitle={2024 IEEE International Parallel and Distributed Processing Symposium (IPDPS)}, + pages={398--407}, + year={2024}, + organization={IEEE} +} + +@inproceedings{meneghin2022neon, + title={Neon: A Multi-{GPU} Programming Model for Grid-based Computations}, + author={Meneghin, Massimiliano and Mahmoud, Ahmed H. and Jayaraman, Pradeep Kumar and Morris, Nigel J. W.}, + booktitle={Proceedings of the 36th IEEE International Parallel and Distributed Processing Symposium}, + pages={817--827}, + year={2022}, + month={june}, + doi={10.1109/IPDPS53621.2022.00084}, + url={https://escholarship.org/uc/item/9fz7k633} +} +``` + ## Key Features +- **Multiple Backend Support:** XLB includes support for JAX, NVIDIA Warp, and Neon backends, providing *state-of-the-art* performance for lattice Boltzmann simulations. The Warp backend targets single-GPU runs, while the Neon backend enables multi-GPU single-resolution and single-GPU multi-resolution simulations. +- **Multi-Resolution Grid Refinement:** Mesh refinement with nested cuboid grids and multiple kernel-fusion strategies for optimal performance on the Neon backend. - **Integration with JAX Ecosystem:** The library can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more. - **Differentiable LBM Kernels:** XLB provides differentiable LBM kernels that can be used in differentiable physics and deep learning applications. -- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems, enabling the execution of large-scale simulations on hundreds of GPUs with billions of cells. +- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems using the JAX backend or the Neon backend, enabling the execution of large-scale simulations on hundreds of GPUs with billions of cells. - **Support for Various LBM Boundary Conditions and Kernels:** XLB supports several LBM boundary conditions and collision kernels. - **User-Friendly Interface:** Written entirely in Python, XLB emphasizes a highly accessible interface that allows users to extend the library with ease and quickly set up and run new simulations. - **Leverages JAX Array and Shardmap:** The library incorporates the new JAX array unified array type and JAX shardmap, providing users with a numpy-like interface. This allows users to focus solely on the semantics, leaving performance optimizations to the compiler. @@ -40,9 +122,16 @@ If you use XLB in your research, please cite the following paper: ## Showcase +

+ Wind Turbine Simulation +

+

+ Simulation of a wind turbine based on the immersed boundary method. +

+

- +

On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells. @@ -50,21 +139,21 @@ If you use XLB in your research, please cite the following paper:

- +

DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells

- +

- Airflow in to, out of, and within a building (~400 million cells) + Airflow into, out of, and within a building (~400 million cells)

- +

The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details) @@ -73,7 +162,7 @@ The stages of a fluid density field from an initial state to the emergence of th

- +

Lid-driven Cavity flow at Re=100,000 (~25 million cells) @@ -85,6 +174,7 @@ The stages of a fluid density field from an initial state to the emergence of th - BGK collision model (Standard LBM collision model) - KBC collision model (unconditionally stable for flows with high Reynolds number) +- Smagorinsky LES sub-grid model for turbulence modelling ### Machine Learning @@ -99,21 +189,26 @@ The stages of a fluid density field from an initial state to the emergence of th - D3Q27 (Must be used for KBC simulation runs) ### Compute Capabilities -- Distributed Multi-GPU support +- Single GPU support for the Warp backend with state-of-the-art performance +- Multi-GPU support using the Neon backend with single-resolution grids +- Grid refinement support on single-GPU using the Neon backend +- Distributed Multi-GPU support using the JAX backend - Mixed-Precision support (store vs compute) +- Multiple kernel-fusion performance strategies for multi-resolution simulations - Out-of-core support (coming soon) ### Output - Binary and ASCII VTK output (based on PyVista library) +- HDF5/XDMF output for multi-resolution data (with gzip compression) - In-situ rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze) library - [Orbax](https://github.com/google/orbax)-based distributed asynchronous checkpointing -- Image Output +- Image Output (including multi-resolution slice images) - 3D mesh voxelizer using trimesh ### Boundary conditions -- **Equilibrium BC:** In this boundary condition, the fluid populations are assumed to be in at equilibrium. Can be used to set prescribed velocity or pressure. +- **Equilibrium BC:** In this boundary condition, the fluid populations are assumed to be at equilibrium. Can be used to set prescribed velocity or pressure. - **Full-Way Bounceback BC:** In this boundary condition, the velocity of the fluid populations is reflected back to the fluid side of the boundary, resulting in zero fluid velocity at the boundary. @@ -125,50 +220,22 @@ The stages of a fluid density field from an initial state to the emergence of th - **Regularized BC:** This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. This BC is more stable than Zouhe BC, but computationally more expensive. - **Extrapolation Outflow BC:** A type of outflow boundary condition that uses extrapolation to avoid strong wave reflections. -- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition due to Bouzidi for a lattice Boltzmann method simulation. - -## Installation Guide +- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition for representing curved boundaries. -To use XLB, you must first install JAX and other dependencies using the following commands: +- **Hybrid BC:** Combines regularized and bounce-back methods with optional wall-distance interpolation for improved accuracy on curved geometries. +## Roadmap -Please refer to https://github.com/google/jax for the latest installation documentation. The following table is taken from [JAX's Github page](https://github.com/google/jax). - -| Hardware | Instructions | -|------------|-----------------------------------------------------------------------------------------------------------------| -| CPU | `pip install -U "jax[cpu]"` | -| NVIDIA GPU on x86_64 | `pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html` | -| Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html` | -| AMD GPU | Use [Docker](https://hub.docker.com/r/rocm/jax) or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). | -| Apple GPU | Follow [Apple's instructions](https://developer.apple.com/metal/jax/). | +### Recently Completed -**Note:** We encountered challenges when executing XLB on Apple GPUs due to the lack of support for certain operations in the Metal backend. We advise using the CPU backend on Mac OS. We will be testing XLB on Apple's GPUs in the future and will update this section accordingly. + - βœ… **Grid Refinement:** Multi-resolution LBM with nested cuboid grids and multiple kernel-fusion strategies via the Neon backend. - -Install dependencies: -```bash -pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor -``` - -Run an example: -```bash -git clone https://github.com/Autodesk/XLB -cd XLB -export PYTHONPATH=. -python3 examples/CFD/cavity2d.py -``` -## Roadmap + - βœ… **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Multi-GPU support through Neon's data structures with Warp-based kernels for single-resolution settings. ### Work in Progress (WIP) *Note: Some of the work-in-progress features can be found in the branches of the XLB repository. For contributions to these features, please reach out.* -- πŸš€ **Warp Backend:** Achieving state-of-the-art performance by leveraging the [Warp](https://github.com/NVIDIA/warp) framework in combination with JAX. - - - 🌐 **Grid Refinement:** Implementing adaptive mesh refinement techniques for enhanced simulation accuracy. - -- ⚑ **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Using Neon's data structure for improved scaling. - -- πŸ’Ύ **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips. + - πŸ’Ύ **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips (coming soon). - πŸ—œοΈ **GPU Accelerated Lossless Compression and Decompression**: Implementing high-performance lossless compression and decompression techniques for larger-scale simulations and improved performance. diff --git a/assets/wind_turbine.gif b/assets/wind_turbine.gif new file mode 100644 index 00000000..62f8d7d1 Binary files /dev/null and b/assets/wind_turbine.gif differ diff --git a/docs/assets/airfoil.png b/docs/assets/airfoil.png deleted file mode 100644 index f0ee4c9c..00000000 Binary files a/docs/assets/airfoil.png and /dev/null differ diff --git a/docs/assets/car.png b/docs/assets/car.png deleted file mode 100644 index ee9f0b03..00000000 Binary files a/docs/assets/car.png and /dev/null differ diff --git a/docs/assets/cavity.gif b/docs/assets/cavity.gif deleted file mode 100644 index 7cd9bdd4..00000000 Binary files a/docs/assets/cavity.gif and /dev/null differ diff --git a/docs/assets/logo-transparent.png b/docs/assets/logo-transparent.png deleted file mode 100644 index 27c85681..00000000 Binary files a/docs/assets/logo-transparent.png and /dev/null differ diff --git a/docs/assets/logo.svg b/docs/assets/logo.svg deleted file mode 100644 index 9efe4992..00000000 --- a/docs/assets/logo.svg +++ /dev/null @@ -1,12 +0,0 @@ - - - - - - - diff --git a/docs/base.md b/docs/base.md deleted file mode 100644 index 7b848bd8..00000000 --- a/docs/base.md +++ /dev/null @@ -1 +0,0 @@ -::: src.base.LBMBase \ No newline at end of file diff --git a/docs/boundary_conditions.md b/docs/boundary_conditions.md deleted file mode 100644 index 3c0d3a86..00000000 --- a/docs/boundary_conditions.md +++ /dev/null @@ -1,18 +0,0 @@ -::: src.boundary_conditions.BoundaryCondition - -::: src.boundary_conditions.BounceBack - -::: src.boundary_conditions.BounceBackMoving - -::: src.boundary_conditions.BounceBackHalfway - -::: src.boundary_conditions.EquilibriumBC - -::: src.boundary_conditions.DoNothing - -::: src.boundary_conditions.ZouHe - -::: src.boundary_conditions.Regularized - -::: src.boundary_conditions.ExtrapolationOutflow - diff --git a/docs/index.md b/docs/index.md index 54e9bbb4..5a74c0d1 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,63 +1,139 @@ +[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) +[![GitHub star chart](https://img.shields.io/github/stars/Autodesk/XLB?style=social)](https://star-history.com/#Autodesk/XLB)

- +

-# XLB: A Hardware-Accelerated Differentiable Lattice Boltzmann Simulation Framework based on JAX for Physics-based Machine Learning +# XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning -XLB (Accelerated LB) is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) solver that leverages hardware acceleration. It's built on top of the [JAX](https://github.com/google/jax) library and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. +πŸŽ‰ **Exciting News!** πŸŽ‰ XLB version 0.2.0 has been released, featuring a complete rewrite of the library and introducing support for the NVIDIA Warp backend! +XLB can now be installed via pip: `pip install xlb`. + +XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax) and [NVIDIA Warp](https://github.com/NVIDIA/warp) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the new Warp backend, XLB now offers state-of-the-art performance for even faster simulations. + +## Getting Started +To get started with XLB, you can install it using pip: +```bash +pip install xlb +``` + +To install the latest development version from source: + +```bash +pip install git+https://github.com/Autodesk/XLB.git +``` + +The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md). + +For examples to get you started please refer to the [examples](https://github.com/Autodesk/XLB/tree/main/examples) folder. + +## Accompanying Paper + +Please refer to the [accompanying paper](https://doi.org/10.1016/j.cpc.2024.109187) for benchmarks, validation, and more details about the library. + +## Citing XLB + +If you use XLB in your research, please cite the following paper: + +``` +@article{ataei2024xlb, + title={{XLB}: A differentiable massively parallel lattice {Boltzmann} library in {Python}}, + author={Ataei, Mohammadmehdi and Salehipour, Hesam}, + journal={Computer Physics Communications}, + volume={300}, + pages={109187}, + year={2024}, + publisher={Elsevier} +} +``` ## Key Features -- **Integration with JAX Ecosystem:** The solver can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more. -- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems, enabling the execution of large-scale simulations with billions of voxels. +- **Multiple Backend Support:** XLB now includes support for multiple backends including JAX and NVIDIA Warp, providing *state-of-the-art* performance for lattice Boltzmann simulations. Currently, only single GPU is supported for the Warp backend. +- **Integration with JAX Ecosystem:** The library can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more. +- **Differentiable LBM Kernels:** XLB provides differentiable LBM kernels that can be used in differentiable physics and deep learning applications. +- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems using the JAX backend, enabling the execution of large-scale simulations on hundreds of GPUs with billions of cells. - **Support for Various LBM Boundary Conditions and Kernels:** XLB supports several LBM boundary conditions and collision kernels. -- **User-Friendly Interface:** Written entirely in Python, XLB emphasizes a highly accessible interface that allows users to extend the solver with ease and quickly set up and run new simulations. -- **Leverages JAX Array and Shardmap:** The solver incorporates the new JAX array unified array type and JAX shardmap, providing users with a numpy-like interface. This allows users to focus solely on the semantics, leaving performance optimizations to the compiler. +- **User-Friendly Interface:** Written entirely in Python, XLB emphasizes a highly accessible interface that allows users to extend the library with ease and quickly set up and run new simulations. +- **Leverages JAX Array and Shardmap:** The library incorporates the new JAX array unified array type and JAX shardmap, providing users with a numpy-like interface. This allows users to focus solely on the semantics, leaving performance optimizations to the compiler. - **Platform Versatility:** The same XLB code can be executed on a variety of platforms including multi-core CPUs, single or multi-GPU systems, TPUs, and it also supports distributed runs on multi-GPU systems or TPU Pod slices. +- **Visualization:** XLB provides a variety of visualization options including in-situ on GPU rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze). -## Documentation -The documentation can be found [here](https://autodesk.github.io/XLB/) (in preparation) ## Showcase -The following examples showcase the capabilities of XLB:

- + +

+

+ On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells. +

+ + +

+ +

+

+ DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells +

+ +

+

- Lid-driven Cavity flow at Re=100,000 (~25 million voxels) + Airflow in to, out of, and within a building (~400 million cells)

- +

- DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million voxels +The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details)

+
+

- +

- Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with approx. 100 million voxels + Lid-driven Cavity flow at Re=100,000 (~25 million cells)

## Capabilities ### LBM + - BGK collision model (Standard LBM collision model) - KBC collision model (unconditionally stable for flows with high Reynolds number) +### Machine Learning + +- Easy integration with JAX's ecosystem of machine learning libraries +- Differentiable LBM kernels +- Differentiable boundary conditions + ### Lattice Models + - D2Q9 - D3Q19 - D3Q27 (Must be used for KBC simulation runs) +### Compute Capabilities +- Single GPU support for the Warp backend with state-of-the-art performance +- Distributed Multi-GPU support using the JAX backend +- Mixed-Precision support (store vs compute) +- Out-of-core support (coming soon) + ### Output + - Binary and ASCII VTK output (based on PyVista library) +- In-situ rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze) library +- [Orbax](https://github.com/google/orbax)-based distributed asynchronous checkpointing - Image Output - 3D mesh voxelizer using trimesh ### Boundary conditions + - **Equilibrium BC:** In this boundary condition, the fluid populations are assumed to be in at equilibrium. Can be used to set prescribed velocity or pressure. - **Full-Way Bounceback BC:** In this boundary condition, the velocity of the fluid populations is reflected back to the fluid side of the boundary, resulting in zero fluid velocity at the boundary. @@ -70,43 +146,46 @@ The following examples showcase the capabilities of XLB: - **Regularized BC:** This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. This BC is more stable than Zouhe BC, but computationally more expensive. - **Extrapolation Outflow BC:** A type of outflow boundary condition that uses extrapolation to avoid strong wave reflections. -### Compute Capabilities -- Distributed Multi-GPU support -- JAX shard-map and JAX Array support -- Mixed-Precision support (store vs compute) +- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition for representing curved boundaries. -## Installation Guide +## Roadmap -To use XLB, you must first install JAX and other dependencies using the following commands: +### Work in Progress (WIP) +*Note: Some of the work-in-progress features can be found in the branches of the XLB repository. For contributions to these features, please reach out.* -```bash -# Please refer to https://github.com/google/jax for the latest installation documentation + - 🌐 **Grid Refinement:** Implementing adaptive mesh refinement techniques for enhanced simulation accuracy. -pip install --upgrade pip + - πŸ’Ύ **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips (coming soon). -# For CPU run -pip install --upgrade "jax[cpu]" -# For GPU run +- ⚑ **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Using Neon's data structure for improved scaling. -# CUDA 12 and cuDNN 8.8 or newer. -pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +- πŸ—œοΈ **GPU Accelerated Lossless Compression and Decompression**: Implementing high-performance lossless compression and decompression techniques for larger-scale simulations and improved performance. -# CUDA 11 and cuDNN 8.6 or newer. -pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +- 🌑️ **Fluid-Thermal Simulation Capabilities:** Incorporating heat transfer and thermal effects into fluid simulations. -# Run dependencies -pip install jmp pyvista numpy matplotlib Rtree trimesh jmp -``` +- 🎯 **Adjoint-based Shape and Topology Optimization:** Implementing gradient-based optimization techniques for design optimization. -Run an example: -```bash -git clone https://github.com/Autodesk/XLB -cd XLB -export PYTHONPATH=. -python3 examples/cavity2d.py -``` -## Citing XLB -Accompanying publication coming soon: +- 🧠 **Machine Learning Accelerated Simulations:** Leveraging machine learning to speed up simulations and improve accuracy. + +- πŸ“‰ **Reduced Order Modeling using Machine Learning:** Developing data-driven reduced-order models for efficient and accurate simulations. + + +### Wishlist +*Contributions to these features are welcome. Please submit PRs for the Wishlist items.* + +- 🌊 **Free Surface Flows:** Simulating flows with free surfaces, such as water waves and droplets. + +- πŸ“‘ **Electromagnetic Wave Propagation:** Simulating the propagation of electromagnetic waves. + +- πŸ›©οΈ **Supersonic Flows:** Simulating supersonic flows. + +- 🌊🧱 **Fluid-Solid Interaction:** Modeling the interaction between fluids and solid objects. + +- 🧩 **Multiphase Flow Simulation:** Simulating flows with multiple immiscible fluids. + +- πŸ”₯ **Combustion:** Simulating combustion processes and reactive flows. + +- πŸͺ¨ **Particle Flows and Discrete Element Method:** Incorporating particle-based methods for granular and particulate flows. -**M. Ataei, H. Salehipour**. XLB: Hardware-Accelerated, Scalable, and Differentiable Lattice Boltzmann Simulation Framework based on JAX. TBA +- πŸ”§ **Better Geometry Processing Pipelines:** Improving the handling and preprocessing of complex geometries for simulations. diff --git a/docs/lattice.md b/docs/lattice.md deleted file mode 100644 index ca1fcb5e..00000000 --- a/docs/lattice.md +++ /dev/null @@ -1,7 +0,0 @@ -::: src.lattice.Lattice - -::: src.lattice.LatticeD2Q9 - -::: src.lattice.LatticeD3Q19 - -::: src.lattice.LatticeD3Q27 \ No newline at end of file diff --git a/docs/models.md b/docs/models.md deleted file mode 100644 index 3adf3fd3..00000000 --- a/docs/models.md +++ /dev/null @@ -1,5 +0,0 @@ -::: src.models.BGKSim - -::: src.models.KBCSim - -::: src.models.AdvectionDiffusionBGK \ No newline at end of file diff --git a/docs/utils.md b/docs/utils.md deleted file mode 100644 index 79637ed7..00000000 --- a/docs/utils.md +++ /dev/null @@ -1,13 +0,0 @@ -::: src.utils.downsample_field - -::: src.utils.save_image - -::: src.utils.save_fields_vtk - -::: src.utils.live_volume_randering - -::: src.utils.save_BCs_vtk - -::: src.utils.rotate_geometry - -::: src.utils.voxelize_stl diff --git a/examples/CFD/airfoil3d.py b/examples/CFD/airfoil3d.py deleted file mode 100644 index 77e4b58c..00000000 --- a/examples/CFD/airfoil3d.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -This is a example for simulating fluid flow around a NACA airfoil using the lattice Boltzmann method (LBM). -The LBM is a computational fluid dynamics method for simulating fluid flow and is particularly effective -for complex geometries and multiphase flow. - -In this example you'll be introduced to the following concepts: - -1. Lattice: The example uses a D3Q27 lattice, which is a three-dimensional lattice model that considers - 27 discrete velocity directions. This allows for a more accurate representation of the fluid flow - in three dimensions. - -2. NACA Airfoil Generation: The example includes a function to generate a NACA airfoil shape, which is - common in aerodynamics. The function allows for customization of the length, thickness, and angle - of the airfoil. - -3. Boundary Conditions: The example includes several boundary conditions. These include a "bounce back" - condition on the airfoil surface and the top and bottom of the domain, a "do nothing" condition - at the outlet (right side of the domain), and an "equilibrium" condition at the inlet - (left side of the domain) to simulate a uniform flow. - -4. Simulation Parameters: The example allows for the setting of various simulation parameters, - including the Reynolds number, inlet velocity, and characteristic length. - -5. In-situ visualization: The example outputs rendering images of the q-criterion using - PhantomGaze library (https://github.com/loliverhennigh/PhantomGaze) without any I/O overhead - while the data is still on the GPU. -""" - - -import numpy as np -# from IPython import display -import matplotlib.pylab as plt -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD3Q19, LatticeD3Q27 -from src.boundary_conditions import * -import numpy as np -from src.utils import * -from jax import config -import os -#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -import jax -import scipy - -# PhantomGaze for in-situ rendering -import phantomgaze as pg - -def makeNacaAirfoil(length, thickness=30, angle=0): - def nacaAirfoil(x, thickness, chordLength): - coeffs = [0.2969, -0.1260, -0.3516, 0.2843, -0.1015] - exponents = [0.5, 1, 2, 3, 4] - yt = [coeff * (x / chordLength) ** exp for coeff, exp in zip(coeffs, exponents)] - yt = 5. * thickness / 100 * chordLength * np.sum(yt) - - return yt - - x = np.linspace(0, length, num=length) - yt = np.array([nacaAirfoil(xi, thickness, length) for xi in x]) - - y_max = int(np.max(yt)) + 1 - domain = np.zeros((2 * y_max, len(x)), dtype=int) - - for i, xi in enumerate(x): - upper_bound = int(y_max + yt[i]) - lower_bound = int(y_max - yt[i]) - domain[lower_bound:upper_bound, i] = 1 - - domain = scipy.ndimage.rotate(domain, angle, reshape=True) - domain = np.where(domain > 0.5, 1, 0) - - return domain - -class Airfoil(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - tx, ty = np.array([self.nx, self.ny], dtype=int) - airfoil.shape - - airfoil_mask = np.pad(airfoil, ((tx // 3, tx - tx // 3), (ty // 2, ty - ty // 2)), 'constant', constant_values=False) - airfoil_mask = np.repeat(airfoil_mask[:, :, np.newaxis], self.nz, axis=2) - - airfoil_indices = np.argwhere(airfoil_mask) - wall = np.concatenate((airfoil_indices, - self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top'])) - self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy)) - - # Store airfoil boundary for visualization - self.visualization_bc = jnp.zeros((self.nx, self.ny, self.nz), dtype=jnp.float32) - self.visualization_bc = self.visualization_bc.at[tuple(airfoil_indices.T)].set(1.0) - - doNothing = self.boundingBoxIndices['right'] - self.BCs.append(DoNothing(tuple(doNothing.T), self.gridInfo, self.precisionPolicy)) - - inlet = self.boundingBoxIndices['left'] - rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_inlet = np.zeros((inlet.shape), dtype=self.precisionPolicy.compute_dtype) - - vel_inlet[:, 0] = prescribed_vel - self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_inlet, vel_inlet)) - - def output_data(self, **kwargs): - # Compute q-criterion and vorticity using finite differences - # Get velocity field - u = kwargs['u'][..., 1:-1, :] - # vorticity and q-criterion - norm_mu, q = q_criterion(u) - - # Make phantomgaze volume - dx = 0.01 - origin = (0.0, 0.0, 0.0) - upper_bound = (self.visualization_bc.shape[0] * dx, self.visualization_bc.shape[1] * dx, self.visualization_bc.shape[2] * dx) - q_volume = pg.objects.Volume( - q, - spacing=(dx, dx, dx), - origin=origin, - ) - norm_mu_volume = pg.objects.Volume( - norm_mu, - spacing=(dx, dx, dx), - origin=origin, - ) - boundary_volume = pg.objects.Volume( - self.visualization_bc, - spacing=(dx, dx, dx), - origin=origin, - ) - - # Make colormap for norm_mu - colormap = pg.Colormap("jet", vmin=0.0, vmax=0.05) - - # Get camera parameters - focal_point = (self.visualization_bc.shape[0] * dx / 2, self.visualization_bc.shape[1] * dx / 2, self.visualization_bc.shape[2] * dx / 2) - radius = 5.0 - angle = kwargs['timestep'] * 0.0001 - camera_position = (focal_point[0] + radius * np.sin(angle), focal_point[1], focal_point[2] + radius * np.cos(angle)) - - # Rotate camera - camera = pg.Camera(position=camera_position, focal_point=focal_point, view_up=(0.0, 1.0, 0.0), max_depth=30.0, height=1080, width=1920, background=pg.SolidBackground(color=(0.0, 0.0, 0.0))) - - # Make wireframe - screen_buffer = pg.render.wireframe(lower_bound=origin, upper_bound=upper_bound, thickness=0.01, camera=camera) - - # Render axes - screen_buffer = pg.render.axes(size=0.1, center=(0.0, 0.0, 1.1), camera=camera, screen_buffer=screen_buffer) - - # Render q-criterion - screen_buffer = pg.render.contour(q_volume, threshold=0.00003, color=norm_mu_volume, colormap=colormap, camera=camera, screen_buffer=screen_buffer) - - # Render boundary - boundary_colormap = pg.Colormap("bone_r", vmin=0.0, vmax=3.0, opacity=np.linspace(0.0, 6.0, 256)) - screen_buffer = pg.render.volume(boundary_volume, camera=camera, colormap=boundary_colormap, screen_buffer=screen_buffer) - - # Show the rendered image - plt.imsave('q_criterion_' + str(kwargs['timestep']).zfill(7) + '.png', np.minimum(screen_buffer.image.get(), 1.0)) - - -if __name__ == '__main__': - airfoil_length = 101 - airfoil_thickness = 30 - airfoil_angle = 20 - airfoil = makeNacaAirfoil(length=airfoil_length, thickness=airfoil_thickness, angle=airfoil_angle).T - precision = 'f32/f32' - - lattice = LatticeD3Q27(precision) - - nx = airfoil.shape[0] - ny = airfoil.shape[1] - - ny = 3 * ny - nx = 5 * nx - nz = 101 - - Re = 30000.0 - prescribed_vel = 0.1 - clength = airfoil_length - - visc = prescribed_vel * clength / Re - omega = 1.0 / (3. * visc + 0.5) - - os.system('rm -rf ./*.vtk && rm -rf ./*.png') - - # Set the parameters for the simulation - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': nz, - 'precision': precision, - 'io_rate': 100, - 'print_info_rate': 100, - } - - sim = Airfoil(**kwargs) - sim.run(20000) diff --git a/examples/CFD/cavity2d.py b/examples/CFD/cavity2d.py deleted file mode 100644 index 28536e41..00000000 --- a/examples/CFD/cavity2d.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -This example implements a 2D Lid-Driven Cavity Flow simulation using the lattice Boltzmann method (LBM). -The Lid-Driven Cavity Flow is a standard test case for numerical schemes applied to fluid dynamics, which involves fluid in a square cavity with a moving lid (top boundary). - -In this example you'll be introduced to the following concepts: - -1. Lattice: The simulation employs a D2Q9 lattice. It's a 2D lattice model with nine discrete velocity directions, which is typically used for 2D simulations. - -2. Boundary Conditions: The code implements two types of boundary conditions: - - BounceBackHalfway: This condition is applied to the stationary walls (left, right, and bottom). It models a no-slip boundary where the velocity of fluid at the wall is zero. - EquilibriumBC: This condition is used for the moving lid (top boundary). It defines a boundary with a set velocity, simulating the "driving" of the cavity by the lid. - -3. Checkpointing: The simulation supports checkpointing. Checkpoints are saved periodically (determined by the 'checkpoint_rate'), allowing the simulation to be stopped and restarted from the last checkpoint. This can be beneficial for long simulations or in case of unexpected interruptions. - -4. Visualization: The simulation outputs data in VTK format for visualization. It also provides images of the velocity field and saves the boundary conditions at each time step. The data can be visualized using software like Paraview. - -""" -from jax import config -import numpy as np -import jax.numpy as jnp -import os - -from src.boundary_conditions import * -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD2Q9 -from src.utils import * - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' - -class Cavity(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - - # concatenate the indices of the left, right, and bottom walls - walls = np.concatenate((self.boundingBoxIndices["left"], self.boundingBoxIndices["right"], self.boundingBoxIndices["bottom"])) - # apply bounce back boundary condition to the walls - self.BCs.append(BounceBackHalfway(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - - # apply inlet equilibrium boundary condition to the top wall - moving_wall = self.boundingBoxIndices["top"] - - rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = prescribed_vel - self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall)) - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs["rho"][1:-1, 1:-1]) - u = np.array(kwargs["u"][1:-1, 1:-1, :]) - timestep = kwargs["timestep"] - - save_image(timestep, u) - fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1]} - save_fields_vtk(timestep, fields) - save_BCs_vtk(timestep, self.BCs, self.gridInfo) - -if __name__ == "__main__": - precision = "f32/f32" - lattice = LatticeD2Q9(precision) - - nx = 200 - ny = 200 - - Re = 200.0 - prescribed_vel = 0.1 - clength = nx - 1 - - checkpoint_rate = 1000 - checkpoint_dir = os.path.abspath("./checkpoints") - - visc = prescribed_vel * clength / Re - omega = 1.0 / (3.0 * visc + 0.5) - - os.system("rm -rf ./*.vtk && rm -rf ./*.png") - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': 0, - 'precision': precision, - 'io_rate': 100, - 'print_info_rate': 100, - 'checkpoint_rate': checkpoint_rate, - 'checkpoint_dir': checkpoint_dir, - 'restore_checkpoint': False, - } - - sim = Cavity(**kwargs) - sim.run(5000) diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py deleted file mode 100644 index 2c30d286..00000000 --- a/examples/CFD/cavity3d.py +++ /dev/null @@ -1,123 +0,0 @@ -""" -This example implements a 3D Lid-Driven Cavity Flow simulation using the lattice Boltzmann method (LBM). -The Lid-Driven Cavity Flow is a standard test case for numerical schemes applied to fluid dynamics, which involves fluid in a square cavity with a moving lid (top boundary). - -In this example you'll be introduced to the following concepts: - -1. Lattice: The simulation employs a D3Q27 lattice. It's a 3D lattice model with 27 discrete velocity directions. - -2. Boundary Conditions: The code implements two types of boundary conditions: - - BounceBack: This condition is applied to the stationary walls, except the top wall. It models a no-slip boundary where the velocity of fluid at the wall is zero. - EquilibriumBC: This condition is used for the moving lid (top boundary). It defines a boundary with a set velocity, simulating the "driving" of the cavity by the lid. - -4. Visualization: The simulation outputs data in VTK format for visualization. The data can be visualized using software like Paraview. - -""" -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' - -import numpy as np -from src.utils import * -from jax import config -import json, codecs - -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD3Q19, LatticeD3Q27 -from src.boundary_conditions import * - - -config.update('jax_enable_x64', True) - -class Cavity(KBCSim): - # Note: We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000 - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # Note: - # We have used halfway BB for Re=(1000, 3200) and regularized BC for Re=10,000 - - # apply inlet boundary condition to the top wall - moving_wall = self.boundingBoxIndices['top'] - vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = prescribed_vel - # self.BCs.append(BounceBackHalfway(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, vel_wall)) - self.BCs.append(Regularized(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall)) - - # concatenate the indices of the left, right, and bottom walls - walls = np.concatenate( - (self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], - self.boundingBoxIndices['front'], self.boundingBoxIndices['back'], - self.boundingBoxIndices['bottom'])) - # apply bounce back boundary condition to the walls - # self.BCs.append(BounceBackHalfway(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - vel_wall = np.zeros(walls.shape, dtype=self.precisionPolicy.compute_dtype) - self.BCs.append(Regularized(tuple(walls.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall)) - return - - def output_data(self, **kwargs): - # 1: -1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs['rho']) - u = np.array(kwargs['u']) - timestep = kwargs['timestep'] - u_prev = kwargs['u_prev'] - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - - err = np.sum(np.abs(u_old - u_new)) - print('error= {:07.6f}'.format(err)) - fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]} - # save_fields_vtk(timestep, fields) - - # output profiles of velocity at mid-plane for benchmarking - output_filename = "./profiles_" + f"{timestep:07d}.json" - ux_mid = 0.5*(u[nx//2, ny//2, :, 0] + u[nx//2+1, ny//2+1, :, 0]) - uz_mid = 0.5*(u[:, ny//2, nz//2, 2] + u[:, ny//2+1, nz//2+1, 2]) - ldc_ref_result = {'ux(x=y=0)': list(ux_mid/prescribed_vel), - 'uz(z=y=0)': list(uz_mid/prescribed_vel)} - json.dump(ldc_ref_result, codecs.open(output_filename, 'w', encoding='utf-8'), - separators=(',', ':'), - sort_keys=True, - indent=4) - - # Calculate the velocity magnitude - # u_mag = np.linalg.norm(u, axis=2) - # live_volume_randering(timestep, u_mag) - -if __name__ == '__main__': - # Note: - # We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000 - precision = 'f64/f64' - lattice = LatticeD3Q27(precision) - - nx = 256 - ny = 256 - nz = 256 - - Re = 10000.0 - prescribed_vel = 0.06 - clength = nx - 2 - - # characteristic time - tc = prescribed_vel/clength - niter_max = int(500//tc) - - visc = prescribed_vel * clength / Re - omega = 1.0 / (3. * visc + 0.5) - os.system("rm -rf ./*.vtk && rm -rf ./*.png") - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': nz, - 'precision': precision, - 'io_rate': int(10//tc), - 'print_info_rate': int(10//tc), - 'downsampling_factor': 1 - } - sim = Cavity(**kwargs) - sim.run(niter_max) \ No newline at end of file diff --git a/examples/CFD/channel3d.py b/examples/CFD/channel3d.py deleted file mode 100644 index 2c7ab739..00000000 --- a/examples/CFD/channel3d.py +++ /dev/null @@ -1,157 +0,0 @@ -""" -This script performs a 3D simulation of turbulent channel flow using the lattice Boltzmann method (LBM). -Turbulent channel flow, also known as plane Couette flow, is a fundamental case in the study of wall-bounded turbulent flows. - -In this example you'll be introduced to the following concepts: - -1. Lattice: A D3Q27 lattice is used, which is a three-dimensional lattice model with 27 discrete velocity directions. This type of lattice allows for a more precise representation of fluid flow in three dimensions. - -2. Initial Conditions: The initial conditions for the flow are randomly generated, and the populations are initialized to be the solution of an advection-diffusion equation. - -3. Boundary Conditions: Bounce back boundary conditions are applied at the top and bottom walls, simulating a no-slip condition typical for wall-bounded flows. - -4. External Force: An external force is applied to drive the flow. - -""" - -from src.boundary_conditions import * -from jax import config -from src.utils import * -import numpy as np -from src.lattice import LatticeD3Q27 -from src.models import KBCSim, AdvectionDiffusionBGK -import jax.numpy as jnp -import os -import matplotlib.pyplot as plt - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -import jax - -# disable JIt compilation - -jax.config.update('jax_enable_x64', True) - -def vonKarman_loglaw_wall(yplus): - vonKarmanConst = 0.41 - cplus = 5.5 - uplus = np.log(yplus)/vonKarmanConst + cplus - return uplus - -def get_dns_data(): - """ - Reference: DNS of Turbulent Channel Flow up to Re_tau=590, 1999, - Physics of Fluids, vol 11, 943-945. - https://turbulence.oden.utexas.edu/data/MKM/chan180/profiles/chan180.means - """ - dns_dic = { - "y":[0,0.000301,0.0012,0.00271,0.00482,0.00752,0.0108,0.0147,0.0192,0.0243,0.03,0.0362,0.0431,0.0505,0.0585,0.067,0.0761,0.0858,0.096,0.107,0.118,0.13,0.142,0.155,0.169,0.182,0.197,0.212,0.227,0.243,0.259,0.276,0.293,0.31,0.328,0.347,0.366,0.385,0.404,0.424,0.444,0.465,0.486,0.507,0.529,0.55,0.572,0.595,0.617,0.64,0.663,0.686,0.71,0.733,0.757,0.781,0.805,0.829,0.853,0.878,0.902,0.926,0.951,0.975,1], - "y+":[0,0.053648,0.21456,0.48263,0.85771,1.3396,1.9279,2.6224,3.4226,4.328,5.3381,6.4523,7.67,8.9902,10.412,11.936,13.559,15.281,17.102,19.019,21.033,23.141,25.342,27.635,30.019,32.492,35.053,37.701,40.432,43.247,46.143,49.118,52.171,55.3,58.503,61.778,65.123,68.536,72.016,75.559,79.164,82.828,86.55,90.327,94.157,98.037,101.97,105.94,109.96,114.02,118.12,122.25,126.42,130.62,134.84,139.1,143.37,147.67,151.99,156.32,160.66,165.02,169.38,173.75,178.12], - "Umean":[0,0.053639,0.21443,0.48197,0.85555,1.3339,1.9148,2.5939,3.3632,4.2095,5.1133,6.0493,6.9892,7.9052,8.7741,9.579,10.311,10.967,11.55,12.066,12.52,12.921,13.276,13.59,13.87,14.121,14.349,14.557,14.75,14.931,15.101,15.264,15.419,15.569,15.714,15.855,15.993,16.128,16.26,16.389,16.515,16.637,16.756,16.872,16.985,17.094,17.2,17.302,17.4,17.494,17.585,17.672,17.756,17.835,17.911,17.981,18.045,18.103,18.154,18.198,18.235,18.264,18.285,18.297,18.301], - "dUmean/dy":[178,178,178,178,177,176,175,173,169,163,155,144,131,116,101,87.1,73.9,62.2,52.2,43.8,36.9,31.1,26.4,22.6,19.4,16.9,14.9,13.3,12,10.9,10.1,9.38,8.79,8.29,7.86,7.49,7.19,6.91,6.63,6.35,6.07,5.81,5.58,5.36,5.14,4.92,4.68,4.45,4.23,4.04,3.85,3.66,3.48,3.28,3.06,2.81,2.54,2.25,1.96,1.67,1.35,1.02,0.673,0.33,0], - "Wmean":[0,0.0000707,0.000283,0.000636,0.00113,0.00176,0.00252,0.00339,0.00435,0.00538,0.00643,0.00751,0.00864,0.00986,0.0112,0.0126,0.0141,0.0156,0.017,0.0181,0.0186,0.0184,0.0176,0.0163,0.0149,0.0135,0.0124,0.0116,0.0107,0.00966,0.00843,0.00695,0.00519,0.00329,0.00145,-0.000284,-0.00177,-0.00292,-0.00377,-0.00445,-0.00497,-0.0054,-0.00594,-0.00681,-0.0082,-0.00996,-0.0119,-0.0139,-0.0163,-0.0191,-0.0225,-0.0263,-0.0306,-0.0354,-0.0405,-0.0455,-0.05,-0.0539,-0.0577,-0.0615,-0.0653,-0.0685,-0.071,-0.0724,-0.0729], - "dWmean/dy":[0.235,0.235,0.235,0.234,0.234,0.232,0.228,0.22,0.208,0.194,0.179,0.168,0.164,0.164,0.166,0.167,0.162,0.148,0.121,0.076,0.0159,-0.0439,-0.087,-0.107,-0.106,-0.0871,-0.0643,-0.0546,-0.061,-0.0707,-0.0818,-0.0958,-0.108,-0.106,-0.0989,-0.0881,-0.0697,-0.0506,-0.0379,-0.0303,-0.0221,-0.0216,-0.0314,-0.0522,-0.0756,-0.0841,-0.0884,-0.0974,-0.114,-0.136,-0.154,-0.172,-0.196,-0.214,-0.215,-0.199,-0.174,-0.156,-0.155,-0.159,-0.147,-0.118,-0.0788,-0.0387,0], - "Pmean":[6.2170e-13,-7.3193e-10,-1.5832e-07,-3.7598e-06,-3.3837e-05,-1.7683e-04,-6.5008e-04,-1.8650e-03,-4.4488e-03,-9.2047e-03,-1.7023e-02,-2.8777e-02,-4.5228e-02,-6.6952e-02,-9.4281e-02,-1.2724e-01,-1.6551e-01,-2.0842e-01,-2.5498e-01,-3.0396e-01,-3.5398e-01,-4.0362e-01,-4.5163e-01,-4.9698e-01,-5.3880e-01,-5.7639e-01,-6.0919e-01,-6.3686e-01,-6.5930e-01,-6.7652e-01,-6.8867e-01,-6.9613e-01,-6.9928e-01,-6.9854e-01,-6.9444e-01,-6.8744e-01,-6.7802e-01,-6.6675e-01,-6.5429e-01,-6.4131e-01,-6.2817e-01,-6.1487e-01,-6.0122e-01,-5.8703e-01,-5.7221e-01,-5.5678e-01,-5.4090e-01,-5.2493e-01,-5.0917e-01,-4.9371e-01,-4.7867e-01,-4.6421e-01,-4.5050e-01,-4.3759e-01,-4.2550e-01,-4.1436e-01,-4.0444e-01,-3.9595e-01,-3.8900e-01,-3.8360e-01,-3.7966e-01,-3.7702e-01,-3.7542e-01,-3.7460e-01,-3.7436e-01] - } - return dns_dic - -class TurbulentChannel(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # top and bottom sides of the channel are no-slip and the other directions are periodic - wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top'])) - self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', np.zeros((wall.shape[0], 3)))) - return - - def initialize_macroscopic_fields(self): - rho = self.precisionPolicy.cast_to_output(1.0) - u = self.distributed_array_init((self.nx, self.ny, self.nz, self.dim), - self.precisionPolicy.compute_dtype, init_val=1e-2 * np.random.random((self.nx, self.ny, self.nz, self.dim))) - u = self.precisionPolicy.cast_to_output(u) - return rho, u - - def initialize_populations(self, rho, u): - omegaADE = 1.0 - lattice = LatticeD3Q27(precision) - - kwargs = {'lattice': lattice, 'nx': self.nx, 'ny': self.ny, 'nz': self.nz, 'precision': precision, 'omega': omegaADE, 'vel': u} - ADE = AdvectionDiffusionBGK(**kwargs) - ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields - print("Initializing the distribution functions using the specified macroscopic fields....") - f = ADE.run(50000) - return f - - def get_force(self): - # define the external force - force = np.zeros((self.nx, self.ny, self.nz, 3)) - force[..., 0] = Re_tau**2 * visc**2 / h**3 - return self.precisionPolicy.cast_to_output(force) - - def output_data(self, **kwargs): - rho = np.array(kwargs["rho"]) - u = np.array(kwargs["u"]) - timestep = kwargs["timestep"] - u_prev = kwargs['u_prev'] - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - - err = np.sum(np.abs(u_old - u_new)) - print("error= {:07.6f}".format(err)) - - # mean streamwise velocity in wall units u^+(z) - uplus = np.mean(u[..., 0], axis=(0,1))/u_tau - uplus_loglaw = vonKarman_loglaw_wall(yplus) - dns_dic = get_dns_data() - plt.clf() - plt.semilogx(yplus, uplus,'r.', yplus, uplus_loglaw, 'k:', dns_dic['y+'], dns_dic['Umean'], 'b-') - ax = plt.gca() - ax.set_xlim([0.1, 300]) - ax.set_ylim([0, 20]) - fname = "uplus_" + str(timestep//10000).zfill(5) + '.pdf' - plt.savefig(fname, format='pdf') - fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]} - save_fields_vtk(timestep, fields) - - - -if __name__ == "__main__": - precision = "f64/f64" - lattice = LatticeD3Q27(precision) - - # h: channel half-width - h = 50 - - # Define channel geometry based on h - nx = 6*h - ny = 3*h - nz = 2*h - - # Define flow regime - Re_tau = 180 - u_tau = 0.001 - DeltaPlus = Re_tau/h # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h - visc = u_tau * h / Re_tau - omega = 1.0 / (3.0 * visc + 0.5) - - # Wall distance in wall units to be used inside output_data - zz = np.arange(nz) - zz = np.minimum(zz, zz.max() - zz) - yplus = zz * u_tau / visc - - os.system("rm -rf ./*.vtk && rm -rf ./*.png") - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': nz, - 'precision': precision, - 'io_rate': 500000, - 'print_info_rate': 100000 - } - sim = turbulentChannel(**kwargs) - sim.run(10000000) diff --git a/examples/CFD/couette2d.py b/examples/CFD/couette2d.py deleted file mode 100644 index 1c15a6a8..00000000 --- a/examples/CFD/couette2d.py +++ /dev/null @@ -1,79 +0,0 @@ -""" -This script performs a 2D simulation of Couette flow using the lattice Boltzmann method (LBM). -""" - -import os -import jax.numpy as jnp -import numpy as np -from src.utils import * -from jax import config - - -from src.models import BGKSim -from src.boundary_conditions import * -from src.lattice import LatticeD2Q9 - -# config.update('jax_disable_jit', True) -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=4' - -class Couette(BGKSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - walls = np.concatenate((self.boundingBoxIndices["top"], self.boundingBoxIndices["bottom"])) - self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - - outlet = self.boundingBoxIndices["right"] - inlet = self.boundingBoxIndices["left"] - - rho_wall = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_wall = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = prescribed_vel - self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall)) - - self.BCs.append(DoNothing(tuple(outlet.T), self.gridInfo, self.precisionPolicy)) - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs["rho"][..., 1:-1, :]) - u = np.array(kwargs["u"][..., 1:-1, :]) - timestep = kwargs["timestep"] - u_prev = kwargs["u_prev"][..., 1:-1, :] - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - err = np.sum(np.abs(u_old - u_new)) - print("error= {:07.6f}".format(err)) - save_image(timestep, u) - fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1]} - save_fields_vtk(timestep, fields) - -if __name__ == "__main__": - precision = "f32/f32" - lattice = LatticeD2Q9(precision) - nx = 501 - ny = 101 - - Re = 100.0 - prescribed_vel = 0.1 - clength = nx - 1 - - visc = prescribed_vel * clength / Re - - omega = 1.0 / (3.0 * visc + 0.5) - assert omega < 1.98, "omega must be less than 2.0" - os.system("rm -rf ./*.vtk && rm -rf ./*.png") - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': 0, - 'precision': precision, - 'io_rate': 100, - 'print_info_rate': 100 - } - sim = Couette(**kwargs) - sim.run(20000) diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py deleted file mode 100644 index 2c9887d5..00000000 --- a/examples/CFD/cylinder2d.py +++ /dev/null @@ -1,148 +0,0 @@ -""" -This script conducts a 2D simulation of flow around a cylinder using the lattice Boltzmann method (LBM). This is a classic problem in fluid dynamics and is often used to examine the behavior of fluid flow over a bluff body. - -In this example you'll be introduced to the following concepts: - -1. Lattice: A D2Q9 lattice is used, which is a two-dimensional lattice model with nine discrete velocity directions. This type of lattice allows for a precise representation of fluid flow in two dimensions. - -2. Boundary Conditions: The script implements several types of boundary conditions: - - BounceBackHalfway: This condition is applied to the cylinder surface, simulating a no-slip condition where the fluid at the cylinder surface has zero velocity. - ExtrapolationOutflow: This condition is applied at the outlet (right boundary), where the fluid is allowed to exit the simulation domain freely. - Regularized: This condition is applied at the inlet (left boundary) and models the inflow of fluid into the domain with a specified velocity profile. Another Regularized condition is used for the stationary top and bottom walls. -3. Velocity Profile: The script uses a Poiseuille flow profile for the inlet velocity. This is a parabolic profile commonly seen in pipe flow. - -4. Drag and lift calculation: The script computes the lift and drag on the cylinder, which are important quantities in fluid dynamics and aerodynamics. - -5. Visualization: The simulation outputs data in VTK format for visualization. It also generates images of the velocity field. The data can be visualized using software like ParaView. - -# To run type: -nohup python3 examples/CFD/cylinder2d.py > logfile.log & -""" -import os -import json -import jax -from time import time -from jax import config -import numpy as np -import jax.numpy as jnp - -from src.utils import * -from src.boundary_conditions import * -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD2Q9 - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -jax.config.update('jax_enable_x64', True) - -class Cylinder(BGKSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # Define the cylinder surface - coord = np.array([(i, j) for i in range(self.nx) for j in range(self.ny)]) - xx, yy = coord[:, 0], coord[:, 1] - cx, cy = 2.*diam, 2.*diam - cylinder = (xx - cx)**2 + (yy-cy)**2 <= (diam/2.)**2 - cylinder = coord[cylinder] - implicit_distance = np.reshape((xx - cx)**2 + (yy-cy)**2 - (diam/2.)**2, (self.nx, self.ny)) - self.BCs.append(InterpolatedBounceBackBouzidi(tuple(cylinder.T), implicit_distance, self.gridInfo, self.precisionPolicy)) - - # Outflow BC - outlet = self.boundingBoxIndices['right'] - rho_outlet = np.ones((outlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy)) - # self.BCs.append(ZouHe(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet)) - - # Inlet BC - inlet = self.boundingBoxIndices['left'] - rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype) - yy_inlet = yy.reshape(self.nx, self.ny)[tuple(inlet.T)] - vel_inlet[:, 0] = poiseuille_profile(yy_inlet, - yy_inlet.min(), - yy_inlet.max()-yy_inlet.min(), 3.0 / 2.0 * prescribed_vel) - self.BCs.append(Regularized(tuple(inlet.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_inlet)) - - # No-slip BC for top and bottom - wall = np.concatenate([self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']]) - vel_wall = np.zeros(wall.shape, dtype=self.precisionPolicy.compute_dtype) - self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall)) - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using bounce-back) - rho = np.array(kwargs["rho"][..., 1:-1, :]) - u = np.array(kwargs["u"][..., 1:-1, :]) - timestep = kwargs["timestep"] - u_prev = kwargs["u_prev"][..., 1:-1, :] - - if timestep == 0: - self.CL_max = 0.0 - self.CD_max = 0.0 - if timestep > 0.5 * niter_max: - # compute lift and drag over the cyliner - cylinder = self.BCs[0] - boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision']) - boundary_force = np.sum(np.array(boundary_force), axis=0) - drag = boundary_force[0] - lift = boundary_force[1] - cd = 2. * drag / (prescribed_vel ** 2 * diam) - cl = 2. * lift / (prescribed_vel ** 2 * diam) - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - err = np.sum(np.abs(u_old - u_new)) - self.CL_max = max(self.CL_max, cl) - self.CD_max = max(self.CD_max, cd) - print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd)) - # save_image(timestep, u) - -# Helper function to specify a parabolic poiseuille profile -poiseuille_profile = lambda x,x0,d,umax: np.maximum(0.,4.*umax/(d**2)*((x-x0)*d-(x-x0)**2)) - -if __name__ == '__main__': - precision = 'f64/f64' - # diam_list = [10, 20, 30, 40, 60, 80] - diam_list = [80] - CL_list, CD_list = [], [] - result_dict = {} - result_dict['resolution_list'] = diam_list - for diam in diam_list: - scale_factor = 80 / diam - prescribed_vel = 0.003 * scale_factor - lattice = LatticeD2Q9(precision) - - nx = int(22*diam) - ny = int(4.1*diam) - - Re = 100.0 - visc = prescribed_vel * diam / Re - omega = 1.0 / (3. * visc + 0.5) - - os.system('rm -rf ./*.vtk && rm -rf ./*.png') - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': 0, - 'precision': precision, - 'io_rate': int(500 / scale_factor), - 'print_info_rate': int(10000 / scale_factor), - 'return_fpost': True # Need to retain fpost-collision for computation of lift and drag - } - # characteristic time - tc = prescribed_vel/diam - niter_max = int(100//tc) - sim = Cylinder(**kwargs) - sim.run(niter_max) - CL_list.append(sim.CL_max) - CD_list.append(sim.CD_max) - - result_dict['CL'] = CL_list - result_dict['CD'] = CD_list - with open('data.json', 'w') as fp: - json.dump(result_dict, fp) diff --git a/examples/CFD/oscilating_cylinder2d.py b/examples/CFD/oscilating_cylinder2d.py deleted file mode 100644 index 97d6746b..00000000 --- a/examples/CFD/oscilating_cylinder2d.py +++ /dev/null @@ -1,146 +0,0 @@ -""" -This script conducts a 2D simulation of flow around a cylinder using the lattice Boltzmann method (LBM). This is a classic problem in fluid dynamics and is often used to examine the behavior of fluid flow over a bluff body. - -In this example you'll be introduced to the following concepts: - -1. Lattice: A D2Q9 lattice is used, which is a two-dimensional lattice model with nine discrete velocity directions. This type of lattice allows for a precise representation of fluid flow in two dimensions. - -2. Boundary Conditions: The script implements several types of boundary conditions: - - BounceBackMoving: This condition is applied to the cylinder surface. Unlike the usual BounceBack condition, this one takes into account the motion of the cylinder. - ExtrapolationOutflow: This condition is applied at the outlet (right boundary), where the fluid is allowed to exit the simulation domain freely. - Regularized: This condition is applied at the inlet (left boundary) and models the inflow of fluid into the domain with a specified velocity profile. Another Regularized condition is used for the stationary top and bottom walls. -3. Velocity Profile: The script uses a Poiseuille flow profile for the inlet velocity. This is a parabolic profile commonly seen in pipe flow. - -4. Drag and lift calculation: The script computes the lift and drag on the cylinder, which are important quantities in fluid dynamics and aerodynamics. - -5. Visualization: The simulation outputs data in VTK format for visualization. It also generates images of the velocity field. The data can be visualized using software like ParaView. - -""" - - -import os -import jax -from time import time -from jax import config -import numpy as np -import jax.numpy as jnp - -from src.utils import * -from src.boundary_conditions import * -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD2Q9 - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -jax.config.update('jax_enable_x64', True) - -class Cylinder(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - wall = np.concatenate([self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']]) - self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy)) - - coord = np.array([np.unravel_index(i, (self.nx, self.ny)) for i in range(self.nx*self.ny)]) - xx, yy = coord[:, 0], coord[:, 1] - cx, cy = 2.*diam, 2.*diam - cyl = ((xx) - cx)**2 + (yy-cy)**2 <= (diam/2.)**2 - cyl = jnp.array(coord[cyl]) - - # Define update rules for boundary conditions - def update_function(time: int): - # Move the cylinder up and down sinusoidally with time - # Define the scale for the sinusoidal motion - scale = 10000 - - # Amplitude of the motion, a quarter of the y-dimension of the grid - A = ny // 4 - - # Calculate the new y-coordinates of the cylinder. The cylinder moves up and down, - # its motion dictated by the sinusoidal function. We use `astype(int)` to ensure - # the indices are integers, as they will be used for array indexing. - new_y_coords = cyl[:, 1] + jnp.array((jnp.sin(time/scale)*A).astype(int)) - - # Define the indices of the grid points occupied by the cylinder - indices = (cyl[:, 0], new_y_coords) - - # Calculate the velocity of the cylinder. The x-component is always 0 (the cylinder - # doesn't move horizontally), and the y-component is the derivative of the sinusoidal - # function governing the cylinder's motion, scaled by the amplitude and the scale factor. - velocity = jnp.array([0., jnp.cos(time/scale)* A / scale], dtype=self.precisionPolicy.compute_dtype) - - return indices, velocity - - self.BCs.append(BounceBackMoving(self.gridInfo, self.precisionPolicy, update_function=update_function)) - - - outlet = self.boundingBoxIndices['right'] - self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy)) - - inlet = self.boundingBoxIndices['left'] - vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype) - yy_inlet = yy.reshape(self.nx, self.ny)[tuple(inlet.T)] - vel_inlet[:, 0] = poiseuille_profile(yy_inlet, - yy_inlet.min(), - yy_inlet.max()-yy_inlet.min(), 3.0 / 2.0 * prescribed_vel) - self.BCs.append(Regularized(tuple(inlet.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_inlet)) - - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs["rho"][..., 1:-1, :]) - u = np.array(kwargs["u"][..., 1:-1, :]) - timestep = kwargs["timestep"] - u_prev = kwargs["u_prev"][..., 1:-1, :] - - # compute lift and drag over the cyliner - cylinder = self.BCs[0] - boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision']) - boundary_force = np.sum(boundary_force, axis=0) - drag = boundary_force[0] - lift = boundary_force[1] - cd = 2. * drag / (prescribed_vel ** 2 * diam) - cl = 2. * lift / (prescribed_vel ** 2 * diam) - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - err = np.sum(np.abs(u_old - u_new)) - print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd)) - save_image(timestep, u) - # u magnitude - fields = {'rho': rho[..., 0], 'u': np.linalg.norm(u, axis=2)} - save_fields_vtk(timestep, fields) - save_BCs_vtk(timestep, self.BCs, self.gridInfo) - -# Helper function to specify a parabolic poiseuille profile -poiseuille_profile = lambda x,x0,d,umax: np.maximum(0.,4.*umax/(d**2)*((x-x0)*d-(x-x0)**2)) - -if __name__ == '__main__': - precision = 'f64/f64' - lattice = LatticeD2Q9(precision) - prescribed_vel = 0.005 - diam = 20 - nx = int(22*diam) - ny = int(4.1*diam) - - Re = 10.0 - visc = prescribed_vel * diam / Re - omega = 1.0 / (3. * visc + 0.5) - - os.system('rm -rf ./*.vtk && rm -rf ./*.png') - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': 0, - 'precision': precision, - 'io_rate': 500, - 'print_info_rate': 500, - 'return_fpost': True, # Need to retain fpost-collision for computation of lift and drag - } - sim = Cylinder(**kwargs) - - sim.run(1000000) diff --git a/examples/CFD/taylor_green_vortex.py b/examples/CFD/taylor_green_vortex.py deleted file mode 100644 index 374c4993..00000000 --- a/examples/CFD/taylor_green_vortex.py +++ /dev/null @@ -1,127 +0,0 @@ -""" -The given script sets up a simulation for the Taylor-Green vortex flow. -The Taylor-Green vortex is a type of two-dimensional, incompressible fluid flow with a known analytical solution, making it an ideal test case for fluid dynamics simulations. -The flow is characterized by a pair of counter-rotating vortices. In this script, the initial fields for the Taylor-Green vortex are set using a known function. -""" - - -import os -import json -import jax -import numpy as np -import matplotlib.pyplot as plt - -from src.utils import * -from src.boundary_conditions import * -from src.models import BGKSim, KBCSim, AdvectionDiffusionBGK -from src.lattice import LatticeD2Q9 - - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -# disable JIt compilation - -jax.config.update('jax_enable_x64', True) - -def taylor_green_initial_fields(xx, yy, u0, rho0, nu, time): - ux = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time) - uy = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time) - rho = 1.0 - rho0 * u0 ** 2 / 12. * (np.cos(2. * xx) + np.cos(2. * yy)) * np.exp(-4 * nu * time) - return ux, uy, np.expand_dims(rho, axis=-1) - -class TaylorGreenVortex(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # no boundary conditions implying periodic BC in all directions - return - - def initialize_macroscopic_fields(self): - ux, uy, rho = taylor_green_initial_fields(xx, yy, vel_ref, 1, 0., 0.) - rho = self.distributed_array_init(rho.shape, self.precisionPolicy.output_dtype, init_val=1.0, sharding=self.sharding) - u = np.stack([ux, uy], axis=-1) - u = self.distributed_array_init(u.shape, self.precisionPolicy.output_dtype, init_val=u, sharding=self.sharding) - return rho, u - - def initialize_populations(self, rho, u): - omegaADE = 1.0 - kwargs = {'lattice': lattice, 'nx': self.nx, 'ny': self.ny, 'nz': self.nz, 'precision': precision, 'omega': omegaADE, 'vel': u, 'print_info_rate': 0, 'io_rate': 0} - ADE = AdvectionDiffusionBGK(**kwargs) - ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields - print("Initializing the distribution functions using the specified macroscopic fields....") - f = ADE.run(int(20000*nx/32)) - return f - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs["rho"]) - u = np.array(kwargs["u"]) - timestep = kwargs["timestep"] - - # theoretical results - time = timestep * (kx**2 + ky**2)/2. - ux_th, uy_th, rho_th = taylor_green_initial_fields(xx, yy, vel_ref, 1, visc, time) - vel_err_L2 = np.sqrt(np.sum((u[..., 0]-ux_th)**2 + (u[..., 1]-uy_th)**2) / np.sum(ux_th**2 + uy_th**2)) - rho_err_L2 = np.sqrt(np.sum((rho - rho_th)**2) / np.sum(rho_th**2)) - print("Vel error= {:07.6f}, Pressure error= {:07.6f}".format(vel_err_L2, rho_err_L2)) - if timestep == endTime: - ErrL2ResList.append(vel_err_L2) - ErrL2ResListRho.append(rho_err_L2) - # save_image(timestep, u) - - -if __name__ == "__main__": - precision_list = ["f32/f32", "f64/f32", "f64/f64"] - resList = [32, 64, 128, 256, 512, 1024] - result_dict = dict.fromkeys(precision_list) - result_dict['resolution_list'] = resList - - for precision in precision_list: - lattice = LatticeD2Q9(precision) - ErrL2ResList = [] - ErrL2ResListRho = [] - result_dict[precision] = dict.fromkeys(['vel_error', 'rho_error']) - for nx in resList: - ny = nx - twopi = 2.0 * np.pi - coord = np.array([(i, j) for i in range(nx) for j in range(ny)]) - xx, yy = coord[:, 0], coord[:, 1] - kx, ky = twopi / nx, twopi / ny - xx = xx.reshape((nx, ny)) * kx - yy = yy.reshape((nx, ny)) * ky - - Re = 1600.0 - vel_ref = 0.04*32/nx - - visc = vel_ref * nx / Re - omega = 1.0 / (3.0 * visc + 0.5) - os.system("rm -rf ./*.vtk && rm -rf ./*.png") - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': 0, - 'precision': precision, - 'io_rate': 5000, - 'print_info_rate': 1000 - } - sim = TaylorGreenVortex(**kwargs) - tc = 2.0/(2. * visc * (kx**2 + ky**2)) - endTime = int(0.05*tc) - sim.run(endTime) - result_dict[precision]['vel_error'] = ErrL2ResList - result_dict[precision]['rho_error'] = ErrL2ResListRho - - with open('data.json', 'w') as fp: - json.dump(result_dict, fp) - - # plt.loglog(resList, ErrL2ResList, '-o') - # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--') - # plt.savefig('ErrorVel.png'); plt.savefig('ErrorVel.pdf', format='pdf') - - # plt.figure() - # plt.loglog(resList, ErrL2ResListRho, '-o') - # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--') - # plt.savefig('ErrorRho.png'); plt.savefig('ErrorRho.pdf', format='pdf') diff --git a/examples/CFD/windtunnel3d.py b/examples/CFD/windtunnel3d.py deleted file mode 100644 index 081f9b53..00000000 --- a/examples/CFD/windtunnel3d.py +++ /dev/null @@ -1,140 +0,0 @@ -""" -This script performs a Lattice Boltzmann Method (LBM) simulation of fluid flow over a car model. Here are the main concepts and steps in the simulation: - -Here are the main concepts introduced simulation: - -1. Lattice: Given the usually high Reynolds number required for these simulations, a D3Q27 lattice is used, which is a three-dimensional lattice model with 27 discrete velocity directions. - -2. Loading geometry and voxelization: The geometry of the car is loaded from a STL file. -This is a file format commonly used for 3D models. The model is then voxelized to a binary matrix which represents the presence or absence of the object in the lattice. We use the DrivAer model, which is a common car model used for aerodynamic simulations. - -3. Output: After each specified number of iterations, the script outputs the state of the simulation. This includes the error (difference between consecutive velocity fields), lift and drag coefficients, and visualization files in the VTK format. -""" - - -import os -import jax -import trimesh -from time import time -import numpy as np -import jax.numpy as jnp -from jax import config - -from src.utils import * -from src.models import BGKSim, KBCSim -from src.lattice import LatticeD3Q19, LatticeD3Q27 -from src.boundary_conditions import * - -# Use 8 CPU devices -# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' - -# disable JIt compilation - -jax.config.update('jax_array', True) - -class Car(KBCSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def voxelize_stl(self, stl_filename, length_lbm_unit): - mesh = trimesh.load_mesh(stl_filename, process=False) - length_phys_unit = mesh.extents.max() - pitch = length_phys_unit/length_lbm_unit - mesh_voxelized = mesh.voxelized(pitch=pitch) - mesh_matrix = mesh_voxelized.matrix - return mesh_matrix, pitch - - def set_boundary_conditions(self): - print('Voxelizing mesh...') - time_start = time() - stl_filename = 'stl-files/DrivAer-Notchback.stl' - car_length_lbm_unit = self.nx / 4 - car_voxelized, pitch = voxelize_stl(stl_filename, car_length_lbm_unit) - car_matrix = car_voxelized.matrix - print('Voxelization time for pitch={}: {} seconds'.format(pitch, time() - time_start)) - print("Car matrix shape: ", car_matrix.shape) - - self.car_area = np.prod(car_matrix.shape[1:]) - tx, ty, tz = np.array([nx, ny, nz]) - car_matrix.shape - shift = [tx//4, ty//2, 0] - car_indices = np.argwhere(car_matrix) + shift - self.BCs.append(BounceBackHalfway(tuple(car_indices.T), self.gridInfo, self.precisionPolicy)) - - wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top'], - self.boundingBoxIndices['front'], self.boundingBoxIndices['back'])) - self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy)) - - doNothing = self.boundingBoxIndices['right'] - self.BCs.append(DoNothing(tuple(doNothing.T), self.gridInfo, self.precisionPolicy)) - self.BCs[-1].implementationStep = 'PostCollision' - # rho_outlet = np.ones(doNothing.shape[0], dtype=self.precisionPolicy.compute_dtype) - # self.BCs.append(ZouHe(tuple(doNothing.T), - # self.gridInfo, - # self.precisionPolicy, - # 'pressure', rho_outlet)) - - inlet = self.boundingBoxIndices['left'] - rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype) - - vel_inlet[:, 0] = prescribed_vel - self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_inlet, vel_inlet)) - # self.BCs.append(ZouHe(tuple(inlet.T), - # self.gridInfo, - # self.precisionPolicy, - # 'velocity', vel_inlet)) - - def output_data(self, **kwargs): - # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back) - rho = np.array(kwargs['rho'][..., 1:-1, 1:-1, :]) - u = np.array(kwargs['u'][..., 1:-1, 1:-1, :]) - timestep = kwargs['timestep'] - u_prev = kwargs['u_prev'][..., 1:-1, 1:-1, :] - - # compute lift and drag over the car - car = self.BCs[0] - boundary_force = car.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision']) - boundary_force = np.sum(boundary_force, axis=0) - drag = np.sqrt(boundary_force[0]**2 + boundary_force[1]**2) #xy-plane - lift = boundary_force[2] #z-direction - cd = 2. * drag / (prescribed_vel ** 2 * self.car_area) - cl = 2. * lift / (prescribed_vel ** 2 * self.car_area) - - u_old = np.linalg.norm(u_prev, axis=2) - u_new = np.linalg.norm(u, axis=2) - - err = np.sum(np.abs(u_old - u_new)) - print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd)) - fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]} - save_fields_vtk(timestep, fields) - -if __name__ == '__main__': - precision = 'f32/f32' - lattice = LatticeD3Q27(precision) - - nx = 601 - ny = 351 - nz = 251 - - Re = 50000.0 - prescribed_vel = 0.05 - clength = nx - 1 - - visc = prescribed_vel * clength / Re - omega = 1.0 / (3. * visc + 0.5) - - os.system('rm -rf ./*.vtk && rm -rf ./*.png') - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': nx, - 'ny': ny, - 'nz': nz, - 'precision': precision, - 'io_rate': 100, - 'print_info_rate': 100, - 'return_fpost': True # Need to retain fpost-collision for computation of lift and drag - } - sim = Car(**kwargs) - sim.run(200000) diff --git a/examples/cfd/data/ahmed.json b/examples/cfd/data/ahmed.json new file mode 100644 index 00000000..6253bc70 --- /dev/null +++ b/examples/cfd/data/ahmed.json @@ -0,0 +1,22 @@ +{ + "_comment": "Ahmed Car Model, slant - angle = 25 degree. Profiles on symmetry plane (y=0) covering entire field. Origin of coordinate system: x=0: end of the car, y=0: symmetry plane, z=0: ground plane S.Becker/H. Lienhart/C Stoots, Institute of Fluid Mechanics, University Erlangen-Nuremberg, Erlangen, Germany, Coordinates in meters need to convert to voxels, Velocity data in m/s", + "data": { + "-1.162" : { "x-velocity" : [26.995,29.825,29.182,28.488,27.703,26.988,26.456,26.163,26.190,26.523,27.083,28.033,29.131,30.429,31.747,33.036,34.268,35.354,36.312,37.083,37.770,38.484,39.033,39.447,39.839,40.086,40.268,40.380,40.451], "height" : [0.028,0.048,0.068,0.088,0.108,0.128,0.148,0.168,0.188,0.208,0.228,0.248,0.268,0.288,0.308,0.328,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.7388]}, + "-1.062" : { "x-velocity" : [30.307,28.962,25.812,21.232,15.848,10.812,7.459,6.080,5.845,6.196,7.428,10.456,15.718,22.129,28.090,32.707,35.888,37.891,39.071,39.840,40.261,40.604,40.767,40.820,40.870,40.890,40.907,40.871,40.853], "height" : [0.028,0.048,0.068,0.088,0.108,0.128,0.148,0.168,0.188,0.208,0.228,0.248,0.268,0.288,0.308,0.328,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.962" : { "x-velocity" : [52.216,51.303,50.196,48.833,47.728,46.790,45.514,44.222,43.379,42.829,42.322,42.056,41.876,41.706,41.584], "height" : [0.363,0.368,0.378,0.388,0.398,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.862" : { "x-velocity" : [46.589,46.538,46.228,46.033,45.810,45.554,45.056,44.369,43.789,43.275,42.789,42.344,42.148,41.913,41.720], "height" : [0.363,0.368,0.378,0.388,0.398,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.562" : { "x-velocity" : [43.237,43.262,43.248,43.225,43.183,43.145,43.083,43.030,42.904,42.776,42.685,42.434,42.358,42.197,42.042], "height" : [0.363,0.368,0.378,0.388,0.398,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.362" : { "x-velocity" : [44.493,44.491,44.443,44.379,44.297,44.215,44.067,43.867,43.577,43.306,43.061,42.689,42.527,42.293,42.105], "height" : [0.363,0.368,0.378,0.388,0.398,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.212" : { "x-velocity" : [49.202,48.429,47.805,46.697,45.883,44.913,44.195,43.650,43.130,42.677,42.432,42.154,41.961], "height" : [0.368,0.378,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.162" : { "x-velocity" : [50.511,49.784,48.894,48.103,47.468,46.322,45.563,44.581,43.933,43.383,42.905,42.505,42.293,42.042,41.863], "height" : [0.348,0.358,0.368,0.378,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.112" : { "x-velocity" : [27.615,35.449,41.526,46.068,46.277,46.038,45.774,45.505,45.237,44.701,44.326,43.765,43.284,42.890,42.529,42.247,42.082,41.880,41.732], "height" : [0.318,0.323,0.328,0.338,0.348,0.358,0.368,0.378,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.062" : { "x-velocity" : [22.891,27.789,32.292,36.568,39.533,41.426,42.371,42.971,43.030,43.081,43.074,43.065,43.039,42.996,42.908,42.665,42.456,42.294,42.105,41.929,41.827,41.660,41.546], "height" : [0.298,0.303,0.308,0.313,0.318,0.323,0.328,0.338,0.348,0.358,0.368,0.378,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "-0.012" : { "x-velocity" : [23.304,26.317,29.429,32.341,34.923,37.106,38.673,39.841,40.447,40.780,40.973,41.085,41.193,41.282,41.359,41.442,41.522,41.699,41.737,41.749,41.724,41.714,41.642,41.574,41.518,41.431,41.366], "height" : [0.278,0.283,0.288,0.293,0.298,0.303,0.308,0.313,0.318,0.323,0.328,0.338,0.348,0.358,0.368,0.378,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.038" : { "x-velocity" : [42.752,37.392,15.320,-4.501,-8.079,-8.892,-8.420,-7.027,-5.143,-2.903,-0.936,0.927,2.200,3.099,3.622,4.026,4.280,4.520,5.620,8.938,13.913,17.872,21.148,24.814,29.075,33.188,36.424,38.490,39.388,39.675,39.794,39.911,40.007,40.219,40.425,40.643,40.757,40.896,40.994,41.058,41.124,41.127,41.143,41.106,41.080], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.088" : { "x-velocity" : [41.859,35.830,22.660,7.745,-5.808,-12.650,-14.748,-13.756,-10.659,-6.484,-2.121,1.303,3.672,5.441,7.066,9.157,11.613,14.620,17.662,20.639,23.565,26.437,29.484,32.441,35.024,36.938,37.938,38.377,38.595,38.728,38.856,38.976,39.133,39.438,39.749,39.975,40.129,40.344,40.499,40.649,40.783,40.853,40.927,40.945,40.960], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.138" : { "x-velocity" : [36.223,32.501,24.752,14.281,2.799,-6.218,-10.908,-11.892,-9.708,-5.258,-0.140,4.331,7.882,10.995,13.961,16.699,19.477,22.063,24.651,27.081,29.524,31.950,34.043,35.594,36.506,37.053,37.386,37.614,37.832,38.032,38.214,38.397,38.575,38.940,39.298,39.533,39.749,40.028,40.206,40.404,40.580,40.691,40.803,40.858,40.921], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.188" : { "x-velocity" : [29.417,27.755,23.967,18.261,11.662,5.405,0.676,-0.652,0.937,4.261,7.958,11.427,14.366,17.138,19.735,22.151,24.577,26.883,29.165,31.111,32.781,34.072,34.893,35.524,35.974,36.329,36.604,36.872,37.138,37.402,37.673,37.900,38.112,38.518,38.829,39.088,39.326,39.639,39.871,40.096,40.275,40.423,40.523,40.603,40.687], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.238" : { "x-velocity" : [24.405,24.168,22.782,20.196,16.970,13.937,12.137,11.757,12.851,14.649,16.780,18.995,21.070,23.335,25.280,27.468,29.262,30.832,32.133,33.102,33.856,34.473,34.922,35.340,35.698,36.039,36.336,36.629,36.906,37.193,37.454,37.691,37.929,38.329,38.611,38.875,39.126,39.414,39.677,39.917,40.097,40.259,40.380,40.478,40.568], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]}, + "0.288" : { "x-velocity" : [21.489,22.225,22.127,21.456,20.404,19.743,19.541,19.909,21.002,22.381,24.018,25.670,27.421,28.998,30.371,31.523,32.406,33.111,33.670,34.155,34.532,34.893,35.240,35.567,35.875,36.158,36.437,36.708,36.974,37.230,37.473,37.709,37.932,38.266,38.515,38.773,39.008,39.270,39.562,39.782,39.962,40.148,40.266,40.369,40.475], "height" : [0.028,0.038,0.048,0.058,0.068,0.078,0.088,0.098,0.108,0.118,0.128,0.138,0.148,0.158,0.168,0.178,0.188,0.198,0.208,0.218,0.228,0.238,0.248,0.258,0.268,0.278,0.288,0.298,0.308,0.318,0.328,0.338,0.348,0.368,0.388,0.408,0.428,0.458,0.488,0.518,0.558,0.598,0.638,0.688,0.738]} + } +} \ No newline at end of file diff --git a/examples/cfd/data/turbulent_channel_dns_data.json b/examples/cfd/data/turbulent_channel_dns_data.json new file mode 100644 index 00000000..bdbee757 --- /dev/null +++ b/examples/cfd/data/turbulent_channel_dns_data.json @@ -0,0 +1 @@ +{"y": [0, 0.000301, 0.0012, 0.00271, 0.00482, 0.00752, 0.0108, 0.0147, 0.0192, 0.0243, 0.03, 0.0362, 0.0431, 0.0505, 0.0585, 0.067, 0.0761, 0.0858, 0.096, 0.107, 0.118, 0.13, 0.142, 0.155, 0.169, 0.182, 0.197, 0.212, 0.227, 0.243, 0.259, 0.276, 0.293, 0.31, 0.328, 0.347, 0.366, 0.385, 0.404, 0.424, 0.444, 0.465, 0.486, 0.507, 0.529, 0.55, 0.572, 0.595, 0.617, 0.64, 0.663, 0.686, 0.71, 0.733, 0.757, 0.781, 0.805, 0.829, 0.853, 0.878, 0.902, 0.926, 0.951, 0.975, 1], "y+": [0, 0.053648, 0.21456, 0.48263, 0.85771, 1.3396, 1.9279, 2.6224, 3.4226, 4.328, 5.3381, 6.4523, 7.67, 8.9902, 10.412, 11.936, 13.559, 15.281, 17.102, 19.019, 21.033, 23.141, 25.342, 27.635, 30.019, 32.492, 35.053, 37.701, 40.432, 43.247, 46.143, 49.118, 52.171, 55.3, 58.503, 61.778, 65.123, 68.536, 72.016, 75.559, 79.164, 82.828, 86.55, 90.327, 94.157, 98.037, 101.97, 105.94, 109.96, 114.02, 118.12, 122.25, 126.42, 130.62, 134.84, 139.1, 143.37, 147.67, 151.99, 156.32, 160.66, 165.02, 169.38, 173.75, 178.12], "Umean": [0, 0.053639, 0.21443, 0.48197, 0.85555, 1.3339, 1.9148, 2.5939, 3.3632, 4.2095, 5.1133, 6.0493, 6.9892, 7.9052, 8.7741, 9.579, 10.311, 10.967, 11.55, 12.066, 12.52, 12.921, 13.276, 13.59, 13.87, 14.121, 14.349, 14.557, 14.75, 14.931, 15.101, 15.264, 15.419, 15.569, 15.714, 15.855, 15.993, 16.128, 16.26, 16.389, 16.515, 16.637, 16.756, 16.872, 16.985, 17.094, 17.2, 17.302, 17.4, 17.494, 17.585, 17.672, 17.756, 17.835, 17.911, 17.981, 18.045, 18.103, 18.154, 18.198, 18.235, 18.264, 18.285, 18.297, 18.301], "dUmean/dy": [178, 178, 178, 178, 177, 176, 175, 173, 169, 163, 155, 144, 131, 116, 101, 87.1, 73.9, 62.2, 52.2, 43.8, 36.9, 31.1, 26.4, 22.6, 19.4, 16.9, 14.9, 13.3, 12, 10.9, 10.1, 9.38, 8.79, 8.29, 7.86, 7.49, 7.19, 6.91, 6.63, 6.35, 6.07, 5.81, 5.58, 5.36, 5.14, 4.92, 4.68, 4.45, 4.23, 4.04, 3.85, 3.66, 3.48, 3.28, 3.06, 2.81, 2.54, 2.25, 1.96, 1.67, 1.35, 1.02, 0.673, 0.33, 0], "Wmean": [0, 7.07e-05, 0.000283, 0.000636, 0.00113, 0.00176, 0.00252, 0.00339, 0.00435, 0.00538, 0.00643, 0.00751, 0.00864, 0.00986, 0.0112, 0.0126, 0.0141, 0.0156, 0.017, 0.0181, 0.0186, 0.0184, 0.0176, 0.0163, 0.0149, 0.0135, 0.0124, 0.0116, 0.0107, 0.00966, 0.00843, 0.00695, 0.00519, 0.00329, 0.00145, -0.000284, -0.00177, -0.00292, -0.00377, -0.00445, -0.00497, -0.0054, -0.00594, -0.00681, -0.0082, -0.00996, -0.0119, -0.0139, -0.0163, -0.0191, -0.0225, -0.0263, -0.0306, -0.0354, -0.0405, -0.0455, -0.05, -0.0539, -0.0577, -0.0615, -0.0653, -0.0685, -0.071, -0.0724, -0.0729], "dWmean/dy": [0.235, 0.235, 0.235, 0.234, 0.234, 0.232, 0.228, 0.22, 0.208, 0.194, 0.179, 0.168, 0.164, 0.164, 0.166, 0.167, 0.162, 0.148, 0.121, 0.076, 0.0159, -0.0439, -0.087, -0.107, -0.106, -0.0871, -0.0643, -0.0546, -0.061, -0.0707, -0.0818, -0.0958, -0.108, -0.106, -0.0989, -0.0881, -0.0697, -0.0506, -0.0379, -0.0303, -0.0221, -0.0216, -0.0314, -0.0522, -0.0756, -0.0841, -0.0884, -0.0974, -0.114, -0.136, -0.154, -0.172, -0.196, -0.214, -0.215, -0.199, -0.174, -0.156, -0.155, -0.159, -0.147, -0.118, -0.0788, -0.0387, 0], "Pmean": [6.217e-13, -7.3193e-10, -1.5832e-07, -3.7598e-06, -3.3837e-05, -0.00017683, -0.00065008, -0.001865, -0.0044488, -0.0092047, -0.017023, -0.028777, -0.045228, -0.066952, -0.094281, -0.12724, -0.16551, -0.20842, -0.25498, -0.30396, -0.35398, -0.40362, -0.45163, -0.49698, -0.5388, -0.57639, -0.60919, -0.63686, -0.6593, -0.67652, -0.68867, -0.69613, -0.69928, -0.69854, -0.69444, -0.68744, -0.67802, -0.66675, -0.65429, -0.64131, -0.62817, -0.61487, -0.60122, -0.58703, -0.57221, -0.55678, -0.5409, -0.52493, -0.50917, -0.49371, -0.47867, -0.46421, -0.4505, -0.43759, -0.4255, -0.41436, -0.40444, -0.39595, -0.389, -0.3836, -0.37966, -0.37702, -0.37542, -0.3746, -0.37436]} \ No newline at end of file diff --git a/examples/cfd/differentiable_lbm.py b/examples/cfd/differentiable_lbm.py new file mode 100644 index 00000000..0aa03784 --- /dev/null +++ b/examples/cfd/differentiable_lbm.py @@ -0,0 +1,570 @@ +""" +Differentiable LBM Example with Configurable Target Shapes + +This example demonstrates gradient-based optimization of initial conditions +to achieve various target density patterns using the Lattice Boltzmann Method. + +Available target shapes: +- 'n_letter': Letter N pattern +- 'circle': Circular pattern +- 'cross': Cross/plus pattern +- 'checkerboard': Checkerboard pattern + +The optimization finds initial conditions (distribution function f) that, +after simulation, produce a density field matching the target pattern. + +Key concepts: +- LBM density stays ~1.0 (physics constraint), so we normalize to [0,1] for loss +- JAX backend is used for automatic differentiation through the stepper +- Simple gradient descent with tuned learning rate + +References: +- Warp example: warp/examples/optim/example_fluid_checkpoint.py +- XLB OOC example: examples/out_of_core/autodiff_lbm.py +""" + +import argparse +import os +from datetime import datetime +import numpy as np +import jax +import jax.numpy as jnp +from jax import value_and_grad + +# Visualization +try: + import matplotlib + matplotlib.use('Agg') # Non-interactive backend for saving + import matplotlib.pyplot as plt + MATPLOTLIB_AVAILABLE = True +except ImportError: + MATPLOTLIB_AVAILABLE = False + print("Warning: matplotlib not available, visualization disabled") + +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.macroscopic import Macroscopic +import xlb.velocity_set + +# For loading XLB logo +try: + from PIL import Image + PIL_AVAILABLE = True +except ImportError: + PIL_AVAILABLE = False + +# Available target shapes +AVAILABLE_SHAPES = ['n_letter', 'circle', 'cross', 'checkerboard'] + + +class DifferentiableLBM: + """ + Differentiable LBM with configurable target shapes. + + Optimizes initial conditions to achieve a target density pattern. + """ + + def __init__( + self, + grid_shape=(128, 128), + Re=100.0, + sim_steps=50, + target_shape='n_letter', + learning_rate=1.0, + target_coverage=0.5, # Fraction of grid covered by target pattern + target_image_path=None, # Path to custom target image (e.g., XLB logo) + ): + self.grid_shape = grid_shape + self.Re = Re + self.sim_steps = sim_steps + self.target_shape = target_shape + self.learning_rate = learning_rate + self.target_coverage = target_coverage + self.output_dir = None # Set by run_optimization if saving + self.target_image_path = target_image_path + + # LBM parameters + self.rho_background = 1.0 + self.rho_variation = 0.1 # Density varies from 0.9 to 1.1 + + # Compute omega from Reynolds number + # Re = u * L / nu, nu = (1/omega - 0.5) / 3 + L = grid_shape[0] + u_ref = 0.1 + nu = u_ref * L / Re + self.omega = 1.0 / (3.0 * nu + 0.5) + self.omega = np.clip(self.omega, 0.5, 1.99) + + # Use JAX backend - required for autodiff through the stepper + # Note: XLB's Warp stepper kernel doesn't have adjoint implementations, + # so gradients are zero when using wp.Tape (verified by test_stepper_autodiff.py). + # JAX uses source transformation which works through the stepper. + self.compute_backend = ComputeBackend.JAX + self.precision_policy = PrecisionPolicy.FP32FP32 + + # Initialize velocity set + self.velocity_set = xlb.velocity_set.D2Q9( + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + + # Store lattice weights and velocities for equilibrium + self.w = jnp.array(self.velocity_set.w, dtype=jnp.float32) + self.c = jnp.array(self.velocity_set.c, dtype=jnp.int32) + + # Initialize XLB + xlb.init( + velocity_set=self.velocity_set, + default_backend=self.compute_backend, + default_precision_policy=self.precision_policy, + ) + + # Create grid and stepper (periodic boundaries) + self.grid = grid_factory(grid_shape, compute_backend=self.compute_backend) + self.stepper = IncompressibleNavierStokesStepper( + grid=self.grid, + boundary_conditions=[], # Periodic + collision_type="BGK", + ) + + # Prepare fields + self.f_0, self.f_1, self.bc_mask, self.missing_mask = self.stepper.prepare_fields() + + # Create macroscopic operator + self.macroscopic = Macroscopic( + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + + # Initialize with uniform density (far from any target) + self._initialize_uniform() + + # Create target pattern + self._create_target() + + print("DifferentiableLBM initialized:") + print(f" Grid: {grid_shape}") + print(f" Re: {Re}, omega: {self.omega:.4f}") + print(f" Sim steps: {sim_steps}") + print(f" Target shape: {target_shape}") + print(f" Learning rate: {learning_rate}") + + def _initialize_uniform(self): + """Initialize with uniform density (normalized = 0).""" + nx, ny = self.grid_shape + rho = np.full((nx, ny), self.rho_background - self.rho_variation, dtype=np.float32) + self.initial_density_normalized = self._normalize_density(rho) + self.f_0 = self._equilibrium(jnp.array(rho), jnp.zeros((2, nx, ny))) + self.f_1 = self.f_0.copy() + + def _normalize_density(self, rho): + """Normalize density to [0, 1] range.""" + rho_min = self.rho_background - self.rho_variation + rho_max = self.rho_background + self.rho_variation + return (rho - rho_min) / (rho_max - rho_min) + + def _equilibrium(self, rho, u): + """Compute equilibrium distribution.""" + cs2 = 1.0 / 3.0 + cu = self.c[0, :, None, None] * u[0] + self.c[1, :, None, None] * u[1] + u_sq = u[0]**2 + u[1]**2 + f_eq = self.w[:, None, None] * rho * ( + 1.0 + cu / cs2 + cu**2 / (2.0 * cs2**2) - u_sq / (2.0 * cs2) + ) + return f_eq + + def _create_target(self): + """Create target pattern based on selected shape.""" + shape_creators = { + 'n_letter': self._create_n_pattern, + 'circle': self._create_circle_pattern, + 'cross': self._create_cross_pattern, + 'checkerboard': self._create_checkerboard_pattern, + } + + # If custom image path provided, use it + if self.target_image_path: + target = self._load_target_image(self.target_image_path) + elif self.target_shape not in shape_creators: + raise ValueError(f"Unknown shape: {self.target_shape}. Available: {AVAILABLE_SHAPES}") + else: + target = shape_creators[self.target_shape]() + self.target_normalized = jnp.array(target) + + coverage = float(jnp.mean(target)) + initial_loss = float(jnp.mean((self.initial_density_normalized - target) ** 2)) + print(f" Target coverage: {coverage*100:.1f}%") + print(f" Expected initial loss: {initial_loss:.6f}") + + def _create_n_pattern(self): + """Create letter N pattern.""" + nx, ny = self.grid_shape + target = np.zeros((nx, ny), dtype=np.float32) + + margin = nx // 10 + bar_width = int(nx * self.target_coverage / 3) # Adjust for coverage + + # Left vertical bar + target[margin:margin+bar_width, margin:ny-margin] = 1.0 + # Right vertical bar + target[nx-margin-bar_width:nx-margin, margin:ny-margin] = 1.0 + # Diagonal + for i in range(nx): + j_center = int(margin + (ny - 2*margin) * (i - margin) / (nx - 2*margin)) + j_start = max(margin, j_center - bar_width//2) + j_end = min(ny - margin, j_center + bar_width//2) + if margin <= i < nx - margin: + target[i, j_start:j_end] = 1.0 + + return target + + def _create_circle_pattern(self): + """Create circular pattern.""" + nx, ny = self.grid_shape + target = np.zeros((nx, ny), dtype=np.float32) + + cx, cy = nx // 2, ny // 2 + # Radius based on coverage: pi*r^2 / (nx*ny) = coverage + radius = np.sqrt(self.target_coverage * nx * ny / np.pi) + + for i in range(nx): + for j in range(ny): + if (i - cx)**2 + (j - cy)**2 < radius**2: + target[i, j] = 1.0 + + return target + + def _create_cross_pattern(self): + """Create cross/plus pattern.""" + nx, ny = self.grid_shape + target = np.zeros((nx, ny), dtype=np.float32) + + # Width based on coverage: 2*w*L - w^2 = coverage * L^2 + # Simplified: w = coverage * L / 2 + width = int(self.target_coverage * nx / 2) + + cx, cy = nx // 2, ny // 2 + margin = nx // 10 + + # Horizontal bar + target[cx-width//2:cx+width//2, margin:ny-margin] = 1.0 + # Vertical bar + target[margin:nx-margin, cy-width//2:cy+width//2] = 1.0 + + return target + + def _create_checkerboard_pattern(self): + """Create checkerboard pattern.""" + nx, ny = self.grid_shape + target = np.zeros((nx, ny), dtype=np.float32) + + # Number of squares based on coverage (checkerboard is always ~50%) + num_squares = 4 # 4x4 checkerboard + sq_size_x = nx // num_squares + sq_size_y = ny // num_squares + + for i in range(num_squares): + for j in range(num_squares): + if (i + j) % 2 == 0: + x_start = i * sq_size_x + x_end = (i + 1) * sq_size_x + y_start = j * sq_size_y + y_end = (j + 1) * sq_size_y + target[x_start:x_end, y_start:y_end] = 1.0 + + return target + + def _load_target_image(self, image_path): + """Load target pattern from image file.""" + if not PIL_AVAILABLE: + raise ImportError("PIL/Pillow is required to load images. Install with: pip install Pillow") + + nx, ny = self.grid_shape + try: + img = Image.open(image_path) + img_resized = img.resize((ny, nx)) # PIL uses (width, height) + # Convert to grayscale and normalize to [0, 1] + img_gray = img_resized.convert('L') + target = np.array(img_gray, dtype=np.float32) / 255.0 + target = np.flipud(target) # Flip vertically to match plot orientation + target = target.T # Transpose to (nx, ny) + print(f" Loaded target image: {image_path}") + return target + except Exception as e: + raise ValueError(f"Could not load image {image_path}: {e}") + + def compute_loss(self, f): + """Compute MSE loss on normalized density.""" + rho, _ = self.macroscopic(f) + rho_norm = self._normalize_density(rho[0]) + rho_norm = jnp.clip(rho_norm, 0.0, 1.0) + loss = jnp.mean((rho_norm - self.target_normalized) ** 2) + return loss + + def forward(self, f_init): + """Run simulation forward.""" + f_curr = f_init + f_next = jnp.zeros_like(f_init) + + for step in range(self.sim_steps): + _, f_next = self.stepper( + f_curr, f_next, self.bc_mask, self.missing_mask, self.omega, step + ) + f_curr, f_next = f_next, f_curr + + return f_curr + + def loss_fn(self, f_init): + """Loss function for optimization.""" + f_final = self.forward(f_init) + return self.compute_loss(f_final) + + def optimize_step(self): + """Perform one gradient descent step.""" + loss_val, grad_f = value_and_grad(self.loss_fn)(self.f_0) + + # Gradient descent update + self.f_0 = self.f_0 - self.learning_rate * grad_f + + # Clamp f to physical range + f_min = 0.01 * self.w[:, None, None] + f_max = 10.0 * self.w[:, None, None] + self.f_0 = jnp.clip(self.f_0, f_min, f_max) + + self.f_1 = self.f_0.copy() + + return float(loss_val) + + def get_initial_density(self): + """Get normalized initial density (from current f_0).""" + rho, _ = self.macroscopic(self.f_0) + rho_norm = self._normalize_density(rho[0]) + return np.array(jnp.clip(rho_norm, 0.0, 1.0)) + + def get_final_density(self): + """Get normalized final density (after simulation).""" + f_final = self.forward(self.f_0) + rho, _ = self.macroscopic(f_final) + rho_norm = self._normalize_density(rho[0]) + return np.array(jnp.clip(rho_norm, 0.0, 1.0)) + + def save_iteration_plot(self, iteration, loss): + """Save plot showing initial, final, and target density for this iteration.""" + if not MATPLOTLIB_AVAILABLE or self.output_dir is None: + return + + initial = self.get_initial_density() + final = self.get_final_density() + target = np.array(self.target_normalized) + + fig, axes = plt.subplots(1, 4, figsize=(16, 4)) + + # Initial density (what we're optimizing) + im0 = axes[0].imshow(initial.T, origin='lower', cmap='viridis', vmin=0, vmax=1) + axes[0].set_title('Initial Condition\n(optimized f_0)') + axes[0].set_xlabel('x') + axes[0].set_ylabel('y') + plt.colorbar(im0, ax=axes[0], shrink=0.8) + + # Final density (after simulation) + im1 = axes[1].imshow(final.T, origin='lower', cmap='viridis', vmin=0, vmax=1) + axes[1].set_title(f'Final Density\n(after {self.sim_steps} steps)') + axes[1].set_xlabel('x') + axes[1].set_ylabel('y') + plt.colorbar(im1, ax=axes[1], shrink=0.8) + + # Target density + im2 = axes[2].imshow(target.T, origin='lower', cmap='viridis', vmin=0, vmax=1) + axes[2].set_title(f'Target\n({self.target_shape})') + axes[2].set_xlabel('x') + axes[2].set_ylabel('y') + plt.colorbar(im2, ax=axes[2], shrink=0.8) + + # Difference (final - target) + diff = np.abs(final - target) + im3 = axes[3].imshow(diff.T, origin='lower', cmap='Reds', vmin=0, vmax=0.5) + axes[3].set_title(f'|Final - Target|\nMSE={loss:.4f}') + axes[3].set_xlabel('x') + axes[3].set_ylabel('y') + plt.colorbar(im3, ax=axes[3], shrink=0.8) + + plt.suptitle(f'Iteration {iteration:05d} - Loss: {loss:.6f}', fontsize=14) + plt.tight_layout() + + filepath = os.path.join(self.output_dir, f'iteration_{iteration:05d}.png') + plt.savefig(filepath, dpi=100) + plt.close(fig) + + def save_convergence_plot(self, losses): + """Save convergence plot showing loss over iterations.""" + if not MATPLOTLIB_AVAILABLE or self.output_dir is None: + return + + fig, ax = plt.subplots(figsize=(10, 6)) + ax.plot(losses, 'b-', linewidth=2) + ax.set_xlabel('Iteration', fontsize=12) + ax.set_ylabel('Loss (MSE)', fontsize=12) + ax.set_title(f'Optimization Convergence - {self.target_shape}', fontsize=14) + ax.grid(True, alpha=0.3) + # Use linear scale with regular numbers (not scientific notation) + ax.ticklabel_format(style='plain', axis='y') + + # Add annotations + ax.axhline(y=losses[-1], color='r', linestyle='--', alpha=0.5, + label=f'Final: {losses[-1]:.4f}') + ax.legend() + + plt.tight_layout() + filepath = os.path.join(self.output_dir, 'convergence.png') + plt.savefig(filepath, dpi=150) + plt.close(fig) + print(f" Saved convergence plot: {filepath}") + + def run_optimization(self, num_iterations=100, verbose=True, save_plots=False, + save_every=10, output_dir=None): + """Run optimization loop with optional visualization. + + Parameters + ---------- + num_iterations : int + Number of optimization iterations + verbose : bool + Print loss each iteration + save_plots : bool + Save density plots to disk + save_every : int + Save plot every N iterations (also saves first and last) + output_dir : str + Directory to save plots (default: output_diff_lbm_) + """ + losses = [] + + # Setup output directory if saving + if save_plots and MATPLOTLIB_AVAILABLE: + if output_dir is None: + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + output_dir = f'output_diff_lbm_{self.target_shape}_{timestamp}' + self.output_dir = output_dir + os.makedirs(self.output_dir, exist_ok=True) + print(f" Saving plots to: {self.output_dir}") + + for i in range(num_iterations): + loss = self.optimize_step() + losses.append(loss) + + if verbose: + print(f"Iteration {i:05d} loss: {loss:.6f}") + + # Save plots at specified intervals + if save_plots and MATPLOTLIB_AVAILABLE: + if i == 0 or i == num_iterations - 1 or (i + 1) % save_every == 0: + self.save_iteration_plot(i, loss) + + # Save convergence plot + if save_plots and MATPLOTLIB_AVAILABLE: + self.save_convergence_plot(losses) + + return losses + + +def main(): + parser = argparse.ArgumentParser( + description="Differentiable LBM with configurable target shapes", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--shape", type=str, default="n_letter", + choices=AVAILABLE_SHAPES, + help="Target shape to optimize towards", + ) + parser.add_argument( + "--grid-size", type=int, default=128, + help="Grid size (NxN)", + ) + parser.add_argument( + "--sim-steps", type=int, default=50, + help="Number of simulation steps per forward pass", + ) + parser.add_argument( + "--iterations", type=int, default=150, + help="Number of optimization iterations", + ) + parser.add_argument( + "--learning-rate", type=float, default=1.0, + help="Learning rate for gradient descent", + ) + parser.add_argument( + "--Re", type=float, default=100.0, + help="Reynolds number", + ) + parser.add_argument( + "--coverage", type=float, default=0.5, + help="Target pattern coverage (0-1)", + ) + parser.add_argument( + "--save-plots", action="store_true", + help="Save density plots to disk", + ) + parser.add_argument( + "--save-every", type=int, default=10, + help="Save plot every N iterations", + ) + parser.add_argument( + "--output-dir", type=str, default=None, + help="Output directory for plots (default: auto-generated)", + ) + parser.add_argument( + "--target-image", type=str, default=None, + help="Path to custom target image (overrides --shape)", + ) + + args = parser.parse_args() + + print("=" * 70) + print("Differentiable LBM - Configurable Target Shapes") + print("=" * 70) + print() + + sim = DifferentiableLBM( + grid_shape=(args.grid_size, args.grid_size), + Re=args.Re, + sim_steps=args.sim_steps, + target_shape=args.shape, + learning_rate=args.learning_rate, + target_coverage=args.coverage, + target_image_path=args.target_image, + ) + + print() + losses = sim.run_optimization( + num_iterations=args.iterations, + verbose=True, + save_plots=args.save_plots, + save_every=args.save_every, + output_dir=args.output_dir, + ) + + print() + print("=" * 70) + print("RESULTS") + print("=" * 70) + print(f"Initial loss: {losses[0]:.6f}") + print(f"Final loss: {losses[-1]:.6f}") + print(f"Improvement: {(losses[0] - losses[-1]) / losses[0] * 100:.2f}%") + + # Check convergence + if len(losses) >= 10: + last_10_change = abs(losses[-10] - losses[-1]) / losses[-10] * 100 + print(f"Last 10 iter change: {last_10_change:.2f}%") + if last_10_change < 1.0: + print("Status: CONVERGED") + else: + print("Status: Still improving (run more iterations)") + + +if __name__ == "__main__": + main() diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py new file mode 100644 index 00000000..1616d32a --- /dev/null +++ b/examples/cfd/flow_past_sphere_3d.py @@ -0,0 +1,187 @@ +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + HalfwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, +) +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_image +import warp as wp +import numpy as np +import jax.numpy as jnp +import time + +# -------------------------- Simulation Setup -------------------------- + +omega = 1.6 +grid_shape = (512 // 2, 128 // 2, 128 // 2) +compute_backend = ComputeBackend.JAX +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend) +u_max = 0.04 +num_steps = 10000 +post_process_interval = 1000 + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Create Grid +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +# Define Boundary Indices +box = grid.bounding_box_indices() +box_no_edge = grid.bounding_box_indices(remove_edges=True) +inlet = box_no_edge["left"] +outlet = box_no_edge["right"] +walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] +walls = np.unique(np.array(walls), axis=-1).tolist() + +sphere_radius = grid_shape[1] // 12 +x = np.arange(grid_shape[0]) +y = np.arange(grid_shape[1]) +z = np.arange(grid_shape[2]) +X, Y, Z = np.meshgrid(x, y, z, indexing="ij") +indices = np.where((X - grid_shape[0] // 6) ** 2 + (Y - grid_shape[1] // 2) ** 2 + (Z - grid_shape[2] // 2) ** 2 < sphere_radius**2) +sphere = [tuple(indices[i].tolist()) for i in range(velocity_set.d)] + + +# Define Boundary Conditions +def bc_profile(): + H_y = float(grid_shape[1] - 1) # Height in y direction + H_z = float(grid_shape[2] - 1) # Height in z direction + + if compute_backend == ComputeBackend.JAX: + + def bc_profile_jax(): + y = jnp.arange(grid_shape[1]) + z = jnp.arange(grid_shape[2]) + Y, Z = jnp.meshgrid(y, z, indexing="ij") + + # Calculate normalized distance from center + y_center = Y - (H_y / 2.0) + z_center = Z - (H_z / 2.0) + r_squared = (2.0 * y_center / H_y) ** 2.0 + (2.0 * z_center / H_z) ** 2.0 + + # Parabolic profile for x velocity, zero for y and z + u_x = u_max * jnp.maximum(0.0, 1.0 - r_squared) + u_y = jnp.zeros_like(u_x) + u_z = jnp.zeros_like(u_x) + + return jnp.stack([u_x, u_y, u_z]) + + return bc_profile_jax + + else: + wp_dtype = precision_policy.compute_precision.wp_dtype + H_y = wp_dtype(grid_shape[1] - 1) # Height in y direction + H_z = wp_dtype(grid_shape[2] - 1) # Height in z direction + two = wp_dtype(2.0) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + # Poiseuille flow profile: parabolic velocity distribution + y = wp_dtype(index[1]) + z = wp_dtype(index[2]) + + # Calculate normalized distance from center + y_center = y - (H_y / two) + z_center = z - (H_z / two) + r_squared = (two * y_center / H_y) ** two + (two * z_center / H_z) ** two + + # Parabolic profile: u = u_max * (1 - rΒ²) + return wp.vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1) + + return bc_profile_warp + + +# Initialize Boundary Conditions +bc_left = RegularizedBC("velocity", profile=bc_profile(), indices=inlet) +# Alternatively, use a prescribed velocity profile +# bc_left = RegularizedBC("velocity", prescribed_value=(u_max, 0.0, 0.0), indices=inlet) +bc_walls = FullwayBounceBackBC(indices=walls) +bc_outlet = ExtrapolationOutflowBC(indices=outlet) +bc_sphere = HalfwayBounceBackBC(indices=sphere) +boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere] + +# Setup Stepper +stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="BGK", +) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + +# Define Macroscopic Calculation +macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), +) +to_jax = xlb.utils.ToJAX("populations", velocity_set.q, grid_shape) + +# Setup Momentum Transfer for Force Calculation +from xlb.operator.force.momentum_transfer import MomentumTransfer + +momentum_transfer = MomentumTransfer(bc_sphere, compute_backend=compute_backend) +sphere_cross_section = np.pi * sphere_radius**2 + + +# Post-Processing Function +def post_process(step, f_0, f_1): + wp.synchronize() + + # Compute lift and drag + boundary_force = momentum_transfer(f_0, f_1, bc_mask, missing_mask) + drag = boundary_force[0] # x-direction + lift = boundary_force[2] + cd = 2.0 * drag / (u_max**2 * sphere_cross_section) + cl = 2.0 * lift / (u_max**2 * sphere_cross_section) + print(f"CD={cd}, CL={cl}") + + # Convert to JAX array if necessary + if not isinstance(f_0, jnp.ndarray): + f_0 = to_jax(f_0) + wp.synchronize() + + rho, u = macro(f_0) + + # Remove boundary cells + u = u[:, 1:-1, 1:-1, 1:-1] + rho = rho[:, 1:-1, 1:-1, 1:-1] + u_magnitude = jnp.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + + fields = { + "u_magnitude": u_magnitude, + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + "rho": rho[0], + } + + # Save the u_magnitude slice at the mid y-plane + save_image(fields["u_magnitude"][:, grid_shape[1] // 2, :], timestep=step) + print(f"Post-processed step {step}: Saved u_magnitude slice at y={grid_shape[1] // 2}") + + +# -------------------------- Simulation Loop -------------------------- + +start_time = time.time() +for step in range(num_steps): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 # Swap the buffers + + if step % post_process_interval == 0 or step == num_steps - 1: + post_process(step, f_0, f_1) + end_time = time.time() + elapsed = end_time - start_time + print(f"Completed step {step}. Time elapsed for {post_process_interval} steps: {elapsed:.6f} seconds.") + start_time = time.time() diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py new file mode 100644 index 00000000..7c0eac7b --- /dev/null +++ b/examples/cfd/lid_driven_cavity_2d.py @@ -0,0 +1,116 @@ +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import HalfwayBounceBackBC, EquilibriumBC +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_fields_vtk, save_image, warp_array_to_jax +import xlb.velocity_set +import jax.numpy as jnp +import numpy as np + + +class LidDrivenCavity2D: + def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, compute_backend, precision_policy): + # initialize compute_backend + xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, + ) + + self.grid_shape = grid_shape + self.velocity_set = velocity_set + self.compute_backend = compute_backend + self.precision_policy = precision_policy + self.omega = omega + self.boundary_conditions = [] + self.prescribed_vel = prescribed_vel + + # Create grid using factory + self.grid = grid_factory(grid_shape, compute_backend=compute_backend) + + # Setup the simulation BC and stepper + self._setup() + + def _setup(self): + self.setup_boundary_conditions() + self.setup_stepper() + # Initialize fields using the stepper + self.f_0, self.f_1, self.bc_mask, self.missing_mask = self.stepper.prepare_fields() + + def define_boundary_indices(self): + box = self.grid.bounding_box_indices() + box_no_edge = self.grid.bounding_box_indices(remove_edges=True) + lid = box_no_edge["top"] + walls = [box["bottom"][i] + box["left"][i] + box["right"][i] for i in range(self.velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + return lid, walls + + def setup_boundary_conditions(self): + lid, walls = self.define_boundary_indices() + bc_top = EquilibriumBC(rho=1.0, u=(self.prescribed_vel, 0.0), indices=lid) + bc_walls = HalfwayBounceBackBC(indices=walls) + self.boundary_conditions = [bc_walls, bc_top] + + def setup_stepper(self): + self.stepper = IncompressibleNavierStokesStepper( + grid=self.grid, + boundary_conditions=self.boundary_conditions, + collision_type="BGK", + ) + + def run(self, num_steps, post_process_interval=100): + for i in range(num_steps): + self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, self.omega, i) + self.f_0, self.f_1 = self.f_1, self.f_0 + + if i % post_process_interval == 0 or i == num_steps - 1: + self.post_process(i) + + def post_process(self, i): + # Write the results. We'll use JAX compute_backend for the post-processing + if not isinstance(self.f_0, jnp.ndarray): + # If the compute_backend is warp, we need to drop the last dimension added by warp for 2D simulations + f_0 = warp_array_to_jax(self.f_0)[..., 0] + else: + f_0 = self.f_0 + + macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=self.precision_policy, + velocity_set=xlb.velocity_set.D2Q9(precision_policy=self.precision_policy, compute_backend=ComputeBackend.JAX), + ) + + rho, u = macro(f_0) + + # remove boundary cells + rho = rho[:, 1:-1, 1:-1] + u = u[:, 1:-1, 1:-1] + u_magnitude = (u[0] ** 2 + u[1] ** 2) ** 0.5 + + fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude} + + save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity") + save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity") + + +if __name__ == "__main__": + # Running the simulation + grid_size = 500 + grid_shape = (grid_size, grid_size) + compute_backend = ComputeBackend.JAX + precision_policy = PrecisionPolicy.FP32FP32 + + velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, compute_backend=compute_backend) + + # Setting fluid viscosity and relaxation parameter. + Re = 200.0 + prescribed_vel = 0.05 + clength = grid_shape[0] - 1 + visc = prescribed_vel * clength / Re + omega = 1.0 / (3.0 * visc + 0.5) + + simulation = LidDrivenCavity2D(omega, prescribed_vel, grid_shape, velocity_set, compute_backend, precision_policy) + simulation.run(num_steps=50000, post_process_interval=1000) diff --git a/examples/cfd/lid_driven_cavity_2d_distributed.py b/examples/cfd/lid_driven_cavity_2d_distributed.py new file mode 100644 index 00000000..06a2822a --- /dev/null +++ b/examples/cfd/lid_driven_cavity_2d_distributed.py @@ -0,0 +1,48 @@ +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.distribute import distribute +from lid_driven_cavity_2d import LidDrivenCavity2D + + +class LidDrivenCavity2D_distributed(LidDrivenCavity2D): + def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, compute_backend, precision_policy): + super().__init__(omega, prescribed_vel, grid_shape, velocity_set, compute_backend, precision_policy) + + def setup_stepper(self): + # Create the base stepper + stepper = IncompressibleNavierStokesStepper( + grid=self.grid, + boundary_conditions=self.boundary_conditions, + collision_type="BGK", + ) + + # Distribute the stepper + self.stepper = distribute( + stepper, + self.grid, + self.velocity_set, + ) + + +if __name__ == "__main__": + # Running the simulation + grid_size = 512 + grid_shape = (grid_size, grid_size) + compute_backend = ( + ComputeBackend.JAX + ) # Must be JAX for distributed multi-GPU computations. Distributed computations on WARP are not supported yet! + precision_policy = PrecisionPolicy.FP32FP32 + + velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, compute_backend=compute_backend) + + # Setting fluid viscosity and relaxation parameter. + Re = 200.0 + prescribed_vel = 0.05 + clength = grid_shape[0] - 1 + visc = prescribed_vel * clength / Re + omega = 1.0 / (3.0 * visc + 0.5) + + simulation = LidDrivenCavity2D_distributed(omega, prescribed_vel, grid_shape, velocity_set, compute_backend, precision_policy) + simulation.run(num_steps=50000, post_process_interval=1000) diff --git a/examples/cfd/multires_flow_past_sphere_3d.py b/examples/cfd/multires_flow_past_sphere_3d.py new file mode 100644 index 00000000..968e9cd4 --- /dev/null +++ b/examples/cfd/multires_flow_past_sphere_3d.py @@ -0,0 +1,300 @@ +""" +3D flow past a sphere with multi-resolution LBM. + +Demonstrates the multi-resolution Neon backend for a 3-D Poiseuille- +inlet flow past a sphere inside a nested cuboid multi-resolution domain. +Uses AABB-Close voxelization with halfway bounce-back on the sphere surface and +computes lift/drag via momentum transfer. +""" + +try: + import neon +except ModuleNotFoundError: + import sys + + raise ModuleNotFoundError( + "The 'neon' module is required for this example (Neon backend). " + "Install with: pip install 'xlb[neon]'. " + f"You are on Python {sys.version_info.major}.{sys.version_info.minor}; " + "the current Neon wheel requires Python 3.11 or 3.12. Use an appropriate environment (e.g. pyenv, conda, or venv)." + ) from None +import warp as wp +import numpy as np +import time + +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import multires_grid_factory +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + HalfwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, + DoNothingBC, + ZouHeBC, + HybridBC, +) +from xlb.operator.boundary_masker import MeshVoxelizationMethod +from xlb.utils.mesher import make_cuboid_mesh, prepare_sparsity_pattern +from xlb.operator.force import MultiresMomentumTransfer + + +def generate_cuboid_mesh(stl_filename, num_finest_voxels_across_part): + """ + Generate a cuboid mesh based on the provided voxel size and domain multipliers. + """ + import trimesh + import os + + # Domain multipliers for each refinement level + # First entry should be full domain size + # Domain multipliers + domainMultiplier = [ + [7, 22, 7, 7, 7, 7], # -x, x, -y, y, -z, z (sphere at 1/4 domain from inlet) + [3, 12, 5, 5, 5, 5], # -x, x, -y, y, -z, z (wake-biased) + [2, 8, 4, 4, 4, 4], + [1, 5, 2, 2, 2, 2], + # [1, 2, 1, 1, 1, 1], + # [0.4, 1, 0.4, 0.4, 0.4, 0.4], + # [0.2, 0.4, 0.2, 0.2, 0.2, 0.2], + ] + + # Load the mesh + mesh = trimesh.load_mesh(stl_filename, process=False) + assert not mesh.is_empty, ValueError("Loaded mesh is empty or invalid.") + + # Compute original bounds + # Find voxel size and sphere radius + min_bound = mesh.vertices.min(axis=0) + max_bound = mesh.vertices.max(axis=0) + partSize = max_bound - min_bound + + # smallest voxel size + voxel_size = min(partSize) / num_finest_voxels_across_part + + # Compute translation to put mesh into first octant of that domainβ€” + shift = np.array( + [ + domainMultiplier[0][0] * partSize[0] - min_bound[0], + domainMultiplier[0][2] * partSize[1] - min_bound[1], + domainMultiplier[0][4] * partSize[2] - min_bound[2], + ], + dtype=float, + ) + + # Apply translation and save out temp stl + mesh.apply_translation(shift) + _ = mesh.vertex_normals + mesh_vertices = np.asarray(mesh.vertices) / voxel_size + mesh.export("temp.stl") + + # Mesh based on temp stl + level_data = make_cuboid_mesh( + voxel_size, + domainMultiplier, + "temp.stl", + ) + grid_shape_finest = tuple([i * 2 ** (len(level_data) - 1) for i in level_data[-1][0].shape]) + print(f"Full shape based on finest voxels size is {grid_shape_finest}") + os.remove("temp.stl") + return level_data, mesh_vertices, tuple([int(a) for a in grid_shape_finest]) + + +# -------------------------- Simulation Setup -------------------------- + +# The following parameters define the resolution of the voxelized grid +sphere_radius = 5 +num_finest_voxels_across_part = 2 * sphere_radius + +# Other setup parameters +Re = 5000.0 +compute_backend = ComputeBackend.NEON +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +u_max = 0.04 +num_steps = 10000 +post_process_interval = 1000 + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Generate the cuboid mesh and sphere vertices +stl_filename = "../stl-files/sphere.stl" +level_data, sphere, grid_shape_finest = generate_cuboid_mesh(stl_filename, num_finest_voxels_across_part) + + +# Define exporter object for hdf5 output +from xlb.utils import MultiresIO + +# Define an exporter for the multiresolution data +exporter = MultiresIO({"velocity": 3, "density": 1}, level_data) + +# Prepare the sparsity pattern and origins from the level data +sparsity_pattern, level_origins = prepare_sparsity_pattern(level_data) + +# get the number of levels +num_levels = len(level_data) + +# Create the multires grid +grid = multires_grid_factory( + grid_shape_finest, + velocity_set=velocity_set, + sparsity_pattern_list=sparsity_pattern, + sparsity_pattern_origins=[neon.Index_3d(*box_origin) for box_origin in level_origins], +) + +# Define Boundary Indices +coarsest_level = grid.count_levels - 1 +box = grid.bounding_box_indices(shape=grid.level_to_shape(coarsest_level)) +box_no_edge = grid.bounding_box_indices(shape=grid.level_to_shape(coarsest_level), remove_edges=True) +inlet = box_no_edge["left"] +outlet = box_no_edge["right"] +walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] +walls = np.unique(np.array(walls), axis=-1).tolist() + + +# Define Boundary Conditions +def bc_profile(): + """Build a Warp function for a Poiseuille parabolic inlet velocity profile.""" + assert compute_backend == ComputeBackend.NEON + # IMPORTANT NOTE: the user defined functional must be defined in terms of the indices at the finest level + _, ny, nz = grid_shape_finest + dtype = precision_policy.compute_precision.wp_dtype + H_y = dtype(ny) # Length in y direction (finest level) + H_z = dtype(nz) # Length in z direction (finest level) + two = dtype(2.0) + one = dtype(1.0) + zero = dtype(0.0) + u_max_wp = dtype(u_max) + _u_vec = wp.vec(velocity_set.d, dtype=dtype) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + # Poiseuille flow profile: parabolic velocity distribution + y = dtype(index[1]) + z = dtype(index[2]) + + # Calculate normalized distance from center + y_center = y - (H_y / two) + z_center = z - (H_z / two) + r_squared = (two * y_center / H_y) ** two + (two * z_center / H_z) ** two + + # Parabolic profile: u = u_max * (1 - rΒ²) + # Note that unlike RegularizedBC and ZouHeBC which only accept normal velocity, hybridBC accepts the full velocity vector + + # For hybridBC + # return _u_vec(u_max_wp * wp.max(zero, one - r_squared), zero, zero) + + # For Regularized and ZouHe + return wp.vec(u_max_wp * wp.max(zero, one - r_squared), length=1) + + return bc_profile_warp + + +# Convert bc indices to a list of list (first entry corresponds to the finest level) +inlet = [[] for _ in range(num_levels - 1)] + [inlet] +outlet = [[] for _ in range(num_levels - 1)] + [outlet] +walls = [[] for _ in range(num_levels - 1)] + [walls] + +# Initialize Boundary Conditions +bc_left = RegularizedBC("velocity", profile=bc_profile(), indices=inlet) +# Alternatives: +# bc_left = HybridBC(bc_method="bounceback_regularized", profile=bc_profile(), indices=inlet) +# bc_left = RegularizedBC("velocity", prescribed_value=(u_max, 0.0, 0.0), indices=inlet) +bc_walls = FullwayBounceBackBC(indices=walls) +bc_outlet = DoNothingBC(indices=outlet) +# bc_outlet = ExtrapolationOutflowBC(indices=outlet) +bc_sphere = HybridBC( + bc_method="nonequilibrium_regularized", mesh_vertices=sphere, voxelization_method=MeshVoxelizationMethod("AABB"), use_mesh_distance=True +) +# bc_sphere = HalfwayBounceBackBC(mesh_vertices=sphere, voxelization_method=MeshVoxelizationMethod('AABB')) + +boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere] + +# Configure the simulation relaxation time +visc = u_max * num_finest_voxels_across_part / Re +omega_finest = 1.0 / (3.0 * visc + 0.5) + +# Make initializer operator +from xlb.helper.initializers import CustomMultiresInitializer + +initializer = CustomMultiresInitializer( + bc_id=bc_outlet.id, + constant_velocity_vector=(u_max, 0.0, 0.0), + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, +) + +# Define a multi-resolution simulation manager +sim = xlb.helper.MultiresSimulationManager( + omega_finest=omega_finest, + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + initializer=initializer, + mres_perf_opt=xlb.mres_perf_optimization_type.MresPerfOptimizationType.FUSION_AT_FINEST, +) + +# Setup Momentum Transfer for Force Calculation +bc_sphere = boundary_conditions[-1] +momentum_transfer = MultiresMomentumTransfer(bc_sphere, mres_perf_opt=sim.mres_perf_opt, compute_backend=compute_backend) + + +def print_lift_drag(sim): + """Compute and print drag and lift coefficients from the simulation state.""" + boundary_force = momentum_transfer(sim.f_0, sim.f_1, sim.bc_mask, sim.missing_mask) + drag = boundary_force[0] # x-direction + lift = boundary_force[2] + sphere_cross_section = np.pi * sphere_radius**2 + u_avg = 0.5 * u_max + cd = 2.0 * drag / (u_avg**2 * sphere_cross_section) + cl = 2.0 * lift / (u_avg**2 * sphere_cross_section) + print(f"\tCD={cd}, CL={cl}") + + +# -------------------------- Simulation Loop -------------------------- + +wp.synchronize() +start_time = time.time() +for step in range(num_steps): + sim.step() + + if step % post_process_interval == 0 or step == num_steps - 1: + # # Export VTK for comparison + # tic_write = time.perf_counter() + # sim.export_macroscopic("multires_flow_over_sphere_3d_") + # toc_write = time.perf_counter() + # print(f"\tVTK file written in {toc_write - tic_write:0.1f} seconds") + + # Call the Macroscopic operator to compute macroscopic fields + wp.synchronize() + sim.macro(sim.f_0, sim.bc_mask, sim.rho, sim.u, streamId=0) + + # Call the exporter to save the current state + nx, ny, nz = grid_shape_finest + filename = f"multires_flow_past_sphere_3d_{step:04d}" + wp.synchronize() + exporter.to_hdf5(filename, {"velocity": sim.u, "density": sim.rho}, compression="gzip", compression_opts=2) + exporter.to_slice_image( + filename, + {"velocity": sim.u}, + plane_point=(nx // 2, ny // 2, nz // 2), + plane_normal=(0, 0, 1), + grid_res=256, + slice_thickness=2 ** (num_levels - 1), + bounds=(0.1, 0.6, 0.3, 0.7), + ) + + # Print lift and drag coefficients + print_lift_drag(sim) + wp.synchronize() + end_time = time.time() + elapsed = end_time - start_time + print(f"\tCompleted step {step}. Time elapsed for {post_process_interval} steps: {elapsed:.6f} seconds.") + start_time = time.time() diff --git a/examples/cfd/multires_windtunnel_3d.py b/examples/cfd/multires_windtunnel_3d.py new file mode 100644 index 00000000..d94d4357 --- /dev/null +++ b/examples/cfd/multires_windtunnel_3d.py @@ -0,0 +1,575 @@ +""" +Ahmed body aerodynamics with multi-resolution LBM. + +Simulates turbulent flow around the Ahmed body (25-degree slant angle) +using the XLB multi-resolution Neon backend. Computes drag and lift +coefficients via momentum transfer and exports HDF5/XDMF data for +post-processing. +""" + +import neon +import warp as wp +import numpy as np +import time +import os +import matplotlib.pyplot as plt +import trimesh +import shutil + +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import multires_grid_factory +from xlb.operator.boundary_condition import ( + DoNothingBC, + HybridBC, + RegularizedBC, +) +from xlb.operator.boundary_masker import MeshVoxelizationMethod +from xlb.utils.mesher import prepare_sparsity_pattern, make_cuboid_mesh, MultiresIO +from xlb.utils import UnitConvertor +from xlb.operator.force import MultiresMomentumTransfer +from xlb.helper.initializers import CustomMultiresInitializer + +wp.clear_kernel_cache() +wp.config.quiet = True + +# User Configuration +# ================= +# Physical and simulation parameters +wind_speed_lbm = 0.05 # Lattice velocity +wind_speed_mps = 38.0 # Physical inlet velocity in m/s (user input) +flow_passes = 2 # Domain flow passes +kinematic_viscosity = 1.508e-5 # Kinematic viscosity of air in m^2/s 1.508e-5 +voxel_size = 0.005 # Finest voxel size in meters + +# STL filename +stl_filename = "../stl-files/Ahmed_25_NoLegs.stl" +script_name = "Ahmed" + +# I/O settings +print_interval_percentage = 1 # Print every 1% of iterations +file_output_crossover_percentage = 10 # Crossover at 50% of iterations +num_file_outputs_pre_crossover = 20 # Outputs before crossover +num_file_outputs_post_crossover = 5 # Outputs after crossover + +# Other setup parameters +compute_backend = ComputeBackend.NEON +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) + + +def generate_cuboid_mesh(stl_filename, voxel_size): + """ + Alternative cuboid mesh generation based on Apolo's method with domain multipliers per level. + """ + # Domain multipliers for each refinement level + domain_multiplier = [ + [3.0, 4.0, 2.5, 2.5, 0.0, 4.0], # -x, x, -y, y, -z, z + [1.2, 1.25, 1.75, 1.75, 0.0, 1.5], + [0.8, 1.0, 1.25, 1.25, 0.0, 1.2], + [0.5, 0.65, 0.6, 0.60, 0.0, 0.6], + [0.25, 0.25, 0.25, 0.25, 0.0, 0.25], + ] + + # Load the mesh + mesh = trimesh.load_mesh(stl_filename, process=False) + if mesh.is_empty: + raise ValueError("Loaded mesh is empty or invalid.") + + # Compute original bounds + min_bound = mesh.vertices.min(axis=0) + max_bound = mesh.vertices.max(axis=0) + partSize = max_bound - min_bound + x0 = max_bound[0] # End of car for Ahmed + + # Compute translation to put mesh into first octant of the domain + stl_shift = np.array( + [ + domain_multiplier[0][0] * partSize[0] - min_bound[0], + domain_multiplier[0][2] * partSize[1] - min_bound[1], + domain_multiplier[0][4] * partSize[2] - min_bound[2], + ], + dtype=float, + ) + + # Apply translation and save out temp STL + mesh.apply_translation(stl_shift) + _ = mesh.vertex_normals + mesh_vertices = np.asarray(mesh.vertices) + mesh.export("temp.stl") + + # Generate mesh using make_cuboid_mesh + level_data = make_cuboid_mesh( + voxel_size, + domain_multiplier, + "temp.stl", + ) + + num_levels = len(level_data) + grid_shape_finest = tuple([int(i * 2 ** (num_levels - 1)) for i in level_data[-1][0].shape]) + print(f"Full shape based on finest voxel size is {grid_shape_finest}") + os.remove("temp.stl") + + return ( + level_data, + mesh_vertices, + tuple([int(a) for a in grid_shape_finest]), + stl_shift, + x0, + ) + + +# Boundary Conditions Setup +# ========================= +def setup_boundary_conditions(grid, level_data, body_vertices, wind_speed_mps): + """ + Set up boundary conditions for the simulation. + """ + # Convert wind speed to lattice units + wind_speed_lbm = unit_convertor.velocity_to_lbm(wind_speed_mps) + + left_indices = grid.boundary_indices_across_levels(level_data, box_side="left", remove_edges=True) + right_indices = grid.boundary_indices_across_levels(level_data, box_side="right", remove_edges=True) + top_indices = grid.boundary_indices_across_levels(level_data, box_side="top", remove_edges=False) + bottom_indices = grid.boundary_indices_across_levels(level_data, box_side="bottom", remove_edges=False) + front_indices = grid.boundary_indices_across_levels(level_data, box_side="front", remove_edges=False) + back_indices = grid.boundary_indices_across_levels(level_data, box_side="back", remove_edges=False) + + # Initialize boundary conditions + bc_inlet = RegularizedBC("velocity", prescribed_value=(wind_speed_lbm, 0.0, 0.0), indices=left_indices) + bc_outlet = DoNothingBC(indices=right_indices) + bc_top = HybridBC(bc_method="nonequilibrium_regularized", indices=top_indices) + bc_bottom = HybridBC(bc_method="nonequilibrium_regularized", indices=bottom_indices) + bc_front = HybridBC(bc_method="nonequilibrium_regularized", indices=front_indices) + bc_back = HybridBC(bc_method="nonequilibrium_regularized", indices=back_indices) + bc_body = HybridBC( + bc_method="nonequilibrium_regularized", + mesh_vertices=unit_convertor.length_to_lbm(body_vertices), + voxelization_method=MeshVoxelizationMethod("AABB_CLOSE", close_voxels=4), + use_mesh_distance=True, + ) + + return [bc_top, bc_bottom, bc_front, bc_back, bc_inlet, bc_outlet, bc_body] + + +# Simulation Initialization +# ========================= +def initialize_simulation( + grid, boundary_conditions, omega_finest, initializer, collision_type="KBC", mres_perf_opt=xlb.MresPerfOptimizationType.FUSION_AT_FINEST +): + """ + Initialize the multiresolution simulation manager. + """ + sim = xlb.helper.MultiresSimulationManager( + omega_finest=omega_finest, + grid=grid, + boundary_conditions=boundary_conditions, + collision_type=collision_type, + initializer=initializer, + mres_perf_opt=mres_perf_opt, + ) + return sim + + +# Utility Functions +# ================= +def compute_force_coefficients(sim, step, momentum_transfer, wind_speed_lbm, reference_area): + """ + Calculate and print lift and drag coefficients. + """ + boundary_force = momentum_transfer(sim.f_0, sim.f_1, sim.bc_mask, sim.missing_mask) + drag = boundary_force[0] + lift = boundary_force[2] + cd = 2.0 * drag / (wind_speed_lbm**2 * reference_area) + cl = 2.0 * lift / (wind_speed_lbm**2 * reference_area) + if np.isnan(cd) or np.isnan(cl): + print(f"NaN detected in coefficients at step {step}") + raise ValueError(f"NaN detected in coefficients at step {step}: Cd={cd}, Cl={cl}") + drag_values.append([cd, cl]) + return cd, cl, drag + + +def plot_force_coefficients(drag_values, output_dir, print_interval, script_name, percentile_range=(15, 85), use_log_scale=False): + """ + Plot CD and CL over time and save the plot to the output directory. + """ + drag_values_array = np.array(drag_values) + steps = np.arange(0, len(drag_values) * print_interval, print_interval) + cd_values = drag_values_array[:, 0] + cl_values = drag_values_array[:, 1] + y_min = min(np.percentile(cd_values, percentile_range[0]), np.percentile(cl_values, percentile_range[0])) + y_max = max(np.percentile(cd_values, percentile_range[1]), np.percentile(cl_values, percentile_range[1])) + padding = (y_max - y_min) * 0.1 + y_min, y_max = y_min - padding, y_max + padding + if use_log_scale: + y_min = max(y_min, 1e-6) + plt.figure(figsize=(10, 6)) + plt.plot(steps, cd_values, label="Drag Coefficient (Cd)", color="blue") + plt.plot(steps, cl_values, label="Lift Coefficient (Cl)", color="red") + plt.xlabel("Simulation Step") + plt.ylabel("Coefficient") + plt.title(f"{script_name}: Drag and Lift Coefficients Over Time") + plt.legend() + plt.grid(True) + plt.ylim(y_min, y_max) + if use_log_scale: + plt.yscale("log") + plt.savefig(os.path.join(output_dir, "drag_lift_plot.png")) + plt.close() + + +def compute_voxel_statistics(sim, bc_mask_exporter, sparsity_pattern, boundary_conditions, unit_convertor): + """ + Compute active/solid voxels, totals, lattice updates, and reference area based on simulation data. + """ + fields_data = bc_mask_exporter.get_fields_data({"bc_mask": sim.bc_mask}) + bc_mask_data = fields_data["bc_mask_0"] + level_id_field = bc_mask_exporter.level_id_field + + # Compute solid voxels per level (assuming 255 is the solid marker) + solid_voxels = [] + for lvl in range(num_levels): + level_mask = level_id_field == lvl + solid_voxels.append(np.sum(bc_mask_data[level_mask] == 255)) + + # Compute active voxels (total non-zero in sparsity minus solids) + active_voxels = [np.count_nonzero(mask) for mask in sparsity_pattern] + active_voxels = [max(0, active_voxels[lvl] - solid_voxels[lvl]) for lvl in range(num_levels)] + + # Totals + total_voxels = sum(active_voxels) + total_lattice_updates_per_step = sum(active_voxels[lvl] * (2 ** (num_levels - 1 - lvl)) for lvl in range(num_levels)) + + # Compute reference area (projected on YZ plane at finest level) + finest_level = 0 + mask_finest = level_id_field == finest_level + bc_mask_finest = bc_mask_data[mask_finest] + active_indices_finest = np.argwhere(sparsity_pattern[0]) + bc_body_id = boundary_conditions[-1].id # Assuming last BC is bc_body + solid_voxels_indices = active_indices_finest[bc_mask_finest == bc_body_id] + unique_jk = np.unique(solid_voxels_indices[:, 1:3], axis=0) + reference_area = unique_jk.shape[0] + reference_area_physical = reference_area * unit_convertor.reference_length**2 + + return { + "active_voxels": active_voxels, + "solid_voxels": solid_voxels, + "total_voxels": total_voxels, + "total_lattice_updates_per_step": total_lattice_updates_per_step, + "reference_area": reference_area, + "reference_area_physical": reference_area_physical, + } + + +def plot_data(x0, output_dir, delta_x_coarse, sim, IOexporter, prefix="Ahmed"): + """ + Ahmed Car Model, slant - angle = 25 degree + Profiles on symmetry plane (y=0) covering entire field + Origin of coordinate system: + x=0: end of the car, y=0: symmetry plane, z=0: ground plane + + S.Becker/H. Lienhart/C.Stoots + Insitute of Fluid Mechanics + University Erlangen-Nuremberg + Erlangen, Germany + Coordaintes in meters need to convert to voxels + Velocity data in m/s + """ + + def _load_sim_line(csv_path): + """ + Read a CSV exported by IOexporter.to_line without pandas. + Returns (z, Ux). + """ + # Read with header as column names + data = np.genfromtxt( + csv_path, + delimiter=",", + names=True, + autostrip=True, + dtype=None, + encoding="utf-8", + ) + if data.size == 0: + raise ValueError(f"No data in {csv_path}") + + z = np.asarray(data["z"], dtype=float) + ux = np.asarray(data["value"], dtype=float) + return z, ux + + # Load reference data + import json + + ref_data_path = "examples/cfd/data/ahmed.json" + with open(ref_data_path, "r") as file: + data = json.load(file) + + for x_str in data["data"].keys(): + # Extract reference horizontal velocity in m/s and its corresponding height in m + refX = np.array(data["data"][x_str]["x-velocity"]) + refY = np.array(data["data"][x_str]["height"]) + + # From reference x0 (rear of body) find x1 for plot + x_pos = float(x_str) + x1 = x0 + x_pos + + print(f" x1 is {x1}") + sim.macro(sim.f_0, sim.bc_mask, sim.rho, sim.u, streamId=0) + filename = os.path.join(output_dir, f"{prefix}_{x_str}") + wp.synchronize() + IOexporter.to_line( + filename, + {"velocity": sim.u}, + start_point=(x1, 0, 0), + end_point=(x1, 0, 0.8), + resolution=250, + component=0, + radius=delta_x_coarse, # needed with model units + ) + # read the CSV written by the exporter + csv_path = filename + "_velocity_0.csv" + print(f"CSV path is {csv_path}") + + try: + sim_z, sim_ux = _load_sim_line(csv_path) + except Exception as e: + print(f"Failed to read {csv_path}: {e}") + continue + + # plot reference vs simulation + plt.figure(figsize=(4.5, 6)) + plt.plot(refX, refY, "o", mfc="none", label="Experimental)") + plt.plot(sim_ux, sim_z, "-", lw=2, label="Simulation") + plt.xlim(np.min(refX) * 0.9, np.max(refX) * 1.1) + plt.ylim(np.min(refY), np.max(refY)) + plt.xlabel("Ux [m/s]") + plt.ylabel("z [m]") + plt.title(f"Velocity Plot at {x_pos:+.3f}") + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + plt.savefig(filename + ".png", dpi=150) + plt.close() + + +# Main Script +# =========== +# Initialize XLB + +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Generate mesh +level_data, body_vertices, grid_shape_zip, stl_shift, x0 = generate_cuboid_mesh(stl_filename, voxel_size) + +# Prepare the sparsity pattern and origins from the level data +sparsity_pattern, level_origins = prepare_sparsity_pattern(level_data) + +# Define a unit convertor +unit_convertor = UnitConvertor( + velocity_lbm_unit=wind_speed_lbm, + velocity_physical_unit=wind_speed_mps, + voxel_size_physical_unit=voxel_size, +) + +# Calculate lattice parameters +num_levels = len(level_data) +delta_x_coarse = voxel_size * 2 ** (num_levels - 1) +nu_lattice = unit_convertor.viscosity_to_lbm(kinematic_viscosity) +omega_finest = 1.0 / (3.0 * nu_lattice + 0.5) + +# Create output directory +current_dir = os.path.join(os.path.dirname(__file__)) +output_dir = os.path.join(current_dir, script_name) +if os.path.exists(output_dir): + shutil.rmtree(output_dir) +os.makedirs(output_dir) + +# Define exporter objects +field_name_cardinality_dict = {"velocity": 3, "density": 1} +h5exporter = MultiresIO( + field_name_cardinality_dict, + level_data, + offset=-stl_shift, + unit_convertor=unit_convertor, +) +bc_mask_exporter = MultiresIO( + {"bc_mask": 1}, + level_data, + offset=-stl_shift, + unit_convertor=unit_convertor, +) + +# Create grid +grid = multires_grid_factory( + grid_shape_zip, + velocity_set=velocity_set, + sparsity_pattern_list=sparsity_pattern, + sparsity_pattern_origins=[neon.Index_3d(*box_origin) for box_origin in level_origins], +) + +# Calculate num_steps +coarsest_level = grid.count_levels - 1 +grid_shape_x_coarsest = grid.level_to_shape(coarsest_level)[0] +num_steps = int(flow_passes * (grid_shape_x_coarsest / wind_speed_lbm)) + +# Calculate print and file output intervals +print_interval = max(1, int(num_steps * (print_interval_percentage / 100.0))) +crossover_step = int(num_steps * (file_output_crossover_percentage / 100.0)) +file_output_interval_pre_crossover = ( + max(1, int(crossover_step / num_file_outputs_pre_crossover)) if num_file_outputs_pre_crossover > 0 else num_steps + 1 +) +file_output_interval_post_crossover = ( + max(1, int((num_steps - crossover_step) / num_file_outputs_post_crossover)) if num_file_outputs_post_crossover > 0 else num_steps + 1 +) + +# Setup boundary conditions +boundary_conditions = setup_boundary_conditions(grid, level_data, body_vertices, wind_speed_mps) + +# Create initializer +wind_speed_lbm = unit_convertor.velocity_to_lbm(wind_speed_mps) +initializer = CustomMultiresInitializer( + bc_id=boundary_conditions[-2].id, # bc_outlet + constant_velocity_vector=(wind_speed_lbm, 0.0, 0.0), + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, +) + +# Initialize simulation +sim = initialize_simulation(grid, boundary_conditions, omega_finest, initializer) + +# Compute voxel statistics and reference area +stats = compute_voxel_statistics(sim, bc_mask_exporter, sparsity_pattern, boundary_conditions, unit_convertor) +active_voxels = stats["active_voxels"] +solid_voxels = stats["solid_voxels"] +total_voxels = stats["total_voxels"] +total_lattice_updates_per_step = stats["total_lattice_updates_per_step"] +reference_area = stats["reference_area"] +reference_area_physical = stats["reference_area_physical"] + +# Save initial bc_mask +filename = os.path.join(output_dir, f"{script_name}_initial_bc_mask") +try: + bc_mask_exporter.to_hdf5(filename, {"bc_mask": sim.bc_mask}, compression="gzip", compression_opts=0) + xmf_filename = f"{filename}.xmf" + hdf5_basename = f"{script_name}_initial_bc_mask.h5" +except Exception as e: + print(f"Error during initial bc_mask output: {e}") +wp.synchronize() + + +# Setup momentum transfer +momentum_transfer = MultiresMomentumTransfer( + boundary_conditions[-1], + mres_perf_opt=xlb.MresPerfOptimizationType.FUSION_AT_FINEST, + compute_backend=compute_backend, +) + +# Print simulation info +print("\n" + "=" * 50 + "\n") +print(f"Number of flow passes: {flow_passes}") +print(f"Calculated iterations: {num_steps:,}") +print(f"Finest voxel size: {voxel_size} meters") +print(f"Coarsest voxel size: {delta_x_coarse} meters") +print(f"Total voxels: {sum(np.count_nonzero(mask) for mask in sparsity_pattern):,}") +print(f"Total active voxels: {total_voxels:,}") +print(f"Active voxels per level: {[int(v) for v in active_voxels]}") +print(f"Solid voxels per level: {[int(v) for v in solid_voxels]}") +print(f"Total lattice updates per global step: {total_lattice_updates_per_step:,}") +print(f"Number of refinement levels: {num_levels}") +print(f"Physical inlet velocity: {wind_speed_mps:.4f} m/s") +print(f"Lattice velocity (ulb): {wind_speed_lbm}") +print(f"Computed reference area (bc_mask): {reference_area} lattice units") +print(f"Physical reference area (bc_mask): {reference_area_physical:.6f} m^2") +print("\n" + "=" * 50 + "\n") + +# -------------------------- Simulation Loop -------------------------- +wp.synchronize() +start_time = time.time() +compute_time = 0.0 +steps_since_last_print = 0 +drag_values = [] + +for step in range(num_steps): + step_start = time.time() + sim.step() + wp.synchronize() + compute_time += time.time() - step_start + steps_since_last_print += 1 + if step % print_interval == 0 or step == num_steps - 1: + sim.macro(sim.f_0, sim.bc_mask, sim.rho, sim.u, streamId=0) + wp.synchronize() + cd, cl, drag = compute_force_coefficients(sim, step, momentum_transfer, wind_speed_lbm, reference_area) + filename = os.path.join(output_dir, f"{script_name}_{step:04d}") + h5exporter.to_hdf5(filename, {"velocity": sim.u, "density": sim.rho}, compression="gzip", compression_opts=0) + h5exporter.to_slice_image( + filename, + {"velocity": sim.u}, + plane_point=(1, 0, 0), + plane_normal=(0, 1, 0), + grid_res=2000, + bounds=(0.25, 0.75, 0, 0.5), + show_axes=False, + show_colorbar=False, + slice_thickness=delta_x_coarse, # needed when using model units + ) + end_time = time.time() + elapsed = end_time - start_time + total_lattice_updates = total_lattice_updates_per_step * steps_since_last_print + MLUPS = total_lattice_updates / compute_time / 1e6 if compute_time > 0 else 0.0 + current_flow_passes = step * wind_speed_lbm / grid_shape_x_coarsest + remaining_steps = num_steps - step - 1 + time_remaining = 0.0 if MLUPS == 0 else (total_lattice_updates_per_step * remaining_steps) / (MLUPS * 1e6) + hours, rem = divmod(time_remaining, 3600) + minutes, seconds = divmod(rem, 60) + time_remaining_str = f"{int(hours):02d}h {int(minutes):02d}m {int(seconds):02d}s" + percent_complete = (step + 1) / num_steps * 100 + print(f"Completed step {step}/{num_steps} ({percent_complete:.2f}% complete)") + print(f" Flow Passes: {current_flow_passes:.2f}") + print(f" Time elapsed: {elapsed:.1f}s, Compute time: {compute_time:.1f}s, ETA: {time_remaining_str}") + print(f" MLUPS: {MLUPS:.1f}") + print(f" Cd={cd:.3f}, Cl={cl:.3f}, Drag Force (lattice units)={drag:.3f}") + start_time = time.time() + compute_time = 0.0 + steps_since_last_print = 0 + file_output_interval = file_output_interval_pre_crossover if step < crossover_step else file_output_interval_post_crossover + if step % file_output_interval == 0 or step == num_steps - 1: + sim.macro(sim.f_0, sim.bc_mask, sim.rho, sim.u, streamId=0) + filename = os.path.join(output_dir, f"{script_name}_{step:04d}") + try: + h5exporter.to_hdf5(filename, {"velocity": sim.u, "density": sim.rho}, compression="gzip", compression_opts=0) + xmf_filename = f"{filename}.xmf" + hdf5_basename = f"{script_name}_{step:04d}.h5" + except Exception as e: + print(f"Error during file output at step {step}: {e}") + wp.synchronize() + if step == num_steps - 1: + plot_data(x0, output_dir, delta_x_coarse, sim, h5exporter, prefix="Ahmed") + +# Save drag and lift data to CSV +if len(drag_values) > 0: + with open(os.path.join(output_dir, "drag_lift.csv"), "w") as fd: + fd.write("Step,Cd,Cl\n") + for i, (cd, cl) in enumerate(drag_values): + fd.write(f"{i * print_interval},{cd},{cl}\n") + plot_force_coefficients(drag_values, output_dir, print_interval, script_name) + +# Calculate and print average Cd and Cl for the last 50% +drag_values_array = np.array(drag_values) +if len(drag_values) > 0: + start_index = len(drag_values) // 2 + last_half = drag_values_array[start_index:, :] + avg_cd = np.mean(last_half[:, 0]) + avg_cl = np.mean(last_half[:, 1]) + print(f"Average Drag Coefficient (Cd) for last 50%: {avg_cd:.6f}") + print(f"Average Lift Coefficient (Cl) for last 50%: {avg_cl:.6f}") + print(f"Experimental Drag Coefficient (Cd): {0.3088}") + print(f"Error Drag Coefficient (Cd): {((avg_cd - 0.3088) / 0.3088) * 100:.2f}%") + +else: + print("No drag or lift data collected.") diff --git a/examples/cfd/rotating_sphere_3d.py b/examples/cfd/rotating_sphere_3d.py new file mode 100644 index 00000000..184408c3 --- /dev/null +++ b/examples/cfd/rotating_sphere_3d.py @@ -0,0 +1,327 @@ +""" +Rotating sphere 3-D example (single-resolution). + +Simulates flow past a sphere rotating about the y-axis using the +halfway bounce-back BC with a prescribed rotational-velocity profile. +Computes drag and lift coefficients over time and saves VTK snapshots. +""" + +import xlb +import trimesh +import time +import warp as wp +import numpy as np +import jax.numpy as jnp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import ( + HalfwayBounceBackBC, + FullwayBounceBackBC, + RegularizedBC, + DoNothingBC, + HybridBC, +) +from xlb.operator.force.momentum_transfer import MomentumTransfer +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_fields_vtk, save_image, warp_array_to_jax +import matplotlib.pyplot as plt +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator import Operator +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.operator.boundary_masker import MeshVoxelizationMethod + +# -------------------------- Simulation Setup -------------------------- + +# Grid parameters +wp.clear_kernel_cache() +diam = 32 +grid_size_x, grid_size_y, grid_size_z = 10 * diam, 7 * diam, 7 * diam +grid_shape = (grid_size_x, grid_size_y, grid_size_z) + +# Simulation Configuration +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 + +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +wind_speed = 0.04 +num_steps = 100000 +print_interval = 1000 +post_process_interval = 1000 + +# Physical Parameters +Re = 200.0 +visc = wind_speed * diam / Re +omega = 1.0 / (3.0 * visc + 0.5) + +# Rotational speed parameters (see [1] which discusses the problem in terms of 2 non-dimensional parameters: Re and Omega) +# [1] J. Fluid Mech. (2016), vol. 807, pp. 62–86. cΒ© Cambridge University Press 2016 doi:10.1017/jfm.2016.596 +# \Omega = \omega * D / (2 U_\infty) where Omega is non-dimensional and omega is dimensional. +rot_rate_nondim = -0.2 +rot_rate = 2.0 * wind_speed * rot_rate_nondim / diam + +# Print simulation info +print("\n" + "=" * 50 + "\n") +print("Simulation Configuration:") +print(f"Grid size: {grid_size_x} x {grid_size_y} x {grid_size_z}") +print(f"Backend: {compute_backend}") +print(f"Velocity set: {velocity_set}") +print(f"Precision policy: {precision_policy}") +print(f"Prescribed velocity: {wind_speed}") +print(f"Reynolds number: {Re}") +print(f"Max iterations: {num_steps}") +print("\n" + "=" * 50 + "\n") + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Create Grid +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +# Bounding box indices +box = grid.bounding_box_indices() +box_no_edge = grid.bounding_box_indices(remove_edges=True) +inlet = box_no_edge["left"] +outlet = box["right"] +walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] +walls = np.unique(np.array(walls), axis=-1).tolist() + +# Load the mesh (replace with your own mesh) +stl_filename = "../stl-files/sphere.stl" +mesh = trimesh.load_mesh(stl_filename, process=False) +mesh_vertices = mesh.vertices + +# Transform the mesh points to be located in the right position in the wind tunnel +mesh_vertices -= mesh_vertices.min(axis=0) +mesh_extents = mesh_vertices.max(axis=0) +length_phys_unit = mesh_extents.max() +length_lbm_unit = grid_shape[1] / 7 +dx = length_phys_unit / length_lbm_unit +mesh_vertices = mesh_vertices / dx +shift = np.array([grid_shape[0] / 3, (grid_shape[1] - mesh_extents[1] / dx) / 2, (grid_shape[2] - mesh_extents[2] / dx) / 2]) +sphere = mesh_vertices + shift +diam = np.max(sphere.max(axis=0) - sphere.min(axis=0)) +sphere_cross_section = np.pi * diam**2 / 4.0 + + +# Define rotating boundary profile +def bc_profile(): + """Build a Warp function returning the rotational wall velocity at a voxel.""" + dtype = precision_policy.compute_precision.wp_dtype + _u_vec = wp.vec(velocity_set.d, dtype=dtype) + angular_velocity = _u_vec(0.0, rot_rate, 0.0) + origin_np = shift + diam / 2 + origin_wp = _u_vec(origin_np[0], origin_np[1], origin_np[2]) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + x = dtype(index[0]) + y = dtype(index[1]) + z = dtype(index[2]) + surface_coord = _u_vec(x, y, z) - origin_wp + return wp.cross(angular_velocity, surface_coord) + + return bc_profile_warp + + +# Define boundary conditions +bc_left = RegularizedBC("velocity", prescribed_value=(wind_speed, 0.0, 0.0), indices=inlet) +bc_do_nothing = DoNothingBC(indices=outlet) +# bc_sphere = HalfwayBounceBackBC(mesh_vertices=sphere, voxelization_method="ray", profile=bc_profile()) +bc_sphere = HybridBC( + bc_method="nonequilibrium_regularized", + mesh_vertices=sphere, + use_mesh_distance=True, + voxelization_method=MeshVoxelizationMethod("RAY"), + profile=bc_profile(), +) +# Not assining BC for walls makes them periodic. +boundary_conditions = [bc_left, bc_do_nothing, bc_sphere] + + +# Setup Stepper +stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", +) + +# Make initializer operator +from xlb.helper.initializers import CustomInitializer + +initializer = CustomInitializer( + bc_id=bc_do_nothing.id, + constant_velocity_vector=(wind_speed, 0.0, 0.0), + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, +) + +# Prepare Fields +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields(initializer=initializer) + + +# -------------------------- Helper Functions -------------------------- + + +def plot_coefficient(time_steps, coefficients, prefix="drag"): + """ + Plot the drag coefficient with various moving averages. + + Args: + time_steps (list): List of time steps. + coefficients (list): List of force coefficients. + """ + # Convert lists to numpy arrays for processing + time_steps_np = np.array(time_steps) + coefficients_np = np.array(coefficients) + + # Define moving average windows + windows = [10, 100, 1000, 10000, 100000] + labels = ["MA 10", "MA 100", "MA 1,000", "MA 10,000", "MA 100,000"] + + plt.figure(figsize=(12, 8)) + plt.plot(time_steps_np, coefficients_np, label="Raw", alpha=0.5) + + for window, label in zip(windows, labels): + if len(coefficients_np) >= window: + ma = np.convolve(coefficients_np, np.ones(window) / window, mode="valid") + plt.plot(time_steps_np[window - 1 :], ma, label=label) + + plt.ylim(-1.0, 1.0) + plt.legend() + plt.xlabel("Time step") + plt.ylabel("Drag coefficient") + plt.title("Drag Coefficient Over Time with Moving Averages") + plt.savefig(prefix + "_ma.png") + plt.close() + + +def post_process( + step, + f_0, + f_1, + grid_shape, + macro, + momentum_transfer, + missing_mask, + bc_mask, + wind_speed, + car_cross_section, + drag_coefficients, + lift_coefficients, + time_steps, +): + """Compute macroscopic fields, force coefficients, and save VTK output.""" + """ + Post-process simulation data: save fields, compute forces, and plot drag coefficient. + + Args: + step (int): Current time step. + f_current: Current distribution function. + grid_shape (tuple): Shape of the grid. + macro: Macroscopic operator object. + momentum_transfer: MomentumTransfer operator object. + missing_mask: Missing mask from stepper. + bc_mask: Boundary condition mask from stepper. + wind_speed (float): Prescribed wind speed. + car_cross_section (float): Cross-sectional area of the car. + drag_coefficients (list): List to store drag coefficients. + lift_coefficients (list): List to store lift coefficients. + time_steps (list): List to store time steps. + """ + wp.synchronize() + # Convert to JAX array if necessary + if not isinstance(f_0, jnp.ndarray): + f_0_jax = warp_array_to_jax(f_0) + else: + f_0_jax = f_0 + + # Compute macroscopic quantities + rho, u = macro(f_0_jax) + + # Remove boundary cells + u = u[:, 1:-1, 1:-1, 1:-1] + u_magnitude = jnp.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + + fields = {"ux": u[0], "uy": u[1], "uz": u[2], "u_magnitude": u_magnitude} + + # Save fields in VTK format + # save_fields_vtk(fields, timestep=step) + + # Save the u_magnitude slice at the mid y-plane + mid_y = grid_shape[1] // 2 + save_image(fields["u_magnitude"][:, mid_y, :], timestep=step) + + # Compute lift and drag + boundary_force = momentum_transfer(f_0, f_1, bc_mask, missing_mask) + drag = boundary_force[0] # x-direction + lift = boundary_force[2] + cd = 2.0 * drag / (wind_speed**2 * car_cross_section) + cl = 2.0 * lift / (wind_speed**2 * car_cross_section) + print(f"CD={cd}, CL={cl}") + drag_coefficients.append(cd) + lift_coefficients.append(cl) + time_steps.append(step) + + # Plot drag coefficient + plot_coefficient(time_steps, drag_coefficients, prefix="drag") + plot_coefficient(time_steps, lift_coefficients, prefix="lift") + + +# Setup Momentum Transfer for Force Calculation +bc_car = boundary_conditions[-1] +momentum_transfer = MomentumTransfer(bc_car, compute_backend=compute_backend) + +# Define Macroscopic Calculation +macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), +) + +# Initialize Lists to Store Coefficients and Time Steps +time_steps = [] +drag_coefficients = [] +lift_coefficients = [] + +# -------------------------- Simulation Loop -------------------------- + +start_time = time.time() +for step in range(num_steps): + # Perform simulation step + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 # Swap the buffers + + # Print progress at intervals + if step % print_interval == 0: + elapsed_time = time.time() - start_time + print(f"Iteration: {step}/{num_steps} | Time elapsed: {elapsed_time:.2f}s") + start_time = time.time() + + # Post-process at intervals and final step + if (step % post_process_interval == 0) or (step == num_steps - 1): + post_process( + step, + f_0, + f_1, + grid_shape, + macro, + momentum_transfer, + missing_mask, + bc_mask, + wind_speed, + sphere_cross_section, + drag_coefficients, + lift_coefficients, + time_steps, + ) + +print("Simulation completed successfully.") diff --git a/examples/cfd/turbulent_channel_3d.py b/examples/cfd/turbulent_channel_3d.py new file mode 100644 index 00000000..e7e32445 --- /dev/null +++ b/examples/cfd/turbulent_channel_3d.py @@ -0,0 +1,217 @@ +import xlb +import time +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import RegularizedBC +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_fields_vtk, save_image, warp_array_to_jax +from xlb.helper import initialize_eq +import warp as wp +import numpy as np +import jax.numpy as jnp +import matplotlib.pyplot as plt +import json + + +# -------------------------- Helper Functions -------------------------- + + +def vonKarman_loglaw_wall(yplus): + vonKarmanConst = 0.41 + cplus = 5.5 + uplus = np.log(yplus) / vonKarmanConst + cplus + return uplus + + +def get_dns_data(): + """ + Reference: DNS of Turbulent Channel Flow up to Re_tau=590, 1999, + Physics of Fluids, vol 11, 943-945. + https://turbulence.oden.utexas.edu/data/MKM/chan180/profiles/chan180.means + """ + file_name = "examples/cfd/data/turbulent_channel_dns_data.json" + with open(file_name, "r") as file: + return json.load(file) + + +# -------------------------- Simulation Setup -------------------------- + +# Channel Parameter +channel_half_width = 50 + +# Define channel geometry based on h +grid_size_x = 6 * channel_half_width +grid_size_y = 3 * channel_half_width +grid_size_z = 2 * channel_half_width + +# Grid parameters +grid_shape = (grid_size_x, grid_size_y, grid_size_z) + +# Define flow regime +Re_tau = 180 +u_tau = 0.001 + +# Compute viscosity and relaxation parameter omega +visc = u_tau * channel_half_width / Re_tau +omega = 1.0 / (3.0 * visc + 0.5) + +# Runtime & compute_backend configurations +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP64FP64 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +num_steps = 10000000 +print_interval = 100000 +post_process_interval = 100000 + +# Print simulation info +print("\n" + "=" * 50 + "\n") +print("Simulation Configuration:") +print(f"Grid size: {grid_size_x} x {grid_size_y} x {grid_size_z}") +print(f"Backend: {compute_backend}") +print(f"Velocity set: {velocity_set}") +print(f"Precision policy: {precision_policy}") +print(f"Reynolds number: {Re_tau}") +print(f"Max iterations: {num_steps}") +print("\n" + "=" * 50 + "\n") + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Create Grid +grid = grid_factory(grid_shape, compute_backend=compute_backend) + + +# Define Force Vector +def get_force(Re_tau, visc, channel_half_width, velocity_set): + shape = (velocity_set.d,) + force = np.zeros(shape) + force[0] = Re_tau**2 * visc**2 / channel_half_width**3 + return force + + +force_vector = get_force(Re_tau, visc, channel_half_width, velocity_set) + + +# Define Boundary Indices +box = grid.bounding_box_indices(remove_edges=True) +walls = [box["bottom"][i] + box["top"][i] for i in range(velocity_set.d)] + + +# Define Boundary Conditions +def setup_boundary_conditions(walls, velocity_set, precision_policy): + # No-slip boundary condition: velocity = (0, 0, 0) + bc_walls = RegularizedBC("velocity", prescribed_value=(0.0, 0.0, 0.0), indices=walls) + boundary_conditions = [bc_walls] + return boundary_conditions + + +boundary_conditions = setup_boundary_conditions(walls, velocity_set, precision_policy) + +# Setup Stepper +stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + force_vector=force_vector, +) + +# Prepare Fields +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + + +# Initialize Fields with Random Velocity +shape = (velocity_set.d,) + grid.shape +np.random.seed(0) +u_init = np.random.random(shape) +if compute_backend == ComputeBackend.JAX: + u_init = jnp.full(shape=shape, fill_value=1e-2 * u_init) +else: + u_init = wp.array(1e-2 * u_init, dtype=precision_policy.compute_precision.wp_dtype) + +f_0 = initialize_eq(f_0, grid, velocity_set, precision_policy, compute_backend, u=u_init) + +# Define Macroscopic Calculation +macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), +) + + +# Post-Processing Function +def post_process(step, f_current, grid_shape, macro): + # Convert to JAX array if necessary + if not isinstance(f_current, jnp.ndarray): + f_current = warp_array_to_jax(f_current) + + rho, u = macro(f_current) + + # Compute velocity magnitude + u_magnitude = jnp.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + fields = { + "rho": rho[0], + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + "u_magnitude": u_magnitude, + } + + # Save the fields in VTK format + save_fields_vtk(fields, timestep=step) + + # Save the u_magnitude slice at the mid y-plane + mid_y = grid_shape[1] // 2 + save_image(fields["u_magnitude"][:, mid_y, :], timestep=step) + + # Save monitor plot + plot_uplus(u, step, grid_shape, u_tau, visc) + + +# Plotting Function +def plot_uplus(u, timestep, grid_shape, u_tau, visc): + # Mean streamwise velocity in wall units u^+(z) + zz = np.arange(grid_shape[-1]) + zz = np.minimum(zz, zz.max() - zz) + yplus = zz * u_tau / visc + uplus = np.mean(u[0], axis=(0, 1)) / u_tau + uplus_loglaw = vonKarman_loglaw_wall(yplus) + dns_dic = get_dns_data() + + plt.clf() + plt.semilogx(yplus, uplus, "r.", label="Simulation") + plt.semilogx(yplus, uplus_loglaw, "k:", label="Von Karman Log Law") + plt.semilogx(dns_dic["y+"], dns_dic["Umean"], "b-", label="DNS Data") + ax = plt.gca() + ax.set_xlim([0.1, 300]) + ax.set_ylim([0, 20]) + plt.xlabel("y+") + plt.ylabel("U+") + plt.title(f"u+ vs y+ at timestep {timestep}") + plt.legend() + fname = f"uplus_{str(timestep).zfill(5)}.png" + plt.savefig(fname, format="png") + plt.close() + + +# -------------------------- Simulation Loop -------------------------- + +start_time = time.time() +for step in range(num_steps): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 # Swap the buffers + + if step % print_interval == 0: + if compute_backend == ComputeBackend.WARP: + wp.synchronize() + elapsed_time = time.time() - start_time + print(f"Iteration: {step}/{num_steps} | Time elapsed: {elapsed_time:.2f}s") + start_time = time.time() + + if (step % post_process_interval == 0) or (step == num_steps - 1): + post_process(step, f_0, grid_shape, macro) diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py new file mode 100644 index 00000000..9909132c --- /dev/null +++ b/examples/cfd/windtunnel_3d.py @@ -0,0 +1,285 @@ +import xlb +import trimesh +import time +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import ( + HalfwayBounceBackBC, + FullwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, + HybridBC, +) +from xlb.operator.force.momentum_transfer import MomentumTransfer +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_fields_vtk, save_image +import warp as wp +import numpy as np +import jax.numpy as jnp +import matplotlib.pyplot as plt +from xlb.operator.boundary_masker import MeshVoxelizationMethod + + +# -------------------------- Simulation Setup -------------------------- + +# Grid parameters +grid_size_x, grid_size_y, grid_size_z = 512, 128, 128 +grid_shape = (grid_size_x, grid_size_y, grid_size_z) + +# Simulation Configuration +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 + +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +wind_speed = 0.02 +num_steps = 100000 +print_interval = 1000 +post_process_interval = 1000 + +# Physical Parameters +Re = 50000.0 +clength = grid_size_x - 1 +visc = wind_speed * clength / Re +omega = 1.0 / (3.0 * visc + 0.5) + +# Print simulation info +print("\n" + "=" * 50 + "\n") +print("Simulation Configuration:") +print(f"Grid size: {grid_size_x} x {grid_size_y} x {grid_size_z}") +print(f"Backend: {compute_backend}") +print(f"Velocity set: {velocity_set}") +print(f"Precision policy: {precision_policy}") +print(f"Prescribed velocity: {wind_speed}") +print(f"Reynolds number: {Re}") +print(f"Max iterations: {num_steps}") +print("\n" + "=" * 50 + "\n") + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Create Grid +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +# Bounding box indices +box = grid.bounding_box_indices() +box_no_edge = grid.bounding_box_indices(remove_edges=True) +inlet = box_no_edge["left"] +outlet = box_no_edge["right"] +walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] +walls = np.unique(np.array(walls), axis=-1).tolist() + +# Load the mesh (replace with your own mesh) +stl_filename = "../stl-files/DrivAer-Notchback.stl" +voxelization_method = MeshVoxelizationMethod("RAY") +mesh = trimesh.load_mesh(stl_filename, process=False) +mesh_vertices = mesh.vertices + +# Transform the mesh points to align with the grid +mesh_vertices -= mesh_vertices.min(axis=0) +mesh_extents = mesh_vertices.max(axis=0) +length_phys_unit = mesh_extents.max() +length_lbm_unit = grid_shape[0] / 4 +dx = length_phys_unit / length_lbm_unit +mesh_vertices = mesh_vertices / dx + +# Depending on the voxelization method, shift_z ensures the bottom ground does not intersect with the voxelized mesh +# Any smaller shift value would lead to large lift computations due to the initial equilibrium distributions. Bigger +# values would be fine but leave a gap between surfaces that are supposed to touch. +if voxelization_method in (MeshVoxelizationMethod("RAY"), MeshVoxelizationMethod("WINDING")): + shift_z = 2 +elif voxelization_method in (MeshVoxelizationMethod("AABB"), MeshVoxelizationMethod("AABB_CLOSE", close_voxels=3)): + shift_z = 3 +shift = np.array([grid_shape[0] / 4, (grid_shape[1] - mesh_extents[1] / dx) / 2, shift_z]) +car_vertices = mesh_vertices + shift +car_cross_section = np.prod(mesh_extents[1:]) / dx**2 + + +bc_left = RegularizedBC("velocity", prescribed_value=(wind_speed, 0.0, 0.0), indices=inlet) +bc_walls = FullwayBounceBackBC(indices=walls) +bc_do_nothing = ExtrapolationOutflowBC(indices=outlet) +bc_car = HalfwayBounceBackBC(mesh_vertices=car_vertices, voxelization_method=voxelization_method) +# bc_car = HybridBC(bc_method="nonequilibrium_regularized", mesh_vertices=car_vertices, +# voxelization_method=voxelization_method, use_mesh_distance=True) +boundary_conditions = [bc_walls, bc_left, bc_do_nothing, bc_car] + + +# Configure backend options: +# backend_config = {"occ": neon.SkeletonConfig.OCC.from_string("standard"), "device_list": [0, 1]} if compute_backend == ComputeBackend.NEON else {} +backend_config = {} + +# Setup Stepper +stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + backend_config=backend_config, +) + +# Prepare Fields +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + + +# -------------------------- Helper Functions -------------------------- + + +def plot_coefficient(time_steps, coefficients, prefix="drag"): + """ + Plot the drag coefficient with various moving averages. + + Args: + time_steps (list): List of time steps. + coefficients (list): List of force coefficients. + """ + # Convert lists to numpy arrays for processing + time_steps_np = np.array(time_steps) + coefficients_np = np.array(coefficients) + + # Define moving average windows + windows = [10, 100, 1000, 10000, 100000] + labels = ["MA 10", "MA 100", "MA 1,000", "MA 10,000", "MA 100,000"] + + plt.figure(figsize=(12, 8)) + plt.plot(time_steps_np, coefficients_np, label="Raw", alpha=0.5) + + for window, label in zip(windows, labels): + if len(coefficients_np) >= window: + ma = np.convolve(coefficients_np, np.ones(window) / window, mode="valid") + plt.plot(time_steps_np[window - 1 :], ma, label=label) + + plt.ylim(-1.0, 1.0) + plt.legend() + plt.xlabel("Time step") + plt.ylabel("Drag coefficient") + plt.title("Drag Coefficient Over Time with Moving Averages") + plt.savefig(prefix + "_ma.png") + plt.close() + + +def post_process( + step, + f_0, + f_1, + grid_shape, + macro, + momentum_transfer, + missing_mask, + bc_mask, + wind_speed, + car_cross_section, + drag_coefficients, + lift_coefficients, + time_steps, +): + """ + Post-process simulation data: save fields, compute forces, and plot drag coefficient. + + Args: + step (int): Current time step. + f_current: Current distribution function. + grid_shape (tuple): Shape of the grid. + macro: Macroscopic operator object. + momentum_transfer: MomentumTransfer operator object. + missing_mask: Missing mask from stepper. + bc_mask: Boundary condition mask from stepper. + wind_speed (float): Prescribed wind speed. + car_cross_section (float): Cross-sectional area of the car. + drag_coefficients (list): List to store drag coefficients. + lift_coefficients (list): List to store lift coefficients. + time_steps (list): List to store time steps. + """ + # Convert to JAX array if necessary + if not isinstance(f_0, jnp.ndarray): + f_0_jax = to_jax(f_0) + else: + f_0_jax = f_0 + + # Compute macroscopic quantities + rho, u = macro(f_0_jax) + + # Remove boundary cells + u = u[:, 1:-1, 1:-1, 1:-1] + u_magnitude = jnp.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + + fields = {"u_magnitude": u_magnitude} + + # Save fields in VTK format + save_fields_vtk(fields, timestep=step) + + # Save the u_magnitude slice at the mid y-plane + mid_y = grid_shape[1] // 2 + save_image(fields["u_magnitude"][:, mid_y, :], timestep=step) + + # Compute lift and drag + boundary_force = momentum_transfer(f_0, f_1, bc_mask, missing_mask) + drag = boundary_force[0] # x-direction + lift = boundary_force[2] + cd = 2.0 * drag / (wind_speed**2 * car_cross_section) + cl = 2.0 * lift / (wind_speed**2 * car_cross_section) + print(f"CD={cd}, CL={cl}") + drag_coefficients.append(cd) + lift_coefficients.append(cl) + time_steps.append(step) + + # Plot drag coefficient + plot_coefficient(time_steps, drag_coefficients, prefix="drag") + plot_coefficient(time_steps, lift_coefficients, prefix="lift") + + +# Setup Momentum Transfer for Force Calculation +bc_car = boundary_conditions[-1] +momentum_transfer = MomentumTransfer(bc_car, compute_backend=compute_backend) + +# Define Macroscopic Calculation +macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), +) +to_jax = xlb.utils.ToJAX("populations", velocity_set.q, grid_shape) + +# Initialize Lists to Store Coefficients and Time Steps +time_steps = [] +drag_coefficients = [] +lift_coefficients = [] + +# -------------------------- Simulation Loop -------------------------- + +start_time = time.time() +for step in range(num_steps): + # Perform simulation step + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 # Swap the buffers + + # Print progress at intervals + if step % print_interval == 0: + if compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + wp.synchronize() + elapsed_time = time.time() - start_time + print(f"Iteration: {step}/{num_steps} | Time elapsed: {elapsed_time:.2f}s") + start_time = time.time() + + # Post-process at intervals and final step + if (step % post_process_interval == 0) or (step == num_steps - 1): + post_process( + step, + f_0, + f_1, + grid_shape, + macro, + momentum_transfer, + missing_mask, + bc_mask, + wind_speed, + car_cross_section, + drag_coefficients, + lift_coefficients, + time_steps, + ) + +print("Simulation completed successfully.") diff --git a/examples/ibm/airfoil_ibm.py b/examples/ibm/airfoil_ibm.py new file mode 100644 index 00000000..62264a22 --- /dev/null +++ b/examples/ibm/airfoil_ibm.py @@ -0,0 +1,363 @@ +import numpy as np +import trimesh +import jax.numpy as jnp +import matplotlib.pyplot as plt +import warp as wp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.operator.stepper import IBMStepper +from xlb.operator.boundary_condition import FullwayBounceBackBC, RegularizedBC, ExtrapolationOutflowBC +from xlb.operator.macroscopic import Macroscopic +from xlb.helper.ibm_helper import prepare_immersed_boundary +from xlb.grid import grid_factory +from xlb.utils import save_image, warp_array_to_jax + + +def generate_naca_profile(chord_length, thickness_ratio, n_points=400): + x = np.linspace(0.0, chord_length, n_points) + x_c = x / chord_length + coeffs = np.array([0.2969, -0.1260, -0.3516, 0.2843, -0.1015], dtype=np.float64) + powers = np.array([0.5, 1.0, 2.0, 3.0, 4.0], dtype=np.float64) + terms = np.stack([x_c**p for p in powers], axis=0) + thickness = 5.0 * thickness_ratio * chord_length * np.tensordot(coeffs, terms, axes=1) + upper = np.stack([x, thickness], axis=1) + lower = np.stack([x[::-1], -thickness[::-1]], axis=1) + profile = np.vstack([upper, lower[1:-1]]) + profile[:, 0] -= chord_length * 0.5 + return profile + + +def extrude_profile_to_mesh(profile, span_length): + lower_z = -0.5 * span_length + upper_z = 0.5 * span_length + lower = np.concatenate([profile, np.full((profile.shape[0], 1), lower_z)], axis=1) + upper = np.concatenate([profile, np.full((profile.shape[0], 1), upper_z)], axis=1) + vertices = np.vstack([lower, upper]) + faces = [] + n = profile.shape[0] + for i in range(1, n - 1): + faces.append([0, i + 1, i]) + top_offset = n + for i in range(1, n - 1): + faces.append([top_offset, top_offset + i, top_offset + i + 1]) + for i in range(n): + j = (i + 1) % n + faces.append([i, j, top_offset + j]) + faces.append([i, top_offset + j, top_offset + i]) + return trimesh.Trimesh(vertices=vertices, faces=np.array(faces, dtype=np.int64), process=False) + + +def create_airfoil_mesh(chord_length, thickness_ratio, span_length, n_points=400): + profile = generate_naca_profile(chord_length, thickness_ratio, n_points) + mesh = extrude_profile_to_mesh(profile, span_length) + return mesh + + +def define_boundary_indices(grid, velocity_set): + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + inlet = box_no_edge["left"] + outlet = box_no_edge["right"] + walls = [box["front"][i] + box["back"][i] + box["top"][i] + box["bottom"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + return inlet, outlet, walls + + +def bc_profile(precision_policy, grid_shape, u_max): + dtype = precision_policy.store_precision.wp_dtype + u_max_d = dtype(u_max) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + return wp.vec(dtype(u_max_d), length=1) + + return bc_profile_warp + + +def setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max): + inlet, outlet, walls = define_boundary_indices(grid, velocity_set) + bc_inlet = RegularizedBC("velocity", indices=inlet, profile=bc_profile(precision_policy, grid_shape, u_max)) + bc_outlet = ExtrapolationOutflowBC(indices=outlet) + bc_walls = FullwayBounceBackBC(indices=walls) + return [bc_walls, bc_inlet, bc_outlet] + + +def setup_stepper(grid, boundary_conditions, ibm_max_iterations=2, ibm_tolerance=1e-5, ibm_relaxation=1.0): + return IBMStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + ibm_max_iterations=ibm_max_iterations, + ibm_tolerance=ibm_tolerance, + ibm_relaxation=ibm_relaxation, + ) + + +def calculate_force_coefficients(lag_forces, areas_np, reference_velocity, reference_area): + forces_np = lag_forces.numpy() + weighted = forces_np * areas_np[:, None] + total_force = -np.sum(weighted, axis=0) + dynamic_pressure = 0.5 * reference_velocity**2 + denom = dynamic_pressure * reference_area if dynamic_pressure * reference_area != 0.0 else 1.0 + cd = total_force[0] / denom + cl = total_force[1] / denom + return cd, cl, total_force + + +def post_process( + step, + post_process_interval, + f_current, + precision_policy, + grid_shape, + lag_forces, + cd_values, + cl_values, + reference_velocity, + reference_area, + areas_np, +): + if not isinstance(f_current, jnp.ndarray): + f_jax = warp_array_to_jax(f_current) + else: + f_jax = f_current + macro_jax = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + ) + rho, u = macro_jax(f_jax) + u = u[:, 1:-1, 1:-1, 1:-1] + fields = { + "u_magnitude": np.sqrt(u[0] ** 2.0 + u[1] ** 2.0 + u[2] ** 2.0), + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + } + slice_idz = grid_shape[2] // 2 + save_image(fields["u_magnitude"][:, :, slice_idz], timestep=step) + cd, cl, total_force = calculate_force_coefficients(lag_forces, areas_np, reference_velocity, reference_area) + cd_values.append((step, float(cd))) + cl_values.append((step, float(cl))) + if step % post_process_interval == 0: + window = 10 + if len(cd_values) >= window: + avg_cd = float(np.mean([v for _, v in cd_values[-window:]])) + avg_cl = float(np.mean([v for _, v in cl_values[-window:]])) + else: + avg_cd = float(np.mean([v for _, v in cd_values])) + avg_cl = float(np.mean([v for _, v in cl_values])) + print( + f"Step {step}: Cd = {cd:.6f}, Cl = {cl:.6f}, Cd(avg{window}) = {avg_cd:.6f}, Cl(avg{window}) = {avg_cl:.6f}, " + f"Fx = {total_force[0]:.6f}, Fy = {total_force[1]:.6f}" + ) + + +def save_force_coefficients(cd_values, cl_values, filename): + with open(filename, "w") as f: + f.write("timestep,cd,cl\n") + for (timestep_cd, cd), (_, cl) in zip(cd_values, cl_values): + f.write(f"{timestep_cd},{cd},{cl}\n") + timesteps = [t for t, _ in cd_values] + cds = [cd for _, cd in cd_values] + cls = [cl for _, cl in cl_values] + plt.figure(figsize=(10, 6)) + plt.plot(timesteps, cds, "r-", label="Cd") + plt.plot(timesteps, cls, "b-", label="Cl") + plt.grid(True, linestyle="--", alpha=0.7) + plt.xlabel("Timestep") + plt.ylabel("Coefficient") + plt.title("Airfoil Force Coefficients") + plt.legend() + plt.tight_layout() + plt.savefig("airfoil_force_coefficients.png", dpi=150) + plt.close() + + +@wp.kernel +def update_airfoil_pose( + step: int, + total_steps: int, + start_angle: float, + total_rotation: float, + origin: wp.vec3, + base_vertices: wp.array(dtype=wp.vec3), + vertices: wp.array(dtype=wp.vec3), + velocities: wp.array(dtype=wp.vec3), +): + idx = wp.tid() + total_span = wp.float32(total_steps - 1) + progress = wp.float32(0.0) + if total_span > 0.0: + progress = wp.float32(step) / total_span + if progress > 1.0: + progress = wp.float32(1.0) + start_angle_f = wp.float32(start_angle) + total_rotation_f = wp.float32(total_rotation) + angle = start_angle_f + total_rotation_f * progress + c = wp.cos(angle) + s = wp.sin(angle) + base = base_vertices[idx] - origin + rotated = wp.vec3( + c * base[0] - s * base[1], + s * base[0] + c * base[1], + base[2], + ) + vertices[idx] = rotated + origin + angular_rate = wp.float32(0.0) + if total_span > 0.0: + angular_rate = total_rotation_f / total_span + velocities[idx] = wp.vec3( + -angular_rate * rotated[1], + angular_rate * rotated[0], + 0.0, + ) + + +if __name__ == "__main__": + chord_length = 60.0 * 1.3 + span_length = 50.0 * 1.3 + thickness_ratio = 0.12 + upstream = int(2 * chord_length) + downstream = int(4 * chord_length) + ly = int(3.0 * chord_length) + lz = int(2.0 * span_length) + lx = upstream + downstream + int(chord_length) + grid_shape = (lx, ly, lz) + u_max = 0.05 + Re = 20000 + start_angle_deg = 0.0 + total_rotation_deg = -45.0 + start_angle_rad = np.deg2rad(start_angle_deg) + total_rotation_rad = np.deg2rad(total_rotation_deg) + num_steps = 30000 + post_process_interval = 100 + print_interval = 100 + ibm_max_iterations = 1 + ibm_tolerance = 1e-5 + ibm_relaxation = 0.5 + compute_backend = ComputeBackend.WARP + precision_policy = PrecisionPolicy.FP32FP32 + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) + xlb.init(velocity_set=velocity_set, default_backend=compute_backend, default_precision_policy=precision_policy) + grid = grid_factory(grid_shape, compute_backend=compute_backend) + print("Airfoil IBM Simulation Configuration:") + print(f" Grid size: {grid_shape}") + print(f" Chord length: {chord_length}") + print(f" Span length: {span_length}") + print(f" Thickness ratio: {thickness_ratio}") + print(f" Inlet velocity: {u_max}") + print(f" Reynolds number: {Re}") + print(f" Start angle: {start_angle_deg}") + print(f" Total rotation: {total_rotation_deg}") + print(f" Max steps: {num_steps}") + print(f" IBM max iterations: {ibm_max_iterations}") + print(f" IBM tolerance: {ibm_tolerance}") + print(f" IBM relaxation: {ibm_relaxation}") + airfoil_mesh = create_airfoil_mesh(chord_length, thickness_ratio, span_length) + airfoil_center = np.array([float(upstream + 0.6 * chord_length), grid_shape[1] * 0.5, grid_shape[2] * 0.5], dtype=np.float64) + translation = airfoil_center - airfoil_mesh.centroid + airfoil_mesh.apply_translation(translation) + vertices_wp, areas_wp, faces_np = prepare_immersed_boundary(airfoil_mesh, max_lbm_length=max(chord_length, span_length)) + vertices_np = vertices_wp.numpy() + base_vertices_wp = wp.array(vertices_np, dtype=wp.vec3) + vertices_wp = wp.array(vertices_np, dtype=wp.vec3) + areas_np = areas_wp.numpy() + leading_edge_x = float(np.min(vertices_np[:, 0])) + rotation_center_y = float(np.mean(vertices_np[:, 1])) + rotation_center_z = float(np.mean(vertices_np[:, 2])) + rotation_origin = np.array( + [ + leading_edge_x + 0.1 * chord_length, + rotation_center_y, + rotation_center_z, + ], + dtype=np.float64, + ) + origin_wp = wp.vec3(float(rotation_origin[0]), float(rotation_origin[1]), float(rotation_origin[2])) + reference_area = chord_length * span_length + bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max) + stepper = setup_stepper(grid, bc_list, ibm_max_iterations, ibm_tolerance, ibm_relaxation) + f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + velocities_wp = wp.zeros(shape=vertices_wp.shape[0], dtype=wp.vec3) + device = vertices_wp.device + wp.launch( + kernel=update_airfoil_pose, + dim=vertices_wp.shape[0], + inputs=[ + 0, + num_steps, + start_angle_rad, + total_rotation_rad, + origin_wp, + base_vertices_wp, + vertices_wp, + velocities_wp, + ], + device=device, + ) + cd_values = [] + cl_values = [] + visc = u_max * chord_length / Re + omega = 1.0 / (3.0 * visc + 0.5) + print(f" Omega: {omega}") + try: + for i in range(num_steps): + f_0, f_1, lag_forces = stepper( + f_0, + f_1, + vertices_wp, + areas_wp, + velocities_wp, + bc_mask, + missing_mask, + omega, + i, + ) + f_0, f_1 = f_1, f_0 + if print_interval > 0 and i % print_interval == 0: + print(f"Step {i}/{num_steps} completed") + if i % post_process_interval == 0 or i == num_steps - 1: + post_process( + i, + post_process_interval, + f_0, + precision_policy, + grid_shape, + lag_forces, + cd_values, + cl_values, + u_max, + reference_area, + areas_np, + ) + next_step = i + 1 + if next_step < num_steps: + wp.launch( + kernel=update_airfoil_pose, + dim=vertices_wp.shape[0], + inputs=[ + next_step, + num_steps, + start_angle_rad, + total_rotation_rad, + origin_wp, + base_vertices_wp, + vertices_wp, + velocities_wp, + ], + device=device, + ) + except KeyboardInterrupt: + print("Simulation interrupted by user.") + if cd_values and cl_values: + save_force_coefficients(cd_values, cl_values, "airfoil_force_coefficients.csv") + print("Force coefficient data saved to airfoil_force_coefficients.csv") + raise + if cd_values and cl_values: + save_force_coefficients(cd_values, cl_values, "airfoil_force_coefficients.csv") + print("Force coefficient data saved to airfoil_force_coefficients.csv") + print(f"Final Cd (avg last 10): {np.mean([cd for _, cd in cd_values[-10:]]):.6f}") + print(f"Final Cl (avg last 10): {np.mean([cl for _, cl in cl_values[-10:]]):.6f}") + print("Simulation finished.") diff --git a/examples/ibm/sphere_ibm.py b/examples/ibm/sphere_ibm.py new file mode 100644 index 00000000..834e5736 --- /dev/null +++ b/examples/ibm/sphere_ibm.py @@ -0,0 +1,270 @@ +""" +Flow past a sphere (IBM) β€” Drag coefficient validation + +References +- Johnson, T. A., & Patel, V. C. (1999). Flow past a sphere up to Re = 300. + Journal of Fluid Mechanics, 378, 19–70. (domain sizing, Cd at Re β‰ˆ 100) +- Uhlmann, M. (2005). An immersed boundary method with direct forcing for + particulate flows. Journal of Computational Physics, 209(2), 448–476. + (IBM forcing and hydrodynamic force evaluation) +- Clift, R., Grace, J. R., & Weber, M. E. (1978). Bubbles, Drops, and + Particles. Academic Press. (Cd correlations vs Reynolds number) +- Achenbach, E. (1972). Experiments on the flow past spheres at very high + Reynolds numbers. Journal of Fluid Mechanics, 54(3), 565–575. + (experimental Cd curve) +""" + +import os +import xlb +import trimesh +import numpy as np +import jax.numpy as jnp +import matplotlib.pyplot as plt +import warp as wp +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.operator.stepper import IBMStepper +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, +) +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_fields_vtk, save_image, warp_array_to_jax +from xlb.helper.ibm_helper import prepare_immersed_boundary +from xlb.grid import grid_factory + + +def create_sphere_mesh(center, radius, subdivisions=3): + sphere = trimesh.creation.icosphere(subdivisions=subdivisions, radius=radius) + sphere.apply_translation(center) + return sphere + + +def define_boundary_indices(grid, velocity_set): + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + inlet = box_no_edge["left"] + outlet = box_no_edge["right"] + walls = [box["front"][i] + box["back"][i] + box["top"][i] + box["bottom"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + return inlet, outlet, walls + + +def bc_profile(precision_policy, grid_shape, u_max): + _dtype = precision_policy.store_precision.wp_dtype + u_max_d = _dtype(u_max) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + return wp.vec(_dtype(u_max_d), length=1) + + return bc_profile_warp + + +def calculate_drag_coefficient(lag_forces, reference_velocity, frontal_area, areas_wp): + forces_np = lag_forces.numpy() + drag_forces = forces_np[:, 0] + # The negative sign is necessary because drag is defined as the force opposing the flow direction. + # In this simulation, the computed force may be positive in the flow direction, so we negate it + # to obtain the drag (force acting against the flow). + total_drag = -np.sum(drag_forces * areas_wp.numpy()) + + dynamic_pressure = 0.5 * reference_velocity**2 + cd = total_drag / (dynamic_pressure * frontal_area) + + return cd, total_drag + + +def setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max): + inlet, outlet, walls = define_boundary_indices(grid, velocity_set) + bc_inlet = RegularizedBC("velocity", indices=inlet, profile=bc_profile(precision_policy, grid_shape, u_max)) + bc_outlet = ExtrapolationOutflowBC(indices=outlet) + bc_walls = FullwayBounceBackBC(indices=walls) + return [bc_walls, bc_inlet, bc_outlet] + + +def setup_stepper(grid, boundary_conditions, ibm_max_iterations=2, ibm_tolerance=1e-5, ibm_relaxation=1.0): + return IBMStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + ibm_max_iterations=ibm_max_iterations, + ibm_tolerance=ibm_tolerance, + ibm_relaxation=ibm_relaxation, + ) + + +def post_process( + i, + post_process_interval, + f_current, + precision_policy, + grid_shape, + lag_forces, + cd_values, + reference_velocity, + frontal_area, + areas_wp, +): + if not isinstance(f_current, jnp.ndarray): + f_jax = warp_array_to_jax(f_current) + else: + f_jax = f_current + + macro_jax = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + ) + rho, u = macro_jax(f_jax) + u = u[:, 1:-1, 1:-1, 1:-1] + + fields = { + "u_magnitude": (u[0] ** 2.0 + u[1] ** 2.0 + u[2] ** 2.0) ** 0.5, + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + } + slice_idy = grid_shape[1] // 2 + save_image(fields["u_magnitude"][:, slice_idy, :], timestep=i) + # save_fields_vtk(fields, i) + + cd, total_drag = calculate_drag_coefficient(lag_forces, reference_velocity, frontal_area, areas_wp) + cd_values.append((i, cd)) + if i % post_process_interval == 0: + window = 10 + if len(cd_values) >= window: + avg_cd = float(np.mean([v for _, v in cd_values[-window:]])) + else: + avg_cd = float(np.mean([v for _, v in cd_values])) + print(f"Step {i}: Cd = {cd:.6f}, Cd(avg{window}) = {avg_cd:.6f}, Total Drag = {total_drag:.6f}") + + +def save_drag_coefficient(cd_values, filename): + with open(filename, "w") as f: + f.write("timestep,cd\n") + for timestep, cd in cd_values: + f.write(f"{timestep},{cd}\n") + + timesteps = [t for t, _ in cd_values] + cds = [cd for _, cd in cd_values] + + plt.figure(figsize=(10, 6)) + plt.plot(timesteps, cds, "b-") + plt.grid(True, linestyle="--", alpha=0.7) + plt.xlabel("Timestep") + plt.ylabel("Drag Coefficient (Cd)") + plt.title("Drag Coefficient vs Time") + plt.tight_layout() + plt.savefig("drag_coefficient_sphere.png", dpi=150) + plt.close() + + +sphere_radius = 25.0 + +diameter = 2.0 * sphere_radius +upstream = int(1.5 * diameter) # 1.5D upstream +downstream = int(7 * diameter) # 2.5D downstream +ly = int(3.0 * diameter) # 3D lateral +lz = int(3.0 * diameter) # 3D vertical +lx = upstream + downstream +grid_shape = (lx, ly, lz) + + +# Uniform inlet velocity +u_max = 0.02 + +# Place sphere at 1/3 from entrance +sphere_center = [float(lx / 3), grid_shape[1] / 2.0, grid_shape[2] / 2.0] + +Re = 300 +visc = u_max * (2.0 * sphere_radius) / Re +omega = 1.0 / (3.0 * visc + 0.5) + +num_steps = 20000 +post_process_interval = 1000 +print_interval = 1000 + +ibm_max_iterations = 4 +ibm_tolerance = 1e-5 +ibm_relaxation = 0.5 + +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +xlb.init(velocity_set=velocity_set, default_backend=compute_backend, default_precision_policy=precision_policy) +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +print("Sphere IBM Simulation Configuration:") +print(f" Grid size: {grid_shape}") +print(f" Sphere radius: {sphere_radius}") +print(f" Sphere center: {sphere_center}") +print(f" Omega: {omega}") +print(f" Inlet velocity: {u_max}") +print(f" Reynolds number: {Re}") +print(f" Max steps: {num_steps}") +print(f" IBM max iterations: {ibm_max_iterations}") +print(f" IBM tolerance: {ibm_tolerance}") +print(f" IBM relaxation: {ibm_relaxation}") + +sphere_mesh = create_sphere_mesh(sphere_center, sphere_radius, subdivisions=4) +vertices_wp, areas_wp, faces_np = prepare_immersed_boundary(sphere_mesh, max_lbm_length=sphere_radius * 2) + +frontal_area = np.pi * sphere_radius**2 +print(f"Frontal area (theoretical): {frontal_area:.2f}") + +bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max) +stepper = setup_stepper(grid, bc_list, ibm_max_iterations, ibm_tolerance, ibm_relaxation) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + +velocities_wp = wp.zeros(shape=vertices_wp.shape[0], dtype=wp.vec3) +cd_values = [] + +try: + for i in range(num_steps): + f_0, f_1, lag_forces = stepper( + f_0, + f_1, + vertices_wp, + areas_wp, + velocities_wp, + bc_mask, + missing_mask, + omega, + i, + ) + f_0, f_1 = f_1, f_0 + + if print_interval > 0 and i % print_interval == 0: + print(f"Step {i}/{num_steps} completed") + + if i % post_process_interval == 0 or i == num_steps - 1: + post_process( + i, + post_process_interval, + f_0, + precision_policy, + grid_shape, + lag_forces, + cd_values, + u_max, + frontal_area, + areas_wp, + ) + +except KeyboardInterrupt: + print("\nSimulation interrupted by user.") + if cd_values: + save_drag_coefficient(cd_values, "drag_coefficient_sphere.csv") + print("Drag coefficient data saved to drag_coefficient_sphere.csv") + import sys + + sys.exit(0) + +if cd_values: + save_drag_coefficient(cd_values, "drag_coefficient_sphere.csv") + print("Drag coefficient data saved to drag_coefficient_sphere.csv") + print(f"Final Cd (average of last 10 values): {np.mean([cd for _, cd in cd_values[-10:]]):.6f}") + +print("Simulation finished.") diff --git a/examples/ibm/wind_turbine_ibm.py b/examples/ibm/wind_turbine_ibm.py new file mode 100644 index 00000000..c8c688e0 --- /dev/null +++ b/examples/ibm/wind_turbine_ibm.py @@ -0,0 +1,461 @@ +import os +import xlb +import trimesh +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.operator.stepper import IBMStepper +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, +) +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import ( + save_fields_vtk, + save_image, + save_usd_vorticity, + save_usd_q_criterion, + update_usd_lagrangian_parts, + plot_object_placement, +) +import warp as wp +import numpy as np +import matplotlib.pyplot as plt +from xlb.helper.ibm_helper import prepare_immersed_boundary +from xlb.grid import grid_factory +from pxr import Usd, UsdGeom, Vt +from xlb.operator.postprocess import QCriterion, Vorticity, GridToPoint + + +def define_boundary_indices(grid, velocity_set): + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + inlet = box_no_edge["front"] + outlet = box_no_edge["back"] + walls = [box["right"][i] + box["left"][i] + box["top"][i] + box["bottom"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + return inlet, outlet, walls + + +def bc_profile(precision_policy, grid_shape, u_max): + _dtype = precision_policy.store_precision.wp_dtype + u_max_d = _dtype(u_max) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + return wp.vec(u_max_d, length=1) + + return bc_profile_warp + + +def setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, inlet_speed): + inlet, outlet, walls = define_boundary_indices(grid, velocity_set) + bc_inlet = RegularizedBC("velocity", indices=inlet, profile=bc_profile(precision_policy, grid_shape, inlet_speed)) + bc_outlet = ExtrapolationOutflowBC(indices=outlet) + bc_walls = FullwayBounceBackBC(indices=walls) + return [bc_inlet, bc_outlet, bc_walls] + + +def setup_stepper(grid, boundary_conditions, lbm_omega): + return IBMStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + ) + + +# You must download the stl files from the following link: https://www.cgtrader.com/free-3d-models/industrial/industrial-machine/offshore-wind-turbine-235-m-radius +# and separate them into two files: turbine_wind_turbine.stl and body_wind_turbine.stl. Put them in the same directory as this script. +def load_and_prepare_meshes_turbine(grid_shape, stl_dir="./examples/ibm/"): + rotor_stl = os.path.join(stl_dir, "turbine_wind_turbine.stl") + body_stl = os.path.join(stl_dir, "body_wind_turbine.stl") + if not os.path.isfile(rotor_stl): + raise FileNotFoundError(f"Cannot find {rotor_stl}") + if not os.path.isfile(body_stl): + raise FileNotFoundError(f"Cannot find {body_stl}") + + rotor_mesh = trimesh.load_mesh(rotor_stl, process=False) + body_mesh = trimesh.load_mesh(body_stl, process=False) + + # Identify rotor bounding box dimension that dictates scaling + rotor_bounds = rotor_mesh.bounds + rotor_size = rotor_bounds[1] - rotor_bounds[0] + rotor_diameter = max(rotor_size) # largest dimension in the rotor + desired_diameter = 150.0 + + scale_factor = desired_diameter / rotor_diameter + print(f"Scale factor: {scale_factor:.4f}") + + # We apply the scale to both rotor and body + rotor_mesh.apply_scale(scale_factor) + body_mesh.apply_scale(scale_factor) + + R_dummy = trimesh.transformations.rotation_matrix(np.radians(0), [1, 0, 0]) + rotor_mesh.apply_transform(R_dummy) + body_mesh.apply_transform(R_dummy) + + combined = trimesh.util.concatenate([rotor_mesh, body_mesh]) + bnds = combined.bounds + min_x, min_y, min_z = bnds[0] + max_x, max_y, max_z = bnds[1] + center_x = 0.5 * grid_shape[0] + center_y = 0.5 * grid_shape[1] + + shift_x = center_x - 0.5 * (min_x + max_x) + shift_y = center_y - 0.5 * (min_y + max_y) + shift_z = -min_z + + combined.apply_translation([shift_x, shift_y, shift_z]) + rotor_mesh = trimesh.load_mesh(rotor_stl, process=False) + body_mesh = trimesh.load_mesh(body_stl, process=False) + rotor_mesh.apply_scale(scale_factor) + body_mesh.apply_scale(scale_factor) + rotor_mesh.apply_transform(R_dummy) + body_mesh.apply_transform(R_dummy) + rotor_mesh.apply_translation([shift_x, shift_y, shift_z]) + body_mesh.apply_translation([shift_x, shift_y, shift_z]) + + Nx, Ny, Nz = grid_shape + rotor_v_wp, rotor_a_wp, rotor_faces = prepare_immersed_boundary(rotor_mesh, max_lbm_length=max(Nx, Ny, Nz)) + body_v_wp, body_a_wp, body_faces = prepare_immersed_boundary(body_mesh, max_lbm_length=max(Nx, Ny, Nz)) + + all_vertices = body_v_wp.numpy() + all_areas = body_a_wp.numpy() + all_faces = body_faces.copy() + current_offset = len(body_v_wp) + + rotor_v_np = rotor_v_wp.numpy() + rotor_a_np = rotor_a_wp.numpy() + rotor_f_offset = rotor_faces + current_offset + all_vertices = np.vstack([all_vertices, rotor_v_np]) + all_areas = np.hstack([all_areas, rotor_a_np]) + all_faces = np.vstack([all_faces, rotor_f_offset]) + + vertices_wp = wp.array(all_vertices, dtype=wp.vec3) + areas_wp = wp.array(all_areas, dtype=wp.float32) + faces_np = all_faces + num_body_vertices = len(body_v_wp) + num_rotor_vertices = len(rotor_v_wp) + body_faces_np = body_faces + rotor_faces_np = faces_np[len(body_faces_np) :] + + print("\nTurbine mesh preparation summary:") + print(f" Scale factor: {scale_factor:.4f}") + print(f" Rotor diameter (desired ~200): {desired_diameter}") + print(f" Total vertices: {len(all_vertices)}") + print(f" Body vertices: {num_body_vertices}") + print(f" Rotor vertices: {num_rotor_vertices}") + + return { + "vertices_wp": vertices_wp, + "areas_wp": areas_wp, + "faces_np": faces_np, + "num_body_vertices": num_body_vertices, + "num_rotor_vertices": num_rotor_vertices, + "body_faces_np": body_faces_np, + "rotor_faces_np": rotor_faces_np, + } + + +@wp.kernel +def rotate_rotor( + timestep: int, + forces: wp.array(dtype=wp.vec3), + vertices: wp.array(dtype=wp.vec3), + velocities: wp.array(dtype=wp.vec3), +): + idx = wp.tid() + + if idx < _num_body_vertices: + velocities[idx] = wp.vec3(0.0, 0.0, 0.0) + return + + # For rotor vertices, rotate about the negative Y axis (axis = -Y). + center = wp.vec3(_rotor_center_x, _rotor_center_y, _rotor_center_z) + pos = vertices[idx] + rel_pos = pos - center + + radius_x = rel_pos[0] + radius_z = rel_pos[2] + r = wp.sqrt(radius_x * radius_x + radius_z * radius_z) + if r < 1e-6: + velocities[idx] = wp.vec3(0.0, 0.0, 0.0) + return + + nx = radius_x / r + nz = radius_z / r + theta = _rotor_speed + c = wp.cos(theta) + s = wp.sin(theta) + + x_new = r * (c * nx - s * nz) + z_new = r * (s * nx + c * nz) + + # Update position + new_rel = wp.vec3(x_new, rel_pos[1], z_new) + vertices[idx] = new_rel + center + + # Tangential velocity from rotation + velocities[idx] = wp.vec3(z_new * _rotor_speed, 0.0, -x_new * _rotor_speed) + + +def post_process( + i, + post_process_interval, + f_current, + bc_mask, + grid, + faces_np, + vertices_wp, + precision_policy, + grid_shape, + usd_mesh_vorticity, + usd_mesh_q_criterion, + vorticity_operator, + q_criterion_operator, + usd_stage, + turbine_body_mesh, + turbine_rotor_mesh, + num_body_vertices, + body_faces_np, + rotor_faces_np, + lag_forces=None, +): + # if not isinstance(f_current, jnp.ndarray): + # f_jax = wp.to_jax(f_current) + # else: + # f_jax = f_current + + # macro_jax = Macroscopic( + # compute_backend=ComputeBackend.JAX, + # precision_policy=precision_policy, + # velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + # ) + # rho, u = macro_jax(f_jax) + # u = u[:, 20:-20, 20:-20, 5:-20] + + # fields = { + # "u_magnitude": (u[0] ** 2.0 + u[1] ** 2.0 + u[2] ** 2.0) ** 0.5, + # "u_x": u[0], + # "u_y": u[1], + # "u_z": u[2], + # } + + # slice_idx = grid_shape[0] // 2 + # slice_idy = grid_shape[1] // 2 + # save_image(fields["u_magnitude"][slice_idx, :, :], timestep=i, prefix="slice_idx") + # save_image(fields["u_magnitude"][:, slice_idy, :], timestep=i, prefix="slice_idy") + + # save_fields_vtk(fields, i) + + save_usd_vorticity( + timestep=i, + post_process_interval=post_process_interval, + bc_mask=bc_mask, + f_current=f_current, + grid_shape=grid_shape, + usd_mesh=usd_mesh_vorticity, + vorticity_operator=vorticity_operator, + precision_policy=precision_policy, + vorticity_threshold=1e-2, + usd_stage=usd_stage, + device="cuda:1", + clip_lower=(20, 20, 5), + clip_upper=(20, 20, 20), + ) + + save_usd_q_criterion( + timestep=i, + post_process_interval=post_process_interval, + bc_mask=bc_mask, + f_current=f_current, + grid_shape=grid_shape, + usd_mesh=usd_mesh_q_criterion, + q_criterion_operator=q_criterion_operator, + precision_policy=precision_policy, + q_threshold=5e-6, + usd_stage=usd_stage, + device="cuda:1", + clip_lower=(20, 20, 5), + clip_upper=(20, 20, 20), + color_range=(0.0, 0.1), + ) + + update_usd_lagrangian_parts( + timestep=i, + post_process_interval=post_process_interval, + vertices_wp=vertices_wp, + parts=[ + { + "start": 0, + "end": num_body_vertices, + "faces": body_faces_np, + "usd_mesh": turbine_body_mesh, + "colorize": True, + }, + { + "start": num_body_vertices, + "end": vertices_wp.shape[0], + "faces": rotor_faces_np, + "usd_mesh": turbine_rotor_mesh, + "colorize": True, + }, + ], + vertex_offset=(20, 20, 5), + lag_forces=lag_forces, + device="cuda:1", + ) + + +# +# Main simulation +# + +grid_shape = (256, 450, 256) # example domain size (Nx, Ny, Nz) +u_inlet = 0.05 # inlet flow speed +num_steps = 25000 +post_process_interval = 100 +print_interval = 100 +turbine_rotation_speed = -0.0005 # user-controlled rotor speed (radians per timestep) +Re = 5e5 + +clength = grid_shape[0] - 1 +visc = u_inlet * clength / Re +omega = 1.0 / (3.0 * visc + 0.5) + +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +xlb.init(velocity_set=velocity_set, default_backend=compute_backend, default_precision_policy=precision_policy) +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +print("Wind Turbine Simulation Configuration:") +print(f" Grid size: {grid_shape}") +print(f" Omega: {omega:.6f}") +print(f" Backend: {compute_backend}") +print(f" Velocity set: {velocity_set}") +print(f" Precision policy: {precision_policy}") +print(f" Inlet velocity: {u_inlet}") +print(f" Reynolds number: {Re}") +print(f" Turbine rotation speed (rad/step): {turbine_rotation_speed}") +print(f" Max steps: {num_steps}") + +usd_output_directory = "usd_output_turbine" +os.makedirs(usd_output_directory, exist_ok=True) +usd_file = os.path.join(usd_output_directory, "turbine_output.usd") +usd_stage = Usd.Stage.CreateNew(usd_file) +usd_mesh_vorticity = UsdGeom.Mesh.Define(usd_stage, "/World/Vorticity") +usd_mesh_q_criterion = UsdGeom.Mesh.Define(usd_stage, "/World/QCriterion") +usd_turbine_body = UsdGeom.Mesh.Define(usd_stage, "/World/TurbineBody") +usd_turbine_rotor = UsdGeom.Mesh.Define(usd_stage, "/World/TurbineRotor") + +mesh_data = load_and_prepare_meshes_turbine(grid_shape) +vertices_wp = mesh_data["vertices_wp"] +areas_wp = mesh_data["areas_wp"] +faces_np = mesh_data["faces_np"] +num_body_vertices = mesh_data["num_body_vertices"] +num_rotor_vertices = mesh_data["num_rotor_vertices"] +body_faces_np = mesh_data["body_faces_np"] +rotor_faces_np = mesh_data["rotor_faces_np"] + +plot_object_placement( + vertices_wp, + grid_shape, + "turbine_placement.png", + "Turbine Placement (X-Y Top View)", + "Turbine bounding box", +) + +# Calculate rotor center (simple approach: average rotor vertex positions) +rotor_center_np = vertices_wp.numpy()[num_body_vertices:].mean(axis=0) +_num_body_vertices = wp.constant(int(num_body_vertices)) +_rotor_speed = wp.constant(float(turbine_rotation_speed)) +_rotor_center_x = wp.constant(float(rotor_center_np[0])) +_rotor_center_y = wp.constant(float(rotor_center_np[1])) +_rotor_center_z = wp.constant(float(rotor_center_np[2])) + +bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_inlet) +stepper = setup_stepper(grid, bc_list, omega) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + +device = "cuda:1" +with wp.ScopedDevice(device): + q_criterion_operator = QCriterion( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + vorticity_operator = Vorticity( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + +velocities_wp = wp.zeros(shape=vertices_wp.shape[0], dtype=wp.vec3) + +try: + for i in range(num_steps): + f_0, f_1, lag_forces = stepper( + f_0, + f_1, + vertices_wp, + areas_wp, + velocities_wp, + bc_mask, + missing_mask, + omega, + i, + ) + f_0, f_1 = f_1, f_0 + + wp.launch( + kernel=rotate_rotor, + dim=vertices_wp.shape[0], + inputs=[ + i, + lag_forces, + vertices_wp, + velocities_wp, + ], + ) + + if i % post_process_interval == 0 or i == num_steps - 1: + post_process( + i, + post_process_interval, + f_0, + bc_mask, + grid, + faces_np, + vertices_wp, + precision_policy, + grid_shape, + usd_mesh_vorticity, + usd_mesh_q_criterion, + vorticity_operator, + q_criterion_operator, + usd_stage, + usd_turbine_body, + usd_turbine_rotor, + num_body_vertices, + body_faces_np, + rotor_faces_np, + lag_forces, + ) +except KeyboardInterrupt: + print("\nSimulation interrupted by user. Saving current USD state...") + current_time_code = i // post_process_interval + usd_stage.SetStartTimeCode(0) + usd_stage.SetEndTimeCode(current_time_code) + usd_stage.SetTimeCodesPerSecond(30) + usd_stage.Save() + print(f"USD file saved with {current_time_code + 1} frames. Exiting.") + import sys + + sys.exit(0) + +usd_stage.SetStartTimeCode(0) +usd_stage.SetEndTimeCode(num_steps // post_process_interval) +usd_stage.SetTimeCodesPerSecond(30) +usd_stage.Save() +print("Simulation finished. USD file saved.") diff --git a/examples/ibm/windtunnel_ibm.py b/examples/ibm/windtunnel_ibm.py new file mode 100644 index 00000000..8f62c9db --- /dev/null +++ b/examples/ibm/windtunnel_ibm.py @@ -0,0 +1,668 @@ +import os +import xlb +import trimesh +import time +from tqdm import tqdm +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.operator.stepper import IBMStepper +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, +) +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import ( + save_fields_vtk, + save_image, + save_usd_vorticity, + save_usd_q_criterion, + update_usd_lagrangian_parts, + plot_object_placement, + warp_array_to_jax, +) +import warp as wp +import numpy as np +import jax.numpy as jnp +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D +from xlb.helper.ibm_helper import prepare_immersed_boundary +from xlb.grid import grid_factory +from pxr import Usd, UsdGeom, Sdf, Vt +from xlb.operator.postprocess import QCriterion, Vorticity, GridToPoint + + +def load_and_prepare_meshes_car(grid_shape, stl_dir): + if not os.path.exists(stl_dir): + raise FileNotFoundError(f"STL directory {stl_dir} does not exist.") + + body_stl = os.path.join(stl_dir, "S550_GT500_BS_5p.stl") + wheel_fr_stl = os.path.join(stl_dir, "S550_GT500_BS_FR_5p.stl") + wheel_fl_stl = os.path.join(stl_dir, "S550_GT500_BS_FL_5p.stl") + wheel_rr_stl = os.path.join(stl_dir, "S550_GT500_BS_RR_5p.stl") + wheel_rl_stl = os.path.join(stl_dir, "S550_GT500_BS_RL_5p.stl") + for stl_file in [body_stl, wheel_fr_stl, wheel_fl_stl, wheel_rr_stl, wheel_rl_stl]: + if not os.path.isfile(stl_file): + raise FileNotFoundError(f"STL file {stl_file} not found.") + + # Load body FIRST to get reference dimensions + body_mesh = trimesh.load_mesh(body_stl, process=False) + body_bounds = body_mesh.bounds + orig_body_width = body_bounds[1][1] - body_bounds[0][1] # Use body width for scaling + + # Now load other components + fr_mesh = trimesh.load_mesh(wheel_fr_stl, process=False) + fl_mesh = trimesh.load_mesh(wheel_fl_stl, process=False) + rr_mesh = trimesh.load_mesh(wheel_rr_stl, process=False) + rl_mesh = trimesh.load_mesh(wheel_rl_stl, process=False) + + # Calculate scale based on BODY width only + Nx, Ny, Nz = grid_shape + target_body_width = Ny / 3.0 + scale_factor = target_body_width / orig_body_width # Use body's original width + + # Apply scale to ALL components + body_mesh.apply_scale(scale_factor) + fr_mesh.apply_scale(scale_factor) + fl_mesh.apply_scale(scale_factor) + rr_mesh.apply_scale(scale_factor) + rl_mesh.apply_scale(scale_factor) + + # Now combine scaled meshes + combined = trimesh.util.concatenate([body_mesh, fr_mesh, fl_mesh, rr_mesh, rl_mesh]) + + R_y = trimesh.transformations.rotation_matrix(np.radians(180), [0, 1, 0]) + R_x = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) + combined.apply_transform(R_y) + combined.apply_transform(R_x) + + bnds = combined.bounds + min_x, min_y, min_z = bnds[0] + max_x, max_y, max_z = bnds[1] + car_length = max_x - min_x + car_height = max_z - min_z + + target_front_x = Nx / 3.0 + shift_x = target_front_x - min_x + + back_x_after_shift = max_x + shift_x + if back_x_after_shift > Nx: + raise ValueError( + f"Car too long to fit in domain with front at 1/3!\n" + f"Car length: {car_length:.1f}\n" + f"Available space: {Nx - target_front_x:.1f}\n" + f"Would need additional {back_x_after_shift - Nx:.1f} units" + ) + + shift_y = (Ny - (max_y - min_y)) / 2.0 - min_y + + ground_clearance = 2.0 + shift_z = ground_clearance - min_z + + combined.apply_translation([shift_x, shift_y, shift_z]) + + bnds = combined.bounds + min_x, min_y, min_z = bnds[0] + max_x, max_y, max_z = bnds[1] + + tolerance = 0.1 + front_tolerance = target_front_x * tolerance + width_tolerance = target_body_width * tolerance + + if not (target_front_x - front_tolerance < min_x < target_front_x + front_tolerance): + raise ValueError(f"Car front not at 1/3 domain length! Front at {min_x:.1f}, should be {target_front_x:.1f}") + + scaled_body_width = body_mesh.bounds[1][1] - body_mesh.bounds[0][1] + if not (target_body_width - width_tolerance < scaled_body_width < target_body_width + width_tolerance): + raise ValueError(f"Car BODY width mismatch! Current: {scaled_body_width:.1f}, Target: {target_body_width:.1f}") + + print("\nScaled dimensions verification:") + print(f"Body width: {scaled_body_width:.1f} (target: {target_body_width:.1f})") + print(f"Total car width (body + wheels): {max_y - min_y:.1f}") + + def apply_transform(mesh): + mesh.apply_scale(scale_factor) + mesh.apply_transform(R_y) + mesh.apply_transform(R_x) + mesh.apply_translation([shift_x, shift_y, shift_z]) + + body_mesh = trimesh.load_mesh(body_stl, process=False) + fr_mesh = trimesh.load_mesh(wheel_fr_stl, process=False) + fl_mesh = trimesh.load_mesh(wheel_fl_stl, process=False) + rr_mesh = trimesh.load_mesh(wheel_rr_stl, process=False) + rl_mesh = trimesh.load_mesh(wheel_rl_stl, process=False) + + apply_transform(body_mesh) + apply_transform(fr_mesh) + apply_transform(fl_mesh) + apply_transform(rr_mesh) + apply_transform(rl_mesh) + + body_v_wp, body_a_wp, body_faces = prepare_immersed_boundary(body_mesh, max_lbm_length=Ny / 2.0) + fr_v_wp, fr_a_wp, fr_faces = prepare_immersed_boundary(fr_mesh, max_lbm_length=Ny / 2.0) + fl_v_wp, fl_a_wp, fl_faces = prepare_immersed_boundary(fl_mesh, max_lbm_length=Ny / 2.0) + rr_v_wp, rr_a_wp, rr_faces = prepare_immersed_boundary(rr_mesh, max_lbm_length=Ny / 2.0) + rl_v_wp, rl_a_wp, rl_faces = prepare_immersed_boundary(rl_mesh, max_lbm_length=Ny / 2.0) + + all_vertices = body_v_wp.numpy() + all_areas = body_a_wp.numpy() + all_faces = body_faces.copy() + current_offset = len(body_v_wp) + + wheels = [ + (fr_v_wp, fr_a_wp, fr_faces), + (fl_v_wp, fl_a_wp, fl_faces), + (rr_v_wp, rr_a_wp, rr_faces), + (rl_v_wp, rl_a_wp, rl_faces), + ] + wheel_ranges = [] + for wv_wp, wa_wp, wf in wheels: + wv_np = wv_wp.numpy() + wa_np = wa_wp.numpy() + wf_offset = wf + current_offset + all_vertices = np.vstack([all_vertices, wv_np]) + all_areas = np.hstack([all_areas, wa_np]) + all_faces = np.vstack([all_faces, wf_offset]) + start_idx = current_offset + end_idx = start_idx + len(wv_np) + wheel_ranges.append((start_idx, end_idx)) + current_offset += len(wv_np) + + vertices_wp = wp.array(all_vertices, dtype=wp.vec3) + areas_wp = wp.array(all_areas, dtype=wp.float32) + faces_np = all_faces + num_body_vertices = len(body_v_wp) + num_wheel_vertices = len(all_vertices) - num_body_vertices + body_faces_np = body_faces + wheels_faces_np = faces_np[len(body_faces_np) :] + + wheel_centers = [] + for start_idx, end_idx in wheel_ranges: + center = np.mean(all_vertices[start_idx:end_idx], axis=0) + wheel_centers.append(center) + + # Calculate frontal area using the bounding box cross-section (YZ plane) + mesh = trimesh.Trimesh(vertices=all_vertices, faces=faces_np) + bounds = mesh.bounds + min_y, min_z = bounds[0][1], bounds[0][2] + max_y, max_z = bounds[1][1], bounds[1][2] + width = max_y - min_y + height = max_z - min_z + frontal_area = width * height + + print(f"Calculated frontal area (bounding box): {frontal_area:.2f} square units") + print(f" Width: {width:.2f}, Height: {height:.2f}") + + print("\nMesh preparation summary:") + print(f"Target size (car width): {target_body_width:.2f}") + print(f"Scale factor applied: {scale_factor:.4f}") + print(f"Total vertices: {len(all_vertices)}") + print(f"Body vertices: {num_body_vertices}") + print(f"Wheel vertices: {num_wheel_vertices}") + print(f"Number of wheels: {len(wheel_ranges)}") + + return { + "vertices_wp": vertices_wp, + "areas_wp": areas_wp, + "faces_np": faces_np, + "wheel_ranges": wheel_ranges, + "wheel_centers": wheel_centers, + "num_body_vertices": num_body_vertices, + "num_wheel_vertices": num_wheel_vertices, + "body_faces_np": body_faces_np, + "wheels_faces_np": wheels_faces_np, + "frontal_area": frontal_area, + } + + +def define_boundary_indices(grid, velocity_set): + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + inlet = box_no_edge["right"] + outlet = box_no_edge["left"] + walls = [box["front"][i] + box["back"][i] + box["top"][i] + box["bottom"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + return inlet, outlet, walls + + +def bc_profile(precision_policy, grid_shape, u_max): + _dtype = precision_policy.store_precision.wp_dtype + u_max_d = _dtype(u_max) + L_y = _dtype(grid_shape[1] - 1) + L_z = _dtype(grid_shape[2] - 1) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + y = _dtype(index[1]) + z = _dtype(index[2]) + y_center = y - (L_y / _dtype(2.0)) + z_center = z - (L_z / _dtype(2.0)) + r_sq = ((_dtype(2.0) * y_center) / L_y) ** _dtype(2.0) + ((_dtype(2.0) * z_center) / L_z) ** _dtype(2.0) + velocity_x = u_max_d * wp.max(_dtype(0.0), _dtype(1.0) - r_sq) + return wp.vec(velocity_x, length=1) + + return bc_profile_warp + + +def calculate_drag_coefficient(lag_forces, reference_velocity, frontal_area, areas_wp): + """ + Calculate the drag coefficient (Cd) from lagrangian forces + + Args: + lag_forces: Warp array of forces on vertices + reference_velocity: Reference velocity (u_max) + frontal_area: Frontal area of the car (from bounding box) + areas_wp: Warp array of vertex areas + + Returns: + cd: Calculated drag coefficient + """ + forces_np = lag_forces.numpy() + drag_forces = forces_np[:, 0] # X-component is the drag direction + + areas_np = areas_wp.numpy() + + drag_forces = drag_forces * areas_np + total_drag = np.sum(drag_forces) + + # Calculate dynamic pressure (rho = 1.0 in lattice units) + dynamic_pressure = 0.5 * reference_velocity**2 + + cd = total_drag / (dynamic_pressure * frontal_area) + + return cd + + +def setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max): + inlet, outlet, walls = define_boundary_indices(grid, velocity_set) + bc_inlet = RegularizedBC("velocity", indices=inlet, profile=bc_profile(precision_policy, grid_shape, u_max)) + bc_outlet = ExtrapolationOutflowBC(indices=outlet) + bc_walls = FullwayBounceBackBC(indices=walls) + return [bc_walls, bc_inlet, bc_outlet] + + +def setup_stepper(grid, boundary_conditions, lbm_omega): + return IBMStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="KBC", + ) + + +@wp.kernel +def rotate_wheels( + timestep: int, + forces: wp.array(dtype=wp.vec3), + vertices: wp.array(dtype=wp.vec3), + velocities: wp.array(dtype=wp.vec3), +): + idx = wp.tid() + + if idx < _num_body_vertices: + velocities[idx] = wp.vec3(0.0, 0.0, 0.0) + return + + wheel_id = wp.int32(-1) + for b in range(_num_wheels): + start = _wheel_starts[b] + end = _wheel_ends[b] + if (idx >= start) and (idx < end): + wheel_id = b + break + + if wheel_id == -1: + velocities[idx] = wp.vec3(0.0, 0.0, 0.0) + return + + center = wp.vec3( + _wheel_centers_x[wheel_id], + _wheel_centers_y[wheel_id], + _wheel_centers_z[wheel_id], + ) + + rel_pos = vertices[idx] - center + + # Use normalized length to prevent shrinking + radius = wp.sqrt(rel_pos[0] * rel_pos[0] + rel_pos[2] * rel_pos[2]) + if radius > 1e-6: # Avoid division by zero + # Normalize x-z components + norm_x = rel_pos[0] / radius + norm_z = rel_pos[2] / radius + + # Apply rotation + theta = _wheel_speed + c = wp.cos(theta) + s = wp.sin(theta) + + # Compute new position while preserving the radius + x_new = radius * (c * norm_x - s * norm_z) + z_new = radius * (s * norm_x + c * norm_z) + + new_rel_pos = wp.vec3(x_new, rel_pos[1], z_new) + vertices[idx] = new_rel_pos + center + velocities[idx] = wp.vec3(-_wheel_speed * rel_pos[2], 0.0, _wheel_speed * rel_pos[0]) + else: + # For points very close to rotation axis, just keep them as is + velocities[idx] = wp.vec3(0.0, 0.0, 0.0) + + +def post_process( + i, + post_process_interval, + f_current, + bc_mask, + vertices_wp, + precision_policy, + grid_shape, + usd_mesh_vorticity, + usd_mesh_q_criterion, + usd_stage, + vorticity_operator, + q_criterion_operator, + lag_forces=None, + cd_values=None, + reference_velocity=None, + frontal_area=None, + areas_wp=None, +): + if not isinstance(f_current, jnp.ndarray): + f_jax = warp_array_to_jax(f_current) + else: + f_jax = f_current + + macro_jax = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + ) + rho, u = macro_jax(f_jax) + u = u[:, 1:-1, 1:-1, 1:-1] + + fields = { + "u_magnitude": (u[0] ** 2.0 + u[1] ** 2.0 + u[2] ** 2.0) ** 0.5, + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + } + slice_idy = grid_shape[1] // 2 + save_image(fields["u_magnitude"][:, slice_idy, :], timestep=i) + save_fields_vtk(fields, i) + + # Calculate and store drag coefficient if forces are available + cd = calculate_drag_coefficient(lag_forces, reference_velocity, frontal_area, areas_wp) + cd_values.append((i, cd)) + if i % post_process_interval == 0: + print(f"Step {i}: Drag Coefficient (Cd) = {cd:.6f}") + # Save current Cd values to file + save_drag_coefficient(cd_values, "drag_coefficient.csv") + + save_usd_vorticity( + timestep=i, + post_process_interval=post_process_interval, + bc_mask=bc_mask, + f_current=f_current, + grid_shape=grid_shape, + usd_mesh=usd_mesh_vorticity, + vorticity_operator=vorticity_operator, + precision_policy=precision_policy, + vorticity_threshold=1e-2, + usd_stage=usd_stage, + device="cuda:1", + clip_lower=(20, 20, 0), + clip_upper=(20, 20, 20), + ) + + save_usd_q_criterion( + timestep=i, + post_process_interval=post_process_interval, + bc_mask=bc_mask, + f_current=f_current, + grid_shape=grid_shape, + usd_mesh=usd_mesh_q_criterion, + q_criterion_operator=q_criterion_operator, + precision_policy=precision_policy, + q_threshold=5e-6, + usd_stage=usd_stage, + device="cuda:1", + clip_lower=(20, 20, 0), + clip_upper=(20, 20, 20), + color_range=(0.0, 0.1), + ) + + update_usd_lagrangian_parts( + timestep=i, + post_process_interval=post_process_interval, + vertices_wp=vertices_wp, + parts=[ + { + "start": 0, + "end": num_body_vertices, + "faces": body_faces_np, + "usd_mesh": usd_car_body, + "colorize": True, + }, + { + "start": num_body_vertices, + "end": vertices_wp.shape[0], + "faces": wheels_faces_np, + "usd_mesh": usd_car_wheels, + "colorize": True, + }, + ], + vertex_offset=(20.0, 20.0, 0.0), + lag_forces=lag_forces, + device="cuda:1", + ) + + +def save_drag_coefficient(cd_values, filename): + """ + Save drag coefficient values to a CSV file + + Args: + cd_values: List of (timestep, cd) tuples + filename: Output CSV filename + """ + with open(filename, "w") as f: + f.write("timestep,cd\n") + for timestep, cd in cd_values: + f.write(f"{timestep},{cd}\n") + + # Also create a plot if matplotlib is available + try: + timesteps = [t for t, _ in cd_values] + cds = [cd for _, cd in cd_values] + + plt.figure(figsize=(10, 6)) + plt.plot(timesteps, cds, "b-") + plt.grid(True, linestyle="--", alpha=0.7) + plt.xlabel("Timestep") + plt.ylabel("Drag Coefficient (Cd)") + plt.title("Drag Coefficient vs Time") + plt.tight_layout() + plt.savefig("drag_coefficient.png", dpi=150) + plt.close() + except Exception as e: + print(f"Could not create Cd plot: {e}") + + +# grid_shape = (1024, 500, 200) +grid_shape = (256, 100, 100) +u_max = 0.02 +iter_per_flow_passes = grid_shape[0] / u_max +num_steps = int(iter_per_flow_passes * 2) +post_process_interval = 100 +print_interval = 100 +num_wheels = 4 +wheel_rotation_speed = -0.002 + +Re = 1e6 +clength = grid_shape[0] - 1 +visc = u_max * clength / Re +omega = 1.0 / (3.0 * visc + 0.5) + +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) +xlb.init(velocity_set=velocity_set, default_backend=compute_backend, default_precision_policy=precision_policy) +grid = grid_factory(grid_shape, compute_backend=compute_backend) + +print("Car Simulation Configuration:") +print(f" Grid size: {grid_shape}") +print(f" Omega: {omega}") +print(f" Backend: {compute_backend}") +print(f" Velocity set: {velocity_set}") +print(f" Precision policy: {precision_policy}") +print(f" Inlet velocity: {u_max}") +print(f" Reynolds number: {Re}") +print(f" Max steps: {num_steps}") + +# Initialize list to store drag coefficient values +cd_values = [] + +usd_output_directory = "usd_output_car" +os.makedirs(usd_output_directory, exist_ok=True) +usd_file = os.path.join(usd_output_directory, "car_output.usd") +usd_stage = Usd.Stage.CreateNew(usd_file) +usd_mesh_vorticity = UsdGeom.Mesh.Define(usd_stage, "/World/Vorticity") +usd_mesh_q_criterion = UsdGeom.Mesh.Define(usd_stage, "/World/QCriterion") +usd_car_body = UsdGeom.Mesh.Define(usd_stage, "/World/CarBody") +usd_car_wheels = UsdGeom.Mesh.Define(usd_stage, "/World/CarWheels") + +mesh_data = load_and_prepare_meshes_car(grid_shape) +vertices_wp = mesh_data["vertices_wp"] +areas_wp = mesh_data["areas_wp"] +faces_np = mesh_data["faces_np"] +wheel_ranges = mesh_data["wheel_ranges"] +wheel_centers = mesh_data["wheel_centers"] +num_body_vertices = mesh_data["num_body_vertices"] +num_wheel_vertices = mesh_data["num_wheel_vertices"] +body_faces_np = mesh_data["body_faces_np"] +wheels_faces_np = mesh_data["wheels_faces_np"] +frontal_area = mesh_data["frontal_area"] + +plot_object_placement( + vertices_wp, + grid_shape, + "car_placement.png", + "Car Placement (Top View)", + "Car", +) + +_num_body_vertices = wp.constant(num_body_vertices) +_num_wheels = wp.constant(num_wheels) +_wheel_speed = wp.constant(wheel_rotation_speed) + +starts_list = [rng[0] for rng in wheel_ranges] +ends_list = [rng[1] for rng in wheel_ranges] +cx_list = [c[0] for c in wheel_centers] +cy_list = [c[1] for c in wheel_centers] +cz_list = [c[2] for c in wheel_centers] + +_wheel_starts = wp.constant(wp.vec(len(starts_list), dtype=int)(starts_list)) +_wheel_ends = wp.constant(wp.vec(len(ends_list), dtype=int)(ends_list)) +_wheel_centers_x = wp.constant(wp.vec(len(cx_list), dtype=float)(cx_list)) +_wheel_centers_y = wp.constant(wp.vec(len(cy_list), dtype=float)(cy_list)) +_wheel_centers_z = wp.constant(wp.vec(len(cz_list), dtype=float)(cz_list)) + +bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max) +stepper = setup_stepper(grid, bc_list, omega) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + +device = "cuda:1" +with wp.ScopedDevice(device): + q_criterion_operator = QCriterion( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + vorticity_operator = Vorticity( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + +velocities_wp = wp.zeros(shape=vertices_wp.shape[0], dtype=wp.vec3) + +try: + for i in range(num_steps): + f_0, f_1, lag_forces = stepper( + f_0, + f_1, + vertices_wp, + areas_wp, + velocities_wp, + bc_mask, + missing_mask, + omega, + i, + ) + # Swap f_0 and f_1 + f_0, f_1 = f_1, f_0 + + # Update solid velocities and positions + wp.launch( + kernel=rotate_wheels, + dim=vertices_wp.shape[0], + inputs=[ + i, + lag_forces, + vertices_wp, + velocities_wp, + ], + ) + + if print_interval > 0 and i % print_interval == 0: + print(f"Step {i}/{num_steps} completed") + + if i % post_process_interval == 0 or i == num_steps - 1: + post_process( + i, + post_process_interval, + f_0, + bc_mask, + vertices_wp, + precision_policy, + grid_shape, + usd_mesh_vorticity, + usd_mesh_q_criterion, + usd_stage, + vorticity_operator, + q_criterion_operator, + lag_forces=lag_forces, + cd_values=cd_values, + reference_velocity=u_max, + frontal_area=frontal_area, + areas_wp=areas_wp, + ) + +except KeyboardInterrupt: + print("\nSimulation interrupted by user. Saving current USD state...") + current_time_code = i // post_process_interval + usd_stage.SetStartTimeCode(0) + usd_stage.SetEndTimeCode(current_time_code) + usd_stage.SetTimeCodesPerSecond(30) + usd_stage.Save() + + # Save drag coefficient data before exiting + if cd_values: + save_drag_coefficient(cd_values, "drag_coefficient.csv") + print("Drag coefficient data saved to drag_coefficient.csv") + + print(f"USD file saved with {current_time_code + 1} frames. Exiting.") + import sys + + sys.exit(0) + + +usd_stage.SetStartTimeCode(0) +usd_stage.SetEndTimeCode(num_steps // post_process_interval) +usd_stage.SetTimeCodesPerSecond(30) +usd_stage.Save() + +# Save final drag coefficient data +if cd_values: + save_drag_coefficient(cd_values, "drag_coefficient.csv") + print("Drag coefficient data saved to drag_coefficient.csv") + +print("Simulation finished. USD file saved.") diff --git a/examples/out_of_core/assets/nvidia_new.stl b/examples/out_of_core/assets/nvidia_new.stl new file mode 100644 index 00000000..d34f05c2 Binary files /dev/null and b/examples/out_of_core/assets/nvidia_new.stl differ diff --git a/examples/out_of_core/autodiff_lbm.py b/examples/out_of_core/autodiff_lbm.py new file mode 100644 index 00000000..a15412b2 --- /dev/null +++ b/examples/out_of_core/autodiff_lbm.py @@ -0,0 +1,418 @@ +# Wind tunnel simulation using the XLB library + +import os +import numpy as np +import warp as wp +from tqdm import tqdm +import logging +import mpi4py # TODO: actually learn how mpi works... + +mpi4py.rc.thread_level = "serialized" # or 'funneled' +import mpi4py.MPI as MPI +import argparse + +wp.init() +wp.clear_kernel_cache() + +# Import xlb +import xlb +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import FullwayBounceBackBC +from xlb.operator.boundary_masker import IndicesBoundaryMasker +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.macroscopic import Macroscopic + +# Local ooc imports +from ds import OOCGrid +from operators import ( + ClampField, + UniformInitializer, + InitializeTargetDensity, + L2Loss, +) +from subroutine import ( + PrepareFieldsSubroutine, + VolumeSaverSubroutine, + ForwardStepperSubroutine, + BackwardStepperSubroutine, + ForwardRhoLossSubroutine, + BackwardRhoLossSubroutine, + GradientDescentSubroutine, + InitializeFieldSubroutine, +) + +# Make command line parser +parser = argparse.ArgumentParser(description="Differential LBM") +parser.add_argument("--output_directory", type=str, default="autodiff_output", help="Output directory") +parser.add_argument("--final_stl_file", type=str, default="assets/nvidia_new.stl", help="STL file to match at the end of the simulation") +parser.add_argument("--max_base_velocity", type=float, default=0.02, help="Base velocity") +parser.add_argument("--shape", type=str, default="(128, 128, 128)", help="Shape") +parser.add_argument("--tau", type=float, default=0.53, help="Tau") +parser.add_argument("--nr_optimization_steps", type=int, default=1024, help="Nr optimization steps") +parser.add_argument("--nr_steps", type=int, default=256, help="Nr steps") +parser.add_argument("--checkpoint_frequency", type=int, default=16, help="Checkpoint frequency") +parser.add_argument("--save_state_frequency", type=int, default=16, help="Save volume frequency") +parser.add_argument("--collision", type=str, default="BGK", help="Collision") +parser.add_argument("--equilibrium", type=str, default="Quadratic", help="Equilibrium") +parser.add_argument("--velocity_set", type=str, default="D3Q19", help="Velocity set") +parser.add_argument("--ooc_block_shape", type=str, default="(64, 64, 64)", help="OOC block shape") +parser.add_argument("--nr_streams", type=int, default=1, help="Nr streams") +parser.add_argument("--comm", type=bool, default=True, help="Comm") +args = parser.parse_args() + + +def forward( + ooc_grid, + loss, + checkpoint_frequency, + nr_checkpoints, + forward_stepper_subroutine, + forward_rho_loss_subroutine, +): + # Zero the loss + loss.zero_() + + # Perform forward pass + for i in range(nr_checkpoints): + # Perform forward step + forward_stepper_subroutine( + ooc_grid, + nr_steps=checkpoint_frequency, + f_input_name=f"f_{str(i).zfill(4)}", + f_output_name=f"f_{str(i + 1).zfill(4)}", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + ) + + # Compute loss + forward_rho_loss_subroutine( + ooc_grid, + f_name=f"f_{str(nr_checkpoints).zfill(4)}", + target_rho_name="target_density", + loss=loss, + ) + + +def backward( + ooc_grid, + loss, + checkpoint_frequency, + nr_checkpoints, + backward_stepper_subroutine, + backward_rho_loss_subroutine, +): + # Set the loss gradient + loss.grad.fill_(1.0) + + # Perform backward pass + backward_rho_loss_subroutine( + ooc_grid, + f_name=f"f_{str(nr_checkpoints).zfill(4)}", + adj_f_name="adj_f", + target_rho_name="target_density", + loss=loss, + ) + + # Perform backward step + for i in range(nr_checkpoints - 1, -1, -1): + # Perform backward step + backward_stepper_subroutine( + ooc_grid, + nr_steps=checkpoint_frequency, + f_input_name=f"f_{str(i).zfill(4)}", + adj_f_name="adj_f", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + ) + + +if __name__ == "__main__": + # Set parameters + output_directory = args.output_directory + final_stl_file = args.final_stl_file + max_base_velocity = args.max_base_velocity + shape = eval(args.shape) + tau = args.tau + nr_optimization_steps = args.nr_optimization_steps + nr_steps = args.nr_steps + checkpoint_frequency = args.checkpoint_frequency + if args.save_state_frequency is None: + save_state_frequency = -1 + else: + save_state_frequency = args.save_state_frequency + collision = args.collision + equilibrium = args.equilibrium + velocity_set = args.velocity_set + ooc_block_shape = eval(args.ooc_block_shape) + ooc_ghost_cell_thickness = args.checkpoint_frequency * 2 + nr_streams = args.nr_streams + if args.comm: + comm = MPI.COMM_WORLD + else: + comm = None + + # Get fluid properties needed for the simulation + omega = 1.0 / tau + density = 1.0 + nr_steps = (nr_steps // checkpoint_frequency) * checkpoint_frequency + nr_checkpoints = nr_steps // checkpoint_frequency + + # Make output directory + os.makedirs(output_directory, exist_ok=True) + + # Make logging + logging.basicConfig(level=logging.INFO) + + # Log the parameters + logging.info(f"Output directory: {output_directory}") + logging.info(f"Final STL file: {final_stl_file}") + logging.info(f"Max base velocity: {max_base_velocity}") + logging.info(f"Shape: {shape}") + logging.info(f"Tau: {tau}") + logging.info(f"Omega: {omega}") + logging.info(f"Nr steps: {nr_steps}") + logging.info(f"Nr optimization steps: {nr_optimization_steps}") + logging.info(f"Save state frequency: {save_state_frequency}") + logging.info(f"Collision: {collision}") + logging.info(f"Equilibrium: {equilibrium}") + logging.info(f"Velocity set: {velocity_set}") + logging.info(f"OOC block shape: {ooc_block_shape}") + logging.info(f"OOC ghost cell thickness: {ooc_ghost_cell_thickness}") + logging.info(f"Nr streams: {nr_streams}") + + # Set the compute backend NOTE: hard coded for now + compute_backend = xlb.ComputeBackend.WARP + + # Set the precision policy NOTE: hard coded for now + precision_policy = xlb.PrecisionPolicy.FP32FP32 + + # Set the velocity set + if velocity_set == "D3Q27": + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) + elif velocity_set == "D3Q19": + velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend) + else: + raise ValueError("Invalid velocity set") + + # Initialize XLB + xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, + ) + + # Make grid for constructing stepper + grid = xlb.grid.WarpGrid(shape=shape) + + # Make boundary conditions + box = grid.bounding_box_indices() + walls = [box["top"][i] + box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + bc_walls = FullwayBounceBackBC( + indices=walls, + ) + boundary_conditions = [bc_walls] + indices_boundary_masker = IndicesBoundaryMasker( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + # Make stepper + stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type=collision, + ) + + # Make other operators + macroscopic = Macroscopic( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + quadratic_equilibrium = QuadraticEquilibrium( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + uniform_initializer = UniformInitializer( + initial_rho=density, + initial_u=(0.0, 0.0, 0.0), + ) + initialize_target_density = InitializeTargetDensity(file_path=final_stl_file, background_density=density, mesh_density=density + 0.0025) + l2_loss = L2Loss() + + # Make subroutines + prepare_fields_subroutine = PrepareFieldsSubroutine( + initializer=uniform_initializer, + equilibrium=quadratic_equilibrium, + boundary_conditions=boundary_conditions, + indices_boundary_masker=indices_boundary_masker, + nr_streams=nr_streams, + ) + forward_stepper_subroutine = ForwardStepperSubroutine( + stepper=stepper, + omega=omega, + nr_streams=nr_streams, + ) + backward_stepper_subroutine = BackwardStepperSubroutine( + stepper=stepper, + omega=omega, + nr_streams=nr_streams, + ) + forward_rho_loss_subroutine = ForwardRhoLossSubroutine( + macroscopic=macroscopic, + loss=l2_loss, + nr_streams=nr_streams, + ) + backward_rho_loss_subroutine = BackwardRhoLossSubroutine( + macroscopic=macroscopic, + loss=l2_loss, + nr_streams=nr_streams, + ) + volume_saver_subroutine = VolumeSaverSubroutine( + nr_streams=1, + ) + initialize_target_density_subroutine = InitializeFieldSubroutine( + initializer=initialize_target_density, + nr_streams=nr_streams, + ) + volume_saver_subroutine = VolumeSaverSubroutine() + gradient_descent_subroutine = GradientDescentSubroutine( + clamp_field=ClampField(), + nr_streams=nr_streams, + ) + + # Make OOC grid + ooc_grid = OOCGrid( + shape=shape, + block_shape=ooc_block_shape, + origin=(-1.0, -1.0, -1.0), + spacing=(2.0 / shape[0], 2.0 / shape[1], 2.0 / shape[2]), + ghost_cell_thickness=ooc_ghost_cell_thickness, + comm=comm, + ) + + # Make loss + loss = wp.zeros((1,), dtype=float, requires_grad=True) + + # Make min and max values for clamping + min_val = wp.from_numpy( + 0.9 * np.array(velocity_set.w), + dtype=wp.float32, + ) + max_val = wp.from_numpy( + 1.1 * np.array(velocity_set.w), + dtype=wp.float32, + ) + + # Initialize boxes for the OOC + ooc_grid.initialize_boxes( + name="target_density", + dtype=wp.float32, + cardinality=1, + ordering="SOA", + ) + ooc_grid.initialize_boxes( + name="boundary_id", + dtype=wp.uint8, + cardinality=1, + ordering="SOA", + ) + ooc_grid.initialize_boxes( + name="missing_mask", + dtype=wp.bool, + cardinality=velocity_set.q, + ordering="SOA", + ) + ooc_grid.initialize_boxes( + name="adj_f", + dtype=wp.float32, + cardinality=velocity_set.q, + ordering="SOA", + ) + for i in range(nr_checkpoints + 1): + ooc_grid.initialize_boxes( + name=f"f_{str(i).zfill(4)}", + dtype=wp.float32, + cardinality=velocity_set.q, + ordering="SOA", + ) + # Allocate ooc + ooc_grid.allocate() + print(f"nr gigs: {ooc_grid.nbytes // 1e9}") + + # Prepare fields + prepare_fields_subroutine( + ooc_grid, + f_name="f_0000", + ) + + # Initialize target velocity norm + initialize_target_density_subroutine( + ooc_grid, + field_name="target_density", + ) + + # Save target density + volume_saver_subroutine( + ooc_grid, + field_names=["target_density", "f_0000", "boundary_id", "missing_mask"], + file_name=os.path.join(output_directory, "target_density"), + ) + + # Start optimization + logging.info("Starting optimization") + for i in tqdm(range(nr_optimization_steps)): + # Perform forward pass + forward( + ooc_grid, + loss, + checkpoint_frequency, + nr_checkpoints, + forward_stepper_subroutine, + forward_rho_loss_subroutine, + ) + + # Perform backward pass + backward( + ooc_grid, + loss, + checkpoint_frequency, + nr_checkpoints, + backward_stepper_subroutine, + backward_rho_loss_subroutine, + ) + + # Perform gradient descent + gradient_descent_subroutine( + ooc_grid, + field_name="f_0000", + adj_field_name="adj_f", + learning_rate=0.001, + min_val=min_val, + max_val=max_val, + ) + + # Print loss + logging.info(f"Loss: {loss.numpy()[0]}") + + # Check if loss is nan + if np.isnan(loss.numpy()[0]): + logging.info("Loss is nan") + break + + # Save volume + if i % save_state_frequency == 0 and i != 0: + volume_saver_subroutine( + ooc_grid, + field_names=["adj_f"], + file_name=os.path.join(output_directory, "adj_f"), + ) + for j in range(nr_checkpoints + 1): + volume_saver_subroutine( + ooc_grid, + field_names=[f"f_{str(j).zfill(4)}"], + file_name=os.path.join(output_directory, f"state_{str(j).zfill(4)}"), + ) diff --git a/examples/out_of_core/ds/__init__.py b/examples/out_of_core/ds/__init__.py new file mode 100644 index 00000000..99934df7 --- /dev/null +++ b/examples/out_of_core/ds/__init__.py @@ -0,0 +1,6 @@ +from ds.ooc_grid import ( + Box, + Block, + OOCGrid, + MemoryPool, +) diff --git a/examples/out_of_core/ds/ooc_grid.py b/examples/out_of_core/ds/ooc_grid.py new file mode 100644 index 00000000..dbf4cef0 --- /dev/null +++ b/examples/out_of_core/ds/ooc_grid.py @@ -0,0 +1,624 @@ +import numpy as np +import warp as wp +from tqdm import tqdm +import logging +import itertools +import pyvista as pv +import gc + + +class Box: + def __init__( + self, + extent, + offset, + origin, + spacing, + cardinality, + ordering, + dtype, + device, + ): + # Check valid ordering + assert ordering in ["AOS", "SOA"], f"Unknown ordering {ordering}, must be 'AOS' or 'SOA'." + + # Parameters + self.extent = extent + self.offset = offset + self.origin = origin + self.spacing = spacing + self.cardinality = cardinality + self.ordering = ordering + self.dtype = dtype + self.device = device + + # Allocate data + self.data = None + + @property + def local_origin(self): + return self.origin + np.array(self.offset) * np.array(self.spacing) + + @property + def local_spacing(self): + return np.array(self.spacing) + + @property + def shape(self): + return tuple(self.extent) + + @property + def data_shape(self): + if self.ordering == "AOS": + return list(self.shape) + [self.cardinality] + elif self.ordering == "SOA": + return [self.cardinality] + list(self.shape) + + @property + def nbytes(self): + if self.data is None: + return 0 + else: + return self.data.capacity + + def allocate( + self, + ): + # Delete data if exists + if self.data is not None: + del self.data + + # Allocate data + if not np.any([s == 0 for s in self.data_shape]): + self.data = wp.zeros( + self.data_shape, + dtype=self.dtype, + pinned=True if self.device == "cpu" else None, + device=self.device, + ) + + @staticmethod + def _box_intersection( + extent_1, + offset_1, + extent_2, + offset_2, + global_extent, + ): + # Get min and max of boxes + min_1 = offset_1 + max_1 = offset_1 + extent_1 + min_2 = offset_2 + max_2 = offset_2 + extent_2 + + # Get intersection + min_intersection = np.maximum(min_1, min_2) + max_intersection = np.minimum(max_1, max_2) + extent_intersection = np.maximum(max_intersection - min_intersection, 0) + offset_intersection = min_intersection + + # Check if intersection is valid + return extent_intersection, offset_intersection + + +class Block: + def __init__( + self, + extent, + offset, + origin, + spacing, + ghost_cell_thickness, + device, + pid=0, + ): + # Set the parameters + self.extent = np.array(extent, dtype=np.int32) + self.offset = np.array(offset, dtype=np.int32) + self.origin = np.array(origin, dtype=np.float32) + self.spacing = np.array(spacing, dtype=np.float32) + self.ghost_cell_thickness = np.array(ghost_cell_thickness, dtype=np.int32) + self.device = device + self.pid = pid + + # Make set for neighbour blocks + self.neighbour_blocks = set() + + # Make dict for boxes + self.boxes = {} + + # Make dict for particles + self.particles = {} + + # Make list for ghost boxes/cells + self.local_ghost_boxes = {} # Ghost boxes from local edges, these will be sent to neighbours + self.neighbour_ghost_boxes = {} # Neighbour blocks + self.neighbour_ghost_boxes_buffer = {} # Buffer for receiving ghost boxes + + @property + def extent_with_ghost(self): + return self.extent + 2 * self.ghost_cell_thickness + + @property + def shape(self): + return tuple(self.extent) + + @property + def shape_with_ghost(self): + return tuple(self.extent_with_ghost) + + @property + def offset_with_ghost(self): + return self.offset - self.ghost_cell_thickness + + @property + def local_origin(self): + return self.origin + np.array(self.offset) * np.array(self.spacing) + + @property + def local_spacing(self): + return np.array(self.spacing) + + @property + def nbytes(self): + nbytes = 0 + for box in self.boxes.values(): + nbytes += box.nbytes + for boxes in self.local_ghost_boxes.values(): + for box in boxes.values(): + nbytes += box.nbytes + for boxes in self.neighbour_ghost_boxes.values(): + for box in boxes.values(): + nbytes += box.nbytes + for boxes in self.neighbour_ghost_boxes_buffer.values(): + for box in boxes.values(): + nbytes += box.nbytes + return nbytes + + def add_neighbour_block(self, block): + self.neighbour_blocks.add(block) + self.local_ghost_boxes[block] = {} + self.neighbour_ghost_boxes[block] = {} + self.neighbour_ghost_boxes_buffer[block] = {} + + def remove_neighbour_block(self, block): + if block in self.neighbour_blocks: + self.neighbour_blocks.remove(block) + self.local_ghost_boxes.pop(block) + self.neighbour_ghost_boxes.pop(block) + self.neighbour_ghost_boxes_buffer.pop(block) + + def initialize_box( + self, + name, + dtype, + cardinality, + ordering, + global_extent, + extent=None, + offset=None, + ): + # Get extent and offset, if None use whole domain + extent = extent if extent is not None else global_extent + offset = offset if offset is not None else np.zeros(len(global_extent), dtype=np.int32) + + # Get intersection + local_extent, local_offset = Box._box_intersection( + extent, + offset, + self.extent, + self.offset, + global_extent, + ) + + # Add box to block + self.boxes[name] = Box( + local_extent, + local_offset, + origin=self.origin, + spacing=self.spacing, + cardinality=cardinality, + ordering=ordering, + dtype=dtype, + device=self.device, + ) + + # Initialize local ghost boxes + for neighbour_block in self.neighbour_blocks: + # Get intersections + local_ghost_extent, local_ghost_offset = Box._box_intersection( + extent, + offset, + neighbour_block.extent_with_ghost, + neighbour_block.offset_with_ghost, + global_extent, + ) + local_ghost_extent, local_ghost_offset = Box._box_intersection( + self.extent, + self.offset, + local_ghost_extent, + local_ghost_offset, + global_extent, + ) + neighbour_ghost_extent, neighbour_ghost_offset = Box._box_intersection( + extent, + offset, + self.extent_with_ghost, + self.offset_with_ghost, + global_extent, + ) + neighbour_ghost_extent, neighbour_ghost_offset = Box._box_intersection( + neighbour_block.extent, + neighbour_block.offset, + neighbour_ghost_extent, + neighbour_ghost_offset, + global_extent, + ) + + # Add local ghost box + self.local_ghost_boxes[neighbour_block][name] = Box( + extent=local_ghost_extent, + offset=local_ghost_offset, + origin=self.origin, + spacing=self.spacing, + cardinality=cardinality, + ordering=ordering, + dtype=dtype, + device=self.device, + ) + + # Add neighbour ghost box + self.neighbour_ghost_boxes[neighbour_block][name] = Box( + extent=neighbour_ghost_extent, + offset=neighbour_ghost_offset, + origin=self.origin, + spacing=self.spacing, + cardinality=cardinality, + ordering=ordering, + dtype=dtype, + device=self.device, + ) + self.neighbour_ghost_boxes_buffer[neighbour_block][name] = Box( + extent=neighbour_ghost_extent, + offset=neighbour_ghost_offset, + origin=self.origin, + spacing=self.spacing, + cardinality=cardinality, + ordering=ordering, + dtype=dtype, + device=self.device, + ) + + def initialize_particles( + self, + ): + pass + + def allocate( + self, + ): + for box in self.boxes.values(): + box.allocate() + for boxes in self.local_ghost_boxes.values(): + for box in boxes.values(): + box.allocate() + for boxes in self.neighbour_ghost_boxes.values(): + for box in boxes.values(): + box.allocate() + for boxes in self.neighbour_ghost_boxes_buffer.values(): + for box in boxes.values(): + box.allocate() + + def send_ghost_boxes(self, comm=None, comm_tag=0, names=None): + # Get names if None + if names is None: + names = list(self.boxes.keys()) + + # Get pid + if comm is not None: + pid = comm.Get_rank() + else: + pid = 0 + + # Make list for send requests + requests = [] + + # Loop over neighbour blocks + for neighbour_block, ghost_boxes in self.local_ghost_boxes.items(): + # Loop over ghost boxes + for name, ghost_box in ghost_boxes.items(): + # Check if required to send + if name not in names: + continue + + # 4 Cases: + # 1. Current pid is the same as block and neighbour + # 2. Current pid is the same as block but different than neighbour + # 3. Current pid is different than block but the same as neighbour + # 4. Current pid is different than block and neighbour + # Case 1 + if (pid == self.pid) and (pid == neighbour_block.pid): + # Swap data + local_data = self.local_ghost_boxes[neighbour_block][name].data + neighbour_data = neighbour_block.neighbour_ghost_boxes_buffer[self][name].data + self.local_ghost_boxes[neighbour_block][name].data = neighbour_data + neighbour_block.neighbour_ghost_boxes_buffer[self][name].data = local_data + + # Case 2 + if (pid == self.pid) and (pid != neighbour_block.pid): + # Send data + requests.append( + comm.Isend( + self.local_ghost_boxes[neighbour_block][name].data, + dest=neighbour_block.pid, + tag=comm_tag, + ) + ) + + # Case 3 + if (pid != self.pid) and (pid == neighbour_block.pid): + # Receive data + requests.append( + comm.Irecv( + neighbour_block.neighbour_ghost_boxes_buffer[self][name].data, + source=self.pid, + tag=comm_tag, + ) + ) + + # Case 4 + if (pid != self.pid) and (pid != neighbour_block.pid): + pass + + # Update tag + comm_tag += 1 + + return requests, comm_tag + + def to_image_data( + self, + include_ghost=False, # Just for debugging + ): + # Return grids + grids = [] + + # Make function for converting data + def _convert_data(box): + if box.data is not None: + if box.ordering == "AOS": + return box.data.numpy().reshape((-1, box.cardinality), order="F") + elif box.ordering == "SOA": + np_data = box.data.numpy() + aos_data = np.stack([np_data[i, ...] for i in range(box.cardinality)], axis=-1) + return aos_data.reshape((-1, box.cardinality), order="F") + + # Make center image data + grid = pv.ImageData( + dimensions=np.array(self.shape) + 1, + origin=self.local_origin, + spacing=self.local_spacing, + ) + + # Add data + for name, box in self.boxes.items(): + grid.cell_data[name] = _convert_data(box) + + # Add grid to grids + grids.append(grid) + + # Add ghost data + if include_ghost: + # Add local ghost data + for ghost_boxes in self.local_ghost_boxes.values(): + for ghost_name, ghost_box in ghost_boxes.items(): + grid = pv.ImageData( + dimensions=np.array(ghost_box.shape) + 1, + origin=ghost_box.local_origin, + spacing=ghost_box.local_spacing, + ) + grid.cell_data[ghost_name + "_local_ghost"] = _convert_data(ghost_box) + grids.append(grid) + for ghost_boxes in self.neighbour_ghost_boxes.values(): + for ghost_name, ghost_box in ghost_boxes.items(): + grid = pv.ImageData( + dimensions=np.array(ghost_box.shape) + 1, + origin=ghost_box.local_origin, + spacing=ghost_box.local_spacing, + ) + grid.cell_data[name + "_neighbour_ghost"] = _convert_data(ghost_box) + grids.append(grid) + return grids + + def swap_buffers(self, names=None): + for neighbour_boxes, neighbour_boxes_buffer in zip(self.neighbour_ghost_boxes.values(), self.neighbour_ghost_boxes_buffer.values()): + if names is None: + names = list(neighbour_boxes.keys()) + for name in names: + neighbour_boxes[name].data, neighbour_boxes_buffer[name].data = neighbour_boxes_buffer[name].data, neighbour_boxes[name].data + + +class OOCGrid: + """An out-of-core Adaptive Mesh Refinement grid data structure.""" + + def __init__( + self, + shape, + block_shape, + origin=None, + spacing=None, + ghost_cell_thickness=1, + comm=None, + pid_device_mapping=None, + ): + """Initialize the out-of-core data structure.""" + + # Set the parameters + self.shape = shape + self.block_shape = block_shape + self.origin = origin if origin is not None else tuple(np.zeros(len(shape))) + self.spacing = spacing if spacing is not None else tuple(np.ones(len(shape))) + if isinstance(ghost_cell_thickness, int): + ghost_cell_thickness = (ghost_cell_thickness,) * len(shape) + self.ghost_cell_thickness = ghost_cell_thickness + self.comm = comm + + # Check that the block shape divides the shape + self.block_dims = tuple([shape[i] // block_shape[i] for i in range(len(shape))]) + + # Get process id and number of processes + if comm is None: + self.pid = 0 + self.size = 1 + else: + self.pid = comm.Get_rank() + self.size = comm.Get_size() + if pid_device_mapping is None: + pid_device_mapping = ["cpu" for _ in range(self.size)] + + # dist = np.arange(self.block_dims[0] * self.block_dims[1] * self.block_dims[2]) + # np.random.shuffle(dist) + + # Initialize blocks and connections + logging.info("Initializing blocks and connections...") + self.blocks = {} + for index, block_index in tqdm(enumerate(itertools.product(*[range(n) for n in self.block_dims]))): + # Get dist + # dist = hilb.distance_from_point([block_index[2], block_index[1], block_index[0]]) + dist = index + + # Get block pid + # block_pid = (dist // self.size) % self.size + block_pid = (dist) % self.size + + # Get device + device = pid_device_mapping[block_pid] + + # Create block + block = Block( + extent=block_shape, + offset=[i * s for i, s in zip(block_index, block_shape)], + origin=self.origin, + spacing=self.spacing, + ghost_cell_thickness=ghost_cell_thickness, + device=device, + pid=block_pid, + ) + + # Add block to blocks + self.blocks[block_index] = block + + # Initialize connections + for block_index in self.blocks.keys(): + # Get neighbour block indices + for direction in itertools.product(*[range(-1, 2) for _ in self.block_dims]): + # Skip if no neighbour + if np.all([d == 0 for d in direction]): + continue + + # Get neighbour block index + # neigh_block_index = tuple([(i + d) % n for i, d, n in zip(block_index, direction, self.block_dims)]) + neigh_block_index = tuple([(i + d) for i, d in zip(block_index, direction)]) + + # Add neighbour block to block + if neigh_block_index in self.blocks: + self.blocks[block_index].add_neighbour_block(self.blocks[neigh_block_index]) + + # Barrier + if self.comm is not None: + self.comm.Barrier() + + @property + def nbytes(self): + nbytes = 0 + for block in self.blocks.values(): + nbytes += block.nbytes + # if self.comm is not None: + # nbytes = self.comm.allreduce(nbytes, op=MPI.SUM) + return nbytes + + def initialize_boxes( + self, + name, + dtype, + cardinality, + ordering="SOA", + extent=None, + offset=None, + ): + # Initialize boxes + for block in self.blocks.values(): + # Initialize box + block.initialize_box( + name=name, + dtype=dtype, + cardinality=cardinality, + ordering=ordering, + global_extent=self.shape, + extent=extent, + offset=offset, + ) + + def initialize_particles( + self, + ): + pass + + def allocate( + self, + ): + for block in self.blocks.values(): + if block.pid == self.pid: + block.allocate() + + def save_vtm( + self, + filename, + ): + # Create multi block dataset + mb = pv.MultiBlock() + + # Loop over blocks + for block in self.blocks.values(): + # Add block to multi block + mb.extend(block.to_image_data()) + + # Save multi block + mb.save(filename) + + +class MemoryPool: + def __init__(self): + self.pool = {} + + def clear(self): + for key in list(self.pool.keys()): + for array in self.pool[key]: + del array + del self.pool[key] + self.pool[key] = [] + wp.synchronize() + gc.collect() + + def get(self, shape, dtype, requires_grad=False): + key = (tuple(shape), dtype, requires_grad) + if key not in self.pool: + self.pool[key] = [] + if len(self.pool[key]) == 0: + self.pool[key].append(wp.zeros(shape, dtype=dtype, requires_grad=requires_grad)) + return self.pool[key].pop() + + def ret(self, array, zero=True): + key = (tuple(array.shape), array.dtype, array.requires_grad) + if zero: + array.zero_() + if array.requires_grad: + array.grad.zero_() + self.pool[key].append(array) + + @property + def nbytes(self): + nbytes = 0 + for key in self.pool.keys(): + for array in self.pool[key]: + nbytes += array.capacity + return nbytes diff --git a/examples/out_of_core/lid_driven_cavity.py b/examples/out_of_core/lid_driven_cavity.py new file mode 100644 index 00000000..65085989 --- /dev/null +++ b/examples/out_of_core/lid_driven_cavity.py @@ -0,0 +1,288 @@ +# Lid Drive Cavity using out-of-core memory with XLB library + +import os +import warp as wp +import numpy as np +from tqdm import tqdm +import logging +import mpi4py # TODO: actually learn how mpi works... + +mpi4py.rc.thread_level = "serialized" # or 'funneled' +import mpi4py.MPI as MPI +import argparse +import math + +wp.init() +wp.clear_kernel_cache() + +# Import xlb +import xlb +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import EquilibriumBC, FullwayBounceBackBC +from xlb.operator.boundary_masker import IndicesBoundaryMasker +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.macroscopic import Macroscopic + +# Local ooc imports +from ds import OOCGrid +from operators import UniformInitializer +from subroutine import ( + PrepareFieldsSubroutine, + StepperSubroutine, + RenderQCriterionSubroutine, + VolumeSaverSubroutine, +) + +# Make command line parser +parser = argparse.ArgumentParser(description="Lid driven cavity simulation") +parser.add_argument("--output_directory", type=str, default="ldc_output", help="Output directory") +parser.add_argument("--base_velocity", type=float, default=0.06, help="Base velocity") +parser.add_argument("--shape", type=str, default="(256, 256, 256)", help="Shape") +parser.add_argument("--tau", type=float, default=0.501, help="Tau") +parser.add_argument("--nr_steps", type=int, default=131072, help="Nr steps") +parser.add_argument("--save_q_criterion_frequency", type=int, default=128, help="Save q criterion frequency") +parser.add_argument("--q_criterion_threshold", type=float, default=1e-6, help="Q criterion threshold") +parser.add_argument("--save_volume_debug", type=bool, default=False, help="Save volume as vtk to debug") +parser.add_argument("--collision", type=str, default="SmagorinskyLESBGK", help="Collision") +parser.add_argument("--equilibrium", type=str, default="Quadratic", help="Equilibrium") +parser.add_argument("--velocity_set", type=str, default="D3Q19", help="Velocity set") +parser.add_argument("--ooc_block_shape", type=str, default="(128, 128, 128)", help="OOC block shape") +parser.add_argument("--ooc_ghost_cell_thickness", type=int, default=16, help="OOC ghost cell thickness") +parser.add_argument("--nr_streams", type=int, default=2, help="Nr streams") +parser.add_argument("--comm", type=bool, default=True, help="Comm") +args = parser.parse_args() + +if __name__ == "__main__": + # Set parameters + output_directory = args.output_directory + base_velocity = args.base_velocity + shape = eval(args.shape) + tau = args.tau + nr_steps = args.nr_steps + if args.save_q_criterion_frequency is None: + save_q_criterion_frequency = -1 + else: + save_q_criterion_frequency = (args.save_q_criterion_frequency // args.ooc_ghost_cell_thickness) * args.ooc_ghost_cell_thickness + q_criterion_threshold = args.q_criterion_threshold + collision = args.collision + equilibrium = args.equilibrium + velocity_set = args.velocity_set + ooc_block_shape = eval(args.ooc_block_shape) + ooc_ghost_cell_thickness = args.ooc_ghost_cell_thickness + nr_streams = args.nr_streams + if args.comm: + comm = MPI.COMM_WORLD + else: + comm = None + + # Get fluid properties needed for the simulation + omega = 1.0 / tau + density = 1.0 + nr_steps = (nr_steps // ooc_ghost_cell_thickness) * ooc_ghost_cell_thickness # Make sure steps is divisible by ghost cell thickness + + # Make output directory + os.makedirs(output_directory, exist_ok=True) + + # Make logging + logging.basicConfig(level=logging.INFO) + + # Log the parameters + logging.info(f"Base velocity: {base_velocity}") + logging.info(f"Shape: {shape}") + logging.info(f"Tau: {tau}") + logging.info(f"Omega: {omega}") + logging.info(f"Nr steps: {nr_steps}") + logging.info(f"Save q criterion frequency: {save_q_criterion_frequency}") + logging.info(f"Collision: {collision}") + logging.info(f"Equilibrium: {equilibrium}") + logging.info(f"Velocity set: {velocity_set}") + logging.info(f"OOC block shape: {ooc_block_shape}") + logging.info(f"OOC ghost cell thickness: {ooc_ghost_cell_thickness}") + logging.info(f"Nr streams: {nr_streams}") + + # Set the compute backend NOTE: hard coded for now + compute_backend = xlb.ComputeBackend.WARP + + # Set the precision policy NOTE: hard coded for now + precision_policy = xlb.PrecisionPolicy.FP32FP32 + + # Set the velocity set + if velocity_set == "D3Q27": + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) + elif velocity_set == "D3Q19": + velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend) + else: + raise ValueError("Invalid velocity set") + + # Initialize XLB + xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, + ) + + # Make grid for constructing stepper + grid = xlb.grid.WarpGrid(shape=shape) + + # Make boundary conditions + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + lid = box_no_edge["top"] + walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + bc_top = EquilibriumBC( + rho=density, + u=(0.0, base_velocity, 0.0), + indices=lid, + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + # bc_walls = HalfwayBounceBackBC( + bc_walls = FullwayBounceBackBC( + indices=walls, + ) + boundary_conditions = [bc_walls, bc_top] + indices_boundary_masker = IndicesBoundaryMasker( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + # Make stepper + stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type=collision, + ) + + # Make other operators + macroscopic = Macroscopic( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + quadratic_equilibrium = QuadraticEquilibrium( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + uniform_initializer = UniformInitializer( + initial_rho=density, + initial_u=(0.0, 0.0, 0.0), + ) + + # Make subroutines + prepare_fields_subroutine = PrepareFieldsSubroutine( + initializer=uniform_initializer, + equilibrium=quadratic_equilibrium, + boundary_conditions=boundary_conditions, + indices_boundary_masker=indices_boundary_masker, + nr_streams=nr_streams, + ) + stepper_subroutine = StepperSubroutine( + stepper=stepper, + omega=omega, + nr_streams=nr_streams, + ) + volume_saver_subroutine = VolumeSaverSubroutine( + nr_streams=1, + ) + render_q_criterion_subroutine = RenderQCriterionSubroutine( + macroscopic=macroscopic, + nr_streams=1, + ) + + # Make OOC + ooc_grid = OOCGrid( + shape=shape, + block_shape=ooc_block_shape, + origin=(0.0, 0.0, 0.0), + spacing=(1.0 / shape[0], 1.0 / shape[1], 1.0 / shape[2]), + ghost_cell_thickness=ooc_ghost_cell_thickness, + comm=comm, + ) + + # Initialize boxes for the OOC + ooc_grid.initialize_boxes( + name="f", + dtype=wp.float32, + cardinality=velocity_set.q, + ordering="SOA", + ) + ooc_grid.initialize_boxes( + name="boundary_id", + dtype=wp.uint8, + cardinality=1, + ordering="SOA", + ) + ooc_grid.initialize_boxes( + name="missing_mask", + dtype=wp.bool, + cardinality=velocity_set.q, + ordering="SOA", + ) + + # Allocate ooc + ooc_grid.allocate() + + # Make pixel buffer (UHD) + pixel_buffer = wp.zeros((2160, 3840, 4), dtype=wp.float32) + depth_buffer = wp.zeros((2160, 3840), dtype=wp.float32) + + # Prepare fields + prepare_fields_subroutine(ooc_grid) + + # Save fields + volume_saver_subroutine( + ooc_grid, + field_names=["f", "boundary_id", "missing_mask"], + file_name=os.path.join(output_directory, "initial"), + ) + + # Start simulation + logging.info("Starting simulation") + for i in tqdm(range(nr_steps // ooc_ghost_cell_thickness)): + # Perform stepper + stepper_subroutine(ooc_grid) + + # Save volume and render q criterion + if (i * ooc_ghost_cell_thickness) % save_q_criterion_frequency == 0 and save_q_criterion_frequency != -1: + # Calculate camera position for orbit + total_frames = nr_steps // save_q_criterion_frequency + current_frame = i * ooc_ghost_cell_thickness // save_q_criterion_frequency + angle = (current_frame / total_frames) * 2 * math.pi # 0 to 2Ο€ + + # Camera parameters + radius = 1.3 # Distance from center + center = (0.5, 0.5, 0.5) # Center of domain + + # Calculate camera position + camera_x = center[0] + radius * math.cos(angle) + camera_y = center[1] - radius * math.sin(angle) + camera_z = center[2] + + pixel_buffer.fill_(0.0) + depth_buffer.fill_(10.0) + render_q_criterion_subroutine( + ooc_grid, + os.path.join(output_directory, f"q_criterion_{i:06d}"), # Zero-padded frame numbers + pixel_buffer, + depth_buffer, + camera_pos=(camera_x, camera_y, camera_z), + camera_target=center, # Look at center + camera_up=(0.0, 0.0, 1.0), # Keep camera upright + fov_degrees=60.0, + ambient_intensity=0.05, + edge_sharpness=1.0, + gamma=1.0, + q_criterion_threshold=q_criterion_threshold, + vmin=0.0, + vmax=0.01, + ) + + # Save + if args.save_volume_debug: + volume_saver_subroutine( + ooc_grid, field_names=["f", "boundary_id", "missing_mask"], file_name=os.path.join(output_directory, f"time_{str(i).zfill(5)}") + ) diff --git a/examples/out_of_core/operators/__init__.py b/examples/out_of_core/operators/__init__.py new file mode 100644 index 00000000..81c5ad68 --- /dev/null +++ b/examples/out_of_core/operators/__init__.py @@ -0,0 +1,25 @@ +from .soa_copy import SOACopy +from .trilinear_interpolation import TrilinearInterpolation +from .mesh_renderer import MeshRenderer +from .color_mapper import ColorMapper +from .transform_mesh import TransformMesh +from .q_criterion import QCriterion +from .uniform_initializer import UniformInitializer +from .gradient_descent import GradientDescent +from .clamp_field import ClampField +from .initialize_target_density import InitializeTargetDensity +from .l2_loss import L2Loss + +__all__ = [ + "SOACopy", + "TrilinearInterpolation", + "MeshRenderer", + "ColorMapper", + "TransformMesh", + "QCriterion", + "UniformInitializer", + "GradientDescent", + "ClampField", + "InitializeTargetDensity", + "L2Loss", +] diff --git a/examples/out_of_core/operators/clamp_field.py b/examples/out_of_core/operators/clamp_field.py new file mode 100644 index 00000000..d81eebb2 --- /dev/null +++ b/examples/out_of_core/operators/clamp_field.py @@ -0,0 +1,39 @@ +from typing import Any +import warp as wp + + +class ClampField: + """ + Clamp field operator. + """ + + @wp.kernel + def clamp_field( + field: wp.array4d(dtype=Any), + min_val: wp.array(dtype=Any), + max_val: wp.array(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + + # Update the field + for ii in range(field.shape[0]): + field[ii, i, j, k] = wp.max(min_val[ii], wp.min(max_val[ii], field[ii, i, j, k])) + + def __call__( + self, + field: wp.array4d(dtype=Any), + min_val: wp.array(dtype=Any), + max_val: wp.array(dtype=Any), + ): + # Launch the warp kernel + wp.launch( + self.clamp_field, + inputs=[ + field, + min_val, + max_val, + ], + dim=field.shape[1:], + ) + return field diff --git a/examples/out_of_core/operators/color_mapper.py b/examples/out_of_core/operators/color_mapper.py new file mode 100644 index 00000000..33f29340 --- /dev/null +++ b/examples/out_of_core/operators/color_mapper.py @@ -0,0 +1,95 @@ +import warp as wp + + +class ColorMapper: + """ + Operator for mapping scalar values to RGB colors using different colormaps. + + Currently supported colormaps: + - 'jet': Blue -> Cyan -> Yellow -> Red + """ + + @staticmethod + @wp.func + def jet_colormap(value: float) -> wp.vec3: + """ + Map a value in [0,1] to RGB colors using the jet colormap. + """ + r = wp.clamp(wp.min(4.0 * value - 1.5, -4.0 * value + 4.5), 0.0, 1.0) + g = wp.clamp(wp.min(4.0 * value - 0.5, -4.0 * value + 3.5), 0.0, 1.0) + b = wp.clamp(wp.min(4.0 * value + 0.5, -4.0 * value + 2.5), 0.0, 1.0) + + return wp.vec3(r, g, b) + + @wp.kernel + def _map_colors( + values: wp.array(dtype=float), + colors: wp.array2d(dtype=float), + vmin: float, + vmax: float, + ): + # Get thread index + idx = wp.tid() + + # Normalize value to [0,1] range + value = values[idx] + normalized = (value - vmin) / (vmax - vmin) + normalized = wp.clamp(normalized, 0.0, 1.0) + + # Map to color using jet colormap + color = ColorMapper.jet_colormap(normalized) + + # Store RGB values + colors[idx, 0] = color[0] + colors[idx, 1] = color[1] + colors[idx, 2] = color[2] + + def __call__( + self, + values: wp.array, + colors: wp.array2d, + vmin: float, + vmax: float, + colormap: str = "jet", + ) -> wp.array2d: + """ + Map scalar values to RGB colors using the specified colormap. + + Parameters + ---------- + values : wp.array(dtype=float) + Input scalar values to map to colors + colors : wp.array2d(dtype=float) + Output RGB colors array with shape (len(values), 3) + vmin : float + Minimum value for normalization + vmax : float + Maximum value for normalization + colormap : str + Name of the colormap to use (currently only 'jet' is supported) + + Returns + ------- + wp.array2d(dtype=float) + Reference to the input colors array + """ + if colormap != "jet": + raise ValueError(f"Colormap '{colormap}' not supported. Use 'jet'.") + + # Verify input shapes + assert len(values.shape) == 1, "Values must be a 1D array" + assert colors.shape == (len(values), 3), f"Colors array must have shape ({len(values)}, 3)" + + # Verify vmin/vmax + assert isinstance(vmin, float), "vmin must be a float" + assert isinstance(vmax, float), "vmax must be a float" + assert vmax > vmin, f"vmax ({vmax}) must be greater than vmin ({vmin})" + + # Launch kernel + wp.launch( + self._map_colors, + dim=len(values), + inputs=[values, colors, vmin, vmax], + ) + + return colors diff --git a/examples/out_of_core/operators/gradient_descent.py b/examples/out_of_core/operators/gradient_descent.py new file mode 100644 index 00000000..f76780c9 --- /dev/null +++ b/examples/out_of_core/operators/gradient_descent.py @@ -0,0 +1,39 @@ +from typing import Any +import warp as wp + + +class GradientDescent: + """ + Gradient descent operator. + """ + + @wp.kernel + def gradient_decent( + field: wp.array4d(dtype=Any), + adj_field: wp.array4d(dtype=Any), + learning_rate: wp.float32, + ): + # Get the global index + i, j, k = wp.tid() + + # Update the field + for ii in range(field.shape[0]): + field[ii, i, j, k] -= learning_rate * adj_field[ii, i, j, k] + + def __call__( + self, + field: wp.array4d(dtype=Any), + adj_field: wp.array4d(dtype=Any), + learning_rate: float, + ): + # Launch the warp kernel + wp.launch( + self.gradient_decent, + inputs=[ + field, + adj_field, + learning_rate, + ], + dim=field.shape[1:], + ) + return field diff --git a/examples/out_of_core/operators/initialize_target_density.py b/examples/out_of_core/operators/initialize_target_density.py new file mode 100644 index 00000000..38778e92 --- /dev/null +++ b/examples/out_of_core/operators/initialize_target_density.py @@ -0,0 +1,81 @@ +from typing import Any +import warp as wp +import numpy as np +import stl as np_mesh + + +class InitializeTargetDensity: + """ + Initialize target density operator. + """ + + def __init__(self, file_path: str, background_density: float, mesh_density: float): + # Load the mesh + mesh = np_mesh.Mesh.from_file(file_path) + mesh_points = mesh.points.reshape(-1, 3) + mesh_indices = np.arange(mesh_points.shape[0]) + self.mesh = wp.Mesh( + points=wp.array(mesh_points, dtype=wp.vec3), + indices=wp.array(mesh_indices, dtype=int), + ) + self.background_density = background_density + self.mesh_density = mesh_density + + @wp.kernel + def _initialize_target_density( + rho: wp.array4d(dtype=Any), + mesh: wp.uint64, + background_density: float, + mesh_density: float, + origin: wp.vec3f, + spacing: wp.vec3f, + ): + # get spatial index + i, j, k = wp.tid() + + # position of voxel (cell center) + ijk = wp.vec3(wp.float32(i), wp.float32(j), wp.float32(k)) + ijk = ijk + wp.vec3(0.5, 0.5, 0.5) # cell center + pos = wp.cw_mul(ijk, spacing) + origin + + # Compute maximum distance to check + max_length = wp.sqrt( + (spacing[0] * wp.float32(rho.shape[0])) ** 2.0 + + (spacing[1] * wp.float32(rho.shape[1])) ** 2.0 + + (spacing[2] * wp.float32(rho.shape[2])) ** 2.0 + ) + + # evaluate distance of point + face_index = int(0) + face_u = float(0.0) + face_v = float(0.0) + sign = float(0.0) + p = wp.mesh_query_point_sign_winding_number(mesh, pos, max_length, sign, face_index, face_u, face_v) + + # set point to be solid + if sign < 0.0: + rho[0, i, j, k] = mesh_density + else: + rho[0, i, j, k] = background_density + + def __call__( + self, + rho, + origin, + spacing, + ): + # Voxelize STL of mesh + wp.launch( + self._initialize_target_density, + inputs=[ + rho, + wp.uint64(self.mesh.id), + self.background_density, + self.mesh_density, + wp.vec3f(origin), + wp.vec3f(spacing), + ], + dim=rho.shape[1:], + ) + + return rho diff --git a/examples/out_of_core/operators/l2_loss.py b/examples/out_of_core/operators/l2_loss.py new file mode 100644 index 00000000..d26b5f13 --- /dev/null +++ b/examples/out_of_core/operators/l2_loss.py @@ -0,0 +1,33 @@ +from typing import Any +import warp as wp + + +class L2Loss: + @wp.kernel + def _l2_loss( + rho: wp.array4d(dtype=Any), + target_rho: wp.array4d(dtype=Any), + boundary_id: wp.array4d(dtype=wp.uint8), + l2_loss: wp.array(dtype=wp.float32), + ): + # Get the global index + i, j, k = wp.tid() + + # Compute the loss + if boundary_id[0, i, j, k] == wp.uint8(0): + wp.atomic_add(l2_loss, 0, (rho[0, i, j, k] - target_rho[0, i, j, k]) ** 2.0) + + def __call__(self, rho, target_rho, boundary_id, l2_loss): + # Launch the warp kernel + wp.launch( + self._l2_loss, + inputs=[ + rho, + target_rho, + boundary_id, + ], + outputs=[l2_loss], + dim=rho.shape[1:], + ) + + return l2_loss diff --git a/examples/out_of_core/operators/mesh_renderer.py b/examples/out_of_core/operators/mesh_renderer.py new file mode 100644 index 00000000..bcdfc177 --- /dev/null +++ b/examples/out_of_core/operators/mesh_renderer.py @@ -0,0 +1,245 @@ +import warp as wp +import math + + +class MeshRenderer: + """ + Operator for rendering a Warp mesh to a pixel and depth buffer using ray tracing. + + This operator takes a Warp mesh and renders it using ray tracing with Blinn-Phong + shading, including diffuse, specular, and fresnel effects. + + Parameters + ---------- + width : int + Width of the output image in pixels + height : int + Height of the output image in pixels + camera_position : wp.vec3 + Position of the camera in world space + + Attributes + ---------- + Buffer Layout + ------------ + pixel_buffer : wp.array3d(dtype=wp.float32) + Shape: (height, width, 4) + RGBA color buffer + depth_buffer : wp.array2d(dtype=wp.float32) + Shape: (height, width) + Depth values in range [0,1] + """ + + @staticmethod + @wp.func + def create_view_matrix(eye: wp.vec3, target: wp.vec3, up: wp.vec3) -> wp.mat44: + """Create a view matrix from camera parameters using right-handed coordinate system.""" + # Forward vector points from eye to target (negative z-axis in view space) + forward = wp.normalize(target - eye) # Note: reversed from before + + # Right vector + right = wp.normalize(wp.cross(forward, up)) + + # Recompute up vector to ensure orthogonality + up = wp.normalize(wp.cross(right, forward)) + + # Construct view matrix - note forward is negated to maintain right-handed system + return wp.mat44( + right[0], up[0], -forward[0], eye[0], right[1], up[1], -forward[1], eye[1], right[2], up[2], -forward[2], eye[2], 0.0, 0.0, 0.0, 1.0 + ) + + @staticmethod + @wp.func + def normal_based_shading( + normal: wp.vec3, + view_dir: wp.vec3, + base_color: wp.vec3, + ambient_intensity: float, + edge_sharpness: float, + ) -> wp.vec3: + """Compute lighting with simple normal-based shading.""" + # Normalize vectors + n = wp.normalize(normal) + v = wp.normalize(view_dir) + + # Check if normal is facing away from view direction + n_dot_v = wp.dot(n, v) + if n_dot_v < 0.0: + # Flip normal if it's facing away + n = wp.vec3(-n[0], -n[1], -n[2]) + n_dot_v = -n_dot_v + + # Use configurable falloff for edge definition + diffuse_factor = wp.pow(n_dot_v, edge_sharpness) + + # Add ambient light to prevent completely black areas + light = wp.vec3(ambient_intensity + diffuse_factor) + + # Apply lighting to base color using component-wise multiplication + return wp.cw_mul(base_color, light) + + @wp.kernel + def _render_mesh( + mesh_id: wp.uint64, + vertex_colors: wp.array2d(dtype=wp.float32), + pixel_buffer: wp.array3d(dtype=wp.float32), + depth_buffer: wp.array2d(dtype=wp.float32), + camera_pos: wp.vec3f, + camera_target: wp.vec3f, + camera_up: wp.vec3f, + fov_degrees: float, + ambient_intensity: float, + edge_sharpness: float, + gamma: float, + ): + # Get pixel coordinates + i, j = wp.tid() + height = pixel_buffer.shape[0] + width = pixel_buffer.shape[1] + + # Get mesh + mesh = wp.mesh_get(mesh_id) + + # Convert FOV to radians and calculate image plane parameters + aspect = float(width) / float(height) + fov = math.radians(fov_degrees) + tan_fov = math.tan(fov * 0.5) + + # Convert to NDC space with proper FOV + sx = (2.0 * float(j) / float(width) - 1.0) * aspect * tan_fov + sy = (1.0 - 2.0 * float(i) / float(height)) * tan_fov + + # Create view matrix + view = MeshRenderer.create_view_matrix(camera_pos, camera_target, camera_up) + + # Create ray in camera space + ray_dir = wp.normalize(wp.vec3(sx, sy, -1.0)) + + # Transform ray to world space + ro = camera_pos + rd = wp.vec3( + ray_dir[0] * view[0, 0] + ray_dir[1] * view[0, 1] + ray_dir[2] * view[0, 2], + ray_dir[0] * view[1, 0] + ray_dir[1] * view[1, 1] + ray_dir[2] * view[1, 2], + ray_dir[0] * view[2, 0] + ray_dir[1] * view[2, 1] + ray_dir[2] * view[2, 2], + ) + rd = wp.normalize(rd) + + # Ray trace against mesh + query = wp.mesh_query_ray(mesh_id, ro, rd, depth_buffer[i, j]) + if query.result: + if query.t < depth_buffer[i, j]: + # Use normal-based coloring for debugging + normal = wp.normalize(query.normal) + + # Get indices of the face + i0 = mesh.indices[3 * query.face + 0] # First vertex + i1 = mesh.indices[3 * query.face + 1] # Second vertex + i2 = mesh.indices[3 * query.face + 2] # Third vertex + + # Get vertex colors + c0 = wp.vec3(vertex_colors[i0, 0], vertex_colors[i0, 1], vertex_colors[i0, 2]) + c1 = wp.vec3(vertex_colors[i1, 0], vertex_colors[i1, 1], vertex_colors[i1, 2]) + c2 = wp.vec3(vertex_colors[i2, 0], vertex_colors[i2, 1], vertex_colors[i2, 2]) + + # Use barycentric coordinates from query + w0 = query.u # Weight for first edge (between v0 and v1) + w1 = query.v # Weight for second edge (between v1 and v2) + w2 = 1.0 - query.u - query.v # Weight for remaining vertex + + # Interpolate vertex colors using barycentric coordinates + base_color = ( + wp.cw_mul(c0, wp.vec3(w0)) # First vertex + + wp.cw_mul(c1, wp.vec3(w1)) # Second vertex + + wp.cw_mul(c2, wp.vec3(w2)) # Third vertex + ) + + # Compute lighting + color = MeshRenderer.normal_based_shading( + normal=normal, + view_dir=rd, + base_color=base_color, + ambient_intensity=ambient_intensity, + edge_sharpness=edge_sharpness, + ) + + # Apply gamma correction (linear to sRGB) + color = wp.vec3( + wp.pow(wp.clamp(color[0], 0.0, 1.0), 1.0 / gamma), + wp.pow(wp.clamp(color[1], 0.0, 1.0), 1.0 / gamma), + wp.pow(wp.clamp(color[2], 0.0, 1.0), 1.0 / gamma), + ) + + # Write results + pixel_buffer[i, j, 0] = color[0] + pixel_buffer[i, j, 1] = color[1] + pixel_buffer[i, j, 2] = color[2] + pixel_buffer[i, j, 3] = 1.0 + depth_buffer[i, j] = query.t + + def __call__( + self, + mesh: wp.Mesh, + vertex_colors: wp.array2d, # Shape: (num_vertices, 3) for RGB colors + pixel_buffer: wp.array3d, + depth_buffer: wp.array2d, + camera_pos: wp.vec3f = wp.vec3f(0.0, 1.0, 2.0), + camera_target: wp.vec3f = wp.vec3f(0.0, 0.0, 0.0), + camera_up: wp.vec3f = wp.vec3f(0.0, 1.0, 0.0), + fov_degrees: float = 60.0, + ambient_intensity: float = 0.05, + edge_sharpness: float = 1.0, + gamma: float = 1.0, + ): + """ + Render a Warp mesh with normal-based shading. + + Parameters + ---------- + mesh : wp.Mesh + Warp mesh to render + vertex_colors : wp.array2d + Vertex colors (num_vertices, 3) for RGB colors + pixel_buffer : wp.array3d + Output pixel buffer (height, width, 4) RGBA + depth_buffer : wp.array2d + Output depth buffer (height, width) + camera_pos : wp.vec3f + Camera position in world space + camera_target : wp.vec3f + Point the camera is looking at + camera_up : wp.vec3f + Camera up vector + fov_degrees : float + Field of view in degrees + ambient_intensity : float + Intensity of ambient light (0.0-1.0) + edge_sharpness : float + Controls edge definition (lower values = softer edges, higher values = sharper edges) + gamma : float + Gamma correction value (typically 1.0-2.2, lower = brighter) + + Returns + ------- + tuple[wp.array3d, wp.array2d] + Updated pixel and depth buffers + """ + # Launch kernel + wp.launch( + self._render_mesh, + dim=(pixel_buffer.shape[0], pixel_buffer.shape[1]), + inputs=[ + mesh.id, + vertex_colors, + pixel_buffer, + depth_buffer, + camera_pos, + camera_target, + camera_up, + fov_degrees, + ambient_intensity, + edge_sharpness, + gamma, + ], + ) + + return pixel_buffer, depth_buffer diff --git a/examples/out_of_core/operators/q_criterion.py b/examples/out_of_core/operators/q_criterion.py new file mode 100644 index 00000000..65b62587 --- /dev/null +++ b/examples/out_of_core/operators/q_criterion.py @@ -0,0 +1,113 @@ +from typing import Any +import warp as wp + + +class QCriterion: + """ + Operator for computing the Q-criterion and vorticity magnitude for vortex identification. + """ + + @wp.kernel + def q_kernel( + u: wp.array4d(dtype=Any), + boundary_id: wp.array4d(dtype=wp.uint8), + norm_mu: wp.array4d(dtype=Any), + q: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + + # Add ghost cells to index + i += 1 + j += 1 + k += 1 + + # Check if anything on edges + b_id_2_1_1 = boundary_id[0, i + 1, j, k] + b_id_1_2_1 = boundary_id[0, i, j + 1, k] + b_id_1_1_2 = boundary_id[0, i, j, k + 1] + b_id_0_1_1 = boundary_id[0, i - 1, j, k] + b_id_1_0_1 = boundary_id[0, i, j - 1, k] + b_id_1_1_0 = boundary_id[0, i, j, k - 1] + if ( + b_id_2_1_1 != wp.uint8(0) + or b_id_1_2_1 != wp.uint8(0) + or b_id_1_1_2 != wp.uint8(0) + or b_id_0_1_1 != wp.uint8(0) + or b_id_1_0_1 != wp.uint8(0) + or b_id_1_1_0 != wp.uint8(0) + ): + return + + # Get derivatives + u_x_dx = (u[0, i + 1, j, k] - u[0, i - 1, j, k]) / 2.0 + u_x_dy = (u[0, i, j + 1, k] - u[0, i, j - 1, k]) / 2.0 + u_x_dz = (u[0, i, j, k + 1] - u[0, i, j, k - 1]) / 2.0 + u_y_dx = (u[1, i + 1, j, k] - u[1, i - 1, j, k]) / 2.0 + u_y_dy = (u[1, i, j + 1, k] - u[1, i, j - 1, k]) / 2.0 + u_y_dz = (u[1, i, j, k + 1] - u[1, i, j, k - 1]) / 2.0 + u_z_dx = (u[2, i + 1, j, k] - u[2, i - 1, j, k]) / 2.0 + u_z_dy = (u[2, i, j + 1, k] - u[2, i, j - 1, k]) / 2.0 + u_z_dz = (u[2, i, j, k + 1] - u[2, i, j, k - 1]) / 2.0 + + # Compute vorticity + mu_x = u_z_dy - u_y_dz + mu_y = u_x_dz - u_z_dx + mu_z = u_y_dx - u_x_dy + mu = wp.sqrt(mu_x**2.0 + mu_y**2.0 + mu_z**2.0) + + # Compute strain rate + s_0_0 = u_x_dx + s_0_1 = 0.5 * (u_x_dy + u_y_dx) + s_0_2 = 0.5 * (u_x_dz + u_z_dx) + s_1_0 = s_0_1 + s_1_1 = u_y_dy + s_1_2 = 0.5 * (u_y_dz + u_z_dy) + s_2_0 = s_0_2 + s_2_1 = s_1_2 + s_2_2 = u_z_dz + s_dot_s = s_0_0**2.0 + s_0_1**2.0 + s_0_2**2.0 + s_1_0**2.0 + s_1_1**2.0 + s_1_2**2.0 + s_2_0**2.0 + s_2_1**2.0 + s_2_2**2.0 + + # Compute omega + omega_0_0 = 0.0 + omega_0_1 = 0.5 * (u_x_dy - u_y_dx) + omega_0_2 = 0.5 * (u_x_dz - u_z_dx) + omega_1_0 = -omega_0_1 + omega_1_1 = 0.0 + omega_1_2 = 0.5 * (u_y_dz - u_z_dy) + omega_2_0 = -omega_0_2 + omega_2_1 = -omega_1_2 + omega_2_2 = 0.0 + omega_dot_omega = ( + omega_0_0**2.0 + + omega_0_1**2.0 + + omega_0_2**2.0 + + omega_1_0**2.0 + + omega_1_1**2.0 + + omega_1_2**2.0 + + omega_2_0**2.0 + + omega_2_1**2.0 + + omega_2_2**2.0 + ) + + # Compute q-criterion + q_value = 0.5 * (omega_dot_omega - s_dot_s) + + # Set the output + norm_mu[0, i, j, k] = mu + q[0, i, j, k] = q_value + + def __call__(self, u, boundary_id, norm_mu, q): + # Launch the warp kernel + wp.launch( + self.q_kernel, + inputs=[ + u, + boundary_id, + norm_mu, + q, + ], + dim=[i - 2 for i in u.shape[1:]], + ) + + return norm_mu, q diff --git a/examples/out_of_core/operators/soa_copy.py b/examples/out_of_core/operators/soa_copy.py new file mode 100644 index 00000000..7e4efee6 --- /dev/null +++ b/examples/out_of_core/operators/soa_copy.py @@ -0,0 +1,61 @@ +from typing import Any +import warp as wp + + +class SOACopy: + """ + SOACopy is an operator for copying data from a source array to a destination array. + This is currently just used to speed up the copy as warps copy is not optimized for + non-contiguous arrays. + """ + + @wp.kernel + def soa_copy_3d( + dest: wp.array4d(dtype=Any), + src: wp.array4d(dtype=Any), + q: wp.int32, + ): + """ + Copy data from a 3D source array to a 3D destination array. + + Parameters + ---------- + dest : wp.array4d + The destination array where data will be copied to. + src : wp.array4d + The source array from which data will be copied. + q : wp.int32 + The number of elements to copy along the first dimension. + """ + # Get the global index + i, j, k = wp.tid() + + # Copy the data + for ii in range(q): + dest[ii, i, j, k] = src[ii, i, j, k] + + def __call__( + self, + dest: wp.array, + src: wp.array, + ): + """ + Parameters + ---------- + dest : wp.array + The destination array where data will be copied to. + src : wp.array + The source array from which data will be copied. + + Returns + ------- + wp.array + The destination array with copied data. + """ + # Launch the warp kernel + wp.launch( + self.soa_copy_3d, + inputs=[dest, src, dest.shape[0]], + dim=dest.shape[1:], + ) + return dest diff --git a/examples/out_of_core/operators/transform_mesh.py b/examples/out_of_core/operators/transform_mesh.py new file mode 100644 index 00000000..696d7f68 --- /dev/null +++ b/examples/out_of_core/operators/transform_mesh.py @@ -0,0 +1,73 @@ +import warp as wp + + +class TransformMesh: + """ + Operator for transforming mesh vertices using translation and scaling. + + The transformation is applied in the following order: + 1. Scale around origin + 2. Translate to new origin + """ + + @wp.kernel + def _transform_vertices( + vertices: wp.array(dtype=wp.vec3), + origin: wp.vec3, + scale: wp.vec3, + ): + # Get thread index + idx = wp.tid() + + # Get vertex + vertex = vertices[idx] + + # Apply scale + vertex = wp.vec3(vertex[0] * scale[0], vertex[1] * scale[1], vertex[2] * scale[2]) + + # Apply translation + vertex = vertex + origin + + # Store result + vertices[idx] = vertex + + def __call__( + self, + mesh: wp.Mesh, + origin: wp.vec3, + scale: wp.vec3, + ) -> wp.Mesh: + """ + Transform mesh vertices using translation and scaling. + + Parameters + ---------- + mesh : wp.Mesh + Input mesh to transform + origin : wp.vec3 + New origin for the mesh (translation) + scale : wp.vec3 + Scale factors for each axis + + Returns + ------- + wp.Mesh + New mesh with transformed vertices + """ + # Create new vertices array + new_vertices = wp.clone(mesh.points) + + # Launch kernel to transform vertices + wp.launch( + self._transform_vertices, + dim=new_vertices.shape[0], + inputs=[new_vertices, origin, scale], + ) + + # Create new mesh with transformed vertices + transformed_mesh = wp.Mesh( + points=new_vertices, + indices=mesh.indices, + ) + + return transformed_mesh diff --git a/examples/out_of_core/operators/trilinear_interpolation.py b/examples/out_of_core/operators/trilinear_interpolation.py new file mode 100644 index 00000000..e5dd4d71 --- /dev/null +++ b/examples/out_of_core/operators/trilinear_interpolation.py @@ -0,0 +1,136 @@ +import warp as wp + + +class TrilinearInterpolation: + """ + Operator for trilinear interpolation from a grid to points in space. + + The grid is assumed to be a 4D array with shape (q, nx, ny, nz) where: + - q: number of quantities to interpolate + - nx, ny, nz: grid dimensions in each direction + + Values are assumed to be cell-centered. + """ + + @wp.kernel + def _trilinear_interpolation( + grid: wp.array4d(dtype=float), + points: wp.array(dtype=wp.vec3), + point_values: wp.array2d(dtype=float), + origin: wp.vec3, + spacing: wp.vec3, + ): + # Get the global index + i = wp.tid() + + # Get the point + point = points[i] + + # Convert point to grid coordinates (cell-centered) + x = (point[0] - origin[0]) / spacing[0] - 0.5 + y = (point[1] - origin[1]) / spacing[1] - 0.5 + z = (point[2] - origin[2]) / spacing[2] - 0.5 + + # Clamp to valid range + nx = grid.shape[1] - 1 + ny = grid.shape[2] - 1 + nz = grid.shape[3] - 1 + + x = wp.clamp(x, 0.0, float(nx)) + y = wp.clamp(y, 0.0, float(ny)) + z = wp.clamp(z, 0.0, float(nz)) + + # Get lower and upper bounds + lower_0_0_0 = wp.vec3i(wp.int32(x), wp.int32(y), wp.int32(z)) + + # Ensure we don't exceed grid bounds + lower_0_0_0[0] = wp.min(lower_0_0_0[0], nx - 1) + lower_0_0_0[1] = wp.min(lower_0_0_0[1], ny - 1) + lower_0_0_0[2] = wp.min(lower_0_0_0[2], nz - 1) + + lower_0_0_1 = lower_0_0_0 + wp.vec3i(0, 0, 1) + lower_0_1_0 = lower_0_0_0 + wp.vec3i(0, 1, 0) + lower_0_1_1 = lower_0_0_0 + wp.vec3i(0, 1, 1) + lower_1_0_0 = lower_0_0_0 + wp.vec3i(1, 0, 0) + lower_1_0_1 = lower_0_0_0 + wp.vec3i(1, 0, 1) + lower_1_1_0 = lower_0_0_0 + wp.vec3i(1, 1, 0) + lower_1_1_1 = lower_0_0_0 + wp.vec3i(1, 1, 1) + + # Compute the interpolation weights + dx = x - wp.float32(lower_0_0_0[0]) + dy = y - wp.float32(lower_0_0_0[1]) + dz = z - wp.float32(lower_0_0_0[2]) + w_000 = (1.0 - dx) * (1.0 - dy) * (1.0 - dz) + w_001 = (1.0 - dx) * (1.0 - dy) * dz + w_010 = (1.0 - dx) * dy * (1.0 - dz) + w_011 = (1.0 - dx) * dy * dz + w_100 = dx * (1.0 - dy) * (1.0 - dz) + w_101 = dx * (1.0 - dy) * dz + w_110 = dx * dy * (1.0 - dz) + w_111 = dx * dy * dz + + # Loop over values to interpolate + for n in range(grid.shape[0]): + # Get grid values + grid_0_0_0 = grid[n, lower_0_0_0[0], lower_0_0_0[1], lower_0_0_0[2]] + grid_0_0_1 = grid[n, lower_0_0_1[0], lower_0_0_1[1], lower_0_0_1[2]] + grid_0_1_0 = grid[n, lower_0_1_0[0], lower_0_1_0[1], lower_0_1_0[2]] + grid_0_1_1 = grid[n, lower_0_1_1[0], lower_0_1_1[1], lower_0_1_1[2]] + grid_1_0_0 = grid[n, lower_1_0_0[0], lower_1_0_0[1], lower_1_0_0[2]] + grid_1_0_1 = grid[n, lower_1_0_1[0], lower_1_0_1[1], lower_1_0_1[2]] + grid_1_1_0 = grid[n, lower_1_1_0[0], lower_1_1_0[1], lower_1_1_0[2]] + grid_1_1_1 = grid[n, lower_1_1_1[0], lower_1_1_1[1], lower_1_1_1[2]] + + # Compute the interpolated value + point_value = ( + w_000 * grid_0_0_0 + + w_001 * grid_0_0_1 + + w_010 * grid_0_1_0 + + w_011 * grid_0_1_1 + + w_100 * grid_1_0_0 + + w_101 * grid_1_0_1 + + w_110 * grid_1_1_0 + + w_111 * grid_1_1_1 + ) + + # Set the output + point_values[n, i] = point_value + + def __call__( + self, + grid: wp.array4d, + points: wp.array, + origin: wp.vec3, + spacing: wp.vec3, + point_values: wp.array2d, + ) -> wp.array2d: + """ + Interpolate values from a grid to points in space. + + Parameters + ---------- + grid : wp.array4d(dtype=float) + Input grid with shape (q, nx, ny, nz) + points : wp.array(dtype=wp.vec3) + Points to interpolate to + origin : wp.vec3 + Origin of the grid (lower corner) + spacing : wp.vec3 + Grid spacing in each direction + point_values : wp.array2d(dtype=float) + Output array with shape (q, num_points). + + Returns + ------- + wp.array2d + Interpolated values at each point + """ + + # Launch the kernel + wp.launch( + self._trilinear_interpolation, + dim=points.shape[0], + inputs=[grid, points, point_values, origin, spacing], + ) + + return point_values diff --git a/examples/out_of_core/operators/uniform_initializer.py b/examples/out_of_core/operators/uniform_initializer.py new file mode 100644 index 00000000..60b6824a --- /dev/null +++ b/examples/out_of_core/operators/uniform_initializer.py @@ -0,0 +1,42 @@ +from typing import Any +import warp as wp + + +class UniformInitializer: + # NOTE: This could be overridden to use more complex initializers + + def __init__( + self, + initial_rho, + initial_u, + ): + self.initial_rho = initial_rho + self.initial_u = initial_u + + @wp.kernel + def uniform_initializer_kernel( + rho: wp.array4d(dtype=Any), + u: wp.array4d(dtype=Any), + boundary_id: wp.array4d(dtype=wp.uint8), + initial_u: wp.vec3f, + initial_rho: float, + ): + # Get the global index + i, j, k = wp.tid() + + # Set the velocity + u[0, i, j, k] = initial_u[0] + u[1, i, j, k] = initial_u[1] + u[2, i, j, k] = initial_u[2] + + # Set the density + rho[0, i, j, k] = initial_rho + + def __call__(self, rho, u, boundary_id): + # Launch the warp kernel + wp.launch( + self.uniform_initializer_kernel, + inputs=[rho, u, boundary_id, wp.vec3f(self.initial_u), self.initial_rho], + dim=rho.shape[1:], + ) + return rho, u diff --git a/examples/out_of_core/subroutine/__init__.py b/examples/out_of_core/subroutine/__init__.py new file mode 100644 index 00000000..1d9ad3a6 --- /dev/null +++ b/examples/out_of_core/subroutine/__init__.py @@ -0,0 +1,23 @@ +from subroutine.subroutine import Subroutine +from subroutine.prepare_fields import PrepareFieldsSubroutine +from subroutine.stepper_subroutine import StepperSubroutine +from subroutine.render_q_criterion import RenderQCriterionSubroutine +from subroutine.volume_saver_subroutine import VolumeSaverSubroutine +from subroutine.autodiff_stepper_subroutine import ForwardStepperSubroutine, BackwardStepperSubroutine +from subroutine.rho_loss_subroutine import ForwardRhoLossSubroutine, BackwardRhoLossSubroutine +from subroutine.gradient_descent import GradientDescentSubroutine +from subroutine.initialize_field import InitializeFieldSubroutine + +__all__ = [ + Subroutine, + PrepareFieldsSubroutine, + StepperSubroutine, + RenderQCriterionSubroutine, + VolumeSaverSubroutine, + ForwardStepperSubroutine, + BackwardStepperSubroutine, + ForwardRhoLossSubroutine, + BackwardRhoLossSubroutine, + GradientDescentSubroutine, + InitializeFieldSubroutine, +] diff --git a/examples/out_of_core/subroutine/autodiff_stepper_subroutine.py b/examples/out_of_core/subroutine/autodiff_stepper_subroutine.py new file mode 100644 index 00000000..b2808341 --- /dev/null +++ b/examples/out_of_core/subroutine/autodiff_stepper_subroutine.py @@ -0,0 +1,555 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.soa_copy import SOACopy + + +class ForwardStepperSubroutine(Subroutine): + def __init__( + self, + stepper: Callable, + omega: float, + my_copy: Callable = SOACopy(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.stepper = stepper + self.omega = omega + self.my_copy = my_copy + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + nr_steps=None, + f_input_name="f_0000", + f_output_name="f_0001", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + clear_memory_pools=True, + ): + # Get number of steps + if nr_steps is None: + nr_steps = min(ooc_grid.ghost_cell_thickness) // 2 + assert nr_steps <= min(ooc_grid.ghost_cell_thickness) // 2 + + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get block cardinality + q = block.boxes[f_input_name].cardinality + + # Get total box offset, extent and shape + offset = block.offset_with_ghost + extent = block.extent_with_ghost + shape = extent + + # Get compute arrays + f0 = self.memory_pools[stream_idx].get((q, *shape), wp.float32) + f1 = self.memory_pools[stream_idx].get((q, *shape), wp.float32) + boundary_id = self.memory_pools[stream_idx].get((1, *shape), wp.uint8) + missing_mask = self.memory_pools[stream_idx].get((q, *shape), wp.bool) + + # Fill boundary id with -1 + boundary_id.fill_(wp.uint8(-1)) + + # Get transmit arrays + f_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + boundary_id_block = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + missing_mask_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.bool) + f_neighbour_ghost = {} + boundary_id_neighbour_ghost = {} + missing_mask_neighbour_ghost = {} + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + f_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[f_input_name].shape), wp.float32) + boundary_id_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get( + (1, *ghost_boxes[boundary_id_name].shape), wp.uint8 + ) + missing_mask_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get( + (q, *ghost_boxes[missing_mask_name].shape), wp.bool + ) + f_local_ghost = {} + boundary_id_local_ghost = {} + missing_mask_local_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + f_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[f_input_name].shape), wp.float32) + boundary_id_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((1, *ghost_boxes[boundary_id_name].shape), wp.uint8) + missing_mask_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[missing_mask_name].shape), wp.bool) + + # Copy from block + wp.copy(f_block, block.boxes[f_input_name].data) + wp.copy(boundary_id_block, block.boxes[boundary_id_name].data) + wp.copy(missing_mask_block, block.boxes[missing_mask_name].data) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + wp.copy(f_neighbour_ghost[ghost_block], ghost_boxes[f_input_name].data) + wp.copy(boundary_id_neighbour_ghost[ghost_block], ghost_boxes[boundary_id_name].data) + wp.copy(missing_mask_neighbour_ghost[ghost_block], ghost_boxes[missing_mask_name].data) + + # Wait for previous event + if event is not None: + self.wp_streams[stream_idx].wait_event(event) + + # Copy to compute arrays + start_1 = int(block.offset[0] - offset[0]) + stop_1 = start_1 + block.extent[0] + start_2 = int(block.offset[1] - offset[1]) + stop_2 = start_2 + block.extent[1] + start_3 = int(block.offset[2] - offset[2]) + stop_3 = start_3 + block.extent[2] + self.my_copy( + f0[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + f_block, + ) + self.my_copy( + boundary_id[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + boundary_id_block, + ) + self.my_copy( + missing_mask[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + missing_mask_block, + ) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + start_1 = int(ghost_boxes[f_input_name].offset[0] - offset[0]) + stop_1 = start_1 + ghost_boxes[f_input_name].extent[0] + start_2 = int(ghost_boxes[f_input_name].offset[1] - offset[1]) + stop_2 = start_2 + ghost_boxes[f_input_name].extent[1] + start_3 = int(ghost_boxes[f_input_name].offset[2] - offset[2]) + stop_3 = start_3 + ghost_boxes[f_input_name].extent[2] + self.my_copy( + f0[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + f_neighbour_ghost[ghost_block], + ) + self.my_copy( + boundary_id[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + boundary_id_neighbour_ghost[ghost_block], + ) + self.my_copy( + missing_mask[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + missing_mask_neighbour_ghost[ghost_block], + ) + + # Perform update + for _ in range(nr_steps): + # Perform steppej + f0, f1 = self.stepper(f0, f1, boundary_id, missing_mask, self.omega, 0) + f0, f1 = f1, f0 + + # Copy from compute arrays + start_1 = int(block.offset[0] - offset[0]) + stop_1 = start_1 + block.extent[0] + start_2 = int(block.offset[1] - offset[1]) + stop_2 = start_2 + block.extent[1] + start_3 = int(block.offset[2] - offset[2]) + stop_3 = start_3 + block.extent[2] + self.my_copy( + f_block, + f0[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + ) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[f_output_name].offset - offset + slice_stop = slice_start + ghost_boxes[f_output_name].extent + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy f + self.my_copy( + f_local_ghost[ghost_block], + f0[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Wait for previous event + if event is None: + event = wp.Event() + self.wp_streams[stream_idx].record_event(event) + + # Copy to block + wp.copy(block.boxes[f_output_name].data, f_block) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[f_output_name].data, f_local_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(f0, zero=False) + self.memory_pools[stream_idx].ret(f1, zero=False) + self.memory_pools[stream_idx].ret(boundary_id, zero=False) + self.memory_pools[stream_idx].ret(missing_mask, zero=False) + self.memory_pools[stream_idx].ret(f_block, zero=False) + self.memory_pools[stream_idx].ret(boundary_id_block, zero=False) + self.memory_pools[stream_idx].ret(missing_mask_block, zero=False) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_neighbour_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(boundary_id_neighbour_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(missing_mask_neighbour_ghost[ghost_block], zero=False) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_local_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(boundary_id_local_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(missing_mask_local_ghost[ghost_block], zero=False) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=[f_output_name], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers(names=[f_output_name]) + + +class BackwardStepperSubroutine(Subroutine): + def __init__( + self, + stepper: Callable, + omega: float, + my_copy: Callable = SOACopy(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.stepper = stepper + self.omega = omega + self.my_copy = my_copy + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + nr_steps=None, + f_input_name="f_0000", + adj_f_name="adj_f", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + clear_memory_pools=True, + ): + # Get number of steps + if nr_steps is None: + nr_steps = min(ooc_grid.ghost_cell_thickness) // 2 + assert nr_steps <= min(ooc_grid.ghost_cell_thickness) // 2 + + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get block cardinality + q = block.boxes[f_input_name].cardinality + + # Get total box offset, extent and shape + offset = block.offset_with_ghost + extent = block.extent_with_ghost + shape = extent + + # Get compute arrays + fs = [self.memory_pools[stream_idx].get((q, *shape), wp.float32, requires_grad=True) for _ in range(nr_steps + 1)] + boundary_id = self.memory_pools[stream_idx].get((1, *shape), wp.uint8) + missing_mask = self.memory_pools[stream_idx].get((q, *shape), wp.bool) + + # Get transmit arrays + f_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + adj_f_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + boundary_id_block = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + missing_mask_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.bool) + f_neighbour_ghost = {} + adj_f_neighbour_ghost = {} + boundary_id_neighbour_ghost = {} + missing_mask_neighbour_ghost = {} + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + f_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[f_input_name].shape), wp.float32) + adj_f_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[adj_f_name].shape), wp.float32) + boundary_id_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get( + (1, *ghost_boxes[boundary_id_name].shape), wp.uint8 + ) + missing_mask_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get( + (q, *ghost_boxes[missing_mask_name].shape), wp.bool + ) + adj_f_local_ghost = {} + boundary_id_local_ghost = {} + missing_mask_local_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + adj_f_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[adj_f_name].shape), wp.float32) + boundary_id_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((1, *ghost_boxes[boundary_id_name].shape), wp.uint8) + missing_mask_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[missing_mask_name].shape), wp.bool) + + # Copy from block + wp.copy(f_block, block.boxes[f_input_name].data) + wp.copy(adj_f_block, block.boxes[adj_f_name].data) + wp.copy(boundary_id_block, block.boxes[boundary_id_name].data) + wp.copy(missing_mask_block, block.boxes[missing_mask_name].data) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + wp.copy(f_neighbour_ghost[ghost_block], ghost_boxes[f_input_name].data) + wp.copy(adj_f_neighbour_ghost[ghost_block], ghost_boxes[adj_f_name].data) + wp.copy(boundary_id_neighbour_ghost[ghost_block], ghost_boxes[boundary_id_name].data) + wp.copy(missing_mask_neighbour_ghost[ghost_block], ghost_boxes[missing_mask_name].data) + + # Wait for previous event + if event is not None: + self.wp_streams[stream_idx].wait_event(event) + + # Copy to compute arrays + start_1 = int(block.offset[0] - offset[0]) + stop_1 = start_1 + block.extent[0] + start_2 = int(block.offset[1] - offset[1]) + stop_2 = start_2 + block.extent[1] + start_3 = int(block.offset[2] - offset[2]) + stop_3 = start_3 + block.extent[2] + self.my_copy( + fs[0][ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + f_block, + ) + self.my_copy( + fs[-1].grad[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + adj_f_block, + ) + self.my_copy( + boundary_id[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + boundary_id_block, + ) + self.my_copy( + missing_mask[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + missing_mask_block, + ) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + start_1 = int(ghost_boxes[f_input_name].offset[0] - offset[0]) + stop_1 = start_1 + ghost_boxes[f_input_name].extent[0] + start_2 = int(ghost_boxes[f_input_name].offset[1] - offset[1]) + stop_2 = start_2 + ghost_boxes[f_input_name].extent[1] + start_3 = int(ghost_boxes[f_input_name].offset[2] - offset[2]) + stop_3 = start_3 + ghost_boxes[f_input_name].extent[2] + self.my_copy( + fs[0][ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + f_neighbour_ghost[ghost_block], + ) + self.my_copy( + fs[-1].grad[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + adj_f_neighbour_ghost[ghost_block], + ) + self.my_copy( + boundary_id[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + boundary_id_neighbour_ghost[ghost_block], + ) + self.my_copy( + missing_mask[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + missing_mask_neighbour_ghost[ghost_block], + ) + + # Perform update + with wp.Tape() as tape: + for _ in range(nr_steps): + # Perform stepper + fs[_ + 1], fs[_] = self.stepper(fs[_], fs[_ + 1], boundary_id, missing_mask, self.omega, 0) + + # Compute gradients + tape.backward() + + # Copy from compute arrays + start_1 = int(block.offset[0] - offset[0]) + stop_1 = start_1 + block.extent[0] + start_2 = int(block.offset[1] - offset[1]) + stop_2 = start_2 + block.extent[1] + start_3 = int(block.offset[2] - offset[2]) + stop_3 = start_3 + block.extent[2] + self.my_copy( + adj_f_block, + fs[0].grad[ + :, + start_1:stop_1, + start_2:stop_2, + start_3:stop_3, + ], + ) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[adj_f_name].offset - offset + slice_stop = slice_start + ghost_boxes[adj_f_name].extent + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy f + self.my_copy( + adj_f_local_ghost[ghost_block], + fs[0].grad[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Wait for previous event + if event is None: + event = wp.Event() + self.wp_streams[stream_idx].record_event(event) + + # Copy to block + wp.copy(block.boxes[adj_f_name].data, adj_f_block) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[adj_f_name].data, adj_f_local_ghost[ghost_block]) + + # Return arrays + for f in fs: + self.memory_pools[stream_idx].ret(f, zero=True) + self.memory_pools[stream_idx].ret(boundary_id, zero=True) + self.memory_pools[stream_idx].ret(missing_mask, zero=True) + self.memory_pools[stream_idx].ret(f_block, zero=True) + self.memory_pools[stream_idx].ret(adj_f_block, zero=True) + self.memory_pools[stream_idx].ret(boundary_id_block, zero=True) + self.memory_pools[stream_idx].ret(missing_mask_block, zero=True) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_neighbour_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(adj_f_neighbour_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(boundary_id_neighbour_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(missing_mask_neighbour_ghost[ghost_block], zero=True) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(adj_f_local_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(boundary_id_local_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(missing_mask_local_ghost[ghost_block], zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=[adj_f_name], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers(names=[adj_f_name]) diff --git a/examples/out_of_core/subroutine/gradient_descent.py b/examples/out_of_core/subroutine/gradient_descent.py new file mode 100644 index 00000000..f529dea4 --- /dev/null +++ b/examples/out_of_core/subroutine/gradient_descent.py @@ -0,0 +1,133 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.gradient_descent import GradientDescent + + +class GradientDescentSubroutine(Subroutine): + def __init__( + self, + gradient_descent: Callable = GradientDescent(), + clamp_field: Callable = None, + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.gradient_descent = gradient_descent + self.clamp_field = clamp_field + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + field_name, + adj_field_name, + learning_rate=1e-3, + min_val=None, + max_val=None, + clear_memory_pools=True, + ): + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get block cardinality + q = block.boxes[field_name].cardinality + + # Get compute arrays + field = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + adj_field = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + field_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + field_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[field_name].shape), wp.float32) + + # Copy from block + wp.copy(field, block.boxes[field_name].data) + wp.copy(adj_field, block.boxes[adj_field_name].data) + + # Perform gradient decent + field = self.gradient_descent(field, adj_field, learning_rate) + + # Clamp field + if self.clamp_field is not None: + field = self.clamp_field(field, min_val, max_val) + + # Copy to local ghost boxes + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[field_name].offset - block.offset + slice_stop = slice_start + ghost_boxes[field_name].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy + wp.copy( + field_ghost[ghost_block], + field[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Copy to block + wp.copy(block.boxes[field_name].data, field) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[field_name].data, field_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(field, zero=True) + self.memory_pools[stream_idx].ret(adj_field, zero=True) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(field_ghost[ghost_block], zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + comm_tag = 0 + requests = [] + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=[field_name], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers( + names=[field_name], + ) + + # Clear memory pools + if clear_memory_pools: + for memory_pool in self.memory_pools: + memory_pool.clear() diff --git a/examples/out_of_core/subroutine/initialize_field.py b/examples/out_of_core/subroutine/initialize_field.py new file mode 100644 index 00000000..e7413973 --- /dev/null +++ b/examples/out_of_core/subroutine/initialize_field.py @@ -0,0 +1,117 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.soa_copy import SOACopy + + +class InitializeFieldSubroutine(Subroutine): + def __init__( + self, + initializer: Callable, + my_copy: Callable = SOACopy(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.initializer = initializer + self.my_copy = my_copy + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + field_name="field", + clear_memory_pools=True, + ): + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Set initial conditions + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get compute arrays + field = self.memory_pools[stream_idx].get(block.boxes[field_name].data_shape, block.boxes[field_name].dtype) + field_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + field_ghost[ghost_block] = self.memory_pools[stream_idx].get( + ghost_boxes[field_name].data_shape, ghost_boxes[field_name].dtype + ) + + # Initialize the field + field = self.initializer( + field, + block.local_origin, + block.local_spacing, + ) + + # Copy to local ghost boxes + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[field_name].offset - block.boxes[field_name].offset + slice_stop = slice_start + ghost_boxes[field_name].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy + self.my_copy( + field_ghost[ghost_block], + field[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Copy to block + wp.copy(block.boxes[field_name].data, field) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[field_name].data, field_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(field, zero=True) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(field_ghost[ghost_block], zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=[field_name], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers( + names=[field_name], + ) + + # Clear memory pools + if clear_memory_pools: + for memory_pool in self.memory_pools: + memory_pool.clear() diff --git a/examples/out_of_core/subroutine/prepare_fields.py b/examples/out_of_core/subroutine/prepare_fields.py new file mode 100644 index 00000000..f4543162 --- /dev/null +++ b/examples/out_of_core/subroutine/prepare_fields.py @@ -0,0 +1,210 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.soa_copy import SOACopy + + +class PrepareFieldsSubroutine(Subroutine): + def __init__( + self, + initializer: Callable, + equilibrium: Callable, + boundary_conditions: List[Callable], + indices_boundary_masker: Callable, + my_copy: Callable = SOACopy(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.initializer = initializer + self.equilibrium = equilibrium + self.boundary_conditions = boundary_conditions + self.indices_boundary_masker = indices_boundary_masker + self.my_copy = my_copy + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + f_name="f", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + clear_memory_pools=True, + ): + # Make stream idx + stream_idx = 0 + + # Set initial conditions + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get q value + q = block.boxes[f_name].cardinality + + # Get total box offset, extent and shape + offset = block.offset_with_ghost + extent = block.extent_with_ghost + + # Get compute arrays + rho = self.memory_pools[stream_idx].get((1, *extent), wp.float32) + u = self.memory_pools[stream_idx].get((3, *extent), wp.float32) + f = self.memory_pools[stream_idx].get((q, *extent), wp.float32) + boundary_id = self.memory_pools[stream_idx].get((1, *extent), wp.uint8) + missing_mask = self.memory_pools[stream_idx].get((q, *extent), wp.bool) + + # Get transmit arrays + f_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + boundary_id_block = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + missing_mask_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.bool) + f_ghost = {} + boundary_id_ghost = {} + missing_mask_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + f_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[f_name].shape), wp.float32) + boundary_id_ghost[ghost_block] = self.memory_pools[stream_idx].get((1, *ghost_boxes[boundary_id_name].shape), wp.uint8) + missing_mask_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[missing_mask_name].shape), wp.bool) + + # Initialize boundary id and missing mask + boundary_id, missing_mask = self.indices_boundary_masker( + self.boundary_conditions, + boundary_id, + missing_mask, + offset, + ) + + # Initialize the flow field + rho, u = self.initializer(rho, u, boundary_id) + f = self.equilibrium(rho, u, f) + + # Copy to block + slice_start = block.offset - offset + slice_stop = slice_start + block.extent + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + self.my_copy( + f_block, + f[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + self.my_copy( + boundary_id_block, + boundary_id[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + self.my_copy( + missing_mask_block, + missing_mask[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Copy to local ghost boxes + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[f_name].offset - offset + slice_stop = slice_start + ghost_boxes[f_name].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy + self.my_copy( + f_ghost[ghost_block], + f[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + self.my_copy( + boundary_id_ghost[ghost_block], + boundary_id[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + self.my_copy( + missing_mask_ghost[ghost_block], + missing_mask[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Copy to block + wp.copy(block.boxes[f_name].data, f_block) + wp.copy(block.boxes[boundary_id_name].data, boundary_id_block) + wp.copy(block.boxes[missing_mask_name].data, missing_mask_block) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[f_name].data, f_ghost[ghost_block]) + wp.copy(ghost_boxes[boundary_id_name].data, boundary_id_ghost[ghost_block]) + wp.copy(ghost_boxes[missing_mask_name].data, missing_mask_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(rho, zero=True) + self.memory_pools[stream_idx].ret(u, zero=True) + self.memory_pools[stream_idx].ret(f, zero=True) + self.memory_pools[stream_idx].ret(boundary_id, zero=True) + self.memory_pools[stream_idx].ret(missing_mask, zero=True) + self.memory_pools[stream_idx].ret(f_block, zero=True) + self.memory_pools[stream_idx].ret(boundary_id_block, zero=True) + self.memory_pools[stream_idx].ret(missing_mask_block, zero=True) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(boundary_id_ghost[ghost_block], zero=True) + self.memory_pools[stream_idx].ret(missing_mask_ghost[ghost_block], zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + comm_tag = 0 + requests = [] + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=[f_name, boundary_id_name, missing_mask_name], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers( + names=[f_name, boundary_id_name, missing_mask_name], + ) + + # Clear memory pools + if clear_memory_pools: + for memory_pool in self.memory_pools: + memory_pool.clear() diff --git a/examples/out_of_core/subroutine/render_q_criterion.py b/examples/out_of_core/subroutine/render_q_criterion.py new file mode 100644 index 00000000..0454813f --- /dev/null +++ b/examples/out_of_core/subroutine/render_q_criterion.py @@ -0,0 +1,226 @@ +from typing import List, Callable +import warp as wp +import numpy as np +from PIL import Image + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.trilinear_interpolation import TrilinearInterpolation +from operators.mesh_renderer import MeshRenderer +from operators.color_mapper import ColorMapper +from operators.transform_mesh import TransformMesh +from operators.q_criterion import QCriterion + + +class RenderQCriterionSubroutine(Subroutine): + def __init__( + self, + macroscopic: Callable, + q_criterion: Callable = QCriterion(), + mesh_renderer: Callable = MeshRenderer(), + color_mapper: Callable = ColorMapper(), + grid_to_point_interpolator: Callable = TrilinearInterpolation(), + mesh_transformer: Callable = TransformMesh(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.macroscopic = macroscopic + self.q_criterion = q_criterion + self.mesh_renderer = mesh_renderer + self.color_mapper = color_mapper + self.grid_to_point_interpolator = grid_to_point_interpolator + self.mesh_transformer = mesh_transformer + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + image_name: str, + pixel_buffer: wp.array(dtype=wp.float32), + depth_buffer: wp.array(dtype=wp.float32), + camera_pos: wp.vec3f, + camera_target: wp.vec3f, + camera_up: wp.vec3f, + fov_degrees: float, + ambient_intensity: float, + edge_sharpness: float, + gamma: float, + q_criterion_threshold: float, + vmin: float, + vmax: float, + boundary_mesh: wp.Mesh = None, + boundary_color: wp.vec3f = None, + f_name="f", + boundary_id_name="boundary_id", + ): + # Make stream idx + stream_idx = 0 + + # Set Perform steps equal to the number of ghost cell thickness + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get block cardinality + q = block.boxes[f_name].cardinality + + # Get compute arrays + f = self.memory_pools[0].get((q, *block.shape), wp.float32) + boundary_id = self.memory_pools[0].get((1, *block.shape), wp.uint8) + rho = self.memory_pools[0].get((1, *block.shape), wp.float32) + u = self.memory_pools[0].get((3, *block.shape), wp.float32) + norm_mu = self.memory_pools[0].get((1, *block.shape), wp.float32) + q = self.memory_pools[0].get((1, *block.shape), wp.float32) + + # Get marching cubes arrays + mc = wp.MarchingCubes( + nx=int(block.extent[0]), + ny=int(block.extent[1]), + nz=int(block.extent[2]), + max_verts=int(block.extent[0]) * int(block.extent[1]) * int(block.extent[2]) * 5, + max_tris=int(block.extent[0]) * int(block.extent[1]) * int(block.extent[2]) * 3, + ) + + # Copy from block + wp.copy(f, block.boxes[f_name].data) + wp.copy(boundary_id, block.boxes[boundary_id_name].data) + + # Compute q criterion + rho, u = self.macroscopic(f, rho, u) + norm_mu, q = self.q_criterion(u, boundary_id, norm_mu, q) + + # Perform marching cubes + mc.surface(q[0], q_criterion_threshold) + + # Check if any vertices found + if mc.verts.shape[0] > 0: + # Make mesh + mesh = wp.Mesh( + points=mc.verts, + indices=mc.indices, + ) + + # Transform mesh + mesh = self.mesh_transformer( + mesh=mesh, + origin=block.local_origin, + scale=block.local_spacing, + ) + + # Get point data + scalars = wp.zeros((1, mc.verts.shape[0]), wp.float32) + scalars = self.grid_to_point_interpolator( + norm_mu, + mesh.points, + origin=block.local_origin, + spacing=block.local_spacing, + point_values=scalars, + ) + + # Map scalars to colors - reshape to 1D array first + vertex_colors = wp.zeros((mc.verts.shape[0], 3), wp.float32) + vertex_colors = self.color_mapper( + scalars[0, :], # Take first channel only since q_criterion is scalar + vertex_colors, + vmin=vmin, + vmax=vmax, + colormap="jet", + ) + + # Render mesh + self.mesh_renderer( + mesh=mesh, + vertex_colors=vertex_colors, + pixel_buffer=pixel_buffer, + depth_buffer=depth_buffer, + camera_pos=camera_pos, + camera_target=camera_target, + camera_up=camera_up, + fov_degrees=fov_degrees, + ambient_intensity=ambient_intensity, + edge_sharpness=edge_sharpness, + gamma=gamma, + ) + + # Return arrays + self.memory_pools[0].ret(f) + self.memory_pools[0].ret(boundary_id) + self.memory_pools[0].ret(rho) + self.memory_pools[0].ret(u) + self.memory_pools[0].ret(norm_mu) + self.memory_pools[0].ret(q) + + # Clear memory pools + for memory_pool in self.memory_pools: + memory_pool.clear() + + # Render boundary mesh + if boundary_mesh is not None: + vertex_colors = wp.full((boundary_mesh.points.shape[0], 3), boundary_color, dtype=wp.float32) + self.mesh_renderer( + mesh=boundary_mesh, + vertex_colors=vertex_colors, + pixel_buffer=pixel_buffer, + depth_buffer=depth_buffer, + camera_pos=camera_pos, + camera_target=camera_target, + camera_up=camera_up, + fov_degrees=fov_degrees, + ambient_intensity=ambient_intensity, + edge_sharpness=edge_sharpness, + gamma=gamma, + ) + + # Get all the files + if ooc_grid.comm is not None: + # Set barrier + ooc_grid.comm.Barrier() + # Send buffers from non-root ranks to root + if ooc_grid.comm.rank != 0: + ooc_grid.comm.Send(pixel_buffer.numpy(), dest=0, tag=2 * ooc_grid.comm.rank) + ooc_grid.comm.Send(depth_buffer.numpy(), dest=0, tag=2 * ooc_grid.comm.rank + 1) + + # Root rank receives and combines buffers + if ooc_grid.comm.rank == 0: + # Initialize with root rank's buffers + np_pixel_buffer = pixel_buffer.numpy() + np_depth_buffer = depth_buffer.numpy() + + # Receive buffers from other ranks + for i in range(1, ooc_grid.comm.size): + # Create receive buffers with same shape as local buffers + other_pixel = np.empty_like(np_pixel_buffer) + other_depth = np.empty_like(np_depth_buffer) + + # Receive pixel and depth buffers + ooc_grid.comm.Recv(other_pixel, source=i, tag=2 * i) + ooc_grid.comm.Recv(other_depth, source=i, tag=2 * i + 1) + + # Update pixels where other depth is smaller + mask = other_depth < np_depth_buffer + np_pixel_buffer[mask] = other_pixel[mask] + np_depth_buffer[mask] = other_depth[mask] + + # Convert float buffer (0-1) to uint8 (0-255) + np_pixel_buffer = (np_pixel_buffer[..., :3] * 255).astype(np.uint8) + + # Ensure correct shape and remove extra dimensions + np_pixel_buffer = np_pixel_buffer.squeeze() + + # Save the combined image + Image.fromarray(np_pixel_buffer).save(f"{image_name}.png") + + else: + np_pixel_buffer = pixel_buffer.numpy()[..., :3] + np_depth_buffer = depth_buffer.numpy() + + # Convert float buffer (0-1) to uint8 (0-255) + np_pixel_buffer = (np_pixel_buffer * 255).astype(np.uint8) + + # Ensure correct shape and remove any extra dimensions + np_pixel_buffer = np_pixel_buffer.squeeze() + + # Save the actual image + Image.fromarray(np_pixel_buffer).save(f"{image_name}.png") diff --git a/examples/out_of_core/subroutine/rho_loss_subroutine.py b/examples/out_of_core/subroutine/rho_loss_subroutine.py new file mode 100644 index 00000000..7dad38ef --- /dev/null +++ b/examples/out_of_core/subroutine/rho_loss_subroutine.py @@ -0,0 +1,219 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.soa_copy import SOACopy + + +class ForwardRhoLossSubroutine(Subroutine): + def __init__( + self, + macroscopic: Callable, + loss: Callable, + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.macroscopic = macroscopic + self.loss = loss + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + amr_grid, + loss, + f_name="f_0000", + boundary_id_name="boundary_id", + target_rho_name="target_rho", + clear_memory_pools=True, + ): + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in amr_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == amr_grid.pid: + # Get block cardinality + q = block.boxes[f_name].cardinality + + # Get compute arrays + rho = self.memory_pools[stream_idx].get((1, *block.shape), wp.float32) + target_rho = self.memory_pools[stream_idx].get((1, *block.shape), wp.float32) + u = self.memory_pools[stream_idx].get((3, *block.shape), wp.float32) + f = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + boundary_id = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + + # Copy from block + wp.copy(f, block.boxes[f_name].data) + wp.copy(target_rho, block.boxes[target_rho_name].data) + wp.copy(boundary_id, block.boxes[boundary_id_name].data) + + # Get rho and u + rho, u = self.macroscopic(f, rho, u) + + # Compute loss + loss = self.loss(rho, target_rho, boundary_id, loss) + + # Return arrays + self.memory_pools[stream_idx].ret(rho, zero=True) + self.memory_pools[stream_idx].ret(u, zero=True) + self.memory_pools[stream_idx].ret(f, zero=True) + self.memory_pools[stream_idx].ret(target_rho, zero=True) + self.memory_pools[stream_idx].ret(boundary_id, zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Synchronize + wp.synchronize() + + # Clear memory pools + if clear_memory_pools: + for memory_pool in self.memory_pools: + memory_pool.clear() + + +class BackwardRhoLossSubroutine(Subroutine): + def __init__( + self, + macroscopic: Callable, + loss: Callable, + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.macroscopic = macroscopic + self.loss = loss + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + amr_grid, + loss, + f_name="f", + adj_f_name="adj_f", + boundary_id_name="boundary_id", + target_rho_name="target_rho", + clear_memory_pools=True, + ): + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in amr_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == amr_grid.pid: + # Get block cardinality + q = block.boxes[f_name].cardinality + + # Get compute arrays + rho = self.memory_pools[stream_idx].get((1, *block.shape), wp.float32, requires_grad=True) + target_rho = self.memory_pools[stream_idx].get((1, *block.shape), wp.float32, requires_grad=True) + u = self.memory_pools[stream_idx].get((3, *block.shape), wp.float32, requires_grad=True) + f = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32, requires_grad=True) + boundary_id = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + adj_f_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + adj_f_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes[f_name].shape), wp.float32) + + # Copy from block + wp.copy(f, block.boxes[f_name].data) + wp.copy(target_rho, block.boxes[target_rho_name].data) + wp.copy(boundary_id, block.boxes[boundary_id_name].data) + + # Make gradient tape + with wp.Tape() as tape: + rho, u = self.macroscopic(f, rho, u) + loss = self.loss(rho, target_rho, boundary_id, loss) + + # Compute gradients + tape.backward() + + # Copy to local ghost boxes + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[f_name].offset - block.offset + slice_stop = slice_start + ghost_boxes[f_name].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy + wp.copy( + adj_f_ghost[ghost_block], + f.grad[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Copy to block + wp.copy(block.boxes[adj_f_name].data, f.grad) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[adj_f_name].data, adj_f_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(rho, zero=True) + self.memory_pools[stream_idx].ret(u, zero=True) + self.memory_pools[stream_idx].ret(f, zero=True) + self.memory_pools[stream_idx].ret(target_rho, zero=True) + self.memory_pools[stream_idx].ret(boundary_id, zero=True) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(adj_f_ghost[ghost_block], zero=True) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + comm_tag = 0 + requests = [] + for block in amr_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + amr_grid.comm, + comm_tag=comm_tag, + names=[adj_f_name], + ) + requests.extend(r) + + # Wait for requests + if amr_grid.comm is not None: + amr_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in amr_grid.blocks.values(): + if block.pid == amr_grid.pid: + block.swap_buffers( + names=[adj_f_name], + ) + + # Clear memory pools + if clear_memory_pools: + for memory_pool in self.memory_pools: + memory_pool.clear() diff --git a/examples/out_of_core/subroutine/stepper_subroutine.py b/examples/out_of_core/subroutine/stepper_subroutine.py new file mode 100644 index 00000000..df135604 --- /dev/null +++ b/examples/out_of_core/subroutine/stepper_subroutine.py @@ -0,0 +1,256 @@ +from typing import List, Callable +from mpi4py import MPI +import warp as wp + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine +from operators.soa_copy import SOACopy + + +class StepperSubroutine(Subroutine): + def __init__( + self, + stepper: Callable, + omega: float, + my_copy: Callable = SOACopy(), + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.stepper = stepper + self.omega = omega + self.my_copy = my_copy + super().__init__(nr_streams, wp_streams, memory_pools) + + def __call__( + self, + ooc_grid, + nr_steps=None, + f_name="f", + boundary_id_name="boundary_id", + missing_mask_name="missing_mask", + clear_memory_pools=True, + ): + # Get number of steps + if nr_steps is None: + nr_steps = min(ooc_grid.ghost_cell_thickness) + assert nr_steps <= min(ooc_grid.ghost_cell_thickness) + + # Make stream idx + stream_idx = 0 + + # MPI communication parameters + comm_tag = 0 + requests = [] + + # Make event + event = None + + # Set Perform steps equal to the number of ghost cell thickness + for block in ooc_grid.blocks.values(): + # Set warp stream + with wp.ScopedStream(self.wp_streams[stream_idx]): + # Check if block matches pid + if block.pid == ooc_grid.pid: + # Get block cardinality + q = block.boxes[f_name].cardinality + + # Get total box offset, extent and shape + offset = block.offset_with_ghost + extent = block.extent_with_ghost + + # Get compute arrays + f0 = self.memory_pools[stream_idx].get((q, *extent), wp.float32) + f1 = self.memory_pools[stream_idx].get((q, *extent), wp.float32) + boundary_id = self.memory_pools[stream_idx].get((1, *extent), wp.uint8) + missing_mask = self.memory_pools[stream_idx].get((q, *extent), wp.bool) + + # Get transmit arrays + f_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.float32) + boundary_id_block = self.memory_pools[stream_idx].get((1, *block.shape), wp.uint8) + missing_mask_block = self.memory_pools[stream_idx].get((q, *block.shape), wp.bool) + f_neighbour_ghost = {} + boundary_id_neighbour_ghost = {} + missing_mask_neighbour_ghost = {} + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + f_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes["f"].shape), wp.float32) + boundary_id_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get((1, *ghost_boxes["boundary_id"].shape), wp.uint8) + missing_mask_neighbour_ghost[ghost_block] = self.memory_pools[stream_idx].get( + (q, *ghost_boxes["missing_mask"].shape), wp.bool + ) + f_local_ghost = {} + boundary_id_local_ghost = {} + missing_mask_local_ghost = {} + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + f_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes["f"].shape), wp.float32) + boundary_id_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((1, *ghost_boxes["boundary_id"].shape), wp.uint8) + missing_mask_local_ghost[ghost_block] = self.memory_pools[stream_idx].get((q, *ghost_boxes["missing_mask"].shape), wp.bool) + + # Copy from block + wp.copy(f_block, block.boxes["f"].data) + wp.copy(boundary_id_block, block.boxes["boundary_id"].data) + wp.copy(missing_mask_block, block.boxes["missing_mask"].data) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + wp.copy(f_neighbour_ghost[ghost_block], ghost_boxes["f"].data) + wp.copy(boundary_id_neighbour_ghost[ghost_block], ghost_boxes["boundary_id"].data) + wp.copy(missing_mask_neighbour_ghost[ghost_block], ghost_boxes["missing_mask"].data) + + # Wait for previous event + if event is not None: + self.wp_streams[stream_idx].wait_event(event) + + # Copy to compute arrays + slice_start = block.offset - offset + slice_stop = slice_start + block.shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + self.my_copy( + f0[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + f_block, + ) + self.my_copy( + boundary_id[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + boundary_id_block, + ) + self.my_copy( + missing_mask[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + missing_mask_block, + ) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + slice_start = ghost_boxes["f"].offset - offset + slice_stop = slice_start + ghost_boxes["f"].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + self.my_copy( + f0[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + f_neighbour_ghost[ghost_block], + ) + self.my_copy( + boundary_id[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + boundary_id_neighbour_ghost[ghost_block], + ) + self.my_copy( + missing_mask[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + missing_mask_neighbour_ghost[ghost_block], + ) + + # Perform update + for _ in range(nr_steps): + # Perform stepper + f0, f1 = self.stepper(f0, f1, boundary_id, missing_mask, self.omega, 0) + f0, f1 = f1, f0 + + # Copy from compute arrays + slice_start = block.offset - offset + slice_stop = slice_start + block.shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + self.my_copy( + f_block, + f0[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + # Get slice start and stop + slice_start = ghost_boxes[f_name].offset - offset + slice_stop = slice_start + ghost_boxes[f_name].shape + slice_start = tuple([int(s) for s in slice_start]) + slice_stop = tuple([int(s) for s in slice_stop]) + + # Copy + self.my_copy( + f_local_ghost[ghost_block], + f0[ + :, + slice_start[0] : slice_stop[0], + slice_start[1] : slice_stop[1], + slice_start[2] : slice_stop[2], + ], + ) + + # Wait for previous event + if event is None: + event = wp.Event() + self.wp_streams[stream_idx].record_event(event) + + # Copy to block + wp.copy(block.boxes["f"].data, f_block) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + wp.copy(ghost_boxes[f_name].data, f_local_ghost[ghost_block]) + + # Return arrays + self.memory_pools[stream_idx].ret(f0, zero=False) + self.memory_pools[stream_idx].ret(f1, zero=False) + self.memory_pools[stream_idx].ret(boundary_id, zero=False) + self.memory_pools[stream_idx].ret(missing_mask, zero=False) + self.memory_pools[stream_idx].ret(f_block, zero=False) + self.memory_pools[stream_idx].ret(boundary_id_block, zero=False) + self.memory_pools[stream_idx].ret(missing_mask_block, zero=False) + for ghost_block, ghost_boxes in block.neighbour_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_neighbour_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(boundary_id_neighbour_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(missing_mask_neighbour_ghost[ghost_block], zero=False) + for ghost_block, ghost_boxes in block.local_ghost_boxes.items(): + self.memory_pools[stream_idx].ret(f_local_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(boundary_id_local_ghost[ghost_block], zero=False) + self.memory_pools[stream_idx].ret(missing_mask_local_ghost[ghost_block], zero=False) + + # Update stream idx + stream_idx = (stream_idx + 1) % self.nr_streams + + # Send blocks + wp.synchronize() + for block in ooc_grid.blocks.values(): + r, comm_tag = block.send_ghost_boxes( + ooc_grid.comm, + comm_tag=comm_tag, + names=["f"], + ) + requests.extend(r) + + # Wait for requests + if ooc_grid.comm is not None: + ooc_grid.comm.Barrier() + MPI.Request.Waitall(requests) + pass + else: + assert len(requests) == 0 + + # Swap neighbour buffers + for block in ooc_grid.blocks.values(): + if block.pid == ooc_grid.pid: + block.swap_buffers(names=["f"]) diff --git a/examples/out_of_core/subroutine/subroutine.py b/examples/out_of_core/subroutine/subroutine.py new file mode 100644 index 00000000..f73f99a5 --- /dev/null +++ b/examples/out_of_core/subroutine/subroutine.py @@ -0,0 +1,21 @@ +# Description: Subroutine class, used to execute complex operations on out of core grids + +from typing import List +import warp as wp + +from ds.ooc_grid import MemoryPool + + +class Subroutine: + def __init__( + self, + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + self.nr_streams = nr_streams + self.wp_streams = wp_streams if wp_streams is not None else [wp.get_stream() for _ in range(nr_streams)] + self.memory_pools = memory_pools if memory_pools is not None else [MemoryPool() for _ in range(nr_streams)] + + def __call__(self, *args): + raise NotImplementedError diff --git a/examples/out_of_core/subroutine/volume_saver_subroutine.py b/examples/out_of_core/subroutine/volume_saver_subroutine.py new file mode 100644 index 00000000..9c85a9f8 --- /dev/null +++ b/examples/out_of_core/subroutine/volume_saver_subroutine.py @@ -0,0 +1,102 @@ +from typing import List +import itertools +import os +import pyvista as pv +import numpy as np +import warp as wp +import xml.etree.ElementTree as ET + +from ds.ooc_grid import MemoryPool +from subroutine.subroutine import Subroutine + + +class VolumeSaverSubroutine(Subroutine): + def __init__( + self, + nr_streams: int = 1, + wp_streams: List[wp.Stream] = None, + memory_pools: List[MemoryPool] = None, + ): + super().__init__(nr_streams, wp_streams, memory_pools) + + @staticmethod + def combine_vtks(files, output_file): + # Create the root element + vtk_file = ET.Element( + "VTKFile", type="vtkMultiBlockDataSet", version="1.0", byte_order="LittleEndian", header_type="UInt32", compressor="vtkZLibDataCompressor" + ) + vtk_multi_block_data_set = ET.SubElement(vtk_file, "vtkMultiBlockDataSet") + + # Create the DataSet elements + for i, file in enumerate(files): + data_set = ET.SubElement(vtk_multi_block_data_set, "DataSet", index=str(i), name=f"Block-{str(i).zfill(5)}", file=file) + + # Create the tree + tree = ET.ElementTree(vtk_file) + + # Write the tree to a file + tree.write(output_file, encoding="utf-8", xml_declaration=True, method="xml") + + def __call__( + self, + ooc_grid, + field_names: List[str], + file_name: str = "initial.vtm", + clear_memory_pools=True, + ): + """ + Save the solid id array. + """ + + # Make directory + os.makedirs(file_name, exist_ok=True) + + # Clear memory pools + for memory_pool in self.memory_pools: + memory_pool.clear() + + # Store the files + files = [] + + # Loop over blocks + for idx, block in enumerate(ooc_grid.blocks.values()): + # Check if block matches pid + if block.pid != ooc_grid.pid: + continue + + # Make grid + grid = pv.ImageData( + dimensions=np.array(block.shape) + 1, + origin=block.local_origin, + spacing=block.local_spacing, + ) + + # Convert data + def _convert_data(data): + np_data = data.numpy() + np_data = np.stack([np_data[i, ...] for i in range(np_data.shape[0])], axis=-1) + return np_data.reshape((-1, np_data.shape[-1]), order="F") + + # Add fields + for field_name in field_names: + np_field = _convert_data(block.boxes[field_name].data) + grid.cell_data[field_name] = np_field + + # Save the grid + post_fix = f"{file_name.split('/')[-1]}" + vtk_file_name = os.path.join(file_name, f"{post_fix}_{idx}.vti") + grid.save(vtk_file_name) + files.append(f"{post_fix}/{post_fix}_{idx}.vti") + + # Get all the files + if ooc_grid.comm is not None: + files = ooc_grid.comm.gather(files, root=0) + if ooc_grid.comm.rank == 0: + files = list(itertools.chain(*files)) + + # Combine the files + if ooc_grid.comm is not None: + if ooc_grid.comm.rank == 0: + self.combine_vtks(files, f"{file_name}.vtm") + else: + self.combine_vtks(files, f"{file_name}.vtm") diff --git a/examples/performance/MLUPS2d.py b/examples/performance/MLUPS2d.py deleted file mode 100644 index 77d32a4b..00000000 --- a/examples/performance/MLUPS2d.py +++ /dev/null @@ -1,67 +0,0 @@ -""" -This script computes the MLUPS (Million Lattice Updates per Second) in 2D by simulating fluid flow inside a 2D cavity. -""" - -import os -import argparse -import jax.numpy as jnp -import numpy as np -from jax import config -from time import time - -from src.utils import * -from src.boundary_conditions import * -from src.lattice import LatticeD2Q9 -from src.models import BGKSim - -class Cavity(BGKSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # concatenate the indices of the left, right, and bottom walls - walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom'])) - # apply bounce back boundary condition to the walls - self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - - # apply inlet equilibrium boundary condition to the top wall - moving_wall = self.boundingBoxIndices['top'] - - rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = u_wall - self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall)) - - -if __name__ == '__main__': - precision = 'f32/f32' - lattice = LatticeD2Q9(precision) - - parser = argparse.ArgumentParser("simple_example") - parser.add_argument("N", help="The total number of voxels will be NxN", type=int) - parser.add_argument("timestep", help="Number of timesteps", type=int) - args = parser.parse_args() - - n = args.N - max_iter = args.timestep - Re = 100.0 - u_wall = 0.1 - clength = n - 1 - - visc = u_wall * clength / Re - omega = 1.0 / (3. * visc + 0.5) - print('omega = ', omega) - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': n, - 'ny': n, - 'nz': 0, - 'precision': precision, - 'compute_MLUPS': True - } - - os.system('rm -rf ./*.vtk && rm -rf ./*.png') - sim = Cavity(**kwargs) - sim.run(max_iter) diff --git a/examples/performance/MLUPS3d.py b/examples/performance/MLUPS3d.py deleted file mode 100644 index 8a9f9e33..00000000 --- a/examples/performance/MLUPS3d.py +++ /dev/null @@ -1,76 +0,0 @@ -""" -This script computes the MLUPS (Million Lattice Updates per Second) in 3D by simulating fluid flow inside a 2D cavity. -""" - -import os -import argparse - -import jax -import jax.numpy as jnp -import numpy as np -from jax import config -from time import time -#config.update('jax_disable_jit', True) -# Use 8 CPU devices -#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -#config.update("jax_enable_x64", True) -from src.utils import * -from src.boundary_conditions import * -from src.models import BGKSim -from src.lattice import LatticeD3Q19 -class Cavity(BGKSim): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # concatenate the indices of the left, right, and bottom walls - walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom'], self.boundingBoxIndices['front'], self.boundingBoxIndices['back'])) - # apply bounce back boundary condition to the walls - self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - - # apply inlet equilibrium boundary condition to the top wall - moving_wall = self.boundingBoxIndices['top'] - - rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = u_wall - self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall)) - -if __name__ == '__main__': - precision = 'f32/f32' - lattice = LatticeD3Q19(precision) - # Create a parser that will read the command line arguments - parser = argparse.ArgumentParser("Calculate MLUPS for a 3D cavity flow simulation") - parser.add_argument("N", help="The total number of voxels all directions. The final dimension will be N*NxN", default=100, type=int) - parser.add_argument("N_ITERS", help="Number of timesteps", default=10000, type=int) - - args = parser.parse_args() - n = args.N - n_iters = args.N_ITERS - - # Store the Reynolds number in the variable Re - Re = 100.0 - # Store the velocity of the lid in the variable u_wall - u_wall = 0.1 - # Store the length of the cavity in the variable clength - clength = n - 1 - - # Compute the viscosity from the Reynolds number, the lid velocity, and the length of the cavity - visc = u_wall * clength / Re - # Compute the relaxation parameter from the viscosity - omega = 1.0 / (3. * visc + 0.5) - - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': n, - 'ny': n, - 'nz': n, - 'precision': precision, - 'compute_MLUPS': True - } - - sim = Cavity(**kwargs) - # Run the simulation - sim.run(n_iters) - \ No newline at end of file diff --git a/examples/performance/MLUPS3d_distributed.py b/examples/performance/MLUPS3d_distributed.py deleted file mode 100644 index 70d53286..00000000 --- a/examples/performance/MLUPS3d_distributed.py +++ /dev/null @@ -1,108 +0,0 @@ -""" -This script computes the MLUPS (Million Lattice Updates per Second) in 3D by simulating fluid flow inside a 2D cavity. -This script is equivalent to MLUPS3d.py, but uses JAX distributed to run the simulation on distributed systems (multi-host, multi-GPUs). -Please refer to https://jax.readthedocs.io/en/latest/multi_process.html for more information on JAX distributed. -""" - - -# Standard Libraries -import argparse -import os -import jax - -import jax.numpy as jnp -import numpy as np - -from jax import config - -from src.boundary_conditions import * -from src.models import BGKSim -from src.lattice import LatticeD3Q19 -from src.utils import * - -#config.update('jax_disable_jit', True) -# Use 8 CPU devices -#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' -#config.update("jax_enable_x64", True) - -class Cavity(BGKSim): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def set_boundary_conditions(self): - # concatenate the indices of the left, right, and bottom walls - walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom'], self.boundingBoxIndices['front'], self.boundingBoxIndices['back'])) - # apply bounce back boundary condition to the walls - self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy)) - - # apply inlet equilibrium boundary condition to the top wall - moving_wall = self.boundingBoxIndices['top'] - - rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype) - vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype) - vel_wall[:, 0] = u_wall - self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall)) - -if __name__ == '__main__': - # Create a parser that will read the command line arguments - parser = argparse.ArgumentParser("Calculate MLUPS for a 3D cavity flow simulation") - parser.add_argument("N", help="The total number of voxels in one direction. The final dimension will be N*NxN", - default=100, type=int) - parser.add_argument("N_ITERS", help="Number of iterations", default=10000, type=int) - parser.add_argument("N_PROCESSES", help="Number of processes. If >1, call jax.distributed.initialize with that number of process. If -1 will call jax.distributed.initialize without any arsgument. So it should pick up the values from SLURM env variable.", - default=1, type=int) - parser.add_argument("IP", help="IP of the master node for multi-node. Useless if using SLURM.", - default='127.0.0.1', type=str, nargs='?') - parser.add_argument("PROCESS_ID_INCREMENT", help="For multi-node only. Useless if using SLURM.", - default=0, type=int, nargs='?') - - args = parser.parse_args() - n = args.N - n_iters = args.N_ITERS - n_processes = args.N_PROCESSES - # Initialize JAX distributed. The IP, number of processes and process id must be set correctly. - print("N processes, ", n_processes) - print("N iter, ", n_iters) - if n_processes > 1: - process_id = int(os.environ.get('CUDA_VISIBLE_DEVICES', 0)) + args.PROCESS_ID_INCREMENT - print("ip, num_processes, process_id, ", args.IP, n_processes, process_id) - jax.distributed.initialize(args.IP, num_processes=n_processes, - process_id=process_id) - elif n_processes == -1: - print("Will call jax.distributed.initialize()") - jax.distributed.initialize() - print("jax.distributed.initialize() ended") - else: - print("No call to jax.distributed.initialize") - print("JAX local devices: ", jax.local_devices()) - - precision = 'f32/f32' - # Create a 3D lattice with the D3Q19 scheme - lattice = LatticeD3Q19(precision) - - # Store the Reynolds number in the variable Re - Re = 100.0 - # Store the velocity of the lid in the variable u_wall - u_wall = 0.1 - # Store the length of the cavity in the variable clength - clength = n - 1 - - # Compute the viscosity from the Reynolds number, the lid velocity, and the length of the cavity - visc = u_wall * clength / Re - # Compute the relaxation parameter from the viscosity - omega = 1.0 / (3. * visc + 0.5) - - # Create a new instance of the Cavity class - kwargs = { - 'lattice': lattice, - 'omega': omega, - 'nx': n, - 'ny': n, - 'nz': n, - 'precision': precision, - 'compute_MLUPS': True - } - - sim = Cavity(**kwargs) # Run the simulation - sim.run(n_iters) diff --git a/examples/performance/mlups_3d.py b/examples/performance/mlups_3d.py new file mode 100644 index 00000000..dbedd2c4 --- /dev/null +++ b/examples/performance/mlups_3d.py @@ -0,0 +1,560 @@ +import xlb +import argparse +import time +import warp as wp +import numpy as np +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import FullwayBounceBackBC, EquilibriumBC +from xlb.distribute import distribute +from xlb.operator.macroscopic import Macroscopic + + +# -------------------------- Simulation Setup -------------------------- + + +def parse_arguments(): + # Define valid options for consistency + COMPUTE_BACKENDS = ["neon", "warp", "jax"] + PRECISION_OPTIONS = ["fp32/fp32", "fp64/fp64", "fp64/fp32", "fp32/fp16"] + VELOCITY_SETS = ["D3Q19", "D3Q27"] + COLLISION_MODELS = ["BGK", "KBC"] + OCC_OPTIONS = ["standard", "none"] + + parser = argparse.ArgumentParser( + description="MLUPS Benchmark for 3D Lattice Boltzmann Method Simulation", + epilog=""" +Examples: + %(prog)s 100 1000 neon fp32/fp32 + %(prog)s 200 500 neon fp64/fp64 --collision_model KBC --velocity_set D3Q27 + %(prog)s 150 2000 neon fp32/fp32 --gpu_devices=[0,1,2] --measure_scalability --report + %(prog)s 100 1000 neon fp32/fp32 --repetitions 5 --export_final_velocity + """, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # Positional arguments + parser.add_argument("cube_edge", type=int, help="Length of the edge of the cubic grid (e.g., 100)") + parser.add_argument("num_steps", type=int, help="Number of timesteps for the simulation (e.g., 1000)") + parser.add_argument("compute_backend", type=str, choices=COMPUTE_BACKENDS, help=f"Backend for the simulation ({', '.join(COMPUTE_BACKENDS)})") + parser.add_argument("precision", type=str, choices=PRECISION_OPTIONS, help=f"Precision for the simulation ({', '.join(PRECISION_OPTIONS)})") + + # Optional arguments + parser.add_argument("--gpu_devices", type=str, default=None, help="CUDA devices to use for Neon backend (e.g., [0,1,2] or [0])") + parser.add_argument( + "--velocity_set", + type=str, + default="D3Q19", + choices=VELOCITY_SETS, + help=f"Lattice velocity set (default: D3Q19, choices: {', '.join(VELOCITY_SETS)})", + ) + parser.add_argument( + "--collision_model", + type=str, + default="BGK", + choices=COLLISION_MODELS, + help=f"Collision model (default: BGK, choices: {', '.join(COLLISION_MODELS)}, KBC requires D3Q27)", + ) + parser.add_argument( + "--occ", + type=str, + default="standard", + choices=OCC_OPTIONS, + help=f"Overlapping Communication and Computation strategy (default: standard, choices: {', '.join(OCC_OPTIONS)})", + ) + parser.add_argument("--report", action="store_true", help="Generate Neon performance report") + parser.add_argument("--export_final_velocity", action="store_true", help="Export final velocity field to VTI file") + parser.add_argument("--measure_scalability", action="store_true", help="Measure performance across different GPU counts") + parser.add_argument( + "--repetitions", type=int, default=1, metavar="N", help="Number of simulation repetitions for statistical analysis (default: 1)" + ) + + args = parser.parse_args() + + # Parse gpu_devices string to list + if args.gpu_devices is not None: + try: + import ast + + args.gpu_devices = ast.literal_eval(args.gpu_devices) + if not isinstance(args.gpu_devices, list): + args.gpu_devices = [args.gpu_devices] # Handle single integer case + except (ValueError, SyntaxError): + raise ValueError("Invalid gpu_devices format. Use format like [0,1,2] or [0]") + + # Validate and convert compute backend + compute_backend_map = { + "jax": ComputeBackend.JAX, + "warp": ComputeBackend.WARP, + "neon": ComputeBackend.NEON, + } + compute_backend = compute_backend_map.get(args.compute_backend) + if compute_backend is None: + raise ValueError(f"Invalid compute backend '{args.compute_backend}'. Use: {', '.join(COMPUTE_BACKENDS)}") + args.compute_backend = compute_backend + + # Handle GPU devices for Neon backend + if args.compute_backend == ComputeBackend.NEON: + if args.gpu_devices is None: + print("[INFO] No GPU devices specified. Using default device 0.") + args.gpu_devices = [0] + + import neon + + occ_enum = neon.SkeletonConfig.OCC.from_string(args.occ) + args.occ_enum = occ_enum # Store the enum for Neon + args.occ_display = args.occ # Store the original string for display + else: + if args.gpu_devices is not None: + raise ValueError(f"--gpu_devices can only be used with Neon backend, not {args.compute_backend.name}") + args.gpu_devices = [0] # Default for non-Neon backends + + # Checking precision policy + precision_policy_map = { + "fp32/fp32": PrecisionPolicy.FP32FP32, + "fp64/fp64": PrecisionPolicy.FP64FP64, + "fp64/fp32": PrecisionPolicy.FP64FP32, + "fp32/fp16": PrecisionPolicy.FP32FP16, + } + precision_policy = precision_policy_map.get(args.precision) + if precision_policy is None: + raise ValueError(f"Invalid precision '{args.precision}'. Use: {', '.join(PRECISION_OPTIONS)}") + args.precision_policy = precision_policy + + # Validate collision model and velocity set compatibility + if args.collision_model == "KBC" and args.velocity_set != "D3Q27": + raise ValueError("KBC collision model requires D3Q27 velocity set. Use --velocity_set D3Q27") + + if args.velocity_set == "D3Q19": + velocity_set = xlb.velocity_set.D3Q19(precision_policy=args.precision_policy, compute_backend=compute_backend) + elif args.velocity_set == "D3Q27": + velocity_set = xlb.velocity_set.D3Q27(precision_policy=args.precision_policy, compute_backend=compute_backend) + args.velocity_set = velocity_set + + print_args(args) + + return args + + +def print_args(args): + """Print simulation configuration in a clean, organized format""" + print("\n" + "=" * 70) + print(" SIMULATION CONFIGURATION") + print("=" * 70) + + # Grid and simulation parameters + print("GRID & SIMULATION:") + print(f" Grid Size: {args.cube_edge}Β³ ({args.cube_edge:,} Γ— {args.cube_edge:,} Γ— {args.cube_edge:,})") + print(f" Total Lattice Points: {args.cube_edge**3:,}") + print(f" Time Steps: {args.num_steps:,}") + print(f" Repetitions: {args.repetitions}") + + # Computational settings + print("\nCOMPUTATIONAL SETTINGS:") + print(f" Compute Backend: {args.compute_backend.name}") + print(f" Precision Policy: {args.precision}") + print(f" Velocity Set: {args.velocity_set.__class__.__name__}") + print(f" Collision Model: {args.collision_model}") + + # Backend-specific settings + if args.compute_backend.name == "NEON": + print("\nNEON BACKEND SETTINGS:") + print(f" GPU Devices: {args.gpu_devices}") + print(f" OCC Strategy: {args.occ_display}") + + # Output options + print("\nOUTPUT OPTIONS:") + print(f" Generate Report: {'Yes' if args.report else 'No'}") + print(f" Measure Scalability: {'Yes' if args.measure_scalability else 'No'}") + print(f" Export Velocity: {'Yes' if args.export_final_velocity else 'No'}") + + print("=" * 70) + print("Starting simulation...\n") + + +def init_xlb(args): + xlb.init( + velocity_set=args.velocity_set, + default_backend=args.compute_backend, + default_precision_policy=args.precision_policy, + ) + options = None + if args.compute_backend == ComputeBackend.NEON: + neon_options = {"occ": args.occ_enum, "device_list": args.gpu_devices} + options = neon_options + return args.compute_backend, args.precision_policy, options + + +def run_simulation( + compute_backend, precision_policy, grid_shape, num_steps, options, export_final_velocity, repetitions, num_devices, collision_model +): + grid = grid_factory(grid_shape, backend_config=options) + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + + lid = box_no_edge["top"] + walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(len(grid.shape))] + walls = np.unique(np.array(walls), axis=-1).tolist() + + boundary_conditions = [ + EquilibriumBC(rho=1.0, u=(0.02, 0.0, 0.0), indices=lid), + FullwayBounceBackBC(indices=walls), + ] + + stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type=collision_model, + backend_config=options, + ) + + # Distribute if using JAX + if compute_backend == ComputeBackend.JAX: + stepper = distribute( + stepper, + grid, + xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend), + ) + + # Initialize fields + omega = 1.0 + f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + + warmup_iterations = 10 + # Warp-up iterations + for i in range(warmup_iterations): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, i) + f_0, f_1 = f_1, f_0 + wp.synchronize() + export_num_steps = warmup_iterations + + elapsed_time_list = [] + for i in range(repetitions): + start_time = time.time() + for i in range(num_steps): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, i) + f_0, f_1 = f_1, f_0 + wp.synchronize() + elapsed_time = time.time() - start_time + elapsed_time_list.append(elapsed_time) + export_num_steps += num_steps + + # Define Macroscopic Calculation + macro = Macroscopic( + compute_backend=compute_backend, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend), + ) + + if compute_backend == ComputeBackend.NEON: + if export_final_velocity: + rho = grid.create_field(cardinality=1, dtype=precision_policy.store_precision) + u = grid.create_field(cardinality=3, dtype=precision_policy.store_precision) + + macro(f_0, rho, u) + wp.synchronize() + u.update_host(0) + wp.synchronize() + u.export_vti(f"mlups_3d_size_{grid_shape[0]}_dev_{num_devices}_step_{export_num_steps}.vti", "u") + + return elapsed_time_list + + +def calculate_mlups(cube_edge, num_steps, elapsed_time): + total_lattice_updates = cube_edge**3 * num_steps + mlups = (total_lattice_updates / elapsed_time) / 1e6 + return mlups + + +def print_summary_with_stats(args, stats): + """Print comprehensive simulation summary with statistics from multiple repetitions""" + total_lattice_points = args.cube_edge**3 + total_lattice_updates = total_lattice_points * args.num_steps + + mean_mlups = stats["mean_mlups"] + std_mlups = stats["std_dev_mlups"] + mean_elapsed_time = stats["mean_elapsed_time"] + std_elapsed_time = stats["std_dev_elapsed_time"] + + print("\n\n\n" + "=" * 70) + print(" SIMULATION SUMMARY") + print("=" * 70) + + # Simulation Parameters + print("SIMULATION PARAMETERS:") + print("-" * 25) + print(f" Grid Size: {args.cube_edge}Β³ ({args.cube_edge:,} Γ— {args.cube_edge:,} Γ— {args.cube_edge:,})") + print(f" Total Lattice Points: {total_lattice_points:,}") + print(f" Time Steps: {args.num_steps:,}") + print(f" Total Lattice Updates: {total_lattice_updates:,}") + print(f" Repetitions: {args.repetitions}") + print(f" Compute Backend: {args.compute_backend.name}") + print(f" Precision Policy: {args.precision}") + print(f" Velocity Set: {args.velocity_set.__class__.__name__}") + print(f" Collision Model: {args.collision_model}") + print(f" Generate Report: {'Yes' if args.report else 'No'}") + print(f" Measure Scalability: {'Yes' if args.measure_scalability else 'No'}") + + if args.compute_backend.name == "NEON": + print(f" GPU Devices: {args.gpu_devices}") + occ_display = args.occ_display + print(f" OCC Strategy: {occ_display}") + + print() + + # Raw Data (if multiple repetitions) + if args.repetitions > 1: + print("RAW MEASUREMENT DATA:") + print("-" * 21) + print(f"{'Run':<6} {'Elapsed Time (s)':<18} {'MLUPs':<12} {'Time/Step (ms)':<15}") + print("-" * 53) + + raw_elapsed_times = stats["raw_elapsed_times"] + raw_mlups = stats["raw_mlups"] + + for i, (elapsed_time, mlups) in enumerate(zip(raw_elapsed_times, raw_mlups)): + time_per_step = elapsed_time / args.num_steps * 1000 + print(f"{i + 1:<6} {elapsed_time:<18.3f} {mlups:<12.2f} {time_per_step:<15.3f}") + + print("-" * 53) + print() + + # Performance Results (Statistical Summary) + print("PERFORMANCE RESULTS:") + print("-" * 20) + if args.repetitions > 1: + print(f" Time in main loop: {mean_elapsed_time:.3f} Β± {std_elapsed_time:.3f} seconds") + print(f" MLUPs: {mean_mlups:.2f} Β± {std_mlups:.2f}") + print(f" Time per LBM step: {mean_elapsed_time / args.num_steps * 1000:.3f} Β± {std_elapsed_time / args.num_steps * 1000:.3f} ms") + else: + print(f" Time in main loop: {mean_elapsed_time:.3f} seconds") + print(f" MLUPs: {mean_mlups:.2f}") + print(f" Time per LBM step: {mean_elapsed_time / args.num_steps * 1000:.3f} ms") + + if args.compute_backend.name == "NEON" and len(args.gpu_devices) > 1: + mlups_per_gpu = mean_mlups / len(args.gpu_devices) + if args.repetitions > 1: + mlups_per_gpu_std = std_mlups / len(args.gpu_devices) + print(f" MLUPs per GPU: {mlups_per_gpu:.2f} Β± {mlups_per_gpu_std:.2f}") + else: + print(f" MLUPs per GPU: {mlups_per_gpu:.2f}") + + print("=" * 70) + + +def print_scalability_summary(args, stats_list): + """Print comprehensive scalability summary with MLUPs statistics for different GPU counts""" + total_lattice_points = args.cube_edge**3 + total_lattice_updates = total_lattice_points * args.num_steps + + print("\n\n\n" + "=" * 95) + print(" SCALABILITY ANALYSIS") + print("=" * 95) + + # Simulation Parameters + print("SIMULATION PARAMETERS:") + print("-" * 25) + print(f" Grid Size: {args.cube_edge}Β³ ({args.cube_edge:,} Γ— {args.cube_edge:,} Γ— {args.cube_edge:,})") + print(f" Total Lattice Points: {total_lattice_points:,}") + print(f" Time Steps: {args.num_steps:,}") + print(f" Total Lattice Updates: {total_lattice_updates:,}") + print(f" Repetitions: {args.repetitions}") + print(f" Compute Backend: {args.compute_backend.name}") + print(f" Precision Policy: {args.precision}") + print(f" Velocity Set: {args.velocity_set.__class__.__name__}") + print(f" Collision Model: {args.collision_model}") + + if args.compute_backend.name == "NEON": + occ_display = args.occ_display + print(f" OCC Strategy: {occ_display}") + print(f" Available GPU Devices: {args.gpu_devices}") + + print() + + # Extract mean MLUPs for calculations + mlups_means = [stats["mean_mlups"] for stats in stats_list] + baseline_mlups = mlups_means[0] if mlups_means else 0 + + # Scalability Results + print("SCALABILITY RESULTS:") + print("-" * 20) + print(f"{'GPUs':<6} {'MLUPs (meanΒ±std)':<18} {'Speedup':<10} {'Efficiency':<12} {'MLUPs/GPU':<12}") + print("-" * 68) + + for i, stats in enumerate(stats_list): + num_gpus = i + 1 + mean_mlups = stats["mean_mlups"] + std_mlups = stats["std_dev_mlups"] + speedup = mean_mlups / baseline_mlups if baseline_mlups > 0 else 0 + efficiency = (speedup / num_gpus) if num_gpus > 0 else 0 + mlups_per_gpu = mean_mlups / num_gpus if num_gpus > 0 else 0 + + # Format MLUPs with standard deviation + if args.repetitions > 1: + mlups_str = f"{mean_mlups:.2f}Β±{std_mlups:.2f}" + else: + mlups_str = f"{mean_mlups:.2f}" + + print(f"{num_gpus:<6} {mlups_str:<18} {speedup:<10.2f} {efficiency:<11.3f} {mlups_per_gpu:<12.2f}") + + print("-" * 68) + + # Summary Statistics + if len(stats_list) > 1: + max_mlups = max(mlups_means) + max_mlups_idx = mlups_means.index(max_mlups) + max_speedup = max_mlups / baseline_mlups if baseline_mlups > 0 else 0 + best_efficiency_idx = 0 + best_efficiency = 0.0 + + for i, mean_mlups in enumerate(mlups_means): + num_gpus = i + 1 + speedup = mean_mlups / baseline_mlups if baseline_mlups > 0 else 0 + efficiency = (speedup / num_gpus) if num_gpus > 0 else 0 + if efficiency > best_efficiency: + best_efficiency = efficiency + best_efficiency_idx = i + + print() + print("SUMMARY STATISTICS:") + print("-" * 19) + print(f" Best Performance: {max_mlups:.2f} MLUPs ({max_mlups_idx + 1} GPUs)") + if args.repetitions > 1: + max_std = stats_list[max_mlups_idx]["std_dev_mlups"] + print(f" Performance Std Dev: Β±{max_std:.2f} MLUPs") + print(f" Maximum Speedup: {max_speedup:.2f}x") + print(f" Best Efficiency: {best_efficiency:.3f} ({best_efficiency_idx + 1} GPUs)") + print(f" Scalability Range: 1-{len(stats_list)} GPUs") + + print("=" * 95) + + +def report(args, stats): + import neon + import sys + + report = neon.Report("LBM MLUPS LDC") + + # Save the full command line + command_line = " ".join(sys.argv) + report.add_member("command_line", command_line) + + report.add_member("velocity_set", args.velocity_set.__class__.__name__) + report.add_member("compute_backend", args.compute_backend.name) + report.add_member("precision_policy", args.precision) + report.add_member("collision_model", args.collision_model) + report.add_member("grid_size", args.cube_edge) + report.add_member("num_steps", args.num_steps) + report.add_member("repetitions", args.repetitions) + + # Statistical measures + report.add_member("mean_elapsed_time", stats["mean_elapsed_time"]) + report.add_member("mean_mlups", stats["mean_mlups"]) + report.add_member("std_dev_elapsed_time", stats["std_dev_elapsed_time"]) + report.add_member("std_dev_mlups", stats["std_dev_mlups"]) + + # Raw data vectors (if multiple repetitions) + if args.repetitions > 1: + report.add_member_vector("raw_elapsed_times", stats["raw_elapsed_times"]) + report.add_member_vector("raw_mlups", stats["raw_mlups"]) + + # Legacy fields for backwards compatibility + report.add_member("elapsed_time", stats["mean_elapsed_time"]) + report.add_member("mlups", stats["mean_mlups"]) + + report.add_member("occ", args.occ_display) + report.add_member_vector("gpu_devices", args.gpu_devices) + report.add_member("num_devices", len(args.gpu_devices)) + report.add_member("measure_scalability", args.measure_scalability) + + # Generate report name following the convention: script_name + parameters + report_name = "mlups_3d" + report_name += f"_velocity_set_{args.velocity_set.__class__.__name__}" + report_name += f"_compute_backend_{args.compute_backend.name}" + report_name += f"_precision_policy_{args.precision.replace('/', '_')}" + report_name += f"_collision_model_{args.collision_model}" + report_name += f"_grid_size_{args.cube_edge}" + report_name += f"_num_steps_{args.num_steps}" + + if args.compute_backend.name == "NEON": + report_name += f"_occ_{args.occ_display}" + report_name += f"_num_devices_{len(args.gpu_devices)}" + + if args.repetitions > 1: + report_name += f"_repetitions_{args.repetitions}" + + report.write(report_name, True) + + +# -------------------------- Simulation Loop -------------------------- + + +def benchmark(args): + compute_backend, precision_policy, options = init_xlb(args) + grid_shape = (args.cube_edge, args.cube_edge, args.cube_edge) + + elapsed_time_list = [] + mlups_list = [] + elapsed_time_list = run_simulation( + compute_backend=compute_backend, + precision_policy=precision_policy, + grid_shape=grid_shape, + num_steps=args.num_steps, + options=options, + export_final_velocity=args.export_final_velocity, + repetitions=args.repetitions, + num_devices=len(args.gpu_devices), + collision_model=args.collision_model, + ) + + for elapsed_time in elapsed_time_list: + mlups = calculate_mlups(args.cube_edge, args.num_steps, elapsed_time) + mlups_list.append(mlups) + + mean_mlups = np.mean(mlups_list) + std_dev_mlups = np.std(mlups_list) + mean_elapsed_time = np.mean(elapsed_time_list) + std_dev_elapsed_time = np.std(elapsed_time_list) + + stats = { + "mean_mlups": mean_mlups, + "std_dev_mlups": std_dev_mlups, + "mean_elapsed_time": mean_elapsed_time, + "std_dev_elapsed_time": std_dev_elapsed_time, + "num_devices": len(args.gpu_devices), + "raw_mlups": mlups_list, + "raw_elapsed_times": elapsed_time_list, + } + # Generate report if requested + if args.report: + report(args, stats) + print("Report generated successfully.") + + return stats + + +def main(): + args = parse_arguments() + if not args.measure_scalability: + stats = benchmark(args) + # For single run, print_summary expects individual values with additional stats + print_summary_with_stats(args, stats) + return + + stats_list = [] + for num_devices in range(1, len(args.gpu_devices) + 1): + import copy + + args_copy = copy.deepcopy(args) + args_copy.gpu_devices = args_copy.gpu_devices[:num_devices] + stats = benchmark(args_copy) + stats_list.append(stats) + + # Print comprehensive scalability analysis + print_scalability_summary(args, stats_list) + + +if __name__ == "__main__": + main() diff --git a/examples/performance/mlups_3d_multires.py b/examples/performance/mlups_3d_multires.py new file mode 100644 index 00000000..ba71985b --- /dev/null +++ b/examples/performance/mlups_3d_multires.py @@ -0,0 +1,402 @@ +""" +MLUPS benchmark for the multi-resolution LBM solver. + +Runs a lid-driven cavity simulation on a multi-resolution Neon grid and +reports the Equivalent Million Lattice Updates Per Second (EMLUPS). + +Usage:: + + python mlups_3d_multires.py neon \\ + [options] + +Example:: + + python mlups_3d_multires.py 100 1000 neon fp32/fp32 2 NAIVE_COLLIDE_STREAM +""" + +import xlb +import argparse +import time +import warp as wp +import numpy as np +import neon + +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import multires_grid_factory +from xlb.operator.stepper import MultiresIncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import FullwayBounceBackBC, EquilibriumBC +from xlb.mres_perf_optimization_type import MresPerfOptimizationType + + +def parse_arguments(): + """Parse and validate command-line arguments.""" + parser = argparse.ArgumentParser( + description="MLUPS for 3D Lattice Boltzmann Method Simulation with Multi-resolution Grid", + epilog=""" +Examples: + %(prog)s 100 1000 neon fp32/fp32 2 NAIVE_COLLIDE_STREAM + %(prog)s 200 500 neon fp64/fp64 3 FUSION_AT_FINEST --report + %(prog)s 50 2000 neon fp32/fp16 2 NAIVE_COLLIDE_STREAM --export_final_velocity + +Valid values: + compute_backend: neon + precision: fp32/fp32, fp64/fp64, fp64/fp32, fp32/fp16 + mres_perf_opt: NAIVE_COLLIDE_STREAM, FUSION_AT_FINEST + velocity_set: D3Q19, D3Q27 + collision_model: BGK, KBC + """, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # Positional arguments + parser.add_argument("cube_edge", type=int, help="Length of the edge of the cubic grid (e.g., 100)") + parser.add_argument("num_steps", type=int, help="Number of timesteps for the simulation (e.g., 1000)") + parser.add_argument("compute_backend", type=str, help="Backend for the simulation (neon)") + parser.add_argument("precision", type=str, help="Precision for the simulation (fp32/fp32, fp64/fp64, fp64/fp32, fp32/fp16)") + parser.add_argument("num_levels", type=int, help="Number of levels for the multiresolution grid (e.g., 2)") + parser.add_argument( + "mres_perf_opt", + type=MresPerfOptimizationType.from_string, + help="Multi-resolution performance optimization strategy (NAIVE_COLLIDE_STREAM, FUSION_AT_FINEST)", + ) + + # Optional arguments + parser.add_argument("--num_devices", type=int, default=0, help="Number of devices for the simulation (default: 0)") + parser.add_argument("--velocity_set", type=str, default="D3Q19", help="Lattice type: D3Q19 or D3Q27 (default: D3Q19)") + parser.add_argument("--collision_model", type=str, default="BGK", help="Collision model: BGK or KBC (default: BGK)") + + parser.add_argument("--report", action="store_true", help="Generate a neon report file (default: disabled)") + parser.add_argument("--export_final_velocity", action="store_true", help="Export the final velocity field to a vti file (default: disabled)") + + try: + args = parser.parse_args() + except SystemExit: + # Re-raise with custom message + print("\n" + "=" * 60) + print("USAGE EXAMPLES:") + print("=" * 60) + print("python mlups_3d_multires.py 100 1000 neon fp32/fp32 2 NAIVE_COLLIDE_STREAM") + print("python mlups_3d_multires.py 200 500 neon fp64/fp64 3 FUSION_AT_FINEST --report") + print("\nVALID VALUES:") + print(" compute_backend: neon") + print(" precision: fp32/fp32, fp64/fp64, fp64/fp32, fp32/fp16") + print(" mres_perf_opt: NAIVE_COLLIDE_STREAM, FUSION_AT_FINEST") + print(" velocity_set: D3Q19, D3Q27") + print(" collision_model: BGK, KBC") + print("=" * 60) + raise + + print_args(args) + + if args.compute_backend != "neon": + raise ValueError("Invalid compute backend specified. Use 'neon' which supports multi-resolution!") + + if args.collision_model not in ["BGK", "KBC"]: + raise ValueError("Invalid collision model specified. Use 'BGK' or 'KBC'.") + + return args + + +def print_args(args): + """Print the simulation configuration to stdout.""" + # Print simulation configuration + print("=" * 60) + print(" 3D LATTICE BOLTZMANN SIMULATION CONFIG") + print("=" * 60) + print(f"Grid Size: {args.cube_edge}Β³ ({args.cube_edge:,} Γ— {args.cube_edge:,} Γ— {args.cube_edge:,})") + print(f"Total Lattice Points: {args.cube_edge**3:,}") + print(f"Time Steps: {args.num_steps:,}") + print(f"Number Levels: {args.num_levels}") + print(f"Compute Backend: {args.compute_backend}") + print(f"Precision Policy: {args.precision}") + print(f"Velocity Set: {args.velocity_set}") + print(f"Collision Model: {args.collision_model}") + print(f"Mres Perf Opt: {args.mres_perf_opt}") + print(f"Generate Report: {'Yes' if args.report else 'No'}") + print(f"Export Velocity: {'Yes' if args.export_final_velocity else 'No'}") + + print("=" * 60) + print("Starting simulation...") + print() + + +def setup_simulation(args): + """Initialize XLB globals (velocity set, backend, precision) from CLI args. + + Returns + ------- + VelocitySet + The configured lattice velocity set. + """ + compute_backend = None + if args.compute_backend == "neon": + compute_backend = ComputeBackend.NEON + else: + raise ValueError("Invalid compute backend specified. Use 'neon' which supports multi-resolution!") + + precision_policy_map = { + "fp32/fp32": PrecisionPolicy.FP32FP32, + "fp64/fp64": PrecisionPolicy.FP64FP64, + "fp64/fp32": PrecisionPolicy.FP64FP32, + "fp32/fp16": PrecisionPolicy.FP32FP16, + } + precision_policy = precision_policy_map.get(args.precision) + if precision_policy is None: + raise ValueError("Invalid precision") + + velocity_set = None + if args.velocity_set == "D3Q19": + velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend) + elif args.velocity_set == "D3Q27": + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=compute_backend) + if velocity_set is None: + raise ValueError("Invalid velocity set") + + xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, + ) + + return velocity_set + + +def ldc_multires_setup(grid_shape, velocity_set, num_levels): + """Lid-driven cavity with refinement peeling inward from the boundary. + + Each finer level covers only the outermost shell of its parent, + concentrating resolution near the walls. + + Parameters + ---------- + grid_shape : tuple of int + Domain size at the finest level. + velocity_set : VelocitySet + Lattice velocity set. + num_levels : int + Number of refinement levels. + + Returns + ------- + grid : NeonMultiresGrid + lid : list of index arrays (per level) + walls : list of index arrays (per level) + """ + + def peel(dim, idx, peel_level, outwards): + if outwards: + xIn = idx.x <= peel_level or idx.x >= dim.x - 1 - peel_level + yIn = idx.y <= peel_level or idx.y >= dim.y - 1 - peel_level + zIn = idx.z <= peel_level or idx.z >= dim.z - 1 - peel_level + return xIn or yIn or zIn + else: + xIn = idx.x >= peel_level and idx.x <= dim.x - 1 - peel_level + yIn = idx.y >= peel_level and idx.y <= dim.y - 1 - peel_level + zIn = idx.z >= peel_level and idx.z <= dim.z - 1 - peel_level + return xIn and yIn and zIn + + dim = neon.Index_3d(grid_shape[0], grid_shape[1], grid_shape[2]) + + def get_peeled_np(level, width): + divider = 2**level + m = neon.Index_3d(dim.x // divider, dim.y // divider, dim.z // divider) + if level == 0: + m = dim + + mask = np.zeros((m.x, m.y, m.z), dtype=int) + mask = np.ascontiguousarray(mask, dtype=np.int32) + # loop over all the elements in mask and set to one any that have x=0 or y=0 or z=0 + for i in range(m.x): + for j in range(m.y): + for k in range(m.z): + idx = neon.Index_3d(i, j, k) + val = 0 + if peel(m, idx, width, True): + val = 1 + mask[i, j, k] = val + return mask + + def get_levels(num_levels): + levels = [] + for i in range(num_levels - 1): + l = get_peeled_np(i, 8) + levels.append(l) + lastLevel = num_levels - 1 + divider = 2**lastLevel + m = neon.Index_3d(dim.x // divider + 1, dim.y // divider + 1, dim.z // divider + 1) + lastLevel = np.ones((m.x, m.y, m.z), dtype=int) + lastLevel = np.ascontiguousarray(lastLevel, dtype=np.int32) + levels.append(lastLevel) + return levels + + levels = get_levels(num_levels) + + grid = multires_grid_factory( + grid_shape, + velocity_set=velocity_set, + sparsity_pattern_list=levels, + sparsity_pattern_origins=[neon.Index_3d(0, 0, 0)] * len(levels), + ) + + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + lid = box_no_edge["top"] + walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(len(grid.shape))] + walls = np.unique(np.array(walls), axis=-1).tolist() + # convert bc indices to a list of list, where the first entry of the list corresponds to the finest level + lid = [lid] + [[] for _ in range(num_levels - 1)] + walls = [walls] + [[] for _ in range(num_levels - 1)] + return grid, lid, walls + + +def run( + velocity_set, + grid_shape, + num_steps, + num_levels, + collision_model, + export_final_velocity, + mres_perf_opt, +): + """Set up and execute the benchmark simulation. + + Returns + ------- + dict + ``{"time": elapsed_seconds, "num_levels": int}`` + """ + # Create grid and setup boundary conditions + grid, lid, walls = ldc_multires_setup(grid_shape, velocity_set, num_levels) + + prescribed_vel = 0.1 + boundary_conditions = [ + EquilibriumBC(rho=1.0, u=(prescribed_vel, 0.0, 0.0), indices=lid), + FullwayBounceBackBC(indices=walls), + ] + + # Problem parameters + Re = 5000.0 + clength = grid_shape[0] - 1 + visc = prescribed_vel * clength / Re + omega_finest = 1.0 / (3.0 * visc + 0.5) + + # Define a multi-resolution simulation manager + sim = xlb.helper.MultiresSimulationManager( + omega_finest=omega_finest, + grid=grid, + boundary_conditions=boundary_conditions, + collision_type=collision_model, + mres_perf_opt=mres_perf_opt, + ) + + # sim.export_macroscopic("Initial_") + # sim.step() + + print("start timing") + wp.synchronize() + start_time = time.time() + + if num_levels == 1: + num_steps = num_steps // 2 + + for i in range(num_steps): + sim.step() + # if i % 1000 == 0: + # print(f"step {i}") + # sim.export_macroscopic("u_lid_driven_cavity_") + wp.synchronize() + t = time.time() - start_time + print(f"Timing {t}") + + if export_final_velocity: + sim.export_macroscopic("u_lid_driven_cavity_") + + # sim.export_macroscopic("u_lid_driven_cavity_") + num_levels = grid.count_levels + return {"time": t, "num_levels": num_levels} + + +def calculate_mlups(cube_edge, num_steps, elapsed_time, num_levels): + """Compute the Equivalent Million Lattice Updates Per Second (EMLUPS). + + The metric accounts for the fact that finer levels are stepped + 2^(num_levels-1) times per coarsest-level step. + + Returns + ------- + dict + ``{"EMLUPS": float, "finer_steps": int}`` + """ + num_step_finer = num_steps * 2 ** (num_levels - 1) + total_lattice_updates = cube_edge**3 * num_step_finer + mlups = (total_lattice_updates / elapsed_time) / 1e6 + return {"EMLUPS": mlups, "finer_steps": num_step_finer} + + # # remove boundary cells + # rho = rho[:, 1:-1, 1:-1, 1:-1] + # u = u[:, 1:-1, 1:-1, 1:-1] + # u_magnitude = (u[0] ** 2 + u[1] ** 2) ** 0.5 + # + # fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude} + # + # # save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity") + # ny=fields["u_magnitude"].shape[1] + # from xlb.utils import save_image + # save_image(fields["u_magnitude"][:, ny//2, :], timestep=i, prefix="lid_driven_cavity") + + +def generate_report(args, stats, mlups_stats): + """Generate a neon report file with simulation parameters and results""" + import neon + import sys + + report = neon.Report("LBM MLUPS Multiresolution LDC") + + # Save the full command line + command_line = " ".join(sys.argv) + report.add_member("command_line", command_line) + + report.add_member("velocity_set", args.velocity_set) + report.add_member("compute_backend", args.compute_backend) + report.add_member("precision_policy", args.precision) + report.add_member("collision_model", args.collision_model) + report.add_member("grid_size", args.cube_edge) + report.add_member("num_steps", args.num_steps) + report.add_member("num_levels", stats["num_levels"]) + report.add_member("finer_steps", mlups_stats["finer_steps"]) + + # Performance metrics + report.add_member("elapsed_time", stats["time"]) + report.add_member("emlups", mlups_stats["EMLUPS"]) + + report_name = f"mlups_3d_multires_size_{args.cube_edge}_levels_{stats['num_levels']}" + report.write(report_name, True) + print("Report generated successfully.") + + +def main(): + args = parse_arguments() + velocity_set = setup_simulation(args) + grid_shape = (args.cube_edge, args.cube_edge, args.cube_edge) + stats = run( + velocity_set, grid_shape, args.num_steps, args.num_levels, args.collision_model, args.export_final_velocity, mres_perf_opt=args.mres_perf_opt + ) + mlups_stats = calculate_mlups(args.cube_edge, args.num_steps, stats["time"], stats["num_levels"]) + + print(f"Simulation completed in {stats['time']:.2f} seconds") + print(f"Number of levels {stats['num_levels']}") + print(f"Cube edge {args.cube_edge}") + print(f"Coarse Iterations {args.num_steps}") + finer_steps = mlups_stats["finer_steps"] + print(f"Fine Iterations {finer_steps}") + EMLUPS = mlups_stats["EMLUPS"] + print(f"EMLUPs: {EMLUPS:.2f}") + + # Generate report if requested + if args.report: + generate_report(args, stats, mlups_stats) + + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index 9ff371a1..70af7180 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -30,7 +30,7 @@ repo_url: https://github.com/Autodesk/XLB edit_uri: edit/master/docs/ watch: -- ./src/ +- ./xlb/ # Can be used to add a meta tag to the HTML header site_description: 'Documentation for project XLB' @@ -147,10 +147,4 @@ extra_javascript: - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js nav: - - XLB's home: index.md - - XLB API: - - XLB base: base.md - - XLB models: models.md - - XLB lattice: lattice.md - - XLB utils: utils.md - - XLB boundary conditions: boundary_conditions.md \ No newline at end of file + - XLB's home: index.md \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 11ee0fd5..da8d50b2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,13 @@ -jax==0.4.20 -jaxlib==0.4.20 -jmp==0.0.4 -matplotlib==3.8.0 -numpy==1.26.1 -pyvista==0.42.3 -Rtree==1.0.1 -trimesh==4.0.0 -orbax-checkpoint==0.4.1 -termcolor==2.3.0 -PhantomGaze @ git+https://github.com/loliverhennigh/PhantomGaze.git \ No newline at end of file +jax[cuda] +matplotlib +numpy +pyvista +Rtree +trimesh +numpy-stl +pydantic +nvtx +pytest +ruff +usd-core +h5py \ No newline at end of file diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 00000000..83c6370e --- /dev/null +++ b/ruff.toml @@ -0,0 +1,42 @@ +# Adopted from tinygrad's ruff.toml thanks @geohot +indent-width = 4 +preview = true +target-version = "py38" + +lint.select = [ + "F", # Pyflakes + "W6", + "E71", + "E72", + "E112", # no-indented-block + "E113", # unexpected-indentation + # "E124", + "E203", # whitespace-before-punctuation + "E272", # multiple-spaces-before-keyword + "E303", # too-many-blank-lines + "E304", # blank-line-after-decorator + "E501", # line-too-long + # "E502", + "E702", # multiple-statements-on-one-line-semicolon + "E703", # useless-semicolon + "E731", # lambda-assignment + "W191", # tab-indentation + "W291", # trailing-whitespace + "W293", # blank-line-with-whitespace + "UP039", # unnecessary-class-parentheses + "C416", # unnecessary-comprehension + "RET506", # superfluous-else-raise + "RET507", # superfluous-else-continue + "A", # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing + "SIM105", # suppressible-exception + "FURB110",# if-exp-instead-of-or-operator +] + +# unused-variable, shadowing a Python builtin module, Module imported but unused +lint.ignore = ["F841", "A005", "F401"] +line-length = 150 + +exclude = [ + "docs/", + "xlb/experimental/", +] \ No newline at end of file diff --git a/setup.py b/setup.py index e69de29b..ed261843 100644 --- a/setup.py +++ b/setup.py @@ -0,0 +1,110 @@ +import os +import platform +import subprocess +import sys + +from setuptools import setup, find_packages +from setuptools.command.install import install + + +def _neon_extra_requested(): + """Best-effort detection of [neon] extra from install invocation.""" + for arg in sys.argv: + if "neon" in arg and ("[" in arg or "xlb" in arg): + return True + return False + + +def _uninstall_warp_lang(*, reason: str) -> None: + """Uninstall the ``warp-lang`` distribution so Neon's bundled warp fork is used.""" + if os.environ.get("XLB_NEON_SKIP_UNINSTALL_WARP", "").lower() in ("1", "true", "yes"): + return + print(f"[xlb] {reason}") + try: + subprocess.run( + [sys.executable, "-m", "pip", "uninstall", "warp-lang", "-y"], + check=False, + capture_output=True, + ) + except Exception as exc: # noqa: BLE001 + print(f"[xlb] Warning: failed to uninstall warp-lang: {exc}") + + +_NEON_VERSION = "0.5.2a1" +_NEON_RELEASE_URL = f"https://github.com/Autodesk/Neon/releases/download/v{_NEON_VERSION}" + + +def _neon_wheel_requirement(): + """Build a direct-reference requirement for the neon_gpu wheel matching the running Python.""" + tag = f"cp{sys.version_info.major}{sys.version_info.minor}" + machine = platform.machine() + plat = "linux_aarch64" if machine == "aarch64" else "linux_x86_64" + wheel = f"neon_gpu-{_NEON_VERSION}-{tag}-{tag}-{plat}.whl" + url = f"{_NEON_RELEASE_URL}/{wheel}" + req = f"neon_gpu @ {url}" + print(f"[xlb] Neon wheel for Python {sys.version_info.major}.{sys.version_info.minor} ({plat}): {url}") + print(f"[xlb] Neon requirement: {req}") + return req + + +class InstallWithNeonHooks(install): + """Uninstall ``warp-lang`` before and after install when ``[neon]`` is requested. + + * **Before** ``pip``/setuptools install dependencies: removes any previously + installed ``warp-lang`` so an older or PyPI build does not linger next to + Neon's fork (``neon_gpu`` ships its own warp). + * **After** install: removes the ``warp-lang`` pulled in by ``install_requires``, + leaving Neon's warp as the one on the path. + + Only runs when installing from source (e.g. sdist or git). Wheel installs + do not run setup.py, so for ``pip install xlb[neon]`` from PyPI you may + need to run ``pip uninstall warp-lang`` first if it is already installed. + Set XLB_NEON_SKIP_UNINSTALL_WARP=1 to disable this behaviour. + """ + + def run(self): + if _neon_extra_requested(): + _uninstall_warp_lang( + reason=("Removing any existing warp-lang before Neon install (neon_gpu provides its own warp fork)."), + ) + install.run(self) + if _neon_extra_requested(): + _uninstall_warp_lang( + reason=("Removing PyPI warp-lang after install (core deps); use the warp bundled with neon_gpu."), + ) + + +setup( + name="xlb", + version="0.3.2", + description="XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML", + long_description=open("README.md").read(), + long_description_content_type="text/markdown", + author="Mehdi Ataei", + url="https://github.com/Autodesk/XLB", + license="Apache License 2.0", + packages=find_packages(), + install_requires=[ + "matplotlib>=3.9.2", + "numpy>=2.1.2", + "pyvista>=0.44.1", + "trimesh>=4.4.9", + "numpy-stl>=3.1.2", + "pydantic>=2.9.1", + "ruff>=0.14.1", + "jax>=0.8.2", # Base JAX CPU-only requirement + "warp-lang>=1.10.0", # Required at import time (core modules import warp) + "nvtx>=0.2.0", # NVTX ranges (e.g. nse_multires_stepper); listed in requirements.txt + ], + extras_require={ + "warp": ["warp-lang>=1.10.0"], # Kept for explicit `pip install xlb[warp]` / Neon uninstall hook docs + "cuda": ["jax[cuda13]>=0.8.2"], # For CUDA installations (pip install -U "jax[cuda13]") + "tpu": ["jax[tpu]>=0.8.2"], # For TPU installations + # h5py: MultiresIO / Neon multi-resolution export to HDF5 (see xlb.utils.mesher). + "neon": [_neon_wheel_requirement(), "h5py>=3.10.0"], + "test": ["pytest>=8.0.0"], + }, + python_requires=">=3.11", + dependency_links=["https://storage.googleapis.com/jax-releases/libtpu_releases.html"], + cmdclass={"install": InstallWithNeonHooks}, +) diff --git a/src/base.py b/src/base.py deleted file mode 100644 index 7fb4f57e..00000000 --- a/src/base.py +++ /dev/null @@ -1,1086 +0,0 @@ -# Standard Libraries -import os -import time - -# Third-Party Libraries -import jax -import jax.numpy as jnp -import jmp -import numpy as np -from termcolor import colored - -# JAX-related imports -from jax import jit, lax, vmap -from jax.experimental import mesh_utils -from jax.experimental.multihost_utils import process_allgather -from jax.experimental.shard_map import shard_map -from jax.sharding import NamedSharding, PartitionSpec, PositionalSharding, Mesh -import orbax.checkpoint as orb - -# functools imports -from functools import partial - -# Local/Custom Libraries -from src.utils import downsample_field - -jax.config.update("jax_spmd_mode", 'allow_all') -# Disables annoying TF warnings -os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' - -class LBMBase(object): - """ - LBMBase: A class that represents a base for Lattice Boltzmann Method simulation. - - Parameters - ---------- - lattice (object): The lattice object that contains the lattice structure and weights. - omega (float): The relaxation parameter for the LBM simulation. - nx (int): Number of grid points in the x-direction. - ny (int): Number of grid points in the y-direction. - nz (int, optional): Number of grid points in the z-direction. Defaults to 0. - precision (str, optional): A string specifying the precision used for the simulation. Defaults to "f32/f32". - """ - - def __init__(self, **kwargs): - self.omega = kwargs.get("omega") - self.nx = kwargs.get("nx") - self.ny = kwargs.get("ny") - self.nz = kwargs.get("nz") - - self.precision = kwargs.get("precision") - computedType, storedType = self.set_precisions(self.precision) - self.precisionPolicy = jmp.Policy(compute_dtype=computedType, - param_dtype=computedType, output_dtype=storedType) - - self.lattice = kwargs.get("lattice") - self.checkpointRate = kwargs.get("checkpoint_rate", 0) - self.checkpointDir = kwargs.get("checkpoint_dir", './checkpoints') - self.downsamplingFactor = kwargs.get("downsampling_factor", 1) - self.printInfoRate = kwargs.get("print_info_rate", 100) - self.ioRate = kwargs.get("io_rate", 0) - self.returnFpost = kwargs.get("return_fpost", False) - self.computeMLUPS = kwargs.get("compute_MLUPS", False) - self.restore_checkpoint = kwargs.get("restore_checkpoint", False) - self.nDevices = jax.device_count() - self.backend = jax.default_backend() - - if self.computeMLUPS: - self.restore_checkpoint = False - self.ioRate = 0 - self.checkpointRate = 0 - self.printInfoRate = 0 - - # Check for distributed mode - if self.nDevices > jax.local_device_count(): - print("WARNING: Running in distributed mode. Make sure that jax.distributed.initialize is called before performing any JAX computations.") - - self.c = self.lattice.c - self.q = self.lattice.q - self.w = self.lattice.w - self.dim = self.lattice.d - - # Set the checkpoint manager - if self.checkpointRate > 0: - mngr_options = orb.CheckpointManagerOptions(save_interval_steps=self.checkpointRate, max_to_keep=1) - self.mngr = orb.CheckpointManager(self.checkpointDir, orb.PyTreeCheckpointer(), options=mngr_options) - else: - self.mngr = None - - # Adjust the number of grid points in the x direction, if necessary. - # If the number of grid points is not divisible by the number of devices - # it increases the number of grid points to the next multiple of the number of devices. - # This is done in order to accommodate the domain sharding per XLA device - nx, ny, nz = kwargs.get("nx"), kwargs.get("ny"), kwargs.get("nz") - if None in {nx, ny, nz}: - raise ValueError("nx, ny, and nz must be provided. For 2D examples, nz must be set to 0.") - self.nx = nx - if nx % self.nDevices: - self.nx = nx + (self.nDevices - nx % self.nDevices) - print("WARNING: nx increased from {} to {} in order to accommodate domain sharding per XLA device.".format(nx, self.nx)) - self.ny = ny - self.nz = nz - - self.show_simulation_parameters() - - # Store grid information - self.gridInfo = { - "nx": self.nx, - "ny": self.ny, - "nz": self.nz, - "dim": self.lattice.d, - "lattice": self.lattice - } - - P = PartitionSpec - - # Define the right permutation - self.rightPerm = [(i, (i + 1) % self.nDevices) for i in range(self.nDevices)] - # Define the left permutation - self.leftPerm = [((i + 1) % self.nDevices, i) for i in range(self.nDevices)] - - # Set up the sharding and streaming for 2D and 3D simulations - if self.dim == 2: - self.devices = mesh_utils.create_device_mesh((self.nDevices, 1, 1)) - self.mesh = Mesh(self.devices, axis_names=("x", "y", "value")) - self.sharding = NamedSharding(self.mesh, P("x", "y", "value")) - - self.streaming = jit(shard_map(self.streaming_m, mesh=self.mesh, - in_specs=P("x", None, None), out_specs=P("x", None, None), check_rep=False)) - - # Set up the sharding and streaming for 2D and 3D simulations - elif self.dim == 3: - self.devices = mesh_utils.create_device_mesh((self.nDevices, 1, 1, 1)) - self.mesh = Mesh(self.devices, axis_names=("x", "y", "z", "value")) - self.sharding = NamedSharding(self.mesh, P("x", "y", "z", "value")) - - self.streaming = jit(shard_map(self.streaming_m, mesh=self.mesh, - in_specs=P("x", None, None, None), out_specs=P("x", None, None, None), check_rep=False)) - - else: - raise ValueError(f"dim = {self.dim} not supported") - - # Compute the bounding box indices for boundary conditions - self.boundingBoxIndices= self.bounding_box_indices() - # Create boundary data for the simulation - self._create_boundary_data() - self.force = self.get_force() - - @property - def lattice(self): - return self._lattice - - @lattice.setter - def lattice(self, value): - if value is None: - raise ValueError("Lattice type must be provided.") - if self.nz == 0 and value.name not in ['D2Q9']: - raise ValueError("For 2D simulations, lattice type must be LatticeD2Q9.") - if self.nz != 0 and value.name not in ['D3Q19', 'D3Q27']: - raise ValueError("For 3D simulations, lattice type must be LatticeD3Q19, or LatticeD3Q27.") - - self._lattice = value - - @property - def omega(self): - return self._omega - - @omega.setter - def omega(self, value): - if value is None: - raise ValueError("omega must be provided") - if not isinstance(value, float): - raise TypeError("omega must be a float") - self._omega = value - - @property - def nx(self): - return self._nx - - @nx.setter - def nx(self, value): - if value is None: - raise ValueError("nx must be provided") - if not isinstance(value, int): - raise TypeError("nx must be an integer") - self._nx = value - - @property - def ny(self): - return self._ny - - @ny.setter - def ny(self, value): - if value is None: - raise ValueError("ny must be provided") - if not isinstance(value, int): - raise TypeError("ny must be an integer") - self._ny = value - - @property - def nz(self): - return self._nz - - @nz.setter - def nz(self, value): - if value is None: - raise ValueError("nz must be provided") - if not isinstance(value, int): - raise TypeError("nz must be an integer") - self._nz = value - - @property - def precision(self): - return self._precision - - @precision.setter - def precision(self, value): - if not isinstance(value, str): - raise TypeError("precision must be a string") - self._precision = value - - @property - def checkpointRate(self): - return self._checkpointRate - - @checkpointRate.setter - def checkpointRate(self, value): - if not isinstance(value, int): - raise TypeError("checkpointRate must be an integer") - self._checkpointRate = value - - @property - def checkpointDir(self): - return self._checkpointDir - - @checkpointDir.setter - def checkpointDir(self, value): - if not isinstance(value, str): - raise TypeError("checkpointDir must be a string") - self._checkpointDir = value - - @property - def downsamplingFactor(self): - return self._downsamplingFactor - - @downsamplingFactor.setter - def downsamplingFactor(self, value): - if not isinstance(value, int): - raise TypeError("downsamplingFactor must be an integer") - self._downsamplingFactor = value - - @property - def printInfoRate(self): - return self._printInfoRate - - @printInfoRate.setter - def printInfoRate(self, value): - if not isinstance(value, int): - raise TypeError("printInfoRate must be an integer") - self._printInfoRate = value - - @property - def ioRate(self): - return self._ioRate - - @ioRate.setter - def ioRate(self, value): - if not isinstance(value, int): - raise TypeError("ioRate must be an integer") - self._ioRate = value - - @property - def returnFpost(self): - return self._returnFpost - - @returnFpost.setter - def returnFpost(self, value): - if not isinstance(value, bool): - raise TypeError("returnFpost must be a boolean") - self._returnFpost = value - - @property - def computeMLUPS(self): - return self._computeMLUPS - - @computeMLUPS.setter - def computeMLUPS(self, value): - if not isinstance(value, bool): - raise TypeError("computeMLUPS must be a boolean") - self._computeMLUPS = value - - @property - def restore_checkpoint(self): - return self._restore_checkpoint - - @restore_checkpoint.setter - def restore_checkpoint(self, value): - if not isinstance(value, bool): - raise TypeError("restore_checkpoint must be a boolean") - self._restore_checkpoint = value - - @property - def nDevices(self): - return self._nDevices - - @nDevices.setter - def nDevices(self, value): - if not isinstance(value, int): - raise TypeError("nDevices must be an integer") - self._nDevices = value - - def show_simulation_parameters(self): - attributes_to_show = [ - 'omega', 'nx', 'ny', 'nz', 'dim', 'precision', 'lattice', - 'checkpointRate', 'checkpointDir', 'downsamplingFactor', - 'printInfoRate', 'ioRate', 'computeMLUPS', - 'restore_checkpoint', 'backend', 'nDevices' - ] - - descriptive_names = { - 'omega': 'Omega', - 'nx': 'Grid Points in X', - 'ny': 'Grid Points in Y', - 'nz': 'Grid Points in Z', - 'dim': 'Dimensionality', - 'precision': 'Precision Policy', - 'lattice': 'Lattice Type', - 'checkpointRate': 'Checkpoint Rate', - 'checkpointDir': 'Checkpoint Directory', - 'downsamplingFactor': 'Downsampling Factor', - 'printInfoRate': 'Print Info Rate', - 'ioRate': 'I/O Rate', - 'computeMLUPS': 'Compute MLUPS', - 'restore_checkpoint': 'Restore Checkpoint', - 'backend': 'Backend', - 'nDevices': 'Number of Devices' - } - simulation_name = self.__class__.__name__ - - print(colored(f'**** Simulation Parameters for {simulation_name} ****', 'green')) - - header = f"{colored('Parameter', 'blue'):>30} | {colored('Value', 'yellow')}" - print(header) - print('-' * 50) - - for attr in attributes_to_show: - value = getattr(self, attr, 'Attribute not set') - descriptive_name = descriptive_names.get(attr, attr) # Use the attribute name as a fallback - row = f"{colored(descriptive_name, 'blue'):>30} | {colored(value, 'yellow')}" - print(row) - - def _create_boundary_data(self): - """ - Create boundary data for the Lattice Boltzmann simulation by setting boundary conditions, - creating grid mask, and preparing local masks and normal arrays. - """ - self.BCs = [] - self.set_boundary_conditions() - # Accumulate the indices of all BCs to create the grid mask with FALSE along directions that - # stream into a boundary voxel. - solid_halo_list = [np.array(bc.indices).T for bc in self.BCs if bc.isSolid] - solid_halo_voxels = np.unique(np.vstack(solid_halo_list), axis=0) if solid_halo_list else None - - # Create the grid mask on each process - start = time.time() - grid_mask = self.create_grid_mask(solid_halo_voxels) - print("Time to create the grid mask:", time.time() - start) - - start = time.time() - for bc in self.BCs: - assert bc.implementationStep in ['PostStreaming', 'PostCollision'] - bc.create_local_mask_and_normal_arrays(grid_mask) - print("Time to create the local masks and normal arrays:", time.time() - start) - - # This is another non-JITed way of creating the distributed arrays. It is not used at the moment. - # def distributed_array_init(self, shape, type, init_val=None): - # sharding_dim = shape[0] // self.nDevices - # sharded_shape = (self.nDevices, sharding_dim, *shape[1:]) - # device_shape = sharded_shape[1:] - # arrays = [] - - # for d, index in self.sharding.addressable_devices_indices_map(sharded_shape).items(): - # jax.default_device = d - # if init_val is None: - # x = jnp.zeros(shape=device_shape, dtype=type) - # else: - # x = jnp.full(shape=device_shape, fill_value=init_val, dtype=type) - # arrays += [jax.device_put(x, d)] - # jax.default_device = jax.devices()[0] - # return jax.make_array_from_single_device_arrays(shape, self.sharding, arrays) - - @partial(jit, static_argnums=(0, 1, 2, 4)) - def distributed_array_init(self, shape, type, init_val=0, sharding=None): - """ - Initialize a distributed array using JAX, with a specified shape, data type, and initial value. - Optionally, provide a custom sharding strategy. - - Parameters - ---------- - shape (tuple): The shape of the array to be created. - type (dtype): The data type of the array to be created. - init_val (scalar, optional): The initial value to fill the array with. Defaults to 0. - sharding (Sharding, optional): The sharding strategy to use. Defaults to `self.sharding`. - - Returns - ------- - jax.numpy.ndarray: A JAX array with the specified shape, data type, initial value, and sharding strategy. - """ - if sharding is None: - sharding = self.sharding - x = jnp.full(shape=shape, fill_value=init_val, dtype=type) - return jax.lax.with_sharding_constraint(x, sharding) - - @partial(jit, static_argnums=(0,)) - def create_grid_mask(self, solid_halo_voxels): - """ - This function creates a mask for the background grid that accounts for the location of the boundaries. - - Parameters - ---------- - solid_halo_voxels: A numpy array representing the voxels in the halo of the solid object. - - Returns - ------- - A JAX array representing the grid mask of the grid. - """ - # Halo width (hw_x is different to accommodate the domain sharding per XLA device) - hw_x = self.nDevices - hw_y = hw_z = 1 - if self.dim == 2: - grid_mask = self.distributed_array_init((self.nx + 2 * hw_x, self.ny + 2 * hw_y, self.lattice.q), jnp.bool_, init_val=True) - grid_mask = grid_mask.at[(slice(hw_x, -hw_x), slice(hw_y, -hw_y), slice(None))].set(False) - if solid_halo_voxels is not None: - solid_halo_voxels = solid_halo_voxels.at[:, 0].add(hw_x) - solid_halo_voxels = solid_halo_voxels.at[:, 1].add(hw_y) - grid_mask = grid_mask.at[tuple(solid_halo_voxels.T)].set(True) - - grid_mask = self.streaming(grid_mask) - return lax.with_sharding_constraint(grid_mask, self.sharding) - - elif self.dim == 3: - grid_mask = self.distributed_array_init((self.nx + 2 * hw_x, self.ny + 2 * hw_y, self.nz + 2 * hw_z, self.lattice.q), jnp.bool_, init_val=True) - grid_mask = grid_mask.at[(slice(hw_x, -hw_x), slice(hw_y, -hw_y), slice(hw_z, -hw_z), slice(None))].set(False) - if solid_halo_voxels is not None: - solid_halo_voxels = solid_halo_voxels.at[:, 0].add(hw_x) - solid_halo_voxels = solid_halo_voxels.at[:, 1].add(hw_y) - solid_halo_voxels = solid_halo_voxels.at[:, 2].add(hw_z) - grid_mask = grid_mask.at[tuple(solid_halo_voxels.T)].set(True) - grid_mask = self.streaming(grid_mask) - return lax.with_sharding_constraint(grid_mask, self.sharding) - - def bounding_box_indices(self): - """ - This function calculates the indices of the bounding box of a 2D or 3D grid. - The bounding box is defined as the set of grid points on the outer edge of the grid. - - Returns - ------- - boundingBox (dict): A dictionary where keys are the names of the bounding box faces - ("bottom", "top", "left", "right" for 2D; additional "front", "back" for 3D), and values - are numpy arrays of indices corresponding to each face. - """ - if self.dim == 2: - # For a 2D grid, the bounding box consists of four edges: bottom, top, left, and right. - # Each edge is represented as an array of indices. For example, the bottom edge includes - # all points where the y-coordinate is 0, so its indices are [[i, 0] for i in range(self.nx)]. - bounding_box = {"bottom": np.array([[i, 0] for i in range(self.nx)], dtype=int), - "top": np.array([[i, self.ny - 1] for i in range(self.nx)], dtype=int), - "left": np.array([[0, i] for i in range(self.ny)], dtype=int), - "right": np.array([[self.nx - 1, i] for i in range(self.ny)], dtype=int)} - - return bounding_box - - elif self.dim == 3: - # For a 3D grid, the bounding box consists of six faces: bottom, top, left, right, front, and back. - # Each face is represented as an array of indices. For example, the bottom face includes all points - # where the z-coordinate is 0, so its indices are [[i, j, 0] for i in range(self.nx) for j in range(self.ny)]. - bounding_box = { - "bottom": np.array([[i, j, 0] for i in range(self.nx) for j in range(self.ny)], dtype=int), - "top": np.array([[i, j, self.nz - 1] for i in range(self.nx) for j in range(self.ny)],dtype=int), - "left": np.array([[0, j, k] for j in range(self.ny) for k in range(self.nz)], dtype=int), - "right": np.array([[self.nx - 1, j, k] for j in range(self.ny) for k in range(self.nz)], dtype=int), - "front": np.array([[i, 0, k] for i in range(self.nx) for k in range(self.nz)], dtype=int), - "back": np.array([[i, self.ny - 1, k] for i in range(self.nx) for k in range(self.nz)], dtype=int)} - - return bounding_box - - def set_precisions(self, precision): - """ - This function sets the precision of the computations. The precision is defined by a pair of values, - representing the precision of the computation and the precision of the storage, respectively. - - Parameters - ---------- - precision (str): A string representing the desired precision. The string should be in the format - "computation/storage", where "computation" and "storage" are either "f64", "f32", or "f16", - representing 64-bit, 32-bit, or 16-bit floating point numbers, respectively. - - Returns - ------- - tuple: A pair of jax.numpy data types representing the computation and storage precisions, respectively. - If the input string does not match any of the predefined options, the function defaults to (jnp.float32, jnp.float32). - """ - return { - "f64/f64": (jnp.float64, jnp.float64), - "f32/f32": (jnp.float32, jnp.float32), - "f32/f16": (jnp.float32, jnp.float16), - "f16/f16": (jnp.float16, jnp.float16), - "f64/f32": (jnp.float64, jnp.float32), - "f64/f16": (jnp.float64, jnp.float16), - }.get(precision, (jnp.float32, jnp.float32)) - - def initialize_macroscopic_fields(self): - """ - This function initializes the macroscopic fields (density and velocity) to their default values. - The default density is 1 and the default velocity is 0. - - Note: This function is a placeholder and should be overridden in a subclass or in an instance of the class - to provide specific initial conditions. - - Returns - ------- - None, None: The default density and velocity, both None. This indicates that the actual values should be set elsewhere. - """ - print("WARNING: Default initial conditions assumed: density = 1, velocity = 0") - print(" To set explicit initial density and velocity, use self.initialize_macroscopic_fields.") - return None, None - - def assign_fields_sharded(self): - """ - This function is used to initialize the simulation by assigning the macroscopic fields and populations. - - The function first initializes the macroscopic fields, which are the density (rho0) and velocity (u0). - Depending on the dimension of the simulation (2D or 3D), it then sets the shape of the array that will hold the - distribution functions (f). - - If the density or velocity are not provided, the function initializes the distribution functions with a default - value (self.w), representing density=1 and velocity=0. Otherwise, it uses the provided density and velocity to initialize the populations. - - Parameters - ---------- - None - - Returns - ------- - f: a distributed JAX array of shape (nx, ny, nz, q) or (nx, ny, q) holding the distribution functions for the simulation. - """ - rho0, u0 = self.initialize_macroscopic_fields() - - if self.dim == 2: - shape = (self.nx, self.ny, self.lattice.q) - if self.dim == 3: - shape = (self.nx, self.ny, self.nz, self.lattice.q) - - if rho0 is None or u0 is None: - f = self.distributed_array_init(shape, self.precisionPolicy.output_dtype, init_val=self.w) - else: - f = self.initialize_populations(rho0, u0) - - return f - - def initialize_populations(self, rho0, u0): - """ - This function initializes the populations (distribution functions) for the simulation. - It uses the equilibrium distribution function, which is a function of the macroscopic - density and velocity. - - Parameters - ---------- - rho0: jax.numpy.ndarray - The initial density field. - u0: jax.numpy.ndarray - The initial velocity field. - - Returns - ------- - f: jax.numpy.ndarray - The array holding the initialized distribution functions for the simulation. - """ - return self.equilibrium(rho0, u0) - - def send_right(self, x, axis_name): - """ - This function sends the data to the right neighboring process in a parallel computing environment. - It uses a permutation operation provided by the LAX library. - - Parameters - ---------- - x: jax.numpy.ndarray - The data to be sent. - axis_name: str - The name of the axis along which the data is sent. - - Returns - ------- - jax.numpy.ndarray - The data after being sent to the right neighboring process. - """ - return lax.ppermute(x, perm=self.rightPerm, axis_name=axis_name) - - def send_left(self, x, axis_name): - """ - This function sends the data to the left neighboring process in a parallel computing environment. - It uses a permutation operation provided by the LAX library. - - Parameters - ---------- - x: jax.numpy.ndarray - The data to be sent. - axis_name: str - The name of the axis along which the data is sent. - - Returns - ------- - The data after being sent to the left neighboring process. - """ - return lax.ppermute(x, perm=self.leftPerm, axis_name=axis_name) - - def streaming_m(self, f): - """ - This function performs the streaming step in the Lattice Boltzmann Method, which is - the propagation of the distribution functions in the lattice. - - To enable multi-GPU/TPU functionality, it extracts the left and right boundary slices of the - distribution functions that need to be communicated to the neighboring processes. - - The function then sends the left boundary slice to the right neighboring process and the right - boundary slice to the left neighboring process. The received data is then set to the - corresponding indices in the receiving domain. - - Parameters - ---------- - f: jax.numpy.ndarray - The array holding the distribution functions for the simulation. - - Returns - ------- - jax.numpy.ndarray - The distribution functions after the streaming operation. - """ - f = self.streaming_p(f) - left_comm, right_comm = f[:1, ..., self.lattice.right_indices], f[-1:, ..., self.lattice.left_indices] - - left_comm, right_comm = self.send_right(left_comm, 'x'), self.send_left(right_comm, 'x') - f = f.at[:1, ..., self.lattice.right_indices].set(left_comm) - f = f.at[-1:, ..., self.lattice.left_indices].set(right_comm) - return f - - @partial(jit, static_argnums=(0,)) - def streaming_p(self, f): - """ - Perform streaming operation on a partitioned (in the x-direction) distribution function. - - The function uses the vmap operation provided by the JAX library to vectorize the computation - over all lattice directions. - - Parameters - ---------- - f: The distribution function. - - Returns - ------- - The updated distribution function after streaming. - """ - def streaming_i(f, c): - """ - Perform individual streaming operation in a direction. - - Parameters - ---------- - f: The distribution function. - c: The streaming direction vector. - - Returns - ------- - jax.numpy.ndarray - The updated distribution function after streaming. - """ - if self.dim == 2: - return jnp.roll(f, (c[0], c[1]), axis=(0, 1)) - elif self.dim == 3: - return jnp.roll(f, (c[0], c[1], c[2]), axis=(0, 1, 2)) - - return vmap(streaming_i, in_axes=(-1, 0), out_axes=-1)(f, self.c.T) - - @partial(jit, static_argnums=(0, 3), inline=True) - def equilibrium(self, rho, u, cast_output=True): - """ - This function computes the equilibrium distribution function in the Lattice Boltzmann Method. - The equilibrium distribution function is a function of the macroscopic density and velocity. - - The function first casts the density and velocity to the compute precision if the cast_output flag is True. - The function finally casts the equilibrium distribution function to the output precision if the cast_output - flag is True. - - Parameters - ---------- - rho: jax.numpy.ndarray - The macroscopic density. - u: jax.numpy.ndarray - The macroscopic velocity. - cast_output: bool, optional - A flag indicating whether to cast the density, velocity, and equilibrium distribution function to the - compute and output precisions. Default is True. - - Returns - ------- - feq: ja.numpy.ndarray - The equilibrium distribution function. - """ - # Cast the density and velocity to the compute precision if the cast_output flag is True - if cast_output: - rho, u = self.precisionPolicy.cast_to_compute((rho, u)) - - # Cast c to compute precision so that XLA call FXX matmul, - # which is faster (it is faster in some older versions of JAX, newer versions are smart enough to do this automatically) - c = jnp.array(self.c, dtype=self.precisionPolicy.compute_dtype) - cu = 3.0 * jnp.dot(u, c) - usqr = 1.5 * jnp.sum(jnp.square(u), axis=-1, keepdims=True) - feq = rho * self.w * (1.0 + cu * (1.0 + 0.5 * cu) - usqr) - - if cast_output: - return self.precisionPolicy.cast_to_output(feq) - else: - return feq - - @partial(jit, static_argnums=(0,)) - def momentum_flux(self, fneq): - """ - This function computes the momentum flux, which is the product of the non-equilibrium - distribution functions (fneq) and the lattice moments (cc). - - The momentum flux is used in the computation of the stress tensor in the Lattice Boltzmann - Method (LBM). - - Parameters - ---------- - fneq: jax.numpy.ndarray - The non-equilibrium distribution functions. - - Returns - ------- - jax.numpy.ndarray - The computed momentum flux. - """ - return jnp.dot(fneq, self.lattice.cc) - - @partial(jit, static_argnums=(0,), inline=True) - def update_macroscopic(self, f): - """ - This function computes the macroscopic variables (density and velocity) based on the - distribution functions (f). - - The density is computed as the sum of the distribution functions over all lattice directions. - The velocity is computed as the dot product of the distribution functions and the lattice - velocities, divided by the density. - - Parameters - ---------- - f: jax.numpy.ndarray - The distribution functions. - - Returns - ------- - rho: jax.numpy.ndarray - Computed density. - u: jax.numpy.ndarray - Computed velocity. - """ - rho =jnp.sum(f, axis=-1, keepdims=True) - c = jnp.array(self.c, dtype=self.precisionPolicy.compute_dtype).T - u = jnp.dot(f, c) / rho - - return rho, u - - @partial(jit, static_argnums=(0, 4), inline=True) - def apply_bc(self, fout, fin, timestep, implementation_step): - """ - This function applies the boundary conditions to the distribution functions. - - It iterates over all boundary conditions (BCs) and checks if the implementation step of the - boundary condition matches the provided implementation step. If it does, it applies the - boundary condition to the post-streaming distribution functions (fout). - - Parameters - ---------- - fout: jax.numpy.ndarray - The post-collision distribution functions. - fin: jax.numpy.ndarray - The post-streaming distribution functions. - implementation_step: str - The implementation step at which the boundary conditions should be applied. - - Returns - ------- - ja.numpy.ndarray - The output distribution functions after applying the boundary conditions. - """ - for bc in self.BCs: - fout = bc.prepare_populations(fout, fin, implementation_step) - if bc.implementationStep == implementation_step: - if bc.isDynamic: - fout = bc.apply(fout, fin, timestep) - else: - fout = fout.at[bc.indices].set(bc.apply(fout, fin)) - - return fout - - @partial(jit, static_argnums=(0, 3), donate_argnums=(1,)) - def step(self, f_poststreaming, timestep, return_fpost=False): - """ - This function performs a single step of the LBM simulation. - - It first performs the collision step, which is the relaxation of the distribution functions - towards their equilibrium values. It then applies the respective boundary conditions to the - post-collision distribution functions. - - The function then performs the streaming step, which is the propagation of the distribution - functions in the lattice. It then applies the respective boundary conditions to the post-streaming - distribution functions. - - Parameters - ---------- - f_poststreaming: jax.numpy.ndarray - The post-streaming distribution functions. - timestep: int - The current timestep of the simulation. - return_fpost: bool, optional - If True, the function also returns the post-collision distribution functions. - - Returns - ------- - f_poststreaming: jax.numpy.ndarray - The post-streaming distribution functions after the simulation step. - f_postcollision: jax.numpy.ndarray or None - The post-collision distribution functions after the simulation step, or None if - return_fpost is False. - """ - f_postcollision = self.collision(f_poststreaming) - f_postcollision = self.apply_bc(f_postcollision, f_poststreaming, timestep, "PostCollision") - f_poststreaming = self.streaming(f_postcollision) - f_poststreaming = self.apply_bc(f_poststreaming, f_postcollision, timestep, "PostStreaming") - - if return_fpost: - return f_poststreaming, f_postcollision - else: - return f_poststreaming, None - - def run(self, t_max): - """ - This function runs the LBM simulation for a specified number of time steps. - - It first initializes the distribution functions and then enters a loop where it performs the - simulation steps (collision, streaming, and boundary conditions) for each time step. - - The function can also print the progress of the simulation, save the simulation data, and - compute the performance of the simulation in million lattice updates per second (MLUPS). - - Parameters - ---------- - t_max: int - The total number of time steps to run the simulation. - Returns - ------- - f: jax.numpy.ndarray - The distribution functions after the simulation. - """ - f = self.assign_fields_sharded() - start_step = 0 - if self.restore_checkpoint: - latest_step = self.mngr.latest_step() - if latest_step is not None: # existing checkpoint present - # Assert that the checkpoint manager is not None - assert self.mngr is not None, "Checkpoint manager does not exist." - state = {'f': f} - shardings = jax.tree_map(lambda x: x.sharding, state) - restore_args = orb.checkpoint_utils.construct_restore_args(state, shardings) - try: - f = self.mngr.restore(latest_step, restore_kwargs={'restore_args': restore_args})['f'] - print(f"Restored checkpoint at step {latest_step}.") - except ValueError: - raise ValueError(f"Failed to restore checkpoint at step {latest_step}.") - - start_step = latest_step + 1 - if not (t_max > start_step): - raise ValueError(f"Simulation already exceeded maximum allowable steps (t_max = {t_max}). Consider increasing t_max.") - if self.computeMLUPS: - start = time.time() - # Loop over all time steps - for timestep in range(start_step, t_max + 1): - io_flag = self.ioRate > 0 and (timestep % self.ioRate == 0 or timestep == t_max) - print_iter_flag = self.printInfoRate> 0 and timestep % self.printInfoRate== 0 - checkpoint_flag = self.checkpointRate > 0 and timestep % self.checkpointRate == 0 - - if io_flag: - # Update the macroscopic variables and save the previous values (for error computation) - rho_prev, u_prev = self.update_macroscopic(f) - rho_prev = downsample_field(rho_prev, self.downsamplingFactor) - u_prev = downsample_field(u_prev, self.downsamplingFactor) - # Gather the data from all processes and convert it to numpy arrays (move to host memory) - rho_prev = process_allgather(rho_prev) - u_prev = process_allgather(u_prev) - - - # Perform one time-step (collision, streaming, and boundary conditions) - f, fstar = self.step(f, timestep, return_fpost=self.returnFpost) - # Print the progress of the simulation - if print_iter_flag: - print(colored("Timestep ", 'blue') + colored(f"{timestep}", 'green') + colored(" of ", 'blue') + colored(f"{t_max}", 'green') + colored(" completed", 'blue')) - - if io_flag: - # Save the simulation data - print(f"Saving data at timestep {timestep}/{t_max}") - rho, u = self.update_macroscopic(f) - rho = downsample_field(rho, self.downsamplingFactor) - u = downsample_field(u, self.downsamplingFactor) - - # Gather the data from all processes and convert it to numpy arrays (move to host memory) - rho = process_allgather(rho) - u = process_allgather(u) - - # Save the data - self.handle_io_timestep(timestep, f, fstar, rho, u, rho_prev, u_prev) - - if checkpoint_flag: - # Save the checkpoint - print(f"Saving checkpoint at timestep {timestep}/{t_max}") - state = {'f': f} - self.mngr.save(timestep, state) - - # Start the timer for the MLUPS computation after the first timestep (to remove compilation overhead) - if self.computeMLUPS and timestep == 1: - jax.block_until_ready(f) - start = time.time() - - if self.computeMLUPS: - # Compute and print the performance of the simulation in MLUPS - jax.block_until_ready(f) - end = time.time() - if self.dim == 2: - print(colored("Domain: ", 'blue') + colored(f"{self.nx} x {self.ny}", 'green') if self.dim == 2 else colored(f"{self.nx} x {self.ny} x {self.nz}", 'green')) - print(colored("Number of voxels: ", 'blue') + colored(f"{self.nx * self.ny}", 'green') if self.dim == 2 else colored(f"{self.nx * self.ny * self.nz}", 'green')) - print(colored("MLUPS: ", 'blue') + colored(f"{self.nx * self.ny * t_max / (end - start) / 1e6}", 'red')) - - elif self.dim == 3: - print(colored("Domain: ", 'blue') + colored(f"{self.nx} x {self.ny} x {self.nz}", 'green')) - print(colored("Number of voxels: ", 'blue') + colored(f"{self.nx * self.ny * self.nz}", 'green')) - print(colored("MLUPS: ", 'blue') + colored(f"{self.nx * self.ny * self.nz * t_max / (end - start) / 1e6}", 'red')) - - return f - - def handle_io_timestep(self, timestep, f, fstar, rho, u, rho_prev, u_prev): - """ - This function handles the input/output (I/O) operations at each time step of the simulation. - - It prepares the data to be saved and calls the output_data function, which can be overwritten - by the user to customize the I/O operations. - - Parameters - ---------- - timestep: int - The current time step of the simulation. - f: jax.numpy.ndarray - The post-streaming distribution functions at the current time step. - fstar: jax.numpy.ndarray - The post-collision distribution functions at the current time step. - rho: jax.numpy.ndarray - The density field at the current time step. - u: jax.numpy.ndarray - The velocity field at the current time step. - """ - kwargs = { - "timestep": timestep, - "rho": rho, - "rho_prev": rho_prev, - "u": u, - "u_prev": u_prev, - "f_poststreaming": f, - "f_postcollision": fstar - } - self.output_data(**kwargs) - - def output_data(self, **kwargs): - """ - This function is intended to be overwritten by the user to customize the input/output (I/O) - operations of the simulation. - - By default, it does nothing. When overwritten, it could save the simulation data to files, - display the simulation results in real time, send the data to another process for analysis, etc. - - Parameters - ---------- - **kwargs: dict - A dictionary containing the simulation data to be outputted. The keys are the names of the - data fields, and the values are the data fields themselves. - """ - pass - - def set_boundary_conditions(self): - """ - This function sets the boundary conditions for the simulation. - - It is intended to be overwritten by the user to specify the boundary conditions according to - the specific problem being solved. - - By default, it does nothing. When overwritten, it could set periodic boundaries, no-slip - boundaries, inflow/outflow boundaries, etc. - """ - pass - - @partial(jit, static_argnums=(0,), donate_argnums=(1,)) - def collision(self, fin): - """ - This function performs the collision step in the Lattice Boltzmann Method. - - It is intended to be overwritten by the user to specify the collision operator according to - the specific LBM model being used. - - By default, it does nothing. When overwritten, it could implement the BGK collision operator, - the MRT collision operator, etc. - - Parameters - ---------- - fin: jax.numpy.ndarray - The pre-collision distribution functions. - - Returns - ------- - fin: jax.numpy.ndarray - The post-collision distribution functions. - """ - pass - - def get_force(self): - """ - This function computes the force to be applied to the fluid in the Lattice Boltzmann Method. - - It is intended to be overwritten by the user to specify the force according to the specific - problem being solved. - - By default, it does nothing and returns None. When overwritten, it could implement a constant - force term. - - Returns - ------- - force: jax.numpy.ndarray - The force to be applied to the fluid. - """ - pass - - @partial(jit, static_argnums=(0,), inline=True) - def apply_force(self, f_postcollision, feq, rho, u): - """ - add force based on exact-difference method due to Kupershtokh - - Parameters - ---------- - f_postcollision: jax.numpy.ndarray - The post-collision distribution functions. - feq: jax.numpy.ndarray - The equilibrium distribution functions. - rho: jax.numpy.ndarray - The density field. - - u: jax.numpy.ndarray - The velocity field. - - Returns - ------- - f_postcollision: jax.numpy.ndarray - The post-collision distribution functions with the force applied. - - References - ---------- - Kupershtokh, A. (2004). New method of incorporating a body force term into the lattice Boltzmann equation. In - Proceedings of the 5th International EHD Workshop (pp. 241-246). University of Poitiers, Poitiers, France. - Chikatamarla, S. S., & Karlin, I. V. (2013). Entropic lattice Boltzmann method for turbulent flow simulations: - Boundary conditions. Physica A, 392, 1925-1930. - KrΓΌger, T., et al. (2017). The lattice Boltzmann method. Springer International Publishing, 10.978-3, 4-15. - """ - delta_u = self.get_force() - feq_force = self.equilibrium(rho, u + delta_u, cast_output=False) - f_postcollision = f_postcollision + feq_force - feq - return f_postcollision - - - diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py deleted file mode 100644 index 6d3352a3..00000000 --- a/src/boundary_conditions.py +++ /dev/null @@ -1,1178 +0,0 @@ -import jax.numpy as jnp -from jax import jit, device_count -from functools import partial -import numpy as np -class BoundaryCondition(object): - """ - Base class for boundary conditions in a LBM simulation. - - This class provides a general structure for implementing boundary conditions. It includes methods for preparing the - boundary attributes and for applying the boundary condition. Specific boundary conditions should be implemented as - subclasses of this class, with the `apply` method overridden as necessary. - - Attributes - ---------- - lattice : Lattice - The lattice used in the simulation. - nx: - The number of nodes in the x direction. - ny: - The number of nodes in the y direction. - nz: - The number of nodes in the z direction. - dim : int - The number of dimensions in the simulation (2 or 3). - precision_policy : PrecisionPolicy - The precision policy used in the simulation. - indices : array-like - The indices of the boundary nodes. - name : str or None - The name of the boundary condition. This should be set in subclasses. - isSolid : bool - Whether the boundary condition is for a solid boundary. This should be set in subclasses. - isDynamic : bool - Whether the boundary condition is dynamic (changes over time). This should be set in subclasses. - needsExtraConfiguration : bool - Whether the boundary condition requires extra configuration. This should be set in subclasses. - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. This should be set in subclasses. - """ - - def __init__(self, indices, gridInfo, precision_policy): - self.lattice = gridInfo["lattice"] - self.nx = gridInfo["nx"] - self.ny = gridInfo["ny"] - self.nz = gridInfo["nz"] - self.dim = gridInfo["dim"] - self.precisionPolicy = precision_policy - self.indices = indices - self.name = None - self.isSolid = False - self.isDynamic = False - self.needsExtraConfiguration = False - self.implementationStep = "PostStreaming" - - def create_local_mask_and_normal_arrays(self, grid_mask): - - """ - Creates local mask and normal arrays for the boundary condition. - - Parameters - ---------- - grid_mask : array-like - The grid mask for the lattice. - - Returns - ------- - None - - Notes - ----- - This method creates local mask and normal arrays for the boundary condition based on the grid mask. - If the boundary condition requires extra configuration, the `configure` method is called. - """ - - if self.needsExtraConfiguration: - boundaryMask = self.get_boundary_mask(grid_mask) - self.configure(boundaryMask) - self.needsExtraConfiguration = False - - boundaryMask = self.get_boundary_mask(grid_mask) - self.normals = self.get_normals(boundaryMask) - self.imissing, self.iknown = self.get_missing_indices(boundaryMask) - self.imissingMask, self.iknownMask, self.imiddleMask = self.get_missing_mask(boundaryMask) - - return - - def get_boundary_mask(self, grid_mask): - """ - Add jax.device_count() to the self.indices in x-direction, and 1 to the self.indices other directions - This is to make sure the boundary condition is applied to the correct nodes as grid_mask is - expanded by (jax.device_count(), 1, 1) - - Parameters - ---------- - grid_mask : array-like - The grid mask for the lattice. - - Returns - ------- - boundaryMask : array-like - """ - shifted_indices = np.array(self.indices) - shifted_indices[0] += device_count() - shifted_indices[1:] += 1 - # Convert back to tuple - shifted_indices = tuple(shifted_indices) - boundaryMask = np.array(grid_mask[shifted_indices]) - - return boundaryMask - - def configure(self, boundaryMask): - """ - Configures the boundary condition. - - Parameters - ---------- - boundaryMask : array-like - The grid mask for the boundary voxels. - - Returns - ------- - None - - Notes - ----- - This method should be overridden in subclasses if the boundary condition requires extra configuration. - """ - return - - @partial(jit, static_argnums=(0, 3), inline=True) - def prepare_populations(self, fout, fin, implementation_step): - """ - Prepares the distribution functions for the boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The incoming distribution functions. - fin : jax.numpy.ndarray - The outgoing distribution functions. - implementation_step : str - The step in the lattice Boltzmann method algorithm at which the preparation is applied. - - Returns - ------- - jax.numpy.ndarray - The prepared distribution functions. - - Notes - ----- - This method should be overridden in subclasses if the boundary condition requires preparation of the distribution functions during post-collision or post-streaming. See ExtrapolationBoundaryCondition for an example. - """ - return fout - - def get_normals(self, boundaryMask): - """ - Calculates the normal vectors at the boundary nodes. - - Parameters - ---------- - boundaryMask : array-like - The boundary mask for the lattice. - - Returns - ------- - array-like - The normal vectors at the boundary nodes. - - Notes - ----- - This method calculates the normal vectors by dotting the boundary mask with the main lattice directions. - """ - main_c = self.lattice.c.T[self.lattice.main_indices] - m = boundaryMask[..., self.lattice.main_indices] - normals = -np.dot(m, main_c) - return normals - - def get_missing_indices(self, boundaryMask): - """ - Returns two int8 arrays the same shape as boundaryMask. The non-zero entries of these arrays indicate missing - directions that require BCs (imissing) as well as their corresponding opposite directions (iknown). - - Parameters - ---------- - boundaryMask : array-like - The boundary mask for the lattice. - - Returns - ------- - tuple of array-like - The missing and known indices for the boundary condition. - - Notes - ----- - This method calculates the missing and known indices based on the boundary mask. The missing indices are the - non-zero entries of the boundary mask, and the known indices are their corresponding opposite directions. - """ - - # Find imissing, iknown 1-to-1 corresponding indices - # Note: the "zero" index is used as default value here and won't affect BC computations - nbd = len(self.indices[0]) - imissing = np.vstack([np.arange(self.lattice.q, dtype='uint8')] * nbd) - iknown = np.vstack([self.lattice.opp_indices] * nbd) - imissing[~boundaryMask] = 0 - iknown[~boundaryMask] = 0 - return imissing, iknown - - def get_missing_mask(self, boundaryMask): - """ - Returns three boolean arrays the same shape as boundaryMask. - Note: these boundary masks are useful for reduction (eg. summation) operators of selected q-directions. - - Parameters - ---------- - boundaryMask : array-like - The boundary mask for the lattice. - - Returns - ------- - tuple of array-like - The missing, known, and middle masks for the boundary condition. - - Notes - ----- - This method calculates the missing, known, and middle masks based on the boundary mask. The missing mask - is the boundary mask, the known mask is the opposite directions of the missing mask, and the middle mask - is the directions that are neither missing nor known. - """ - # Find masks for imissing, iknown and imiddle - imissingMask = boundaryMask - iknownMask = imissingMask[:, self.lattice.opp_indices] - imiddleMask = ~(imissingMask | iknownMask) - return imissingMask, iknownMask, imiddleMask - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - None - - Notes - ----- - This method should be overridden in subclasses to implement the specific boundary condition. The method should - modify the output distribution functions in place to apply the boundary condition. - """ - pass - - @partial(jit, static_argnums=(0,)) - def equilibrium(self, rho, u): - """ - Compute equilibrium distribution function. - - Parameters - ---------- - rho : jax.numpy.ndarray - The density at each node in the lattice. - u : jax.numpy.ndarray - The velocity at each node in the lattice. - - Returns - ------- - jax.numpy.ndarray - The equilibrium distribution function at each node in the lattice. - - Notes - ----- - This method computes the equilibrium distribution function based on the density and velocity. The computation is - performed in the compute precision specified by the precision policy. The result is not cast to the output precision as - this is function is used inside other functions that require the compute precision. - """ - rho, u = self.precisionPolicy.cast_to_compute((rho, u)) - c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype) - cu = 3.0 * jnp.dot(u, c) - usqr = 1.5 * jnp.sum(u**2, axis=-1, keepdims=True) - feq = rho * self.lattice.w * (1.0 + 1.0 * cu + 0.5 * cu**2 - usqr) - - return feq - - @partial(jit, static_argnums=(0,)) - def momentum_flux(self, fneq): - """ - Compute the momentum flux. - - Parameters - ---------- - fneq : jax.numpy.ndarray - The non-equilibrium distribution function at each node in the lattice. - - Returns - ------- - jax.numpy.ndarray - The momentum flux at each node in the lattice. - - Notes - ----- - This method computes the momentum flux by dotting the non-equilibrium distribution function with the lattice - direction vectors. - """ - return jnp.dot(fneq, self.lattice.cc) - - @partial(jit, static_argnums=(0,)) - def momentum_exchange_force(self, f_poststreaming, f_postcollision): - """ - Using the momentum exchange method to compute the boundary force vector exerted on the solid geometry - based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a - bounce-back method (e.g. Bouzidi) that accounts for curved boundaries. - NOTE: this function should be called after BC's are imposed. - [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation. - Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339. - [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving - curved geometry, Phys. Rev. E 65 (2002) 041203. - [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange - algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423. - - Parameters - ---------- - f_poststreaming : jax.numpy.ndarray - The post-streaming distribution function at each node in the lattice. - f_postcollision : jax.numpy.ndarray - The post-collision distribution function at each node in the lattice. - - Returns - ------- - jax.numpy.ndarray - The force exerted on the solid geometry at each boundary node. - - Notes - ----- - This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. - The force is computed based on the post-streaming and post-collision distribution functions. This method - should be called after the boundary conditions are imposed. - """ - c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype) - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - phi = f_postcollision[self.indices][bindex, self.iknown] + \ - f_poststreaming[self.indices][bindex, self.imissing] - force = jnp.sum(c[:, self.iknown] * phi, axis=-1).T - return force - -class BounceBack(BoundaryCondition): - """ - Bounce-back boundary condition for a lattice Boltzmann method simulation. - - This class implements a full-way bounce-back boundary condition, where particles hitting the boundary are reflected - back in the direction they came from. The boundary condition is applied after the collision step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "BounceBackFullway". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostCollision". - """ - def __init__(self, indices, gridInfo, precision_policy): - super().__init__(indices, gridInfo, precision_policy) - self.name = "BounceBackFullway" - self.implementationStep = "PostCollision" - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the bounce-back boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - - Notes - ----- - This method applies the bounce-back boundary condition by reflecting the input distribution functions at the - boundary nodes in the opposite direction. - """ - - return fin[self.indices][..., self.lattice.opp_indices] - -class BounceBackMoving(BoundaryCondition): - """ - Moving bounce-back boundary condition for a lattice Boltzmann method simulation. - - This class implements a moving bounce-back boundary condition, where particles hitting the boundary are reflected - back in the direction they came from, with an additional velocity due to the movement of the boundary. The boundary - condition is applied after the collision step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "BounceBackFullwayMoving". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostCollision". - isDynamic : bool - Whether the boundary condition is dynamic (changes over time). For this class, it is True. - update_function : function - A function that updates the boundary condition. For this class, it is a function that updates the boundary - condition based on the current time step. The signature of the function is `update_function(time) -> (indices, vel)`, - - """ - def __init__(self, gridInfo, precision_policy, update_function=None): - # We get the indices at time zero to pass to the parent class for initialization - indices, _ = update_function(0) - super().__init__(indices, gridInfo, precision_policy) - self.name = "BounceBackFullwayMoving" - self.implementationStep = "PostCollision" - self.isDynamic = True - self.update_function = jit(update_function) - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin, time): - """ - Applies the moving bounce-back boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - time : int - The current time step. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - """ - indices, vel = self.update_function(time) - c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype) - cu = 6.0 * self.lattice.w * jnp.dot(vel, c) - return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu) - - -class BounceBackHalfway(BoundaryCondition): - """ - Halfway bounce-back boundary condition for a lattice Boltzmann method simulation. - - This class implements a halfway bounce-back boundary condition. The boundary condition is applied after - the streaming step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "BounceBackHalfway". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostStreaming". - needsExtraConfiguration : bool - Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True. - isSolid : bool - Whether the boundary condition represents a solid boundary. For this class, it is True. - vel : array-like - The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default). - """ - def __init__(self, indices, gridInfo, precision_policy, vel=None): - super().__init__(indices, gridInfo, precision_policy) - self.name = "BounceBackHalfway" - self.implementationStep = "PostStreaming" - self.needsExtraConfiguration = True - self.isSolid = True - self.vel = vel - - def configure(self, boundaryMask): - """ - Configures the boundary condition. - - Parameters - ---------- - boundaryMask : array-like - The grid mask for the boundary voxels. - - Returns - ------- - None - - Notes - ----- - This method performs an index shift for the halfway bounce-back boundary condition. It updates the indices of - the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes. - """ - # Perform index shift for halfway BB. - hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices] - nbd_orig = len(self.indices[0]) - idx = np.array(self.indices).T - idx_trg = [] - for i in range(self.lattice.q): - idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i]) - indices_new = np.unique(np.vstack(idx_trg), axis=0) - self.indices = tuple(indices_new.T) - nbd_modified = len(self.indices[0]) - if (nbd_orig != nbd_modified) and self.vel is not None: - vel_avg = np.mean(self.vel, axis=0) - self.vel = jnp.zeros(indices_new.shape, dtype=self.precisionPolicy.compute_dtype) + vel_avg - print("WARNING: assuming a constant averaged velocity vector is imposed at all BC cells!") - - return - - @partial(jit, static_argnums=(0,)) - def impose_boundary_vel(self, fbd, bindex): - c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype) - cu = 6.0 * self.lattice.w * jnp.dot(self.vel, c) - fbd = fbd.at[bindex, self.imissing].add(-cu[bindex, self.iknown]) - return fbd - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the halfway bounce-back boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - """ - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fbd = fout[self.indices] - - fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown]) - if self.vel is not None: - fbd = self.impose_boundary_vel(fbd, bindex) - return fbd - -class EquilibriumBC(BoundaryCondition): - """ - Equilibrium boundary condition for a lattice Boltzmann method simulation. - - This class implements an equilibrium boundary condition, where the distribution function at the boundary nodes is - set to the equilibrium distribution function. The boundary condition is applied after the streaming step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "EquilibriumBC". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostStreaming". - out : jax.numpy.ndarray - The equilibrium distribution function at the boundary nodes. - """ - - def __init__(self, indices, gridInfo, precision_policy, rho, u): - super().__init__(indices, gridInfo, precision_policy) - self.out = self.precisionPolicy.cast_to_output(self.equilibrium(rho, u)) - self.name = "EquilibriumBC" - self.implementationStep = "PostStreaming" - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the equilibrium boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - - Notes - ----- - This method applies the equilibrium boundary condition by setting the output distribution functions at the - boundary nodes to the equilibrium distribution function. - """ - return self.out - -class DoNothing(BoundaryCondition): - def __init__(self, indices, gridInfo, precision_policy): - """ - Do-nothing boundary condition for a lattice Boltzmann method simulation. - - This class implements a do-nothing boundary condition, where no action is taken at the boundary nodes. The boundary - condition is applied after the streaming step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "DoNothing". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostStreaming". - - Notes - ----- - This boundary condition enforces skipping of streaming altogether as it sets post-streaming equal to post-collision - populations (so no streaming at this BC voxels). The problem with returning post-streaming values or "fout[self.indices] - is that the information that exit the domain on the opposite side of this boundary, would "re-enter". This is because - we roll the entire array and so the boundary condition acts like a one-way periodic BC. If EquilibriumBC is used as - the BC for that opposite boundary, then the rolled-in values are taken from the initial condition at equilibrium. - Otherwise if ZouHe is used for example the simulation looks like a run-down simulation at low-Re. The opposite boundary - may be even a wall (consider pipebend example). If we correct imissing directions and assign "fin", this method becomes - much less stable and also one needs to correctly take care of corner cases. - """ - super().__init__(indices, gridInfo, precision_policy) - self.name = "DoNothing" - self.implementationStep = "PostStreaming" - - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the do-nothing boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - - Notes - ----- - This method applies the do-nothing boundary condition by simply returning the input distribution functions at the - boundary nodes. - """ - return fin[self.indices] - - -class ZouHe(BoundaryCondition): - """ - Zou-He boundary condition for a lattice Boltzmann method simulation. - - This class implements the Zou-He boundary condition, which is a non-equilibrium bounce-back boundary condition. - It can be used to set inflow and outflow boundary conditions with prescribed pressure or velocity. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "ZouHe". - implementationStep : str - The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, - it is "PostStreaming". - type : str - The type of the boundary condition. It can be either 'velocity' for a prescribed velocity boundary condition, - or 'pressure' for a prescribed pressure boundary condition. - prescribed : float or array-like - The prescribed values for the boundary condition. It can be either the prescribed velocities for a 'velocity' - boundary condition, or the prescribed pressures for a 'pressure' boundary condition. - - References - ---------- - Zou, Q., & He, X. (1997). On pressure and velocity boundary conditions for the lattice Boltzmann BGK model. - Physics of Fluids, 9(6), 1591-1598. doi:10.1063/1.869307 - """ - def __init__(self, indices, gridInfo, precision_policy, type, prescribed): - super().__init__(indices, gridInfo, precision_policy) - self.name = "ZouHe" - self.implementationStep = "PostStreaming" - self.type = type - self.prescribed = prescribed - self.needsExtraConfiguration = True - - def configure(self, boundaryMask): - """ - Correct boundary indices to ensure that only voxelized surfaces with normal vectors along main cartesian axes - are assigned this type of BC. - """ - nv = np.dot(self.lattice.c, ~boundaryMask.T) - corner_voxels = np.count_nonzero(nv, axis=0) > 1 - # removed_voxels = np.array(self.indices)[:, corner_voxels] - self.indices = tuple(np.array(self.indices)[:, ~corner_voxels]) - self.prescribed = self.prescribed[~corner_voxels] - return - - @partial(jit, static_argnums=(0,), inline=True) - def calculate_vel(self, fpop, rho): - """ - Calculate velocity based on the prescribed pressure/density (Zou/He BC) - """ - unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) + - 2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True)) - - # Return the above unormal as a normal vector which sets the tangential velocities to zero - vel = unormal * self.normals - return vel - - @partial(jit, static_argnums=(0,), inline=True) - def calculate_rho(self, fpop, vel): - """ - Calculate density based on the prescribed velocity (Zou/He BC) - """ - unormal = np.sum(self.normals*vel, axis=1) - - rho = (1.0/(1.0 + unormal))[..., None] * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) + - 2.*jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True)) - return rho - - @partial(jit, static_argnums=(0,), inline=True) - def calculate_equilibrium(self, fpop): - """ - This is the ZouHe method of calculating the missing macroscopic variables at the boundary. - """ - if self.type == 'velocity': - vel = self.prescribed - rho = self.calculate_rho(fpop, vel) - elif self.type == 'pressure': - rho = self.prescribed - vel = self.calculate_vel(fpop, rho) - else: - raise ValueError(f"type = {self.type} not supported! Use \'pressure\' or \'velocity\'.") - - # compute feq at the boundary - feq = self.equilibrium(rho, vel) - return feq - - @partial(jit, static_argnums=(0,), inline=True) - def bounceback_nonequilibrium(self, fpop, feq): - """ - Calculate unknown populations using bounce-back of non-equilibrium populations - a la original Zou & He formulation - """ - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fbd = fpop[self.indices] - fknown = fpop[self.indices][bindex, self.iknown] + feq[bindex, self.imissing] - feq[bindex, self.iknown] - fbd = fbd.at[bindex, self.imissing].set(fknown) - return fbd - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, _): - """ - Applies the Zou-He boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - _ : jax.numpy.ndarray - The input distribution functions. This is not used in this method. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - - Notes - ----- - This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based - on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions - based on the non-equilibrium bounce-back method. - Tangential velocity is not ensured to be zero by adding transverse contributions based on - Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher - Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds. - """ - # compute the equilibrium based on prescribed values and the type of BC - feq = self.calculate_equilibrium(fout) - - # set the unknown f populations based on the non-equilibrium bounce-back method - fbd = self.bounceback_nonequilibrium(fout, feq) - - - return fbd - -class Regularized(ZouHe): - """ - Regularized boundary condition for a lattice Boltzmann method simulation. - - This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition - with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure - or velocity. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "Regularized". - Qi : numpy.ndarray - The Qi tensor, which is used in the regularization of the distribution functions. - - References - ---------- - Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva. - Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the - lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703 - """ - - def __init__(self, indices, gridInfo, precision_policy, type, prescribed): - super().__init__(indices, gridInfo, precision_policy, type, prescribed) - self.name = "Regularized" - #TODO for Hesam: check to understand why corner cases cause instability here. - # self.needsExtraConfiguration = False - self.construct_symmetric_lattice_moment() - - def construct_symmetric_lattice_moment(self): - """ - Construct the symmetric lattice moment Qi. - - The Qi tensor is used in the regularization of the distribution functions. It is defined as Qi = cc - cs^2*I, - where cc is the tensor of lattice velocities, cs is the speed of sound, and I is the identity tensor. - """ - Qi = self.lattice.cc - if self.dim == 3: - diagonal = (0, 3, 5) - offdiagonal = (1, 2, 4) - elif self.dim == 2: - diagonal = (0, 2) - offdiagonal = (1,) - else: - raise ValueError(f"dim = {self.dim} not supported") - - # Qi = cc - cs^2*I - Qi = Qi.at[:, diagonal].set(self.lattice.cc[:, diagonal] - 1./3.) - - # multiply off-diagonal elements by 2 because the Q tensor is symmetric - Qi = Qi.at[:, offdiagonal].set(self.lattice.cc[:, offdiagonal] * 2.0) - - self.Qi = Qi.T - return - - @partial(jit, static_argnums=(0,), inline=True) - def regularize_fpop(self, fpop, feq): - """ - Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop. - - Parameters - ---------- - fpop : jax.numpy.ndarray - The distribution functions. - feq : jax.numpy.ndarray - The equilibrium distribution functions. - - Returns - ------- - jax.numpy.ndarray - The regularized distribution functions. - """ - - # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq} - f_neq = fpop - feq - PiNeq = self.momentum_flux(f_neq) - # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq) - - # Compute double dot product Qi:Pi1 - # QiPi1 = np.zeros_like(fpop) - # Pi1 = PiNeq - # QiPi1 = jnp.dot(Qi, Pi1) - QiPi1 = jnp.dot(PiNeq, self.Qi) - - # assign all populations based on eq 45 of Latt et al (2008) - # fneq ~ f^1 - fpop1 = 9. / 2. * self.lattice.w[None, :] * QiPi1 - fpop_regularized = feq + fpop1 - - return fpop_regularized - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, _): - """ - Applies the regularized boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - _ : jax.numpy.ndarray - The input distribution functions. This is not used in this method. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - - Notes - ----- - This method applies the regularized boundary condition by first computing the equilibrium distribution functions based - on the prescribed values and the type of boundary condition, then setting the unknown distribution functions - based on the non-equilibrium bounce-back method, and finally regularizing the distribution functions. - """ - - # compute the equilibrium based on prescribed values and the type of BC - feq = self.calculate_equilibrium(fout) - - # set the unknown f populations based on the non-equilibrium bounce-back method - fbd = self.bounceback_nonequilibrium(fout, feq) - - # Regularize the boundary fpop - fbd = self.regularize_fpop(fbd, feq) - return fbd - - -class ExtrapolationOutflow(BoundaryCondition): - """ - Extrapolation outflow boundary condition for a lattice Boltzmann method simulation. - - This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition - that uses extrapolation to avoid strong wave reflections. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "ExtrapolationOutflow". - sound_speed : float - The speed of sound in the simulation. - - References - ---------- - Geier, M., SchΓΆnherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three - dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547. - doi:10.1016/j.camwa.2015.05.001. - """ - - def __init__(self, indices, gridInfo, precision_policy): - super().__init__(indices, gridInfo, precision_policy) - self.name = "ExtrapolationOutflow" - self.needsExtraConfiguration = True - self.sound_speed = 1./jnp.sqrt(3.) - - def configure(self, boundaryMask): - """ - Configure the boundary condition by finding neighbouring voxel indices. - - Parameters - ---------- - boundaryMask : np.ndarray - The grid mask for the boundary voxels. - """ - hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices] - idx = np.array(self.indices).T - idx_trg = [] - for i in range(self.lattice.q): - idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i]) - indices_nbr = np.unique(np.vstack(idx_trg), axis=0) - self.indices_nbr = tuple(indices_nbr.T) - - return - - @partial(jit, static_argnums=(0, 3), inline=True) - def prepare_populations(self, fout, fin, implementation_step): - """ - Prepares the distribution functions for the boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The incoming distribution functions. - fin : jax.numpy.ndarray - The outgoing distribution functions. - implementation_step : str - The step in the lattice Boltzmann method algorithm at which the preparation is applied. - - Returns - ------- - jax.numpy.ndarray - The prepared distribution functions. - - Notes - ----- - Because this function is called "PostCollision", f_poststreaming refers to previous time step or t-1 - """ - f_postcollision = fout - f_poststreaming = fin - if implementation_step == 'PostStreaming': - return f_postcollision - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fps_bdr = f_poststreaming[self.indices] - fps_nbr = f_poststreaming[self.indices_nbr] - fpc_bdr = f_postcollision[self.indices] - fpop = fps_bdr[bindex, self.imissing] - fpop_neighbour = fps_nbr[bindex, self.imissing] - fpop_extrapolated = self.sound_speed * fpop_neighbour + (1. - self.sound_speed) * fpop - - # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data - fpc_bdr = fpc_bdr.at[bindex, self.iknown].set(fpop_extrapolated) - f_postcollision = f_postcollision.at[self.indices].set(fpc_bdr) - return f_postcollision - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the extrapolation outflow boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - """ - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fbd = fout[self.indices] - fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown]) - return fbd - - -class InterpolatedBounceBackBouzidi(BounceBackHalfway): - """ - A local single-node version of the interpolated bounce-back boundary condition due to Bouzidi for a lattice - Boltzmann method simulation. - - This class implements a interpolated bounce-back boundary condition. The boundary condition is applied after - the streaming step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "InterpolatedBounceBackBouzidi". - implicit_distances : array-like - An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls - weights : array-like - An array of shape (number_of_bc_cells, q) initialized as None and constructed using implicit_distances array - during runtime. These "weights" are associated with the fractional distance of fluid cell to the boundary - position defined as: weights(dir_i) = |x_fluid - x_boundary(dir_i)| / |x_fluid - x_solid(dir_i)|. - """ - - def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None): - - super().__init__(indices, grid_info, precision_policy, vel=vel) - self.name = "InterpolatedBounceBackBouzidi" - self.implicit_distances = implicit_distances - self.weights = None - - def set_proximity_ratio(self): - """ - Creates the interpolation data needed for the boundary condition. - - Returns - ------- - None. The function updates the object's weights attribute in place. - """ - epsilon = 1e-12 - nbd = len(self.indices[0]) - idx = np.array(self.indices).T - bindex = np.arange(nbd)[:, None] - weights = np.full((idx.shape[0], self.lattice.q), 0.5) - c = np.array(self.lattice.c) - sdf_f = self.implicit_distances[self.indices] - for q in range(1, self.lattice.q): - solid_indices = idx + c[:, q] - solid_indices_tuple = tuple(map(tuple, solid_indices.T)) - sdf_s = self.implicit_distances[solid_indices_tuple] - weights[:, q] = sdf_f / (sdf_f - sdf_s + epsilon) - self.weights = weights[bindex, self.iknown] - return - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the halfway bounce-back boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - """ - if self.weights is None: - self.set_proximity_ratio() - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fbd = fout[self.indices] - f_postcollision_iknown = fin[self.indices][bindex, self.iknown] - f_postcollision_imissing = fin[self.indices][bindex, self.imissing] - f_poststreaming_iknown = fout[self.indices][bindex, self.iknown] - - # if weights<0.5 - fs_near = 2. * self.weights * f_postcollision_iknown + \ - (1.0 - 2.0 * self.weights) * f_poststreaming_iknown - - # if weights>=0.5 - fs_far = 1.0 / (2. * self.weights) * f_postcollision_iknown + \ - (2.0 * self.weights - 1.0) / (2. * self.weights) * f_postcollision_imissing - - # combine near and far contributions - fmissing = jnp.where(self.weights < 0.5, fs_near, fs_far) - fbd = fbd.at[bindex, self.imissing].set(fmissing) - - if self.vel is not None: - fbd = self.impose_boundary_vel(fbd, bindex) - return fbd - - -class InterpolatedBounceBackDifferentiable(InterpolatedBounceBackBouzidi): - """ - A differentiable variant of the "InterpolatedBounceBackBouzidi" BC scheme. This BC is now differentiable at - self.weight = 0.5 unlike the original Bouzidi scheme which switches between 2 equations at weight=0.5. Refer to - [1] (their Appendix E) for more information. - - References - ---------- - [1] Geier, M., SchΓΆnherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three - dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547. - doi:10.1016/j.camwa.2015.05.001. - - - This class implements a interpolated bounce-back boundary condition. The boundary condition is applied after - the streaming step. - - Attributes - ---------- - name : str - The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable". - """ - - def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None): - - super().__init__(indices, implicit_distances, grid_info, precision_policy, vel=vel) - self.name = "InterpolatedBounceBackDifferentiable" - - - @partial(jit, static_argnums=(0,)) - def apply(self, fout, fin): - """ - Applies the halfway bounce-back boundary condition. - - Parameters - ---------- - fout : jax.numpy.ndarray - The output distribution functions. - fin : jax.numpy.ndarray - The input distribution functions. - - Returns - ------- - jax.numpy.ndarray - The modified output distribution functions after applying the boundary condition. - """ - if self.weights is None: - self.set_proximity_ratio() - nbd = len(self.indices[0]) - bindex = np.arange(nbd)[:, None] - fbd = fout[self.indices] - f_postcollision_iknown = fin[self.indices][bindex, self.iknown] - f_postcollision_imissing = fin[self.indices][bindex, self.imissing] - f_poststreaming_iknown = fout[self.indices][bindex, self.iknown] - fmissing = ((1. - self.weights) * f_poststreaming_iknown + - self.weights * (f_postcollision_imissing + f_postcollision_iknown)) / (1.0 + self.weights) - fbd = fbd.at[bindex, self.imissing].set(fmissing) - - if self.vel is not None: - fbd = self.impose_boundary_vel(fbd, bindex) - return fbd \ No newline at end of file diff --git a/src/lattice.py b/src/lattice.py deleted file mode 100644 index 788796bb..00000000 --- a/src/lattice.py +++ /dev/null @@ -1,281 +0,0 @@ -import re -import numpy as np -import jax.numpy as jnp - - -class Lattice(object): - """ - This class represents a lattice in the Lattice Boltzmann Method. - - It stores the properties of the lattice, including the dimensions, the number of - velocities, the velocity vectors, the weights, the moments, and the indices of the - opposite, main, right, and left velocities. - - The class also provides methods to construct these properties based on the name of the - lattice. - - Parameters - ---------- - name: str - The name of the lattice, which specifies the dimensions and the number of velocities. - For example, "D2Q9" represents a 2D lattice with 9 velocities. - precision: str, optional - The precision of the computations. It can be "f32/f32", "f32/f16", "f64/f64", - "f64/f32", or "f64/f16". The first part before the slash is the precision of the - computations, and the second part after the slash is the precision of the outputs. - """ - def __init__(self, name, precision="f32/f32") -> None: - self.name = name - dq = re.findall(r"\d+", name) - self.precision = precision - self.d = int(dq[0]) - self.q = int(dq[1]) - if precision == "f32/f32" or precision == "f32/f16": - self.precisionPolicy = jnp.float32 - elif precision == "f64/f64" or precision == "f64/f32" or precision == "f64/f16": - self.precisionPolicy = jnp.float64 - elif precision == "f16/f16": - self.precisionPolicy = jnp.float16 - else: - raise ValueError("precision not supported") - - # Construct the properties of the lattice - self.c = jnp.array(self.construct_lattice_velocity(), dtype=jnp.int8) - self.w = jnp.array(self.construct_lattice_weight(), dtype=self.precisionPolicy) - self.cc = jnp.array(self.construct_lattice_moment(), dtype=self.precisionPolicy) - self.opp_indices = jnp.array(self.construct_opposite_indices(), dtype=jnp.int8) - self.main_indices = jnp.array(self.construct_main_indices(), dtype=jnp.int8) - self.right_indices = np.array(self.construct_right_indices(), dtype=jnp.int8) - self.left_indices = np.array(self.construct_left_indices(), dtype=jnp.int8) - - def construct_opposite_indices(self): - """ - This function constructs the indices of the opposite velocities for each velocity. - - The opposite velocity of a velocity is the velocity that has the same magnitude but the - opposite direction. - - Returns - ------- - opposite: numpy.ndarray - The indices of the opposite velocities. - """ - c = self.c.T - opposite = np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)]) - return opposite - - def construct_right_indices(self): - """ - This function constructs the indices of the velocities that point in the positive - x-direction. - - Returns - ------- - numpy.ndarray - The indices of the right velocities. - """ - c = self.c.T - return np.nonzero(c[:, 0] == 1)[0] - - def construct_left_indices(self): - """ - This function constructs the indices of the velocities that point in the negative - x-direction. - - Returns - ------- - numpy.ndarray - The indices of the left velocities. - """ - c = self.c.T - return np.nonzero(c[:, 0] == -1)[0] - - def construct_main_indices(self): - """ - This function constructs the indices of the main velocities. - - The main velocities are the velocities that have a magnitude of 1 in lattice units. - - Returns - ------- - numpy.ndarray - The indices of the main velocities. - """ - c = self.c.T - if self.d == 2: - return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0] - - elif self.d == 3: - return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0] - - def construct_lattice_velocity(self): - """ - This function constructs the velocity vectors of the lattice. - - The velocity vectors are defined based on the name of the lattice. For example, for a D2Q9 - lattice, there are 9 velocities: (0,0), (1,0), (-1,0), (0,1), (0,-1), (1,1), (-1,-1), - (1,-1), and (-1,1). - - Returns - ------- - c.T: numpy.ndarray - The velocity vectors of the lattice. - """ - if self.name == "D2Q9": # D2Q9 - cx = [0, 0, 0, 1, -1, 1, -1, 1, -1] - cy = [0, 1, -1, 0, 1, -1, 0, 1, -1] - c = np.array(tuple(zip(cx, cy))) - elif self.name == "D3Q19": # D3Q19 - c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]] - c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5]) - elif self.name == "D3Q27": # D3Q27 - c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]] - # c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5]) - c = np.array(c) - else: - raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27") - - return c.T - - def construct_lattice_weight(self): - """ - This function constructs the weights of the lattice. - - The weights are defined based on the name of the lattice. For example, for a D2Q9 lattice, - the weights are 4/9 for the rest velocity, 1/9 for the main velocities, and 1/36 for the - diagonal velocities. - - Returns - ------- - w: numpy.ndarray - The weights of the lattice. - """ - # Get the transpose of the lattice vector - c = self.c.T - - # Initialize the weights to be 1/36 - w = 1.0 / 36.0 * np.ones(self.q) - - # Update the weights for 2D and 3D lattices - if self.name == "D2Q9": - w[np.linalg.norm(c, axis=1) < 1.1] = 1.0 / 9.0 - w[0] = 4.0 / 9.0 - elif self.name == "D3Q19": - w[np.linalg.norm(c, axis=1) < 1.1] = 2.0 / 36.0 - w[0] = 1.0 / 3.0 - elif self.name == "D3Q27": - cl = np.linalg.norm(c, axis=1) - w[np.isclose(cl, 1.0, atol=1e-8)] = 2.0 / 27.0 - w[(cl > 1) & (cl <= np.sqrt(2))] = 1.0 / 54.0 - w[(cl > np.sqrt(2)) & (cl <= np.sqrt(3))] = 1.0 / 216.0 - w[0] = 8.0 / 27.0 - else: - raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27") - - # Return the weights - return w - - def construct_lattice_moment(self): - """ - This function constructs the moments of the lattice. - - The moments are the products of the velocity vectors, which are used in the computation of - the equilibrium distribution functions and the collision operator in the Lattice Boltzmann - Method (LBM). - - Returns - ------- - cc: numpy.ndarray - The moments of the lattice. - """ - c = self.c.T - # Counter for the loop - cntr = 0 - - # nt: number of independent elements of a symmetric tensor - nt = self.d * (self.d + 1) // 2 - - cc = np.zeros((self.q, nt)) - for a in range(0, self.d): - for b in range(a, self.d): - cc[:, cntr] = c[:, a] * c[:, b] - cntr += 1 - - return cc - - def __str__(self): - return self.name - -class LatticeD2Q9(Lattice): - """ - Lattice class for 2D D2Q9 lattice. - - D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the - Lat tice Boltzmann Method for simulating fluid flows in two dimensions. - - Parameters - ---------- - precision: str, optional - The precision of the lattice. The default is "f32/f32" - """ - def __init__(self, precision="f32/f32"): - super().__init__("D2Q9", precision) - self._set_constants() - - def _set_constants(self): - self.cs = jnp.sqrt(3) / 3.0 - self.cs2 = 1.0 / 3.0 - self.inv_cs2 = 3.0 - self.i_s = jnp.asarray(list(range(9))) - self.im = 3 # Number of imiddles (includes center) - self.ik = 3 # Number of iknowns or iunknowns - - -class LatticeD3Q19(Lattice): - """ - Lattice class for 3D D3Q19 lattice. - - D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the - Lattice Boltzmann Method for simulating fluid flows in three dimensions. - - Parameters - ---------- - precision: str, optional - The precision of the lattice. The default is "f32/f32" - """ - def __init__(self, precision="f32/f32"): - super().__init__("D3Q19", precision) - self._set_constants() - - def _set_constants(self): - self.cs = jnp.sqrt(3) / 3.0 - self.cs2 = 1.0 / 3.0 - self.inv_cs2 = 3.0 - self.i_s = jnp.asarray(list(range(19)), dtype=jnp.int8) - - self.im = 9 # Number of imiddles (includes center) - self.ik = 5 # Number of iknowns or iunknowns - - -class LatticeD3Q27(Lattice): - """ - Lattice class for 3D D3Q27 lattice. - - D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the - Lattice Boltzmann Method for simulating fluid flows in three dimensions. - - Parameters - ---------- - precision: str, optional - The precision of the lattice. The default is "f32/f32" - """ - - def __init__(self, precision="f32/f32"): - super().__init__("D3Q27", precision) - self._set_constants() - - def _set_constants(self): - self.cs = jnp.sqrt(3) / 3.0 - self.cs2 = 1.0 / 3.0 - self.inv_cs2 = 3.0 - self.i_s = jnp.asarray(list(range(27)), dtype=jnp.int8) \ No newline at end of file diff --git a/src/models.py b/src/models.py deleted file mode 100644 index a0500c85..00000000 --- a/src/models.py +++ /dev/null @@ -1,260 +0,0 @@ -import jax.numpy as jnp -from jax import jit -from functools import partial -from src.base import LBMBase -""" -Collision operators are defined in this file for different models. -""" - -class BGKSim(LBMBase): - """ - BGK simulation class. - - This class implements the Bhatnagar-Gross-Krook (BGK) approximation for the collision step in the Lattice Boltzmann Method. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - @partial(jit, static_argnums=(0,), donate_argnums=(1,)) - def collision(self, f): - """ - BGK collision step for lattice. - - The collision step is where the main physics of the LBM is applied. In the BGK approximation, - the distribution function is relaxed towards the equilibrium distribution function. - """ - f = self.precisionPolicy.cast_to_compute(f) - rho, u = self.update_macroscopic(f) - feq = self.equilibrium(rho, u, cast_output=False) - fneq = f - feq - fout = f - self.omega * fneq - if self.force is not None: - fout = self.apply_force(fout, feq, rho, u) - return self.precisionPolicy.cast_to_output(fout) - -class KBCSim(LBMBase): - """ - KBC simulation class. - - This class implements the Karlin-BΓΆsch-Chikatamarla (KBC) model for the collision step in the Lattice Boltzmann Method. - """ - def __init__(self, **kwargs): - if kwargs.get('lattice').name != 'D3Q27' and kwargs.get('nz') > 0: - raise ValueError("KBC collision operator in 3D must only be used with D3Q27 lattice.") - super().__init__(**kwargs) - - @partial(jit, static_argnums=(0,), donate_argnums=(1,)) - def collision(self, f): - """ - KBC collision step for lattice. - """ - f = self.precisionPolicy.cast_to_compute(f) - tiny = 1e-32 - beta = self.omega * 0.5 - rho, u = self.update_macroscopic(f) - feq = self.equilibrium(rho, u, cast_output=False) - fneq = f - feq - if self.dim == 2: - deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0 - else: - deltaS = self.fdecompose_shear_d3q27(fneq) * rho - deltaH = fneq - deltaS - invBeta = 1.0 / beta - gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (tiny + self.entropic_scalar_product(deltaH, deltaH, feq)) - - fout = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH) - - # add external force - if self.force is not None: - fout = self.apply_force(fout, feq, rho, u) - return self.precisionPolicy.cast_to_output(fout) - - @partial(jit, static_argnums=(0,), donate_argnums=(1,)) - def collision_modified(self, f): - """ - Alternative KBC collision step for lattice. - Note: - At low Reynolds number the orignal KBC collision above produces inaccurate results because - it does not check for the entropy increase/decrease. The KBC stabalizations should only be - applied in principle to cells whose entropy decrease after a regular BGK collision. This is - the case in most cells at higher Reynolds numbers and hence a check may not be needed. - Overall the following alternative collision is more reliable and may replace the original - implementation. The issue at the moment is that it is about 60-80% slower than the above method. - """ - f = self.precisionPolicy.cast_to_compute(f) - tiny = 1e-32 - beta = self.omega * 0.5 - rho, u = self.update_macroscopic(f) - feq = self.equilibrium(rho, u, castOutput=False) - - # Alternative KBC: only stabalizes for voxels whose entropy decreases after BGK collision. - f_bgk = f - self.omega * (f - feq) - H_fin = jnp.sum(f * jnp.log(f / self.w), axis=-1, keepdims=True) - H_fout = jnp.sum(f_bgk * jnp.log(f_bgk / self.w), axis=-1, keepdims=True) - - # the rest is identical to collision_deprecated - fneq = f - feq - if self.dim == 2: - deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0 - else: - deltaS = self.fdecompose_shear_d3q27(fneq) * rho - deltaH = fneq - deltaS - invBeta = 1.0 / beta - gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (tiny + self.entropic_scalar_product(deltaH, deltaH, feq)) - - f_kbc = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH) - fout = jnp.where(H_fout > H_fin, f_kbc, f_bgk) - - # add external force - if self.force is not None: - fout = self.apply_force(fout, feq, rho, u) - return self.precisionPolicy.cast_to_output(fout) - - @partial(jit, static_argnums=(0,), inline=True) - def entropic_scalar_product(self, x, y, feq): - """ - Compute the entropic scalar product of x and y to approximate gamma in KBC. - - Returns - ------- - jax.numpy.array - Entropic scalar product of x, y, and feq. - """ - return jnp.sum(x * y / feq, axis=-1) - - @partial(jit, static_argnums=(0,), inline=True) - def fdecompose_shear_d2q9(self, fneq): - """ - Decompose fneq into shear components for D2Q9 lattice. - - Parameters - ---------- - fneq : jax.numpy.array - Non-equilibrium distribution function. - - Returns - ------- - jax.numpy.array - Shear components of fneq. - """ - Pi = self.momentum_flux(fneq) - N = Pi[..., 0] - Pi[..., 2] - s = jnp.zeros_like(fneq) - s = s.at[..., 6].set(N) - s = s.at[..., 3].set(N) - s = s.at[..., 2].set(-N) - s = s.at[..., 1].set(-N) - s = s.at[..., 8].set(Pi[..., 1]) - s = s.at[..., 4].set(-Pi[..., 1]) - s = s.at[..., 5].set(-Pi[..., 1]) - s = s.at[..., 7].set(Pi[..., 1]) - - return s - - @partial(jit, static_argnums=(0,), inline=True) - def fdecompose_shear_d3q27(self, fneq): - """ - Decompose fneq into shear components for D3Q27 lattice. - - Parameters - ---------- - fneq : jax.numpy.ndarray - Non-equilibrium distribution function. - - Returns - ------- - jax.numpy.ndarray - Shear components of fneq. - """ - # if self.grid.dim == 3: - # diagonal = (0, 3, 5) - # offdiagonal = (1, 2, 4) - # elif self.grid.dim == 2: - # diagonal = (0, 2) - # offdiagonal = (1,) - - # c= - # array([[0, 0, 0],-----0 - # [0, 0, -1],----1 - # [0, 0, 1],-----2 - # [0, -1, 0],----3 - # [0, -1, -1],---4 - # [0, -1, 1],----5 - # [0, 1, 0],-----6 - # [0, 1, -1],----7 - # [0, 1, 1],-----8 - # [-1, 0, 0],----9 - # [-1, 0, -1],--10 - # [-1, 0, 1],---11 - # [-1, -1, 0],--12 - # [-1, -1, -1],-13 - # [-1, -1, 1],--14 - # [-1, 1, 0],---15 - # [-1, 1, -1],--16 - # [-1, 1, 1],---17 - # [1, 0, 0],----18 - # [1, 0, -1],---19 - # [1, 0, 1],----20 - # [1, -1, 0],---21 - # [1, -1, -1],--22 - # [1, -1, 1],---23 - # [1, 1, 0],----24 - # [1, 1, -1],---25 - # [1, 1, 1]])---26 - Pi = self.momentum_flux(fneq) - Nxz = Pi[..., 0] - Pi[..., 5] - Nyz = Pi[..., 3] - Pi[..., 5] - - # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k) - s = jnp.zeros_like(fneq) - s = s.at[..., 9].set((2.0 * Nxz - Nyz) / 6.0) - s = s.at[..., 18].set((2.0 * Nxz - Nyz) / 6.0) - s = s.at[..., 3].set((-Nxz + 2.0 * Nyz) / 6.0) - s = s.at[..., 6].set((-Nxz + 2.0 * Nyz) / 6.0) - s = s.at[..., 1].set((-Nxz - Nyz) / 6.0) - s = s.at[..., 2].set((-Nxz - Nyz) / 6.0) - - # For c = (i, j, 0) - s = s.at[..., 12].set(Pi[..., 1] / 4.0) - s = s.at[..., 24].set(Pi[..., 1] / 4.0) - s = s.at[..., 21].set(-Pi[..., 1] / 4.0) - s = s.at[..., 15].set(-Pi[..., 1] / 4.0) - - # For c = (i, 0, k) - s = s.at[..., 10].set(Pi[..., 2] / 4.0) - s = s.at[..., 20].set(Pi[..., 2] / 4.0) - s = s.at[..., 19].set(-Pi[..., 2] / 4.0) - s = s.at[..., 11].set(-Pi[..., 2] / 4.0) - - # For c = (0, j, k) - s = s.at[..., 8].set(Pi[..., 4] / 4.0) - s = s.at[..., 4].set(Pi[..., 4] / 4.0) - s = s.at[..., 7].set(-Pi[..., 4] / 4.0) - s = s.at[..., 5].set(-Pi[..., 4] / 4.0) - - return s - - -class AdvectionDiffusionBGK(LBMBase): - """ - Advection Diffusion Model based on the BGK model. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.vel = kwargs.get("vel", None) - if self.vel is None: - raise ValueError("Velocity must be specified for AdvectionDiffusionBGK.") - - @partial(jit, static_argnums=(0,), donate_argnums=(1,)) - def collision(self, f): - """ - BGK collision step for lattice. - """ - f = self.precisionPolicy.cast_to_compute(f) - rho =jnp.sum(f, axis=-1, keepdims=True) - feq = self.equilibrium(rho, self.vel, cast_output=False) - fneq = f - feq - fout = f - self.omega * fneq - return self.precisionPolicy.cast_to_output(fout) \ No newline at end of file diff --git a/src/utils.py b/src/utils.py deleted file mode 100644 index 1720ca60..00000000 --- a/src/utils.py +++ /dev/null @@ -1,414 +0,0 @@ -import numpy as np -import matplotlib.pylab as plt -from matplotlib import cm -import numpy as np -from time import time -import pyvista as pv -from jax.image import resize -from jax import jit -import jax.numpy as jnp -from functools import partial -import trimesh - -import os -import __main__ - - -@partial(jit, static_argnums=(1, 2)) -def downsample_field(field, factor, method='bicubic'): - """ - Downsample a JAX array by a factor of `factor` along each axis. - - Parameters - ---------- - field : jax.numpy.ndarray - The input vector field to be downsampled. This should be a 3D or 4D JAX array where the last dimension is 2 or 3 (vector components). - factor : int - The factor by which to downsample the field. The dimensions of the field will be divided by this factor. - method : str, optional - The method to use for downsampling. Default is 'bicubic'. - - Returns - ------- - jax.numpy.ndarray - The downsampled field. - """ - if factor == 1: - return field - else: - new_shape = tuple(dim // factor for dim in field.shape[:-1]) - downsampled_components = [] - for i in range(field.shape[-1]): # Iterate over the last dimension (vector components) - resized = resize(field[..., i], new_shape, method=method) - downsampled_components.append(resized) - - return jnp.stack(downsampled_components, axis=-1) - -def save_image(timestep, fld, prefix=None): - """ - Save an image of a field at a given timestep. - - Parameters - ---------- - timestep : int - The timestep at which the field is being saved. - fld : jax.numpy.ndarray - The field to be saved. This should be a 2D or 3D JAX array. If the field is 3D, the magnitude of the field will be calculated and saved. - prefix : str, optional - A prefix to be added to the filename. The filename will be the name of the main script file by default. - - Returns - ------- - None - - Notes - ----- - This function saves the field as an image in the PNG format. The filename is based on the name of the main script file, the provided prefix, and the timestep number. - If the field is 3D, the magnitude of the field is calculated and saved. The image is saved with the 'nipy_spectral' colormap and the origin set to 'lower'. - """ - fname = os.path.basename(__main__.__file__) - fname = os.path.splitext(fname)[0] - if prefix is not None: - fname = prefix + fname - fname = fname + "_" + str(timestep).zfill(4) - - if len(fld.shape) > 3: - raise ValueError("The input field should be 2D!") - elif len(fld.shape) == 3: - fld = np.sqrt(fld[..., 0] ** 2 + fld[..., 1] ** 2) - - plt.clf() - plt.imsave(fname + '.png', fld.T, cmap=cm.nipy_spectral, origin='lower') - -def save_fields_vtk(timestep, fields, output_dir='.', prefix='fields'): - """ - Save VTK fields to the specified directory. - - Parameters - ---------- - timestep (int): The timestep number to be associated with the saved fields. - fields (Dict[str, np.ndarray]): A dictionary of fields to be saved. Each field must be an array-like object - with dimensions (nx, ny) for 2D fields or (nx, ny, nz) for 3D fields, where: - - nx : int, number of grid points along the x-axis - - ny : int, number of grid points along the y-axis - - nz : int, number of grid points along the z-axis (for 3D fields only) - The key value for each field in the dictionary must be a string containing the name of the field. - output_dir (str, optional, default: '.'): The directory in which to save the VTK files. Defaults to the current directory. - prefix (str, optional, default: 'fields'): A prefix to be added to the filename. Defaults to 'fields'. - - Returns - ------- - None - - Notes - ----- - This function saves the VTK fields in the specified directory, with filenames based on the provided timestep number - and the filename. For example, if the timestep number is 10 and the file name is fields, the VTK file - will be saved as 'fields_0000010.vtk'in the specified directory. - - """ - # Assert that all fields have the same dimensions except for the last dimension assuming fields is a dictionary - for key, value in fields.items(): - if key == list(fields.keys())[0]: - dimensions = value.shape - else: - assert value.shape == dimensions, "All fields must have the same dimensions!" - - output_filename = os.path.join(output_dir, prefix + "_" + f"{timestep:07d}.vtk") - - # Add 1 to the dimensions tuple as we store cell values - dimensions = tuple([dim + 1 for dim in dimensions]) - - # Create a uniform grid - if value.ndim == 2: - dimensions = dimensions + (1,) - - grid = pv.ImageData(dimensions=dimensions) - - # Add the fields to the grid - for key, value in fields.items(): - grid[key] = value.flatten(order='F') - - # Save the grid to a VTK file - start = time() - grid.save(output_filename, binary=True) - print(f"Saved {output_filename} in {time() - start:.6f} seconds.") - -def live_volume_randering(timestep, field): - # WORK IN PROGRESS - """ - Live rendering of a 3D volume using pyvista. - - Parameters - ---------- - field (np.ndarray): A 3D array containing the field to be rendered. - - Returns - ------- - None - - Notes - ----- - This function uses pyvista to render a 3D volume. The volume is rendered with a colormap based on the field values. - The colormap is updated every 0.1 seconds to reflect changes to the field. - - """ - # Create a uniform grid (Note that the field must be 3D) otherwise raise error - if field.ndim != 3: - raise ValueError("The input field must be 3D!") - dimensions = field.shape - grid = pv.ImageData(dimensions=dimensions) - - # Add the field to the grid - grid['field'] = field.flatten(order='F') - - # Create the rendering scene - if timestep == 0: - plt.ion() - plt.figure(figsize=(10, 10)) - plt.axis('off') - plt.title("Live rendering of the field") - pl = pv.Plotter(off_screen=True) - pl.add_volume(grid, cmap='nipy_spectral', opacity='sigmoid_10', shade=False) - plt.imshow(pl.screenshot()) - - else: - pl = pv.Plotter(off_screen=True) - pl.add_volume(grid, cmap='nipy_spectral', opacity='sigmoid_10', shade=False) - # Update the rendering scene every 0.1 seconds - plt.imshow(pl.screenshot()) - plt.pause(0.1) - -def save_BCs_vtk(timestep, BCs, gridInfo, output_dir='.'): - """ - Save boundary conditions as VTK format to the specified directory. - - Parameters - ---------- - timestep (int): The timestep number to be associated with the saved fields. - BCs (List[BC]): A list of boundary conditions to be saved. Each boundary condition must be an object of type BC. - - Returns - ------- - None - - Notes - ----- - This function saves the boundary conditions in the specified directory, with filenames based on the provided timestep number - and the filename. For example, if the timestep number is 10, the VTK file - will be saved as 'BCs_0000010.vtk'in the specified directory. - """ - - # Create a uniform grid - if gridInfo['nz'] == 0: - gridDimensions = (gridInfo['nx'] + 1, gridInfo['ny'] + 1, 1) - fieldDimensions = (gridInfo['nx'], gridInfo['ny'], 1) - else: - gridDimensions = (gridInfo['nx'] + 1, gridInfo['ny'] + 1, gridInfo['nz'] + 1) - fieldDimensions = (gridInfo['nx'], gridInfo['ny'], gridInfo['nz']) - - grid = pv.ImageData(dimensions=gridDimensions) - - # Dictionary to keep track of encountered BC names - bcNamesCount = {} - - for bc in BCs: - bcName = bc.name - if bcName in bcNamesCount: - bcNamesCount[bcName] += 1 - else: - bcNamesCount[bcName] = 0 - bcName += f"_{bcNamesCount[bcName]}" - - if bc.isDynamic: - bcIndices, _ = bc.update_function(timestep) - else: - bcIndices = bc.indices - - # Convert indices to 1D indices - if gridInfo['dim'] == 2: - bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions[:-1], order='F') - else: - bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions, order='F') - - grid[bcName] = np.zeros(fieldDimensions, dtype=bool).flatten(order='F') - grid[bcName][bcIndices] = True - - # Save the grid to a VTK file - output_filename = os.path.join(output_dir, "BCs_" + f"{timestep:07d}.vtk") - - start = time() - grid.save(output_filename, binary=True) - print(f"Saved {output_filename} in {time() - start:.6f} seconds.") - - -def rotate_geometry(indices, origin, axis, angle): - """ - Rotates a voxelized mesh around a given axis. - - Parameters - ---------- - indices : array-like - The indices of the voxels in the mesh. - origin : array-like - The coordinates of the origin of the rotation axis. - axis : array-like - The direction vector of the rotation axis. This should be a 3-element sequence. - angle : float - The angle by which to rotate the mesh, in radians. - - Returns - ------- - tuple - The indices of the voxels in the rotated mesh. - - Notes - ----- - This function rotates the mesh by applying a rotation matrix to the voxel indices. The rotation matrix is calculated - using the axis-angle representation of rotations. The origin of the rotation axis is assumed to be at (0, 0, 0). - """ - indices_rotated = (jnp.array(indices).T - origin) @ axangle2mat(axis, angle) + origin - return tuple(jnp.rint(indices_rotated).astype('int32').T) - -def voxelize_stl(stl_filename, length_lbm_unit=None, tranformation_matrix=None, pitch=None): - """ - Converts an STL file to a voxelized mesh. - - Parameters - ---------- - stl_filename : str - The name of the STL file to be voxelized. - length_lbm_unit : float, optional - The unit length in LBM. Either this or 'pitch' must be provided. - tranformation_matrix : array-like, optional - A transformation matrix to be applied to the mesh before voxelization. - pitch : float, optional - The pitch of the voxel grid. Either this or 'length_lbm_unit' must be provided. - - Returns - ------- - trimesh.VoxelGrid, float - The voxelized mesh and the pitch of the voxel grid. - - Notes - ----- - This function uses the trimesh library to load the STL file and voxelized the mesh. If a transformation matrix is - provided, it is applied to the mesh before voxelization. The pitch of the voxel grid is calculated based on the - maximum extent of the mesh and the provided lattice Boltzmann unit length, unless a pitch is provided directly. - """ - if length_lbm_unit is None and pitch is None: - raise ValueError("Either 'length_lbm_unit' or 'pitch' must be provided!") - mesh = trimesh.load_mesh(stl_filename, process=False) - length_phys_unit = mesh.extents.max() - if tranformation_matrix is not None: - mesh.apply_transform(tranformation_matrix) - if pitch is None: - pitch = length_phys_unit / length_lbm_unit - mesh_voxelized = mesh.voxelized(pitch=pitch) - return mesh_voxelized, pitch - - -def axangle2mat(axis, angle, is_normalized=False): - ''' Rotation matrix for rotation angle `angle` around `axis` - Parameters - ---------- - axis : 3 element sequence - vector specifying axis for rotation. - angle : scalar - angle of rotation in radians. - is_normalized : bool, optional - True if `axis` is already normalized (has norm of 1). Default False. - Returns - ------- - mat : array shape (3,3) - rotation matrix for specified rotation - Notes - ----- - From : https://github.com/matthew-brett/transforms3d - Ref : http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle - ''' - x, y, z = axis - if not is_normalized: - n = jnp.sqrt(x * x + y * y + z * z) - x = x / n - y = y / n - z = z / n - c = jnp.cos(angle) - s = jnp.sin(angle) - C = 1 - c - xs = x * s - ys = y * s - zs = z * s - xC = x * C - yC = y * C - zC = z * C - xyC = x * yC - yzC = y * zC - zxC = z * xC - return jnp.array([ - [x * xC + c, xyC - zs, zxC + ys], - [xyC + zs, y * yC + c, yzC - xs], - [zxC - ys, yzC + xs, z * zC + c]]) - -@partial(jit) -def q_criterion(u): - # Compute derivatives - u_x = u[..., 0] - u_y = u[..., 1] - u_z = u[..., 2] - - # Compute derivatives - u_x_dx = (u_x[2:, 1:-1, 1:-1] - u_x[:-2, 1:-1, 1:-1]) / 2 - u_x_dy = (u_x[1:-1, 2:, 1:-1] - u_x[1:-1, :-2, 1:-1]) / 2 - u_x_dz = (u_x[1:-1, 1:-1, 2:] - u_x[1:-1, 1:-1, :-2]) / 2 - u_y_dx = (u_y[2:, 1:-1, 1:-1] - u_y[:-2, 1:-1, 1:-1]) / 2 - u_y_dy = (u_y[1:-1, 2:, 1:-1] - u_y[1:-1, :-2, 1:-1]) / 2 - u_y_dz = (u_y[1:-1, 1:-1, 2:] - u_y[1:-1, 1:-1, :-2]) / 2 - u_z_dx = (u_z[2:, 1:-1, 1:-1] - u_z[:-2, 1:-1, 1:-1]) / 2 - u_z_dy = (u_z[1:-1, 2:, 1:-1] - u_z[1:-1, :-2, 1:-1]) / 2 - u_z_dz = (u_z[1:-1, 1:-1, 2:] - u_z[1:-1, 1:-1, :-2]) / 2 - - # Compute vorticity - mu_x = u_z_dy - u_y_dz - mu_y = u_x_dz - u_z_dx - mu_z = u_y_dx - u_x_dy - norm_mu = jnp.sqrt(mu_x ** 2 + mu_y ** 2 + mu_z ** 2) - - # Compute strain rate - s_0_0 = u_x_dx - s_0_1 = 0.5 * (u_x_dy + u_y_dx) - s_0_2 = 0.5 * (u_x_dz + u_z_dx) - s_1_0 = s_0_1 - s_1_1 = u_y_dy - s_1_2 = 0.5 * (u_y_dz + u_z_dy) - s_2_0 = s_0_2 - s_2_1 = s_1_2 - s_2_2 = u_z_dz - s_dot_s = ( - s_0_0 ** 2 + s_0_1 ** 2 + s_0_2 ** 2 + - s_1_0 ** 2 + s_1_1 ** 2 + s_1_2 ** 2 + - s_2_0 ** 2 + s_2_1 ** 2 + s_2_2 ** 2 - ) - - # Compute omega - omega_0_0 = 0.0 - omega_0_1 = 0.5 * (u_x_dy - u_y_dx) - omega_0_2 = 0.5 * (u_x_dz - u_z_dx) - omega_1_0 = -omega_0_1 - omega_1_1 = 0.0 - omega_1_2 = 0.5 * (u_y_dz - u_z_dy) - omega_2_0 = -omega_0_2 - omega_2_1 = -omega_1_2 - omega_2_2 = 0.0 - omega_dot_omega = ( - omega_0_0 ** 2 + omega_0_1 ** 2 + omega_0_2 ** 2 + - omega_1_0 ** 2 + omega_1_1 ** 2 + omega_1_2 ** 2 + - omega_2_0 ** 2 + omega_2_1 ** 2 + omega_2_2 ** 2 - ) - - # Compute q-criterion - q = 0.5 * (omega_dot_omega - s_dot_s) - - return norm_mu, q - - diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py new file mode 100644 index 00000000..94a66e5e --- /dev/null +++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py @@ -0,0 +1,94 @@ +import pytest +import numpy as np +import jax.numpy as jnp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb import DefaultConfig +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (100, 100)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + ], +) +def test_bc_equilibrium_jax(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + + equilibrium_bc = xlb.operator.boundary_condition.EquilibriumBC( + rho=1.0, + u=(0.0, 0.0, 0.0) if dim == 3 else (0.0, 0.0), + equilibrium_operator=xlb.operator.equilibrium.QuadraticEquilibrium(), + indices=indices, + ) + + bc_mask, missing_mask = indices_boundary_masker([equilibrium_bc], bc_mask, missing_mask, start_index=None) + + f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32) + + f_post = my_grid.create_field( + cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0 + ) # Arbitrary value so that we can check if the values are changed outside the boundary + + f = equilibrium_bc(f_pre, f_post, bc_mask, missing_mask) + + assert f.shape == (velocity_set.q,) + grid_shape + + # Assert that the values are correct in the indices of the sphere + weights = velocity_set.w + for i, weight in enumerate(weights): + if dim == 2: + assert jnp.allclose(f[i, indices[0], indices[1]], weight), f"Direction {i} in f does not match the expected weight" + else: + assert jnp.allclose(f[i, indices[0], indices[1], indices[2]], weight), f"Direction {i} in f does not match the expected weight" + + # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values + mask_outside = np.ones(grid_shape, dtype=bool) + mask_outside[indices] = False # Mark boundary as false + if dim == 2: + for i in range(velocity_set.q): + assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside]) + else: + for i in range(velocity_set.q): + assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside]) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py new file mode 100644 index 00000000..711c34bf --- /dev/null +++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py @@ -0,0 +1,97 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb import DefaultConfig +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (100, 100)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + ], +) +def test_bc_equilibrium_warp(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.UINT8) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + equilibrium = xlb.operator.equilibrium.QuadraticEquilibrium() + + equilibrium_bc = xlb.operator.boundary_condition.EquilibriumBC( + rho=1.0, + u=(0.0, 0.0, 0.0) if dim == 3 else (0.0, 0.0), + equilibrium_operator=equilibrium, + indices=indices, + ) + + bc_mask, missing_mask = indices_boundary_masker([equilibrium_bc], bc_mask, missing_mask, start_index=None) + + f = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32) + f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32) + f_post = my_grid.create_field( + cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0 + ) # Arbitrary value so that we can check if the values are changed outside the boundary + + f = equilibrium_bc(f_pre, f_post, bc_mask, missing_mask) + + f = f.numpy() + f_post = f_post.numpy() + + assert f.shape == (velocity_set.q,) + grid_shape if dim == 3 else (velocity_set.q, grid_shape[0], grid_shape[1], 1) + + # Assert that the values are correct in the indices of the sphere + weights = velocity_set.w + for i, weight in enumerate(weights): + if dim == 2: + assert np.allclose(f[i, indices[0], indices[1]], weight), f"Direction {i} in f does not match the expected weight" + else: + assert np.allclose(f[i, indices[0], indices[1], indices[2]], weight), f"Direction {i} in f does not match the expected weight" + + # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values + mask_outside = np.ones(grid_shape, dtype=bool) + mask_outside[indices] = False # Mark boundary as false + if dim == 2: + for i in range(velocity_set.q): + assert np.allclose(f[i, mask_outside], f_post[i, mask_outside]) + else: + for i in range(velocity_set.q): + assert np.allclose(f[i, mask_outside], f_post[i, mask_outside]) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py new file mode 100644 index 00000000..b4bc797b --- /dev/null +++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py @@ -0,0 +1,90 @@ +import pytest +import numpy as np +import jax.numpy as jnp +import xlb +import jax +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb import DefaultConfig +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_fullway_bounce_back_jax(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices) + + bc_mask, missing_mask = indices_boundary_masker([fullway_bc], bc_mask, missing_mask, start_index=None) + + f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=0.0) + # Generate a random field with the same shape + key = jax.random.PRNGKey(0) + random_field = jax.random.uniform(key, f_pre.shape) + # Add the random field to f_pre + f_pre += random_field + + f_post = my_grid.create_field( + cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0 + ) # Arbitrary value so that we can check if the values are changed outside the boundary + + f = fullway_bc(f_pre, f_post, bc_mask, missing_mask) + + assert f.shape == (velocity_set.q,) + grid_shape + + for i in range(velocity_set.q): + jnp.allclose( + f[velocity_set.opp_indices[i]][tuple(indices)], + f_pre[i][tuple(indices)], + ) + + # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values + mask_outside = np.ones(grid_shape, dtype=bool) + mask_outside[indices] = False # Mark boundary as false + if dim == 2: + for i in range(velocity_set.q): + assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside]) + else: + for i in range(velocity_set.q): + assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside]) diff --git a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py new file mode 100644 index 00000000..4f5d6757 --- /dev/null +++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py @@ -0,0 +1,93 @@ +import pytest +import numpy as np +import warp as wp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb import DefaultConfig +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.UINT8) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices) + + bc_mask, missing_mask = indices_boundary_masker([fullway_bc], bc_mask, missing_mask, start_index=None) + + # Generate a random field with the same shape + if dim == 2: + random_field = np.random.rand(velocity_set.q, grid_shape[0], grid_shape[1], 1).astype(np.float32) + else: + random_field = np.random.rand(velocity_set.q, grid_shape[0], grid_shape[1], grid_shape[2]).astype(np.float32) + # Add the random field to f_pre + f_pre = wp.array(random_field) + + f_post = my_grid.create_field( + cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0 + ) # Arbitrary value so that we can check if the values are changed outside the boundary + + f_pre = fullway_bc(f_pre, f_post, bc_mask, missing_mask) + + f = f_pre.numpy() + f_post = f_post.numpy() + + assert f.shape == (velocity_set.q,) + grid_shape if dim == 3 else (velocity_set.q, grid_shape[0], grid_shape[1], 1) + + for i in range(velocity_set.q): + np.allclose( + f[velocity_set.opp_indices[i]][tuple(indices)], + f_post[i][tuple(indices)], + ) + + # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values + mask_outside = np.ones(grid_shape, dtype=bool) + mask_outside[indices] = False # Mark boundary as false + if dim == 2: + for i in range(velocity_set.q): + assert np.allclose(f[i, mask_outside], f_post[i, mask_outside]) + else: + for i in range(velocity_set.q): + assert np.allclose(f[i, mask_outside], f_post[i, mask_outside]) diff --git a/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py new file mode 100644 index 00000000..eb81eda3 --- /dev/null +++ b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py @@ -0,0 +1,83 @@ +import pytest +import jax.numpy as jnp +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb import DefaultConfig +from xlb.grid import grid_factory + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (4, 4)), + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_indices_masker_jax(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + + assert len(indices) == dim + test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices) + test_bc.id = 5 + bc_mask, missing_mask = indices_boundary_masker([test_bc], bc_mask, missing_mask, start_index=None) + + assert missing_mask.dtype == xlb.Precision.BOOL.jax_dtype + + assert bc_mask.dtype == xlb.Precision.UINT8.jax_dtype + + assert bc_mask.shape == (1,) + grid_shape + + assert missing_mask.shape == (velocity_set.q,) + grid_shape + + if dim == 2: + assert jnp.all(bc_mask[0, indices[0], indices[1]] == test_bc.id) + # assert that the rest of the bc_mask is zero + bc_mask = bc_mask.at[0, indices[0], indices[1]].set(0) + assert jnp.all(bc_mask == 0) + if dim == 3: + assert jnp.all(bc_mask[0, indices[0], indices[1], indices[2]] == test_bc.id) + # assert that the rest of the bc_mask is zero + bc_mask = bc_mask.at[0, indices[0], indices[1], indices[2]].set(0) + assert jnp.all(bc_mask == 0) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py new file mode 100644 index 00000000..cb012cee --- /dev/null +++ b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py @@ -0,0 +1,91 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb import DefaultConfig +from xlb.grid import grid_factory +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_indices_masker_warp(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.UINT8) + + bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8) + + indices_boundary_masker = IndicesBoundaryMasker() + + # Make indices for boundary conditions (sphere) + sphere_radius = grid_shape[0] // 4 + nr = grid_shape[0] + x = np.arange(nr) + y = np.arange(nr) + z = np.arange(nr) + if dim == 2: + X, Y = np.meshgrid(x, y) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2) + else: + X, Y, Z = np.meshgrid(x, y, z) + indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2) + + indices = [tuple(indices[i]) for i in range(velocity_set.d)] + + assert len(indices) == dim + test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices) + test_bc.id = 5 + bc_mask, missing_mask = indices_boundary_masker( + [test_bc], + bc_mask, + missing_mask, + ) + assert missing_mask.dtype == xlb.Precision.UINT8.wp_dtype + + assert bc_mask.dtype == xlb.Precision.UINT8.wp_dtype + + bc_mask = bc_mask.numpy() + missing_mask = missing_mask.numpy() + + if len(grid_shape) == 2: + assert bc_mask.shape == (1,) + grid_shape + (1,), "bc_mask shape is incorrect got {}".format(bc_mask.shape) + assert missing_mask.shape == (velocity_set.q,) + grid_shape + (1,), "missing_mask shape is incorrect got {}".format(missing_mask.shape) + else: + assert bc_mask.shape == (1,) + grid_shape, "bc_mask shape is incorrect got {}".format(bc_mask.shape) + assert missing_mask.shape == (velocity_set.q,) + grid_shape, "missing_mask shape is incorrect got {}".format(missing_mask.shape) + + if dim == 2: + assert np.all(bc_mask[0, indices[0], indices[1]] == test_bc.id) + # assert that the rest of the bc_mask is zero + bc_mask[0, indices[0], indices[1]] = 0 + assert np.all(bc_mask == 0) + if dim == 3: + assert np.all(bc_mask[0, indices[0], indices[1], indices[2]] == test_bc.id) + # assert that the rest of the bc_mask is zero + bc_mask[0, indices[0], indices[1], indices[2]] = 0 + assert np.all(bc_mask == 0) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/grids/__init__.py b/tests/grids/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/grids/test_grid_jax.py b/tests/grids/test_grid_jax.py new file mode 100644 index 00000000..7255b7f6 --- /dev/null +++ b/tests/grids/test_grid_jax.py @@ -0,0 +1,74 @@ +import pytest +import jax +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from jax.sharding import Mesh +from jax.experimental import mesh_utils +import jax.numpy as jnp + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize("grid_size", [50, 100, 150]) +def test_jax_2d_grid_initialization(grid_size): + init_xlb_env(xlb.velocity_set.D2Q9) + grid_shape = (grid_size, grid_size) + my_grid = grid_factory(grid_shape) + f = my_grid.create_field(cardinality=9) + n_devices = jax.device_count() + + device_mesh = mesh_utils.create_device_mesh((1, n_devices, 1)) + expected_mesh = Mesh(device_mesh, axis_names=("cardinality", "x", "y")) + + assert f.shape == (9,) + grid_shape, "Field shape is incorrect" + assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect" + assert f.sharding.spec == ("cardinality", "x", "y"), "PartitionSpec is incorrect" + + +@pytest.mark.parametrize("grid_size", [50, 100, 150]) +def test_jax_3d_grid_initialization(grid_size): + init_xlb_env(xlb.velocity_set.D3Q19) + grid_shape = (grid_size, grid_size, grid_size) + my_grid = grid_factory(grid_shape) + f = my_grid.create_field(cardinality=9) + n_devices = jax.device_count() + + device_mesh = mesh_utils.create_device_mesh((1, n_devices, 1, 1)) + expected_mesh = Mesh(device_mesh, axis_names=("cardinality", "x", "y", "z")) + + assert f.shape == (9,) + grid_shape, "Field shape is incorrect" + assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect" + assert f.sharding.spec == ( + "cardinality", + "x", + "y", + "z", + ), "PartitionSpec is incorrect" + + +def test_jax_grid_create_field_fill_value(): + init_xlb_env(xlb.velocity_set.D2Q9) + grid_shape = (100, 100) + fill_value = 3.14 + my_grid = grid_factory(grid_shape) + + f = my_grid.create_field(cardinality=9, fill_value=fill_value) + assert f.shape == (9,) + grid_shape, "Field shape is incorrect" + assert jnp.allclose(f, fill_value), "Field not properly initialized with fill_value" + + +@pytest.fixture(autouse=True) +def setup_xlb_env(): + init_xlb_env(xlb.velocity_set.D2Q9) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/grids/test_grid_warp.py b/tests/grids/test_grid_warp.py new file mode 100644 index 00000000..2f1ab232 --- /dev/null +++ b/tests/grids/test_grid_warp.py @@ -0,0 +1,51 @@ +import pytest +import warp as wp +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb.precision_policy import Precision + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize("grid_size", [50, 100, 150]) +def test_warp_grid_create_field(grid_size): + for grid_shape in [(grid_size, grid_size), (grid_size, grid_size, grid_size)]: + init_xlb_env(xlb.velocity_set.D3Q19) + my_grid = grid_factory(grid_shape) + f = my_grid.create_field(cardinality=9, dtype=Precision.FP32) + if len(grid_shape) == 2: + assert f.shape == (9,) + grid_shape + (1,), "Field shape is incorrect got {}".format(f.shape) + else: + assert f.shape == (9,) + grid_shape, "Field shape is incorrect got {}".format(f.shape) + assert isinstance(f, wp.array), "Field should be a Warp ndarray" + + +def test_warp_grid_create_field_fill_value(): + init_xlb_env(xlb.velocity_set.D2Q9) + grid_shape = (100, 100) + fill_value = 3.14 + my_grid = grid_factory(grid_shape) + + f = my_grid.create_field(cardinality=9, dtype=Precision.FP32, fill_value=fill_value) + assert isinstance(f, wp.array), "Field should be a Warp ndarray" + + f = f.numpy() + assert np.allclose(f, fill_value), "Field not properly initialized with fill_value" + + +@pytest.fixture(autouse=True) +def setup_xlb_env(): + init_xlb_env(xlb.velocity_set.D2Q9) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/install/flow_past_sphere_3d_test.py b/tests/install/flow_past_sphere_3d_test.py new file mode 100644 index 00000000..6ae302cc --- /dev/null +++ b/tests/install/flow_past_sphere_3d_test.py @@ -0,0 +1,248 @@ +""" +Flow past a sphere (3D) β€” smoke test for JAX / WARP backends. + +NEON is **not** run: ``HalfwayBounceBackBC`` (sphere) has no NEON implementation in +XLB yet; NEON is listed under "Skipped (unsupported)" when the package is +installed. + +Run from the repository root:: + + python tests/install/flow_past_sphere_3d_test.py + +Domain and step counts are kept small for fast CI / install verification. +""" + +from __future__ import annotations + +import sys +import time +import traceback +from pathlib import Path +from typing import Any + +import numpy as np +import jax.numpy as jnp +import warp as wp +import xlb +import xlb.velocity_set +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb.operator.boundary_condition import ( + FullwayBounceBackBC, + HalfwayBounceBackBC, + RegularizedBC, + ExtrapolationOutflowBC, +) +from xlb.operator.force.momentum_transfer import MomentumTransfer +from xlb.operator.macroscopic import Macroscopic +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.precision_policy import PrecisionPolicy +from xlb.utils import save_image + + +def _run_flow_past_sphere_for_backend(compute_backend: ComputeBackend) -> None: + # Small domain for install / smoke tests (original example uses 256Γ—64Γ—64). + grid_shape = (32, 16, 16) + omega = 1.6 + precision_policy = PrecisionPolicy.FP32FP32 + velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=compute_backend) + u_max = 0.04 + num_steps = 20 + post_process_interval = 10 + + xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, + ) + + grid = grid_factory(grid_shape, compute_backend=compute_backend) + + box = grid.bounding_box_indices() + box_no_edge = grid.bounding_box_indices(remove_edges=True) + inlet = box_no_edge["left"] + outlet = box_no_edge["right"] + walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(velocity_set.d)] + walls = np.unique(np.array(walls), axis=-1).tolist() + + sphere_radius = max(grid_shape[1] // 12, 1) + x = np.arange(grid_shape[0]) + y = np.arange(grid_shape[1]) + z = np.arange(grid_shape[2]) + X, Y, Z = np.meshgrid(x, y, z, indexing="ij") + indices = np.where((X - grid_shape[0] // 6) ** 2 + (Y - grid_shape[1] // 2) ** 2 + (Z - grid_shape[2] // 2) ** 2 < sphere_radius**2) + sphere = [tuple(indices[i].tolist()) for i in range(velocity_set.d)] + + def bc_profile(): + H_y = float(grid_shape[1] - 1) + H_z = float(grid_shape[2] - 1) + + if compute_backend == ComputeBackend.JAX: + + def bc_profile_jax(): + yy = jnp.arange(grid_shape[1]) + zz = jnp.arange(grid_shape[2]) + Y, Z = jnp.meshgrid(yy, zz, indexing="ij") + y_center = Y - (H_y / 2.0) + z_center = Z - (H_z / 2.0) + r_squared = (2.0 * y_center / H_y) ** 2.0 + (2.0 * z_center / H_z) ** 2.0 + u_x = u_max * jnp.maximum(0.0, 1.0 - r_squared) + u_y = jnp.zeros_like(u_x) + u_z = jnp.zeros_like(u_x) + return jnp.stack([u_x, u_y, u_z]) + + return bc_profile_jax + + wp_dtype = precision_policy.compute_precision.wp_dtype + H_y_w = wp_dtype(grid_shape[1] - 1) + H_z_w = wp_dtype(grid_shape[2] - 1) + two = wp_dtype(2.0) + + @wp.func + def bc_profile_warp(index: wp.vec3i): + y = wp_dtype(index[1]) + z = wp_dtype(index[2]) + y_center = y - (H_y_w / two) + z_center = z - (H_z_w / two) + r_squared = (two * y_center / H_y_w) ** two + (two * z_center / H_z_w) ** two + return wp.vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1) + + return bc_profile_warp + + bc_left = RegularizedBC("velocity", profile=bc_profile(), indices=inlet) + bc_walls = FullwayBounceBackBC(indices=walls) + bc_outlet = ExtrapolationOutflowBC(indices=outlet) + bc_sphere = HalfwayBounceBackBC(indices=sphere) + boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere] + + stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="BGK", + ) + f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + + macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + ) + to_jax = xlb.utils.ToJAX("populations", velocity_set.q, grid_shape) + + momentum_transfer = MomentumTransfer(bc_sphere, compute_backend=compute_backend) + sphere_cross_section = float(np.pi * sphere_radius**2) + + prefix = f"flow_past_sphere_{compute_backend.name.lower()}" + + def post_process(step: int, f_0, f_1) -> None: + if compute_backend in (ComputeBackend.WARP, ComputeBackend.NEON): + wp.synchronize() + + boundary_force = momentum_transfer(f_0, f_1, bc_mask, missing_mask) + drag = boundary_force[0] + lift = boundary_force[2] + cd = 2.0 * drag / (u_max**2 * sphere_cross_section) + cl = 2.0 * lift / (u_max**2 * sphere_cross_section) + print(f"CD={cd}, CL={cl}") + + if not isinstance(f_0, jnp.ndarray): + f_0 = to_jax(f_0) + if compute_backend in (ComputeBackend.WARP, ComputeBackend.NEON): + wp.synchronize() + + rho, u = macro(f_0) + + u = u[:, 1:-1, 1:-1, 1:-1] + rho = rho[:, 1:-1, 1:-1, 1:-1] + u_magnitude = jnp.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + + fields = { + "u_magnitude": u_magnitude, + "u_x": u[0], + "u_y": u[1], + "u_z": u[2], + "rho": rho[0], + } + + save_image(fields["u_magnitude"][:, grid_shape[1] // 2, :], timestep=step, prefix=prefix) + print(f"Post-processed step {step}: saved u_magnitude slice (prefix={prefix})") + + start_time = time.time() + for step in range(num_steps): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 + + if step % post_process_interval == 0 or step == num_steps - 1: + post_process(step, f_0, f_1) + end_time = time.time() + elapsed = end_time - start_time + print(f"Completed step {step}. Elapsed for last chunk: {elapsed:.6f} s.") + start_time = time.time() + + +def run_flow_past_sphere_smoke() -> dict[str, Any]: + """Run JAX and WARP; skip NEON (unsupported BC). Missing NEON package -> ImportError path unused.""" + backends_order: tuple[ComputeBackend, ...] = ( + ComputeBackend.WARP, + ComputeBackend.JAX, + ComputeBackend.NEON, + ) + + executed: list[str] = [] + skipped_not_installed: list[str] = [] + skipped_unsupported: list[str] = [] + failed: list[tuple[str, str]] = [] + + for backend in backends_order: + print(f"\n--- Backend: {backend.name} ---") + if backend == ComputeBackend.NEON: + reason = "HalfwayBounceBackBC on the sphere has no NEON implementation in XLB (see bc_halfway_bounce_back.neon_implementation)." + print(f"SKIP (unsupported): NEON β€” {reason}") + skipped_unsupported.append(f"NEON ({reason})") + continue + + try: + _run_flow_past_sphere_for_backend(backend) + executed.append(backend.name) + print(f"OK: {backend.name} finished.") + except ImportError: + skipped_not_installed.append(backend.name) + except Exception as exc: + failed.append((backend.name, str(exc))) + print(f"FAIL {backend.name}:") + traceback.print_exc() + + print("\n=== Summary ===") + print(f"Executed: {', '.join(executed) if executed else '(none)'}") + if skipped_not_installed: + print("Skipped (not installed): " + ", ".join(skipped_not_installed) + " β€” required package not available.") + else: + print("Skipped (not installed): (none)") + if skipped_unsupported: + print("Skipped (unsupported configuration):") + for s in skipped_unsupported: + print(f" - {s}") + else: + print("Skipped (unsupported configuration): (none)") + if failed: + print("Failed:") + for name, msg in failed: + print(f" - {name}: {msg}") + else: + print("Failed: (none)") + + return { + "executed": executed, + "skipped_not_installed": skipped_not_installed, + "skipped_unsupported": skipped_unsupported, + "failed": failed, + } + + +def main() -> int: + result = run_flow_past_sphere_smoke() + return 1 if result["failed"] else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/install/run_install_matrix.py b/tests/install/run_install_matrix.py new file mode 100644 index 00000000..9a5241a3 --- /dev/null +++ b/tests/install/run_install_matrix.py @@ -0,0 +1,244 @@ +#!/usr/bin/env python3 +""" +Create isolated virtualenvs, install XLB under several dependency profiles, run +the 3D flow-past-sphere install smoke test, and print a full summary. + +Usage (from repository root):: + + python tests/install/run_install_matrix.py + +Options:: + + --repo-root PATH Repository root (default: parent of tests/install) + --reuse-venvs Do not recreate venvs if they already exist + --skip-install Only run tests (assume venvs already populated) + +Environment:: + + XLB_INSTALL_VENV_ROOT Override directory for venvs (default: .xlb_install_test_venvs) +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + + +def _venv_python(venv_dir: Path) -> Path: + if sys.platform == "win32": + return venv_dir / "Scripts" / "python.exe" + return venv_dir / "bin" / "python" + + +@dataclass +class ScenarioResult: + name: str + label: str + pip_spec: str + pip_ok: bool = False + pip_error: str = "" + sphere_ok: bool | None = None + sphere_exit: int | None = None + sphere_error: str = "" + + +def _run( + cmd: list[str], + cwd: Path, + *, + env: dict[str, str] | None = None, +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + cmd, + cwd=cwd, + env=env, + capture_output=True, + text=True, + check=False, + ) + + +def ensure_venv(venv_dir: Path, *, reuse: bool) -> Path: + py = _venv_python(venv_dir) + if reuse and py.is_file(): + return py + if venv_dir.exists() and not reuse: + shutil.rmtree(venv_dir) + venv_dir.parent.mkdir(parents=True, exist_ok=True) + subprocess.run([sys.executable, "-m", "venv", str(venv_dir)], check=True) + return _venv_python(venv_dir) + + +def pip_install_editable(venv_py: Path, repo_root: Path, extras: str) -> tuple[bool, str]: + """Run ``pip install -e .[extras]`` from *repo_root*.""" + spec = f".{extras}" if extras else "." + proc = _run( + [str(venv_py), "-m", "pip", "install", "--upgrade", "pip", "wheel"], + cwd=repo_root, + ) + if proc.returncode != 0: + return False, proc.stderr + proc.stdout + proc = _run([str(venv_py), "-m", "pip", "install", "-e", spec], cwd=repo_root) + if proc.returncode != 0: + return False, proc.stderr + proc.stdout + return True, "" + + +def run_test_script(venv_py: Path, repo_root: Path, test_script: Path) -> tuple[int, str]: + proc = _run([str(venv_py), str(test_script)], cwd=repo_root) + out = (proc.stdout or "") + (proc.stderr or "") + return proc.returncode, out + + +def main() -> int: + parser = argparse.ArgumentParser(description="XLB install matrix + 3D sphere smoke test") + parser.add_argument( + "--repo-root", + type=Path, + default=Path(__file__).resolve().parents[2], + help="XLB repository root", + ) + parser.add_argument("--reuse-venvs", action="store_true", help="Reuse existing venv directories") + parser.add_argument("--skip-install", action="store_true", help="Skip pip install; only run tests") + args = parser.parse_args() + + repo_root: Path = args.repo_root.resolve() + sphere_script = repo_root / "tests" / "install" / "flow_past_sphere_3d_test.py" + if not sphere_script.is_file(): + print(f"ERROR: test script not found: {sphere_script}", file=sys.stderr) + return 2 + + venv_root = Path(os.environ.get("XLB_INSTALL_VENV_ROOT", repo_root / ".xlb_install_test_venvs")).resolve() + + scenarios: list[tuple[str, str, str, str]] = [ + ( + "jax-cpu", + "JAX (CPU): editable install with test extra", + "[test]", + "Base JAX CPU jaxlib + test deps (pytest). Warp-lang is still pulled in as a core dependency.", + ), + ( + "jax-cuda", + "JAX[cuda]: editable install with cuda + test extras", + "[cuda,test]", + "Adds the cuda extra from setup (jax[cuda13] per setup.py). Requires a matching CUDA stack.", + ), + ( + "warp", + "WARP: explicit [warp,test] extras", + "[warp,test]", + "Explicit WARP extra plus test deps; core install already includes warp-lang.", + ), + ( + "neon", + "NEON: editable install with neon + test extras", + "[neon,test]", + "Installs neon_gpu wheel per setup.py (Linux x86_64/aarch64, Python 3.11+). " + "Includes h5py for Neon multires HDF5 export. " + "Uninstalls any existing warp-lang before install and PyPI warp-lang after (Neon's fork).", + ), + ] + + results: list[ScenarioResult] = [] + + for folder_name, label, extras, note in scenarios: + sr = ScenarioResult(name=folder_name, label=label, pip_spec=f"pip install -e .{extras}") + venv_dir = venv_root / folder_name + print(f"\n{'=' * 72}\nScenario: {folder_name}\n{label}\nNote: {note}\n{'=' * 72}") + + try: + venv_py = ensure_venv(venv_dir, reuse=args.reuse_venvs) + except Exception as exc: + sr.pip_ok = False + sr.pip_error = f"venv creation failed: {exc}" + results.append(sr) + print(sr.pip_error) + continue + + if not args.skip_install: + ok, err = pip_install_editable(venv_py, repo_root, extras) + sr.pip_ok = ok + sr.pip_error = err + if not ok: + print("PIP INSTALL FAILED\n", err[-4000:] if len(err) > 4000 else err) + results.append(sr) + continue + print("pip install: OK") + else: + sr.pip_ok = True + print("pip install: skipped (--skip-install)") + + code, combined = run_test_script(venv_py, repo_root, sphere_script) + sr.sphere_exit = code + sr.sphere_ok = code == 0 + sr.sphere_error = combined if code != 0 else "" + status: Literal["OK", "FAIL"] = "OK" if code == 0 else "FAIL" + print(f"flow_past_sphere_3d_test exit code: {code} ({status})") + if code != 0 and combined: + print("--- test output (tail) ---\n", combined[-2500:] if len(combined) > 2500 else combined) + + results.append(sr) + + # ------------------------------------------------------------------ summary + print("\n") + print("=" * 72) + print("FULL SUMMARY β€” XLB install matrix + flow_past_sphere_3d_test.py") + print("=" * 72) + print(f"Repository: {repo_root}") + print(f"Venv root: {venv_root}") + print() + + col_w = max(len(r.name) for r in results) if results else 12 + hdr = f"{'Scenario':<{col_w}} {'pip':^5} {'3d':^5}" + print(hdr) + print("-" * len(hdr)) + + any_test_failed = False + any_pip_failed = False + + for r in results: + pip_s = "ok" if r.pip_ok else "FAIL" + if not r.pip_ok: + any_pip_failed = True + if r.sphere_ok is None: + sph_s = "β€”" + elif r.sphere_ok: + sph_s = "ok" + else: + sph_s = "FAIL" + any_test_failed = True + print(f"{r.name:<{col_w}} {pip_s:^5} {sph_s:^5}") + print(f"{'':<{col_w}} {r.label}") + + print() + print("Details per scenario") + print("-" * 72) + for r in results: + print(f"\n* {r.name} β€” {r.label}") + print(f" pip ok: {r.pip_ok}") + if r.pip_error and not r.pip_ok: + print(f" pip error (truncated): {r.pip_error[:1500]}...") + print(f" sphere test exit: {r.sphere_exit}") + if r.sphere_error and r.sphere_exit not in (0, None): + print(f" sphere output (truncated): {r.sphere_error[:2000]}...") + + print() + print("=" * 72) + if any_pip_failed: + print("Overall: FAILURE (at least one pip install failed).") + return 1 + if any_test_failed: + print("Overall: FAILURE (at least one 3D smoke test returned non-zero).") + return 1 + print("Overall: SUCCESS β€” all configured installs and tests completed.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/install/test_neon_install_warp_cleanup.py b/tests/install/test_neon_install_warp_cleanup.py new file mode 100644 index 00000000..927af050 --- /dev/null +++ b/tests/install/test_neon_install_warp_cleanup.py @@ -0,0 +1,88 @@ +""" +Verify that installing XLB with the ``[neon]`` extra uninstalls a pre-existing +``warp-lang`` (Neon ships its own fork via ``neon_gpu``). + +Uses an isolated virtualenv and ``pip install -e``, same style as +``run_install_matrix.py``. Skips on platforms where Neon wheels are not built +(see ``setup.py``). +""" + +from __future__ import annotations + +import os +import platform +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _venv_python(venv_dir: Path) -> Path: + if sys.platform == "win32": + return venv_dir / "Scripts" / "python.exe" + return venv_dir / "bin" / "python" + + +def _neon_wheels_supported() -> bool: + if sys.platform != "linux": + return False + return platform.machine() in ("x86_64", "aarch64") + + +def _run(cmd: list[str], *, cwd: Path, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: + return subprocess.run( + cmd, + cwd=cwd, + env=env, + capture_output=True, + text=True, + check=False, + ) + + +@pytest.mark.skipif(sys.version_info < (3, 11), reason="XLB requires Python >= 3.11") +@pytest.mark.skipif(not _neon_wheels_supported(), reason="Neon wheels: Linux x86_64 / aarch64 only") +def test_neon_editable_install_after_prior_warp_lang(tmp_path: Path) -> None: + """Pre-install ``warp-lang``, then ``pip install -e .[neon,test]``; imports must work.""" + venv_dir = tmp_path / "venv" + subprocess.run([sys.executable, "-m", "venv", str(venv_dir)], check=True) + venv_py = _venv_python(venv_dir) + assert venv_py.is_file(), f"missing venv python: {venv_py}" + + env = {**os.environ, "XLB_NEON_SKIP_UNINSTALL_WARP": ""} + + proc = _run([str(venv_py), "-m", "pip", "install", "--upgrade", "pip", "wheel"], cwd=REPO_ROOT, env=env) + assert proc.returncode == 0, proc.stdout + proc.stderr + + proc = _run( + [str(venv_py), "-m", "pip", "install", "warp-lang==1.10.0"], + cwd=REPO_ROOT, + env=env, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + + proc = _run( + [str(venv_py), "-m", "pip", "install", "-e", ".[neon,test]"], + cwd=REPO_ROOT, + env=env, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + + proc = _run( + [str(venv_py), "-c", "import neon; import warp; print('ok', neon.__file__, warp.__file__)"], + cwd=REPO_ROOT, + env=env, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "ok" in proc.stdout + + # Repeat editable install: pre/post uninstall hooks must stay safe on re-run. + proc = _run( + [str(venv_py), "-m", "pip", "install", "-e", ".[neon,test]"], + cwd=REPO_ROOT, + env=env, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr diff --git a/tests/kernels/__init__.py b/tests/kernels/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/kernels/collision/test_bgk_collision_jax.py b/tests/kernels/collision/test_bgk_collision_jax.py new file mode 100644 index 00000000..3d6ea901 --- /dev/null +++ b/tests/kernels/collision/test_bgk_collision_jax.py @@ -0,0 +1,54 @@ +import pytest +import jax.numpy as jnp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.collision import BGK +from xlb.grid import grid_factory +from xlb import DefaultConfig + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape,omega", + [ + (2, xlb.velocity_set.D2Q9, (100, 100), 0.6), + (2, xlb.velocity_set.D2Q9, (100, 100), 1.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 0.6), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 0.6), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0), + ], +) +def test_bgk_collision(dim, velocity_set, grid_shape, omega): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho = my_grid.create_field(cardinality=1, fill_value=1.0) + u = my_grid.create_field(cardinality=dim, fill_value=0.0) + + # Compute equilibrium + compute_macro = QuadraticEquilibrium() + f_eq = compute_macro(rho, u) + + # Compute collision + + compute_collision = BGK() + + f_orig = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + + f_out = compute_collision(f_orig, f_eq, omega) + + assert jnp.allclose(f_out, f_orig - omega * (f_orig - f_eq)) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/collision/test_bgk_collision_warp.py b/tests/kernels/collision/test_bgk_collision_warp.py new file mode 100644 index 00000000..fa6884b2 --- /dev/null +++ b/tests/kernels/collision/test_bgk_collision_warp.py @@ -0,0 +1,57 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.collision import BGK +from xlb.grid import grid_factory +from xlb import DefaultConfig + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape,omega", + [ + (2, xlb.velocity_set.D2Q9, (100, 100), 0.6), + (2, xlb.velocity_set.D2Q9, (100, 100), 1.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 0.6), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 0.6), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0), + ], +) +def test_bgk_collision_warp(dim, velocity_set, grid_shape, omega): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho = my_grid.create_field(cardinality=1, fill_value=1.0) + u = my_grid.create_field(cardinality=dim, fill_value=0.0) + + compute_macro = QuadraticEquilibrium() + + f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + f_eq = compute_macro(rho, u, f_eq) + + compute_collision = BGK() + f_orig = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + + f_out = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + f_out = compute_collision(f_orig, f_eq, f_out, omega) + + f_eq = f_eq.numpy() + f_out = f_out.numpy() + f_orig = f_orig.numpy() + + assert np.allclose(f_out, f_orig - omega * (f_orig - f_eq), atol=1e-5) + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/equilibrium/test_equilibrium_jax.py b/tests/kernels/equilibrium/test_equilibrium_jax.py new file mode 100644 index 00000000..e0451845 --- /dev/null +++ b/tests/kernels/equilibrium/test_equilibrium_jax.py @@ -0,0 +1,52 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.grid import grid_factory +from xlb import DefaultConfig + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_quadratic_equilibrium_jax(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho = my_grid.create_field(cardinality=1, fill_value=1.0) + u = my_grid.create_field(cardinality=dim, fill_value=0.0) + + # Compute equilibrium + compute_macro = QuadraticEquilibrium() + f_eq = compute_macro(rho, u) + + # Test sum of f_eq across cardinality at each point + sum_f_eq = np.sum(f_eq, axis=0) + assert np.allclose(sum_f_eq, 1.0), "Sum of f_eq should be 1.0 across all directions at each grid point" + + # Test that each direction matches the expected weights + weights = DefaultConfig.velocity_set.w + for i, weight in enumerate(weights): + assert np.allclose(f_eq[i, ...], weight), f"Direction {i} in f_eq does not match the expected weight" + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/equilibrium/test_equilibrium_warp.py b/tests/kernels/equilibrium/test_equilibrium_warp.py new file mode 100644 index 00000000..8b7e8f44 --- /dev/null +++ b/tests/kernels/equilibrium/test_equilibrium_warp.py @@ -0,0 +1,58 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.grid import grid_factory +from xlb import DefaultConfig + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_quadratic_equilibrium_warp(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho = my_grid.create_field(cardinality=1, fill_value=1.0) + u = my_grid.create_field(cardinality=dim, fill_value=0.0) + + f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + + compute_macro = QuadraticEquilibrium() + f_eq = compute_macro(rho, u, f_eq) + + f_eq_np = f_eq.numpy() + + sum_f_eq = np.sum(f_eq_np, axis=0) + assert np.allclose(sum_f_eq, 1.0), "Sum of f_eq should be 1.0 across all directions at each grid point" + + weights = DefaultConfig.velocity_set.w + for i, weight in enumerate(weights): + assert np.allclose(f_eq_np[i, ...], weight), f"Direction {i} in f_eq does not match the expected weight" + + +# @pytest.fixture(autouse=True) +# def setup_xlb_env(request): +# dim, velocity_set, grid_shape = request.param +# init_xlb_env(velocity_set) + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/macroscopic/test_macroscopic_jax.py b/tests/kernels/macroscopic/test_macroscopic_jax.py new file mode 100644 index 00000000..a40a4520 --- /dev/null +++ b/tests/kernels/macroscopic/test_macroscopic_jax.py @@ -0,0 +1,54 @@ +import pytest +import numpy as np +import xlb +import jax +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.macroscopic import Macroscopic +from xlb.grid import grid_factory + +# Set JAX to use highest precision for matmul operations +jax.config.update("jax_default_matmul_precision", "highest") + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape,rho,velocity", + [ + (2, xlb.velocity_set.D2Q9, (100, 100), 1.0, 0.0), + (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 1.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0, 0.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 1.0), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0, 0.0), + (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.1, 1.0), + ], +) +def test_macroscopic_jax(dim, velocity_set, grid_shape, rho, velocity): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho_field = my_grid.create_field(cardinality=1, fill_value=rho) + velocity_field = my_grid.create_field(cardinality=dim, fill_value=velocity) + + # Compute equilibrium + f_eq = QuadraticEquilibrium()(rho_field, velocity_field) + + compute_macro = Macroscopic() + + rho_calc, u_calc = compute_macro(f_eq) + + # Test sum of f_eq which should be 1.0 for rho and 0.0 for u + assert np.allclose(rho_calc, rho), "Sum of f_eq should be {rho} for rho" + assert np.allclose(u_calc, velocity, atol=1e-06), "Sum of f_eq should be {velocity} for u" + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/macroscopic/test_macroscopic_warp.py b/tests/kernels/macroscopic/test_macroscopic_warp.py new file mode 100644 index 00000000..9ffe849d --- /dev/null +++ b/tests/kernels/macroscopic/test_macroscopic_warp.py @@ -0,0 +1,53 @@ +import pytest +import numpy as np +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.macroscopic import Macroscopic +from xlb.grid import grid_factory +from xlb import DefaultConfig + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape,rho,velocity", + [ + (2, xlb.velocity_set.D2Q9, (100, 100), 1.0, 0.0), + (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 1.0), + (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 2.0), + (2, xlb.velocity_set.D2Q9, (50, 50), 1.1, 2.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0, 0.0), + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 1.0), # TODO: Uncommenting will cause a Warp error. Needs investigation. + (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 2.0), # TODO: Uncommenting will cause a Warp error. Needs investigation. + ], +) +def test_macroscopic_warp(dim, velocity_set, grid_shape, rho, velocity): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + + rho_field = my_grid.create_field(cardinality=1, fill_value=rho) + velocity_field = my_grid.create_field(cardinality=dim, fill_value=velocity) + + f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q) + f_eq = QuadraticEquilibrium()(rho_field, velocity_field, f_eq) + + compute_macro = Macroscopic() + rho_calc = my_grid.create_field(cardinality=1) + u_calc = my_grid.create_field(cardinality=dim) + + rho_calc, u_calc = compute_macro(f_eq, rho_calc, u_calc) + + assert np.allclose(rho_calc.numpy(), rho), f"Computed density should be close to initialized density {rho}" + assert np.allclose(u_calc.numpy(), velocity, atol=1e-06), f"Computed velocity should be close to initialized velocity {velocity}" + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/stream/test_stream_jax.py b/tests/kernels/stream/test_stream_jax.py new file mode 100644 index 00000000..015c21c6 --- /dev/null +++ b/tests/kernels/stream/test_stream_jax.py @@ -0,0 +1,69 @@ +import pytest +import jax.numpy as jnp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.stream import Stream +from xlb import DefaultConfig +from xlb.grid import grid_factory + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.JAX) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.JAX, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_stream_operator_jax(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid = grid_factory(grid_shape) + velocity_set = DefaultConfig.velocity_set + + stream_op = Stream() + + f_initial = my_grid.create_field(cardinality=velocity_set.q) + f_initial = f_initial.at[..., f_initial.shape[-1] // 2].set(1) + + f_streamed = stream_op(f_initial) + + expected = [] + + if dim == 2: + for i in range(velocity_set.q): + expected.append( + jnp.roll( + f_initial[i, ...], + (velocity_set.c[0][i], velocity_set.c[1][i]), + axis=(0, 1), + ) + ) + elif dim == 3: + for i in range(velocity_set.q): + expected.append( + jnp.roll( + f_initial[i, ...], + (velocity_set.c[0][i], velocity_set.c[1][i], velocity_set.c[2][i]), + axis=(0, 1, 2), + ) + ) + + expected = jnp.stack(expected, axis=0) + + assert jnp.allclose(f_streamed, expected), "Streaming did not occur as expected" + + +if __name__ == "__main__": + pytest.main() diff --git a/tests/kernels/stream/test_stream_warp.py b/tests/kernels/stream/test_stream_warp.py new file mode 100644 index 00000000..e10cc467 --- /dev/null +++ b/tests/kernels/stream/test_stream_warp.py @@ -0,0 +1,81 @@ +import pytest +import jax.numpy as jnp +import numpy as np +import warp as wp +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.operator.stream import Stream +from xlb import DefaultConfig +from xlb.grid import grid_factory + + +def init_xlb_env(velocity_set): + vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, compute_backend=ComputeBackend.WARP) + xlb.init( + default_precision_policy=xlb.PrecisionPolicy.FP32FP32, + default_backend=ComputeBackend.WARP, + velocity_set=vel_set, + ) + + +@pytest.mark.parametrize( + "dim,velocity_set,grid_shape", + [ + (2, xlb.velocity_set.D2Q9, (50, 50)), + (2, xlb.velocity_set.D2Q9, (100, 100)), + (3, xlb.velocity_set.D3Q19, (50, 50, 50)), + (3, xlb.velocity_set.D3Q19, (100, 100, 100)), + (3, xlb.velocity_set.D3Q27, (50, 50, 50)), + (3, xlb.velocity_set.D3Q27, (100, 100, 100)), + ], +) +def test_stream_operator_warp(dim, velocity_set, grid_shape): + init_xlb_env(velocity_set) + my_grid_jax = grid_factory(grid_shape, compute_backend=ComputeBackend.JAX) + velocity_set = DefaultConfig.velocity_set + + f_initial = my_grid_jax.create_field(cardinality=velocity_set.q) + f_initial = f_initial.at[..., f_initial.shape[-1] // 2].set(1) + + expected = [] + + if dim == 2: + for i in range(velocity_set.q): + expected.append( + jnp.roll( + f_initial[i, ...], + (velocity_set.c[0][i], velocity_set.c[1][i]), + axis=(0, 1), + ) + ) + elif dim == 3: + for i in range(velocity_set.q): + expected.append( + jnp.roll( + f_initial[i, ...], + (velocity_set.c[0][i], velocity_set.c[1][i], velocity_set.c[2][i]), + axis=(0, 1, 2), + ) + ) + + expected = jnp.stack(expected, axis=0) + + if dim == 2: + f_initial_warp = wp.array(f_initial[..., np.newaxis]) + + elif dim == 3: + f_initial_warp = wp.array(f_initial) + + stream_op = Stream() + my_grid_warp = grid_factory(grid_shape, compute_backend=ComputeBackend.WARP) + f_streamed = my_grid_warp.create_field(cardinality=velocity_set.q) + f_streamed = stream_op(f_initial_warp, f_streamed) + + if len(grid_shape) == 2: + assert jnp.allclose(f_streamed.numpy()[..., 0], np.array(expected)), "Streaming did not occur as expected" + else: + assert jnp.allclose(f_streamed.numpy(), np.array(expected)), "Streaming did not occur as expected" + + +if __name__ == "__main__": + pytest.main() diff --git a/xlb/__init__.py b/xlb/__init__.py new file mode 100644 index 00000000..02b7a994 --- /dev/null +++ b/xlb/__init__.py @@ -0,0 +1,38 @@ +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("xlb") +except PackageNotFoundError: + __version__ = "0.0.0" + +# Enum classes +from xlb.compute_backend import ComputeBackend as ComputeBackend +from xlb.precision_policy import PrecisionPolicy as PrecisionPolicy, Precision as Precision +from xlb.physics_type import PhysicsType as PhysicsType +from xlb.mres_perf_optimization_type import MresPerfOptimizationType as MresPerfOptimizationType + +# Config +from .default_config import init as init, DefaultConfig as DefaultConfig + +# Velocity Set +import xlb.velocity_set + +# Operators +import xlb.operator.equilibrium +import xlb.operator.collision +import xlb.operator.stream +import xlb.operator.boundary_condition +import xlb.operator.macroscopic +import xlb.operator.postprocess + +# Grids +import xlb.grid + +# Solvers +import xlb.helper + +# Utils +import xlb.utils + +# Distributed computing +import xlb.distribute diff --git a/xlb/cell_type.py b/xlb/cell_type.py new file mode 100644 index 00000000..ea082e10 --- /dev/null +++ b/xlb/cell_type.py @@ -0,0 +1,11 @@ +# Boundary-mask constants for the bc_mask field. +# Each voxel in the domain carries a uint8 tag in bc_mask that encodes its role: +# BC_NONE β€” regular fluid voxel (no boundary condition) +# BC_SFV β€” Simple Fluid Voxel: fluid cell not involved in any BC, +# explosion, or coalescence (used for fast-path kernels) +# BC_SOLID β€” solid / obstacle voxel (skipped by all LBM operators) +# Registered boundary conditions receive IDs in the range [1, 253]. + +BC_NONE = 0 +BC_SFV = 254 +BC_SOLID = 255 diff --git a/xlb/compute_backend.py b/xlb/compute_backend.py new file mode 100644 index 00000000..cbd19918 --- /dev/null +++ b/xlb/compute_backend.py @@ -0,0 +1,18 @@ +""" +Compute-backend enumeration for XLB. +""" + +from enum import Enum, auto + + +class ComputeBackend(Enum): + """Available compute backends. + + ``JAX`` β€” single-res, multi-GPU/TPU via JAX. + ``WARP`` β€” single-res, single-GPU CUDA via NVIDIA Warp. + ``NEON`` β€” single-res, multi-GPU or multi-res single-GPU via Neon (uses Warp kernels internally). + """ + + JAX = auto() + WARP = auto() + NEON = auto() diff --git a/xlb/default_config.py b/xlb/default_config.py new file mode 100644 index 00000000..ab6a331a --- /dev/null +++ b/xlb/default_config.py @@ -0,0 +1,126 @@ +""" +Global configuration for XLB. + +Call :func:`init` once at the start of every script to select the velocity +set, compute backend, and precision policy. All operators read their +defaults from :class:`DefaultConfig` when explicit arguments are omitted. +""" + +import os + +from xlb.compute_backend import ComputeBackend +from dataclasses import dataclass +from xlb.precision_policy import PrecisionPolicy + + +@dataclass +class DefaultConfig: + """Singleton holding the active global configuration. + + Attributes are set by :func:`init` and read by operators, grids, and + helpers throughout XLB. + + Attributes + ---------- + default_precision_policy : PrecisionPolicy or None + Active precision policy (compute / store dtype pair). + velocity_set : VelocitySet or None + Active lattice velocity set. + default_backend : ComputeBackend or None + Active compute backend. + """ + + default_precision_policy = None + velocity_set = None + default_backend = None + + +def _warp_init_and_select_cuda_device(): + """Initialize Warp and pin the default CUDA device for single-GPU XLB runs. + + With multiple GPUs, Warp's default device for allocations and launches can + otherwise diverge. Set ``XLB_WARP_DEVICE`` (e.g. ``cuda:0`` or ``cuda:1``) + to choose which GPU Warp uses; defaults to ``cuda:0`` when unset. + """ + import warp as wp + + wp.init() # TODO: Must be removed in the future versions of WARP + if wp.get_cuda_device_count() == 0: + return + choice = os.environ.get("XLB_WARP_DEVICE", "cuda:0").strip() + try: + wp.set_device(choice) + except Exception: + try: + wp.set_device("cuda:0") + except Exception: + pass + + +def init(velocity_set, default_backend, default_precision_policy): + """Initialize the global XLB configuration. + + Must be called before creating any grid, operator, or field. + + Parameters + ---------- + velocity_set : VelocitySet + Lattice velocity set (e.g. ``D3Q19``). + default_backend : ComputeBackend + Compute backend to use (JAX, WARP, or NEON). + default_precision_policy : PrecisionPolicy + Precision policy for compute and storage dtypes. + """ + DefaultConfig.velocity_set = velocity_set + DefaultConfig.default_backend = default_backend + DefaultConfig.default_precision_policy = default_precision_policy + + if default_backend == ComputeBackend.WARP: + _warp_init_and_select_cuda_device() + elif default_backend == ComputeBackend.NEON: + import warp as wp + import neon + + # wp.config.mode = "release" + # wp.config.llvm_cuda = False + # wp.config.verbose = True + # wp.verbose_warnings = True + + _warp_init_and_select_cuda_device() + + # It's a good idea to always clear the kernel cache when developing new native or codegen features + wp.build.clear_kernel_cache() + + # !!! DO THIS BEFORE DEFINING/USING ANY KERNELS WITH CUSTOM TYPES + neon.init() + + elif default_backend == ComputeBackend.JAX: + check_backend_support() + else: + raise ValueError(f"Unsupported compute backend: {default_backend}") + + +def default_backend() -> ComputeBackend: + """Return the currently configured compute backend.""" + return DefaultConfig.default_backend + + +def check_backend_support(): + """Print a summary of available JAX hardware accelerators.""" + import jax + + if jax.devices()[0].platform == "gpu": + gpus = jax.devices("gpu") + if len(gpus) > 1: + print("Multi-GPU support is available: {} GPUs detected.".format(len(gpus))) + elif len(gpus) == 1: + print("Single-GPU support is available: 1 GPU detected.") + + elif jax.devices()[0].platform == "tpu": + tpus = jax.devices("tpu") + if len(tpus) > 1: + print("Multi-TPU support is available: {} TPUs detected.".format(len(tpus))) + elif len(tpus) == 1: + print("Single-TPU support is available: 1 TPU detected.") + else: + print("No GPU support is available; CPU fallback will be used.") diff --git a/xlb/distribute/__init__.py b/xlb/distribute/__init__.py new file mode 100644 index 00000000..dd9f33dd --- /dev/null +++ b/xlb/distribute/__init__.py @@ -0,0 +1 @@ +from .distribute import distribute diff --git a/xlb/distribute/distribute.py b/xlb/distribute/distribute.py new file mode 100644 index 00000000..1fc9138e --- /dev/null +++ b/xlb/distribute/distribute.py @@ -0,0 +1,105 @@ +from jax.sharding import PartitionSpec as P +from xlb.operator import Operator +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition.boundary_condition import ImplementationStep +from jax import lax +from jax import shard_map +from jax import jit + + +def distribute_operator( + operator: Operator, + grid, + velocity_set, + num_results=1, + ops="permute", +) -> Operator: + # Define the sharded operator + def _sharded_operator(*args): + result = operator(*args) + + if ops == "permute": + # Define permutation rules for right and left communication + rightPerm = [(i, (i + 1) % grid.nDevices) for i in range(grid.nDevices)] + leftPerm = [((i + 1) % grid.nDevices, i) for i in range(grid.nDevices)] + + left_comm, right_comm = ( + result[velocity_set.right_indices, :1, ...], + result[velocity_set.left_indices, -1:, ...], + ) + + left_comm = lax.ppermute( + left_comm, + perm=rightPerm, + axis_name="x", + ) + + right_comm = lax.ppermute( + right_comm, + perm=leftPerm, + axis_name="x", + ) + + result = result.at[velocity_set.right_indices, :1, ...].set(left_comm) + result = result.at[velocity_set.left_indices, -1:, ...].set(right_comm) + + return result + else: + raise NotImplementedError(f"Operation {ops} not implemented") + + # Build sharding_flags and in_specs based on args + def build_specs(grid, *args): + sharding_flags = [] + in_specs = [] + for arg in args: + if arg.shape[1:] == grid.shape: + sharding_flags.append(True) + else: + sharding_flags.append(False) + + in_specs = tuple(P(*((None, "x") + (grid.dim - 1) * (None,))) if flag else P() for flag in sharding_flags) + out_specs = tuple(P(*((None, "x") + (grid.dim - 1) * (None,))) for _ in range(num_results)) + return tuple(sharding_flags), in_specs, out_specs + + def _wrapped_operator(*args): + sharding_flags, in_specs, out_specs = build_specs(grid, *args) + + if len(out_specs) == 1: + out_specs = out_specs[0] + + distributed_operator = shard_map( + _sharded_operator, + mesh=grid.global_mesh, + in_specs=in_specs, + out_specs=out_specs, + check_vma=False, + ) + return distributed_operator(*args) + + return jit(_wrapped_operator) + + +def distribute(operator, grid, velocity_set, num_results=1, ops="permute"): + """ + Distribute an operator or a stepper. + If the operator is a stepper, check for post-streaming boundary conditions + before deciding how to distribute. + """ + if isinstance(operator, IncompressibleNavierStokesStepper): + # Check for post-streaming boundary conditions + has_post_streaming_bc = any(bc.implementation_step == ImplementationStep.STREAMING for bc in operator.boundary_conditions) + + if has_post_streaming_bc: + # If there are post-streaming BCs, only distribute the stream operator + distributed_stream = distribute_operator(operator.stream, grid, velocity_set) + operator.stream = distributed_stream + else: + # If no post-streaming BCs, distribute the whole operator + distributed_op = distribute_operator(operator, grid, velocity_set, num_results=num_results, ops=ops) + return distributed_op + + return operator + else: + # For other operators, apply the original distribution logic + distributed_op = distribute_operator(operator, grid, velocity_set, num_results=num_results, ops=ops) + return distributed_op diff --git a/xlb/experimental/__init__.py b/xlb/experimental/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/xlb/experimental/__init__.py @@ -0,0 +1 @@ + diff --git a/xlb/experimental/ooc/__init__.py b/xlb/experimental/ooc/__init__.py new file mode 100644 index 00000000..5206cc18 --- /dev/null +++ b/xlb/experimental/ooc/__init__.py @@ -0,0 +1,2 @@ +from xlb.experimental.ooc.out_of_core import OOCmap +from xlb.experimental.ooc.ooc_array import OOCArray diff --git a/xlb/experimental/ooc/ooc_array.py b/xlb/experimental/ooc/ooc_array.py new file mode 100644 index 00000000..11aedc3c --- /dev/null +++ b/xlb/experimental/ooc/ooc_array.py @@ -0,0 +1,440 @@ +import numpy as np +import cupy as cp + +# from mpi4py import MPI +import itertools + +from xlb.experimental.ooc.tiles.dense_tile import DenseTile, DenseGPUTile, DenseCPUTile +from xlb.experimental.ooc.tiles.compressed_tile import ( + CompressedTile, + CompressedGPUTile, + CompressedCPUTile, +) + + +class OOCArray: + """An out-of-core distributed array class. + + Parameters + ---------- + shape : tuple + The shape of the array. + dtype : cp.dtype + The data type of the array. + tile_shape : tuple + The shape of the tiles. Should be a factor of the shape. + padding : int or tuple + The padding of the tiles. + comm : MPI communicator + The MPI communicator. + devices : list of cp.cuda.Device + The list of GPU devices to use. + codec : Codec + The codec to use for compression. None for no compression (Dense tiles). + nr_compute_tiles : int + The number of compute tiles used for asynchronous copies. + TODO currently only 1 is supported when using JAX. + """ + + def __init__( + self, + shape, + dtype, + tile_shape, + padding=1, + comm=None, + devices=[cp.cuda.Device(0)], + codec=None, + nr_compute_tiles=1, + ): + self.shape = shape + self.tile_shape = tile_shape + self.dtype = dtype + if isinstance(padding, int): + padding = (padding,) * len(shape) + self.padding = padding + self.comm = comm + self.devices = devices + self.codec = codec + self.nr_compute_tiles = nr_compute_tiles + + # Set tile class + if self.codec is None: + self.Tile = DenseTile + self.DeviceTile = DenseGPUTile + self.HostTile = DenseCPUTile # TODO: Possibly make HardDiskTile or something + + else: + self.Tile = CompressedTile + self.DeviceTile = CompressedGPUTile + self.HostTile = CompressedCPUTile + + # Get process id and number of processes + self.pid = self.comm.Get_rank() + self.nr_proc = self.comm.Get_size() + + # Check that the tile shape divides the array shape. + if any([shape[i] % tile_shape[i] != 0 for i in range(len(shape))]): + raise ValueError(f"Tile shape {tile_shape} does not divide shape {shape}.") + self.tile_dims = tuple([shape[i] // tile_shape[i] for i in range(len(shape))]) + self.nr_tiles = np.prod(self.tile_dims) + + # Get number of tiles per process + if self.nr_tiles % self.nr_proc != 0: + raise ValueError(f"Number of tiles {self.nr_tiles} does not divide number of processes {self.nr_proc}.") + self.nr_tiles_per_proc = self.nr_tiles // self.nr_proc + + # Make the tile mapppings + self.tile_process_map = {} + self.tile_device_map = {} + for i, tile_index in enumerate(itertools.product(*[range(n) for n in self.tile_dims])): + self.tile_process_map[tile_index] = i % self.nr_proc + self.tile_device_map[tile_index] = devices[i % len(devices)] # Checkoboard pattern, TODO: may not be optimal + + # Get my device + if self.nr_proc != len(self.devices): + raise ValueError(f"Number of processes {self.nr_proc} does not equal number of devices {len(self.devices)}.") + self.device = self.devices[self.pid] + + # Make the tiles + self.tiles = {} + for tile_index in self.tile_process_map.keys(): + if self.pid == self.tile_process_map[tile_index]: + self.tiles[tile_index] = self.HostTile(self.tile_shape, self.dtype, self.padding, self.codec) + + # Make GPU tiles for copying data between CPU and GPU + if self.nr_tiles % self.nr_compute_tiles != 0: + raise ValueError( + f"Number of tiles {self.nr_tiles} does not divide number of compute tiles {self.nr_compute_tiles}. This is used for asynchronous copies." + ) + compute_array_shape = [s + 2 * p for (s, p) in zip(self.tile_shape, self.padding)] + self.compute_tiles_htd = [] + self.compute_tiles_dth = [] + self.compute_streams_htd = [] + self.compute_streams_dth = [] + self.compute_arrays = [] + self.current_compute_index = 0 + with cp.cuda.Device(self.device): + for i in range(self.nr_compute_tiles): + # Make compute tiles for copying data + compute_tile = self.DeviceTile(self.tile_shape, self.dtype, self.padding, self.codec) + self.compute_tiles_htd.append(compute_tile) + compute_tile = self.DeviceTile(self.tile_shape, self.dtype, self.padding, self.codec) + self.compute_tiles_dth.append(compute_tile) + + # Make cupy stream + self.compute_streams_htd.append(cp.cuda.Stream(non_blocking=True)) + self.compute_streams_dth.append(cp.cuda.Stream(non_blocking=True)) + + # Make compute array + + self.compute_arrays.append(cp.empty(compute_array_shape, self.dtype)) + + # Make compute tile mappings + self.compute_tile_mapping_htd = {} + self.compute_tile_mapping_dth = {} + self.compute_stream_mapping_htd = {} + + def size(self): + """Return number of allocated bytes for all host tiles.""" + return sum([tile.size() for tile in self.tiles.values()]) + + def nbytes(self): + """Return number of bytes for all host tiles.""" + return sum([tile.nbytes for tile in self.tiles.values()]) + + def compression_ratio(self): + """Return the compression ratio for all host tiles.""" + return self.nbytes() / self.size() + + def compression_ratio(self): + """Return the compression ratio aggregated over all tiles.""" + + if self.codec is None: + return 1.0 + else: + total_bytes = 0 + total_uncompressed_bytes = 0 + for tile in self.tiles.values(): + ( + tile_total_bytes_uncompressed, + tile_total_bytes_compressed, + ) = tile.compression_ratio() + total_bytes += tile_total_bytes_compressed + total_uncompressed_bytes += tile_total_bytes_uncompressed + return total_uncompressed_bytes / total_bytes + + def update_compute_index(self): + """Update the current compute index.""" + self.current_compute_index = (self.current_compute_index + 1) % self.nr_compute_tiles + + def _guess_next_tile_index(self, tile_index): + """Guess the next tile index to use for the compute array.""" + # TODO: This assumes access is sequential + tile_indices = list(self.tiles.keys()) + current_ind = tile_indices.index(tile_index) + next_ind = current_ind + 1 + if next_ind >= len(tile_indices): + return None + else: + return tile_indices[next_ind] + + def reset_queue_htd(self): + """Reset the queue for host to device copies.""" + + self.compute_tile_mapping_htd = {} + self.compute_stream_mapping_htd = {} + self.current_compute_index = 0 + + def managed_compute_tiles_htd(self, tile_index): + """Get the compute tiles needed for computation. + + Parameters + ---------- + tile_index : tuple + The tile index. + + Returns + ------- + compute_tile : ComputeTile + The compute tile needed for computation. + """ + + ################################################### + # TODO: This assumes access is sequential for tiles + ################################################### + + # Que up the next tiles + cur_tile_index = tile_index + cur_compute_index = self.current_compute_index + for i in range(self.nr_compute_tiles): + # Check if already in compute tile map and if not que it + if cur_tile_index not in self.compute_tile_mapping_htd.keys(): + # Get the store tile + tile = self.tiles[cur_tile_index] + + # Get the compute tile + compute_tile = self.compute_tiles_htd[cur_compute_index] + + # Get the compute stream + compute_stream = self.compute_streams_htd[cur_compute_index] + + # Copy the tile to the compute tile using the compute stream + with compute_stream: + tile.to_gpu_tile(compute_tile) + tile.to_gpu_tile(compute_tile) + + # Set the compute tile mapping + self.compute_tile_mapping_htd[cur_tile_index] = compute_tile + self.compute_stream_mapping_htd[cur_tile_index] = compute_stream + + # Update the tile index and compute index + cur_tile_index = self._guess_next_tile_index(cur_tile_index) + if cur_tile_index is None: + break + cur_compute_index = (cur_compute_index + 1) % self.nr_compute_tiles + + # Get the compute tile + self.compute_stream_mapping_htd[tile_index].synchronize() + compute_tile = self.compute_tile_mapping_htd[tile_index] + + # Pop the tile from the compute tile map + self.compute_tile_mapping_htd.pop(tile_index) + self.compute_stream_mapping_htd.pop(tile_index) + + # Return the compute tile + return compute_tile + + def get_compute_array(self, tile_index): + """Given a tile index, copy the tile to the compute array. + + Parameters + ---------- + tile_index : tuple + The tile index. + + Returns + ------- + compute_array : array + The compute array. + global_index : tuple + The lower bound index that the compute array corresponds to in the global array. + For example, if the compute array is the 0th tile and has padding 1, then the + global index will be (-1, -1, ..., -1). + """ + + # Get the compute tile + compute_tile = self.managed_compute_tiles_htd(tile_index) + + # Concatenate the sub-arrays to make the compute array + compute_tile.to_array(self.compute_arrays[self.current_compute_index]) + + # Return the compute array index in global array + global_index = tuple([i * s - p for (i, s, p) in zip(tile_index, self.tile_shape, self.padding)]) + + return self.compute_arrays[self.current_compute_index], global_index + + def set_tile(self, compute_array, tile_index): + """Given a tile index, copy the compute array to the tile. + + Parameters + ---------- + compute_array : array + The compute array. + tile_index : tuple + The tile index. + """ + + # Syncronize the current stream dth stream + stream = self.compute_streams_dth[self.current_compute_index] + stream.synchronize() + cp.cuda.get_current_stream().synchronize() + + # Set the compute tile to the correct one + compute_tile = self.compute_tiles_dth[self.current_compute_index] + + # Split the compute array into a tile + compute_tile.from_array(compute_array) + + # Syncronize the current stream and the compute stream + cp.cuda.get_current_stream().synchronize() + + # Copy the tile from the compute tile to the store tile + with stream: + compute_tile.to_cpu_tile(self.tiles[tile_index]) + compute_tile.to_cpu_tile(self.tiles[tile_index]) + + def update_padding(self): + """Perform a padding swap between neighboring tiles.""" + + # Get padding indices + pad_ind = self.compute_tiles_htd[0].pad_ind + + # Loop over tiles + comm_tag = 0 + for tile_index in self.tile_process_map.keys(): + # Loop over all padding + for pad_index in pad_ind: + # Get neighboring tile index + neigh_tile_index = tuple([(i + p) % s for (i, p, s) in zip(tile_index, pad_index, self.tile_dims)]) + neigh_pad_index = tuple([-p for p in pad_index]) # flip + + # 4 cases: + # 1. the tile and neighboring tile are on the same process + # 2. the tile is on this process and the neighboring tile is on another process + # 3. the tile is on another process and the neighboring tile is on this process + # 4. the tile and neighboring tile are on different processes + + # Case 1: the tile and neighboring tile are on the same process + if self.pid == self.tile_process_map[tile_index] and self.pid == self.tile_process_map[neigh_tile_index]: + # Get the tile and neighboring tile + tile = self.tiles[tile_index] + neigh_tile = self.tiles[neigh_tile_index] + + # Get pointer to padding and neighboring padding + padding = tile._padding[pad_index] + neigh_padding = neigh_tile._buf_padding[neigh_pad_index] + + # Swap padding + tile._padding[pad_index] = neigh_padding + neigh_tile._buf_padding[neigh_pad_index] = padding + + # Case 2: the tile is on this process and the neighboring tile is on another process + if self.pid == self.tile_process_map[tile_index] and self.pid != self.tile_process_map[neigh_tile_index]: + # Get the tile and padding + tile = self.tiles[tile_index] + padding = tile._padding[pad_index] + + # Send padding to neighboring process + self.comm.Send( + padding, + dest=self.tile_process_map[neigh_tile_index], + tag=comm_tag, + ) + + # Case 3: the tile is on another process and the neighboring tile is on this process + if self.pid != self.tile_process_map[tile_index] and self.pid == self.tile_process_map[neigh_tile_index]: + # Get the neighboring tile and padding + neigh_tile = self.tiles[neigh_tile_index] + neigh_padding = neigh_tile._buf_padding[neigh_pad_index] + + # Receive padding from neighboring process + self.comm.Recv( + neigh_padding, + source=self.tile_process_map[tile_index], + tag=comm_tag, + ) + + # Case 4: the tile and neighboring tile are on different processes + if self.pid != self.tile_process_map[tile_index] and self.pid != self.tile_process_map[neigh_tile_index]: + pass + + # Increment the communication tag + comm_tag += 1 + + # Shuffle padding with buffers + for tile in self.tiles.values(): + tile.swap_buf_padding() + + def get_array(self): + """Get the full array out from all the sub-arrays. This should only be used for testing.""" + + # Get the full array + if self.comm.rank == 0: + array = np.ones(self.shape, dtype=self.dtype) + else: + array = None + + # Loop over tiles + comm_tag = 0 + for tile_index in self.tile_process_map.keys(): + # Set the center array in the full array + slice_index = tuple([slice(i * s, (i + 1) * s) for (i, s) in zip(tile_index, self.tile_shape)]) + + # if tile on this process compute the center array + if self.comm.rank == self.tile_process_map[tile_index]: + # Get the tile + tile = self.tiles[tile_index] + + # Copy the tile to the compute tile + tile.to_gpu_tile(self.compute_tiles_htd[0]) + + # Get the compute array + self.compute_tiles_htd[0].to_array(self.compute_arrays[0]) + + # Get the center array + center_array = self.compute_arrays[0][tile._slice_center].get() + + # 4 cases: + # 1. the tile is on rank 0 and this process is rank 0 + # 2. the tile is on another rank and this process is rank 0 + # 3. the tile is on this rank and this process is not rank 0 + # 4. the tile is not on rank 0 and this process is not rank 0 + + # Case 1: the tile is on rank 0 + if self.comm.rank == 0 and self.tile_process_map[tile_index] == 0: + # Set the center array in the full array + array[slice_index] = center_array + + # Case 2: the tile is on another rank and this process is rank 0 + if self.comm.rank == 0 and self.tile_process_map[tile_index] != 0: + # Get the data from the other rank + center_array = np.empty(self.tile_shape, dtype=self.dtype) + self.comm.Recv(center_array, source=self.tile_process_map[tile_index], tag=comm_tag) + + # Set the center array in the full array + array[slice_index] = center_array + + # Case 3: the tile is on this rank and this process is not rank 0 + if self.comm.rank != 0 and self.tile_process_map[tile_index] == self.comm.rank: + # Send the data to rank 0 + self.comm.Send(center_array, dest=0, tag=comm_tag) + + # Case 4: the tile is not on rank 0 and this process is not rank 0 + if self.comm.rank != 0 and self.tile_process_map[tile_index] != 0: + pass + + # Update the communication tag + comm_tag += 1 + + return array diff --git a/xlb/experimental/ooc/out_of_core.py b/xlb/experimental/ooc/out_of_core.py new file mode 100644 index 00000000..bc42fab3 --- /dev/null +++ b/xlb/experimental/ooc/out_of_core.py @@ -0,0 +1,104 @@ +# Out-of-core decorator for functions that take a lot of memory + +import cupy as cp + +from xlb.experimental.ooc.ooc_array import OOCArray +from xlb.experimental.ooc.utils import ( + _cupy_to_backend, + _backend_to_cupy, +) + + +def OOCmap(comm, ref_args, add_index=False, backend="jax"): + """Decorator for out-of-core functions. + + Parameters + ---------- + comm : MPI communicator + The MPI communicator. (TODO add functionality) + ref_args : List[int] + The indices of the arguments that are OOC arrays to be written to by outputs of the function. + add_index : bool, optional + Whether to add the index of the global array to the arguments of the function. Default is False. + If true the function will take in a tuple of (array, index) instead of just the array. + backend : str, optional + The backend to use for the function. Default is 'jax'. + Options are 'jax' and 'warp'. + give_stream : bool, optional + Whether to give the function a stream to run on. Default is False. + If true the function will take in a last argument of the stream to run on. + """ + + def decorator(func): + def wrapper(*args): + # Get list of OOC arrays + ooc_array_args = [] + for arg in args: + if isinstance(arg, OOCArray): + ooc_array_args.append(arg) + + # Check that all ooc arrays are compatible + # TODO: Add better checks + for ooc_array in ooc_array_args: + if ooc_array_args[0].tile_dims != ooc_array.tile_dims: + raise ValueError(f"Tile dimensions of ooc arrays do not match. {ooc_array_args[0].tile_dims} != {ooc_array.tile_dims}") + + # Apply the function to each of the ooc arrays + for tile_index in ooc_array_args[0].tiles.keys(): + # Run through args and kwargs and replace ooc arrays with their compute arrays + new_args = [] + for arg in args: + if isinstance(arg, OOCArray): + # Get the compute array (this performs all the memory copies) + compute_array, global_index = arg.get_compute_array(tile_index) + + # Convert to backend array + compute_array = _cupy_to_backend(compute_array, backend) + + # Add index to the arguments if requested + if add_index: + compute_array = (compute_array, global_index) + + new_args.append(compute_array) + else: + new_args.append(arg) + + # Run the function + results = func(*new_args) + + # Convert the results to a tuple if not already + if not isinstance(results, tuple): + results = (results,) + + # Convert the results back to cupy arrays + results = tuple([_backend_to_cupy(result, backend) for result in results]) + + # Write the results back to the ooc array + for arg_index, result in zip(ref_args, results): + args[arg_index].set_tile(result, tile_index) + + # Update the ooc arrays compute tile index + for ooc_array in ooc_array_args: + ooc_array.update_compute_index() + + # Syncronize all processes + cp.cuda.Device().synchronize() + comm.Barrier() + + # Update the ooc arrays padding + for i, ooc_array in enumerate(ooc_array_args): + if i in ref_args: + ooc_array.update_padding() + + # Reset que + ooc_array.reset_queue_htd() + + # Return OOC arrays + if len(ref_args) == 1: + return ooc_array_args[ref_args[0]] + else: + return tuple([args[arg_index] for arg_index in ref_args]) + + return wrapper + + return decorator diff --git a/xlb/experimental/ooc/tiles/__init__.py b/xlb/experimental/ooc/tiles/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xlb/experimental/ooc/tiles/compressed_tile.py b/xlb/experimental/ooc/tiles/compressed_tile.py new file mode 100644 index 00000000..ccdd2bba --- /dev/null +++ b/xlb/experimental/ooc/tiles/compressed_tile.py @@ -0,0 +1,248 @@ +import numpy as np +import cupy as cp +import warnings + +try: + from kvikio._lib.arr import asarray +except ImportError: + warnings.warn("kvikio not installed. Compression will not work.") + +from xlb.experimental.ooc.tiles.tile import Tile +from xlb.experimental.ooc.tiles.dense_tile import DenseGPUTile +from xlb.experimental.ooc.tiles.dynamic_array import DynamicPinnedArray + + +def _decode(comp_array, dest_array, codec): + """ + Decompresses comp_array into dest_array. + + Parameters + ---------- + comp_array : cupy array + The compressed array to be decompressed. Data type uint8. + dest_array : cupy array + The storage array for the decompressed data. + codec : Codec + The codec to use for decompression. For example, `kvikio.nvcomp.CascadedManager`. + + """ + + # Store needed information + dtype = dest_array.dtype + shape = dest_array.shape + + # Reshape dest_array to match to make into buffer + dest_array = dest_array.view(cp.uint8).reshape(-1) + + # Decompress + codec._manager.decompress(asarray(dest_array), asarray(comp_array)) + + return dest_array.view(dtype).reshape(shape) + + +def _encode(array, dest_array, codec): + """ + Compresses array into dest_array. + + Parameters + ---------- + array : cupy array + The array to be compressed. + dest_array : cupy array + The storage array for the compressed data. Data type uint8. + codec : Codec + The codec to use for compression. For example, `kvikio.nvcomp.CascadedManager`. + """ + + # Make sure array is contiguous + array = cp.ascontiguousarray(array) + + # Configure compression + codec._manager.configure_compression(array.nbytes) + + # Compress + size = codec._manager.compress(asarray(array), asarray(dest_array)) + return size + + +class CompressedTile(Tile): + """A Tile where the data is stored in compressed form.""" + + def __init__(self, shape, dtype, padding, codec): + super().__init__(shape, dtype, padding, codec) + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + raise NotImplementedError + + def to_array(self, array): + """Copy a tile to a full array.""" + # Only implemented for GPU tiles + raise NotImplementedError + + def from_array(self, array): + """Copy a full array to tile.""" + # Only implemented for GPU tiles + raise NotImplementedError + + def compression_ratio(self): + """Returns the compression ratio of the tile.""" + # Get total number of bytes in tile + total_bytes = self._array.size() + for pad_ind in self.pad_ind: + total_bytes += self._padding[pad_ind].size() + + # Get total number of bytes in uncompressed tile + total_bytes_uncompressed = np.prod(self.shape) * self.dtype_itemsize + for pad_ind in self.pad_ind: + total_bytes_uncompressed += np.prod(self._padding_shape[pad_ind]) * self.dtype_itemsize + + # Return compression ratio + return total_bytes_uncompressed, total_bytes + + +class CompressedCPUTile(CompressedTile): + """A tile with cells on the CPU.""" + + def __init__(self, shape, dtype, padding, codec): + super().__init__(shape, dtype, padding, codec) + + def size(self): + """Returns the size of the tile in bytes.""" + size = self._array.size() + for pad_ind in self.pad_ind: + size += self._padding[pad_ind].size() + return size + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + # Make zero array + cp_array = cp.zeros(shape, dtype=self.dtype) + + # compress array + codec = self.codec() + compressed_cp_array = codec.compress(cp_array) + + # Allocate array on CPU + array = DynamicPinnedArray(compressed_cp_array.nbytes) + + # Copy array + compressed_cp_array.get(out=array.array) + + # add nbytes + self.nbytes += cp_array.nbytes + + # delete GPU arrays + del compressed_cp_array + del cp_array + + return array + + def to_gpu_tile(self, dst_gpu_tile): + """Copy tile to a GPU tile.""" + + # Check tile is Compressed + assert isinstance(dst_gpu_tile, CompressedGPUTile), "Destination tile must be a CompressedGPUTile" + + # Copy array + dst_gpu_tile._array[: len(self._array.array)].set(self._array.array) + dst_gpu_tile._array_bytes = self._array.nbytes + + # Copy padding + for pad_ind in self.pad_ind: + dst_gpu_tile._padding[pad_ind][: len(self._padding[pad_ind].array)].set(self._padding[pad_ind].array) + dst_gpu_tile._padding_bytes[pad_ind] = self._padding[pad_ind].nbytes + + +class CompressedGPUTile(CompressedTile): + """A sub-array with ghost cells on the GPU.""" + + def __init__(self, shape, dtype, padding, codec): + super().__init__(shape, dtype, padding, codec) + + # Allocate dense GPU tile + self.dense_gpu_tile = DenseGPUTile(shape, dtype, padding) + + # Set bytes for each array and padding + self._array_bytes = -1 + self._padding_bytes = {} + for pad_ind in self.pad_ind: + self._padding_bytes[pad_ind] = -1 + + # Set codec for each array and padding + self._array_codec = None + self._padding_codec = {} + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + nbytes = np.prod(shape) * self.dtype_itemsize + codec = self.codec() + max_compressed_buffer = codec._manager.configure_compression(nbytes)["max_compressed_buffer_size"] + array = cp.zeros((max_compressed_buffer,), dtype=np.uint8) + return array + + def to_array(self, array): + """Copy a tile to a full array.""" + + # Copy center array + if self._array_codec is None: + self._array_codec = self.codec() + self._array_codec._manager.configure_decompression_with_compressed_buffer(asarray(self._array[: self._array_bytes])) + self._array_codec.decompression_config = self._array_codec._manager.configure_decompression_with_compressed_buffer( + asarray(self._array[: self._array_bytes]) + ) + self.dense_gpu_tile._array = _decode( + self._array[: self._array_bytes], + self.dense_gpu_tile._array, + self._array_codec, + ) + array[self._slice_center] = self.dense_gpu_tile._array + + # Copy padding + for pad_ind in self.pad_ind: + if pad_ind not in self._padding_codec: + self._padding_codec[pad_ind] = self.codec() + self._padding_codec[pad_ind].decompression_config = self._padding_codec[ + pad_ind + ]._manager.configure_decompression_with_compressed_buffer(asarray(self._padding[pad_ind][: self._padding_bytes[pad_ind]])) + self.dense_gpu_tile._padding[pad_ind] = _decode( + self._padding[pad_ind][: self._padding_bytes[pad_ind]], + self.dense_gpu_tile._padding[pad_ind], + self._padding_codec[pad_ind], + ) + array[self._slice_padding_to_array[pad_ind]] = self.dense_gpu_tile._padding[pad_ind] + + def from_array(self, array): + """Copy a full array to tile.""" + + # Copy center array + if self._array_codec is None: + self._array_codec = self.codec() + self._array_codec.configure_compression(self._array.nbytes) + self._array_bytes = _encode(array[self._slice_center], self._array, self._array_codec) + + # Copy padding + for pad_ind in self.pad_ind: + if pad_ind not in self._padding_codec: + self._padding_codec[pad_ind] = self.codec() + self._padding_codec[pad_ind].configure_compression(self._padding[pad_ind].nbytes) + self._padding_bytes[pad_ind] = _encode( + array[self._slice_array_to_padding[pad_ind]], + self._padding[pad_ind], + self._padding_codec[pad_ind], + ) + + def to_cpu_tile(self, dst_cpu_tile): + """Copy tile to a CPU tile.""" + + # Check tile is Compressed + assert isinstance(dst_cpu_tile, CompressedCPUTile), "Destination tile must be a CompressedCPUTile" + + # Copy array + dst_cpu_tile._array.resize(self._array_bytes) + self._array[: self._array_bytes].get(out=dst_cpu_tile._array.array) + + # Copy padding + for pad_ind in self.pad_ind: + dst_cpu_tile._padding[pad_ind].resize(self._padding_bytes[pad_ind]) + self._padding[pad_ind][: self._padding_bytes[pad_ind]].get(out=dst_cpu_tile._padding[pad_ind].array) diff --git a/xlb/experimental/ooc/tiles/dense_tile.py b/xlb/experimental/ooc/tiles/dense_tile.py new file mode 100644 index 00000000..41fc1291 --- /dev/null +++ b/xlb/experimental/ooc/tiles/dense_tile.py @@ -0,0 +1,88 @@ +import numpy as np +import cupy as cp + +from xlb.experimental.ooc.tiles.tile import Tile + + +class DenseTile(Tile): + """A Tile where the data is stored in a dense array of the requested dtype.""" + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + raise NotImplementedError + + def to_array(self, array): + """Copy a tile to a full array.""" + # TODO: This can be done with a single kernel call, profile to see if it is faster and needs to be done. + + # Copy center array + array[self._slice_center] = self._array + + # Copy padding + for pad_ind in self.pad_ind: + array[self._slice_padding_to_array[pad_ind]] = self._padding[pad_ind] + + def from_array(self, array): + """Copy a full array to tile.""" + # TODO: This can be done with a single kernel call, profile to see if it is faster and needs to be done. + + # Copy center array + self._array[...] = array[self._slice_center] + + # Copy padding + for pad_ind in self.pad_ind: + self._padding[pad_ind][...] = array[self._slice_array_to_padding[pad_ind]] + + +class DenseCPUTile(DenseTile): + """A dense tile with cells on the CPU.""" + + def __init__(self, shape, dtype, padding, codec=None): + super().__init__(shape, dtype, padding, None) + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + # TODO: Seems hacky, but it works. Is there a better way? + mem = cp.cuda.alloc_pinned_memory(np.prod(shape) * self.dtype_itemsize) + array = np.frombuffer(mem, dtype=self.dtype, count=np.prod(shape)).reshape(shape) + self.nbytes += mem.size() + return array + + def to_gpu_tile(self, dst_gpu_tile): + """Copy tile to a GPU tile.""" + + # Check that the destination tile is on the GPU + assert isinstance(dst_gpu_tile, DenseGPUTile), "Destination tile must be on GPU" + + # Copy array + dst_gpu_tile._array.set(self._array) + + # Copy padding + for src_array, dst_gpu_array in zip(self._padding.values(), dst_gpu_tile._padding.values()): + dst_gpu_array.set(src_array) + + +class DenseGPUTile(DenseTile): + """A sub-array with ghost cells on the GPU.""" + + def __init__(self, shape, dtype, padding, codec=None): + super().__init__(shape, dtype, padding, None) + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + array = cp.zeros(shape, dtype=self.dtype) + self.nbytes += array.nbytes + return array + + def to_cpu_tile(self, dst_cpu_tile): + """Copy tile to a CPU tile.""" + + # Check that the destination tile is on the CPU + assert isinstance(dst_cpu_tile, DenseCPUTile), "Destination tile must be on CPU" + + # Copy arra + self._array.get(out=dst_cpu_tile._array) + + # Copy padding + for src_array, dst_array in zip(self._padding.values(), dst_cpu_tile._padding.values()): + src_array.get(out=dst_array) diff --git a/xlb/experimental/ooc/tiles/dynamic_array.py b/xlb/experimental/ooc/tiles/dynamic_array.py new file mode 100644 index 00000000..403d164d --- /dev/null +++ b/xlb/experimental/ooc/tiles/dynamic_array.py @@ -0,0 +1,66 @@ +# Dynamic array class for pinned memory allocation + +import math +import cupy as cp +import numpy as np + + +class DynamicArray: + """ + Dynamic pinned memory array class. + + Attributes + ---------- + nbytes : int + The number of bytes in the array. + bytes_resize : int + The number of bytes to resize the array by if the number of bytes requested exceeds the allocated number of bytes. + """ + + def __init__(self, nbytes, bytes_resize_factor=0.025): + # Set the number of bytes + self.nbytes = nbytes + self.bytes_resize_factor = bytes_resize_factor + self.bytes_resize = math.ceil(bytes_resize_factor * nbytes) + + # Set the number of bytes + self.allocated_bytes = math.ceil(nbytes / self.bytes_resize) * self.bytes_resize + + +class DynamicPinnedArray(DynamicArray): + def __init__(self, nbytes, bytes_resize_factor=0.05): + super().__init__(nbytes, bytes_resize_factor) + + # Allocate the memory + self.mem = cp.cuda.alloc_pinned_memory(self.allocated_bytes) + + # Make np array that points to the pinned memory + self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes)) + + def size(self): + return self.mem.size() + + def resize(self, nbytes): + # Set the new number of bytes + self.nbytes = nbytes + + # Check if the number of bytes requested is less than 2xbytes_resize or if the number of bytes requested exceeds the allocated number of bytes + if nbytes < (self.allocated_bytes - 2 * self.bytes_resize) or nbytes > self.allocated_bytes: + ## Free the memory + # del self.mem + + # Set the new number of allocated bytes + self.allocated_bytes = math.ceil(nbytes / self.bytes_resize) * self.bytes_resize + + # Allocate the memory + self.mem = cp.cuda.alloc_pinned_memory(self.allocated_bytes) + + # Make np array that points to the pinned memory + self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes)) + + # Set new resize number of bytes + self.bytes_resize = math.ceil(self.bytes_resize_factor * nbytes) + + # Otherwise change numpy array size + else: + self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes)) diff --git a/xlb/experimental/ooc/tiles/tile.py b/xlb/experimental/ooc/tiles/tile.py new file mode 100644 index 00000000..9bb347b9 --- /dev/null +++ b/xlb/experimental/ooc/tiles/tile.py @@ -0,0 +1,109 @@ +import cupy as cp +import itertools + + +class Tile: + """Base class for Tile with ghost cells. This tile is used to build a distributed array. + + Attributes + ---------- + shape : tuple + Shape of the tile. This will be the shape of the array without padding/ghost cells. + dtype : cp.dtype + Data type the tile represents. Note that the data data may be stored in a different + data type. For example, if it is stored in compressed form. + padding : tuple + Number of padding/ghost cells in each dimension. + """ + + def __init__(self, shape, dtype, padding, codec=None): + # Store parameters + self.shape = shape + self.dtype = dtype + self.padding = padding + self.dtype_itemsize = cp.dtype(self.dtype).itemsize + self.nbytes = 0 # Updated when array is allocated + self.codec = codec # Codec to use for compression TODO: Find better abstraction for this + + # Make center array + self._array = self.allocate_array(self.shape) + + # Make padding indices + pad_dir = [] + for i in range(len(self.shape)): + if self.padding[i] == 0: + pad_dir.append((0,)) + else: + pad_dir.append((-1, 0, 1)) + self.pad_ind = list(itertools.product(*pad_dir)) + self.pad_ind.remove((0,) * len(self.shape)) + + # Make padding and padding buffer arrays + self._padding = {} + self._buf_padding = {} + for ind in self.pad_ind: + # determine array shape + shape = [] + for i in range(len(self.shape)): + if ind[i] == -1 or ind[i] == 1: + shape.append(self.padding[i]) + else: + shape.append(self.shape[i]) + + # Make padding and padding buffer + self._padding[ind] = self.allocate_array(shape) + self._buf_padding[ind] = self.allocate_array(shape) + + # Get slicing for array copies + self._slice_center = tuple([slice(pad, pad + shape) for (pad, shape) in zip(self.padding, self.shape)]) + self._slice_padding_to_array = {} + self._slice_array_to_padding = {} + self._padding_shape = {} + for pad_ind in self.pad_ind: + slice_padding_to_array = [] + slice_array_to_padding = [] + padding_shape = [] + for pad, ind, s in zip(self.padding, pad_ind, self.shape): + if ind == -1: + slice_padding_to_array.append(slice(0, pad)) + slice_array_to_padding.append(slice(pad, 2 * pad)) + padding_shape.append(pad) + elif ind == 0: + slice_padding_to_array.append(slice(pad, s + pad)) + slice_array_to_padding.append(slice(pad, s + pad)) + padding_shape.append(s) + else: + slice_padding_to_array.append(slice(s + pad, s + 2 * pad)) + slice_array_to_padding.append(slice(s, s + pad)) + padding_shape.append(pad) + self._slice_padding_to_array[pad_ind] = tuple(slice_padding_to_array) + self._slice_array_to_padding[pad_ind] = tuple(slice_array_to_padding) + self._padding_shape[pad_ind] = tuple(padding_shape) + + def size(self): + """Returns the number of bytes allocated for the tile.""" + raise NotImplementedError + + def allocate_array(self, shape): + """Returns a cupy array with the given shape.""" + raise NotImplementedError + + def copy_tile(self, dst_tile): + """Copy a tile from one tile to another.""" + raise NotImplementedError + + def to_array(self, array): + """Copy a tile to a full array.""" + raise NotImplementedError + + def from_array(self, array): + """Copy a full array to a tile.""" + raise NotImplementedError + + def swap_buf_padding(self): + """Swap the padding buffer pointer with the padding pointer.""" + for index in self.pad_ind: + (self._buf_padding[index], self._padding[index]) = ( + self._padding[index], + self._buf_padding[index], + ) diff --git a/xlb/experimental/ooc/utils.py b/xlb/experimental/ooc/utils.py new file mode 100644 index 00000000..1179c769 --- /dev/null +++ b/xlb/experimental/ooc/utils.py @@ -0,0 +1,79 @@ +import warp as wp +import cupy as cp +import jax.dlpack as jdlpack +import jax + + +def _cupy_to_backend(cupy_array, backend): + """ + Convert cupy array to backend array + + Parameters + ---------- + cupy_array : cupy.ndarray + Input cupy array + backend : str + Backend to convert to. Options are "jax", "warp", or "cupy" + """ + + # Convert cupy array to backend array + dl_array = cupy_array.toDlpack() + if backend == "jax": + backend_array = jdlpack.from_dlpack(dl_array) + elif backend == "warp": + backend_array = wp.from_dlpack(dl_array) + elif backend == "cupy": + backend_array = cupy_array + else: + raise ValueError(f"Backend {backend} not supported") + return backend_array + + +def _backend_to_cupy(backend_array, backend): + """ + Convert backend array to cupy array + + Parameters + ---------- + backend_array : backend.ndarray + Input backend array + backend : str + Backend to convert from. Options are "jax", "warp", or "cupy" + """ + + # Convert backend array to cupy array + if backend == "jax": + (jax.device_put(0.0) + 0).block_until_ready() + dl_array = jdlpack.to_dlpack(backend_array) + elif backend == "warp": + dl_array = wp.to_dlpack(backend_array) + elif backend == "cupy": + return backend_array + else: + raise ValueError(f"Backend {backend} not supported") + cupy_array = cp.fromDlpack(dl_array) + return cupy_array + + +def _stream_to_backend(stream, backend): + """ + Convert cupy stream to backend stream + + Parameters + ---------- + stream : cupy.cuda.Stream + Input cupy stream + backend : str + Backend to convert to. Options are "jax", "warp", or "cupy" + """ + + # Convert stream to backend stream + if backend == "jax": + raise ValueError("Jax currently does not support streams") + if backend == "warp": + backend_stream = wp.Stream(cuda_stream=stream.ptr) + elif backend == "cupy": + backend_stream = stream + else: + raise ValueError(f"Backend {backend} not supported") + return backend_stream diff --git a/xlb/grid/__init__.py b/xlb/grid/__init__.py new file mode 100644 index 00000000..c4a80a8b --- /dev/null +++ b/xlb/grid/__init__.py @@ -0,0 +1,6 @@ +from xlb.grid.grid import grid_factory as grid_factory +from xlb.grid.grid import multires_grid_factory as multires_grid_factory +from xlb.grid.warp_grid import WarpGrid +from xlb.grid.jax_grid import JaxGrid + +__all__ = ["grid_factory", "WarpGrid", "JaxGrid"] diff --git a/xlb/grid/grid.py b/xlb/grid/grid.py new file mode 100644 index 00000000..9043baba --- /dev/null +++ b/xlb/grid/grid.py @@ -0,0 +1,191 @@ +""" +Grid abstraction and factory functions for XLB. + +Defines the :class:`Grid` abstract base class that every backend-specific +grid must implement, plus two factory helpers: + +* :func:`grid_factory` β€” creates a single-resolution grid for any backend. +* :func:`multires_grid_factory` β€” creates a multi-resolution grid (Neon only). +""" + +from abc import ABC, abstractmethod +from typing import Tuple, List +import numpy as np + +from xlb import DefaultConfig +from xlb.compute_backend import ComputeBackend + + +def grid_factory( + shape: Tuple[int, ...], + compute_backend: ComputeBackend = None, + velocity_set=None, + backend_config=None, +): + """Create a single-resolution grid for the specified backend. + + Parameters + ---------- + shape : tuple of int + Domain dimensions, e.g. ``(nx, ny, nz)``. + compute_backend : ComputeBackend, optional + Backend to use. Defaults to ``DefaultConfig.default_backend``. + velocity_set : VelocitySet, optional + Required for the Neon backend. + backend_config : dict, optional + Backend-specific configuration (Neon only). + + Returns + ------- + Grid + A backend-specific grid instance. + """ + compute_backend = compute_backend or DefaultConfig.default_backend + velocity_set = velocity_set or DefaultConfig.velocity_set + if compute_backend == ComputeBackend.WARP: + from xlb.grid.warp_grid import WarpGrid + + return WarpGrid(shape) + elif compute_backend == ComputeBackend.NEON: + from xlb.grid.neon_grid import NeonGrid + + return NeonGrid(shape=shape, velocity_set=velocity_set, backend_config=backend_config) + elif compute_backend == ComputeBackend.JAX: + from xlb.grid.jax_grid import JaxGrid + + return JaxGrid(shape) + + raise ValueError(f"Compute backend {compute_backend} is not supported") + + +def multires_grid_factory( + shape: Tuple[int, ...], + compute_backend: ComputeBackend = None, + velocity_set=None, + sparsity_pattern_list: List[np.ndarray] = None, + sparsity_pattern_origins=None, +): + import neon + + """Create a multi-resolution grid (Neon backend only). + + Parameters + ---------- + shape : tuple of int + Bounding-box dimensions at the finest level. + compute_backend : ComputeBackend, optional + Must be ``ComputeBackend.NEON``. + velocity_set : VelocitySet, optional + Lattice velocity set. + sparsity_pattern_list : list of np.ndarray + Active-voxel masks, one per level (finest first). + sparsity_pattern_origins : list of neon.Index_3d + Origin of each level's pattern in finest-level coordinates. + + Returns + ------- + NeonMultiresGrid + A multi-resolution Neon grid. + """ + compute_backend = compute_backend or DefaultConfig.default_backend + velocity_set = velocity_set or DefaultConfig.velocity_set + if compute_backend == ComputeBackend.NEON: + from xlb.grid.multires_grid import NeonMultiresGrid + + return NeonMultiresGrid( + shape=shape, velocity_set=velocity_set, sparsity_pattern_list=sparsity_pattern_list, sparsity_pattern_origins=sparsity_pattern_origins + ) + else: + raise ValueError(f"Compute backend {compute_backend} is not supported for multires grid") + + +class Grid(ABC): + """Abstract base class for all XLB computational grids. + + Subclasses must implement :meth:`_initialize_backend` to set up the + backend-specific data structures and :meth:`create_field` (not + enforced by ABC but expected by all operators). + + Parameters + ---------- + shape : tuple of int + Domain dimensions. + compute_backend : ComputeBackend + The compute backend this grid is associated with. + """ + + def __init__( + self, + shape: Tuple[int, ...], + compute_backend: ComputeBackend, + ): + self.shape = shape + self.dim = len(shape) + self.compute_backend = compute_backend + self._initialize_backend() + + @abstractmethod + def _initialize_backend(self): + pass + + def get_compute_backend(self): + """Return the compute backend associated with this grid.""" + return self.compute_backend + + def bounding_box_indices(self, shape=None, remove_edges=False): + """ + This function calculates the indices of the bounding box of a 2D or 3D grid. + The bounding box is defined as the set of grid points on the outer edge of the grid. + + Parameters + ---------- + remove_edges : bool, optional + If True, the nodes along the edges (not just the corners) are removed from the bounding box indices. + Default is False. + + Returns + ------- + boundingBox (dict): A dictionary where keys are the names of the bounding box faces + ("bottom", "top", "left", "right" for 2D; additional "front", "back" for 3D), and values + are numpy arrays of indices corresponding to each face. + """ + + # If shape is not give, use self.shape + if shape is None: + shape = self.shape + + # Get the shape of the grid + origin = np.array([0, 0, 0]) + bounds = np.array(shape) + if remove_edges: + origin += 1 + bounds -= 1 + slice_x = slice(origin[0], bounds[0]) + slice_y = slice(origin[1], bounds[1]) + dim = len(bounds) + + # Generate bounding box indices for each face + grid = np.indices(shape) + boundingBoxIndices = {} + + if dim == 2: + nx, ny = shape + boundingBoxIndices = { + "bottom": grid[:, slice_x, 0], + "top": grid[:, slice_x, ny - 1], + "left": grid[:, 0, slice_y], + "right": grid[:, nx - 1, slice_y], + } + elif dim == 3: + nx, ny, nz = shape + slice_z = slice(origin[2], bounds[2]) + boundingBoxIndices = { + "bottom": grid[:, slice_x, slice_y, 0].reshape(3, -1), + "top": grid[:, slice_x, slice_y, nz - 1].reshape(3, -1), + "left": grid[:, 0, slice_y, slice_z].reshape(3, -1), + "right": grid[:, nx - 1, slice_y, slice_z].reshape(3, -1), + "front": grid[:, slice_x, 0, slice_z].reshape(3, -1), + "back": grid[:, slice_x, ny - 1, slice_z].reshape(3, -1), + } + + return {k: v.tolist() for k, v in boundingBoxIndices.items()} diff --git a/xlb/grid/jax_grid.py b/xlb/grid/jax_grid.py new file mode 100644 index 00000000..9002cec3 --- /dev/null +++ b/xlb/grid/jax_grid.py @@ -0,0 +1,59 @@ +from typing import Literal +from jax.sharding import PartitionSpec as P +from jax.sharding import NamedSharding, Mesh +from jax.experimental import mesh_utils +from xlb.compute_backend import ComputeBackend + +import jax.numpy as jnp +import jax + +from xlb import DefaultConfig + + +from .grid import Grid +from xlb.precision_policy import Precision + + +class JaxGrid(Grid): + def __init__(self, shape): + super().__init__(shape, ComputeBackend.JAX) + + def _initialize_backend(self): + self.nDevices = jax.device_count() + self.compute_backend = jax.default_backend() + self.device_mesh = ( + mesh_utils.create_device_mesh((1, self.nDevices, 1)) if self.dim == 2 else mesh_utils.create_device_mesh((1, self.nDevices, 1, 1)) + ) + self.global_mesh = ( + Mesh(self.device_mesh, axis_names=("cardinality", "x", "y")) + if self.dim == 2 + else Mesh(self.device_mesh, axis_names=("cardinality", "x", "y", "z")) + ) + self.sharding = ( + NamedSharding(self.global_mesh, P("cardinality", "x", "y")) + if self.dim == 2 + else NamedSharding(self.global_mesh, P("cardinality", "x", "y", "z")) + ) + + def create_field( + self, + cardinality: int, + dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16, Precision.BOOL] = None, + fill_value=None, + ): + sharding_dim = self.shape[0] // self.nDevices + device_shape = (cardinality, sharding_dim, *self.shape[1:]) + full_shape = (cardinality, *self.shape) + arrays = [] + + dtype = dtype.jax_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.jax_dtype + + for d, index in self.sharding.addressable_devices_indices_map(full_shape).items(): + jax.default_device = d + if fill_value: + x = jnp.full(device_shape, fill_value, dtype=dtype) + else: + x = jnp.zeros(shape=device_shape, dtype=dtype) + arrays += [jax.device_put(x, d)] + jax.default_device = jax.devices()[0] + return jax.make_array_from_single_device_arrays(full_shape, self.sharding, arrays) diff --git a/xlb/grid/multires_grid.py b/xlb/grid/multires_grid.py new file mode 100644 index 00000000..8ffb695d --- /dev/null +++ b/xlb/grid/multires_grid.py @@ -0,0 +1,224 @@ +""" +Multi-resolution sparse grid backed by the Neon ``mGrid`` runtime. + +This module wraps ``neon.multires.mGrid`` and exposes it through the +:class:`Grid` interface. The grid is hierarchical: level 0 is the finest +and level *N-1* is the coarsest. Each coarser level has half the +resolution of the level below it (refinement factor 2). +""" + +import numpy as np +import warp as wp +import neon +from .grid import Grid +from xlb.precision_policy import Precision +from xlb.compute_backend import ComputeBackend +from typing import Literal, List +from xlb import DefaultConfig + + +class NeonMultiresGrid(Grid): + """Hierarchical multi-resolution grid on the Neon backend. + + Wraps ``neon.multires.mGrid``. Each level is described by a boolean + sparsity pattern (active-voxel mask) and an integer origin that + places it within the finest-level coordinate system. + + Parameters + ---------- + shape : tuple of int + Bounding-box dimensions at the **finest** level ``(nx, ny, nz)``. + velocity_set : VelocitySet + Lattice velocity set defining neighbour connectivity. + sparsity_pattern_list : list of np.ndarray + One boolean/int array per level indicating which voxels are active. + Index 0 = finest level, index *N-1* = coarsest. + sparsity_pattern_origins : list of neon.Index_3d + Origin offset for each level's pattern in the finest-level + coordinate system. + """ + + def __init__( + self, + shape, + velocity_set, + sparsity_pattern_list: List[np.ndarray], + sparsity_pattern_origins: List[neon.Index_3d], + ): + self.bk = None + self.dim = None + self.grid = None + self.velocity_set = velocity_set + self.sparsity_pattern_list = sparsity_pattern_list + self.sparsity_pattern_origins = sparsity_pattern_origins + self.count_levels = len(sparsity_pattern_list) + self.refinement_factor = 2 + + super().__init__(shape, ComputeBackend.NEON) + + def _get_velocity_set(self): + return self.velocity_set + + def _initialize_backend(self): + num_devs = 1 + dev_idx_list = list(range(num_devs)) + + if len(self.shape) == 2: + import py_neon + + self.dim = py_neon.Index_3d(self.shape[0], 1, self.shape[1]) + self.neon_stencil = [] + for q in range(self.velocity_set.q): + xval, yval = self.velocity_set._c[:, q] + self.neon_stencil.append([xval, 0, yval]) + + else: + self.dim = neon.Index_3d(self.shape[0], self.shape[1], self.shape[2]) + + self.neon_stencil = [] + for q in range(self.velocity_set.q): + xval, yval, zval = self.velocity_set._c[:, q] + self.neon_stencil.append([xval, yval, zval]) + + self.bk = neon.Backend(runtime=neon.Backend.Runtime.stream, dev_idx_list=dev_idx_list) + + self.grid = neon.multires.mGrid( + backend=self.bk, + dim=self.dim, + sparsity_pattern_list=self.sparsity_pattern_list, + sparsity_pattern_origins=self.sparsity_pattern_origins, + stencil=self.neon_stencil, + ) + # Print grid stats about voxel distribution between levels. + self.grid.print_info() + pass + + def create_field( + self, + cardinality: int, + dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16] = None, + fill_value=None, + neon_memory_type: neon.MemoryType = neon.MemoryType.host_device(), + ): + """Allocate a new multi-resolution Neon field. + + The field spans all grid levels. Each level is either zero-filled + or filled with *fill_value*. + + Parameters + ---------- + cardinality : int + Number of components per voxel. + dtype : Precision, optional + Element precision. Defaults to the store precision from the + global config. + fill_value : float, optional + Value to fill every element with. ``None`` means zero. + neon_memory_type : neon.MemoryType + Memory residency (host, device, or both). + + Returns + ------- + neon.multires.mField + The newly allocated multi-resolution field. + """ + dtype = dtype.wp_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.wp_dtype + field = self.grid.new_field( + cardinality=cardinality, + dtype=dtype, + memory_type=neon_memory_type, + ) + for l in range(self.count_levels): + if fill_value is None: + field.zero_run(l, stream_idx=0) + else: + field.fill_run(level=l, value=fill_value, stream_idx=0) + return field + + def get_neon_backend(self): + """Return the underlying ``neon.Backend`` instance.""" + return self.bk + + def level_to_shape(self, level): + """Return the bounding-box shape at the given grid level. + + Level 0 is the finest and has shape ``self.shape``. Each subsequent + level halves each dimension. + """ + # level = 0 corresponds to the finest level + return tuple(x // self.refinement_factor**level for x in self.shape) + + def boundary_indices_across_levels(self, level_data, box_side: str = "front", remove_edges: bool = False): + """ + Get indices for creating a boundary condition on the specified box side that crosses multiples levels of a multiresolution grid. + The indices are returned as a list of lists, where each sublist corresponds to a level + + Parameters + ---------- + - level_data: Level data containing the origins and sparsity patterns for each level as prepared by mesher/make_cuboid_mesh function! + - box_side: The side of the bounding box to get indices for (default is "front"). + returns: + - A list of lists, where each sublist contains the indices for the boundary condition at that level. + """ + num_levels = len(level_data) + bc_indices_list = [] + d = self.velocity_set.d # Dimensionality (2 or 3) + + # Define side configurations (adjust if your conventions differ) + if d == 3: + side_config = { + "left": {"dim": 0, "value": 0}, + "right": {"dim": 0, "value": lambda s: s[0] - 1}, + "front": {"dim": 1, "value": 0}, + "back": {"dim": 1, "value": lambda s: s[1] - 1}, + "bottom": {"dim": 2, "value": 0}, + "top": {"dim": 2, "value": lambda s: s[2] - 1}, + } + elif d == 2: + side_config = { + "left": {"dim": 0, "value": 0}, + "right": {"dim": 0, "value": lambda s: s[0] - 1}, + "bottom": {"dim": 1, "value": 0}, + "top": {"dim": 1, "value": lambda s: s[1] - 1}, + } + else: + raise ValueError(f"Unsupported dimensionality: {d}") + + if box_side not in side_config: + raise ValueError(f"Unsupported box_side: {box_side}") + + for level in range(num_levels): + mask = level_data[level][0] + origin = level_data[level][2] # Assume np.array of shape (d,) + grid_shape = self.level_to_shape(level) # tuple of length d + + conf = side_config[box_side] + dim_idx = conf["dim"] + grid_bounds = conf["value"](grid_shape) if callable(conf["value"]) else conf["value"] + + # Get local indices of active voxels + local_coords = np.nonzero(mask) # Tuple of d arrays, each of length num_active + if not local_coords[0].size: + bc_indices_list.append([]) + continue + + # Compute global coords (list of d arrays) + global_coords = [local_coords[i] + origin[i] for i in range(d)] + + # Filter: must match grid_bounds along the dimension associated with the selected box_side + cond = global_coords[dim_idx] == grid_bounds + + # If remove_edges, exclude perimeter of the face + if remove_edges: + for i in range(d): + if i != dim_idx: + cond &= (global_coords[i] > 0) & (global_coords[i] < grid_shape[i] - 1) + + # Collect filtered indices + if np.any(cond): + active_bc = [gc[cond] for gc in global_coords] + bc_indices_list.append([arr.tolist() for arr in active_bc]) + else: + bc_indices_list.append([]) + + return bc_indices_list diff --git a/xlb/grid/neon_grid.py b/xlb/grid/neon_grid.py new file mode 100644 index 00000000..eaaef9aa --- /dev/null +++ b/xlb/grid/neon_grid.py @@ -0,0 +1,136 @@ +""" +Single-resolution dense grid backed by the Neon multi-GPU runtime. + +This module wraps ``neon.dense.dGrid`` and exposes it through the +:class:`Grid` interface so that XLB operators can allocate and operate on +fields transparently. +""" + +import neon +from .grid import Grid +from xlb.precision_policy import Precision +from xlb.compute_backend import ComputeBackend +from typing import Literal +from xlb import DefaultConfig + + +class NeonGrid(Grid): + """Dense single-resolution grid on the Neon backend. + + Wraps a ``neon.dense.dGrid``. The grid is initialized with the LBM + stencil derived from the provided *velocity_set* so that Neon can + set up the correct halo exchanges for neighbour communication. + + Parameters + ---------- + shape : tuple of int + Bounding-box dimensions of the domain ``(nx, ny, nz)`` (or + ``(nx, ny)`` for 2-D). + velocity_set : VelocitySet + Lattice velocity set whose stencil defines neighbour connectivity. + backend_config : dict, optional + Neon backend configuration. Must contain ``"device_list"`` (list + of GPU device indices). Defaults to ``{"device_list": [0]}``. + """ + + def __init__( + self, + shape, + velocity_set, + backend_config=None, + ): + from .warp_grid import WarpGrid + + if backend_config is None: + backend_config = { + "device_list": [0], + "skeleton_config": neon.SkeletonConfig.OCC.none(), + } + + # check that the config dictionary has the required keys + required_keys = ["device_list"] + for key in required_keys: + if key not in backend_config: + raise ValueError(f"backend_config must contain a '{key}' key") + + # check that the device list is a list of integers + if not isinstance(backend_config["device_list"], list): + raise ValueError("backend_config['device_list'] must be a list of integers") + for device in backend_config["device_list"]: + if not isinstance(device, int): + raise ValueError("backend_config['device_list'] must be a list of integers") + + self.config = backend_config + self.bk = None + self.dim = None + self.grid = None + self.velocity_set = velocity_set + + super().__init__(shape, ComputeBackend.NEON) + + def _get_velocity_set(self): + return self.velocity_set + + def _initialize_backend(self): + dev_idx_list = self.config["device_list"] + + if len(self.shape) == 2: + import py_neon + + self.dim = py_neon.Index_3d(self.shape[0], 1, self.shape[1]) + self.neon_stencil = [] + for q in range(self.velocity_set.q): + xval, yval = self.velocity_set._c[:, q] + self.neon_stencil.append([xval, 0, yval]) + + else: + self.dim = neon.Index_3d(self.shape[0], self.shape[1], self.shape[2]) + + self.neon_stencil = [] + for q in range(self.velocity_set.q): + xval, yval, zval = self.velocity_set._c[:, q] + self.neon_stencil.append([xval, yval, zval]) + + self.bk = neon.Backend(runtime=neon.Backend.Runtime.stream, dev_idx_list=dev_idx_list) + # self.bk.info_print() + self.grid = neon.dense.dGrid(backend=self.bk, dim=self.dim, sparsity=None, stencil=self.neon_stencil) + + def create_field( + self, + cardinality: int, + dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16] = None, + fill_value=None, + ): + """Allocate a new Neon field on this grid. + + Parameters + ---------- + cardinality : int + Number of components per voxel (e.g. ``q`` for populations). + dtype : Precision, optional + Element precision. Defaults to the store precision from the + global config. + fill_value : float, optional + If provided every element is set to this value; otherwise the + field is zero-initialized. + + Returns + ------- + neon.dense.dField + The newly allocated field. + """ + dtype = dtype.wp_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.wp_dtype + field = self.grid.new_field( + cardinality=cardinality, + dtype=dtype, + ) + + if fill_value is None: + field.zero_run(stream_idx=0) + else: + field.fill_run(value=fill_value, stream_idx=0) + return field + + def get_neon_backend(self): + """Return the underlying ``neon.Backend`` instance.""" + return self.bk diff --git a/xlb/grid/warp_grid.py b/xlb/grid/warp_grid.py new file mode 100644 index 00000000..b772f0a1 --- /dev/null +++ b/xlb/grid/warp_grid.py @@ -0,0 +1,35 @@ +import warp as wp + +from .grid import Grid +from xlb.precision_policy import Precision +from xlb.compute_backend import ComputeBackend +from typing import Literal +from xlb import DefaultConfig + + +class WarpGrid(Grid): + def __init__(self, shape): + super().__init__(shape, ComputeBackend.WARP) + + def _initialize_backend(self): + pass + + def create_field( + self, + cardinality: int, + dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16] = None, + fill_value=None, + ): + dtype = dtype.wp_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.wp_dtype + + # Check if shape is 2D, and if so, append a singleton dimension to the shape + shape = (cardinality,) + (self.shape if len(self.shape) != 2 else self.shape + (1,)) + + # Pin allocations to the active Warp device (important when multiple CUDA GPUs exist). + dev = wp.get_device() + + if fill_value is None: + f = wp.zeros(shape, dtype=dtype, device=dev) + else: + f = wp.full(shape, fill_value, dtype=dtype, device=dev) + return f diff --git a/xlb/grid_backend.py b/xlb/grid_backend.py new file mode 100644 index 00000000..42cd0225 --- /dev/null +++ b/xlb/grid_backend.py @@ -0,0 +1,9 @@ +# Enum used to keep track of the compute backends + +from enum import Enum, auto + + +class GridBackend(Enum): + JAX = auto() + WARP = auto() + OOC = auto() diff --git a/xlb/helper/__init__.py b/xlb/helper/__init__.py new file mode 100644 index 00000000..452cfbce --- /dev/null +++ b/xlb/helper/__init__.py @@ -0,0 +1,10 @@ +from xlb.helper.nse_fields import create_nse_fields +from xlb.helper.initializers import initialize_eq, initialize_multires_eq, CustomInitializer, CustomMultiresInitializer +from xlb.helper.check_boundary_overlaps import check_bc_overlaps +from xlb.helper.simulation_manager import MultiresSimulationManager +from xlb.helper.ibm_helper import ( + reconstruct_mesh_from_vertices_and_faces, + transform_mesh, + prepare_immersed_boundary, + calculate_voronoi_areas, +) diff --git a/xlb/helper/check_boundary_overlaps.py b/xlb/helper/check_boundary_overlaps.py new file mode 100644 index 00000000..831939f0 --- /dev/null +++ b/xlb/helper/check_boundary_overlaps.py @@ -0,0 +1,24 @@ +import numpy as np +from xlb.compute_backend import ComputeBackend + + +def check_bc_overlaps(bclist, dim, compute_backend): + index_list = [[] for _ in range(dim)] + for bc in bclist: + if bc.indices is None: + continue + # Detect duplicates within bc.indices + index_arr = np.unique(bc.indices, axis=-1) + if index_arr.shape[-1] != len(bc.indices[0]): + if compute_backend == ComputeBackend.WARP: + raise ValueError(f"Boundary condition {bc.__class__.__name__} has duplicate indices!") + print(f"WARNING: there are duplicate indices in {bc.__class__.__name__} and hence the order in bc list matters!") + for d in range(dim): + index_list[d] += bc.indices[d] + + # Detect duplicates within bclist + index_arr = np.unique(index_list, axis=-1) + if index_arr.shape[-1] != len(index_list[0]): + if compute_backend == ComputeBackend.WARP: + raise ValueError("Boundary condition list containes duplicate indices!") + print("WARNING: there are duplicate indices in the boundary condition list and hence the order in this list matters!") diff --git a/xlb/helper/ibm_helper.py b/xlb/helper/ibm_helper.py new file mode 100644 index 00000000..8dc87469 --- /dev/null +++ b/xlb/helper/ibm_helper.py @@ -0,0 +1,239 @@ +from xlb import DefaultConfig +from xlb.grid import grid_factory +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import Precision +from typing import Tuple +import warp as wp +import trimesh +import numpy as np + + +def create_ibm_fields(grid_shape: Tuple[int, int, int], velocity_set=None, precision_policy=None): + velocity_set = velocity_set or DefaultConfig.velocity_set + compute_backend = ComputeBackend.WARP + precision_policy = precision_policy or DefaultConfig.default_precision_policy + grid = grid_factory(grid_shape, compute_backend=compute_backend) + + # Create fields + f_0 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision) + f_1 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision) + velocity_eulerian = grid.create_field(cardinality=3, dtype=precision_policy.store_precision) + missing_mask = grid.create_field(cardinality=velocity_set.q, dtype=Precision.BOOL) + bc_mask = grid.create_field(cardinality=1, dtype=Precision.UINT8) + + return grid, f_0, f_1, missing_mask, bc_mask + + +def transform_mesh(mesh, translation=None, rotation=None, rotation_order="xyz", scale=None): + """ + Transform a mesh using translation, rotation, and scaling. + + Parameters + ---------- + mesh : trimesh.Trimesh + Input triangle mesh + translation : array-like or None, shape (3,) + Translation vector [x, y, z]. If None, no translation is applied + rotation : array-like or None, shape (3,) + Rotation angles in degrees [rx, ry, rz]. If None, no rotation is applied + rotation_order : str, default='xyz' + Order of rotations. Valid options: 'xyz', 'xzy', 'yxz', 'yzx', 'zxy', 'zyx' + scale : float or array-like or None, shape (3,) + Scale factor. If float, uniform scaling is applied. + If array-like, [sx, sy, sz] for non-uniform scaling. + If None, no scaling is applied + + Returns + ------- + trimesh.Trimesh + Transformed mesh + """ + # Create a copy of the mesh to avoid modifying the original + transformed_mesh = mesh.copy() + + # Apply scaling + if scale is not None: + if isinstance(scale, (int, float)): + scale = [scale, scale, scale] + transformed_mesh.apply_scale(scale) + + # Apply rotation + if rotation is not None: + # Convert degrees to radians + rotation = np.array(rotation) * np.pi / 180.0 + + # Create rotation matrix based on the specified order + matrix = trimesh.transformations.euler_matrix(rotation[0], rotation[1], rotation[2], axes=f"r{rotation_order}") + transformed_mesh.apply_transform(matrix) + + # Apply translation + if translation is not None: + translation_matrix = trimesh.transformations.translation_matrix(translation) + transformed_mesh.apply_transform(translation_matrix) + + return transformed_mesh + + +def prepare_immersed_boundary(mesh, max_lbm_length, translation=None, rotation=None, rotation_order="xyz", scale=None): + """ + Prepare an immersed boundary from an STL file with optional transformations. + + Parameters + ---------- + mesh : trimesh.Trimesh + Input triangle mesh + max_lbm_length : float + Desired maximum length in lattice units + translation : array-like or None, shape (3,) + Translation vector [x, y, z]. If None, no translation is applied + rotation : array-like or None, shape (3,) + Rotation angles in degrees [rx, ry, rz]. If None, no rotation is applied + rotation_order : str, default='xyz' + Order of rotations. Valid options: 'xyz', 'xzy', 'yxz', 'yzx', 'zxy', 'zyx' + scale : float or array-like or None, shape (3,) + Additional scale factor applied after normalization. + If float, uniform scaling is applied. + If array-like, [sx, sy, sz] for non-uniform scaling. + If None, no additional scaling is applied + + Returns + ------- + tuple + (vertices_wp, vertex_areas_wp, faces_np) + - vertices_wp: Warp array containing vertex coordinates + - vertex_areas_wp: Warp array containing Voronoi areas for each vertex + - faces_np: NumPy array containing face indices + """ + + # Subdivide to ensure at least one vertex per cell + mesh = mesh.subdivide_to_size(max_edge=1.0, max_iter=200) + + # Calculate vertices and voronoi areas + vertices_wp, vertex_areas_wp = calculate_voronoi_areas(mesh) + + # Return the faces along with vertices and areas + return vertices_wp, vertex_areas_wp, mesh.faces + + +def calculate_voronoi_areas(mesh, check_area=True): + """ + Calculate Voronoi areas for vertices in a triangle mesh using Warp. + + Parameters + ---------- + mesh : trimesh.Trimesh + Input triangle mesh + check_area : bool, optional + Whether to check if the sum of the Voronoi areas matches the mesh area + + Returns + ------- + tuple + (vertex_areas_wp, vertex_areas) + - vertex_areas_wp: Warp array containing Voronoi areas for each vertex + - vertex_areas: NumPy array containing Voronoi areas for each vertex + """ + # Get face areas and vertices of each face + face_areas = mesh.area_faces + faces = mesh.faces + vertices = mesh.vertices + + # Define the number of vertices and faces + num_vertices = len(vertices) + num_faces = len(faces) + + @wp.kernel + def voronoi_area_kernel( + faces: wp.array2d(dtype=int), vertices: wp.array(dtype=wp.vec3), face_areas: wp.array1d(dtype=float), vertex_areas: wp.array1d(dtype=float) + ): + tid = wp.tid() + + # Get vertex indices of the face + v0 = faces[tid, 0] + v1 = faces[tid, 1] + v2 = faces[tid, 2] + + # Get vertex positions + p0 = wp.vec3(vertices[v0][0], vertices[v0][1], vertices[v0][2]) + p1 = wp.vec3(vertices[v1][0], vertices[v1][1], vertices[v1][2]) + p2 = wp.vec3(vertices[v2][0], vertices[v2][1], vertices[v2][2]) + + # Compute edge lengths of the triangle + a = wp.length(p1 - p2) + b = wp.length(p0 - p2) + c = wp.length(p0 - p1) + + # Compute area and cotangent weights + face_area = face_areas[tid] + + cot_alpha = (b**2.0 + c**2.0 - a**2.0) / (4.0 * face_area) + cot_beta = (a**2.0 + c**2.0 - b**2.0) / (4.0 * face_area) + cot_gamma = (a**2.0 + b**2.0 - c**2.0) / (4.0 * face_area) + + # Normalize the cotangent weights + total_cot = cot_alpha + cot_beta + cot_gamma + if total_cot > 0: + cot_alpha /= total_cot + cot_beta /= total_cot + cot_gamma /= total_cot + + # Distribute the face area to each vertex based on the normalized weights + wp.atomic_add(vertex_areas, v0, face_area * cot_beta / 2.0 + face_area * cot_gamma / 2.0) + wp.atomic_add(vertex_areas, v1, face_area * cot_alpha / 2.0 + face_area * cot_gamma / 2.0) + wp.atomic_add(vertex_areas, v2, face_area * cot_alpha / 2.0 + face_area * cot_beta / 2.0) + + # Convert data to Warp arrays + faces_wp = wp.array(faces, dtype=wp.int32) + vertices_wp = wp.array(vertices, dtype=wp.vec3) + face_areas_wp = wp.array(face_areas, dtype=wp.float32) + vertex_areas_wp = wp.zeros(num_vertices, dtype=wp.float32) + + # Launch the kernel + wp.launch(kernel=voronoi_area_kernel, dim=num_faces, inputs=[faces_wp, vertices_wp, face_areas_wp, vertex_areas_wp], device="cuda") + + # Validate the result + if check_area: + vertex_areas_np = vertex_areas_wp.numpy() + if abs(vertex_areas_np.sum() - mesh.area) > 1e-2: + # Copy the result back to the CPU + print("Warning: Sum of Voronoi areas does not match mesh area") + else: + print("Voronoi areas calculated successfully") + print(f"Sum of Voronoi areas: {vertex_areas_np.sum()}") + print(f"Mesh area: {mesh.area}") + + return vertices_wp, vertex_areas_wp + + +def reconstruct_mesh_from_vertices_and_faces(vertices_wp, faces_np, save_path=None): + """ + Reconstruct a trimesh from Warp vertices and NumPy-based faces. + + Parameters + ---------- + vertices_wp : wp.array + Warp array containing vertex coordinates + faces_np : np.ndarray + NumPy array containing face indices (each row is a triangle with 3 vertex indices) + save_path : str or None, optional + If provided, saves the mesh to this path with .stl extension + + Returns + ------- + trimesh.Trimesh + Reconstructed mesh + """ + # Convert Warp vertices to numpy + vertices = vertices_wp.numpy() + + # Create the mesh using the provided faces + mesh = trimesh.Trimesh(vertices=vertices, faces=faces_np) + + # Save if path is provided + if save_path is not None: + if not save_path.endswith(".stl"): + save_path += ".stl" + mesh.export(save_path) + print(f"Mesh saved to: {save_path}") + + return mesh diff --git a/xlb/helper/initializers.py b/xlb/helper/initializers.py new file mode 100644 index 00000000..3fbc3f1b --- /dev/null +++ b/xlb/helper/initializers.py @@ -0,0 +1,315 @@ +""" +Initializers for distribution function fields. + +Provides helper functions and Operator subclasses that populate +distribution-function fields with equilibrium values. Two usage patterns +are supported: + +* **Functional helpers** (`initialize_eq`, `initialize_multires_eq`) β€” + one-shot initialization used during simulation setup. +* **Operator classes** (`CustomInitializer`, `CustomMultiresInitializer`) β€” + reusable operators that can target the whole domain or a single boundary + condition region, with support for JAX, Warp, and Neon backends. +""" + +import warp as wp +from typing import Any +from xlb import DefaultConfig +from xlb.operator import Operator +from xlb.velocity_set import VelocitySet +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.equilibrium import MultiresQuadraticEquilibrium + + +def initialize_eq(f, grid, velocity_set, precision_policy, compute_backend, rho=None, u=None): + """Initialize a distribution-function field to equilibrium. + + Computes the quadratic equilibrium for the given density and velocity + fields and writes it into *f*. When *rho* or *u* are ``None`` the + defaults are uniform density 1 and zero velocity. + + Parameters + ---------- + f : field + Distribution-function field to populate (modified in-place for + Warp / Neon backends; replaced for JAX). + grid : Grid + Computational grid used to allocate temporary fields. + velocity_set : VelocitySet + Lattice velocity set (e.g. D3Q19). + precision_policy : PrecisionPolicy + Precision policy for compute / store dtypes. + compute_backend : ComputeBackend + Active compute backend (JAX, WARP, or NEON). + rho : field, optional + Density field. Defaults to uniform 1.0. + u : field, optional + Velocity field. Defaults to uniform 0.0. + + Returns + ------- + field + The initialized distribution-function field. + """ + if rho is None: + rho = grid.create_field(cardinality=1, fill_value=1.0, dtype=precision_policy.compute_precision) + if u is None: + u = grid.create_field(cardinality=velocity_set.d, fill_value=0.0, dtype=precision_policy.compute_precision) + equilibrium = QuadraticEquilibrium() + + if compute_backend == ComputeBackend.JAX: + f = equilibrium(rho, u) + elif compute_backend == ComputeBackend.WARP: + f = equilibrium(rho, u, f) + elif compute_backend == ComputeBackend.NEON: + f = equilibrium(rho, u, f) + else: + raise NotImplementedError(f"Backend {compute_backend} not implemented") + + del rho, u + + return f + + +def initialize_multires_eq(f, grid, velocity_set, precision_policy, backend, rho, u): + """Initialize a multi-resolution distribution-function field to equilibrium. + + Parameters + ---------- + f : field + Multi-resolution distribution-function field to populate. + grid : NeonMultiresGrid + Multi-resolution grid. + velocity_set : VelocitySet + Lattice velocity set. + precision_policy : PrecisionPolicy + Precision policy. + backend : ComputeBackend + Compute backend (expected to be NEON). + rho : field + Density field across all grid levels. + u : field + Velocity field across all grid levels. + + Returns + ------- + field + The initialized multi-resolution distribution-function field. + """ + equilibrium = MultiresQuadraticEquilibrium() + return equilibrium(rho, u, f, stream=0) + + +class CustomInitializer(Operator): + """Operator that initializes distribution functions to equilibrium. + + When ``bc_id == -1`` (default) the entire domain is initialized with the + given constant velocity and density. Otherwise only voxels whose + ``bc_mask`` matches *bc_id* are set while the rest receive the + weight-only equilibrium (zero velocity, unit density). + + Supports JAX, Warp, and Neon backends. + + Parameters + ---------- + constant_velocity_vector : list of float + Macroscopic velocity [ux, uy, uz] used for initialization. + constant_density : float + Macroscopic density used for initialization. + bc_id : int + Boundary-condition ID to target. ``-1`` means the whole domain. + initialization_operator : Operator, optional + Equilibrium operator to use. Defaults to ``QuadraticEquilibrium``. + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + """ + + def __init__( + self, + constant_velocity_vector=[0.0, 0.0, 0.0], + constant_density: float = 1.0, + bc_id: int = -1, + initialization_operator=None, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + ): + self.bc_id = bc_id + self.constant_velocity_vector = constant_velocity_vector + self.constant_density = constant_density + if initialization_operator is None: + compute_backend = compute_backend or DefaultConfig.default_backend + self.initialization_operator = QuadraticEquilibrium( + velocity_set=velocity_set or DefaultConfig.velocity_set, + precision_policy=precision_policy or DefaultConfig.precision_policy, + compute_backend=compute_backend if compute_backend == ComputeBackend.JAX else ComputeBackend.WARP, + ) + super().__init__(velocity_set, precision_policy, compute_backend) + + @Operator.register_backend(ComputeBackend.JAX) + def jax_implementation(self, bc_mask, f_field): + from xlb.grid import grid_factory + import jax.numpy as jnp + + grid_shape = f_field.shape[1:] + grid = grid_factory(grid_shape) + rho_init = grid.create_field(cardinality=1, fill_value=self.constant_density, dtype=self.precision_policy.compute_precision) + u_init = grid.create_field(cardinality=self.velocity_set.d, fill_value=0.0, dtype=self.precision_policy.compute_precision) + _vel = jnp.array(self.constant_velocity_vector)[(...,) + (None,) * self.velocity_set.d] + if self.bc_id == -1: + u_init += _vel + else: + u_init = jnp.where(bc_mask[0] == self.bc_id, u_init + _vel, u_init) + return self.initialization_operator(rho_init, u_init) + + def _construct_warp(self): + _q = self.velocity_set.q + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + _u = _u_vec(self.constant_velocity_vector[0], self.constant_velocity_vector[1], self.constant_velocity_vector[2]) + _rho = self.compute_dtype(self.constant_density) + _w = self.velocity_set.w + bc_id = self.bc_id + + @wp.func + def functional_local(index: Any, bc_mask: Any, f_field: Any): + # Check if the index corresponds to the outlet + if self.read_field(bc_mask, index, 0) == bc_id: + _f_init = self.initialization_operator.warp_functional(_rho, _u) + for l in range(_q): + self.write_field(f_field, index, l, self.store_dtype(_f_init[l])) + else: + # In the rest of the domain, we assume zero velocity and equilibrium distribution. + for l in range(_q): + self.write_field(f_field, index, l, self.store_dtype(_w[l])) + + @wp.func + def functional_domain(index: Any, bc_mask: Any, f_field: Any): + # If bc_id is -1, initialize the entire domain according to the custom initialization operator for the given velocity + _f_init = self.initialization_operator.warp_functional(_rho, _u) + for l in range(_q): + self.write_field(f_field, index, l, self.store_dtype(_f_init[l])) + + # Set the functional based on whether we are initializing a specific BC or the entire domain + functional = functional_local if self.bc_id != -1 else functional_domain + + # Construct the warp kernel + @wp.kernel + def kernel( + bc_mask: wp.array4d(dtype=wp.uint8), + f_field: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Set the velocity at the outlet (i.e. where i = nx-1) + functional(index, bc_mask, f_field) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, bc_mask, f_field): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[bc_mask, f_field], + dim=f_field.shape[1:], + ) + return f_field + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="CustomInitializer") + def container( + bc_mask: Any, + f_field: Any, + ): + def launcher(loader: neon.Loader): + loader.set_grid(f_field.get_grid()) + f_field_pn = loader.get_write_handle(f_field) + bc_mask_pn = loader.get_read_handle(bc_mask) + + @wp.func + def kernel(index: Any): + # apply the functional + functional(index, bc_mask_pn, f_field_pn) + + loader.declare_kernel(kernel) + + return launcher + + return _, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, bc_mask, f_field, stream=0): + import neon + + # Launch the neon container + c = self.neon_container(bc_mask, f_field) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return f_field + + +class CustomMultiresInitializer(CustomInitializer): + """Multi-resolution variant of :class:`CustomInitializer`. + + Iterates over all grid levels and initializes distribution functions + using the Neon multi-resolution container API. + """ + + def __init__( + self, + constant_velocity_vector=[0.0, 0.0, 0.0], + constant_density: float = 1.0, + bc_id: int = -1, + initialization_operator=None, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + ): + super().__init__(constant_velocity_vector, constant_density, bc_id, initialization_operator, velocity_set, precision_policy, compute_backend) + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="CustomMultiresInitializer") + def container( + bc_mask: Any, + f_field: Any, + level: Any, + ): + def launcher(loader: neon.Loader): + loader.set_mres_grid(f_field.get_grid(), level) + f_field_pn = loader.get_mres_write_handle(f_field) + bc_mask_pn = loader.get_mres_read_handle(bc_mask) + + @wp.func + def kernel(index: Any): + # apply the functional + functional(index, bc_mask_pn, f_field_pn) + + loader.declare_kernel(kernel) + + return launcher + + return _, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, bc_mask, f_field, stream=0): + import neon + + grid = bc_mask.get_grid() + for level in range(grid.num_levels): + # Launch the neon container + c = self.neon_container(bc_mask, f_field, level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return f_field diff --git a/xlb/helper/nse_fields.py b/xlb/helper/nse_fields.py new file mode 100644 index 00000000..81e01006 --- /dev/null +++ b/xlb/helper/nse_fields.py @@ -0,0 +1,55 @@ +""" +Factory function for creating the standard Navier-Stokes field arrays. + +Returns the distribution-function pair (*f_0*, *f_1*), the boundary- +condition mask, and the missing-population mask, all allocated on the +given grid and backend. +""" + +from xlb import DefaultConfig +from xlb.grid import grid_factory +from xlb.precision_policy import Precision +from xlb.compute_backend import ComputeBackend +from typing import Tuple + + +def create_nse_fields( + grid_shape: Tuple[int, int, int] = None, + grid=None, + velocity_set=None, + compute_backend=None, + precision_policy=None, +): + """Create fields for Navier-Stokes equation solver. + + Args: + grid_shape: Tuple of grid dimensions. Required if grid is not provided. + grid: Optional Grid object. If provided, will be used instead of creating new grid. + velocity_set: Optional velocity set. Defaults to DefaultConfig.velocity_set. + compute_backend: Optional compute backend. Defaults to DefaultConfig.default_backend. + precision_policy: Optional precision policy. Defaults to DefaultConfig.default_precision_policy. + + Returns: + Tuple of (grid, f_0, f_1, missing_mask, bc_mask) + """ + velocity_set = velocity_set or DefaultConfig.velocity_set + compute_backend = compute_backend or DefaultConfig.default_backend + precision_policy = precision_policy or DefaultConfig.default_precision_policy + + if grid is None: + if grid_shape is None: + raise ValueError("grid_shape must be provided when grid is None") + grid = grid_factory(grid_shape, compute_backend=compute_backend, velocity_set=velocity_set) + + # Create fields + f_0 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision) + f_1 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision) + bc_mask = grid.create_field(cardinality=1, dtype=Precision.UINT8) + if compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + # For WARP and NEON, we use UINT8 for missing mask + missing_mask = grid.create_field(cardinality=velocity_set.q, dtype=Precision.UINT8) + else: + # For JAX, we use bool for missing mask + missing_mask = grid.create_field(cardinality=velocity_set.q, dtype=Precision.BOOL) + + return grid, f_0, f_1, missing_mask, bc_mask diff --git a/xlb/helper/simulation_manager.py b/xlb/helper/simulation_manager.py new file mode 100644 index 00000000..96cb1a4f --- /dev/null +++ b/xlb/helper/simulation_manager.py @@ -0,0 +1,244 @@ +""" +High-level simulation manager for multi-resolution LBM on the Neon backend. + +:class:`MultiresSimulationManager` orchestrates the complete simulation +lifecycle: field allocation, boundary-condition setup, coalescence-factor +precomputation, and the recursive time-stepping skeleton that correctly +interleaves coarse and fine grid updates. +""" + +import warp as wp +from xlb.operator.stepper import MultiresIncompressibleNavierStokesStepper +from xlb.operator.macroscopic import MultiresMacroscopic +from xlb.mres_perf_optimization_type import MresPerfOptimizationType + + +class MultiresSimulationManager(MultiresIncompressibleNavierStokesStepper): + """Orchestrates multi-resolution LBM simulations on the Neon backend. + + Inherits from :class:`MultiresIncompressibleNavierStokesStepper` and + adds field management, omega computation across levels, and the + recursive skeleton builder that encodes the multi-resolution + time-stepping order. + + Parameters + ---------- + omega_finest : float + Relaxation parameter at the finest grid level. + grid : NeonMultiresGrid + Multi-resolution grid. + boundary_conditions : list of BoundaryCondition + Boundary conditions to apply. + collision_type : str + ``"BGK"`` or ``"KBC"``. + forcing_scheme : str + Forcing scheme (used only when *force_vector* is given). + force_vector : array-like, optional + External body force. + initializer : Operator, optional + Custom initializer for distribution functions. If ``None`` + the default equilibrium initialization is used. + mres_perf_opt : MresPerfOptimizationType + Performance optimization strategy. + """ + + def __init__( + self, + omega_finest, + grid, + boundary_conditions=[], + collision_type="BGK", + forcing_scheme="exact_difference", + force_vector=None, + initializer=None, + mres_perf_opt: MresPerfOptimizationType = MresPerfOptimizationType.NAIVE_COLLIDE_STREAM, + ): + super().__init__(grid, boundary_conditions, collision_type, forcing_scheme, force_vector) + + self.initializer = initializer + self.count_levels = grid.count_levels + self.omega_list = [self.compute_omega(omega_finest, level) for level in range(self.count_levels)] + self.mres_perf_opt = mres_perf_opt + # Create fields + self.rho = grid.create_field(cardinality=1, dtype=self.precision_policy.store_precision) + self.u = grid.create_field(cardinality=3, dtype=self.precision_policy.store_precision) + self.coalescence_factor = grid.create_field(cardinality=self.velocity_set.q, dtype=self.precision_policy.store_precision) + + for level in range(self.count_levels): + self.u.fill_run(level, 0.0, 0) + self.rho.fill_run(level, 1.0, 0) + self.coalescence_factor.fill_run(level, 0.0, 0) + + # Prepare fields + self.f_0, self.f_1, self.bc_mask, self.missing_mask = self.prepare_fields(self.rho, self.u, self.initializer) + self.prepare_coalescence_count(coalescence_factor=self.coalescence_factor, bc_mask=self.bc_mask) + + self.iteration_idx = -1 + self.macro = MultiresMacroscopic( + compute_backend=self.compute_backend, + precision_policy=self.precision_policy, + velocity_set=self.velocity_set, + ) + + # Construct the stepper skeleton + self._construct_stepper_skeleton() + + def compute_omega(self, omega_finest, level): + """ + Compute the relaxation parameter omega at a given grid level based on the finest level omega. + We select a refinement ratio of 2 where a coarse cell at level L is uniformly divided into 2^d cells + where d is the dimension. to arrive at level L - 1, or in other words βˆ†x_{L-1} = βˆ†x_L/2. + For neighboring cells that interface two grid levels, a maximum jump in grid level of βˆ†L = 1 is + allowed. Due to acoustic scaling which requires the speed of sound cs to remain constant across various grid levels, + βˆ†tL ∝ βˆ†xL and hence βˆ†t_{L-1} = βˆ†t_{L}/2. In addition, the fluid viscosity \nu must also remain constant on each + grid level which leads to the following relationship for the relaxation parameter omega at grid level L base + on the finest grid level omega_finest. + + Args: + omega_finest: Relaxation parameter at the finest grid level. + level: Current grid level (0-indexed, with 0 being the finest level). + + Returns: + Relaxation parameter omega at the specified grid level. + """ + omega0 = omega_finest + return 2 ** (level + 1) * omega0 / ((2**level - 1.0) * omega0 + 2.0) + + def export_macroscopic(self, fname_prefix): + """Compute macroscopic fields and export velocity to a VTI file. + + Parameters + ---------- + fname_prefix : str + Output filename prefix. The iteration index is appended + automatically (e.g. ``"u_"`` β†’ ``"u_42.vti"``). + """ + print(f"exporting macroscopic: #levels {self.count_levels}") + self.macro(self.f_0, self.bc_mask, self.rho, self.u, streamId=0) + + wp.synchronize() + self.u.update_host(0) + wp.synchronize() + self.u.export_vti(f"{fname_prefix}{self.iteration_idx}.vti", "u") + print("DONE exporting macroscopic") + + return + + def step(self): + """Advance the simulation by one coarsest-level timestep. + + Internally this executes the pre-compiled Neon skeleton which + performs the correct number of sub-steps at each finer level + according to the acoustic-scaling time refinement ratio. + """ + self.iteration_idx = self.iteration_idx + 1 + self.sk.run() + + def _build_recursion(self, level, app, config): + """Unified multi-resolution recursion builder. + + config keys: + finest_ops: list of (op_name, swap_fields, extra_kwargs) for level 0, + or None to treat level 0 like any coarse level. + coarse_collide_ops: list of op_names for coarse collision. + coarse_stream_ops: list of (op_name, extra_kwargs) for coarse streaming. + fuse_finest: if True, recurse once (not twice) when child is at level 0. + """ + if level < 0: + return + + omega = self.omega_list[level] + fields = dict(f_0_fd=self.f_0, f_1_fd=self.f_1, bc_mask_fd=self.bc_mask, missing_mask_fd=self.missing_mask) + fields_swapped = dict(f_0_fd=self.f_1, f_1_fd=self.f_0, bc_mask_fd=self.bc_mask, missing_mask_fd=self.missing_mask) + + if level == 0 and config["finest_ops"] is not None: + for op_name, swap, extra in config["finest_ops"]: + base = fields_swapped if swap else fields + self.add_to_app(app=app, op_name=op_name, level=level, **base, omega=omega, **extra) + return + + for op_name in config["coarse_collide_ops"]: + self.add_to_app(app=app, op_name=op_name, level=level, **fields, omega=omega, timestep=0) + + if config["fuse_finest"] and level - 1 == 0: + self._build_recursion(level - 1, app, config) + else: + self._build_recursion(level - 1, app, config) + self._build_recursion(level - 1, app, config) + + for op_name, extra in config["coarse_stream_ops"]: + self.add_to_app(app=app, op_name=op_name, level=level, **fields_swapped, **extra) + + def _construct_stepper_skeleton(self): + import neon + + """Build the Neon skeleton that encodes the recursive time-stepping order. + + The skeleton is a list of Neon container invocations that, when + executed in sequence, perform one coarsest-level timestep with the + correct sub-cycling at finer levels. The structure depends on + ``self.mres_perf_opt``. + """ + self.app = [] + + stream_abc = {"coalescence_factor": self.coalescence_factor, "timestep": 0} + + # Finest-level op descriptors: (op_name, swap_f0_f1, extra_kwargs) + fused_pull_finest = [ + ("finest_fused_pull", False, {"timestep": 0, "is_f1_the_explosion_src_field": True}), + ("finest_fused_pull", True, {"timestep": 0, "is_f1_the_explosion_src_field": False}), + ] + sfv_fused_pull_finest = [ + ("CFV_finest_fused_pull", False, {"timestep": 0, "is_f1_the_explosion_src_field": True}), + ("SFV_finest_fused_pull", False, {}), + ("CFV_finest_fused_pull", True, {"timestep": 0, "is_f1_the_explosion_src_field": False}), + ("SFV_finest_fused_pull", True, {}), + ] + + configs = { + MresPerfOptimizationType.NAIVE_COLLIDE_STREAM: { + "finest_ops": None, + "coarse_collide_ops": ["collide_coarse"], + "coarse_stream_ops": [("stream_coarse_step_ABC", stream_abc)], + "fuse_finest": False, + }, + MresPerfOptimizationType.FUSION_AT_FINEST: { + "finest_ops": fused_pull_finest, + "coarse_collide_ops": ["collide_coarse"], + "coarse_stream_ops": [("stream_coarse_step_ABC", stream_abc)], + "fuse_finest": True, + }, + MresPerfOptimizationType.FUSION_AT_FINEST_SFV: { + "finest_ops": sfv_fused_pull_finest, + "coarse_collide_ops": ["collide_coarse"], + "coarse_stream_ops": [("stream_coarse_step_ABC", stream_abc)], + "fuse_finest": True, + }, + MresPerfOptimizationType.FUSION_AT_FINEST_SFV_ALL: { + "finest_ops": sfv_fused_pull_finest, + "coarse_collide_ops": ["CFV_collide_coarse", "SFV_collide_coarse"], + "coarse_stream_ops": [("SFV_stream_coarse_step_ABC", stream_abc), ("SFV_stream_coarse_step", {})], + "fuse_finest": True, + }, + } + + config = configs.get(self.mres_perf_opt) + if config is None: + raise ValueError(f"Unknown optimization level: {self.mres_perf_opt}") + + # Pre-recursion SFV mask setup + if self.mres_perf_opt == MresPerfOptimizationType.FUSION_AT_FINEST_SFV: + wp.synchronize() + self.neon_container["SFV_reset_bc_mask"](0, self.f_0, self.f_1, self.bc_mask, self.bc_mask).run(0) + wp.synchronize() + elif self.mres_perf_opt == MresPerfOptimizationType.FUSION_AT_FINEST_SFV_ALL: + wp.synchronize() + for l in range(self.f_0.get_grid().num_levels): + self.neon_container["SFV_reset_bc_mask"](l, self.f_0, self.f_1, self.bc_mask, self.bc_mask).run(0) + wp.synchronize() + + self._build_recursion(self.count_levels - 1, self.app, config) + + bk = self.grid.get_neon_backend() + self.sk = neon.Skeleton(backend=bk) + self.sk.sequence("mres_nse_stepper", self.app) diff --git a/xlb/mres_perf_optimization_type.py b/xlb/mres_perf_optimization_type.py new file mode 100644 index 00000000..797699f5 --- /dev/null +++ b/xlb/mres_perf_optimization_type.py @@ -0,0 +1,78 @@ +""" +Multi-resolution performance-optimization strategies. + +Defines the kernel-fusion levels available for the multi-resolution LBM +stepper and provides CLI argument parsing helpers. +""" + +import argparse +from enum import Enum, auto + + +class MresPerfOptimizationType(Enum): + """ + Enumeration of available optimization strategies for the LBM solver. + + Supports parsing from either the enum member name (case-insensitive) + or its integer value, and provides a method to build the CLI parser. + """ + + NAIVE_COLLIDE_STREAM = auto() + FUSION_AT_FINEST = auto() + FUSION_AT_FINEST_SFV = auto() + FUSION_AT_FINEST_SFV_ALL = auto() + + @staticmethod + def from_string(value: str) -> "MresPerfOptimizationType": + """ + Parse a string to an OptimizationType. + + Accepts either the enum member name (case-insensitive) or its integer value. + + Args: + value: The enum name (e.g. 'naive_collide_stream') or integer value (e.g. '0'). + + Returns: + An OptimizationType member. + + Raises: + argparse.ArgumentTypeError: If the input is invalid. + """ + # Attempt to parse by name (case-insensitive) + key = value.strip().upper() + if key in MresPerfOptimizationType.__members__: + return MresPerfOptimizationType[key] + + # Attempt to parse by integer value + try: + int_value = int(value) + return MresPerfOptimizationType(int_value) + except (ValueError, KeyError): + valid_options = ", ".join(f"{member.name}({member.value})" for member in MresPerfOptimizationType) + raise argparse.ArgumentTypeError(f"Invalid OptimizationType {value!r}. Choose from: {valid_options}.") + + def __str__(self) -> str: + """ + Return a human-readable string for the enum member. + """ + return self.name + + @staticmethod + def build_arg_parser() -> argparse.ArgumentParser: + """ + Create and configure the argument parser with optimization option. + + Returns: + A configured ArgumentParser instance. + """ + parser = argparse.ArgumentParser(description="Run the LBM multiresolution simulation with specified optimizations.") + # Dynamically generate help text from enum members + valid_options = ", ".join(f"{member.name}({member.value})" for member in MresPerfOptimizationType) + parser.add_argument( + "-o", + "--optimization", + type=MresPerfOptimizationType.from_string, + default=MresPerfOptimizationType.NAIVE_COLLIDE_STREAM, + help=f"Select optimization strategy: {valid_options}", + ) + return parser diff --git a/xlb/operator/__init__.py b/xlb/operator/__init__.py new file mode 100644 index 00000000..02b8a590 --- /dev/null +++ b/xlb/operator/__init__.py @@ -0,0 +1,2 @@ +from xlb.operator.operator import Operator +from xlb.operator.parallel_operator import ParallelOperator diff --git a/xlb/operator/boundary_condition/__init__.py b/xlb/operator/boundary_condition/__init__.py new file mode 100644 index 00000000..8be2f226 --- /dev/null +++ b/xlb/operator/boundary_condition/__init__.py @@ -0,0 +1,11 @@ +from xlb.operator.boundary_condition.helper_functions_bc import HelperFunctionsBC, EncodeAuxiliaryData, MultiresEncodeAuxiliaryData +from xlb.operator.boundary_condition.boundary_condition import BoundaryCondition +from xlb.operator.boundary_condition.boundary_condition_registry import BoundaryConditionRegistry +from xlb.operator.boundary_condition.bc_equilibrium import EquilibriumBC +from xlb.operator.boundary_condition.bc_do_nothing import DoNothingBC +from xlb.operator.boundary_condition.bc_halfway_bounce_back import HalfwayBounceBackBC +from xlb.operator.boundary_condition.bc_fullway_bounce_back import FullwayBounceBackBC +from xlb.operator.boundary_condition.bc_zouhe import ZouHeBC +from xlb.operator.boundary_condition.bc_regularized import RegularizedBC +from xlb.operator.boundary_condition.bc_extrapolation_outflow import ExtrapolationOutflowBC +from xlb.operator.boundary_condition.bc_hybrid import HybridBC diff --git a/xlb/operator/boundary_condition/bc_do_nothing.py b/xlb/operator/boundary_condition/bc_do_nothing.py new file mode 100644 index 00000000..7c5ae0b5 --- /dev/null +++ b/xlb/operator/boundary_condition/bc_do_nothing.py @@ -0,0 +1,91 @@ +""" +Do-nothing boundary condition. + +Skips the streaming step at tagged boundary voxels, leaving the +populations unchanged. +""" + +import jax.numpy as jnp +from jax import jit +from functools import partial +import warp as wp +from typing import Any + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class DoNothingBC(BoundaryCondition): + """ + Do nothing boundary condition. This boundary condition skips the streaming step for the + boundary nodes. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + boundary = bc_mask == self.id + return jnp.where(boundary, f_pre, f_post) + + def _construct_warp(self): + # Construct the functional for this BC + @wp.func + def functional( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + return f_pre + + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/boundary_condition/bc_equilibrium.py b/xlb/operator/boundary_condition/bc_equilibrium.py new file mode 100644 index 00000000..85ebe92b --- /dev/null +++ b/xlb/operator/boundary_condition/bc_equilibrium.py @@ -0,0 +1,124 @@ +""" +Base class for boundary conditions in a LBM simulation. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Tuple, Any + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import Equilibrium, QuadraticEquilibrium +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class EquilibriumBC(BoundaryCondition): + """Equilibrium boundary condition. + + Sets populations at tagged voxels to the equilibrium distribution + computed from the prescribed macroscopic density *rho* and velocity + *u*. Commonly used as an inlet or outlet condition. + + Parameters + ---------- + rho : float + Prescribed macroscopic density. + u : tuple of float + Prescribed macroscopic velocity ``(ux, uy, uz)``. + equilibrium_operator : Operator, optional + Equilibrium operator. Defaults to ``QuadraticEquilibrium``. + """ + + def __init__( + self, + rho: float, + u: Tuple[float, float, float], + equilibrium_operator: Operator = None, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + # Store the equilibrium information + self.rho = rho + self.u = u + self.equilibrium_operator = QuadraticEquilibrium() if equilibrium_operator is None else equilibrium_operator + # Raise error if equilibrium operator is not a subclass of Equilibrium + if not issubclass(type(self.equilibrium_operator), Equilibrium): + raise ValueError("Equilibrium operator must be a subclass of Equilibrium") + + # Call the parent constructor + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + feq = self.equilibrium_operator(jnp.array([self.rho]), jnp.array(self.u)) + new_shape = feq.shape + (1,) * self.velocity_set.d + feq = lax.broadcast_in_dim(feq, new_shape, [0]) + boundary = bc_mask == self.id + + return jnp.where(boundary, feq, f_post) + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + _rho = self.compute_dtype(self.rho) + _u = _u_vec(self.u[0], self.u[1], self.u[2]) if self.velocity_set.d == 3 else _u_vec(self.u[0], self.u[1]) + + # Construct the functional for this BC + @wp.func + def functional( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + _f = self.equilibrium_operator.warp_functional(_rho, _u) + return _f + + # Use the parent class's kernel and pass the functional + kernel = self._construct_kernel(functional) + + return functional, kernel + + def _construct_neon(self): + # Redefine the equilibrium operators for the neon backend + # This is because the neon backend relies on the warp functionals for its operations. + self.equilibrium_operator = QuadraticEquilibrium(compute_backend=ComputeBackend.WARP) + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.WARP) + def warp_launch(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post diff --git a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py new file mode 100644 index 00000000..a1a26e0b --- /dev/null +++ b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py @@ -0,0 +1,267 @@ +""" +Extrapolation outflow boundary condition. + +Uses first-order extrapolation from the interior to set the unknown +populations at outflow boundaries, avoiding strong wave reflections. + +Reference +--------- +Geier, M. et al. (2015). "The cumulant lattice Boltzmann equation in +three dimensions: Theory and validation." *Computers & Mathematics +with Applications*, 70(4), 507-547. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Any +from collections import Counter +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class ExtrapolationOutflowBC(BoundaryCondition): + """ + Extrapolation outflow boundary condition for a lattice Boltzmann method simulation. + + This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition + that uses extrapolation to avoid strong wave reflections. + + References + ---------- + Geier, M., SchΓΆnherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three + dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507-547. + doi:10.1016/j.camwa.2015.05.001. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + # Call the parent constructor + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + # find and store the normal vector using indices + if self.compute_backend == ComputeBackend.JAX: + self._get_normal_vectors(indices) + + # Unpack the two warp functionals needed for this BC! + if self.compute_backend == ComputeBackend.WARP: + self.warp_functional, self.assemble_auxiliary_data = self.warp_functional + elif self.compute_backend == ComputeBackend.NEON: + self.neon_functional, self.assemble_auxiliary_data = self.neon_functional + + def _get_normal_vectors(self, indices): + # Get the frequency count and most common element directly + freq_counts = [Counter(coord).most_common(1)[0] for coord in indices] + + # Extract counts and elements + counts = np.array([count for _, count in freq_counts]) + elements = np.array([element for element, _ in freq_counts]) + + # Normalize the counts + self.normal = counts // counts.max() + + # Reverse the normal vector if the most frequent element is 0 + if elements[np.argmax(counts)] == 0: + self.normal *= -1 + return + + @partial(jit, static_argnums=(0,), inline=True) + def _roll(self, fld, vec): + """ + Perform rolling operation of a field with dimentions [q, nx, ny, nz] in a direction + given by vec. All q-directions are rolled at the same time. + # TODO: how to improve this for multi-gpu runs? + """ + if self.velocity_set.d == 2: + return jnp.roll(fld, (vec[0], vec[1]), axis=(1, 2)) + elif self.velocity_set.d == 3: + return jnp.roll(fld, (vec[0], vec[1], vec[2]), axis=(1, 2, 3)) + + @partial(jit, static_argnums=(0,), inline=True) + def assemble_auxiliary_data(self, f_pre, f_post, bc_mask, missing_mask): + """ + Prepare time-dependent dynamic data for imposing the boundary condition in the next iteration after streaming. + We use directions that leave the domain for storing this prepared data. + Since this function is called post-collisiotn: f_pre = f_post_stream and f_post = f_post_collision + """ + sound_speed = 1.0 / jnp.sqrt(3.0) + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + + # Roll boundary mask in the opposite of the normal vector to mask its next immediate neighbour + neighbour = self._roll(boundary, -self.normal) + + # gather post-streaming values associated with previous time-step to construct the required data for BC + fpop = jnp.where(boundary, f_pre, f_post) + fpop_neighbour = jnp.where(neighbour, f_pre, f_post) + + # With fpop_neighbour isolated, now roll it back to be positioned at the boundary for subsequent operations + fpop_neighbour = self._roll(fpop_neighbour, self.normal) + fpop_extrapolated = sound_speed * fpop_neighbour + (1.0 - sound_speed) * fpop + + # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data + opp = self.velocity_set.opp_indices + known_mask = missing_mask[opp] + f_post = jnp.where(jnp.logical_and(boundary, known_mask), fpop_extrapolated[opp], f_post) + return f_post + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + return jnp.where( + jnp.logical_and(missing_mask, boundary), + f_pre[self.velocity_set.opp_indices], + f_post, + ) + + def _construct_warp(self): + # Set local constants + sound_speed = self.compute_dtype(1.0 / wp.sqrt(3.0)) + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + + @wp.func + def get_normal_vectors( + missing_mask: Any, + ): + if wp.static(self.velocity_set.d == 3): + for l in range(_q): + if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1: + return -wp.vec3i(_c[0, l], _c[1, l], _c[2, l]) + else: + for l in range(_q): + if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1: + return -wp.vec2i(_c[0, l], _c[1, l]) + + # Construct the functionals for this BC + @wp.func + def functional( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + _f_pre: Any, + _f_post: Any, + ): + # Post-streaming values are only modified at missing direction + _f = _f_post + for l in range(self.velocity_set.q): + # If the mask is missing then take the opposite index + if missing_mask[l] == wp.uint8(1): + _f[l] = _f_pre[_opp_indices[l]] + return _f + + @wp.func + def assemble_auxiliary_data_warp( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + _f_pre: Any, + _f_post: Any, + ): + # Prepare time-dependent dynamic data for imposing the boundary condition in the next iteration after streaming. + # We use directions that leave the domain for storing this prepared data. + # Since this function is called post-collisiotn: f_pre = f_post_stream and f_post = f_post_collision + _f = _f_post + nv = get_normal_vectors(missing_mask) + for l in range(self.velocity_set.q): + if missing_mask[l] == wp.uint8(1): + # f_0 is the post-collision values of the current time-step + # Get pull index associated with the "neighbours" pull_index + pull_index = type(index)() + for d in range(self.velocity_set.d): + pull_index[d] = index[d] - (_c[d, l] + nv[d]) + # The following is the post-streaming values of the neighbor cell + f_aux = self.compute_dtype(f_0[l, pull_index[0], pull_index[1], pull_index[2]]) + _f[_opp_indices[l]] = (self.compute_dtype(1.0) - sound_speed) * _f_pre[l] + sound_speed * f_aux + return _f + + @wp.func + def assemble_auxiliary_data_neon( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + _f_pre: Any, + _f_post: Any, + level: Any = 0, + ): + # Prepare time-dependent dynamic data for imposing the boundary condition in the next iteration after streaming. + # We use directions that leave the domain for storing this prepared data. + # Since this function is called post-collisiotn: f_pre = f_post_stream and f_post = f_post_collision + _f = _f_post + nv = get_normal_vectors(missing_mask) + for lattice_dir in range(self.velocity_set.q): + if missing_mask[lattice_dir] == wp.uint8(1): + # f_0 is the post-collision values of the current time-step + # Get pull index associated with the "neighbours" pull_index + offset = wp.vec3i(-_c[0, lattice_dir], -_c[1, lattice_dir], -_c[2, lattice_dir]) + for d in range(self.velocity_set.d): + offset[d] = offset[d] - nv[d] + offset_pull_index = wp.neon_ngh_idx(wp.int8(offset[0]), wp.int8(offset[1]), wp.int8(offset[2])) + + # The following is the post-streaming values of the neighbor cell + # This function reads a field value at a given neighboring index and direction. + unused_is_valid = wp.bool(False) + f_aux = self.compute_dtype(wp.neon_read_ngh(f_0, index, offset_pull_index, lattice_dir, self.store_dtype(0.0), unused_is_valid)) + _f[_opp_indices[lattice_dir]] = (self.compute_dtype(1.0) - sound_speed) * _f_pre[lattice_dir] + sound_speed * f_aux + return _f + + kernel = self._construct_kernel(functional) + assemble_auxiliary_data = assemble_auxiliary_data_warp if self.compute_backend == ComputeBackend.WARP else assemble_auxiliary_data_neon + + return (functional, assemble_auxiliary_data), kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, _f_pre, _f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[_f_pre, _f_post, bc_mask, missing_mask], + dim=_f_pre.shape[1:], + ) + return _f_post + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py new file mode 100644 index 00000000..4b7f8f0f --- /dev/null +++ b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py @@ -0,0 +1,96 @@ +""" +Full-way bounce-back boundary condition. + +Reverses every population at tagged solid voxels, effectively +imposing a no-slip wall located *on* the grid node. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Any + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + BoundaryCondition, + ImplementationStep, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class FullwayBounceBackBC(BoundaryCondition): + """ + Full Bounce-back boundary condition for a lattice Boltzmann method simulation. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + super().__init__( + ImplementationStep.COLLISION, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + return jnp.where(boundary, f_pre[self.velocity_set.opp_indices, ...], f_post) + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _opp_indices = self.velocity_set.opp_indices + _q = wp.constant(self.velocity_set.q) + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + # Construct the functional for this BC + @wp.func + def functional( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + fliped_f = _f_vec() + for l in range(_q): + fliped_f[l] = f_pre[_opp_indices[l]] + return fliped_f + + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None diff --git a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py new file mode 100644 index 00000000..c01259e3 --- /dev/null +++ b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py @@ -0,0 +1,192 @@ +""" +Halfway bounce-back boundary condition. + +Implements the standard halfway bounce-back scheme where the no-slip +wall is located halfway between a solid node and a fluid node. +Optionally supports prescribed wall velocity (moving walls) and +interpolated variants that use wall-distance data. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Any, Union, Tuple, Callable +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, + HelperFunctionsBC, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class HalfwayBounceBackBC(BoundaryCondition): + """ + Halfway Bounce-back boundary condition for a lattice Boltzmann method simulation. + + TODO: Implement moving boundary conditions for this + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + profile: Callable = None, + prescribed_value: Union[float, Tuple[float, ...], np.ndarray] = None, + ): + # Call the parent constructor + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + # This BC needs padding for finding missing directions when imposed on a geometry that is in the domain interior + self.needs_padding = True + + # This BC class accepts both constant prescribed values of velocity with keyword "prescribed_value" or + # velocity profiles given by keyword "profile" which must be a callable function. + self.profile = profile + + # A flag to enable moving wall treatment when either "prescribed_value" or "profile" are provided. + self.needs_moving_wall_treatment = False + + if (profile is not None) or (prescribed_value is not None): + self.needs_moving_wall_treatment = True + + # Handle no-slip BCs if neither prescribed_value or profile are provided. + if prescribed_value is None and profile is None: + print(f"WARNING! Assuming no-slip condition for BC type = {self.__class__.__name__}!") + prescribed_value = [0] * self.velocity_set.d + + # Handle prescribed value if provided + if prescribed_value is not None: + if profile is not None: + raise ValueError("Cannot specify both profile and prescribed_value") + + # Ensure prescribed_value is a NumPy array of floats + if isinstance(prescribed_value, (tuple, list, np.ndarray)): + prescribed_value = np.asarray(prescribed_value, dtype=np.float64) + else: + raise ValueError("Velocity prescribed_value must be a tuple, list, or array") + + # Create a constant prescribed profile function + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + if self.velocity_set.d == 2: + prescribed_value = np.array([prescribed_value[0], prescribed_value[1], 0.0], dtype=np.float64) + prescribed_value = wp.vec(3, dtype=self.compute_dtype)(prescribed_value) + self.profile = self._create_constant_prescribed_profile(prescribed_value) + + def _create_constant_prescribed_profile(self, prescribed_value): + # JAX uses jnp dtypes; wp.vec requires Warp dtypes β€” build Warp helpers only for WARP/NEON. + if self.compute_backend == ComputeBackend.JAX: + + def prescribed_profile_jax(): + return jnp.array(prescribed_value, dtype=self.precision_policy.store_precision.jax_dtype).reshape(-1, 1) + + return prescribed_profile_jax + + _u_vec = wp.vec(3, dtype=self.compute_dtype) + + @wp.func + def prescribed_profile_warp(index: Any, time: Any): + return _u_vec(prescribed_value[0], prescribed_value[1], prescribed_value[2]) + + if self.compute_backend == ComputeBackend.WARP: + return prescribed_profile_warp + if self.compute_backend == ComputeBackend.NEON: + return prescribed_profile_warp + raise ValueError(f"Constant prescribed profile unsupported for backend {self.compute_backend}") + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + + # Add contribution due to moving_wall to f_missing + moving_wall_component = 0.0 + if self.needs_moving_wall_treatment: + u_wall = self.profile() + cu = self.velocity_set.w[:, None] * jnp.tensordot(self.velocity_set.c, u_wall, axes=(0, 0)) + cu = cu.reshape((-1,) + (1,) * (len(f_post[1:].shape) - 1)) + moving_wall_component = 6.0 * cu + + # Apply the halfway bounce-back condition + f_post = jnp.where(jnp.logical_and(missing_mask, boundary), f_pre[self.velocity_set.opp_indices] + moving_wall_component, f_post) + + return f_post + + def _construct_warp(self): + # load helper functions. Explicitly using the WARP backend for helper functions as it may also be called by the Neon backend. + bc_helper = HelperFunctionsBC(velocity_set=self.velocity_set, precision_policy=self.precision_policy, compute_backend=ComputeBackend.WARP) + + # Set local constants + _opp_indices = self.velocity_set.opp_indices + + # Construct the functional for this BC + @wp.func + def functional( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # Get wall velocity + u_wall = self.profile(index, timestep) + + # Post-streaming values are only modified at missing direction + _f = f_post + for l in range(self.velocity_set.q): + # If the mask is missing then take the opposite index + if missing_mask[l] == wp.uint8(1): + # Get the pre-streaming distribution function in oppisite direction + _f[l] = f_pre[_opp_indices[l]] + + # Add contribution due to moving_wall to f_missing + if wp.static(self.needs_moving_wall_treatment): + _f[l] += bc_helper.moving_wall_fpop_correction(u_wall, l) + + return _f + + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/boundary_condition/bc_hybrid.py b/xlb/operator/boundary_condition/bc_hybrid.py new file mode 100644 index 00000000..62584160 --- /dev/null +++ b/xlb/operator/boundary_condition/bc_hybrid.py @@ -0,0 +1,391 @@ +""" +Hybrid boundary condition combining interpolated bounce-back with regularization. + +Provides three wall-treatment strategies, selectable via *bc_method*: + +* ``"bounceback_regularized"`` β€” interpolated bounce-back + Latt regularization. +* ``"bounceback_grads"`` β€” interpolated bounce-back + Grad's approximation. +* ``"nonequilibrium_regularized"`` β€” Tao non-equilibrium bounce-back + Latt + regularization. + +All variants optionally support: + +* Moving walls (via *prescribed_value* or *profile*). +* Curved boundaries with fractional distance to the mesh surface (via + *use_mesh_distance*). +""" + +import inspect +from jax import jit +from functools import partial +import warp as wp +from typing import Any, Union, Tuple, Callable +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.macroscopic import Macroscopic +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, + HelperFunctionsBC, +) +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class HybridBC(BoundaryCondition): + """ + The hybrid BC methods in this boundary condition have been originally developed by H. Salehipour and are inspired from + various previous publications, in particular [1]. The reformulations are aimed to provide local formulations that are + computationally efficient and numerically stable at high Reynolds numbers. + + [1] Dorschner, B., Chikatamarla, S. S., BΓΆsch, F., & Karlin, I. V. (2015). Grad's approximation for moving and + stationary walls in entropic lattice Boltzmann simulations. Journal of Computational Physics, 295, 340-354. + """ + + def __init__( + self, + bc_method, + profile: Callable = None, + prescribed_value: Union[float, Tuple[float, ...], np.ndarray] = None, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + use_mesh_distance=False, + ): + """ + Parameters + ---------- + bc_method : str + Wall-treatment strategy. One of ``"bounceback_regularized"``, + ``"bounceback_grads"``, or ``"nonequilibrium_regularized"``. + profile : callable, optional + Warp function ``(index) -> u_vec`` or ``(index, timestep) -> u_vec`` + defining the wall velocity. Mutually exclusive with *prescribed_value*. + prescribed_value : float or array-like, optional + Constant wall velocity vector. Mutually exclusive with *profile*. + If neither is given, a no-slip wall is assumed. + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + indices : list of array-like, optional + Boundary voxel indices (use this **or** *mesh_vertices*, not both). + mesh_vertices : np.ndarray, optional + Mesh triangle vertices for mesh-based voxelization. + voxelization_method : MeshVoxelizationMethod, optional + Voxelization strategy (AABB, RAY, AABB_CLOSE, etc.). + use_mesh_distance : bool + If ``True``, fractional distances to the mesh surface are + computed and stored for interpolated boundary schemes. + """ + assert bc_method in [ + "bounceback_regularized", + "bounceback_grads", + "nonequilibrium_regularized", + ], f"type = {bc_method} not supported! Use 'bounceback_regularized', 'bounceback_grads' or 'nonequilibrium_regularized'." + self.bc_method = bc_method + + # Call the parent constructor + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + # Raise error if used for 2d examples: + if self.velocity_set.d == 2: + raise NotImplementedError("This BC is not implemented in 2D!") + + # Check if the compute backend is Warp + assert self.compute_backend in (ComputeBackend.WARP, ComputeBackend.NEON), "This BC is currently not supported by JAX backend!" + + # Instantiate the operator for computing macroscopic values + # Explicitly using the WARP backend for these operators as they may also be called by the Neon backend. + self.macroscopic = Macroscopic(compute_backend=ComputeBackend.WARP) + self.equilibrium = QuadraticEquilibrium(compute_backend=ComputeBackend.WARP) + + # Define BC helper functions. Explicitly using the WARP backend for helper functions as it may also be called by the Neon backend. + self.bc_helper = HelperFunctionsBC( + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=ComputeBackend.WARP, + distance_decoder_function=self._construct_distance_decoder_function(), + ) + + # A flag to enable moving wall treatment when either "prescribed_value" or "profile" are provided. + self.needs_moving_wall_treatment = False + + if (profile is not None) or (prescribed_value is not None): + self.needs_moving_wall_treatment = True + + # Handle no-slip BCs if neither prescribed_value or profile are provided. + if prescribed_value is None and profile is None: + print(f"WARNING! Assuming no-slip condition for BC type = {self.__class__.__name__}_{self.bc_method}!") + prescribed_value = [0, 0, 0] + + # Handle prescribed value if provided + if prescribed_value is not None: + assert profile is None, "Cannot specify both profile and prescribed_value" + + # Ensure prescribed_value is a NumPy array of floats + if isinstance(prescribed_value, (tuple, list, np.ndarray)): + prescribed_value = np.asarray(prescribed_value, dtype=np.float64) + else: + raise ValueError("Velocity prescribed_value must be a tuple, list, or array") + + # Handle 2D velocity sets + if self.velocity_set.d == 2: + assert len(prescribed_value) == 2, "For 2D velocity set, prescribed_value must be a tuple or array of length 2!" + prescribed_value = np.array([prescribed_value[0], prescribed_value[1], 0.0], dtype=np.float64) + + # create a constant prescribed profile + _u_vec = wp.vec(3, dtype=self.compute_dtype) + prescribed_value = _u_vec(prescribed_value) + + @wp.func + def prescribed_profile_warp(index: Any): + return _u_vec(prescribed_value[0], prescribed_value[1], prescribed_value[2]) + + profile = prescribed_profile_warp + + # Inspect the function signature and add time parameter if needed + self.is_time_dependent = False + sig = inspect.signature(profile) + if len(sig.parameters) > 1: + # We assume the profile function takes only the index as input and is hence time-independent. + # In case it is defined with more than 1 input, we assume the second input is time and create + # a wrapper function that also accepts time as a parameter. + self.is_time_dependent = True + + # This BC class accepts both constant prescribed values of velocity with keyword "prescribed_value" or + # velocity profiles given by keyword "profile" which must be a callable function. + self.profile = profile + + # Set whether this BC needs mesh distance + self.needs_mesh_distance = use_mesh_distance + + # This BC needs normalized distance to the mesh + if self.needs_mesh_distance: + # This BC needs auxiliary data recovery after streaming + self.needs_aux_recovery = True + + # If this BC is defined using indices, it would need padding in order to find missing directions + # when imposed on a geometry that is in the domain interior + if self.mesh_vertices is None: + assert self.indices is not None + assert self.needs_mesh_distance is False, 'To use mesh distance, please provide the mesh vertices using keyword "mesh_vertices"!' + assert self.voxelization_method is None, "Voxelization method is only applicable when using mesh vertices!" + self.needs_padding = True + else: + assert self.indices is None, "Cannot use indices with mesh vertices! Please provide mesh vertices only." + + # Define the profile functional + self.profile_functional = self._construct_profile_functional() + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + raise NotImplementedError(f"Operation {self.__class__.__name__} not implemented in JAX!") + + def _construct_distance_decoder_function(self): + """ + Constructs the distance decoder function for this BC. + """ + # Get the opposite indices for the velocity set + _opp_indices = self.velocity_set.opp_indices + + # Define the distance decoder function for this BC + @wp.func + def distance_decoder_function(f_1: Any, index: Any, direction: Any): + return self.read_field(f_1, index, _opp_indices[direction]) + + return distance_decoder_function + + def _construct_profile_functional(self): + """ + Get the profile functional for this BC. + TODO@Hesam: + Right now, we can impose a profile on a boundary which requires mesh-distance only if that boundary lives on the finest level. + In order to extract "level" from the "neon_field_hdl" we can use the function wp.neon_level(neon_field_hdl). This will allow us + to do the following and get rid of the above limitation. + cIdx = wp.neon_global_idx(field_neon_hdl, index) + gx = wp.neon_get_x(cIdx) // 2 ** level + gy = wp.neon_get_y(cIdx) // 2 ** level + gz = wp.neon_get_z(cIdx) // 2 ** level + """ + + @wp.func + def profile_functional_neon(f_1: Any, index: Any, timestep: Any): + # Convert neon index to warp index + warp_index = self.bc_helper.neon_index_to_warp(f_1, index) + if wp.static(self.is_time_dependent): + return self.profile(warp_index, timestep) + else: + return self.profile(warp_index) + + @wp.func + def profile_functional_warp(f_1: Any, index: Any, timestep: Any): + if wp.static(self.is_time_dependent): + return self.profile(index, timestep) + else: + return self.profile(index) + + return profile_functional_warp if self.compute_backend == ComputeBackend.WARP else profile_functional_neon + + def _construct_warp(self): + # Construct the functionals for this BC + @wp.func + def hybrid_bounceback_regularized( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # Using regularization technique [1] to represent fpop using macroscopic values derived from interpolated bounceback scheme of [2]. + # missing data in lattice Boltzmann. + # [1] Latt, J., Chopard, B., Malaspinas, O., Deville, M., Michler, A., 2008. Straight velocity + # boundaries in the lattice Boltzmann method. Physical Review E 77, 056703. + # [2] Yu, D., Mei, R., Shyy, W., 2003. A unified boundary treatment in lattice boltzmann method, + # in: 41st aerospace sciences meeting and exhibit, p. 953. + + # Apply interpolated bounceback first to find missing populations at the boundary + u_wall = self.profile_functional(f_1, index, timestep) + f_post = self.bc_helper.interpolated_bounceback( + index, + _missing_mask, + f_0, + f_1, + f_pre, + f_post, + u_wall, + wp.static(self.needs_moving_wall_treatment), + wp.static(self.needs_mesh_distance), + ) + + # Compute density, velocity using all f_post-streaming values + rho, u = self.macroscopic.warp_functional(f_post) + + # Regularize the resulting populations + feq = self.equilibrium.warp_functional(rho, u) + f_post = self.bc_helper.regularize_fpop(f_post, feq) + return f_post + + @wp.func + def hybrid_bounceback_grads( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # Using Grad's approximation [1] to represent fpop using macroscopic values derived from interpolated bounceback scheme of [2]. + # missing data in lattice Boltzmann. + # [1] Dorschner, B., Chikatamarla, S. S., BΓΆsch, F., & Karlin, I. V. (2015). Grad's approximation for moving and + # stationary walls in entropic lattice Boltzmann simulations. Journal of Computational Physics, 295, 340-354. + # [2] Yu, D., Mei, R., Shyy, W., 2003. A unified boundary treatment in lattice boltzmann method, + # in: 41st aerospace sciences meeting and exhibit, p. 953. + + # Apply interpolated bounceback first to find missing populations at the boundary + u_wall = self.profile_functional(f_1, index, timestep) + f_post = self.bc_helper.interpolated_bounceback( + index, + _missing_mask, + f_0, + f_1, + f_pre, + f_post, + u_wall, + wp.static(self.needs_moving_wall_treatment), + wp.static(self.needs_mesh_distance), + ) + + # Compute density, velocity using all f_post-streaming values + rho, u = self.macroscopic.warp_functional(f_post) + + # Compute Grad's approximation using full equation as in Eq (10) of Dorschner et al. + f_post = self.bc_helper.grads_approximate_fpop(_missing_mask, rho, u, f_post) + return f_post + + @wp.func + def hybrid_nonequilibrium_regularized( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # This boundary condition uses the method of Tao et al (2018) [1] to get unknown populations on curved boundaries (denoted here by + # interpolated_nonequilibrium_bounceback method). To further stabilize this BC, we add regularization technique of [2]. + # [1] Tao, Shi, et al. "One-point second-order curved boundary condition for lattice Boltzmann simulation of suspended particles." + # Computers & Mathematics with Applications 76.7 (2018): 1593-1607. + # [2] Latt, J., Chopard, B., Malaspinas, O., Deville, M., Michler, A., 2008. Straight velocity + # boundaries in the lattice Boltzmann method. Physical Review E 77, 056703. + + # Apply interpolated bounceback first to find missing populations at the boundary + u_wall = self.profile_functional(f_1, index, timestep) + f_post = self.bc_helper.interpolated_nonequilibrium_bounceback( + index, + _missing_mask, + f_0, + f_1, + f_pre, + f_post, + u_wall, + wp.static(self.needs_moving_wall_treatment), + wp.static(self.needs_mesh_distance), + ) + + # Compute density, velocity using all f_post-streaming values + rho, u = self.macroscopic.warp_functional(f_post) + + # Regularize the resulting populations + feq = self.equilibrium.warp_functional(rho, u) + f_post = self.bc_helper.regularize_fpop(f_post, feq) + return f_post + + if self.bc_method == "bounceback_regularized": + functional = hybrid_bounceback_regularized + elif self.bc_method == "bounceback_grads": + functional = hybrid_bounceback_grads + elif self.bc_method == "nonequilibrium_regularized": + functional = hybrid_nonequilibrium_regularized + + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, _missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, _missing_mask], + dim=f_pre.shape[1:], + ) + return f_post + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/boundary_condition/bc_regularized.py b/xlb/operator/boundary_condition/bc_regularized.py new file mode 100644 index 00000000..f5f682a4 --- /dev/null +++ b/xlb/operator/boundary_condition/bc_regularized.py @@ -0,0 +1,229 @@ +""" +Regularized boundary condition. + +A non-equilibrium bounce-back scheme with additional regularization of the +distribution function. Applicable as velocity or pressure boundary conditions. + +Reference +--------- +Latt, J. et al. (2008). "Straight velocity boundaries in the lattice +Boltzmann method." *Physical Review E*, 77(5), 056703. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Any, Union, Tuple, Callable +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition import ZouHeBC, HelperFunctionsBC +from xlb.operator.macroscopic import SecondMoment as MomentumFlux +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class RegularizedBC(ZouHeBC): + """ + Regularized boundary condition for a lattice Boltzmann method simulation. + + This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition + with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure + or velocity. + + Attributes + ---------- + name : str + The name of the boundary condition. For this class, it is "Regularized". + Qi : numpy.ndarray + The Qi tensor, which is used in the regularization of the distribution functions. + + References + ---------- + Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva. + Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the + lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703 + """ + + def __init__( + self, + bc_type, + profile: Callable = None, + prescribed_value: Union[float, Tuple[float, ...], np.ndarray] = None, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + # Call the parent constructor + super().__init__( + bc_type, + profile, + prescribed_value, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + self.momentum_flux = MomentumFlux() + + @partial(jit, static_argnums=(0,), inline=True) + def regularize_fpop(self, fpop, feq): + """ + Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop. + + Parameters + ---------- + fpop : jax.numpy.ndarray + The distribution functions. + feq : jax.numpy.ndarray + The equilibrium distribution functions. + + Returns + ------- + jax.numpy.ndarray + The regularized distribution functions. + """ + # Qi = cc - cs^2*I + dim = self.velocity_set.d + weights = self.velocity_set.w[(slice(None),) + (None,) * dim] + Qi = jnp.array(self.velocity_set.qi, dtype=self.compute_dtype) + + # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq} + f_neq = fpop - feq + PiNeq = self.momentum_flux(f_neq) + # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq) + + # Compute double dot product Qi:Pi1 + # QiPi1 = np.zeros_like(fpop) + # Pi1 = PiNeq + QiPi1 = jnp.tensordot(Qi, PiNeq, axes=(1, 0)) + + # assign all populations based on eq 45 of Latt et al (2008) + # fneq ~ f^1 + fpop1 = 9.0 / 2.0 * weights * QiPi1 + fpop_regularized = feq + fpop1 + return fpop_regularized + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # creat a mask to slice boundary cells + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + + # compute the equilibrium based on prescribed values and the type of BC + feq = self.calculate_equilibrium(f_post, missing_mask) + + # set the unknown f populations based on the non-equilibrium bounce-back method + f_post_bd = self.bounceback_nonequilibrium(f_post, feq, missing_mask) + + # Regularize the boundary fpop + f_post_bd = self.regularize_fpop(f_post_bd, feq) + + # apply bc + f_post = jnp.where(boundary, f_post_bd, f_post) + return f_post + + def _construct_warp(self): + # load helper functions. Always use warp backend for helper functions as it may also be called by the Neon backend. + bc_helper = HelperFunctionsBC(velocity_set=self.velocity_set, precision_policy=self.precision_policy, compute_backend=ComputeBackend.WARP) + # Set local constants + _d = self.velocity_set.d + + @wp.func + def functional_velocity( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # Post-streaming values are only modified at missing direction + _f = f_post + + # Find normal vector + normals = bc_helper.get_normal_vectors(_missing_mask) + + # Find the value of u from the missing directions + # Since we are only considering normal velocity, we only need to find one value (stored at the center of f_1) + # Create velocity vector by multiplying the prescribed value with the normal vector + prescribed_value = self.decoder_functional(f_1, index, _missing_mask)[0] + _u = -prescribed_value * normals + + # calculate rho + fsum = bc_helper.get_bc_fsum(_f, _missing_mask) + unormal = self.compute_dtype(0.0) + for d in range(_d): + unormal += _u[d] * normals[d] + _rho = fsum / (self.compute_dtype(1.0) + unormal) + + # impose non-equilibrium bounceback + feq = self.equilibrium_operator.warp_functional(_rho, _u) + _f = bc_helper.bounceback_nonequilibrium(_f, feq, _missing_mask) + + # Regularize the boundary fpop + _f = bc_helper.regularize_fpop(_f, feq) + return _f + + @wp.func + def functional_pressure( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + # Post-streaming values are only modified at missing direction + _f = f_post + + # Find normal vector + normals = bc_helper.get_normal_vectors(_missing_mask) + + # Find the value of rho from the missing directions + # Since we need only one scalar value, we only need to find one value (stored at the center of f_1) + _rho = self.decoder_functional(f_1, index, _missing_mask)[0] + + # calculate velocity + fsum = bc_helper.get_bc_fsum(_f, _missing_mask) + unormal = -self.compute_dtype(1.0) + fsum / _rho + _u = unormal * normals + + # impose non-equilibrium bounceback + feq = self.equilibrium_operator.warp_functional(_rho, _u) + _f = bc_helper.bounceback_nonequilibrium(_f, feq, _missing_mask) + + # Regularize the boundary fpop + _f = bc_helper.regularize_fpop(_f, feq) + return _f + + if self.bc_type == "velocity": + functional = functional_velocity + elif self.bc_type == "pressure": + functional = functional_pressure + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post diff --git a/xlb/operator/boundary_condition/bc_zouhe.py b/xlb/operator/boundary_condition/bc_zouhe.py new file mode 100644 index 00000000..8e5688c5 --- /dev/null +++ b/xlb/operator/boundary_condition/bc_zouhe.py @@ -0,0 +1,408 @@ +""" +Zou-He boundary condition. + +Sets unknown populations at velocity or pressure boundaries using +mass and momentum conservation combined with non-equilibrium +bounce-back. Commonly used for inlets and outlets. + +Reference +--------- +Zou, Q. & He, X. (1997). "On pressure and velocity boundary conditions +for the lattice Boltzmann BGK model." *Physics of Fluids*, 9(6), 1591. +""" + +import jax.numpy as jnp +from jax import jit +import jax.lax as lax +from functools import partial +import warp as wp +from typing import Any, Union, Tuple, Callable +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_condition.boundary_condition import ( + ImplementationStep, + BoundaryCondition, +) +from xlb.operator.boundary_condition import HelperFunctionsBC +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod +from xlb.operator.boundary_condition.helper_functions_bc import EncodeAuxiliaryData + + +class ZouHeBC(BoundaryCondition): + """ + Zou-He boundary condition for a lattice Boltzmann method simulation. + + This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based + on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions + based on the non-equilibrium bounce-back method. + Tangential velocity is not ensured to be zero by adding transverse contributions based on + Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher + Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds. + """ + + def __init__( + self, + bc_type, + profile: Callable = None, + prescribed_value: Union[float, Tuple[float, ...], np.ndarray] = None, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + # Important Note: it is critical to add id inside __init__ for this BC because different instantiations of this BC + # may have different types (velocity or pressure). + assert bc_type in ["velocity", "pressure"], f"type = {bc_type} not supported! Use 'pressure' or 'velocity'." + self.bc_type = bc_type + self.equilibrium_operator = QuadraticEquilibrium() + + # Call the parent constructor + super().__init__( + ImplementationStep.STREAMING, + velocity_set, + precision_policy, + compute_backend, + indices, + mesh_vertices, + voxelization_method, + ) + + # This BC class accepts both constant prescribed values of velocity with keyword "prescribed_value" or + # velocity profiles given by keyword "profile" which must be a callable function. + self.profile = profile + + # Handle prescribed value if provided + if prescribed_value is not None: + if profile is not None: + raise ValueError("Cannot specify both profile and prescribed_value") + + # Ensure prescribed_value is a NumPy array of floats + if bc_type == "velocity": + if isinstance(prescribed_value, (tuple, list, np.ndarray)): + prescribed_value = np.asarray(prescribed_value, dtype=np.float64) + else: + raise ValueError("Velocity prescribed_value must be a tuple, list, or array-like") + elif bc_type == "pressure": + if isinstance(prescribed_value, (int, float)): + prescribed_value = float(prescribed_value) + else: + raise ValueError("Pressure prescribed_value must be a scalar (int or float)") + + # Check for non-zero elements - only one element should be non-zero + non_zero_count = np.count_nonzero(prescribed_value) + if non_zero_count > 1: + raise ValueError("This BC only supports normal prescribed values (only one non-zero element allowed)") + + # Prescribed value for this BC must be: + # a single non-zero number associated with normal velocity magnitude for velocity BC OR + # a single non-zero number associated with pressure BC OR + # a vector of zeros associated with no-slip BC. + # Accounting for all scenarios here. + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + if bc_type == "velocity": + # Collapse the velocity vector down to its single non-zero + # normal component (or 0.0 for no-slip). + idx = np.nonzero(prescribed_value)[0] + prescribed_value = prescribed_value[idx][0] if idx.size else 0.0 + # Pressure already arrives as a Python float; nothing to collapse. + prescribed_value = self.precision_policy.store_precision.wp_dtype(prescribed_value) + self.prescribed_value = prescribed_value + self.profile = self._create_constant_prescribed_profile() + + if self.compute_backend == ComputeBackend.JAX: + self.prescribed_values = self.profile() + else: + # This BC needs auxiliary data initialization before streaming + self.needs_aux_init = True + + # This BC needs auxiliary data recovery after streaming + self.needs_aux_recovery = True + + # This BC needs one auxiliary data for the density or normal velocity + self.num_of_aux_data = 1 + + # Create the encoder operator for storing the auxiliary data + encode_auxiliary_data = EncodeAuxiliaryData( + self.id, + self.num_of_aux_data, + self.profile, + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + + # get decoder functional + functional_dict, _ = encode_auxiliary_data._construct_warp() + self.decoder_functional = functional_dict["decoder"] + + # This BC needs padding for finding missing directions when imposed on a geometry that is in the domain interior + self.needs_padding = True + + def _create_constant_prescribed_profile(self): + _prescribed_value = self.prescribed_value + + @wp.func + def prescribed_profile_warp(index: wp.vec3i): + return wp.vec(_prescribed_value, length=1) + + def prescribed_profile_jax(): + return jnp.array(_prescribed_value, dtype=self.precision_policy.store_precision.jax_dtype).reshape(-1, 1) + + if self.compute_backend == ComputeBackend.JAX: + return prescribed_profile_jax + elif self.compute_backend == ComputeBackend.WARP: + return prescribed_profile_warp + elif self.compute_backend == ComputeBackend.NEON: + return prescribed_profile_warp + + @partial(jit, static_argnums=(0,), inline=True) + def _get_known_middle_mask(self, missing_mask): + known_mask = missing_mask[self.velocity_set.opp_indices] + middle_mask = ~(missing_mask | known_mask) + return known_mask, middle_mask + + @partial(jit, static_argnums=(0,), inline=True) + def _get_normal_vec(self, missing_mask): + main_c = self.velocity_set.c[:, self.velocity_set.main_indices] + m = missing_mask[self.velocity_set.main_indices] + normals = -jnp.tensordot(main_c, m, axes=(-1, 0)) + return normals + + @partial(jit, static_argnums=(0, 2, 3), inline=True) + def _broadcast_prescribed_values(self, prescribed_values, prescribed_values_shape, target_shape): + """ + Broadcasts `prescribed_values` to `target_shape` following specific rules: + + - If `prescribed_values_shape` is (2, 1) or (3, 1) (for constant profiles), + broadcast along the last 2 or 3 dimensions of `target_shape` respectively. + - For other shapes, identify mismatched dimensions and broadcast only in that direction. + """ + # Determine the number of dimensions to match + num_dims_prescribed = len(prescribed_values_shape) + num_dims_target = len(target_shape) + + if num_dims_prescribed > num_dims_target: + raise ValueError("prescribed_values has more dimensions than target_shape") + + # Insert singleton dimensions after the first dimension to match target_shape + if num_dims_prescribed < num_dims_target: + # Number of singleton dimensions to add + num_singleton = num_dims_target - num_dims_prescribed + + if num_dims_prescribed == 0: + # If prescribed_values is scalar, reshape to all singleton dimensions + prescribed_values_shape = (1,) * num_dims_target + else: + # Insert singleton dimensions after the first dimension + prescribed_values_shape = (prescribed_values_shape[0], *(1,) * num_singleton, *prescribed_values_shape[1:]) + prescribed_values = prescribed_values.reshape(prescribed_values_shape) + + # Create broadcast shape based on the rules + broadcast_shape = [] + for pv_dim, tgt_dim in zip(prescribed_values_shape, target_shape): + if pv_dim == 1 or pv_dim == tgt_dim: + broadcast_shape.append(tgt_dim) + else: + raise ValueError(f"Cannot broadcast dimension {pv_dim} to {tgt_dim}") + + return jnp.broadcast_to(prescribed_values, target_shape) + + @partial(jit, static_argnums=(0,), inline=True) + def get_rho(self, fpop, missing_mask): + if self.bc_type == "velocity": + target_shape = (self.velocity_set.d,) + fpop.shape[1:] + vel = self._broadcast_prescribed_values(self.prescribed_values, self.prescribed_values.shape, target_shape) + rho = self.calculate_rho(fpop, vel, missing_mask) + elif self.bc_type == "pressure": + rho = self.prescribed_values + else: + raise ValueError(f"type = {self.bc_type} not supported! Use 'pressure' or 'velocity'.") + return rho + + @partial(jit, static_argnums=(0,), inline=True) + def get_vel(self, fpop, missing_mask): + if self.bc_type == "velocity": + target_shape = (self.velocity_set.d,) + fpop.shape[1:] + vel = self._broadcast_prescribed_values(self.prescribed_values, self.prescribed_values.shape, target_shape) + elif self.bc_type == "pressure": + rho = self.prescribed_values + vel = self.calculate_vel(fpop, rho, missing_mask) + else: + raise ValueError(f"type = {self.bc_type} not supported! Use 'pressure' or 'velocity'.") + return vel + + @partial(jit, static_argnums=(0,), inline=True) + def calculate_vel(self, fpop, rho, missing_mask): + """ + Calculate velocity based on the prescribed pressure/density (Zou/He BC) + """ + + normals = self._get_normal_vec(missing_mask) + known_mask, middle_mask = self._get_known_middle_mask(missing_mask) + fsum = jnp.sum(fpop * middle_mask, axis=0, keepdims=True) + 2.0 * jnp.sum(fpop * known_mask, axis=0, keepdims=True) + unormal = -1.0 + fsum / rho + + # Return the above unormal as a normal vector which sets the tangential velocities to zero + vel = unormal * normals + return vel + + @partial(jit, static_argnums=(0,), inline=True) + def calculate_rho(self, fpop, vel, missing_mask): + """ + Calculate density based on the prescribed velocity (Zou/He BC) + """ + normals = self._get_normal_vec(missing_mask) + known_mask, middle_mask = self._get_known_middle_mask(missing_mask) + unormal = jnp.sum(normals * vel, keepdims=True, axis=0) + fsum = jnp.sum(fpop * middle_mask, axis=0, keepdims=True) + 2.0 * jnp.sum(fpop * known_mask, axis=0, keepdims=True) + rho = fsum / (1.0 + unormal) + return rho + + @partial(jit, static_argnums=(0,), inline=True) + def calculate_equilibrium(self, f_post, missing_mask): + """ + This is the ZouHe method of calculating the missing macroscopic variables at the boundary. + """ + rho = self.get_rho(f_post, missing_mask) + vel = self.get_vel(f_post, missing_mask) + + feq = self.equilibrium_operator(rho, vel) + return feq + + @partial(jit, static_argnums=(0,), inline=True) + def bounceback_nonequilibrium(self, fpop, feq, missing_mask): + """ + Calculate unknown populations using bounce-back of non-equilibrium populations + a la original Zou & He formulation + """ + opp = self.velocity_set.opp_indices + fknown = fpop[opp] + feq - feq[opp] + fpop = jnp.where(missing_mask, fknown, fpop) + return fpop + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # creat a mask to slice boundary cells + boundary = bc_mask == self.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + + # compute the equilibrium based on prescribed values and the type of BC + feq = self.calculate_equilibrium(f_post, missing_mask) + + # set the unknown f populations based on the non-equilibrium bounce-back method + f_post_bd = self.bounceback_nonequilibrium(f_post, feq, missing_mask) + f_post = jnp.where(boundary, f_post_bd, f_post) + return f_post + + def _construct_warp(self): + # load helper functions. Always use warp backend for helper functions as it may also be called by the Neon backend. + bc_helper = HelperFunctionsBC(velocity_set=self.velocity_set, precision_policy=self.precision_policy, compute_backend=ComputeBackend.WARP) + + # Set local constants + _d = self.velocity_set.d + + @wp.func + def functional_velocity( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + _f_pre: Any, + _f_post: Any, + ): + # Post-streaming values are only modified at missing direction + _f = _f_post + + # Find normal vector + normals = bc_helper.get_normal_vectors(_missing_mask) + + # calculate rho + fsum = bc_helper.get_bc_fsum(_f, _missing_mask) + unormal = self.compute_dtype(0.0) + + # Find the value of u from the missing directions + # Since we are only considering normal velocity, we only need to find one value (stored at the center of f_1) + # Create velocity vector by multiplying the prescribed value with the normal vector + prescribed_value = self.decoder_functional(f_1, index, _missing_mask)[0] + _u = -prescribed_value * normals + + for d in range(_d): + unormal += _u[d] * normals[d] + + _rho = fsum / (self.compute_dtype(1.0) + unormal) + + # impose non-equilibrium bounceback + _feq = self.equilibrium_operator.warp_functional(_rho, _u) + _f = bc_helper.bounceback_nonequilibrium(_f, _feq, _missing_mask) + return _f + + @wp.func + def functional_pressure( + index: Any, + timestep: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + _f_pre: Any, + _f_post: Any, + ): + # Post-streaming values are only modified at missing direction + _f = _f_post + + # Find normal vector + normals = bc_helper.get_normal_vectors(_missing_mask) + + # Find the value of rho from the missing directions + # Since we need only one scalar value, we only need to find one value (stored at the center of f_1) + _rho = self.decoder_functional(f_1, index, _missing_mask)[0] + + # calculate velocity + fsum = bc_helper.get_bc_fsum(_f, _missing_mask) + unormal = -self.compute_dtype(1.0) + fsum / _rho + _u = unormal * normals + + # impose non-equilibrium bounceback + feq = self.equilibrium_operator.warp_functional(_rho, _u) + _f = bc_helper.bounceback_nonequilibrium(_f, feq, _missing_mask) + return _f + + if self.bc_type == "velocity": + functional = functional_velocity + elif self.bc_type == "pressure": + functional = functional_pressure + + kernel = self._construct_kernel(functional) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_pre, f_post, bc_mask, missing_mask], + dim=f_pre.shape[1:], + ) + return f_post + + def _construct_neon(self): + # Redefine the quadratic eq operator for the neon backend + # This is because the neon backend relies on the warp functionals for its operations. + self.equilibrium_operator = QuadraticEquilibrium(compute_backend=ComputeBackend.WARP) + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_pre, f_post, bc_mask, missing_mask): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/boundary_condition/boundary_condition.py b/xlb/operator/boundary_condition/boundary_condition.py new file mode 100644 index 00000000..4ac2a96a --- /dev/null +++ b/xlb/operator/boundary_condition/boundary_condition.py @@ -0,0 +1,180 @@ +""" +Base class for boundary conditions in a Lattice Boltzmann simulation. + +Every concrete BC inherits from :class:`BoundaryCondition`, which provides +a registration mechanism, helper-function access, and the boilerplate +needed to encode auxiliary data into the ``f_1`` buffer. +""" + +from enum import Enum, auto +import warp as wp +from typing import Any +from jax import jit +from functools import partial +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb import DefaultConfig +from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry +from xlb.operator.boundary_condition import HelperFunctionsBC +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod + + +class ImplementationStep(Enum): + """At which algorithmic stage the boundary condition is applied.""" + + COLLISION = auto() + STREAMING = auto() + + +class BoundaryCondition(Operator): + """Abstract base class for all LBM boundary conditions. + + Each BC is registered with a unique numeric *id* and annotated with: + + * ``implementation_step`` - whether it executes after streaming or after + collision. + * ``needs_aux_recovery`` / ``needs_aux_init`` - whether the BC stores + auxiliary data in the ``f_1`` distribution buffer. + + Parameters + ---------- + implementation_step : ImplementationStep + Phase in the LBM algorithm where this BC is applied. + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + indices : array-like, optional + Explicit voxel indices for this BC. + mesh_vertices : array-like, optional + Mesh vertices for geometry-based BCs. + voxelization_method : MeshVoxelizationMethod, optional + Voxelization strategy when *mesh_vertices* is provided. + """ + + def __init__( + self, + implementation_step: ImplementationStep, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + indices=None, + mesh_vertices=None, + voxelization_method: MeshVoxelizationMethod = None, + ): + self.id = boundary_condition_registry.register_boundary_condition(self.__class__.__name__ + "_" + str(hash(self))) + velocity_set = velocity_set or DefaultConfig.velocity_set + precision_policy = precision_policy or DefaultConfig.default_precision_policy + compute_backend = compute_backend or DefaultConfig.default_backend + + super().__init__(velocity_set, precision_policy, compute_backend) + + # Set the BC indices + self.indices = indices + self.mesh_vertices = mesh_vertices + + # Set the implementation step + self.implementation_step = implementation_step + + # A flag to indicate whether bc indices need to be padded in both normal directions to identify missing directions + # when inside/outside of the geometry is not known + self.needs_padding = False + + # A flag for BCs that need normalized distance between the grid and a mesh (to be set to True if applicable inside each BC) + self.needs_mesh_distance = False + + # A flag for BCs that need auxiliary data initialization before stepper + self.needs_aux_init = False + + # A flag to track if the BC is initialized with auxiliary data + self.is_initialized_with_aux_data = False + + # Number of auxiliary data needed for the BC (for prescribed values) + self.num_of_aux_data = 0 + + # A flag for BCs that need auxiliary data recovery after streaming + self.needs_aux_recovery = False + + # Voxelization method. For BC's specified on a mesh, the user can specify the voxelization scheme. + # Currently we support three methods based on (a) aabb method (b) ray casting and (c) winding number. + self.voxelization_method = voxelization_method + + # Construct a default warp functional for assembling auxiliary data if needed + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + + @wp.func + def assemble_auxiliary_data( + index: Any, + timestep: Any, + missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + level: Any = 0, + ): + return f_post + + self.assemble_auxiliary_data = assemble_auxiliary_data + + def pad_indices(self): + """ + This method pads the indices to ensure that the boundary condition can be applied correctly. + It is used to find missing directions in indices_boundary_masker when the BC is imposed on a + geometry that is in the domain interior. + """ + _d = self.velocity_set.d + bc_indices = np.array(self.indices) + lattice_velocity_np = self.velocity_set._c + if self.needs_padding: + bc_indices_padded = bc_indices[:, :, None] + lattice_velocity_np[:, None, :] + return np.unique(bc_indices_padded.reshape(_d, -1), axis=1) + else: + return bc_indices + + @partial(jit, static_argnums=(0,), inline=True) + def assemble_auxiliary_data(self, f_pre, f_post, bc_mask, missing_mask): + """ + A placeholder function for prepare the auxiliary distribution functions for the boundary condition. + currently being called after collision only. + """ + return f_post + + def _construct_kernel(self, functional): + """ + Constructs the warp kernel for the boundary condition. + The functional is specific to each boundary condition and should be passed as an argument. + """ + bc_helper = HelperFunctionsBC(velocity_set=self.velocity_set, precision_policy=self.precision_policy, compute_backend=self.compute_backend) + _id = wp.uint8(self.id) + + # Construct the warp kernel + @wp.kernel + def kernel( + f_pre: wp.array4d(dtype=Any), + f_post: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # read tid data + _f_pre, _f_post, _boundary_id, _missing_mask = bc_helper.get_bc_thread_data(f_pre, f_post, bc_mask, missing_mask, index) + + # Apply the boundary condition + if _boundary_id == _id: + timestep = 0 + _f = functional(index, timestep, _missing_mask, f_pre, f_post, _f_pre, _f_post) + else: + _f = _f_post + + # Write the result + for l in range(self.velocity_set.q): + f_post[l, index[0], index[1], index[2]] = self.store_dtype(_f[l]) + + return kernel diff --git a/xlb/operator/boundary_condition/boundary_condition_registry.py b/xlb/operator/boundary_condition/boundary_condition_registry.py new file mode 100644 index 00000000..6238fc58 --- /dev/null +++ b/xlb/operator/boundary_condition/boundary_condition_registry.py @@ -0,0 +1,30 @@ +""" +Registry for boundary conditions in a LBM simulation. +""" + + +class BoundaryConditionRegistry: + """ + Registry for boundary conditions in a LBM simulation. + """ + + def __init__( + self, + ): + self.id_to_bc = {} # Maps id number to boundary condition + self.bc_to_id = {} # Maps boundary condition to id number + self.next_id = 1 # 0 is reserved for no boundary condition + + def register_boundary_condition(self, boundary_condition): + """ + Register a boundary condition. + """ + _id = self.next_id + self.next_id += 1 + self.id_to_bc[_id] = boundary_condition + self.bc_to_id[boundary_condition] = _id + print(f"registered bc {boundary_condition} with id {_id}") + return _id + + +boundary_condition_registry = BoundaryConditionRegistry() diff --git a/xlb/operator/boundary_condition/helper_functions_bc.py b/xlb/operator/boundary_condition/helper_functions_bc.py new file mode 100644 index 00000000..be25db1c --- /dev/null +++ b/xlb/operator/boundary_condition/helper_functions_bc.py @@ -0,0 +1,647 @@ +""" +Warp/Neon helper functions shared by multiple boundary conditions. + +:class:`HelperFunctionsBC` exposes ``@wp.func`` helpers for bounce-back, +regularization, Grad's approximation, moving-wall corrections, +interpolated BCs, and BC thread-data loading. These are used as building +blocks by the concrete BC classes. + +Also contains :class:`EncodeAuxiliaryData` and +:class:`MultiresEncodeAuxiliaryData` operators for writing user-prescribed +BC profiles into the ``f_1`` buffer during initialization. +""" + +import inspect +from typing import Any, Callable + +import warp as wp + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb import DefaultConfig, ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.macroscopic import SecondMoment as MomentumFlux +from xlb.operator.macroscopic import Macroscopic +from xlb.operator.equilibrium import QuadraticEquilibrium + + +class HelperFunctionsBC(object): + """Collection of Warp/Neon ``@wp.func`` helpers for boundary conditions. + + Parameters + ---------- + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + Must be ``WARP`` or ``NEON`` (JAX not supported). + distance_decoder_function : callable, optional + Function to decode wall-distance data for interpolated BCs. + """ + + def __init__(self, velocity_set=None, precision_policy=None, compute_backend=None, distance_decoder_function=None): + if compute_backend == ComputeBackend.JAX: + raise ValueError("This helper class contains helper functions only for the WARP implementation of some BCs not JAX!") + + # Set the default values from the global config + self.velocity_set = velocity_set or DefaultConfig.velocity_set + self.precision_policy = precision_policy or DefaultConfig.default_precision_policy + self.compute_backend = compute_backend or DefaultConfig.default_backend + self.distance_decoder_function = distance_decoder_function + + # Set the compute and Store dtypes + compute_dtype = self.precision_policy.compute_precision.wp_dtype + store_dtype = self.precision_policy.store_precision.wp_dtype + + # Set local constants + _d = self.velocity_set.d + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + _w = self.velocity_set.w + _c = self.velocity_set.c + _c_float = self.velocity_set.c_float + _qi = self.velocity_set.qi + _u_vec = wp.vec(_d, dtype=compute_dtype) + _f_vec = wp.vec(_q, dtype=compute_dtype) + _missing_mask_vec = wp.vec(_q, dtype=wp.uint8) # TODO fix vec bool + + # Define the operator needed for computing equilibrium + equilibrium = QuadraticEquilibrium(velocity_set, precision_policy, compute_backend) + + # Define the operator needed for computing macroscopic variables + macroscopic = Macroscopic(velocity_set, precision_policy, compute_backend) + + # Define the operator needed for computing the momentum flux + momentum_flux = MomentumFlux(velocity_set, precision_policy, compute_backend) + + @wp.func + def get_bc_thread_data( + f_pre: wp.array4d(dtype=Any), + f_post: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + index: wp.vec3i, + ): + # Get the boundary id and missing mask + _f_pre = _f_vec() + _f_post = _f_vec() + _boundary_id = bc_mask[0, index[0], index[1], index[2]] + _missing_mask = _missing_mask_vec() + for l in range(_q): + # q-sized vector of populations + _f_pre[l] = compute_dtype(f_pre[l, index[0], index[1], index[2]]) + _f_post[l] = compute_dtype(f_post[l, index[0], index[1], index[2]]) + + # TODO fix vec bool + if missing_mask[l, index[0], index[1], index[2]]: + _missing_mask[l] = wp.uint8(1) + else: + _missing_mask[l] = wp.uint8(0) + return _f_pre, _f_post, _boundary_id, _missing_mask + + @wp.func + def neon_get_bc_thread_data( + f_pre_pn: Any, + f_post_pn: Any, + bc_mask_pn: Any, + missing_mask_pn: Any, + index: Any, + ): + # Get the boundary id and missing mask + _f_pre = _f_vec() + _f_post = _f_vec() + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + _missing_mask = _missing_mask_vec() + for l in range(_q): + # q-sized vector of populations + _f_pre[l] = compute_dtype(wp.neon_read(f_pre_pn, index, l)) + _f_post[l] = compute_dtype(wp.neon_read(f_post_pn, index, l)) + _missing_mask[l] = wp.neon_read(missing_mask_pn, index, l) + + return _f_pre, _f_post, _boundary_id, _missing_mask + + @wp.func + def get_bc_fsum( + fpop: Any, + _missing_mask: Any, + ): + fsum_known = compute_dtype(0.0) + fsum_middle = compute_dtype(0.0) + for l in range(_q): + if _missing_mask[_opp_indices[l]] == wp.uint8(1): + fsum_known += compute_dtype(2.0) * fpop[l] + elif _missing_mask[l] != wp.uint8(1): + fsum_middle += fpop[l] + return fsum_known + fsum_middle + + @wp.func + def get_normal_vectors( + _missing_mask: Any, + ): + if wp.static(_d == 3): + for l in range(_q): + if _missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1: + return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l]) + else: + for l in range(_q): + if _missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1: + return -_u_vec(_c_float[0, l], _c_float[1, l]) + + @wp.func + def bounceback_nonequilibrium( + fpop: Any, + feq: Any, + _missing_mask: Any, + ): + for l in range(_q): + if _missing_mask[l] == wp.uint8(1): + fpop[l] = fpop[_opp_indices[l]] + feq[l] - feq[_opp_indices[l]] + return fpop + + @wp.func + def regularize_fpop( + fpop: Any, + feq: Any, + ): + """ + Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop. + """ + # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq} + f_neq = fpop - feq + PiNeq = momentum_flux.warp_functional(f_neq) + + # Compute double dot product Qi:Pi1 (where Pi1 = PiNeq) + nt = _d * (_d + 1) // 2 + for l in range(_q): + QiPi1 = compute_dtype(0.0) + for t in range(nt): + QiPi1 += _qi[l, t] * PiNeq[t] + + # assign all populations based on eq 45 of Latt et al (2008) + # fneq ~ f^1 + fpop1 = compute_dtype(4.5) * _w[l] * QiPi1 + fpop[l] = feq[l] + fpop1 + return fpop + + @wp.func + def grads_approximate_fpop( + _missing_mask: Any, + rho: Any, + u: Any, + f_post: Any, + ): + # Purpose: Using Grad's approximation to represent fpop based on macroscopic inputs used for outflow [1] and + # Dirichlet BCs [2] + # [1] S. Chikatax`marla, S. Ansumali, and I. Karlin, "Grad's approximation for missing data in lattice Boltzmann + # simulations", Europhys. Lett. 74, 215 (2006). + # [2] Dorschner, B., Chikatamarla, S. S., BΓΆsch, F., & Karlin, I. V. (2015). Grad's approximation for moving and + # stationary walls in entropic lattice Boltzmann simulations. Journal of Computational Physics, 295, 340-354. + + # Note: See also self.regularize_fpop function which is somewhat similar. + + # Compute pressure tensor Pi using all f_post-streaming values + Pi = momentum_flux.warp_functional(f_post) + + # Compute double dot product Qi:Pi1 (where Pi1 = PiNeq) + nt = _d * (_d + 1) // 2 + for l in range(_q): + if _missing_mask[l] == wp.uint8(1): + # compute dot product of qi and Pi + QiPi = compute_dtype(0.0) + for t in range(nt): + if t == 0 or t == 3 or t == 5: + QiPi += _qi[l, t] * (Pi[t] - rho / compute_dtype(3.0)) + else: + QiPi += _qi[l, t] * Pi[t] + + # Compute c.u + cu = compute_dtype(0.0) + for d in range(_d): + if _c[d, l] == 1: + cu += u[d] + elif _c[d, l] == -1: + cu -= u[d] + cu *= compute_dtype(3.0) + + # change f_post using the Grad's approximation + f_post[l] = rho * _w[l] * (compute_dtype(1.0) + cu) + _w[l] * compute_dtype(4.5) * QiPi + + return f_post + + @wp.func + def moving_wall_fpop_correction( + u_wall: Any, + lattice_direction: Any, + ): + # Add forcing term necessary to account for the local density changes caused by the mass displacement + # as the object moves with velocity u_wall. + # [1] L.-S. Luo, Unified theory of lattice Boltzmann models for nonideal gases, Phys. Rev. Lett. 81 (1998) 1618-1621. + # [2] L.-S. Luo, Theory of the lattice Boltzmann method: Lattice Boltzmann models for nonideal gases, Phys. Rev. E 62 (2000) 4982-4996. + # + # Note: this function must be called within a for-loop over all lattice directions and the populations to be modified must + # be only those in the missing direction (the check for missing direction must be outside of this function). + cu = compute_dtype(0.0) + l = lattice_direction + for d in range(_d): + if _c[d, l] == 1: + cu += u_wall[d] + elif _c[d, l] == -1: + cu -= u_wall[d] + cu *= compute_dtype(6.0) * _w[l] + return cu + + @wp.func + def interpolated_bounceback( + index: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + u_wall: Any, + needs_moving_wall_treatment: bool, + needs_mesh_distance: bool, + ): + # A local single-node version of the interpolated bounce-back boundary condition due to Bouzidi for a lattice + # Boltzmann method simulation. + # Ref: + # [1] Yu, D., Mei, R., Shyy, W., 2003. A unified boundary treatment in lattice boltzmann method, + # in: 41st aerospace sciences meeting and exhibit, p. 953. + + one = compute_dtype(1.0) + for l in range(_q): + # If the mask is missing then take the opposite index + if _missing_mask[l] == wp.uint8(1): + # The normalized distance to the mesh or "weights" have been stored in known directions of f_1 + if needs_mesh_distance: + # use weights associated with curved boundaries that are properly stored in f_1. + weight = compute_dtype(self.distance_decoder_function(f_1, index, l)) + + # Use differentiable interpolated BB to find f_missing: + f_post[l] = ((one - weight) * f_post[_opp_indices[l]] + weight * (f_pre[l] + f_pre[_opp_indices[l]])) / (one + weight) + else: + # Use regular halfway bounceback + f_post[l] = f_pre[_opp_indices[l]] + + if _missing_mask[_opp_indices[l]] == wp.uint8(1): + # These are cases where the boundary is sandwiched between 2 solid cells and so both opposite directions are missing. + f_post[l] = f_pre[_opp_indices[l]] + + # Add contribution due to moving_wall to f_missing as is usual in regular Bouzidi BC + if needs_moving_wall_treatment: + f_post[l] += moving_wall_fpop_correction(u_wall, l) + return f_post + + @wp.func + def interpolated_nonequilibrium_bounceback( + index: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + u_wall: Any, + needs_moving_wall_treatment: bool, + needs_mesh_distance: bool, + ): + # Compute density, velocity using all f_post-collision values + rho, u = macroscopic.warp_functional(f_pre) + feq = equilibrium.warp_functional(rho, u) + + # Compute equilibrium distribution at the wall + if needs_moving_wall_treatment: + feq_wall = equilibrium.warp_functional(rho, u_wall) + else: + feq_wall = _f_vec() + + # Apply method in Tao et al (2018) [1] to find missing populations at the boundary + one = compute_dtype(1.0) + for l in range(_q): + # If the mask is missing then take the opposite index + if _missing_mask[l] == wp.uint8(1): + # The normalized distance to the mesh or "weights" have been stored in known directions of f_1 + if needs_mesh_distance: + # use weights associated with curved boundaries that are properly stored in f_1. + weight = compute_dtype(self.distance_decoder_function(f_1, index, l)) + else: + weight = compute_dtype(0.5) + + # Use non-equilibrium bounceback to find f_missing: + fneq = f_pre[_opp_indices[l]] - feq[_opp_indices[l]] + + # Compute equilibrium distribution at the wall + # Same quadratic equilibrium but accounting for zero velocity (no-slip) + if not needs_moving_wall_treatment: + feq_wall[l] = _w[l] * rho + + # Assemble wall population for doing interpolation at the boundary + f_wall = feq_wall[l] + fneq + f_post[l] = (f_wall + weight * f_pre[l]) / (one + weight) + + return f_post + + @wp.func + def neon_index_to_warp(neon_field_hdl: Any, index: Any): + # Unpack the global index in Neon at the finest level and convert it to a warp vector + cIdx = wp.neon_global_idx(neon_field_hdl, index) + gx = wp.neon_get_x(cIdx) + gy = wp.neon_get_y(cIdx) + gz = wp.neon_get_z(cIdx) + + # XLB is flattening the z dimension in 3D, while neon uses the y dimension + if _d == 2: + gy, gz = gz, gy + + # Get warp indices + index_wp = wp.vec3i(gx, gy, gz) + return index_wp + + self.get_bc_thread_data = get_bc_thread_data + self.get_bc_fsum = get_bc_fsum + self.get_normal_vectors = get_normal_vectors + self.bounceback_nonequilibrium = bounceback_nonequilibrium + self.regularize_fpop = regularize_fpop + self.grads_approximate_fpop = grads_approximate_fpop + self.moving_wall_fpop_correction = moving_wall_fpop_correction + self.interpolated_bounceback = interpolated_bounceback + self.interpolated_nonequilibrium_bounceback = interpolated_nonequilibrium_bounceback + self.neon_get_bc_thread_data = neon_get_bc_thread_data + self.neon_index_to_warp = neon_index_to_warp + + +class EncodeAuxiliaryData(Operator): + """ + Operator for encoding boundary auxiliary data during initialization. + """ + + def __init__( + self, + boundary_id: int, + num_of_aux_data: int, + user_defined_functional: Callable, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + self.user_defined_functional = user_defined_functional + self.boundary_id = wp.uint8(boundary_id) + self.num_of_aux_data = num_of_aux_data + + super().__init__(velocity_set, precision_policy, compute_backend) + + # Inspect the signature of the user-defined functional. + # We assume the profile function takes only the index as input and is hence time-independent. + sig = inspect.signature(user_defined_functional) + assert self.compute_backend != ComputeBackend.JAX, "Encoding/decoding of auxiliary data are not required for boundary conditions in JAX" + assert len(sig.parameters) == 1, f"User-defined functional must take exactly one argument (the index), it received {len(sig.parameters)}." + + # Define a HelperFunctionsBC instance + self.bc_helper = HelperFunctionsBC( + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + + # TODO: Somehow raise an error if the number of prescribed values does not match the number of missing directions + + def _construct_warp(self): + """ + Constructs the warp kernel for the auxiliary data recovery. + """ + # Find velocity index for (0, 0, 0) + lattice_central_index = self.velocity_set.center_index + _opp_indices = self.velocity_set.opp_indices + _id = self.boundary_id + _num_of_aux_data = self.num_of_aux_data + _aux_vec = wp.vec(_num_of_aux_data, dtype=self.compute_dtype) + + @wp.func + def encoder_functional( + index: Any, + _missing_mask: Any, + field_storage: Any, + prescribed_values: Any, + ): + if len(prescribed_values) != _num_of_aux_data: + wp.printf("Error: User-defined profile must return a vector of size %d\n", _num_of_aux_data) + return + + # Write the result for all q directions, but only store up to _num_of_aux_data + counter = wp.int32(0) + for l in range(self.velocity_set.q): + # Only store up to _num_of_aux_data + if counter == _num_of_aux_data: + return + + if l == lattice_central_index: + # The first BC auxiliary data is stored in the zero'th index of f_1 associated with its center. + self.write_field(field_storage, index, l, self.store_dtype(prescribed_values[l])) + counter += 1 + elif _missing_mask[l] == wp.uint8(1): + # The other remaining BC auxiliary data are stored in missing directions of f_1. + self.write_field(field_storage, index, _opp_indices[l], self.store_dtype(prescribed_values[l])) + counter += 1 + + @wp.func + def decoder_functional( + field_storage: Any, + index: Any, + _missing_mask: Any, + ): + """ + Decode the encoded values needed for the boundary condition treatment from the center location in field_storage. + """ + + # Define a vector to hold prescribed_values + prescribed_values = _aux_vec() + + # Read all q directions, but only retrieve up to _num_of_aux_data + counter = wp.int32(0) + for l in range(self.velocity_set.q): + # Only retrieve up to _num_of_aux_data + if counter == _num_of_aux_data: + return prescribed_values + + if l == lattice_central_index: + # The first BC auxiliary data is stored in the zero'th index of f_1 associated with its center. + value = self.read_field(field_storage, index, l) + prescribed_values[counter] = self.compute_dtype(value) + counter += 1 + elif _missing_mask[l] == wp.uint8(1): + # The other remaining BC auxiliary data are stored in missing directions of f_1. + value = self.read_field(field_storage, index, _opp_indices[l]) + prescribed_values[counter] = self.compute_dtype(value) + counter += 1 + + # Construct the warp kernel + @wp.kernel + def kernel( + f_1: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # read tid data + _, _, _boundary_id, _missing_mask = self.bc_helper.get_bc_thread_data(f_1, f_1, bc_mask, missing_mask, index) + + # Apply the functional + # change this to use central location + if _boundary_id == _id: + # prescribed_values is a q-sized vector of type wp.vec + prescribed_values = self.user_defined_functional(index) + + # call the functional + encoder_functional(index, _missing_mask, f_1, prescribed_values) + + functional_dict = {"encoder": encoder_functional, "decoder": decoder_functional} + return functional_dict, kernel + + def _construct_neon(self): + import neon + + """ + Constructs the Neon container for encoding auxiliary data recovery. + """ + # Use the warp functional for the Neon backend + functional_dict, _ = self._construct_warp() + encoder_functional = functional_dict["encoder"] + _id = self.boundary_id + + # Construct the Neon container + @neon.Container.factory(name="EncodingAuxData_" + str(_id)) + def aux_data_init_container( + f_1: Any, + bc_mask: Any, + missing_mask: Any, + ): + def aux_data_init_ll(loader: neon.Loader): + loader.set_grid(f_1.get_grid()) + + f_1_pn = loader.get_write_handle(f_1) + bc_mask_pn = loader.get_read_handle(bc_mask) + missing_mask_pn = loader.get_read_handle(missing_mask) + + @wp.func + def aux_data_init_cl(index: Any): + # read tid data + _, _, _boundary_id, _missing_mask = self.bc_helper.neon_get_bc_thread_data(f_1_pn, f_1_pn, bc_mask_pn, missing_mask_pn, index) + + # Apply the functional + if _boundary_id == _id: + warp_index = self.bc_helper.neon_index_to_warp(f_1_pn, index) + prescribed_values = self.user_defined_functional(warp_index) + + # Call the functional + encoder_functional(index, _missing_mask, f_1_pn, prescribed_values) + + # Declare the kernel in the Neon loader + loader.declare_kernel(aux_data_init_cl) + + return aux_data_init_ll + + return functional_dict, aux_data_init_container + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_1, bc_mask, missing_mask): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_1, bc_mask, missing_mask], + dim=f_1.shape[1:], + ) + return f_1 + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_1, bc_mask, missing_mask): + import neon + + c = self.neon_container(f_1, bc_mask, missing_mask) + c.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + return f_1 + + +class MultiresEncodeAuxiliaryData(EncodeAuxiliaryData): + """ + Operator for encoding boundary auxiliary data during initialization. + """ + + def __init__( + self, + boundary_id: int, + num_of_aux_data: int, + user_defined_functional: Callable, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + super().__init__( + boundary_id=boundary_id, + num_of_aux_data=num_of_aux_data, + user_defined_functional=user_defined_functional, + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + assert self.compute_backend == ComputeBackend.NEON, f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend." + + def _construct_neon(self): + """ + Constructs the Neon container for encoding auxiliary data recovery. + """ + import neon + + # Borrow the functional from the warp implementation + functional_dict, _ = self._construct_warp() + encoder_functional = functional_dict["encoder"] + _id = self.boundary_id + + # Construct the Neon container + @neon.Container.factory(name="MultiresEncodingAuxData_" + str(_id)) + def aux_data_init_container( + f_1: Any, + bc_mask: Any, + missing_mask: Any, + level: Any, + ): + def aux_data_init_ll(loader: neon.Loader): + loader.set_mres_grid(f_1.get_grid(), level) + + f_1_pn = loader.get_mres_write_handle(f_1) + bc_mask_pn = loader.get_mres_read_handle(bc_mask) + missing_mask_pn = loader.get_mres_read_handle(missing_mask) + + @wp.func + def aux_data_init_cl(index: Any): + # read tid data + _, _, _boundary_id, _missing_mask = self.bc_helper.neon_get_bc_thread_data(f_1_pn, f_1_pn, bc_mask_pn, missing_mask_pn, index) + + # Apply the functional + if _boundary_id == _id: + # IMPORTANT: XLB assumes the user_defined_functional in multi-res + # simulations uses finest-level indices, enabling BCs that span + # multiple levels. + warp_index = self.bc_helper.neon_index_to_warp(f_1_pn, index) + prescribed_values = self.user_defined_functional(warp_index) + + # Call the functional + encoder_functional(index, _missing_mask, f_1_pn, prescribed_values) + + # Declare the kernel in the Neon loader + loader.declare_kernel(aux_data_init_cl) + + return aux_data_init_ll + + return functional_dict, aux_data_init_container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_1, bc_mask, missing_mask, stream): + import neon + + grid = bc_mask.get_grid() + for level in range(grid.num_levels): + c = self.neon_container(f_1, bc_mask, missing_mask, level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return f_1 diff --git a/xlb/operator/boundary_masker/__init__.py b/xlb/operator/boundary_masker/__init__.py new file mode 100644 index 00000000..5a1ceb75 --- /dev/null +++ b/xlb/operator/boundary_masker/__init__.py @@ -0,0 +1,12 @@ +from xlb.operator.boundary_masker.helper_functions_masker import HelperFunctionsMasker +from xlb.operator.boundary_masker.indices_boundary_masker import IndicesBoundaryMasker +from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker +from xlb.operator.boundary_masker.aabb import MeshMaskerAABB +from xlb.operator.boundary_masker.ray import MeshMaskerRay +from xlb.operator.boundary_masker.winding import MeshMaskerWinding +from xlb.operator.boundary_masker.aabb_close import MeshMaskerAABBClose +from xlb.operator.boundary_masker.mesh_voxelization_method import MeshVoxelizationMethod +from xlb.operator.boundary_masker.multires_aabb import MultiresMeshMaskerAABB +from xlb.operator.boundary_masker.multires_aabb_close import MultiresMeshMaskerAABBClose +from xlb.operator.boundary_masker.multires_indices_boundary_masker import MultiresIndicesBoundaryMasker +from xlb.operator.boundary_masker.multires_ray import MultiresMeshMaskerRay diff --git a/xlb/operator/boundary_masker/aabb.py b/xlb/operator/boundary_masker/aabb.py new file mode 100644 index 00000000..9e69e9c4 --- /dev/null +++ b/xlb/operator/boundary_masker/aabb.py @@ -0,0 +1,200 @@ +""" +AABB mesh-based boundary masker. + +Voxelizes an STL mesh using ``warp.mesh_query_aabb`` for approximate +one-voxel-thick surface detection around the geometry. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker +from xlb.operator.operator import Operator +from xlb.cell_type import BC_SOLID + + +class MeshMaskerAABB(MeshBoundaryMasker): + """ + Operator for creating boundary missing_mask from mesh using Axis-Aligned Bounding Box (AABB) voxelization. + + This implementation uses warp.mesh_query_aabb for efficient mesh-voxel intersection testing, + providing approximate 1-voxel thick surface detection around the mesh geometry. + Suitable for scenarios where fast, approximate boundary detection is sufficient. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + + def _construct_warp(self): + # Make constants for warp + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + + # Set local constants + lattice_central_index = self.velocity_set.center_index + + @wp.func + def functional( + index: Any, + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + ): + # position of the point + cell_center_pos = self.helper_masker.index_to_position(bc_mask, index) + HALF_VOXEL = wp.vec3(0.5, 0.5, 0.5) + + if self.read_field(bc_mask, index, 0) == wp.uint8(BC_SOLID) or self.mesh_voxel_intersect( + mesh_id=mesh_id, low=cell_center_pos - HALF_VOXEL + ): + # Make solid voxel + self.write_field(bc_mask, index, 0, wp.uint8(BC_SOLID)) + else: + # Find the boundary voxels and their missing directions + for direction_idx in range(_q): + if direction_idx == lattice_central_index: + # Skip the central index as it is not relevant for boundary masking + continue + + # Get the lattice direction vector + direction_vec = wp.vec3f(wp.float32(_c[0, direction_idx]), wp.float32(_c[1, direction_idx]), wp.float32(_c[2, direction_idx])) + + # Check to see if this neighbor is solid + if self.mesh_voxel_intersect(mesh_id=mesh_id, low=cell_center_pos + direction_vec - HALF_VOXEL): + # We know we have a solid neighbor + # Set the boundary id and missing_mask + self.write_field(bc_mask, index, 0, wp.uint8(id_number)) + self.write_field(missing_mask, index, _opp_indices[direction_idx], wp.uint8(True)) + + # If we don't need the mesh distance, we can return early + if not needs_mesh_distance: + continue + + # Find the fractional distance to the mesh in each direction + # We increase max_length to find intersections in neighboring cells + max_length = wp.length(direction_vec) + query = wp.mesh_query_ray(mesh_id, cell_center_pos, direction_vec / max_length, 1.5 * max_length) + if query.result: + # get position of the mesh triangle that intersects with the ray + pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) + # We reduce the distance to give some wall thickness + dist = wp.length(pos_mesh - cell_center_pos) - 0.5 * max_length + weight = dist / max_length + self.write_field(distances, index, direction_idx, self.store_dtype(weight)) + else: + # Expected an intersection in this direction but none was found. + # Assume the solid extends one lattice unit beyond the BC voxel leading to a distance fraction of 1. + self.write_field(distances, index, direction_idx, self.store_dtype(1.0)) + + @wp.kernel + def kernel( + mesh_id: wp.uint64, + id_number: wp.int32, + distances: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + needs_mesh_distance: bool, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # apply the functional + functional( + index, + mesh_id, + id_number, + distances, + bc_mask, + missing_mask, + needs_mesh_distance, + ) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + return self.warp_implementation_base( + bc, + distances, + bc_mask, + missing_mask, + ) + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MeshMaskerAABB") + def container( + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + ): + def aabb_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + missing_mask_pn = loader.get_write_handle(missing_mask) + distances_pn = loader.get_write_handle(distances) + + @wp.func + def aabb_kernel(index: Any): + # apply the functional + functional( + index, + mesh_id, + id_number, + distances_pn, + bc_mask_pn, + missing_mask_pn, + needs_mesh_distance, + ) + + loader.declare_kernel(aabb_kernel) + + return aabb_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + import neon + + # Prepare inputs + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + # Launch the appropriate neon container + c = self.neon_container(mesh_id, bc_id, distances, bc_mask, missing_mask, wp.static(bc.needs_mesh_distance)) + c.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/aabb_close.py b/xlb/operator/boundary_masker/aabb_close.py new file mode 100644 index 00000000..21ee31f9 --- /dev/null +++ b/xlb/operator/boundary_masker/aabb_close.py @@ -0,0 +1,365 @@ +""" +AABB-Close boundary masker with morphological close operation. + +Identifies solid voxels via axis-aligned bounding-box (AABB) intersection, +then applies a morphological *close* (dilate followed by erode) to fill +thin gaps and small cavities in the mesh surface. The resulting solid +mask is used to determine boundary voxels and their missing population +directions. + +Supports both Warp (single-resolution) and Neon (multi-resolution) +backends. +""" + +import numpy as np +import warp as wp +import jax +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker +from xlb.cell_type import BC_SOLID + + +class MeshMaskerAABBClose(MeshBoundaryMasker): + """Boundary masker using AABB voxelization with morphological close. + + The *close* operation (dilate then erode) thickens the raw solid mask + by ``close_voxels`` layers before shrinking it back, sealing small + holes and thin slits in the mesh surface. + + Parameters + ---------- + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + close_voxels : int + Half-width of the morphological structuring element. Must be + provided explicitly. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + close_voxels: int = None, + ): + assert close_voxels is not None, ( + "Please provide the number of close voxels using the 'close_voxels' argument! e.g., MeshVoxelizationMethod('AABB_CLOSE', close_voxels=3)" + ) + self.close_voxels = close_voxels + # Call super + self.tile_half = close_voxels + self.tile_size = self.tile_half * 2 + 1 + super().__init__(velocity_set, precision_policy, compute_backend) + + def _construct_warp(self): + # Make constants for warp + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + TILE_SIZE = wp.constant(self.tile_size) + TILE_HALF = wp.constant(self.tile_half) + lattice_central_index = self.velocity_set.center_index + + # Erode the solid mask in mask_field, removing a layer of outer solid voxels, storing output in mask_field_out + @wp.kernel + def erode_tile(mask_field: wp.array3d(dtype=Any), mask_field_out: wp.array3d(dtype=Any)): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + if not self.helper_masker.is_in_bounds(index, wp.vec3i(mask_field.shape[0], mask_field.shape[1], mask_field.shape[2]), TILE_HALF): + mask_field_out[i, j, k] = mask_field[i, j, k] + return + t = wp.tile_load(mask_field, shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE), offset=(i - TILE_HALF, j - TILE_HALF, k - TILE_HALF)) + min_val = wp.tile_min(t) + mask_field_out[i, j, k] = min_val[0] + + # Dilate the solid mask in mask_field, adding a layer of outer solid voxels, storing output in mask_field_out + @wp.kernel + def dilate_tile(mask_field: wp.array3d(dtype=Any), mask_field_out: wp.array3d(dtype=Any)): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + if not self.helper_masker.is_in_bounds(index, wp.vec3i(mask_field.shape[0], mask_field.shape[1], mask_field.shape[2]), TILE_HALF): + mask_field_out[i, j, k] = mask_field[i, j, k] + return + t = wp.tile_load(mask_field, shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE), offset=(i - TILE_HALF, j - TILE_HALF, k - TILE_HALF)) + max_val = wp.tile_max(t) + mask_field_out[i, j, k] = max_val[0] + + # Erode the solid mask in mask_field, removing a layer of outer solid voxels, storing output in mask_field_out + @wp.func + def functional_erode(index: Any, mask_field: Any, mask_field_out: Any): + min_val = wp.uint8(BC_SOLID) + for l in range(_q): + if l == lattice_central_index: + continue + is_valid = wp.bool(False) + ngh = wp.neon_ngh_idx(wp.int8(_c[0, l]), wp.int8(_c[1, l]), wp.int8(_c[2, l])) + ngh_val = wp.neon_read_ngh(mask_field, index, ngh, 0, wp.uint8(0), is_valid) + if is_valid: + # Take the min value of all neighbors in bounds + min_val = wp.min(min_val, ngh_val) + self.write_field(mask_field_out, index, 0, min_val) + + # Dilate the solid mask in mask_field, adding a layer of outer solid voxels, storing output in mask_field_out + @wp.func + def functional_dilate(index: Any, mask_field: Any, mask_field_out: Any): + max_val = wp.uint8(0) + for l in range(_q): + if l == lattice_central_index: + continue + is_valid = wp.bool(False) + ngh = wp.neon_ngh_idx(wp.int8(_c[0, l]), wp.int8(_c[1, l]), wp.int8(_c[2, l])) + ngh_val = wp.neon_read_ngh(mask_field, index, ngh, 0, wp.uint8(0), is_valid) + if is_valid: + max_val = wp.max(max_val, ngh_val) + self.write_field(mask_field_out, index, 0, max_val) + + # Construct the warp kernel + # Find solid voxels that intersect the mesh + @wp.func + def functional_solid(index: Any, mesh_id: Any, solid_mask: Any, offset: Any): + # position of the point + cell_center_pos = self.helper_masker.index_to_position(solid_mask, index) + offset + half = wp.vec3(0.5, 0.5, 0.5) + + if self.mesh_voxel_intersect(mesh_id=mesh_id, low=cell_center_pos - half): + # Make solid voxel + self.write_field(solid_mask, index, 0, wp.uint8(BC_SOLID)) + + @wp.kernel + def kernel_solid( + mesh_id: wp.uint64, + solid_mask: wp.array3d(dtype=wp.int32), + offset: wp.vec3f, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + functional_solid(index, mesh_id, solid_mask, offset) + + return + + @wp.func + def functional_aabb( + index: Any, + mesh_id: wp.uint64, + id_number: wp.int32, + distances: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + solid_mask: wp.array3d(dtype=wp.uint8), + needs_mesh_distance: bool, + ): + # position of the point + cell_center_pos = self.helper_masker.index_to_position(bc_mask, index) + HALF_VOXEL = wp.vec3(0.5, 0.5, 0.5) + + if self.read_field(solid_mask, index, 0) == wp.uint8(BC_SOLID) or self.read_field(bc_mask, index, 0) == wp.uint8(BC_SOLID): + # Make solid voxel + self.write_field(bc_mask, index, 0, wp.uint8(BC_SOLID)) + else: + # Find the boundary voxels and their missing directions + for direction_idx in range(_q): + if direction_idx == lattice_central_index: + # Skip the central index as it is not relevant for boundary masking + continue + + # Get the lattice direction vector + direction_vec = wp.vec3f(wp.float32(_c[0, direction_idx]), wp.float32(_c[1, direction_idx]), wp.float32(_c[2, direction_idx])) + + # Check to see if this neighbor is solid + if self.helper_masker.is_in_bounds(index, wp.vec3i(solid_mask.shape[0], solid_mask.shape[1], solid_mask.shape[2]), 1): + if self.read_field(solid_mask, index + direction_idx, 0) == wp.uint8(BC_SOLID): + # We know we have a solid neighbor + # Set the boundary id and missing_mask + self.write_field(bc_mask, index, 0, wp.uint8(id_number)) + self.write_field(missing_mask, index, _opp_indices[direction_idx], wp.uint8(True)) + + # If we don't need the mesh distance, we can return early + if not needs_mesh_distance: + continue + + # Find the fractional distance to the mesh in each direction + # We increase max_length to find intersections in neighboring cells + max_length = wp.length(direction_vec) + query = wp.mesh_query_ray(mesh_id, cell_center_pos, direction_vec / max_length, 1.5 * max_length) + if query.result: + # get position of the mesh triangle that intersects with the ray + pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) + # We reduce the distance to give some wall thickness + dist = wp.length(pos_mesh - cell_center_pos) - 0.5 * max_length + weight = dist / max_length + self.write_field(distances, index, direction_idx, self.store_dtype(weight)) + else: + # Expected an intersection in this direction but none was found. + # Assume the solid extends one lattice unit beyond the BC voxel leading to a distance fraction of 1. + self.write_field(distances, index, direction_idx, self.store_dtype(1.0)) + + # Assign the bc_mask and distances based on the solid_mask we already computed + @wp.kernel + def kernel( + mesh_id: wp.uint64, + id_number: wp.int32, + distances: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + solid_mask: wp.array3d(dtype=wp.uint8), + needs_mesh_distance: bool, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # position of the point + cell_center_pos = self.helper_masker.index_to_position(bc_mask, index) + + if solid_mask[i, j, k] == wp.uint8(BC_SOLID) or bc_mask[0, index[0], index[1], index[2]] == wp.uint8(BC_SOLID): + # Make solid voxel + bc_mask[0, index[0], index[1], index[2]] = wp.uint8(BC_SOLID) + else: + # Find the boundary voxels and their missing directions + for direction_idx in range(_q): + if direction_idx == lattice_central_index: + # Skip the central index as it is not relevant for boundary masking + continue + direction_vec = wp.vec3f(wp.float32(_c[0, direction_idx]), wp.float32(_c[1, direction_idx]), wp.float32(_c[2, direction_idx])) + + # Check to see if this neighbor is solid - this is super inefficient TODO: make it way better + # if solid_mask[i,j,k] == wp.uint8(BC_SOLID): + if solid_mask[i + _c[0, direction_idx], j + _c[1, direction_idx], k + _c[2, direction_idx]] == wp.uint8(BC_SOLID): + # We know we have a solid neighbor + # Set the boundary id and missing_mask + bc_mask[0, index[0], index[1], index[2]] = wp.uint8(id_number) + missing_mask[_opp_indices[direction_idx], index[0], index[1], index[2]] = wp.uint8(True) + + # If we don't need the mesh distance, we can return early + if not needs_mesh_distance: + continue + + # Find the fractional distance to the mesh in each direction + # We increase max_length to find intersections in neighboring cells + max_length = wp.length(direction_vec) + query = wp.mesh_query_ray(mesh_id, cell_center_pos, direction_vec / max_length, 1.5 * max_length) + if query.result: + # get position of the mesh triangle that intersects with the ray + pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) + # We reduce the distance to give some wall thickness + dist = wp.length(pos_mesh - cell_center_pos) - 0.5 * max_length + weight = self.store_dtype(dist / max_length) + distances[direction_idx, index[0], index[1], index[2]] = weight + else: + # We didn't have an intersection in the given direction but we know we should so we assume the solid is slightly thicker + # and one lattice direction away from the BC voxel + distances[direction_idx, index[0], index[1], index[2]] = self.store_dtype(1.0) + + functional_dict = { + "functional_erode": functional_erode, + "functional_dilate": functional_dilate, + "functional_solid": functional_solid, + "functional_aabb": functional_aabb, + } + kernel_dict = { + "kernel": kernel, + "kernel_solid": kernel_solid, + "erode_tile": erode_tile, + "dilate_tile": dilate_tile, + } + return functional_dict, kernel_dict + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + assert bc.mesh_vertices is not None, f'Please provide the mesh vertices for {bc.__class__.__name__} BC using keyword "mesh_vertices"!' + assert bc.indices is None, f"Please use IndicesBoundaryMasker operator if {bc.__class__.__name__} is imposed on known indices of the grid!" + assert bc.mesh_vertices.shape[1] == self.velocity_set.d, ( + "Mesh points must be reshaped into an array (N, 3) where N indicates number of points!" + ) + + domain_shape = bc_mask.shape[1:] # (nx, ny, nz) + mesh_vertices = bc.mesh_vertices + mesh_min = np.min(mesh_vertices, axis=0) + mesh_max = np.max(mesh_vertices, axis=0) + + if any(mesh_min < 0) or any(mesh_max >= domain_shape): + raise ValueError( + f"Mesh extents ({mesh_min}, {mesh_max}) exceed domain dimensions {domain_shape}. The mesh must be fully contained within the domain." + ) + + # We are done with bc.mesh_vertices. Remove them from BC objects + bc.__dict__.pop("mesh_vertices", None) + + mesh_indices = np.arange(mesh_vertices.shape[0]) + mesh = wp.Mesh( + points=wp.array(mesh_vertices, dtype=wp.vec3), + indices=wp.array(mesh_indices, dtype=wp.int32), + ) + mesh_id = wp.uint64(mesh.id) + bc_id = bc.id + + # Create a padded mask for the solid voxels to account for the tile size + # It needs to be padded by twice the tile size on each side since we run two tile operations + tile_length = 2 * self.tile_half + offset = wp.vec3f(-tile_length, -tile_length, -tile_length) + pad = 2 * tile_length + nx, ny, nz = domain_shape + solid_mask = wp.zeros((nx + pad, ny + pad, nz + pad), dtype=wp.int32) + solid_mask_out = wp.zeros((nx + pad, ny + pad, nz + pad), dtype=wp.int32) + + # Prepare the warp kernel dictionary + kernel_dict = self.warp_kernel + + # Launch all required kernels for creating the solid mask + wp.launch( + kernel=kernel_dict["kernel_solid"], + inputs=[ + mesh_id, + solid_mask, + offset, + ], + dim=solid_mask.shape, + ) + wp.launch_tiled( + kernel=kernel_dict["dilate_tile"], + dim=solid_mask.shape, + block_dim=32, + inputs=[solid_mask, solid_mask_out], + ) + wp.launch_tiled( + kernel=kernel_dict["erode_tile"], + dim=solid_mask.shape, + block_dim=32, + inputs=[solid_mask_out, solid_mask], + ) + solid_mask_cropped = wp.array( + solid_mask[tile_length:-tile_length, tile_length:-tile_length, tile_length:-tile_length], + dtype=wp.uint8, + ) + + # Launch the main kernel for boundary masker + wp.launch( + kernel_dict["kernel"], + inputs=[mesh_id, bc_id, distances, bc_mask, missing_mask, solid_mask_cropped, wp.static(bc.needs_mesh_distance)], + dim=bc_mask.shape[1:], + ) + + # Resolve out of bound indices + wp.launch( + self.resolve_out_of_bound_kernel, + inputs=[bc_id, bc_mask, missing_mask], + dim=bc_mask.shape[1:], + ) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/helper_functions_masker.py b/xlb/operator/boundary_masker/helper_functions_masker.py new file mode 100644 index 00000000..565455fe --- /dev/null +++ b/xlb/operator/boundary_masker/helper_functions_masker.py @@ -0,0 +1,135 @@ +""" +Warp/Neon helper functions shared by boundary masker operators. +""" + +import warp as wp +from typing import Any +from xlb import DefaultConfig, ComputeBackend + + +class HelperFunctionsMasker(object): + """Warp ``@wp.func`` helpers for boundary masker operators. + + Provides coordinate-conversion, bounds-checking, pull-index + computation, and BC-index membership tests used by the mesh and + indices boundary maskers on both Warp and Neon backends. + """ + + def __init__(self, velocity_set=None, precision_policy=None, compute_backend=None): + if compute_backend == ComputeBackend.JAX: + raise ValueError("This helper class contains helper functions only for the WARP implementation of some BCs not JAX!") + + # Set the default values from the global config + self.velocity_set = velocity_set or DefaultConfig.velocity_set + self.precision_policy = precision_policy or DefaultConfig.default_precision_policy + self.compute_backend = compute_backend or DefaultConfig.default_backend + + # Set local constants + _d = self.velocity_set.d + _c = self.velocity_set.c + + @wp.func + def neon_index_to_warp(neon_field_hdl: Any, index: Any): + # Unpack the global index in Neon at the finest level and convert it to a warp vector + cIdx = wp.neon_global_idx(neon_field_hdl, index) + gx = wp.neon_get_x(cIdx) + gy = wp.neon_get_y(cIdx) + gz = wp.neon_get_z(cIdx) + + # XLB is flattening the z dimension in 3D, while neon uses the y dimension + if _d == 2: + gy, gz = gz, gy + + # Get warp indices + index_wp = wp.vec3i(gx, gy, gz) + return index_wp + + @wp.func + def index_to_position_warp(field: Any, index: wp.vec3i): + # position of the point + ijk = wp.vec3(wp.float32(index[0]), wp.float32(index[1]), wp.float32(index[2])) + pos = ijk + wp.vec3(0.5, 0.5, 0.5) # cell center + return pos + + @wp.func + def index_to_position_neon(field: Any, index: Any): + # position of the point + index_wp = neon_index_to_warp(field, index) + return index_to_position_warp(field, index_wp) + + @wp.func + def is_in_bounds(index: wp.vec3i, grid_shape: wp.vec3i, SHIFT: Any = 0): + return ( + index[0] >= SHIFT + and index[0] < grid_shape[0] - SHIFT + and index[1] >= SHIFT + and index[1] < grid_shape[1] - SHIFT + and index[2] >= SHIFT + and index[2] < grid_shape[2] - SHIFT + ) + + @wp.func + def get_pull_index_warp( + field: Any, + lattice_dir: wp.int32, + index: wp.vec3i, + level: Any, + ): + pull_index = wp.vec3i() + offset = wp.vec3i() + for d in range(self.velocity_set.d): + offset[d] = -_c[d, lattice_dir] + for _ in range(level): + offset[d] *= 2 + pull_index[d] = index[d] + offset[d] + + return pull_index, offset + + @wp.func + def get_pull_index_neon( + field: Any, + lattice_dir: wp.int32, + index: Any, + level: Any, + ): + # Convert the index to warp + index_wp = neon_index_to_warp(field, index) + pull_index_wp, _ = get_pull_index_warp(field, lattice_dir, index_wp, level) + offset = wp.neon_ngh_idx(wp.int8(-_c[0, lattice_dir]), wp.int8(-_c[1, lattice_dir]), wp.int8(-_c[2, lattice_dir])) + return pull_index_wp, offset + + @wp.func + def is_in_bc_indices_warp( + field: Any, + index: Any, + bc_indices: wp.array2d(dtype=wp.int32), + ii: wp.int32, + ): + return bc_indices[0, ii] == index[0] and bc_indices[1, ii] == index[1] and bc_indices[2, ii] == index[2] + + @wp.func + def is_in_bc_indices_neon( + field: Any, + index: Any, + bc_indices: wp.array2d(dtype=wp.int32), + ii: wp.int32, + ): + index_wp = neon_index_to_warp(field, index) + return is_in_bc_indices_warp(field, index_wp, bc_indices, ii) + + # Construct some helper warp functions + self.is_in_bounds = is_in_bounds + self.index_to_position = index_to_position_warp if self.compute_backend == ComputeBackend.WARP else index_to_position_neon + self.get_pull_index = get_pull_index_warp if self.compute_backend == ComputeBackend.WARP else get_pull_index_neon + self.is_in_bc_indices = is_in_bc_indices_warp if self.compute_backend == ComputeBackend.WARP else is_in_bc_indices_neon + + def get_grid_shape(self, field): + """ + Get the grid shape from the boundary mask. This is a CPU function that returns the shape of the grid + """ + if self.compute_backend == ComputeBackend.WARP: + return field.shape[1:] + elif self.compute_backend == ComputeBackend.NEON: + return wp.vec3i(field.get_grid().dim.x, field.get_grid().dim.y, field.get_grid().dim.z) + else: + raise ValueError(f"Unsupported compute backend: {self.compute_backend}") diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py new file mode 100644 index 00000000..439fa7d9 --- /dev/null +++ b/xlb/operator/boundary_masker/indices_boundary_masker.py @@ -0,0 +1,572 @@ +""" +Indices-based boundary masker. + +Creates boundary masks from explicit arrays of voxel indices, computing +missing-population masks via pull-index tests for each tagged voxel. +""" + +from typing import Any +import copy + +import jax +import jax.numpy as jnp +import numpy as np +import warp as wp + +from xlb.compute_backend import ComputeBackend +from xlb.grid import grid_factory +from xlb.operator.operator import Operator +from xlb.operator.stream.stream import Stream +from xlb.precision_policy import Precision +from xlb.operator.boundary_masker.helper_functions_masker import HelperFunctionsMasker +from xlb.cell_type import BC_SOLID + + +class IndicesBoundaryMasker(Operator): + """ + Operator for creating a boundary mask + """ + + def __init__( + self, + velocity_set=None, + precision_policy=None, + compute_backend=None, + grid=None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + self.grid = grid + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + # Define masker helper functions + self.helper_masker = HelperFunctionsMasker( + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + else: + # Make stream operator + self.stream = Stream(velocity_set, precision_policy, compute_backend) + + def are_indices_in_interior(self, indices, shape): + """ + Check if each 2D or 3D index is inside the bounds of the domain with the given shape and not + at its boundary. + + :param indices: Array of indices, where each column contains indices for each dimension. + :param shape: Tuple representing the shape of the domain (nx, ny) for 2D or (nx, ny, nz) for 3D. + :return: Array of boolean flags where each flag indicates whether the corresponding index is inside the bounds. + """ + _d = self.velocity_set.d + shape_array = np.array(shape) + return np.all((indices[:_d] > 0) & (indices[:_d] < shape_array[:_d, np.newaxis] - 1), axis=0) + + def _find_bclist_interior(self, bclist, grid_shape): + bc_interior = [] + for bc in bclist: + if any(self.are_indices_in_interior(np.array(bc.indices), grid_shape)): + bc_copy = copy.copy(bc) # shallow copy of the whole object + bc_copy.indices = copy.deepcopy(bc.pad_indices()) # deep copy only the modified part + bc_interior.append(bc_copy) + return bc_interior + + @Operator.register_backend(ComputeBackend.JAX) + # TODO HS: figure out why uncommenting the line below fails unlike other operators! + # @partial(jit, static_argnums=(0)) + def jax_implementation(self, bclist, bc_mask, missing_mask, start_index=None): + # Extend the missing mask by padding to identify out of bound boundaries + # Set padded region to True (i.e. boundary) + dim = missing_mask.ndim - 1 + grid_shape = bc_mask[0].shape + nDevices = jax.device_count() + pad_x, pad_y, pad_z = nDevices, 1, 1 + + # Shift indices due to padding + shift = np.array((pad_x, pad_y) if dim == 2 else (pad_x, pad_y, pad_z))[:, np.newaxis] + if start_index is None: + start_index = (0,) * dim + + # TODO MEHDI: There is sometimes a halting problem here when padding is used in a multi-GPU setting since we're not jitting this function. + # For now, we compute the bc_mask_extended on GPU zero. + if dim == 2: + bc_mask_extended = jnp.pad(bc_mask[0], ((pad_x, pad_x), (pad_y, pad_y)), constant_values=0) + missing_mask_extended = jnp.pad(missing_mask, ((0, 0), (pad_x, pad_x), (pad_y, pad_y)), constant_values=True) + if dim == 3: + bc_mask_extended = jnp.pad(bc_mask[0], ((pad_x, pad_x), (pad_y, pad_y), (pad_z, pad_z)), constant_values=0) + missing_mask_extended = jnp.pad(missing_mask, ((0, 0), (pad_x, pad_x), (pad_y, pad_y), (pad_z, pad_z)), constant_values=True) + + # Iterate over boundary conditions and set the mask + for bc in bclist: + assert bc.indices is not None, f"Please specify indices associated with the {bc.__class__.__name__} BC!" + assert bc.mesh_vertices is None, ( + f"Please use operators based on MeshBoundaryMasker if {bc.__class__.__name__} is imposed on a mesh (e.g. STL)!" + ) + id_number = bc.id + bc_indices = np.array(bc.indices) + indices_origin = np.array(start_index)[:, np.newaxis] + if any(self.are_indices_in_interior(bc_indices, grid_shape)): + # If the indices are in the interior, we assume the usre specified indices are solid indices + solid_indices = bc_indices - indices_origin + solid_indices_shifted = solid_indices + shift + + # We obtain the boundary indices by padding the solid indices in all lattice directions + indices_padded = bc.pad_indices() - indices_origin + indices_shifted = indices_padded + shift + + # The missing mask is set to True meaning (exterior or solid nodes) using the original indices. + # This is because of the following streaming step which will assign missing directions for the boundary nodes. + if dim == 2: + missing_mask_extended = missing_mask_extended.at[:, solid_indices_shifted[0], solid_indices_shifted[1]].set(True) + else: + missing_mask_extended = missing_mask_extended.at[ + :, solid_indices_shifted[0], solid_indices_shifted[1], solid_indices_shifted[2] + ].set(True) + else: + indices_shifted = bc_indices - indices_origin + shift + + # Assign the boundary id to the shifted (and possibly padded) indices + bc_mask_extended = bc_mask_extended.at[tuple(indices_shifted)].set(id_number) + + # We are done with bc.indices. Remove them from BC objects + bc.__dict__.pop("indices", None) + + # Stream the missing mask to identify missing directions + missing_mask_extended = self.stream(missing_mask_extended) + + # Crop the extended masks to remove padding + if dim == 2: + missing_mask = missing_mask_extended[:, pad_x:-pad_x, pad_y:-pad_y] + bc_mask = bc_mask.at[0].set(bc_mask_extended[pad_x:-pad_x, pad_y:-pad_y]) + if dim == 3: + missing_mask = missing_mask_extended[:, pad_x:-pad_x, pad_y:-pad_y, pad_z:-pad_z] + bc_mask = bc_mask.at[0].set(bc_mask_extended[pad_x:-pad_x, pad_y:-pad_y, pad_z:-pad_z]) + return bc_mask, missing_mask + + def _construct_warp(self): + # Make constants for warp + _q = self.velocity_set.q + + @wp.func + def functional_domain_bounds( + index: Any, + bc_indices: Any, + id_number: Any, + is_interior: Any, + bc_mask: Any, + missing_mask: Any, + grid_shape: Any, + level: Any = 0, + ): + for ii in range(bc_indices.shape[1]): + # If the current index does not match the boundary condition index, we skip it + if not self.helper_masker.is_in_bc_indices(bc_mask, index, bc_indices, ii): + continue + + if is_interior[ii] == wp.uint8(True): + # If the index is in the interior, we set that index to be a solid node (identified by BC_SOLID) + # This information will be used in the next kernel to identify missing directions using the + # padded indices of the solid node that are associated with the boundary condition. + self.write_field(bc_mask, index, 0, wp.uint8(BC_SOLID)) + return + + # Set bc_mask for all bc indices + self.write_field(bc_mask, index, 0, wp.uint8(id_number[ii])) + + # Stream indices + for l in range(_q): + # Get the pull index which is the index of the neighboring node where information is pulled from + pull_index, _ = self.helper_masker.get_pull_index(bc_mask, l, index, level) + + # Check if pull index is out of bound + # These directions will have missing information after streaming + if not self.helper_masker.is_in_bounds(pull_index, grid_shape): + # Set the missing mask + self.write_field(missing_mask, index, l, wp.uint8(True)) + + @wp.func + def functional_interior_bc_mask( + index: Any, + bc_indices: Any, + id_number: Any, + bc_mask: Any, + ): + for ii in range(bc_indices.shape[1]): + # If the current index does not match the boundary condition index, we skip it + if not self.helper_masker.is_in_bc_indices(bc_mask, index, bc_indices, ii): + continue + # Set bc_mask for all interior bc indices + self.write_field(bc_mask, index, 0, wp.uint8(id_number[ii])) + + @wp.func + def functional_interior_missing_mask( + index: Any, + bc_indices: Any, + bc_mask: Any, + missing_mask: Any, + grid_shape: Any, + level: Any = 0, + ): + for ii in range(bc_indices.shape[1]): + # If the current index does not match the boundary condition index, we skip it + if not self.helper_masker.is_in_bc_indices(bc_mask, index, bc_indices, ii): + continue + for l in range(_q): + # Get the index of the streaming direction + pull_index, offset = self.helper_masker.get_pull_index(bc_mask, l, index, level) + + # Check if pull index is a fluid node (bc_mask is zero for fluid nodes) + bc_mask_ngh = self.read_field_neighbor(bc_mask, index, offset, 0) + if (self.helper_masker.is_in_bounds(pull_index, grid_shape)) and (bc_mask_ngh == wp.uint8(BC_SOLID)): + self.write_field(missing_mask, index, l, wp.uint8(True)) + + # Construct the warp 3D kernel + @wp.kernel + def kernel_domain_bounds( + bc_indices: wp.array2d(dtype=wp.int32), + id_number: wp.array1d(dtype=wp.uint8), + is_interior: wp.array1d(dtype=wp.uint8), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + grid_shape: wp.vec3i, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # Call the functional + functional_domain_bounds( + index, + bc_indices, + id_number, + is_interior, + bc_mask, + missing_mask, + grid_shape, + ) + + @wp.kernel + def kernel_interior_bc_mask( + bc_indices: wp.array2d(dtype=wp.int32), + id_number: wp.array1d(dtype=wp.uint8), + bc_mask: wp.array4d(dtype=wp.uint8), + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # Set bc_mask for all interior bc indices + functional_interior_bc_mask( + index, + bc_indices, + id_number, + bc_mask, + ) + return + + @wp.kernel + def kernel_interior_missing_mask( + bc_indices: wp.array2d(dtype=wp.int32), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + grid_shape: wp.vec3i, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + functional_interior_missing_mask(index, bc_indices, bc_mask, missing_mask, grid_shape) + + functional_dict = { + "functional_domain_bounds": functional_domain_bounds, + "functional_interior_bc_mask": functional_interior_bc_mask, + "functional_interior_missing_mask": functional_interior_missing_mask, + } + kernel_dict = { + "kernel_domain_bounds": kernel_domain_bounds, + "kernel_interior_bc_mask": kernel_interior_bc_mask, + "kernel_interior_missing_mask": kernel_interior_missing_mask, + } + return functional_dict, kernel_dict + + def _prepare_kernel_inputs(self, bclist, grid_shape, start_index=None): + """ + Prepare the inputs for the warp kernel by pre-allocating arrays and filling them with boundary condition information. + """ + + # Pre-allocate arrays with maximum possible size + max_size = sum( + len(bc.indices[0]) if isinstance(bc.indices, (list, tuple)) else bc.indices.shape[1] for bc in bclist if bc.indices is not None + ) + indices = np.zeros((3, max_size), dtype=np.int32) + id_numbers = np.zeros(max_size, dtype=np.uint8) + is_interior = np.zeros(max_size, dtype=np.uint8) + + current_index = 0 + for bc in bclist: + assert bc.indices is not None, f'Please specify indices associated with the {bc.__class__.__name__} BC using keyword "indices"!' + assert bc.mesh_vertices is None, ( + f"Please use operators based on MeshBoundaryMasker if {bc.__class__.__name__} is imposed on a mesh (e.g. STL)!" + ) + bc_indices = np.asarray(bc.indices) + num_indices = bc_indices.shape[1] + + # Normalize indices with respect to the start index if applicable + if start_index is not None: + bc_indices = bc_indices - np.array(start_index)[:, np.newaxis] + + # Ensure indices are 3D + if bc_indices.shape[0] == 2: + bc_indices = np.vstack([bc_indices, np.zeros(num_indices, dtype=int)]) + + # Add indices to the pre-allocated array + indices[:, current_index : current_index + num_indices] = bc_indices + + # Set id numbers + id_numbers[current_index : current_index + num_indices] = bc.id + + # Set is_interior flags + is_interior[current_index : current_index + num_indices] = self.are_indices_in_interior(bc_indices, grid_shape) + + current_index += num_indices + + # Remove indices from BC objects + # bc.__dict__.pop("indices", None) + + # Trim arrays to actual size + total_index = current_index + indices = indices[:, :total_index] + id_numbers = id_numbers[:total_index] + is_interior = is_interior[:total_index] + + # Convert to Warp arrays + def _to_wp_arrays(indices, id_numbers, is_interior, device=None): + return ( + wp.array(indices, dtype=wp.int32, device=device), + wp.array(id_numbers, dtype=wp.uint8, device=device), + wp.array(is_interior, dtype=wp.uint8, device=device), + ) + + if self.compute_backend == ComputeBackend.NEON: + grid = self.grid + ndevice = 1 if grid is None else grid.bk.get_num_devices() + + if ndevice == 1: + return _to_wp_arrays(indices, id_numbers, is_interior) + else: + # For multi-device, we need to split the indices across devices + wp_bc_indices = [] + wp_id_numbers = [] + wp_is_interior = [] + for i in range(ndevice): + device_name = grid.bk.get_device_name(i) + wp_bc_indices.append(wp.array(indices, dtype=wp.int32, device=device_name)) + wp_id_numbers.append(wp.array(id_numbers, dtype=wp.uint8, device=device_name)) + wp_is_interior.append(wp.array(is_interior, dtype=wp.uint8, device=device_name)) + return wp_bc_indices, wp_id_numbers, wp_is_interior + else: + return _to_wp_arrays(indices, id_numbers, is_interior) + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, bclist, bc_mask, missing_mask, start_index=None): + # get the grid shape + grid_shape = self.helper_masker.get_grid_shape(bc_mask) + + # Find interior boundary conditions + bc_interior = self._find_bclist_interior(bclist, grid_shape) + + # Prepare the first kernel inputs for all items in boundary condition list + wp_bc_indices, wp_id_numbers, wp_is_interior = self._prepare_kernel_inputs(bclist, grid_shape, start_index) + + # Launch the warp kernel + wp.launch( + self.warp_kernel["kernel_domain_bounds"], + dim=bc_mask.shape[1:], + inputs=[wp_bc_indices, wp_id_numbers, wp_is_interior, bc_mask, missing_mask, grid_shape], + ) + + # If there are no interior boundary conditions, skip the rest and retun early + if not bc_interior: + return bc_mask, missing_mask + + # Prepare the second and third kernel inputs for only a subset of boundary conditions associated with the interior + # Note 1: launching order of the following kernels are important here! + # Note 2: Due to race conditioning, the two kernels cannot be fused together. + wp_bc_indices, wp_id_numbers, _ = self._prepare_kernel_inputs(bc_interior, grid_shape) + wp.launch( + self.warp_kernel["kernel_interior_missing_mask"], + dim=bc_mask.shape[1:], + inputs=[wp_bc_indices, bc_mask, missing_mask, grid_shape], + ) + wp.launch( + self.warp_kernel["kernel_interior_bc_mask"], + dim=bc_mask.shape[1:], + inputs=[ + wp_bc_indices, + wp_id_numbers, + bc_mask, + ], + ) + + return bc_mask, missing_mask + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional_dict, _ = self._construct_warp() + functional_domain_bounds = functional_dict.get("functional_domain_bounds") + functional_interior_bc_mask = functional_dict.get("functional_interior_bc_mask") + functional_interior_missing_mask = functional_dict.get("functional_interior_missing_mask") + + @neon.Container.factory(name="IndicesBoundaryMasker_DomainBounds") + def container_domain_bounds( + wp_bc_indices_, + wp_id_numbers_, + wp_is_interior_, + bc_mask, + missing_mask, + grid_shape, + ): + def domain_bounds_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + missing_mask_pn = loader.get_write_handle(missing_mask) + grid = bc_mask.get_grid() + bk = grid.backend + if bk.get_num_devices() == 1: + # If there is only one device, we can use the warp arrays directly + wp_bc_indices = wp_bc_indices_ + wp_id_numbers = wp_id_numbers_ + wp_is_interior = wp_is_interior_ + else: + device_id = loader.get_device_id() + wp_bc_indices = wp_bc_indices_[device_id] + wp_id_numbers = wp_id_numbers_[device_id] + wp_is_interior = wp_is_interior_[device_id] + + @wp.func + def domain_bounds_kernel(index: Any): + # apply the functional + functional_domain_bounds( + index, + wp_bc_indices, + wp_id_numbers, + wp_is_interior, + bc_mask_pn, + missing_mask_pn, + grid_shape, + ) + + loader.declare_kernel(domain_bounds_kernel) + + return domain_bounds_launcher + + @neon.Container.factory(name="IndicesBoundaryMasker_InteriorBcMask") + def container_interior_bc_mask( + wp_bc_indices, + wp_id_numbers, + bc_mask, + ): + def interior_bc_mask_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + + @wp.func + def interior_bc_mask_kernel(index: Any): + # apply the functional + functional_interior_bc_mask( + index, + wp_bc_indices, + wp_id_numbers, + bc_mask_pn, + ) + + loader.declare_kernel(interior_bc_mask_kernel) + + return interior_bc_mask_launcher + + @neon.Container.factory(name="IndicesBoundaryMasker_InteriorMissingMask") + def container_interior_missing_mask( + wp_bc_indices, + bc_mask, + missing_mask, + grid_shape, + ): + def interior_bc_mask_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + missing_mask_pn = loader.get_write_handle(missing_mask) + + @wp.func + def interior_missing_mask_kernel(index: Any): + # apply the functional + functional_interior_missing_mask( + index, + wp_bc_indices, + bc_mask_pn, + missing_mask_pn, + grid_shape, + ) + + loader.declare_kernel(interior_missing_mask_kernel) + + return interior_bc_mask_launcher + + container_dict = { + "container_domain_bounds": container_domain_bounds, + "container_interior_bc_mask": container_interior_bc_mask, + "container_interior_missing_mask": container_interior_missing_mask, + } + + return functional_dict, container_dict + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, bclist, bc_mask, missing_mask, start_index=None): + import neon + + # get the grid shape + grid_shape = self.helper_masker.get_grid_shape(bc_mask) + + # Find interior boundary conditions + bc_interior = self._find_bclist_interior(bclist, grid_shape) + + # Prepare the first kernel inputs for all items in boundary condition list + wp_bc_indices, wp_id_numbers, wp_is_interior = self._prepare_kernel_inputs(bclist, grid_shape, start_index) + + # Launch the first container + container_domain_bounds = self.neon_container["container_domain_bounds"]( + wp_bc_indices, + wp_id_numbers, + wp_is_interior, + bc_mask, + missing_mask, + grid_shape, + ) + container_domain_bounds.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + # If there are no interior boundary conditions, skip the rest and retun early + if not bc_interior: + return bc_mask, missing_mask + + # Prepare the second and third kernel inputs for only a subset of boundary conditions associated with the interior + # Note 1: launching order of the following kernels are important here! + # Note 2: Due to race conditioning, the two kernels cannot be fused together. + wp_bc_indices, wp_id_numbers, _ = self._prepare_kernel_inputs(bc_interior, grid_shape) + container_interior_missing_mask = self.neon_container["container_interior_missing_mask"](wp_bc_indices, bc_mask, missing_mask, grid_shape) + container_interior_missing_mask.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + # Launch the third container + container_interior_bc_mask = self.neon_container["container_interior_bc_mask"]( + wp_bc_indices, + wp_id_numbers, + bc_mask, + ) + container_interior_bc_mask.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + return bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/mesh_boundary_masker.py b/xlb/operator/boundary_masker/mesh_boundary_masker.py new file mode 100644 index 00000000..c6fb778d --- /dev/null +++ b/xlb/operator/boundary_masker/mesh_boundary_masker.py @@ -0,0 +1,245 @@ +""" +Abstract base class for mesh-based boundary maskers. + +Provides shared input preparation logic (mesh construction, kernel arrays) +used by AABB, Ray, Winding, and AABB-Close masker subclasses. +""" + +import numpy as np +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.boundary_masker.helper_functions_masker import HelperFunctionsMasker + + +class MeshBoundaryMasker(Operator): + """ + Operator for creating a boundary missing_mask from a mesh file + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + + assert self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON], ( + f"MeshBoundaryMasker is only implemented for {ComputeBackend.WARP} and {ComputeBackend.NEON} backends!" + ) + + assert self.velocity_set.d == 3, "MeshBoundaryMasker is only implemented for 3D velocity sets!" + # Raise error if used for 2d examples: + if self.velocity_set.d == 2: + raise NotImplementedError("This Operator is not implemented in 2D!") + + # Make constants for warp + _c = self.velocity_set.c + _q = self.velocity_set.q + + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + # Define masker helper functions + self.helper_masker = HelperFunctionsMasker( + velocity_set=self.velocity_set, + precision_policy=self.precision_policy, + compute_backend=self.compute_backend, + ) + + @wp.func + def out_of_bound_pull_index( + lattice_dir: wp.int32, + index: wp.vec3i, + field: wp.array4d(dtype=wp.uint8), + grid_shape: wp.vec3i, + ): + # Get the index of the streaming direction + pull_index = wp.vec3i() + for d in range(self.velocity_set.d): + pull_index[d] = index[d] - _c[d, lattice_dir] + + # check if pull index is out of bound + # These directions will have missing information after streaming + missing = not self.helper_masker.is_in_bounds(pull_index, grid_shape) + return missing + + # Function to precompute useful values per triangle, assuming spacing is (1,1,1) + # inputs: verts: triangle vertices, normal: triangle unit normal + # outputs: dist1, dist2, normal_edge0, normal_edge1, dist_edge + @wp.func + def pre_compute( + verts: wp.mat33f, # triangle vertices + normal: wp.vec3f, # triangle normal + ): + corner = wp.vec3f(float(normal[0] > 0.0), float(normal[1] > 0.0), float(normal[2] > 0.0)) + + dist1 = wp.dot(normal, corner - verts[0]) + dist2 = wp.dot(normal, wp.vec3f(1.0, 1.0, 1.0) - corner - verts[0]) + + edges = wp.transpose(wp.mat33(verts[1] - verts[0], verts[2] - verts[1], verts[0] - verts[2])) + normal_edge0 = wp.mat33f(0.0) + normal_edge1 = wp.mat33f(0.0) + dist_edge = wp.mat33f(0.0) + + for axis0 in range(0, 3): + axis1 = (axis0 + 1) % 3 + axis2 = (axis0 + 2) % 3 + + sgn = 1.0 + if normal[axis2] < 0.0: + sgn = -1.0 + + for i in range(0, 3): + normal_edge0[i, axis0] = -1.0 * sgn * edges[i, axis1] + normal_edge1[i, axis0] = sgn * edges[i, axis0] + + dist_edge[i, axis0] = ( + -1.0 * (normal_edge0[i, axis0] * verts[i, axis0] + normal_edge1[i, axis0] * verts[i, axis1]) + + wp.max(0.0, normal_edge0[i, axis0]) + + wp.max(0.0, normal_edge1[i, axis0]) + ) + + return dist1, dist2, normal_edge0, normal_edge1, dist_edge + + # Check whether this triangle intersects the unit cube at position low + @wp.func + def triangle_box_intersect( + low: wp.vec3f, + normal: wp.vec3f, + dist1: wp.float32, + dist2: wp.float32, + normal_edge0: wp.mat33f, + normal_edge1: wp.mat33f, + dist_edge: wp.mat33f, + ): + if (wp.length(normal) > 0.0) and (wp.dot(normal, low) + dist1) * (wp.dot(normal, low) + dist2) <= 0.0: + intersect = True + # Loop over primary axis for projection + for ax0 in range(0, 3): + ax1 = (ax0 + 1) % 3 + for i in range(0, 3): + intersect = intersect and (normal_edge0[i, ax0] * low[ax0] + normal_edge1[i, ax0] * low[ax1] + dist_edge[i, ax0] >= 0.0) + + return intersect + else: + return False + + # Check whether the unit voxel at position low intersects the warp mesh, assumes mesh has valid normals + # inputs: mesh_id: mesh id, low: position of the voxel + # outputs: True if intersection, False otherwise + @wp.func + def mesh_voxel_intersect(mesh_id: wp.uint64, low: wp.vec3): + query = wp.mesh_query_aabb(mesh_id, low, low + wp.vec3f(1.0, 1.0, 1.0)) + + for f in query: + v0 = wp.mesh_eval_position(mesh_id, f, 1.0, 0.0) + v1 = wp.mesh_eval_position(mesh_id, f, 0.0, 1.0) + v2 = wp.mesh_eval_position(mesh_id, f, 0.0, 0.0) + normal = wp.mesh_eval_face_normal(mesh_id, f) + + v = wp.transpose(wp.mat33f(v0, v1, v2)) + + # TODO: run this on triangles in advance + dist1, dist2, normal_edge0, normal_edge1, dist_edge = pre_compute(verts=v, normal=normal) + + if triangle_box_intersect( + low=low, normal=normal, dist1=dist1, dist2=dist2, normal_edge0=normal_edge0, normal_edge1=normal_edge1, dist_edge=dist_edge + ): + return True + + return False + + @wp.kernel + def resolve_out_of_bound_kernel( + id_number: wp.int32, + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # domain shape to check for out of bounds + grid_shape = wp.vec3i(bc_mask.shape[1], bc_mask.shape[2], bc_mask.shape[3]) + + # Find the fractional distance to the mesh in each direction + if bc_mask[0, index[0], index[1], index[2]] == wp.uint8(id_number): + for l in range(1, _q): + # Ensuring out of bound pull indices are properly considered in the missing_mask + if out_of_bound_pull_index(l, index, missing_mask, grid_shape): + missing_mask[l, index[0], index[1], index[2]] = wp.uint8(True) + + # Construct some helper warp functions + self.mesh_voxel_intersect = mesh_voxel_intersect + self.resolve_out_of_bound_kernel = resolve_out_of_bound_kernel + + def _prepare_kernel_inputs( + self, + bc, + bc_mask, + ): + assert bc.mesh_vertices is not None, f'Please provide the mesh vertices for {bc.__class__.__name__} BC using keyword "mesh_vertices"!' + assert bc.indices is None, f"Please use IndicesBoundaryMasker operator if {bc.__class__.__name__} is imposed on known indices of the grid!" + assert bc.mesh_vertices.shape[1] == self.velocity_set.d, ( + "Mesh points must be reshaped into an array (N, 3) where N indicates number of points!" + ) + + grid_shape = self.helper_masker.get_grid_shape(bc_mask) # (nx, ny, nz) + mesh_vertices = bc.mesh_vertices + mesh_min = np.min(mesh_vertices, axis=0) + mesh_max = np.max(mesh_vertices, axis=0) + + if any(mesh_min < 0) or any(mesh_max >= grid_shape): + raise ValueError( + f"Mesh extents ({mesh_min}, {mesh_max}) exceed domain dimensions {grid_shape}. The mesh must be fully contained within the domain." + ) + + # We are done with bc.mesh_vertices. Remove them from BC objects + bc.__dict__.pop("mesh_vertices", None) + + mesh_indices = np.arange(mesh_vertices.shape[0]) + mesh = wp.Mesh( + points=wp.array(mesh_vertices, dtype=wp.vec3), + indices=wp.array(mesh_indices, dtype=wp.int32), + ) + mesh_id = wp.uint64(mesh.id) + bc_id = bc.id + return mesh_id, bc_id + + @Operator.register_backend(ComputeBackend.JAX) + def jax_implementation( + self, + bc, + bc_mask, + missing_mask, + ): + raise NotImplementedError(f"Operation {self.__class__.__name__} not implemented in JAX!") + + def warp_implementation_base( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + # Prepare inputs + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + # Launch the appropriate warp kernel + wp.launch( + self.warp_kernel, + inputs=[mesh_id, bc_id, distances, bc_mask, missing_mask, wp.static(bc.needs_mesh_distance)], + dim=bc_mask.shape[1:], + ) + wp.launch( + self.resolve_out_of_bound_kernel, + inputs=[bc_id, bc_mask, missing_mask], + dim=bc_mask.shape[1:], + ) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/mesh_voxelization_method.py b/xlb/operator/boundary_masker/mesh_voxelization_method.py new file mode 100644 index 00000000..b0162de7 --- /dev/null +++ b/xlb/operator/boundary_masker/mesh_voxelization_method.py @@ -0,0 +1,55 @@ +""" +Mesh voxelization method registry. + +Defines the available voxelization strategies (AABB, Ray, AABB-Close, +Winding) and provides a factory function to create the corresponding +:class:`VoxelizationMethod` data object. +""" + +from dataclasses import dataclass + + +METHODS = { + "AABB": 1, + "RAY": 2, + "AABB_CLOSE": 3, + "WINDING": 4, +} + + +@dataclass +class VoxelizationMethod: + """Describes a mesh voxelization strategy. + + Attributes + ---------- + id : int + Numeric identifier for the method. + name : str + Human-readable name (``"AABB"``, ``"RAY"``, etc.). + options : dict + Extra options (e.g. ``close_voxels`` for AABB_CLOSE). + """ + + id: int + name: str + options: dict + + +def MeshVoxelizationMethod(name: str, **options): + """Create a :class:`VoxelizationMethod` by name. + + Parameters + ---------- + name : str + One of ``"AABB"``, ``"RAY"``, ``"AABB_CLOSE"``, ``"WINDING"``. + **options + Additional keyword arguments forwarded to + ``VoxelizationMethod.options``. + + Returns + ------- + VoxelizationMethod + """ + assert name in METHODS.keys(), f"Unsupported voxelization method: {name}" + return VoxelizationMethod(METHODS[name], name, options) diff --git a/xlb/operator/boundary_masker/multires_aabb.py b/xlb/operator/boundary_masker/multires_aabb.py new file mode 100644 index 00000000..7d6df6c6 --- /dev/null +++ b/xlb/operator/boundary_masker/multires_aabb.py @@ -0,0 +1,101 @@ +""" +Multi-resolution AABB mesh-based boundary masker for the Neon backend. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker import MeshMaskerAABB +from xlb.operator.operator import Operator + + +class MultiresMeshMaskerAABB(MeshMaskerAABB): + """ + Operator for creating boundary missing_mask from mesh using Axis-Aligned Bounding Box (AABB) voxelization in multiresolution simulations. + + This implementation uses warp.mesh_query_aabb for efficient mesh-voxel intersection testing, + providing approximate 1-voxel thick surface detection around the mesh geometry. + Suitable for scenarios where fast, approximate boundary detection is sufficient. + TODO@Hesam: + Right now, we cannot properly mask a mesh file if it lives on any level other than the finest. This issue can be easily solved by adding + gx = wp.neon_get_x(cIdx) // 2 ** level + gy = wp.neon_get_y(cIdx) // 2 ** level + gz = wp.neon_get_z(cIdx) // 2 ** level + to the "neon_index_to_warp" and subsequently add "level" to the arguments of "index_to_position_neon", "get_pull_index_neon" and + "is_in_bc_indices_neon". In order to extract "level" from the "neon_field_hdl" we can use the function wp.neon_level(neon_field_hdl). + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MeshMaskerAABB") + def container( + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + level: Any, + ): + def aabb_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + distances_pn = loader.get_mres_write_handle(distances) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + + @wp.func + def aabb_kernel(index: Any): + # apply the functional + functional( + index, + mesh_id, + id_number, + distances_pn, + bc_mask_pn, + missing_mask_pn, + needs_mesh_distance, + ) + + loader.declare_kernel(aabb_kernel) + + return aabb_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + stream=0, + ): + import neon + + # Prepare inputs + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + grid = bc_mask.get_grid() + for level in range(grid.num_levels): + # Launch the neon container + c = self.neon_container(mesh_id, bc_id, distances, bc_mask, missing_mask, wp.static(bc.needs_mesh_distance), level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/multires_aabb_close.py b/xlb/operator/boundary_masker/multires_aabb_close.py new file mode 100644 index 00000000..31e4da29 --- /dev/null +++ b/xlb/operator/boundary_masker/multires_aabb_close.py @@ -0,0 +1,275 @@ +""" +Multi-resolution AABB-Close boundary masker with morphological closing. + +Extends the AABB-Close masker for Neon multi-resolution grids, applying +dilate-then-erode operations to fill narrow channels with solid voxels. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker import MeshMaskerAABBClose +from xlb.operator.operator import Operator +from xlb.cell_type import BC_SOLID + + +class MultiresMeshMaskerAABBClose(MeshMaskerAABBClose): + """ + Operator for creating boundary missing_mask from mesh using Axis-Aligned Bounding Box (AABB) voxelization + in multiresolution simulations (NEON backend). It takes in a number of close_voxels to perform morphological + operations (dilate followed by erode) to ensure small channels are filled with solid voxels. + + This version provides NEON-specific functionals working on multires partitions (mPartition) and bIndex. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + close_voxels: int = None, + ): + super().__init__(velocity_set, precision_policy, compute_backend, close_voxels) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + # Build and store NEON dicts + self.neon_functional_dict, self.neon_container_dict = self._construct_neon() + + def _construct_neon(self): + import neon + + # Use the warp functionals from the base (for reference), but implement NEON variants here + functional_dict_warp, _ = self._construct_warp() + functional_erode_warp = functional_dict_warp.get("functional_erode") + functional_dilate_warp = functional_dict_warp.get("functional_dilate") + functional_solid = functional_dict_warp.get("functional_solid") + # We will not directly reuse functional_solid / functional_aabb from warp; we write NEON-specific ones. + + # We also need lattice info for neighbor iteration + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + + # Set local constants + lattice_central_index = self.velocity_set.center_index + + # Main AABB close: sets bc_mask, missing_mask, distances based on solid_mask + # bc_mask: wp.uint8, missing_mask: wp.uint8, distances: dtype from precision policy (float) + @wp.func + def mres_functional_aabb( + index: Any, + mesh_id: wp.uint64, + id_number: wp.int32, + distances_pn: Any, # mPartition(dtype=distance type), cardinality=_q + bc_mask_pn: Any, # mPartition_uint8, cardinality=1 + missing_mask_pn: Any, # mPartition_uint8, cardinality=_q + solid_mask_pn: Any, # mPartition_uint8, cardinality=1 + needs_mesh_distance: bool, + ): + # Cell center from bc_mask partition + cell_center = self.helper_masker.index_to_position(bc_mask_pn, index) + + # If already solid or bc, mark solid + solid_val = wp.neon_read(solid_mask_pn, index, 0) + bc_val = wp.neon_read(bc_mask_pn, index, 0) + if solid_val == wp.uint8(BC_SOLID) or bc_val == wp.uint8(BC_SOLID): + wp.neon_write(bc_mask_pn, index, 0, wp.uint8(BC_SOLID)) + return + + # loop lattice directions + for direction_idx in range(_q): + # skip central if provided by velocity set + if direction_idx == lattice_central_index: + continue + + # If neighbor index is valid at this resolution level + ngh = wp.neon_ngh_idx(wp.int8(_c[0, direction_idx]), wp.int8(_c[1, direction_idx]), wp.int8(_c[2, direction_idx])) + is_valid = wp.bool(False) + nval = wp.neon_read_ngh(solid_mask_pn, index, ngh, 0, wp.uint8(0), is_valid) + if is_valid: + if nval == wp.uint8(BC_SOLID): + # Found solid neighbor -> boundary cell + self.write_field(bc_mask_pn, index, 0, wp.uint8(id_number)) + self.write_field(missing_mask_pn, index, _opp_indices[direction_idx], wp.uint8(True)) + + if not needs_mesh_distance: + # No distance needed; continue to next direction + continue + + # Compute mesh distance along lattice direction + dir_vec = wp.vec3f( + wp.float32(_c[0, direction_idx]), + wp.float32(_c[1, direction_idx]), + wp.float32(_c[2, direction_idx]), + ) + max_length = wp.length(dir_vec) + # Avoid division by zero for any pathological dir (shouldn't happen) + norm_dir = dir_vec / (max_length if max_length > 0.0 else 1.0) + query = wp.mesh_query_ray(mesh_id, cell_center, norm_dir, 1.5 * max_length) + if query.result: + pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) + dist = wp.length(pos_mesh - cell_center) - 0.5 * max_length + weight = dist / (max_length if max_length > 0.0 else 1.0) + # distances has cardinality _q; store into this channel + self.write_field(distances_pn, index, direction_idx, self.store_dtype(weight)) + else: + self.write_field(distances_pn, index, direction_idx, self.store_dtype(1.0)) + + # Containers + + # Erode: f_field -> f_field_out + @neon.Container.factory(name="Erode") + def container_erode(f_field: wp.array3d(dtype=Any), f_field_out: wp.array3d(dtype=Any), level: int): + def erode_launcher(loader: neon.Loader): + loader.set_mres_grid(f_field.get_grid(), level) + f_field_pn = loader.get_mres_read_handle(f_field) + f_field_out_pn = loader.get_mres_write_handle(f_field_out) + + @wp.func + def erode_kernel(index: Any): + functional_erode_warp(index, f_field_pn, f_field_out_pn) + + loader.declare_kernel(erode_kernel) + + return erode_launcher + + # Dilate: f_field -> f_field_out + @neon.Container.factory(name="Dilate") + def container_dilate(f_field: wp.array3d(dtype=Any), f_field_out: wp.array3d(dtype=Any), level: int): + def dilate_launcher(loader: neon.Loader): + loader.set_mres_grid(f_field.get_grid(), level) + f_field_pn = loader.get_mres_read_handle(f_field) + f_field_out_pn = loader.get_mres_write_handle(f_field_out) + + @wp.func + def dilate_kernel(index: Any): + functional_dilate_warp(index, f_field_pn, f_field_out_pn) + + loader.declare_kernel(dilate_kernel) + + return dilate_launcher + + # Solid mask: voxelize mesh into solid_mask + @neon.Container.factory(name="Solid") + def container_solid(mesh_id: wp.uint64, solid_mask: wp.array3d(dtype=wp.uint8), level: int): + def solid_launcher(loader: neon.Loader): + loader.set_mres_grid(solid_mask.get_grid(), level) + solid_mask_pn = loader.get_mres_write_handle(solid_mask) + + @wp.func + def solid_kernel(index: Any): + # apply the functional + functional_solid(index, mesh_id, solid_mask_pn, wp.vec3f(0.0, 0.0, 0.0)) + + loader.declare_kernel(solid_kernel) + + return solid_launcher + + # Main AABB container + @neon.Container.factory(name="MeshMaskerAABBClose") + def container( + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + solid_mask: Any, + needs_mesh_distance: Any, + level: Any, + ): + def aabb_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + distances_pn = loader.get_mres_write_handle(distances) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + solid_mask_pn = loader.get_mres_write_handle(solid_mask) + + @wp.func + def aabb_kernel(index: Any): + mres_functional_aabb( + index, + mesh_id, + id_number, + distances_pn, + bc_mask_pn, + missing_mask_pn, + solid_mask_pn, + needs_mesh_distance, + ) + + loader.declare_kernel(aabb_kernel) + + return aabb_launcher + + container_dict = { + "container_erode": container_erode, + "container_dilate": container_dilate, + "container_solid": container_solid, + "container_aabb": container, + } + + # Expose NEON functionals too (in case callers want to reuse) + functional_dict = { + "mres_functional_aabb": mres_functional_aabb, + } + + return functional_dict, container_dict + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + stream=0, + ): + import neon + + # Prepare inputs + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + grid = bc_mask.get_grid() + # Create fields using new_field + solid_mask = grid.new_field(cardinality=1, dtype=wp.uint8, memory_type=neon.MemoryType.device()) + solid_mask_out = grid.new_field( + cardinality=1, + dtype=wp.uint8, + memory_type=neon.MemoryType.device(), + # memory_type=neon.MemoryType.host_device() + ) + + for level in range(grid.num_levels): + # Initialize to 0 + solid_mask.fill_run(level=level, value=wp.uint8(0), stream_idx=stream) + solid_mask_out.fill_run(level=level, value=wp.uint8(0), stream_idx=stream) + + # Launch the neon containers + container_solid = self.neon_container_dict["container_solid"](mesh_id, solid_mask, level) + container_solid.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + for _ in range(self.close_voxels): + container_dilate = self.neon_container_dict["container_dilate"](solid_mask, solid_mask_out, level) + container_dilate.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + solid_mask, solid_mask_out = solid_mask_out, solid_mask + + if self.close_voxels % 2 > 0: + solid_mask, solid_mask_out = solid_mask_out, solid_mask + + for _ in range(self.close_voxels): + container_erode = self.neon_container_dict["container_erode"](solid_mask_out, solid_mask, level) + container_erode.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + solid_mask, solid_mask_out = solid_mask_out, solid_mask + + if self.close_voxels % 2 > 0: + solid_mask, solid_mask_out = solid_mask_out, solid_mask + + container_aabb = self.neon_container_dict["container_aabb"]( + mesh_id, bc_id, distances, bc_mask, missing_mask, solid_mask, wp.static(bc.needs_mesh_distance), level + ) + container_aabb.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/multires_indices_boundary_masker.py b/xlb/operator/boundary_masker/multires_indices_boundary_masker.py new file mode 100644 index 00000000..4ee67a16 --- /dev/null +++ b/xlb/operator/boundary_masker/multires_indices_boundary_masker.py @@ -0,0 +1,212 @@ +""" +Multi-resolution indices-based boundary masker for the Neon backend. + +Creates boundary masks from explicit voxel indices on multi-resolution +grids, computing missing-population masks for each tagged voxel. +""" + +from typing import Any +import copy +import numpy as np + +import warp as wp + +from xlb.operator.operator import Operator +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker import IndicesBoundaryMasker + + +class MultiresIndicesBoundaryMasker(IndicesBoundaryMasker): + """ + Operator for creating a boundary mask using indices of boundary conditions in a multi-resolution setting. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional_dict, _ = self._construct_warp() + functional_domain_bounds = functional_dict.get("functional_domain_bounds") + functional_interior_bc_mask = functional_dict.get("functional_interior_bc_mask") + functional_interior_missing_mask = functional_dict.get("functional_interior_missing_mask") + + @neon.Container.factory(name="IndicesBoundaryMasker_DomainBounds") + def container_domain_bounds( + wp_bc_indices, + wp_id_numbers, + wp_is_interior, + bc_mask, + missing_mask, + grid_shape, + level, + ): + def domain_bounds_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + + @wp.func + def domain_bounds_kernel(index: Any): + # apply the functional + functional_domain_bounds( + index, + wp_bc_indices, + wp_id_numbers, + wp_is_interior, + bc_mask_pn, + missing_mask_pn, + grid_shape, + level, + ) + + loader.declare_kernel(domain_bounds_kernel) + + return domain_bounds_launcher + + @neon.Container.factory(name="IndicesBoundaryMasker_InteriorBcMask") + def container_interior_bc_mask( + wp_bc_indices, + wp_id_numbers, + bc_mask, + level, + ): + def interior_bc_mask_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + + @wp.func + def interior_bc_mask_kernel(index: Any): + # apply the functional + functional_interior_bc_mask( + index, + wp_bc_indices, + wp_id_numbers, + bc_mask_pn, + ) + + loader.declare_kernel(interior_bc_mask_kernel) + + return interior_bc_mask_launcher + + @neon.Container.factory(name="IndicesBoundaryMasker_InteriorMissingMask") + def container_interior_missing_mask( + wp_bc_indices, + bc_mask, + missing_mask, + grid_shape, + level, + ): + def interior_bc_mask_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + + @wp.func + def interior_missing_mask_kernel(index: Any): + # apply the functional + functional_interior_missing_mask( + index, + wp_bc_indices, + bc_mask_pn, + missing_mask_pn, + grid_shape, + level, + ) + + loader.declare_kernel(interior_missing_mask_kernel) + + return interior_bc_mask_launcher + + container_dict = { + "container_domain_bounds": container_domain_bounds, + "container_interior_bc_mask": container_interior_bc_mask, + "container_interior_missing_mask": container_interior_missing_mask, + } + + return functional_dict, container_dict + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, bclist, bc_mask, missing_mask, start_index=None): + import neon + + grid = bc_mask.get_grid() + num_levels = grid.num_levels + grid_shape_finest = self.helper_masker.get_grid_shape(bc_mask) + for level in range(num_levels): + # Create a copy of the boundary condition list for the current level if the indices at that level are not empty + bclist_at_level = [] + for bc in bclist: + if bc.indices is not None and bc.indices[level]: + bc_copy = copy.copy(bc) # shallow copy of the whole object + indices = copy.deepcopy(bc.indices[level]) # deep copy only the modified part + indices = np.array(indices) * 2**level # TODO: This is a hack + bc_copy.indices = tuple(indices.tolist()) # convert to tuple + bclist_at_level.append(bc_copy) + + # If the boundary condition list is empty, skip to the next level + if not bclist_at_level: + continue + + # find grid shape at current level + # TODO: this is a hack. Should be corrected in the helper function when getting neon global indices + grid_shape_at_level = tuple([shape // 2**level for shape in grid_shape_finest]) + grid_shape_finest_warp = wp.vec3i(*grid_shape_finest) + + # find interior boundary conditions + bc_interior = self._find_bclist_interior(bclist_at_level, grid_shape_at_level) + + # Prepare the first kernel inputs for all items in boundary condition list + wp_bc_indices, wp_id_numbers, wp_is_interior = self._prepare_kernel_inputs(bclist_at_level, grid_shape_at_level) + + # Launch the first container + container_domain_bounds = self.neon_container["container_domain_bounds"]( + wp_bc_indices, + wp_id_numbers, + wp_is_interior, + bc_mask, + missing_mask, + grid_shape_finest_warp, + level, + ) + container_domain_bounds.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + # If there are no interior boundary conditions, skip the rest of the processing for this level + if not bc_interior: + continue + + # Prepare the second and third kernel inputs for only a subset of boundary conditions associated with the interior + # Note 1: launching order of the following kernels are important here! + # Note 2: Due to race conditioning, the two kernels cannot be fused together. + wp_bc_indices, wp_id_numbers, _ = self._prepare_kernel_inputs(bc_interior, grid_shape_at_level) + container_interior_missing_mask = self.neon_container["container_interior_missing_mask"]( + wp_bc_indices, + bc_mask, + missing_mask, + grid_shape_finest_warp, + level, + ) + container_interior_missing_mask.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + # Launch the third container + container_interior_bc_mask = self.neon_container["container_interior_bc_mask"]( + wp_bc_indices, + wp_id_numbers, + bc_mask, + level, + ) + container_interior_bc_mask.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + return bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/multires_ray.py b/xlb/operator/boundary_masker/multires_ray.py new file mode 100644 index 00000000..6e46a954 --- /dev/null +++ b/xlb/operator/boundary_masker/multires_ray.py @@ -0,0 +1,92 @@ +""" +Multi-resolution ray-cast mesh-based boundary masker for the Neon backend. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker import MeshMaskerRay +from xlb.operator.operator import Operator + + +class MultiresMeshMaskerRay(MeshMaskerRay): + """ + Operator for creating a boundary missing_mask from an STL file in multiresolution simulations. + + This implementation uses warp.mesh_query_ray for efficient mesh-voxel intersection testing. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MeshMaskerRay") + def container( + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + level: Any, + ): + def ray_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + distances_pn = loader.get_mres_write_handle(distances) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + + @wp.func + def ray_kernel(index: Any): + # apply the functional + functional( + index, + mesh_id, + id_number, + distances_pn, + bc_mask_pn, + missing_mask_pn, + needs_mesh_distance, + ) + + loader.declare_kernel(ray_kernel) + + return ray_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + stream=0, + ): + import neon + + # Prepare inputs + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + grid = bc_mask.get_grid() + for level in range(grid.num_levels): + # Launch the neon container + c = self.neon_container(mesh_id, bc_id, distances, bc_mask, missing_mask, wp.static(bc.needs_mesh_distance), level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/ray.py b/xlb/operator/boundary_masker/ray.py new file mode 100644 index 00000000..5a78d91c --- /dev/null +++ b/xlb/operator/boundary_masker/ray.py @@ -0,0 +1,177 @@ +""" +Ray-cast mesh-based boundary masker. + +Voxelizes a mesh file by casting rays along each lattice direction using +``warp.mesh_query_ray`` to detect surface crossings. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker +from xlb.operator.operator import Operator + + +class MeshMaskerRay(MeshBoundaryMasker): + """ + Operator for creating a boundary missing_mask from a mesh file + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + + def _construct_warp(self): + # Make constants for warp + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + + # Set local constants + lattice_central_index = self.velocity_set.center_index + + @wp.func + def functional( + index: Any, + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + ): + # position of the point + cell_center_pos = self.helper_masker.index_to_position(bc_mask, index) + + # Find the fractional distance to the mesh in each direction + for direction_idx in range(_q): + if direction_idx == lattice_central_index: + # Skip the central index as it is not relevant for boundary masking + continue + + direction_vec = wp.vec3f(wp.float32(_c[0, direction_idx]), wp.float32(_c[1, direction_idx]), wp.float32(_c[2, direction_idx])) + # Max length depends on ray direction (diagonals are longer) + max_length = wp.length(direction_vec) + query = wp.mesh_query_ray(mesh_id, cell_center_pos, direction_vec / max_length, max_length) + if query.result: + # Set the boundary id and missing_mask + self.write_field(bc_mask, index, 0, wp.uint8(id_number)) + self.write_field(missing_mask, index, _opp_indices[direction_idx], wp.uint8(True)) + + # If we don't need the mesh distance, we can return early + if not needs_mesh_distance: + continue + + # get position of the mesh triangle that intersects with the ray + pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) + dist = wp.length(pos_mesh - cell_center_pos) + weight = self.store_dtype(dist / max_length) + self.write_field(distances, index, direction_idx, self.store_dtype(weight)) + + @wp.kernel + def kernel( + mesh_id: wp.uint64, + id_number: wp.int32, + distances: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + needs_mesh_distance: bool, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # apply the functional + functional( + index, + mesh_id, + id_number, + distances, + bc_mask, + missing_mask, + needs_mesh_distance, + ) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + return self.warp_implementation_base( + bc, + distances, + bc_mask, + missing_mask, + ) + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MeshMaskerRay") + def container( + mesh_id: Any, + id_number: Any, + distances: Any, + bc_mask: Any, + missing_mask: Any, + needs_mesh_distance: Any, + ): + def ray_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + missing_mask_pn = loader.get_write_handle(missing_mask) + distances_pn = loader.get_write_handle(distances) + + @wp.func + def ray_kernel(index: Any): + # apply the functional + functional( + index, + mesh_id, + id_number, + distances_pn, + bc_mask_pn, + missing_mask_pn, + needs_mesh_distance, + ) + + loader.declare_kernel(ray_kernel) + + return ray_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + # Prepare inputs + import neon + + mesh_id, bc_id = self._prepare_kernel_inputs(bc, bc_mask) + + # Launch the appropriate neon container + c = self.neon_container(mesh_id, bc_id, distances, bc_mask, missing_mask, wp.static(bc.needs_mesh_distance)) + c.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + return distances, bc_mask, missing_mask diff --git a/xlb/operator/boundary_masker/winding.py b/xlb/operator/boundary_masker/winding.py new file mode 100644 index 00000000..1510f3a5 --- /dev/null +++ b/xlb/operator/boundary_masker/winding.py @@ -0,0 +1,115 @@ +""" +Winding-number mesh-based boundary masker. + +Uses the generalized winding-number test (``warp.mesh_query_point``) to +classify voxels as inside or outside the mesh, providing a +solid-detection method even for non-watertight geometries. +""" + +import warp as wp +from typing import Any +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker +from xlb.operator.operator import Operator +from xlb.cell_type import BC_SOLID + + +class MeshMaskerWinding(MeshBoundaryMasker): + """ + Operator for creating a boundary missing_mask from a mesh file + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call super + super().__init__(velocity_set, precision_policy, compute_backend) + assert self.compute_backend != ComputeBackend.NEON, ( + 'MeshVoxelizationMethod("WINDING") is not implemented in Neon yet! Please use a different method of mesh voxelization!' + ) + + def _construct_warp(self): + # Make constants for warp + _c = self.velocity_set.c + _q = self.velocity_set.q + _opp_indices = self.velocity_set.opp_indices + + @wp.kernel + def kernel( + mesh_id: wp.uint64, + id_number: wp.int32, + distances: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + needs_mesh_distance: bool, + ): + # get index + i, j, k = wp.tid() + + # Get local indices + index = wp.vec3i(i, j, k) + + # position of the point + pos_cell = self.helper_masker.index_to_position(bc_mask, index) + + # Compute the maximum length + max_length = wp.sqrt( + (wp.float32(bc_mask.shape[1])) ** 2.0 + (wp.float32(bc_mask.shape[2])) ** 2.0 + (wp.float32(bc_mask.shape[3])) ** 2.0 + ) + + # evaluate if point is inside mesh + query = wp.mesh_query_point_sign_winding_number(mesh_id, pos_cell, max_length) + if query.result: + # set point to be solid + if query.sign <= 0: # TODO: fix this + # Make solid voxel + bc_mask[0, index[0], index[1], index[2]] = wp.uint8(BC_SOLID) + + # Find the fractional distance to the mesh in each direction + for direction_idx in range(1, _q): + direction_vec = wp.vec3f(wp.float32(_c[0, direction_idx]), wp.float32(_c[1, direction_idx]), wp.float32(_c[2, direction_idx])) + # Max length depends on ray direction (diagonals are longer) + max_length = wp.length(direction_vec) + query_dir = wp.mesh_query_ray(mesh_id, pos_cell, direction_vec / max_length, max_length) + if query_dir.result: + # Get the index of the streaming direction + push_index = wp.vec3i() + for d in range(self.velocity_set.d): + push_index[d] = index[d] + _c[d, direction_idx] + + # Set the boundary id and missing_mask + bc_mask[0, push_index[0], push_index[1], push_index[2]] = wp.uint8(id_number) + missing_mask[direction_idx, push_index[0], push_index[1], push_index[2]] = wp.uint8(True) + + # If we don't need the mesh distance, we can return early + if not needs_mesh_distance: + continue + + # get position of the mesh triangle that intersects with the ray + pos_mesh = wp.mesh_eval_position(mesh_id, query_dir.face, query_dir.u, query_dir.v) + cell_center_pos = self.helper_masker.index_to_position(bc_mask, push_index) + dist = wp.length(pos_mesh - cell_center_pos) + weight = self.store_dtype(dist / max_length) + distances[_opp_indices[direction_idx], push_index[0], push_index[1], push_index[2]] = weight + + return None, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation( + self, + bc, + distances, + bc_mask, + missing_mask, + ): + return self.warp_implementation_base( + bc, + distances, + bc_mask, + missing_mask, + ) diff --git a/xlb/operator/collision/__init__.py b/xlb/operator/collision/__init__.py new file mode 100644 index 00000000..7e09facb --- /dev/null +++ b/xlb/operator/collision/__init__.py @@ -0,0 +1,5 @@ +from xlb.operator.collision.collision import Collision +from xlb.operator.collision.bgk import BGK +from xlb.operator.collision.kbc import KBC +from xlb.operator.collision.smagorinsky_les_bgk import SmagorinskyLESBGK +from xlb.operator.collision.forced_collision import ForcedCollision diff --git a/xlb/operator/collision/bgk.py b/xlb/operator/collision/bgk.py new file mode 100644 index 00000000..29331dc7 --- /dev/null +++ b/xlb/operator/collision/bgk.py @@ -0,0 +1,91 @@ +""" +Bhatnagar-Gross-Krook (BGK) single-relaxation-time collision operator. +""" + +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.collision.collision import Collision +from xlb.operator import Operator +from functools import partial + + +class BGK(Collision): + """Single-relaxation-time BGK collision operator. + + Relaxes the distribution function toward equilibrium at a rate + controlled by the relaxation parameter *omega*:: + + f_out = f - omega * (f - f_eq) + + Supports JAX, Warp, and Neon backends. + """ + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,)) + def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, omega): + fneq = f - feq + fout = f - self.compute_dtype(omega) * fneq + return fout + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _w = self.velocity_set.w + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + # Construct the functional + @wp.func + def functional(f: Any, feq: Any, omega: Any): + fneq = f - feq + fout = f - self.compute_dtype(omega) * fneq + return fout + + # Construct the warp kernel + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + feq: wp.array4d(dtype=Any), + fout: wp.array4d(dtype=Any), + omega: Any, + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) # TODO: Warp needs to fix this + + # Load needed values + _f = _f_vec() + _feq = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _feq[l] = feq[l, index[0], index[1], index[2]] + + # Compute the collision + _fout = functional(_f, _feq, omega) + + # Write the result + for l in range(self.velocity_set.q): + fout[l, index[0], index[1], index[2]] = self.store_dtype(_fout[l]) + + return functional, kernel + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, feq, fout, omega): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + f, + feq, + fout, + omega, + ], + dim=f.shape[1:], + ) + return fout diff --git a/xlb/operator/collision/collision.py b/xlb/operator/collision/collision.py new file mode 100644 index 00000000..c5dffe75 --- /dev/null +++ b/xlb/operator/collision/collision.py @@ -0,0 +1,22 @@ +""" +Base class for Collision operators +""" + +from xlb.velocity_set import VelocitySet +from xlb.operator import Operator + + +class Collision(Operator): + """ + Base class for collision operators. + + This class defines the collision step for the Lattice Boltzmann Method. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + ): + super().__init__(velocity_set, precision_policy, compute_backend) diff --git a/xlb/operator/collision/forced_collision.py b/xlb/operator/collision/forced_collision.py new file mode 100644 index 00000000..2c2cfa07 --- /dev/null +++ b/xlb/operator/collision/forced_collision.py @@ -0,0 +1,125 @@ +""" +Collision operator with external body-force correction. +""" + +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.collision.collision import Collision +from xlb.operator.macroscopic import Macroscopic +from xlb.operator import Operator +from functools import partial +from xlb.operator.force import ExactDifference + + +class ForcedCollision(Collision): + """Collision operator that wraps another collision with a forcing term. + + After the inner collision the forcing operator is applied to + incorporate the effect of an external body force. + + Parameters + ---------- + collision_operator : Operator + The base collision operator (e.g. :class:`BGK`). + forcing_scheme : str + Forcing scheme. Currently only ``"exact_difference"`` is supported. + force_vector : array-like + External force vector of length ``d`` (number of spatial dimensions). + """ + + def __init__( + self, + collision_operator: Operator, + forcing_scheme="exact_difference", + force_vector=None, + ): + assert collision_operator is not None + self.collision_operator = collision_operator + self.macroscopic = Macroscopic() + super().__init__() + + assert forcing_scheme == "exact_difference", NotImplementedError(f"Force model {forcing_scheme} not implemented!") + assert force_vector.shape[0] == self.velocity_set.d, "Check the dimensions of the input force!" + self.force_vector = force_vector + if forcing_scheme == "exact_difference": + self.forcing_operator = ExactDifference(force_vector) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,)) + def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, omega): + fout = self.collision_operator(f, feq, omega) + rho, u = self.macroscopic(fout) + fout = self.forcing_operator(fout, feq, rho, u) + return fout + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + # Construct the functional + @wp.func + def functional(f: Any, feq: Any, omega: Any): + fout = self.collision_operator.warp_functional(f, feq, omega) + rho, u = self.macroscopic.warp_functional(fout) + fout = self.forcing_operator.warp_functional(fout, feq, rho, u) + return fout + + # Construct the warp kernel + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + feq: wp.array4d(dtype=Any), + fout: wp.array4d(dtype=Any), + omega: Any, + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) # TODO: Warp needs to fix this + + # Load needed values + _f = _f_vec() + _feq = _f_vec() + _d = self.velocity_set.d + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _feq[l] = feq[l, index[0], index[1], index[2]] + + # Compute the collision + _fout = functional(_f, _feq, omega) + + # Write the result + for l in range(self.velocity_set.q): + fout[l, index[0], index[1], index[2]] = _fout[l] + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, feq, fout, omega): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + f, + feq, + fout, + omega, + ], + dim=f.shape[1:], + ) + return fout + + def _construct_neon(self): + # Construct the functional + @wp.func + def functional(f: Any, feq: Any, omega: Any): + fout = self.collision_operator.neon_functional(f, feq, omega) + rho, u = self.macroscopic.neon_functional(fout) + fout = self.forcing_operator.neon_functional(fout, feq, rho, u) + return fout + + return functional, None diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py new file mode 100644 index 00000000..d814a7cf --- /dev/null +++ b/xlb/operator/collision/kbc.py @@ -0,0 +1,339 @@ +""" +KBC collision operator for LBM. +""" + +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any +from functools import partial + +from xlb.velocity_set import VelocitySet, D2Q9, D3Q27 +from xlb.compute_backend import ComputeBackend +from xlb.operator.collision.collision import Collision +from xlb.operator import Operator +from xlb.operator.macroscopic import SecondMoment as MomentumFlux + + +class KBC(Collision): + """ + KBC collision operator for LBM. + + This class implements the Karlin-BΓΆsch-Chikatamarla (KBC) model for the collision step in the Lattice Boltzmann Method. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + ): + self.momentum_flux = MomentumFlux() + self.epsilon = 1e-32 + + super().__init__( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,), donate_argnums=(1, 2, 3)) + def jax_implementation( + self, + f: jnp.ndarray, + feq: jnp.ndarray, + omega, + ): + """ + KBC collision step for lattice. + + Parameters + ---------- + f : jax.numpy.array + Distribution function. + feq : jax.numpy.array + Equilibrium distribution function. + """ + fneq = f - feq + if isinstance(self.velocity_set, D2Q9): + shear = self.decompose_shear_d2q9_jax(fneq) + delta_s = shear / 4.0 + elif isinstance(self.velocity_set, D3Q27): + shear = self.decompose_shear_d3q27_jax(fneq) + delta_s = shear + else: + raise NotImplementedError("Velocity set not supported: {}".format(type(self.velocity_set))) + + # Compute required constants based on the input omega (omega is the inverse relaxation time) + beta = self.compute_dtype(0.5) * self.compute_dtype(omega) + inv_beta = 1.0 / beta + + # Perform collision + delta_h = fneq - delta_s + sp1, sp2 = self.compute_entropic_scalar_products(delta_s, delta_h, feq) + gamma = inv_beta - (2.0 - inv_beta) * sp1 / (self.epsilon + sp2) + + fout = f - beta * (2.0 * delta_s + gamma[None, ...] * delta_h) + + return fout + + @partial(jit, static_argnums=(0,), inline=True) + def compute_entropic_scalar_products(self, delta_s: jnp.ndarray, delta_h: jnp.ndarray, feq: jnp.ndarray): + """ + Compute the entropic scalar products to approximate gamma in KBC. + + Returns + ------- + jax.numpy.array + sp1 and sp2: Entropic scalar products of delta_s, delta_h, and feq. + """ + temp = delta_h / feq + sp1 = jnp.sum(temp * delta_s, axis=0) + sp2 = jnp.sum(temp * delta_h, axis=0) + return sp1, sp2 + + @partial(jit, static_argnums=(0,), inline=True) + def decompose_shear_d3q27_jax(self, fneq): + """ + Decompose fneq into shear components for D3Q27 lattice. + + Parameters + ---------- + fneq : jax.numpy.ndarray + Non-equilibrium distribution function. + + Returns + ------- + jax.numpy.ndarray + Shear components of fneq. + """ + + # Calculate the momentum flux + Pi = self.momentum_flux(fneq) + # Calculating Nxz and Nyz with indices moved to the first dimension + Nxz = Pi[0, ...] - Pi[5, ...] + Nyz = Pi[3, ...] - Pi[5, ...] + + # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k) + s = jnp.zeros_like(fneq) + s = s.at[9, ...].set((2.0 * Nxz - Nyz) / 6.0) + s = s.at[18, ...].set((2.0 * Nxz - Nyz) / 6.0) + s = s.at[3, ...].set((-Nxz + 2.0 * Nyz) / 6.0) + s = s.at[6, ...].set((-Nxz + 2.0 * Nyz) / 6.0) + s = s.at[1, ...].set((-Nxz - Nyz) / 6.0) + s = s.at[2, ...].set((-Nxz - Nyz) / 6.0) + + # For c = (i, j, 0) + s = s.at[12, ...].set(Pi[1, ...] / 4.0) + s = s.at[24, ...].set(Pi[1, ...] / 4.0) + s = s.at[21, ...].set(-Pi[1, ...] / 4.0) + s = s.at[15, ...].set(-Pi[1, ...] / 4.0) + + # For c = (i, 0, k) + s = s.at[10, ...].set(Pi[2, ...] / 4.0) + s = s.at[20, ...].set(Pi[2, ...] / 4.0) + s = s.at[19, ...].set(-Pi[2, ...] / 4.0) + s = s.at[11, ...].set(-Pi[2, ...] / 4.0) + + # For c = (0, j, k) + s = s.at[8, ...].set(Pi[4, ...] / 4.0) + s = s.at[4, ...].set(Pi[4, ...] / 4.0) + s = s.at[7, ...].set(-Pi[4, ...] / 4.0) + s = s.at[5, ...].set(-Pi[4, ...] / 4.0) + + return s + + @partial(jit, static_argnums=(0,), inline=True) + def decompose_shear_d2q9_jax(self, fneq): + """ + Decompose fneq into shear components for D2Q9 lattice. + + Parameters + ---------- + fneq : jax.numpy.array + Non-equilibrium distribution function. + + Returns + ------- + jax.numpy.array + Shear components of fneq. + """ + Pi = self.momentum_flux(fneq) + N = Pi[0, ...] - Pi[2, ...] + s = jnp.zeros_like(fneq) + s = s.at[3, ...].set(N) + s = s.at[6, ...].set(N) + s = s.at[2, ...].set(-N) + s = s.at[1, ...].set(-N) + s = s.at[8, ...].set(Pi[1, ...]) + s = s.at[4, ...].set(-Pi[1, ...]) + s = s.at[5, ...].set(-Pi[1, ...]) + s = s.at[7, ...].set(Pi[1, ...]) + + return s + + def _construct_warp(self): + # Raise error if velocity set is not supported + if not (isinstance(self.velocity_set, D3Q27) or isinstance(self.velocity_set, D2Q9)): + raise NotImplementedError("Velocity set not supported for warp backend: {}".format(type(self.velocity_set))) + + # Set local constants TODO: This is a hack and should be fixed with warp update + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _epsilon = wp.constant(self.compute_dtype(self.epsilon)) + + @wp.func + def decompose_shear_d2q9(fneq: Any): + pi = self.momentum_flux.warp_functional(fneq) + N = pi[0] - pi[2] + s = _f_vec() + s[3] = N + s[6] = N + s[2] = -N + s[1] = -N + s[8] = pi[1] + s[4] = -pi[1] + s[5] = -pi[1] + s[7] = pi[1] + return s + + # Construct functional for decomposing shear + @wp.func + def decompose_shear_d3q27( + fneq: Any, + ): + # Get momentum flux + pi = self.momentum_flux.warp_functional(fneq) + nxz = pi[0] - pi[5] + nyz = pi[3] - pi[5] + + # set shear components + s = _f_vec() + + # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k) + two = self.compute_dtype(2.0) + four = self.compute_dtype(4.0) + six = self.compute_dtype(6.0) + + s[9] = (two * nxz - nyz) / six + s[18] = (two * nxz - nyz) / six + s[3] = (-nxz + two * nyz) / six + s[6] = (-nxz + two * nyz) / six + s[1] = (-nxz - nyz) / six + s[2] = (-nxz - nyz) / six + + # For c = (i, j, 0) + s[12] = pi[1] / four + s[24] = pi[1] / four + s[21] = -pi[1] / four + s[15] = -pi[1] / four + + # For c = (i, 0, k) + s[10] = pi[2] / four + s[20] = pi[2] / four + s[19] = -pi[2] / four + s[11] = -pi[2] / four + + # For c = (0, j, k) + s[8] = pi[4] / four + s[4] = pi[4] / four + s[7] = -pi[4] / four + s[5] = -pi[4] / four + + return s + + # Construct functional for computing entropic scalar products + @wp.func + def compute_entropic_scalar_products( + delta_s: Any, + delta_h: Any, + feq: Any, + ): + temp = wp.cw_div(delta_h, feq) + sp1 = self.compute_dtype(0.0) + sp2 = self.compute_dtype(0.0) + for i in range(self.velocity_set.q): + sp1 += temp[i] * delta_s[i] + sp2 += temp[i] * delta_h[i] + return sp1, sp2 + + # Construct the functional + @wp.func + def functional( + f: Any, + feq: Any, + omega: Any, + ): + # Compute shear and delta_s + fneq = f - feq + if wp.static(self.velocity_set.d == 3): + shear = decompose_shear_d3q27(fneq) + delta_s = shear + else: + shear = decompose_shear_d2q9(fneq) + delta_s = shear / self.compute_dtype(4.0) + + # Compute required constants based on the input omega (omega is the inverse relaxation time) + _beta = self.compute_dtype(0.5) * self.compute_dtype(omega) + _inv_beta = self.compute_dtype(1.0) / _beta + + # Perform collision + delta_h = fneq - delta_s + two = self.compute_dtype(2.0) + sp1, sp2 = compute_entropic_scalar_products(delta_s, delta_h, feq) + gamma = _inv_beta - (two - _inv_beta) * sp1 / (_epsilon + sp2) + fout = f - _beta * (two * delta_s + gamma * delta_h) + + return fout + + # Construct the warp kernel + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + feq: wp.array4d(dtype=Any), + fout: wp.array4d(dtype=Any), + omega: Any, + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) # TODO: Warp needs to fix this + + # Load needed values + _f = _f_vec() + _feq = _f_vec() + _d = self.velocity_set.d + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _feq[l] = feq[l, index[0], index[1], index[2]] + + # Compute the collision + _fout = functional(_f, _feq, omega) + + # Write the result + for l in range(self.velocity_set.q): + fout[l, index[0], index[1], index[2]] = self.store_dtype(_fout[l]) + + return functional, kernel + + def _construct_neon(self): + # Redefine the momentum flux operator for the neon backend + # This is because the neon backend relies on the warp functionals for its operations. + self.momentum_flux = MomentumFlux(compute_backend=ComputeBackend.WARP) + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, feq, fout, omega): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + f, + feq, + fout, + omega, + ], + dim=f.shape[1:], + ) + return fout diff --git a/xlb/operator/collision/smagorinsky_les_bgk.py b/xlb/operator/collision/smagorinsky_les_bgk.py new file mode 100644 index 00000000..7921105e --- /dev/null +++ b/xlb/operator/collision/smagorinsky_les_bgk.py @@ -0,0 +1,158 @@ +""" +BGK collision operator with Smagorinsky large-eddy-simulation sub-grid model. +""" + +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any +import numpy as np + +from xlb.velocity_set import VelocitySet +from xlb.compute_backend import ComputeBackend +from xlb.operator.collision.collision import Collision +from xlb.operator import Operator +from functools import partial + + +class SmagorinskyLESBGK(Collision): + """BGK collision with Smagorinsky LES turbulence modelling. + + Adjusts the effective relaxation time based on the local strain rate + estimated from the non-equilibrium stress tensor, using the + Smagorinsky model constant *C_s*. + + Parameters + ---------- + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + smagorinsky_coef : float + Smagorinsky model constant (default 0.17). + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + smagorinsky_coef: float = 0.17, + ): + self.smagorinsky_coef = smagorinsky_coef + super().__init__(velocity_set, precision_policy, compute_backend) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,)) + def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, omega): + fneq = f - feq + + pi_neq = jnp.tensordot(self.velocity_set.cc, fneq, axes=(0, 0)) + + if self.velocity_set.d == 3: + diag = pi_neq[(0, 3, 5), ...] + offdiag = pi_neq[(1, 2, 4), ...] + else: + diag = pi_neq[(0, 2), ...] + offdiag = pi_neq[(1,), ...] + + strain = jnp.sum(diag * diag, axis=0) + self.compute_dtype(2.0) * jnp.sum(offdiag * offdiag, axis=0) + + tau0 = self.compute_dtype(1.0) / self.compute_dtype(omega) + cs = self.compute_dtype(self.smagorinsky_coef) + tau = self.compute_dtype(0.5) * (tau0 + jnp.sqrt(tau0 * tau0 + self.compute_dtype(36.0) * (cs * cs) * jnp.sqrt(strain))) + + omega_eff = self.compute_dtype(1.0) / tau + fout = f - omega_eff[None, ...] * fneq + return fout + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _d = self.velocity_set.d + _cc = self.velocity_set.cc + _smagorinsky_coef = wp.constant(self.compute_dtype(self.smagorinsky_coef)) + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _pi_dim = self.velocity_set.d * (self.velocity_set.d + 1) // 2 + _pi_vec = wp.vec(_pi_dim, dtype=self.compute_dtype) + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + + # Construct the functional + @wp.func + def functional( + f: Any, + feq: Any, + omega: Any, + ): + # Compute the non-equilibrium distribution + fneq = f - feq + + # Compute strain + pi_neq = _pi_vec() + for a in range(_pi_dim): + pi_neq[a] = self.compute_dtype(0.0) + for l in range(self.velocity_set.q): + pi_neq[a] += _cc[l, a] * fneq[l] + + strain = self.compute_dtype(0.0) + if wp.static(_d == 3): + strain += pi_neq[0] * pi_neq[0] + pi_neq[3] * pi_neq[3] + pi_neq[5] * pi_neq[5] + strain += self.compute_dtype(2.0) * (pi_neq[1] * pi_neq[1] + pi_neq[2] * pi_neq[2] + pi_neq[4] * pi_neq[4]) + else: + strain += pi_neq[0] * pi_neq[0] + pi_neq[2] * pi_neq[2] + strain += self.compute_dtype(2.0) * (pi_neq[1] * pi_neq[1]) + + # Compute the Smagorinsky model + _tau = self.compute_dtype(1.0) / self.compute_dtype(omega) + tau = _tau + ( + self.compute_dtype(0.5) * (wp.sqrt(_tau * _tau + self.compute_dtype(36.0) * (_smagorinsky_coef**2.0) * wp.sqrt(strain)) - _tau) + ) + + # Compute the collision + fout = f - (self.compute_dtype(1.0) / tau) * fneq + return fout + + # Construct the warp kernel + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + feq: wp.array4d(dtype=Any), + fout: wp.array4d(dtype=Any), + omega: wp.float32, + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) # TODO: Warp needs to fix this + + # Load needed values + _f = _f_vec() + _feq = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _feq[l] = feq[l, index[0], index[1], index[2]] + + # Compute the collision + _fout = functional(_f, _feq, omega) + + # Write the result + for l in range(self.velocity_set.q): + fout[l, index[0], index[1], index[2]] = self.store_dtype(_fout[l]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, feq, fout, omega): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + f, + feq, + fout, + omega, + ], + dim=f.shape[1:], + ) + return fout + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None diff --git a/xlb/operator/equilibrium/__init__.py b/xlb/operator/equilibrium/__init__.py new file mode 100644 index 00000000..beb7bb5e --- /dev/null +++ b/xlb/operator/equilibrium/__init__.py @@ -0,0 +1,3 @@ +from xlb.operator.equilibrium.equilibrium import Equilibrium +from xlb.operator.equilibrium.quadratic_equilibrium import QuadraticEquilibrium +from xlb.operator.equilibrium.multires_quadratic_equilibrium import MultiresQuadraticEquilibrium diff --git a/xlb/operator/equilibrium/equilibrium.py b/xlb/operator/equilibrium/equilibrium.py new file mode 100644 index 00000000..11f4155a --- /dev/null +++ b/xlb/operator/equilibrium/equilibrium.py @@ -0,0 +1,17 @@ +# Base class for all equilibriums +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.operator.operator import Operator + + +class Equilibrium(Operator): + """ + Base class for all equilibriums + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy=None, + compute_backend=None, + ): + super().__init__(velocity_set, precision_policy, compute_backend) diff --git a/xlb/operator/equilibrium/multires_quadratic_equilibrium.py b/xlb/operator/equilibrium/multires_quadratic_equilibrium.py new file mode 100644 index 00000000..58ed96cc --- /dev/null +++ b/xlb/operator/equilibrium/multires_quadratic_equilibrium.py @@ -0,0 +1,77 @@ +""" +Multi-resolution quadratic equilibrium operator for the Neon backend. +""" + +import warp as wp +from typing import Any +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator import Operator + + +class MultiresQuadraticEquilibrium(QuadraticEquilibrium): + """Quadratic equilibrium operator for multi-resolution grids (Neon only). + + Computes the second-order Hermite-polynomial equilibrium distribution + from density and velocity at every active cell on each grid level. + Cells that have child refinement (halo cells) are zeroed out. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + # Set local constants TODO: This is a hack and should be fixed with warp update + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + + @neon.Container.factory(name="QuadraticEquilibrium") + def container( + rho: Any, + u: Any, + f: Any, + level, + ): + def quadratic_equilibrium_ll(loader: neon.Loader): + loader.set_mres_grid(rho.get_grid(), level) + + rho_pn = loader.get_mres_read_handle(rho) + u_pn = loader.get_mres_read_handle(u) + f_pn = loader.get_mres_write_handle(f) + + @wp.func + def quadratic_equilibrium_cl(index: Any): + _u = _u_vec() + for d in range(self.velocity_set.d): + _u[d] = self.compute_dtype(wp.neon_read(u_pn, index, d)) + _rho = self.compute_dtype(wp.neon_read(rho_pn, index, 0)) + feq = functional(_rho, _u) + + if wp.neon_has_child(f_pn, index): + for l in range(self.velocity_set.q): + feq[l] = self.compute_dtype(0.0) + # Set the output + for l in range(self.velocity_set.q): + wp.neon_write(f_pn, index, l, self.store_dtype(feq[l])) + + loader.declare_kernel(quadratic_equilibrium_cl) + + return quadratic_equilibrium_ll + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, rho, u, f, stream=0): + import neon + + grid = f.get_grid() + for level in range(grid.num_levels): + c = self.neon_container(rho, u, f, level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return f diff --git a/xlb/operator/equilibrium/quadratic_equilibrium.py b/xlb/operator/equilibrium/quadratic_equilibrium.py new file mode 100644 index 00000000..30ae6dee --- /dev/null +++ b/xlb/operator/equilibrium/quadratic_equilibrium.py @@ -0,0 +1,150 @@ +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +import os + +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.equilibrium import Equilibrium +from xlb.operator import Operator + + +class QuadraticEquilibrium(Equilibrium): + """ + Quadratic equilibrium of Boltzmann equation using hermite polynomials. + Standard equilibrium model for LBM. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, rho, u): + cu = 3.0 * jnp.tensordot(self.velocity_set.c, u, axes=(0, 0)) + usqr = 1.5 * jnp.sum(jnp.square(u), axis=0, keepdims=True) + w = self.velocity_set.w.reshape((-1,) + (1,) * (len(u.shape) - 1)) + feq = rho * w * (1.0 + cu * (1.0 + 0.5 * cu) - usqr) + return feq + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _c = self.velocity_set.c + _w = self.velocity_set.w + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + + # Construct the equilibrium functional + @wp.func + def functional( + rho: Any, + u: Any, + ): + # Allocate the equilibrium + feq = _f_vec() + + # Compute the equilibrium + for l in range(self.velocity_set.q): + # Compute cu + cu = self.compute_dtype(0.0) + for d in range(self.velocity_set.d): + if _c[d, l] == 1: + cu += u[d] + elif _c[d, l] == -1: + cu -= u[d] + cu *= self.compute_dtype(3.0) + + # Compute usqr + usqr = self.compute_dtype(1.5) * wp.dot(u, u) + + # Compute feq + feq[l] = rho * _w[l] * (self.compute_dtype(1.0) + cu * (self.compute_dtype(1.0) + self.compute_dtype(0.5) * cu) - usqr) + + return feq + + # Construct the warp kernel + @wp.kernel + def kernel( + rho: wp.array4d(dtype=Any), + u: wp.array4d(dtype=Any), + f: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Get the equilibrium + _u = _u_vec() + for d in range(self.velocity_set.d): + _u[d] = u[d, index[0], index[1], index[2]] + _rho = rho[0, index[0], index[1], index[2]] + feq = functional(_rho, _u) + + # Set the output + for l in range(self.velocity_set.q): + f[l, index[0], index[1], index[2]] = self.store_dtype(feq[l]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, rho, u, f): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + rho, + u, + f, + ], + dim=rho.shape[1:], + ) + return f + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + # Set local constants TODO: This is a hack and should be fixed with warp update + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + + @neon.Container.factory(name="QuadraticEquilibrium") + def container( + rho: Any, + u: Any, + f: Any, + ): + def quadratic_equilibrium_ll(loader: neon.Loader): + loader.set_grid(rho.get_grid()) + rho_pn = loader.get_read_handle(rho) + u_pn = loader.get_read_handle(u) + f_pn = loader.get_write_handle(f) + + @wp.func + def quadratic_equilibrium_cl(index: Any): + _u = _u_vec() + for d in range(self.velocity_set.d): + _u[d] = self.compute_dtype(wp.neon_read(u_pn, index, d)) + _rho = self.compute_dtype(wp.neon_read(rho_pn, index, 0)) + feq = functional(_rho, _u) + + # Set the output + for l in range(self.velocity_set.q): + wp.neon_write(f_pn, index, l, self.store_dtype(feq[l])) + + loader.declare_kernel(quadratic_equilibrium_cl) + + return quadratic_equilibrium_ll + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, rho, u, f): + import neon + + c = self.neon_container(rho, u, f) + c.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + return f diff --git a/xlb/operator/force/__init__.py b/xlb/operator/force/__init__.py new file mode 100644 index 00000000..f3ceec57 --- /dev/null +++ b/xlb/operator/force/__init__.py @@ -0,0 +1,3 @@ +from xlb.operator.force.momentum_transfer import MomentumTransfer +from xlb.operator.force.exact_difference_force import ExactDifference +from xlb.operator.force.multires_momentum_transfer import MultiresMomentumTransfer diff --git a/xlb/operator/force/exact_difference_force.py b/xlb/operator/force/exact_difference_force.py new file mode 100644 index 00000000..760a9ef1 --- /dev/null +++ b/xlb/operator/force/exact_difference_force.py @@ -0,0 +1,135 @@ +from functools import partial +from jax import jit, lax +import warp as wp +from typing import Any + +from xlb import DefaultConfig +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.equilibrium import QuadraticEquilibrium + + +class ExactDifference(Operator): + """ + Add external body force based on the exact-difference method due to Kupershtokh (2004) + + References + ---------- + Kupershtokh, A. (2004). New method of incorporating a body force term into the lattice Boltzmann equation. In + Proceedings of the 5th International EHD Workshop (pp. 241-246). University of Poitiers, Poitiers, France. + Chikatamarla, S. S., & Karlin, I. V. (2013). Entropic lattice Boltzmann method for turbulent flow simulations: + Boundary conditions. Physica A, 392, 1925-1930. + KrΓΌger, T., et al. (2017). The lattice Boltzmann method. Springer International Publishing, 10.978-3, 4-15. + """ + + def __init__( + self, + force_vector, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # TODO: currently we are limited to a single force vector not a spatially dependent forcing field + self.force_vector = force_vector + + # Resolve compute_backend the same way Operator.__init__ does, so we + # know which equilibrium backend to use before super().__init__ runs. + # Neon kernels reuse Warp functionals, so sub-operators on the Neon + # backend are built on Warp. + resolved_backend = compute_backend or DefaultConfig.default_backend + eq_backend = ComputeBackend.WARP if resolved_backend == ComputeBackend.NEON else resolved_backend + self.equilibrium = QuadraticEquilibrium(compute_backend=eq_backend) + + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_postcollision, feq, rho, u): + """ + Parameters + ---------- + f_postcollision: jax.numpy.ndarray + The post-collision distribution functions. + feq: jax.numpy.ndarray + The equilibrium distribution functions. + rho: jax.numpy.ndarray + The density field. + + u: jax.numpy.ndarray + The velocity field. + + Returns + ------- + f_postcollision: jax.numpy.ndarray + The post-collision distribution functions with the force applied. + """ + delta_u = lax.broadcast_in_dim(self.force_vector, u.shape, (0,)) + feq_force = self.equilibrium(rho, u + delta_u) + f_postcollision += feq_force - feq + return f_postcollision + + def _construct_warp(self): + _d = self.velocity_set.d + _u_vec = wp.vec(_d, dtype=self.compute_dtype) + if _d == 2: + _force = _u_vec(self.force_vector[0], self.force_vector[1]) + else: + _force = _u_vec(self.force_vector[0], self.force_vector[1], self.force_vector[2]) + + # Construct the functional + @wp.func + def functional(f_postcollision: Any, feq: Any, rho: Any, u: Any): + delta_u = _force + feq_force = self.equilibrium.warp_functional(rho, u + delta_u) + f_postcollision += feq_force - feq + return f_postcollision + + # Construct the warp kernel + @wp.kernel + def kernel( + f_postcollision: Any, + feq: Any, + fout: wp.array4d(dtype=Any), + rho: wp.array4d(dtype=Any), + u: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) # TODO: Warp needs to fix this + + # Load needed values + _u = _u_vec() + for l in range(_d): + _u[l] = u[l, index[0], index[1], index[2]] + _rho = rho[0, index[0], index[1], index[2]] + + # Compute the collision + _fout = functional(f_postcollision, feq, _rho, _u) + + # Write the result + for l in range(self.velocity_set.q): + fout[l, index[0], index[1], index[2]] = self.store_dtype(_fout[l]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_postcollision, feq, fout, rho, u): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_postcollision, feq, fout, rho, u], + dim=f_postcollision.shape[1:], + ) + return fout + + def _construct_neon(self): + # The neon backend relies on the warp functionals for its operations. + functional, _ = self._construct_warp() + return functional, None diff --git a/xlb/operator/force/momentum_transfer.py b/xlb/operator/force/momentum_transfer.py new file mode 100644 index 00000000..9032abb1 --- /dev/null +++ b/xlb/operator/force/momentum_transfer.py @@ -0,0 +1,364 @@ +from functools import partial +import jax.numpy as jnp +from jax import jit, lax +import warp as wp +from typing import Any +from enum import Enum, auto + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.stream import Stream + + +# Enum used to keep track of LBM operations +class LBMOperationSequence(Enum): + """ + Note that for dense and single resolution simulations in XLB, the order of operations in the stepper is "stream-then-collide". + For MultiRes stepper however the order of operations is always "collide-then-stream" except at the finest level when the FUSION_AT_FINEST + optimization is used. + In that case the order of operations is "stream-then-collide" ONLY at the finest level. + """ + + STREAM_THEN_COLLIDE = auto() + COLLIDE_THEN_STREAM = auto() + + +class FetchPopulations(Operator): + """ + This operator is used to get the post-collision and post-streaming populations + Note that for dense and single resolution simulations in XLB, the order of operations in the stepper is "stream-then-collide". + Therefore, f_0 represents the post-collision values and post_streaming values of the current time step need to be reconstructed + by applying the streaming and boundary conditions. These populations are readily available in XLB when using multi-resolution + grids because the mres stepper relies on "collide-then-stream". + """ + + def __init__( + self, + no_slip_bc_instance, + operation_sequence: LBMOperationSequence = LBMOperationSequence.STREAM_THEN_COLLIDE, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + self.no_slip_bc_instance = no_slip_bc_instance + self.stream = Stream(velocity_set, precision_policy, compute_backend) + self.operation_sequence = operation_sequence + + if compute_backend == ComputeBackend.WARP: + self.stream_functional = self.stream.warp_functional + self.bc_functional = self.no_slip_bc_instance.warp_functional + elif compute_backend == ComputeBackend.NEON: + self.stream_functional = self.stream.neon_functional + self.bc_functional = self.no_slip_bc_instance.neon_functional + + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_0, f_1, bc_mask, missing_mask): + # Give the input post-collision populations, streaming once and apply the BC the find post-stream values. + f_post_collision = f_0 + f_post_stream = self.stream(f_post_collision) + f_post_stream = self.no_slip_bc_instance(f_post_collision, f_post_stream, bc_mask, missing_mask) + return f_post_collision, f_post_stream + + def _construct_warp(self): + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + @wp.func + def functional_stream_then_collide( + index: Any, + f_0: Any, + f_1: Any, + _missing_mask: Any, + ): + # Get the distribution function + f_post_collision = _f_vec() + for l in range(self.velocity_set.q): + f_post_collision[l] = self.compute_dtype(self.read_field(f_0, index, l)) + + # Apply streaming (pull method) + timestep = 0 + f_post_stream = self.stream_functional(f_0, index) + f_post_stream = self.bc_functional(index, timestep, _missing_mask, f_0, f_1, f_post_collision, f_post_stream) + return f_post_collision, f_post_stream + + @wp.func + def functional_collide_then_stream( + index: Any, + f_0: Any, + f_1: Any, + _missing_mask: Any, + ): + # Get the distribution function + f_post_collision = _f_vec() + f_post_stream = _f_vec() + for l in range(self.velocity_set.q): + f_post_stream[l] = self.compute_dtype(self.read_field(f_0, index, l)) + f_post_collision[l] = self.compute_dtype(self.read_field(f_1, index, l)) + return f_post_collision, f_post_stream + + if self.operation_sequence == LBMOperationSequence.STREAM_THEN_COLLIDE: + return functional_stream_then_collide, None + elif self.operation_sequence == LBMOperationSequence.COLLIDE_THEN_STREAM: + return functional_collide_then_stream, None + else: + raise ValueError(f"Unknown operation sequence: {self.operation_sequence}") + + def _construct_neon(self): + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + return functional, None + + +class MomentumTransfer(Operator): + """ + An opertor for the momentum exchange method to compute the boundary force vector exerted on the solid geometry + based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a + bounce-back method (e.g. Bouzidi) that accounts for curved boundaries. + NOTE: this function should be called after BC's are imposed. + [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation. + Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339. + [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving + curved geometry, Phys. Rev. E 65 (2002) 041203. + [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange + algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423. + + Notes + ----- + This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. + The force is computed based on the post-streaming and post-collision distribution functions. This method + should be called after the boundary conditions are imposed. + """ + + def __init__( + self, + no_slip_bc_instance, + operation_sequence: LBMOperationSequence = LBMOperationSequence.STREAM_THEN_COLLIDE, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Assign the no-slip boundary condition instance + self.no_slip_bc_instance = no_slip_bc_instance + self.operation_sequence = operation_sequence + + # Define the needed for the momentum transfer + self.fetcher = FetchPopulations( + no_slip_bc_instance=self.no_slip_bc_instance, + operation_sequence=self.operation_sequence, + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + if self.compute_backend != ComputeBackend.JAX: + # Allocate the force vector (the total integral value will be computed) + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + self.force = wp.zeros((1), dtype=_u_vec) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_0, f_1, bc_mask, missing_mask): + """ + Parameters + ---------- + f_0 : jax.numpy.ndarray + The post-collision distribution function at each node in the grid. + f_1 : jax.numpy.ndarray + The buffer field the same size as f_0 (only given as input for consistency with the WARP backened API.) + bc_mask : jax.numpy.ndarray + A grid field with 0 everywhere except for boundary nodes which are designated + by their respective boundary id's. + missing_mask : jax.numpy.ndarray + A grid field with lattice cardinality that specifies missing lattice directions + for each boundary node. + + Returns + ------- + jax.numpy.ndarray + The force exerted on the solid geometry at each boundary node. + """ + # Give the input post-collision populations, streaming once and apply the BC the find post-stream values. + f_post_collision, f_post_stream = self.fetcher(f_0, f_1, bc_mask, missing_mask) + + # Compute momentum transfer + boundary = bc_mask == self.no_slip_bc_instance.id + new_shape = (self.velocity_set.q,) + boundary.shape[1:] + boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1))) + + # the following will return force as a grid-based field with zero everywhere except for boundary nodes. + is_edge = jnp.logical_and(boundary, ~missing_mask[0]) + opp = self.velocity_set.opp_indices + phi = f_post_collision[opp] + f_post_stream + phi = jnp.where(jnp.logical_and(missing_mask, is_edge), phi, 0.0) + force = jnp.tensordot(self.velocity_set.c[:, opp], phi, axes=(-1, 0)) + force_net = jnp.sum(force, axis=(i + 1 for i in range(self.velocity_set.d))) + return force_net + + def _construct_warp(self): + # Set local constants + _c = self.velocity_set.c + _opp_indices = self.velocity_set.opp_indices + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8) + _no_slip_id = self.no_slip_bc_instance.id + + # Find velocity index for (0, 0, 0) + lattice_central_index = self.velocity_set.center_index + + @wp.func + def functional( + index: Any, + f_0: Any, + f_1: Any, + bc_mask: Any, + missing_mask: Any, + force: Any, + ): + # Get the boundary id + _boundary_id = self.read_field(bc_mask, index, 0) + _missing_mask = _missing_mask_vec() + for l in range(self.velocity_set.q): + _missing_mask[l] = self.read_field(missing_mask, index, l) + + # Determin if boundary is an edge by checking if center is missing + is_edge = wp.bool(False) + if _boundary_id == wp.uint8(_no_slip_id): + if _missing_mask[lattice_central_index] == wp.uint8(0): + is_edge = wp.bool(True) + + # If the boundary is an edge then add the momentum transfer + m = _u_vec() + if is_edge: + # fetch the post-collision and post-streaming populations + f_post_collision, f_post_stream = self.fetcher_functional(index, f_0, f_1, _missing_mask) + + # Compute the momentum transfer + for d in range(self.velocity_set.d): + m[d] = self.compute_dtype(0.0) + for l in range(self.velocity_set.q): + if _missing_mask[l] == wp.uint8(1): + phi = f_post_collision[_opp_indices[l]] + f_post_stream[l] + if _c[d, _opp_indices[l]] == 1: + m[d] += phi + elif _c[d, _opp_indices[l]] == -1: + m[d] -= phi + # Atomic sum to get the total force vector + wp.atomic_add(force, 0, m) + + # Construct the warp kernel + @wp.kernel + def kernel( + f_0: wp.array4d(dtype=Any), + f_1: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + missing_mask: wp.array4d(dtype=wp.uint8), + force: wp.array(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Call the functional to compute the force + functional( + index, + f_0, + f_1, + bc_mask, + missing_mask, + force, + ) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_0, f_1, bc_mask, missing_mask): + # Ensure the force is initialized to zero + self.force *= self.compute_dtype(0.0) + + # Define the warp functionals needed for this operation + self.fetcher_functional = self.fetcher.warp_functional + + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[f_0, f_1, bc_mask, missing_mask, self.force], + dim=f_0.shape[1:], + ) + return self.force.numpy()[0] + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MomentumTransfer") + def container( + f_0: Any, + f_1: Any, + bc_mask: Any, + missing_mask: Any, + force: Any, + ): + def container_launcher(loader: neon.Loader): + loader.set_grid(bc_mask.get_grid()) + bc_mask_pn = loader.get_write_handle(bc_mask) + missing_mask_pn = loader.get_write_handle(missing_mask) + f_0_pn = loader.get_write_handle(f_0) + f_1_pn = loader.get_write_handle(f_1) + + @wp.func + def container_kernel(index: Any): + # apply the functional + functional( + index, + f_0_pn, + f_1_pn, + bc_mask_pn, + missing_mask_pn, + force, + ) + + loader.declare_kernel(container_kernel) + + return container_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + f_0, + f_1, + bc_mask, + missing_mask, + stream=0, + ): + import neon + + # Ensure the force is initialized to zero + self.force *= self.compute_dtype(0.0) + + # Define the neon functionals needed for this operation + self.fetcher_functional = self.fetcher.neon_functional + + # Launch the neon container + c = self.neon_container(f_0, f_1, bc_mask, missing_mask, self.force) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return self.force.numpy()[0] diff --git a/xlb/operator/force/multires_momentum_transfer.py b/xlb/operator/force/multires_momentum_transfer.py new file mode 100644 index 00000000..2c762494 --- /dev/null +++ b/xlb/operator/force/multires_momentum_transfer.py @@ -0,0 +1,139 @@ +""" +Multi-resolution momentum-transfer force operator for the Neon backend. +""" + +from typing import Any + +import warp as wp + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.force import MomentumTransfer +from xlb.mres_perf_optimization_type import MresPerfOptimizationType + + +class MultiresMomentumTransfer(MomentumTransfer): + """Momentum-transfer force computation on a multi-resolution grid. + + Extends :class:`MomentumTransfer` with Neon-specific container code that + iterates over all grid levels. The LBM operation sequence (collide-then- + stream vs. stream-then-collide) is inferred from the performance + optimization type. + + Parameters + ---------- + no_slip_bc_instance : BoundaryCondition + The no-slip BC whose tagged voxels define the force integration + surface. + mres_perf_opt : MresPerfOptimizationType + Multi-resolution performance strategy. + velocity_set : VelocitySet, optional + precision_policy : PrecisionPolicy, optional + compute_backend : ComputeBackend, optional + """ + + def __init__( + self, + no_slip_bc_instance, + mres_perf_opt=MresPerfOptimizationType.NAIVE_COLLIDE_STREAM, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + from xlb.operator.force.momentum_transfer import LBMOperationSequence + + if compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {compute_backend} backend.") + + # Set the sequence of operations based on the performance optimization type + if mres_perf_opt == MresPerfOptimizationType.NAIVE_COLLIDE_STREAM: + operation_sequence = LBMOperationSequence.COLLIDE_THEN_STREAM + elif mres_perf_opt in ( + MresPerfOptimizationType.FUSION_AT_FINEST, + MresPerfOptimizationType.FUSION_AT_FINEST_SFV, + MresPerfOptimizationType.FUSION_AT_FINEST_SFV_ALL, + ): + operation_sequence = LBMOperationSequence.STREAM_THEN_COLLIDE + else: + raise ValueError(f"Unknown performance optimization type: {mres_perf_opt}") + + # Check if the performance optimization type is compatible with the use of mesh distance + if operation_sequence != LBMOperationSequence.STREAM_THEN_COLLIDE: + assert not no_slip_bc_instance.needs_mesh_distance, ( + "Mesh distance is only supported in the MultiresMomentumTransfer operator when the LBM operation sequence is STREAM_THEN_COLLIDE." + ) + + # Print a warning to the user about the boundary voxels + print( + "WARNING! make sure boundary voxels are all at the same level and not among the transition regions from one level to another. " + "Otherwise, the results of force calculation are not correct!\n" + ) + + # Call super + super().__init__(no_slip_bc_instance, operation_sequence, velocity_set, precision_policy, compute_backend) + + def _construct_neon(self): + import neon + + # Use the warp functional for the NEON backend + functional, _ = self._construct_warp() + + @neon.Container.factory(name="MomentumTransfer") + def container( + f_0: Any, + f_1: Any, + bc_mask: Any, + missing_mask: Any, + force: Any, + level: Any, + ): + def container_launcher(loader: neon.Loader): + loader.set_mres_grid(bc_mask.get_grid(), level) + bc_mask_pn = loader.get_mres_write_handle(bc_mask) + missing_mask_pn = loader.get_mres_write_handle(missing_mask) + f_0_pn = loader.get_mres_write_handle(f_0) + f_1_pn = loader.get_mres_write_handle(f_1) + + @wp.func + def container_kernel(index: Any): + # apply the functional + functional( + index, + f_0_pn, + f_1_pn, + bc_mask_pn, + missing_mask_pn, + force, + ) + + loader.declare_kernel(container_kernel) + + return container_launcher + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation( + self, + f_0, + f_1, + bc_mask, + missing_mask, + stream=0, + ): + import neon + + # Ensure the force is initialized to zero + self.force *= self.compute_dtype(0.0) + + # Define the neon functionals needed for this operation + self.fetcher_functional = self.fetcher.neon_functional + + grid = bc_mask.get_grid() + for level in range(grid.num_levels): + # Launch the neon container + c = self.neon_container(f_0, f_1, bc_mask, missing_mask, self.force, level) + c.run(stream, container_runtime=neon.Container.ContainerRuntime.neon) + return self.force.numpy()[0] diff --git a/xlb/operator/macroscopic/__init__.py b/xlb/operator/macroscopic/__init__.py new file mode 100644 index 00000000..75eacee6 --- /dev/null +++ b/xlb/operator/macroscopic/__init__.py @@ -0,0 +1,5 @@ +from xlb.operator.macroscopic.macroscopic import Macroscopic +from xlb.operator.macroscopic.second_moment import SecondMoment +from xlb.operator.macroscopic.zero_moment import ZeroMoment +from xlb.operator.macroscopic.first_moment import FirstMoment +from xlb.operator.macroscopic.multires_macroscopic import MultiresMacroscopic diff --git a/xlb/operator/macroscopic/first_moment.py b/xlb/operator/macroscopic/first_moment.py new file mode 100644 index 00000000..626767fd --- /dev/null +++ b/xlb/operator/macroscopic/first_moment.py @@ -0,0 +1,90 @@ +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class FirstMoment(Operator): + """A class to compute the first moment (velocity) of distribution functions.""" + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0), inline=True) + def jax_implementation(self, f, rho): + u = jnp.tensordot(self.velocity_set.c, f, axes=(-1, 0)) / rho + return u + + def _construct_warp(self): + _c = self.velocity_set.c + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) + + @wp.func + def neumaier_sum_component(d: int, f: _f_vec): + total = self.compute_dtype(0.0) + compensation = self.compute_dtype(0.0) + for l in range(self.velocity_set.q): + # Get contribution based on the sign of _c[d, l] + if _c[d, l] == 1: + val = f[l] + elif _c[d, l] == -1: + val = -f[l] + else: + val = self.compute_dtype(0.0) + t = total + val + if wp.abs(total) >= wp.abs(val): + compensation = compensation + ((total - t) + val) + else: + compensation = compensation + ((val - t) + total) + total = t + return total + compensation + + @wp.func + def functional(f: _f_vec, rho: Any): + u = _u_vec() + # Use Neumaier summation for each spatial component + for d in range(self.velocity_set.d): + u[d] = neumaier_sum_component(d, f) + u /= rho + return u + + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + rho: wp.array4d(dtype=Any), + u: wp.array4d(dtype=Any), + ): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + _f = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _rho = rho[0, index[0], index[1], index[2]] + _u = functional(_f, _rho) + + for d in range(self.velocity_set.d): + u[d, index[0], index[1], index[2]] = self.store_dtype(_u[d]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, rho, u): + wp.launch( + self.warp_kernel, + inputs=[f, rho, u], + dim=u.shape[1:], + ) + return u + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f, rho): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/macroscopic/macroscopic.py b/xlb/operator/macroscopic/macroscopic.py new file mode 100644 index 00000000..cb477b9a --- /dev/null +++ b/xlb/operator/macroscopic/macroscopic.py @@ -0,0 +1,113 @@ +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.macroscopic.zero_moment import ZeroMoment +from xlb.operator.macroscopic.first_moment import FirstMoment + + +class Macroscopic(Operator): + """A class to compute both zero and first moments of distribution functions (rho, u).""" + + def __init__(self, *args, **kwargs): + self.zero_moment = ZeroMoment(*args, **kwargs) + self.first_moment = FirstMoment(*args, **kwargs) + super().__init__(*args, **kwargs) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0), inline=True) + def jax_implementation(self, f, rho=None, u=None): + rho = self.zero_moment(f) + u = self.first_moment(f, rho) + return rho, u + + def _construct_warp(self): + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + @wp.func + def functional(f: _f_vec): + rho = self.zero_moment.warp_functional(f) + u = self.first_moment.warp_functional(f, rho) + return rho, u + + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + rho: wp.array4d(dtype=Any), + u: wp.array4d(dtype=Any), + ): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + _f = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _rho, _u = functional(_f) + + rho[0, index[0], index[1], index[2]] = self.store_dtype(_rho) + for d in range(self.velocity_set.d): + u[d, index[0], index[1], index[2]] = self.store_dtype(_u[d]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, rho, u): + wp.launch( + self.warp_kernel, + inputs=[f, rho, u], + dim=rho.shape[1:], + ) + return rho, u + + def _construct_neon(self): + import neon + + # Redefine the zero and first moment operators for the neon backend + # This is because the neon backend relies on the warp functionals for its operations. + self.zero_moment = ZeroMoment(compute_backend=ComputeBackend.WARP) + self.first_moment = FirstMoment(compute_backend=ComputeBackend.WARP) + functional, _ = self._construct_warp() + + # Set local vectors + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + @neon.Container.factory("macroscopic") + def container( + f_field: Any, + rho_field: Any, + u_fild: Any, + ): + _d = self.velocity_set.d + + def macroscopic_ll(loader: neon.Loader): + loader.set_grid(f_field.get_grid()) + + rho = loader.get_read_handle(rho_field) + u = loader.get_read_handle(u_fild) + f = loader.get_read_handle(f_field) + + @wp.func + def macroscopic_cl(gIdx: Any): + _f = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = self.compute_dtype(wp.neon_read(f, gIdx, l)) + _rho, _u = functional(_f) + wp.neon_write(rho, gIdx, 0, self.store_dtype(_rho)) + for d in range(_d): + wp.neon_write(u, gIdx, d, self.store_dtype(_u[d])) + + loader.declare_kernel(macroscopic_cl) + + return macroscopic_ll + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f, rho, u): + c = self.neon_container(f, rho, u) + c.run(0) + return rho, u diff --git a/xlb/operator/macroscopic/multires_macroscopic.py b/xlb/operator/macroscopic/multires_macroscopic.py new file mode 100644 index 00000000..abd8d6f2 --- /dev/null +++ b/xlb/operator/macroscopic/multires_macroscopic.py @@ -0,0 +1,89 @@ +""" +Multi-resolution macroscopic moment computation for the Neon backend. +""" + +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator +from xlb.operator.macroscopic import Macroscopic, ZeroMoment, FirstMoment +from xlb.cell_type import BC_SOLID + + +class MultiresMacroscopic(Macroscopic): + """Compute density and velocity on a multi-resolution grid (Neon only). + + Iterates over all grid levels, computing zero-th and first moments of + the distribution function. Solid voxels and voxels that have child + refinement (halo cells) are set to zero. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.compute_backend in [ComputeBackend.JAX, ComputeBackend.WARP]: + raise NotImplementedError(f"Operator {self.__class__.__name__} not supported in {self.compute_backend} backend.") + + def _construct_neon(self): + import neon + + # Redefine the zero and first moment operators for the neon backend + # This is because the neon backend relies on the warp functionals for its operations. + self.zero_moment = ZeroMoment(compute_backend=ComputeBackend.WARP) + self.first_moment = FirstMoment(compute_backend=ComputeBackend.WARP) + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + functional, _ = self._construct_warp() + + @neon.Container.factory("macroscopic") + def container( + level: int, + f_field: Any, + bc_mask: Any, + rho_field: Any, + u_fild: Any, + ): + _d = self.velocity_set.d + + def macroscopic_ll(loader: neon.Loader): + loader.set_mres_grid(f_field.get_grid(), level) + + rho = loader.get_mres_write_handle(rho_field) + u = loader.get_mres_write_handle(u_fild) + f = loader.get_mres_read_handle(f_field) + bc_mask_pn = loader.get_mres_read_handle(bc_mask) + + @wp.func + def macroscopic_cl(gIdx: Any): + _f = _f_vec() + _boundary_id = wp.neon_read(bc_mask_pn, gIdx, 0) + + for l in range(self.velocity_set.q): + _f[l] = self.compute_dtype(wp.neon_read(f, gIdx, l)) + + _rho, _u = functional(_f) + + if _boundary_id == wp.uint8(BC_SOLID) or wp.neon_has_child(f, gIdx): + _rho = self.compute_dtype(0.0) + for d in range(_d): + _u[d] = self.compute_dtype(0.0) + + wp.neon_write(rho, gIdx, 0, self.store_dtype(_rho)) + for d in range(_d): + wp.neon_write(u, gIdx, d, self.store_dtype(_u[d])) + + loader.declare_kernel(macroscopic_cl) + + return macroscopic_ll + + return functional, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f, bc_mask, rho, u, streamId=0): + grid = f.get_grid() + for level in range(grid.num_levels): + c = self.neon_container(level, f, bc_mask, rho, u) + c.run(streamId) + return rho, u diff --git a/xlb/operator/macroscopic/second_moment.py b/xlb/operator/macroscopic/second_moment.py new file mode 100644 index 00000000..1a0a0f07 --- /dev/null +++ b/xlb/operator/macroscopic/second_moment.py @@ -0,0 +1,115 @@ +# Base class for all equilibriums + +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class SecondMoment(Operator): + """ + Operator to calculate the second moment of distribution functions. + + The second moment may be used to compute the momentum flux in the computation of + the stress tensor in the Lattice Boltzmann Method (LBM). + + Important Note: + Note that this rank 2 symmetric tensor (dim*dim) has been converted into a rank one + vector where the diagonal and off-diagonal components correspond to the following elements of + the vector: + if self.grid.dim == 3: + diagonal = (0, 3, 5) + offdiagonal = (1, 2, 4) + elif self.grid.dim == 2: + diagonal = (0, 2) + offdiagonal = (1,) + + ** For any reduction operation on the full tensor it is crucial to account for the full tensor by + considering all diagonal and off-diagonal components. + """ + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,), donate_argnums=(1,)) + def jax_implementation( + self, + fneq: jnp.ndarray, + ): + """ + This function computes the second order moment, which is the product of the + distribution functions (f) and the lattice moments (cc). + + Parameters + ---------- + fneq: jax.numpy.ndarray + The distribution functions. + + Returns + ------- + jax.numpy.ndarray + The computed second moment. + """ + return jnp.tensordot(self.velocity_set.cc, fneq, axes=(0, 0)) + + def _construct_warp(self): + # Make constants for warp + _cc = self.velocity_set.cc + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _pi_dim = self.velocity_set.d * (self.velocity_set.d + 1) // 2 + _pi_vec = wp.vec( + _pi_dim, + dtype=self.compute_dtype, + ) + + # Construct functional for computing second moment + @wp.func + def functional( + fneq: Any, + ): + # Get second order moment (a symmetric tensore shaped into a vector) + pi = _pi_vec() + for d in range(_pi_dim): + pi[d] = self.compute_dtype(0.0) + for q in range(self.velocity_set.q): + pi[d] += _cc[q, d] * fneq[q] + return pi + + # Construct the kernel + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + pi: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Get the equilibrium + _f = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _pi = functional(_f) + + # Set the output + for d in range(_pi_dim): + pi[d, index[0], index[1], index[2]] = self.store_dtype(_pi[d]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, pi): + # Launch the warp kernel + wp.launch(self.warp_kernel, inputs=[f, pi], dim=pi.shape[1:]) + return pi + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f, rho): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/macroscopic/zero_moment.py b/xlb/operator/macroscopic/zero_moment.py new file mode 100644 index 00000000..f536f8d7 --- /dev/null +++ b/xlb/operator/macroscopic/zero_moment.py @@ -0,0 +1,70 @@ +from functools import partial +import jax.numpy as jnp +from jax import jit +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class ZeroMoment(Operator): + """A class to compute the zeroth moment (density) of distribution functions.""" + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0), inline=True) + def jax_implementation(self, f): + return jnp.sum(f, axis=0, keepdims=True) + + def _construct_warp(self): + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + @wp.func + def neumaier_sum(f: _f_vec): + total = self.compute_dtype(0.0) + compensation = self.compute_dtype(0.0) + for l in range(self.velocity_set.q): + x = f[l] + t = total + x + # Using wp.abs to compute absolute value + if wp.abs(total) >= wp.abs(x): + compensation = compensation + ((total - t) + x) + else: + compensation = compensation + ((x - t) + total) + total = t + return total + compensation + + @wp.func + def functional(f: _f_vec): + return neumaier_sum(f) + + @wp.kernel + def kernel( + f: wp.array4d(dtype=Any), + rho: wp.array4d(dtype=Any), + ): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + _f = _f_vec() + for l in range(self.velocity_set.q): + _f[l] = f[l, index[0], index[1], index[2]] + _rho = functional(_f) + + rho[0, index[0], index[1], index[2]] = _rho + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f, rho): + wp.launch(self.warp_kernel, inputs=[f, rho], dim=rho.shape[1:]) + return rho + + def _construct_neon(self): + functional, _ = self._construct_warp() + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f, rho): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/operator/operator.py b/xlb/operator/operator.py new file mode 100644 index 00000000..4405708c --- /dev/null +++ b/xlb/operator/operator.py @@ -0,0 +1,316 @@ +""" +Base operator module for XLB. + +Every LBM operator (collision, streaming, equilibrium, boundary condition, +masker, stepper, etc.) inherits from :class:`Operator`. The class provides: + +* **Backend dispatch** β€” ``__call__`` automatically selects the registered + implementation for the active compute backend. +* **Precision management** β€” ``compute_dtype`` and ``store_dtype`` properties + return the correct type for the active backend and precision policy. +* **Kernel construction hooks** β€” ``_construct_warp()`` / ``_construct_neon()`` + are called at init time to compile backend-specific kernels and functionals. +""" + +import inspect +import traceback +import jax +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb import DefaultConfig +from xlb.precision_policy import PrecisionPolicy + + +class Operator: + """ + Base class for all operators, collision, streaming, equilibrium, etc. + + This class is responsible for handling compute backends. + """ + + _backends = {} + + def __init__(self, velocity_set=None, precision_policy=None, compute_backend=None): + """Initialize the operator. + + Parameters + ---------- + velocity_set : VelocitySet, optional + Lattice velocity set. Defaults to ``DefaultConfig.velocity_set``. + precision_policy : PrecisionPolicy, optional + Precision policy. Defaults to ``DefaultConfig.default_precision_policy``. + compute_backend : ComputeBackend, optional + Compute backend. Defaults to ``DefaultConfig.default_backend``. + """ + # Set the default values from the global config + self.velocity_set = velocity_set or DefaultConfig.velocity_set + self.precision_policy = precision_policy or DefaultConfig.default_precision_policy + self.compute_backend = compute_backend or DefaultConfig.default_backend + + # Check if the compute compute_backend is supported + if self.compute_backend not in ComputeBackend: + raise ValueError(f"Compute_backend {compute_backend} is not supported") + + # Construct read/write functions for the compute backend + if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]: + self.read_field, self.write_field = self._construct_read_write_functions() + self.read_field_neighbor = self._construct_read_field_neighbor() + + # Construct the kernel based compute_backend functions TODO: Maybe move this to the register or something + if self.compute_backend == ComputeBackend.WARP: + self.warp_functional, self.warp_kernel = self._construct_warp() + + if self.compute_backend == ComputeBackend.NEON: + self.neon_functional, self.neon_container = self._construct_neon() + + # Updating JAX config in case fp64 is requested + if self.compute_backend == ComputeBackend.JAX and ( + precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32 + ): + jax.config.update("jax_enable_x64", True) + + @classmethod + def register_backend(cls, backend_name): + """ + Decorator to register a compute_backend for the operator. + """ + + def decorator(func): + subclass_name = func.__qualname__.split(".")[0] + signature = inspect.signature(func) + key = (subclass_name, backend_name, str(signature)) + cls._backends[key] = func + return func + + return decorator + + def __call__(self, *args, callback=None, **kwargs): + """Dispatch to the registered backend implementation. + + Iterates over all registered implementations for this operator class + and the active backend, attempts to bind the provided arguments, and + executes the first matching signature. An optional *callback* is + invoked with the result after successful execution. + + Raises + ------ + NotImplementedError + If no implementation is registered for the active backend. + Exception + If all candidate implementations raise errors. + """ + method_candidates = [ + (key, method) for key, method in self._backends.items() if key[0] == self.__class__.__name__ and key[1] == self.compute_backend + ] + if not method_candidates: + supported = [key for key in self._backends.keys() if key[0] == self.__class__.__name__] + raise NotImplementedError( + f"No implementation found for operator {self.__class__.__name__} with backend {self.compute_backend}. " + f"Available implementations: {supported}" + ) + + bound_arguments = None + key = None + error = None + traceback_str = None + for key, backend_method in method_candidates: + try: + # This attempts to bind the provided args and kwargs to the compute_backend method's signature + bound_arguments = inspect.signature(backend_method).bind(self, *args, **kwargs) + bound_arguments.apply_defaults() # This fills in any default values + result = backend_method(self, *args, **kwargs) + callback_arg = result if result is not None else (args, kwargs) + if callback and callable(callback): + callback(callback_arg) + return result + except Exception as e: + error = e + traceback_str = traceback.format_exc() + continue # This skips to the next candidate if binding fails + method_candidates = [(key, method) for key, method in self._backends.items() if key[1] == self.compute_backend] + raise Exception(f"Error captured for backend with key {key} for operator {self.__class__.__name__}: {error}\n {traceback_str}") + + @property + def supported_compute_backend(self): + """ + Returns the supported compute backend for the operator + """ + return list(self._backends.keys()) + + def _is_method_overridden(self, method_name): + """ + Helper method to check if a method is overridden in a subclass. + """ + method = getattr(self, method_name, None) + if method is None: + return False + return method.__func__ is not getattr(Operator, method_name, None).__func__ + + def __repr__(self): + return f"{self.__class__.__name__}()" + + @property + def backend(self): + """ + Returns the compute backend object for the operator (e.g. jax, warp) + This should be used with caution as all backends may not have the same API. + """ + if self.compute_backend == ComputeBackend.JAX: + import jax.numpy as compute_backend + elif self.compute_backend == ComputeBackend.WARP: + import warp as compute_backend + return compute_backend + + @property + def compute_dtype(self): + """ + Returns the compute dtype + """ + if self.compute_backend == ComputeBackend.JAX: + return self.precision_policy.compute_precision.jax_dtype + elif self.compute_backend == ComputeBackend.WARP: + return self.precision_policy.compute_precision.wp_dtype + elif self.compute_backend == ComputeBackend.NEON: + return self.precision_policy.compute_precision.wp_dtype + + @property + def store_dtype(self): + """ + Returns the store dtype + """ + if self.compute_backend == ComputeBackend.JAX: + return self.precision_policy.store_precision.jax_dtype + elif self.compute_backend == ComputeBackend.WARP: + return self.precision_policy.store_precision.wp_dtype + elif self.compute_backend == ComputeBackend.NEON: + return self.precision_policy.store_precision.wp_dtype + + def get_precision_policy(self): + """ + Returns the precision policy + """ + return self.precision_policy + + def get_grid(self): + """ + Returns the grid object + """ + return self.grid + + def _construct_warp(self): + """ + Construct the warp functional and kernel of the operator + TODO: Maybe a better way to do this? + Maybe add this to the compute backend decorator? + Leave it for now, as it is not clear how the warp compute backend will evolve + """ + return None, None + + def _construct_neon(self): + """ + Construct the Neon functional and Neon container of the operator + TODO: Maybe a better way to do this? + Maybe add this to the backend decorator? + Leave it for now, as it is not clear how the neon backend will evolve + """ + return None, None + + def _construct_read_write_functions(self): + """Build backend-specific ``read_field`` / ``write_field`` helpers. + + For the Warp backend these are direct 4-D array accesses. For the + Neon backend they wrap ``wp.neon_read`` / ``wp.neon_write``. + + Returns + ------- + tuple of wp.func + ``(read_field, write_field)`` + """ + if self.compute_backend == ComputeBackend.WARP: + + @wp.func + def read_field( + field: Any, + index: Any, + direction: Any, + ): + # This function reads a field value at a given index and direction. + return field[direction, index[0], index[1], index[2]] + + @wp.func + def write_field( + field: Any, + index: Any, + direction: Any, + value: Any, + ): + # This function writes a value to a field at a given index and direction. + field[direction, index[0], index[1], index[2]] = value + + elif self.compute_backend == ComputeBackend.NEON: + import neon + + @wp.func + def read_field( + field: Any, + index: Any, + direction: Any, + ): + # This function reads a field value at a given index and direction. + return wp.neon_read(field, index, direction) + + @wp.func + def write_field( + field: Any, + index: Any, + direction: Any, + value: Any, + ): + # This function writes a value to a field at a given index and direction. + wp.neon_write(field, index, direction, value) + + else: + raise ValueError(f"Unsupported compute backend: {self.compute_backend}") + + return read_field, write_field + + def _construct_read_field_neighbor(self): + """ + Construct a function to read a field value at a neighboring index along a given direction. + """ + + if self.compute_backend == ComputeBackend.WARP: + + @wp.func + def read_field_neighbor( + field: Any, + index: Any, + offset: Any, + direction: Any, + ): + # This function reads a field value at a given neighboring index and direction. + neighbor = index + offset + return field[direction, neighbor[0], neighbor[1], neighbor[2]] + + elif self.compute_backend == ComputeBackend.NEON: + import neon + # from neon.multires.mPartition import neon_get_type + + @wp.func + def read_field_neighbor( + field: Any, + index: Any, + offset: Any, + direction: Any, + ): + # This function reads a field value at a given neighboring index and direction. + unused_is_valid = wp.bool(False) + # dtype = neon_get_type(field) # This is a placeholder to ensure the dtype is set correctly + return wp.neon_read_ngh(field, index, offset, direction, wp.uint8(0.0), unused_is_valid) + + else: + raise ValueError(f"Unsupported compute backend: {self.compute_backend}") + + return read_field_neighbor diff --git a/xlb/operator/parallel_operator.py b/xlb/operator/parallel_operator.py new file mode 100644 index 00000000..e0505f03 --- /dev/null +++ b/xlb/operator/parallel_operator.py @@ -0,0 +1,82 @@ +from jax import shard_map +from jax.sharding import PartitionSpec as P +from jax import lax + + +class ParallelOperator: + """ + A generic class for parallelizing operations across multiple GPUs/TPUs. + """ + + def __init__(self, grid, func, velocity_set): + """ + Initialize the ParallelOperator. + + Parameters + ---------- + grid : Grid object + The computational grid. + func : function + The function to be parallelized. + velocity_set : VelocitySet object + The velocity set used in the Lattice Boltzmann Method. + """ + self.grid = grid + self.func = func + self.velocity_set = velocity_set + + def __call__(self, f): + """ + Execute the parallel operation. + + Parameters + ---------- + f : jax.numpy.ndarray + The input data for the operation. + + Returns + ------- + jax.numpy.ndarray + The result after applying the parallel operation. + """ + in_specs = P(*((None, "x") + (self.grid.dim - 1) * (None,))) + out_specs = in_specs + + f = shard_map( + self._parallel_func, + mesh=self.grid.global_mesh, + in_specs=in_specs, + out_specs=out_specs, + check_vma=False, + )(f) + return f + + def _parallel_func(self, f): + """ + Internal function to handle data communication and apply the given function in parallel. + + Parameters + ---------- + f : jax.numpy.ndarray + The input data. + + Returns + ------- + jax.numpy.ndarray + The processed data. + """ + rightPerm = [(i, (i + 1) % self.grid.nDevices) for i in range(self.grid.nDevices)] + leftPerm = [((i + 1) % self.grid.nDevices, i) for i in range(self.grid.nDevices)] + f = self.func(f) + left_comm, right_comm = ( + f[self.velocity_set.right_indices, :1, ...], + f[self.velocity_set.left_indices, -1:, ...], + ) + left_comm, right_comm = ( + lax.ppermute(left_comm, perm=rightPerm, axis_name="x"), + lax.ppermute(right_comm, perm=leftPerm, axis_name="x"), + ) + f = f.at[self.velocity_set.right_indices, :1, ...].set(left_comm) + f = f.at[self.velocity_set.left_indices, -1:, ...].set(right_comm) + + return f diff --git a/xlb/operator/postprocess/__init__.py b/xlb/operator/postprocess/__init__.py new file mode 100644 index 00000000..386dd917 --- /dev/null +++ b/xlb/operator/postprocess/__init__.py @@ -0,0 +1,3 @@ +from xlb.operator.postprocess.q_criterion import QCriterion +from xlb.operator.postprocess.grid_to_point import GridToPoint +from xlb.operator.postprocess.vorticity import Vorticity diff --git a/xlb/operator/postprocess/grid_to_point.py b/xlb/operator/postprocess/grid_to_point.py new file mode 100644 index 00000000..a30b79e4 --- /dev/null +++ b/xlb/operator/postprocess/grid_to_point.py @@ -0,0 +1,109 @@ +import warp as wp +from typing import Any +from functools import partial +from jax import jit + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class GridToPoint(Operator): + """ + Interpolate values from a grid to arbitrary points using trilinear interpolation. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + def _construct_warp(self): + # Construct the warp kernel + @wp.kernel + def kernel( + grid: wp.array4d(dtype=Any), + points: wp.array(dtype=wp.vec3), + point_values: wp.array(dtype=Any), + ): + # Get the global index + i = wp.tid() + + # Get the point + point = points[i] + + # Get lower and upper bounds + lower_0_0_0 = wp.vec3i(wp.int32(point[0]), wp.int32(point[1]), wp.int32(point[2])) + lower_0_0_1 = lower_0_0_0 + wp.vec3i(0, 0, 1) + lower_0_1_0 = lower_0_0_0 + wp.vec3i(0, 1, 0) + lower_0_1_1 = lower_0_0_0 + wp.vec3i(0, 1, 1) + lower_1_0_0 = lower_0_0_0 + wp.vec3i(1, 0, 0) + lower_1_0_1 = lower_0_0_0 + wp.vec3i(1, 0, 1) + lower_1_1_0 = lower_0_0_0 + wp.vec3i(1, 1, 0) + lower_1_1_1 = lower_0_0_0 + wp.vec3i(1, 1, 1) + + # Get grid values + grid_0_0_0 = grid[0, lower_0_0_0[0], lower_0_0_0[1], lower_0_0_0[2]] + grid_0_0_1 = grid[0, lower_0_0_1[0], lower_0_0_1[1], lower_0_0_1[2]] + grid_0_1_0 = grid[0, lower_0_1_0[0], lower_0_1_0[1], lower_0_1_0[2]] + grid_0_1_1 = grid[0, lower_0_1_1[0], lower_0_1_1[1], lower_0_1_1[2]] + grid_1_0_0 = grid[0, lower_1_0_0[0], lower_1_0_0[1], lower_1_0_0[2]] + grid_1_0_1 = grid[0, lower_1_0_1[0], lower_1_0_1[1], lower_1_0_1[2]] + grid_1_1_0 = grid[0, lower_1_1_0[0], lower_1_1_0[1], lower_1_1_0[2]] + grid_1_1_1 = grid[0, lower_1_1_1[0], lower_1_1_1[1], lower_1_1_1[2]] + + # Compute the interpolation weights + dx = point[0] - wp.float32(lower_0_0_0[0]) + dy = point[1] - wp.float32(lower_0_0_0[1]) + dz = point[2] - wp.float32(lower_0_0_0[2]) + w_000 = (1.0 - dx) * (1.0 - dy) * (1.0 - dz) + w_001 = (1.0 - dx) * (1.0 - dy) * dz + w_010 = (1.0 - dx) * dy * (1.0 - dz) + w_011 = (1.0 - dx) * dy * dz + w_100 = dx * (1.0 - dy) * (1.0 - dz) + w_101 = dx * (1.0 - dy) * dz + w_110 = dx * dy * (1.0 - dz) + w_111 = dx * dy * dz + + # Compute the interpolated value + # Trilinear interpolation: sum contributions from each corner of the cube + point_value = ( + (w_000 * grid_0_0_0) # (0,0,0) corner + + (w_001 * grid_0_0_1) # (0,0,1) corner + + (w_010 * grid_0_1_0) # (0,1,0) corner + + (w_011 * grid_0_1_1) # (0,1,1) corner + + (w_100 * grid_1_0_0) # (1,0,0) corner + + (w_101 * grid_1_0_1) # (1,0,1) corner + + (w_110 * grid_1_1_0) # (1,1,0) corner + + (w_111 * grid_1_1_1) # (1,1,1) corner + ) + + # Set the output + point_values[i] = point_value + + return None, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, grid, points, point_values): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[grid, points, point_values], + dim=[points.shape[0]], + ) + return point_values + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, grid, points, point_values): + # TODO: Implement JAX version + raise NotImplementedError("JAX implementation not yet available") diff --git a/xlb/operator/postprocess/q_criterion.py b/xlb/operator/postprocess/q_criterion.py new file mode 100644 index 00000000..227b3655 --- /dev/null +++ b/xlb/operator/postprocess/q_criterion.py @@ -0,0 +1,143 @@ +# Credit https://github.com/loliverhennigh/XLB Original Author: Oliver Hennigh +import warp as wp +from typing import Any +from functools import partial +from jax import jit + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class QCriterion(Operator): + """ + Compute Q-criterion and vorticity magnitude for flow visualization and analysis. + + The Q-criterion is the second invariant of the velocity gradient tensor, + defined as Q = 1/2(|Ξ©|^2 - |S|^2) where Ξ© is the vorticity tensor and + S is the rate-of-strain tensor. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + def _construct_warp(self): + # Construct the warp kernel + @wp.kernel + def kernel( + u: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + norm_mu: wp.array4d(dtype=Any), + q: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + + # Add ghost cells to index + i += 1 + j += 1 + k += 1 + + # Check if anything on edges + b_id_2_1_1 = bc_mask[0, i + 1, j, k] + b_id_1_2_1 = bc_mask[0, i, j + 1, k] + b_id_1_1_2 = bc_mask[0, i, j, k + 1] + b_id_0_1_1 = bc_mask[0, i - 1, j, k] + b_id_1_0_1 = bc_mask[0, i, j - 1, k] + b_id_1_1_0 = bc_mask[0, i, j, k - 1] + if ( + b_id_2_1_1 != wp.uint8(0) + or b_id_1_2_1 != wp.uint8(0) + or b_id_1_1_2 != wp.uint8(0) + or b_id_0_1_1 != wp.uint8(0) + or b_id_1_0_1 != wp.uint8(0) + or b_id_1_1_0 != wp.uint8(0) + ): + return + + # Get derivatives + u_x_dx = (u[0, i + 1, j, k] - u[0, i - 1, j, k]) / 2.0 + u_x_dy = (u[0, i, j + 1, k] - u[0, i, j - 1, k]) / 2.0 + u_x_dz = (u[0, i, j, k + 1] - u[0, i, j, k - 1]) / 2.0 + u_y_dx = (u[1, i + 1, j, k] - u[1, i - 1, j, k]) / 2.0 + u_y_dy = (u[1, i, j + 1, k] - u[1, i, j - 1, k]) / 2.0 + u_y_dz = (u[1, i, j, k + 1] - u[1, i, j, k - 1]) / 2.0 + u_z_dx = (u[2, i + 1, j, k] - u[2, i - 1, j, k]) / 2.0 + u_z_dy = (u[2, i, j + 1, k] - u[2, i, j - 1, k]) / 2.0 + u_z_dz = (u[2, i, j, k + 1] - u[2, i, j, k - 1]) / 2.0 + + # Compute vorticity + mu_x = u_z_dy - u_y_dz + mu_y = u_x_dz - u_z_dx + mu_z = u_y_dx - u_x_dy + mu = wp.sqrt(mu_x**2.0 + mu_y**2.0 + mu_z**2.0) + + # Compute strain rate + s_0_0 = u_x_dx + s_0_1 = 0.5 * (u_x_dy + u_y_dx) + s_0_2 = 0.5 * (u_x_dz + u_z_dx) + s_1_0 = s_0_1 + s_1_1 = u_y_dy + s_1_2 = 0.5 * (u_y_dz + u_z_dy) + s_2_0 = s_0_2 + s_2_1 = s_1_2 + s_2_2 = u_z_dz + s_dot_s = s_0_0**2.0 + s_0_1**2.0 + s_0_2**2.0 + s_1_0**2.0 + s_1_1**2.0 + s_1_2**2.0 + s_2_0**2.0 + s_2_1**2.0 + s_2_2**2.0 + + # Compute omega + omega_0_0 = 0.0 + omega_0_1 = 0.5 * (u_x_dy - u_y_dx) + omega_0_2 = 0.5 * (u_x_dz - u_z_dx) + omega_1_0 = -omega_0_1 + omega_1_1 = 0.0 + omega_1_2 = 0.5 * (u_y_dz - u_z_dy) + omega_2_0 = -omega_0_2 + omega_2_1 = -omega_1_2 + omega_2_2 = 0.0 + omega_dot_omega = ( + omega_0_0**2.0 + + omega_0_1**2.0 + + omega_0_2**2.0 + + omega_1_0**2.0 + + omega_1_1**2.0 + + omega_1_2**2.0 + + omega_2_0**2.0 + + omega_2_1**2.0 + + omega_2_2**2.0 + ) + + # Compute q-criterion + q_value = 0.5 * (omega_dot_omega - s_dot_s) + + # Set the output + norm_mu[0, i, j, k] = mu + q[0, i, j, k] = q_value + + return None, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, u, bc_mask, norm_mu, q): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[u, bc_mask, norm_mu, q], + dim=[i - 2 for i in u.shape[1:]], + ) + return norm_mu, q + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, u, bc_mask, norm_mu, q): + # TODO: Implement JAX version + raise NotImplementedError("JAX implementation not yet available") diff --git a/xlb/operator/postprocess/vorticity.py b/xlb/operator/postprocess/vorticity.py new file mode 100644 index 00000000..55382d02 --- /dev/null +++ b/xlb/operator/postprocess/vorticity.py @@ -0,0 +1,101 @@ +import warp as wp +from typing import Any +from functools import partial +from jax import jit + +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class Vorticity(Operator): + """ + Compute vorticity vector and magnitude for flow visualization and analysis. + """ + + def __init__( + self, + velocity_set: VelocitySet = None, + precision_policy: PrecisionPolicy = None, + compute_backend: ComputeBackend = None, + ): + # Call the parent constructor + super().__init__( + velocity_set, + precision_policy, + compute_backend, + ) + + def _construct_warp(self): + # Construct the warp kernel + @wp.kernel + def kernel( + u: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=wp.uint8), + vorticity: wp.array4d(dtype=Any), + vorticity_magnitude: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + + # Add ghost cells to index + i += 1 + j += 1 + k += 1 + + # Check if anything on edges + b_id_2_1_1 = bc_mask[0, i + 1, j, k] + b_id_1_2_1 = bc_mask[0, i, j + 1, k] + b_id_1_1_2 = bc_mask[0, i, j, k + 1] + b_id_0_1_1 = bc_mask[0, i - 1, j, k] + b_id_1_0_1 = bc_mask[0, i, j - 1, k] + b_id_1_1_0 = bc_mask[0, i, j, k - 1] + if ( + b_id_2_1_1 != wp.uint8(0) + or b_id_1_2_1 != wp.uint8(0) + or b_id_1_1_2 != wp.uint8(0) + or b_id_0_1_1 != wp.uint8(0) + or b_id_1_0_1 != wp.uint8(0) + or b_id_1_1_0 != wp.uint8(0) + ): + return + + # Get derivatives using central differences + u_x_dy = (u[0, i, j + 1, k] - u[0, i, j - 1, k]) / 2.0 + u_x_dz = (u[0, i, j, k + 1] - u[0, i, j, k - 1]) / 2.0 + u_y_dx = (u[1, i + 1, j, k] - u[1, i - 1, j, k]) / 2.0 + u_y_dz = (u[1, i, j, k + 1] - u[1, i, j, k - 1]) / 2.0 + u_z_dx = (u[2, i + 1, j, k] - u[2, i - 1, j, k]) / 2.0 + u_z_dy = (u[2, i, j + 1, k] - u[2, i, j - 1, k]) / 2.0 + + # Compute vorticity components (curl of velocity) + vort_x = u_z_dy - u_y_dz + vort_y = u_x_dz - u_z_dx + vort_z = u_y_dx - u_x_dy + + # Store vorticity vector components + vorticity[0, i, j, k] = vort_x + vorticity[1, i, j, k] = vort_y + vorticity[2, i, j, k] = vort_z + + # Compute and store vorticity magnitude + vorticity_magnitude[0, i, j, k] = wp.sqrt(vort_x * vort_x + vort_y * vort_y + vort_z * vort_z) + + return None, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, u, bc_mask, vorticity, vorticity_magnitude): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[u, bc_mask, vorticity, vorticity_magnitude], + dim=[i - 2 for i in u.shape[1:]], + ) + return vorticity, vorticity_magnitude + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, u, bc_mask, vorticity, vorticity_magnitude): + # TODO: Implement JAX version + raise NotImplementedError("JAX implementation not yet available") diff --git a/xlb/operator/precision_caster/__init__.py b/xlb/operator/precision_caster/__init__.py new file mode 100644 index 00000000..a027c524 --- /dev/null +++ b/xlb/operator/precision_caster/__init__.py @@ -0,0 +1 @@ +from xlb.operator.precision_caster.precision_caster import PrecisionCaster diff --git a/xlb/operator/precision_caster/precision_caster.py b/xlb/operator/precision_caster/precision_caster.py new file mode 100644 index 00000000..5427cba9 --- /dev/null +++ b/xlb/operator/precision_caster/precision_caster.py @@ -0,0 +1,96 @@ +""" +Base class for casting precision of the input data to the desired precision +""" + +import jax.numpy as jnp +from jax import jit +import warp as wp +from functools import partial + +from xlb.operator.operator import Operator +from xlb.velocity_set import VelocitySet +from xlb.precision_policy import Precision, PrecisionPolicy +from xlb.compute_backend import ComputeBackend + + +class PrecisionCaster(Operator): + """ + Class that handles the construction of lattice boltzmann precision casting operator. + """ + + def __init__( + self, + input_precision: Precision, + output_precision: Precision, + velocity_set: VelocitySet, + precision_policy: PrecisionPolicy, + compute_backend: ComputeBackend, + ): + super().__init__( + velocity_set=velocity_set, + precision_policy=precision_policy, + compute_backend=compute_backend, + ) + + # Set the input and output precision based on the backend + self.input_precision = self._precision_to_dtype(input_precision) + self.output_precision = self._precision_to_dtype(output_precision) + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,)) + def jax_implementation(self, f: jnp.ndarray) -> jnp.ndarray: + return self.output_precision(f) + + def _construct_warp(self): + # Construct needed types and constants + from_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.input_precision) + to_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.output_precision) + from_array_type = wp.array4d(dtype=self.input_precision) + to_array_type = wp.array4d(dtype=self.output_precision) + _q = wp.constant(self.velocity_set.q) + + # Construct the functional + @wp.func + def functional( + from_f: from_lattice_vec, + ) -> to_lattice_vec: + to_f = to_lattice_vec() + for i in range(self.velocity_set.q): + to_f[i] = self.output_precision(from_f[i]) + return to_f + + # Construct the warp kernel + @wp.kernel + def kernel( + from_f: from_array_type, + to_f: to_array_type, + ): + # Get the global index + i, j, k = wp.tid() + + # Get f + _from_f = from_lattice_vec() + for l in range(_q): + _from_f[l] = from_f[l, i, j, k] + + # Cast the precision + _to_f = functional(_from_f) + + # Set f + for l in range(_q): + to_f[l, i, j, k] = _to_f[l] + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, from_f, to_f): + # Launch the warp kernel + wp.launch( + self._kernel, + inputs=[ + from_f, + to_f, + ], + dim=from_f.shape[1:], + ) + return to_f diff --git a/xlb/operator/stepper/__init__.py b/xlb/operator/stepper/__init__.py new file mode 100644 index 00000000..87c1274f --- /dev/null +++ b/xlb/operator/stepper/__init__.py @@ -0,0 +1,4 @@ +from xlb.operator.stepper.stepper import Stepper +from xlb.operator.stepper.nse_stepper import IncompressibleNavierStokesStepper +from xlb.operator.stepper.nse_multires_stepper import MultiresIncompressibleNavierStokesStepper +from xlb.operator.stepper.ibm_stepper import IBMStepper diff --git a/xlb/operator/stepper/ibm_stepper.py b/xlb/operator/stepper/ibm_stepper.py new file mode 100644 index 00000000..8507b7bf --- /dev/null +++ b/xlb/operator/stepper/ibm_stepper.py @@ -0,0 +1,476 @@ +from functools import partial +from jax import jit +import warp as wp +from typing import Any +from contextlib import nullcontext + +from xlb.compute_backend import ComputeBackend +from xlb.operator import Operator +from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry +from xlb.operator.stepper.nse_stepper import IncompressibleNavierStokesStepper +from warp import ScopedTimer + + +class IBMStepper(IncompressibleNavierStokesStepper): + """ + Incompressible Navier-Stokes stepper with immersed boundary coupling. + + Note: + The iterative IBM loop follows the spirit of multi-direct forcing schemes + such as those discussed by Inamuro (2012) and Ataei et al. (2022). The flow + field is corrected multiple times within a single fluid step so that the + no-slip constraint on immersed surfaces is satisfied more accurately. This + implementation differs in three ways: + + 1) Velocity-based target field instead of direct force spreading. Classical + IBM spreads the Lagrangian forces F_k with + + f(x_i) ~= sum_k F_k delta_h(x_i - X_k) * DeltaA_k. + + Here we accumulate a target Eulerian velocity using a partition of unity: + + eul_forces[i] <- sum_k w_ik A_k U*_k + eul_weights[i] <- sum_k w_ik + target_u[i] = eul_forces[i] / eul_weights[i] (if the weight sum is > 0) + correction_force[i] = target_u[i] - u[i] + + Normalizing by eul_weights keeps the constraint consistent for nonuniform + marker spacing and avoids over-forcing when many markers map to one node. + + 2) Relaxed fixed point iteration. Lagrangian forces are updated from the + mismatch between solid and interpolated fluid velocities, and the + resulting IBM correction on the Eulerian grid is under-relaxed via + ibm_relaxation inside compute_velocity_and_correct. This turns the IBM + loop into a relaxed fixed point iteration which improves robustness for + for fine Lagrangian resolution and high Reynolds number flows. + + 3) Residual-based stopping instead of a fixed iteration count. The loop + monitors the maximum incremental change in Lagrangian forces and stops + early when it falls below ibm_tolerance, avoiding unnecessary sweeps when + the constraint is already well satisfied. A pinned host flag is aliased on + the device so kernels can report residual breaches without extra copies, + and the host only synchronizes when that flag is read, keeping the final + check inexpensive while still guaranteeing correctness. + + The Eulerian-Lagrangian coupling is implemented with a hash grid so that + interpolation and spreading share the same neighbor search, which keeps the + method efficient on GPUs and compatible with multi-resolution meshes. + """ + + def __init__( + self, + grid, + boundary_conditions=[], + collision_type="BGK", + use_scoped_timer=False, + ibm_max_iterations=4, + ibm_tolerance=1e-5, + ibm_relaxation=1.0, + ): + super().__init__(grid, boundary_conditions, collision_type) + self.timer_context = ( + ScopedTimer("IBM_Stepper", use_nvtx=True, synchronize=True, cuda_filter=wp.TIMING_ALL) if use_scoped_timer else nullcontext() + ) + + self.ibm_max_iterations = ibm_max_iterations + self.ibm_tolerance = ibm_tolerance + self.ibm_relaxation = ibm_relaxation + + self.grid_dim = grid.shape + dim_x, dim_y, dim_z = self.grid_dim + + # Initialize Eulerian points array + self.f_eulerian_points = wp.zeros(shape=(dim_x * dim_y * dim_z), dtype=wp.vec3) + self.f_eulerian_forces = wp.zeros(shape=(dim_x * dim_y * dim_z), dtype=wp.vec3) + self.f_eulerian_velocities = wp.zeros(shape=(dim_x * dim_y * dim_z), dtype=wp.vec3) + self.f_eulerian_weights = wp.zeros(shape=(dim_x * dim_y * dim_z), dtype=self.compute_dtype) + + @wp.func + def hash_to_grid_idx(hash_idx: int, dim_x: int, dim_y: int) -> wp.vec3i: + """Convert hash grid index to 3D grid coordinates""" + k = hash_idx // (dim_x * dim_y) + j = (hash_idx % (dim_x * dim_y)) // dim_x + i = hash_idx % dim_x + return wp.vec3i(i, j, k) + + @wp.func + def grid_to_hash_idx(i: int, j: int, k: int, dim_x: int, dim_y: int) -> int: + """Convert 3D grid coordinates to hash grid index""" + return k * (dim_x * dim_y) + j * dim_x + i + + @wp.kernel + def init_eulerian_points(points: wp.array(dtype=wp.vec3), dim_x: int, dim_y: int, dim_z: int): + idx = wp.tid() + grid_pos = hash_to_grid_idx(idx, dim_x, dim_y) + points[idx] = wp.vec3(float(grid_pos[0]) + 0.5, float(grid_pos[1]) + 0.5, float(grid_pos[2]) + 0.5) + + # Launch kernel to initialize points + wp.launch(kernel=init_eulerian_points, dim=dim_x * dim_y * dim_z, inputs=[self.f_eulerian_points, dim_x, dim_y, dim_z]) + + self.hash_grid = wp.HashGrid(dim_x=dim_x, dim_y=dim_y, dim_z=dim_z) + self.hash_grid.build(self.f_eulerian_points, 2.0) # 2.0 is the radius + + self.s_lagr_forces_initialized = False + + self._construct_ibm_warp() + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep): + raise NotImplementedError("IBM stepper is not implemented in JAX backend. Please use WARP backend.") + + def _construct_ibm_warp(self): + # Set local constants + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8) + _opp_indices = self.velocity_set.opp_indices + _weights = self.velocity_set.w + _c = self.velocity_set.c + _dim_x = self.grid_dim[0] + _dim_y = self.grid_dim[1] + + # Read the list of bc_to_id created upon instantiation + bc_to_id = boundary_condition_registry.bc_to_id + + # Gather IDs of ExtrapolationOutflowBC boundary conditions + extrapolation_outflow_bc_ids = [] + for bc_name, bc_id in bc_to_id.items(): + if bc_name.startswith("ExtrapolationOutflowBC"): + extrapolation_outflow_bc_ids.append(bc_id) + # Group active boundary conditions + active_bcs = set(boundary_condition_registry.id_to_bc[bc.id] for bc in self.boundary_conditions) + + @wp.func + def hash_to_grid_idx(hash_idx: int, dim_x: int, dim_y: int) -> wp.vec3i: + """Convert hash grid index to 3D grid coordinates""" + k = hash_idx // (dim_x * dim_y) + j = (hash_idx % (dim_x * dim_y)) // dim_x + i = hash_idx % dim_x + return wp.vec3i(i, j, k) + + @wp.func + def grid_to_hash_idx(i: int, j: int, k: int, dim_x: int, dim_y: int) -> int: + """Convert 3D grid coordinates to hash grid index""" + return k * (dim_x * dim_y) + j * dim_x + i + + # Smoothing function as proposed by Peskin + @wp.func + def peskin_weight(r: float): + abs_r = wp.abs(r) + if abs_r <= 1.0: + return self.compute_dtype(0.125) * ( + self.compute_dtype(3.0) + - 2.0 * abs_r + + wp.sqrt(self.compute_dtype(1.0) + self.compute_dtype(4.0) * abs_r - self.compute_dtype(4.0) * abs_r * abs_r) + ) + elif abs_r <= 2.0: + return self.compute_dtype(0.125) * ( + self.compute_dtype(5.0) + - 2.0 * abs_r + - wp.sqrt(self.compute_dtype(-7.0) + self.compute_dtype(12.0) * abs_r - self.compute_dtype(4.0) * abs_r * abs_r) + ) + else: + return self.compute_dtype(0.0) + + @wp.func + def weight(x: wp.vec3, Xk: wp.vec3): + r = x - Xk + return peskin_weight(r[0]) * peskin_weight(r[1]) * peskin_weight(r[2]) + + # Kernel to initialize the force on Lagrangian points (Step 1) + @wp.kernel + def initialize_lagr_force( + solid_lagr_velocities: wp.array(dtype=wp.vec3), + fluid_lagr_velocities: wp.array(dtype=wp.vec3), + lag_forces: wp.array(dtype=wp.vec3), + ): + tid = wp.tid() + vk = solid_lagr_velocities[tid] + u_Xk = fluid_lagr_velocities[tid] + + # Initialize force + lag_forces[tid] = vk - u_Xk + + # Kernel to interpolate force from Lagrangian to Eulerian grid (Step 2) + @wp.kernel + def interpolate_force_to_eulerian_atomic( + lag_positions: wp.array(dtype=wp.vec3), + lag_forces: wp.array(dtype=wp.vec3), + lag_areas: wp.array(dtype=Any), + eul_positions: wp.array(dtype=wp.vec3), + eul_forces: wp.array(dtype=wp.vec3), + eul_weights: wp.array(dtype=Any), # Accumulator for weights + grid: wp.uint64, + ): + tid = wp.tid() + Xk = lag_positions[tid] + Fk = lag_forces[tid] + Ak = lag_areas[tid] + + # Query neighboring Eulerian points + query = wp.hash_grid_query(grid, Xk, 2.0) + index = int(0) + + while wp.hash_grid_query_next(query, index): + x_pos = eul_positions[index] + w = weight(x_pos, Xk) + # First accumulate the weight + wp.atomic_add(eul_weights, index, w) + # Then accumulate the weighted force + delta_f = Fk * w * Ak + wp.atomic_add(eul_forces, index, delta_f) + + @wp.kernel + def compute_eulerian_velocity_from_f_1(f_1: wp.array4d(dtype=Any), eul_velocities: wp.array(dtype=wp.vec3)): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + # Read from thread local memory + _f_1_thread = _f_vec() + + for l in range(self.velocity_set.q): + _f_1_thread[l] = self.compute_dtype(f_1[l, index[0], index[1], index[2]]) + + _rho, _u = self.macroscopic.warp_functional(_f_1_thread) + + eul_velocities[grid_to_hash_idx(i, j, k, _dim_x, _dim_y)] = _u + + @wp.kernel + def correct_population_ibm(f_1: wp.array4d(dtype=Any), eul_forces: wp.array(dtype=wp.vec3)): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Initialize thread-local storage for populations + _f1_thread = _f_vec() + + # Retrieve f_1 values for the current grid point + for l in range(self.velocity_set.q): + _f1_thread[l] = self.compute_dtype(f_1[l, index[0], index[1], index[2]]) + + # Compute macroscopic quantities (rho, u) from f_1 + _rho, _u = self.macroscopic.warp_functional(_f1_thread) + + # Retrieve the force at the current grid point + force = eul_forces[grid_to_hash_idx(i, j, k, _dim_x, _dim_y)] + + # Compute equilibrium with force applied + feq_force = self.equilibrium.warp_functional(_rho, _u + force) + feq = self.equilibrium.warp_functional(_rho, _u) + + # Update f_1 with the new post-collision population + for l in range(self.velocity_set.q): + f_1[l, index[0], index[1], index[2]] += self.store_dtype(feq_force[l] - feq[l]) + + # Add a new kernel that combines force interpolation and conservation in one step + @wp.kernel + def improved_interpolate_force_to_eulerian( + lag_positions: wp.array(dtype=wp.vec3), + lag_forces: wp.array(dtype=wp.vec3), + lag_areas: wp.array(dtype=Any), + eul_positions: wp.array(dtype=wp.vec3), + eul_forces: wp.array(dtype=wp.vec3), # Will store desired velocity, not force directly + eul_weights: wp.array(dtype=Any), # For normalization + grid: wp.uint64, + ): + tid = wp.tid() + Xk = lag_positions[tid] + Fk = lag_forces[tid] # Fk here represents the desired velocity change at the Lagrangian point + Ak = lag_areas[tid] + + # Query neighboring Eulerian points + query = wp.hash_grid_query(grid, Xk, 2.0) + index = int(0) + + while wp.hash_grid_query_next(query, index): + x_pos = eul_positions[index] + w = weight(x_pos, Xk) + + # The weight represents how much this Lagrangian point influences this Eulerian point + wp.atomic_add(eul_weights, index, w) + + # We accumulate the weighted desired velocity from each Lagrangian point + # Each Lagrangian point contributes according to its weight and area + target_velocity = Fk * w * Ak + wp.atomic_add(eul_forces, index, target_velocity) + + # Add this to the constructor + self.improved_interpolate_force_to_eulerian = improved_interpolate_force_to_eulerian + + _ibm_relaxation = self.compute_dtype(self.ibm_relaxation) + + @wp.kernel + def compute_velocity_and_correct( + f_1: wp.array4d(dtype=Any), + eul_forces: wp.array(dtype=wp.vec3), + eul_weights: wp.array(dtype=Any), + eul_velocities: wp.array(dtype=wp.vec3), + ): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + _f1_thread = _f_vec() + + for l in range(self.velocity_set.q): + _f1_thread[l] = self.compute_dtype(f_1[l, index[0], index[1], index[2]]) + + _rho, _u = self.macroscopic.warp_functional(_f1_thread) + + hash_idx = grid_to_hash_idx(i, j, k, _dim_x, _dim_y) + eul_velocities[hash_idx] = _u + + weight_sum = eul_weights[hash_idx] + + if weight_sum > self.compute_dtype(0.0): + target_velocity = eul_forces[hash_idx] / weight_sum + correction_force = target_velocity - _u + eul_forces[hash_idx] = _ibm_relaxation * correction_force + + @wp.kernel + def interpolate_velocity_and_update_force( + lag_positions: wp.array(dtype=wp.vec3), + eul_positions: wp.array(dtype=wp.vec3), + eul_velocities: wp.array(dtype=wp.vec3), + solid_lagr_velocities: wp.array(dtype=wp.vec3), + lag_forces: wp.array(dtype=wp.vec3), + lag_forces_prev: wp.array(dtype=wp.vec3), + convergence_flag: wp.array(dtype=wp.int32), + compute_residual: int, + tolerance_sq: Any, + grid: wp.uint64, + ): + tid = wp.tid() + Xk = lag_positions[tid] + + numerator = wp.vec3(self.compute_dtype(0.0), self.compute_dtype(0.0), self.compute_dtype(0.0)) + denominator = self.compute_dtype(0.0) + + query = wp.hash_grid_query(grid, Xk, 2.0) + index = int(0) + + while wp.hash_grid_query_next(query, index): + x_pos = eul_positions[index] + u = eul_velocities[index] + w_val = weight(x_pos, Xk) + numerator += u * w_val + denominator += w_val + + if denominator > self.compute_dtype(0.0): + u_interp = numerator / denominator + else: + u_interp = wp.vec3(self.compute_dtype(0.0), self.compute_dtype(0.0), self.compute_dtype(0.0)) + + delta_F = solid_lagr_velocities[tid] - u_interp + lag_forces[tid] += delta_F + + if compute_residual != 0: + diff = lag_forces[tid] - lag_forces_prev[tid] + squared_norm = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2] + if squared_norm > tolerance_sq: + wp.atomic_max(convergence_flag, 0, 1) + + self.compute_velocity_and_correct = compute_velocity_and_correct + self.interpolate_velocity_and_update_force = interpolate_velocity_and_update_force + + self.initialize_lagr_force = initialize_lagr_force + self.interpolate_force_to_eulerian_atomic = interpolate_force_to_eulerian_atomic + self.compute_eulerian_velocity_from_f_1 = compute_eulerian_velocity_from_f_1 + self.correct_population_ibm = correct_population_ibm + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation( + self, + f_0, + f_1, + s_lagr_vertices_wp, + lagr_solid_vertex_areas_wp, + lagr_solid_velocities_wp, + bc_mask, + missing_mask, + omega, + timestep, + ): + self.s_lagr_forces = wp.zeros(shape=(s_lagr_vertices_wp.shape[0]), dtype=wp.vec3) + s_lagr_forces_prev = wp.zeros(shape=(s_lagr_vertices_wp.shape[0]), dtype=wp.vec3) + convergence_flag_host_wp = wp.zeros(1, dtype=wp.int32, device="cpu", pinned=True) + device = f_0.device + if device.is_cuda: + # Warp recommends zero-copy by aliasing pinned host memory instead of issuing explicit copies + convergence_flag = wp.array( + ptr=convergence_flag_host_wp.ptr, + dtype=wp.int32, + shape=convergence_flag_host_wp.shape, + strides=convergence_flag_host_wp.strides, + device=device, + ) + else: + convergence_flag = convergence_flag_host_wp + convergence_flag_host = convergence_flag_host_wp.numpy() + flag_pending = False # Tracks whether an async convergence check is outstanding + tolerance_sq = self.compute_dtype(self.ibm_tolerance * self.ibm_tolerance) + + with self.timer_context: + wp.launch(kernel=self.warp_kernel, dim=f_0.shape[1:], inputs=[f_0, f_1, bc_mask, missing_mask, omega, timestep]) + for iteration in range(self.ibm_max_iterations): + if flag_pending: + # Complete any in-flight convergence update before using the host flag + wp.synchronize_stream() + needs_more_iterations = bool(convergence_flag_host[0]) + flag_pending = False + if not needs_more_iterations: + break + + wp.copy(s_lagr_forces_prev, self.s_lagr_forces) + + self.f_eulerian_forces.zero_() + self.f_eulerian_weights.zero_() + + wp.launch( + kernel=self.improved_interpolate_force_to_eulerian, + dim=s_lagr_vertices_wp.shape[0], + inputs=[ + s_lagr_vertices_wp, + self.s_lagr_forces, + lagr_solid_vertex_areas_wp, + self.f_eulerian_points, + self.f_eulerian_forces, + self.f_eulerian_weights, + wp.uint64(self.hash_grid.id), + ], + ) + + wp.launch( + kernel=self.compute_velocity_and_correct, + dim=f_1.shape[1:], + inputs=[f_1, self.f_eulerian_forces, self.f_eulerian_weights, self.f_eulerian_velocities], + ) + + compute_residual_flag = 1 if (iteration > 0 and self.ibm_tolerance > 0) else 0 + if compute_residual_flag: + convergence_flag.zero_() + + wp.launch( + kernel=self.interpolate_velocity_and_update_force, + dim=s_lagr_vertices_wp.shape[0], + inputs=[ + s_lagr_vertices_wp, + self.f_eulerian_points, + self.f_eulerian_velocities, + lagr_solid_velocities_wp, + self.s_lagr_forces, + s_lagr_forces_prev, + convergence_flag, + compute_residual_flag, + tolerance_sq, + wp.uint64(self.hash_grid.id), + ], + ) + + if compute_residual_flag: + # Device flag will be read at the start of the next iteration + flag_pending = True + + if flag_pending: + wp.synchronize_stream() + + wp.launch(kernel=self.correct_population_ibm, dim=f_1.shape[1:], inputs=[f_1, self.f_eulerian_forces]) + + return f_0, f_1, self.s_lagr_forces diff --git a/xlb/operator/stepper/nse_multires_stepper.py b/xlb/operator/stepper/nse_multires_stepper.py new file mode 100644 index 00000000..a9ee8b30 --- /dev/null +++ b/xlb/operator/stepper/nse_multires_stepper.py @@ -0,0 +1,1193 @@ +""" +Multi-Resolution Navier-Stokes Stepper for the NEON Backend + +This module implements the multi-resolution LBM stepper using Warp kernels on the +Neon multi-GPU runtime. It uses several programming patterns specific to Warp's +compile-time code generation model. + +Compile-Time Specialization Pattern +----------------------------------- +Warp's @wp.func decorator traces Python code at kernel compilation time, not runtime. +This means runtime boolean parameters cause Warp to emit branching code for both paths, +increasing register pressure even when only one path is ever taken. + +To generate optimized, branch-free kernels, we use a **factory pattern** that captures +boolean configuration at function-definition time: + + def make_specialized_func(do_feature: bool): + @wp.func + def impl(...): + if wp.static(do_feature): # Evaluated at compile time + # This code is only emitted when do_feature=True + ... + else: + # This code is only emitted when do_feature=False + ... + return impl + + # Generate specialized variants + func_with_feature = make_specialized_func(do_feature=True) + func_without_feature = make_specialized_func(do_feature=False) + +The `wp.static()` call evaluates its argument during Warp's tracing phase. Since +`do_feature` is a Python bool captured in the closure, Warp sees a constant and +eliminates the dead branch entirely. + +This pattern is used for: +- `apply_bc_post_streaming` / `apply_bc_post_collision`: Specialized BC application + for streaming vs collision implementation steps +- `collide_bc_accum` / `collide_simple`: Collision pipeline variants with/without + BC application and multi-resolution accumulation + +Closure Capture for Self Attributes +----------------------------------- +Warp cannot resolve `self.X` in plain assignments inside @wp.func bodies (e.g., +`_c = self.velocity_set.c` fails with "Invalid external reference type"). However, +it can resolve `self.X` in: +- Function call contexts: `self.stream.neon_functional(...)` +- Range arguments: `range(self.velocity_set.q)` +- Type casts: `self.compute_dtype(0)` + +For other uses, we pre-capture attributes at the Python level before defining the +@wp.func, making them available as simple closure variables: + + _c = self.velocity_set.c # Captured in Python scope + + @wp.func + def my_kernel(...): + # Use _c directly β€” Warp sees it as a closure variable + direction = wp.neon_ngh_idx(wp.int8(_c[0, l]), ...) + +Cell Type Constants +------------------- +Cell types are defined in `xlb.cell_type`: +- BC_SFV (254): Simple Fluid Voxel β€” no BC, no explosion/coalescence +- BC_SOLID (255): Solid obstacle voxel +- BC_NONE (0): Regular fluid voxel with potential BCs or multi-res interactions +""" + +import nvtx +import warp as wp +from typing import Any + +from xlb import DefaultConfig +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import Precision +from xlb.operator import Operator +from xlb.operator.stream import Stream +from xlb.operator.collision import BGK, KBC, SmagorinskyLESBGK +from xlb.operator.equilibrium import MultiresQuadraticEquilibrium +from xlb.operator.macroscopic import MultiresMacroscopic +from xlb.operator.stepper import Stepper +from xlb.operator.boundary_condition.boundary_condition import ImplementationStep +from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry +from xlb.operator.collision import ForcedCollision +from xlb.helper import check_bc_overlaps +from xlb.operator.boundary_masker import ( + MeshVoxelizationMethod, + MultiresMeshMaskerAABB, + MultiresMeshMaskerAABBClose, + MultiresIndicesBoundaryMasker, + MultiresMeshMaskerRay, +) +from xlb.operator.boundary_condition.helper_functions_bc import MultiresEncodeAuxiliaryData +from xlb.cell_type import BC_SFV, BC_SOLID + +""" +SFV = Simple Fluid Voxel: a fluid voxel that is not a BC nor is involved in explosion or coalescence +CFV = Complex Fluid Voxel: a fluid voxel that is not a SFV +""" + + +class MultiresIncompressibleNavierStokesStepper(Stepper): + """Multi-resolution incompressible Navier-Stokes stepper for the Neon backend. + + Implements the full LBM step (stream, collide, boundary conditions) across + a hierarchy of grid levels using Neon containers. Each container is a + compile-time specialized Warp kernel wrapped in a Neon execution-graph + node. + + The stepper supports several performance optimization strategies (see + :class:`MresPerfOptimizationType`): + + * **NAIVE_COLLIDE_STREAM** β€” separate collide and stream containers at + every level. + * **FUSION_AT_FINEST** β€” fused stream+collide at the finest level. + * **FUSION_AT_FINEST_SFV** β€” additionally splits SFV / CFV voxels at + the finest level for reduced branching. + * **FUSION_AT_FINEST_SFV_ALL** β€” SFV / CFV splitting at all levels. + + Parameters + ---------- + grid : NeonMultiresGrid + The multi-resolution grid. + boundary_conditions : list of BoundaryCondition + Boundary conditions to apply. + collision_type : str + Collision operator type: ``"BGK"`` or ``"KBC"`` or ``"SmagorinskyLESBGK"``. + forcing_scheme : str + Forcing scheme name (only used when *force_vector* is given). + force_vector : array-like, optional + External body force vector. + """ + + def __init__( + self, + grid, + boundary_conditions=[], + collision_type="BGK", + forcing_scheme="exact_difference", + force_vector=None, + ): + super().__init__(grid, boundary_conditions) + + # Construct the collision operator + if collision_type == "BGK": + self.collision = BGK(self.velocity_set, self.precision_policy, self.compute_backend) + elif collision_type == "KBC": + self.collision = KBC(self.velocity_set, self.precision_policy, self.compute_backend) + elif collision_type == "SmagorinskyLESBGK": + self.collision = SmagorinskyLESBGK(self.velocity_set, self.precision_policy, self.compute_backend) + + if force_vector is not None: + self.collision = ForcedCollision(collision_operator=self.collision, forcing_scheme=forcing_scheme, force_vector=force_vector) + + # Construct the operators + self.stream = Stream(self.velocity_set, self.precision_policy, self.compute_backend) + self.equilibrium = MultiresQuadraticEquilibrium(self.velocity_set, self.precision_policy, self.compute_backend) + self.macroscopic = MultiresMacroscopic(self.velocity_set, self.precision_policy, self.compute_backend) + + def prepare_fields(self, rho, u, initializer=None): + import neon + + """Prepare the fields required for the stepper. + + Args: + initializer: Optional operator to initialize the distribution functions. + If provided, it should be a callable that takes (grid, velocity_set, + precision_policy, compute_backend) as arguments and returns initialized f_0. + If None, default equilibrium initialization is used with rho=1 and u=0. + + Returns: + Tuple of (f_0, f_1, bc_mask, missing_mask): + - f_0: Initial distribution functions + - f_1: Copy of f_0 for double-buffering + - bc_mask: Boundary condition mask indicating which BC applies to each node + - missing_mask: Mask indicating which populations are missing at boundary nodes + """ + + f_0 = self.grid.create_field( + cardinality=self.velocity_set.q, dtype=self.precision_policy.store_precision, neon_memory_type=neon.MemoryType.device() + ) + + f_1 = self.grid.create_field( + cardinality=self.velocity_set.q, dtype=self.precision_policy.store_precision, neon_memory_type=neon.MemoryType.device() + ) + + missing_mask = self.grid.create_field(cardinality=self.velocity_set.q, dtype=Precision.UINT8) + bc_mask = self.grid.create_field(cardinality=1, dtype=Precision.UINT8) + + for level in range(self.grid.count_levels): + f_1.copy_from_run(level, f_0, 0) + + # Process boundary conditions and update masks + f_1, bc_mask, missing_mask = self._process_boundary_conditions(self.boundary_conditions, f_1, bc_mask, missing_mask) + # Initialize auxiliary data if needed + f_1 = self._initialize_auxiliary_data(self.boundary_conditions, f_1, bc_mask, missing_mask) + + # Initialize distribution functions if initializer is provided + if initializer is not None: + # Refer to xlb.helper.initializers for available initializers + f_0 = initializer(bc_mask, f_0) + else: + from xlb.helper.initializers import initialize_multires_eq + + f_0 = initialize_multires_eq(f_0, self.grid, self.velocity_set, self.precision_policy, self.compute_backend, rho=rho, u=u) + + return f_0, f_1, bc_mask, missing_mask + + def prepare_coalescence_count(self, coalescence_factor, bc_mask): + """Precompute coalescence weighting factors for multi-resolution streaming. + + For each non-halo voxel at every level, this method accumulates + the number of finer neighbours that contribute populations via + coalescence (child-to-parent transfer), then inverts the count + so that the streaming kernel can apply the correct averaging weight. + + Parameters + ---------- + coalescence_factor : field + Multi-resolution field to store the per-direction coalescence + weights (modified in-place). + bc_mask : field + Boundary-condition mask used to skip solid voxels. + """ + import neon + + lattice_central_index = self.velocity_set.center_index + num_levels = coalescence_factor.get_grid().num_levels + + @neon.Container.factory(name="sum_kernel_by_level") + def sum_kernel_by_level(level): + def ll_coalescence_count(loader: neon.Loader): + loader.set_mres_grid(coalescence_factor.get_grid(), level) + + coalescence_factor_pn = loader.get_mres_read_handle(coalescence_factor) + bc_mask_pn = loader.get_mres_read_handle(bc_mask) + + _c = self.velocity_set.c + _w = self.velocity_set.w + + @wp.func + def cl_collide_coarse(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if not wp.neon_has_child(coalescence_factor_pn, index): + for l in range(self.velocity_set.q): + if level < num_levels - 1: + push_direction = wp.neon_ngh_idx(wp.int8(_c[0, l]), wp.int8(_c[1, l]), wp.int8(_c[2, l])) + val = self.store_dtype(1) + wp.neon_mres_lbm_store_op(coalescence_factor_pn, index, l, push_direction, val) + + loader.declare_kernel(cl_collide_coarse) + + return ll_coalescence_count + + for level in range(num_levels): + sum_kernel = sum_kernel_by_level(level) + sum_kernel.run(0) + + @neon.Container.factory(name="sum_kernel_by_level") + def invert_count(level): + def loading(loader: neon.Loader): + loader.set_mres_grid(coalescence_factor.get_grid(), level) + + coalescence_factor_pn = loader.get_mres_read_handle(coalescence_factor) + bc_mask_pn = loader.get_mres_read_handle(bc_mask) + + _c = self.velocity_set.c + _w = self.velocity_set.w + + @wp.func + def compute(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + + if wp.neon_has_child(coalescence_factor_pn, index): + # we are a halo cell so we just exit + return + + for l in range(self.velocity_set.q): + if l == lattice_central_index: + continue + + pull_direction = wp.neon_ngh_idx(wp.int8(-_c[0, l]), wp.int8(-_c[1, l]), wp.int8(-_c[2, l])) + + has_ngh_at_same_level = wp.bool(False) + coalescence_factor = self.compute_dtype( + wp.neon_read_ngh(coalescence_factor_pn, index, pull_direction, l, self.store_dtype(0), has_ngh_at_same_level) + ) + + if not wp.neon_has_finer_ngh(coalescence_factor_pn, index, pull_direction): + pass + else: + # Finer neighbour exists in the pull direction (opposite of l). + # Read from the halo sitting on top of that finer neighbour. + if has_ngh_at_same_level: + # Finer ngh in pull direction: YES + # Same-level ngh: YES + # Compute coalescence factor + if coalescence_factor > self.compute_dtype(0): + coalescence_factor = self.compute_dtype(1) / (self.compute_dtype(2) * coalescence_factor) + wp.neon_write(coalescence_factor_pn, index, l, self.store_dtype(coalescence_factor)) + + loader.declare_kernel(compute) + + return loading + + for level in range(num_levels): + sum_kernel = invert_count(level) + sum_kernel.run(0) + return + + @classmethod + def _process_boundary_conditions(cls, boundary_conditions, f_1, bc_mask, missing_mask): + """Process boundary conditions and update boundary masks.""" + + # Check for boundary condition overlaps + # TODO! check_bc_overlaps(boundary_conditions, DefaultConfig.velocity_set.d, DefaultConfig.default_backend) + + # Create boundary maskers + indices_masker = MultiresIndicesBoundaryMasker( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + + # Split boundary conditions by type + bc_with_vertices = [bc for bc in boundary_conditions if bc.mesh_vertices is not None] + bc_with_indices = [bc for bc in boundary_conditions if bc.indices is not None] + + # Process indices-based boundary conditions + if bc_with_indices: + bc_mask, missing_mask = indices_masker(bc_with_indices, bc_mask, missing_mask) + + # Process mesh-based boundary conditions for 3D + if DefaultConfig.velocity_set.d == 3 and bc_with_vertices: + for bc in bc_with_vertices: + if bc.voxelization_method.id is MeshVoxelizationMethod("AABB").id: + mesh_masker = MultiresMeshMaskerAABB( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + elif bc.voxelization_method.id is MeshVoxelizationMethod("RAY").id: + mesh_masker = MultiresMeshMaskerRay( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + elif bc.voxelization_method.id is MeshVoxelizationMethod("AABB_CLOSE").id: + mesh_masker = MultiresMeshMaskerAABBClose( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + close_voxels=bc.voxelization_method.options.get("close_voxels"), + ) + else: + raise ValueError(f"Unsupported voxelization method for multi-res: {bc.voxelization_method}") + # Apply the mesh masker to the boundary condition + f_1, bc_mask, missing_mask = mesh_masker(bc, f_1, bc_mask, missing_mask) + + return f_1, bc_mask, missing_mask + + @staticmethod + def _initialize_auxiliary_data(boundary_conditions, f_1, bc_mask, missing_mask): + """Initialize auxiliary data for boundary conditions that require it.""" + for bc in boundary_conditions: + if bc.needs_aux_init and not bc.is_initialized_with_aux_data: + # Create the encoder operator for storing the auxiliary data + encode_auxiliary_data = MultiresEncodeAuxiliaryData( + bc.id, + bc.num_of_aux_data, + bc.profile, + velocity_set=bc.velocity_set, + precision_policy=bc.precision_policy, + compute_backend=bc.compute_backend, + ) + + # Encode the auxiliary data in f_1 + f_1 = encode_auxiliary_data(f_1, bc_mask, missing_mask, stream=0) + bc.is_initialized_with_aux_data = True + return f_1 + + def _construct_neon(self): + import neon + + # Pre-capture self attributes that Warp cannot resolve inside @wp.func bodies. + # Warp rejects `self` as an "Invalid external reference type" when it appears + # in a plain assignment (e.g. `_c = self.velocity_set.c`). Capturing here + # makes these values available as simple closure variables. + lattice_central_index = self.velocity_set.center_index + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8) + _opp_indices = self.velocity_set.opp_indices + _c = self.velocity_set.c + + # Read the list of bc_to_id created upon instantiation + bc_to_id = boundary_condition_registry.bc_to_id + + # Gather IDs of ExtrapolationOutflowBC boundary conditions + extrapolation_outflow_bc_ids = [] + for bc_name, bc_id in bc_to_id.items(): + if bc_name.startswith("ExtrapolationOutflowBC"): + extrapolation_outflow_bc_ids.append(bc_id) + + # Factory for apply_bc: generates compile-time specialized variants + def make_apply_bc(is_post_streaming: bool): + @wp.func + def apply_bc_impl( + index: Any, + timestep: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + ): + f_result = f_post + + for i in range(wp.static(len(self.boundary_conditions))): + if wp.static(is_post_streaming): + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.STREAMING): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].neon_functional)( + index, timestep, _missing_mask, f_0, f_1, f_pre, f_post + ) + else: + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.COLLISION): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].neon_functional)( + index, timestep, _missing_mask, f_0, f_1, f_pre, f_post + ) + if wp.static(self.boundary_conditions[i].id in extrapolation_outflow_bc_ids): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].assemble_auxiliary_data)( + index, timestep, _missing_mask, f_0, f_1, f_pre, f_post + ) + return f_result + + return apply_bc_impl + + # Compile-time specialized BC application variants + apply_bc_post_streaming = make_apply_bc(is_post_streaming=True) + apply_bc_post_collision = make_apply_bc(is_post_streaming=False) + + @wp.func + def neon_get_thread_data( + f0_pn: Any, + missing_mask_pn: Any, + index: Any, + ): + # Read thread data for populations + _f0_thread = _f_vec() + _missing_mask = _missing_mask_vec() + for l in range(self.velocity_set.q): + # q-sized vector of pre-streaming populations + _f0_thread[l] = self.compute_dtype(wp.neon_read(f0_pn, index, l)) + _missing_mask[l] = wp.neon_read(missing_mask_pn, index, l) + + return _f0_thread, _missing_mask + + @wp.func + def neon_apply_aux_recovery_bc( + index: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0_pn: Any, + f_1_pn: Any, + ): + # Note: + # In XLB, the BC auxiliary data (e.g. prescribed values of pressure or normal velocity) are stored in (i) central index of f_1 and/or + # (ii) missing directions of f_1. Some BCs may or may not need all these available storage space. This function checks whether + # the BC needs recovery of auxiliary data and then recovers the information for the next iteration (due to buffer swapping) by + # writting the values of f_1 into f_0. + + # Unroll the loop over boundary conditions + for i in range(wp.static(len(self.boundary_conditions))): + if wp.static(self.boundary_conditions[i].needs_aux_recovery): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + for l in range(self.velocity_set.q): + # Perform the swapping of data + if l == lattice_central_index: + # (i) Recover the values stored in the central index of f_1 + _f1_thread = wp.neon_read(f_1_pn, index, l) + wp.neon_write(f_0_pn, index, l, self.store_dtype(_f1_thread)) + elif _missing_mask[l] == wp.uint8(1): + # (ii) Recover the values stored in the missing directions of f_1 + _f1_thread = wp.neon_read(f_1_pn, index, _opp_indices[l]) + wp.neon_write(f_0_pn, index, _opp_indices[l], self.store_dtype(_f1_thread)) + + # Factory for neon_collide_pipeline: generates compile-time specialized variants + def make_collide_pipeline(do_bc: bool, do_accumulation: bool): + @wp.func + def collide_pipeline_impl( + index: Any, + timestep: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0_pn: Any, + f_1_pn: Any, + _f_post_stream: Any, + omega: Any, + num_levels: int, + level: int, + accumulation_pn: Any, + ): + _rho, _u = self.macroscopic.neon_functional(_f_post_stream) + _feq = self.equilibrium.neon_functional(_rho, _u) + _f_post_collision = self.collision.neon_functional(_f_post_stream, _feq, omega) + + if wp.static(do_bc): + _f_post_collision = apply_bc_post_collision( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_stream, _f_post_collision + ) + neon_apply_aux_recovery_bc(index, _boundary_id, _missing_mask, f_0_pn, f_1_pn) + + if wp.static(do_accumulation): + for l in range(self.velocity_set.q): + push_direction = wp.neon_ngh_idx(wp.int8(_c[0, l]), wp.int8(_c[1, l]), wp.int8(_c[2, l])) + if level < num_levels - 1: + wp.neon_mres_lbm_store_op(accumulation_pn, index, l, push_direction, self.store_dtype(_f_post_collision[l])) + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_collision[l])) + else: + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_collision[l])) + + return _f_post_collision + + return collide_pipeline_impl + + # Compile-time specialized collision pipeline variants + collide_bc_accum = make_collide_pipeline(do_bc=True, do_accumulation=True) + collide_bc_only = make_collide_pipeline(do_bc=True, do_accumulation=False) + collide_simple = make_collide_pipeline(do_bc=False, do_accumulation=False) + + @wp.func + def neon_stream_explode_coalesce( + index: Any, + f_0_pn: Any, + coalescence_factor_pn: Any, + ): + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + + for l in range(self.velocity_set.q): + if l == lattice_central_index: + continue + + pull_direction = wp.neon_ngh_idx(wp.int8(-_c[0, l]), wp.int8(-_c[1, l]), wp.int8(-_c[2, l])) + + has_ngh_at_same_level = wp.bool(False) + accumulated = wp.neon_read_ngh(f_0_pn, index, pull_direction, l, self.store_dtype(0), has_ngh_at_same_level) + + if not wp.neon_has_finer_ngh(f_0_pn, index, pull_direction): + # No finer ngh in the pull direction (opposite of l) + if not has_ngh_at_same_level: + # No same-level ngh β€” could we have a coarser-level ngh? + if wp.neon_has_parent(f_0_pn, index): + # Halo cell on top of us (parent exists) + has_a_coarser_ngh = wp.bool(False) + exploded_pop = wp.neon_lbm_read_coarser_ngh(f_0_pn, index, pull_direction, l, self.store_dtype(0), has_a_coarser_ngh) + if has_a_coarser_ngh: + # No finer ngh in pull direction, no same-level ngh, + # but a parent (ghost cell) exists with a coarser ngh + # -> Explosion: read the exploded population from the + # coarser level's halo. + _f_post_stream[l] = self.compute_dtype(exploded_pop) + else: + # Finer ngh exists in the pull direction (opposite of l). + # Read from the halo on top of that finer ngh. + if has_ngh_at_same_level: + # Finer ngh in pull direction: YES + # Same-level ngh: YES + # -> Coalescence + coalescence_factor = wp.neon_read(coalescence_factor_pn, index, l) + accumulated = accumulated * coalescence_factor + _f_post_stream[l] = self.compute_dtype(accumulated) + + return _f_post_stream + + @neon.Container.factory(name="collide_coarse") + def collide_coarse(level: int, f_0_fd: Any, f_1_fd: Any, bc_mask_fd: Any, missing_mask_fd: Any, omega: Any, timestep: int): + num_levels = f_0_fd.get_grid().num_levels + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + if level + 1 < f_0_fd.get_grid().num_levels: + f_0_pn = loader.get_mres_write_handle(f_0_fd, neon.Loader.Operation.stencil_up) + f_1_pn = loader.get_mres_write_handle(f_1_fd, neon.Loader.Operation.stencil_up) + else: + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if not wp.neon_has_child(f_0_pn, index): + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + collide_bc_accum( + index, + timestep, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f0_thread, + omega, + num_levels, + level, + f_1_pn, + ) + else: + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(0)) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="SFV_collide_coarse") + def SFV_collide_coarse(level: int, f_0_fd: Any, f_1_fd: Any, bc_mask_fd: Any, missing_mask_fd: Any, omega: Any, timestep: int): + """Collision on SFV voxels only β€” no BCs, no multi-resolution accumulation.""" + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id != wp.uint8(BC_SFV): + return + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + collide_simple( + index, + 0, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f0_thread, + omega, + 0, + level, + f_1_pn, + ) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="CFV_collide_coarse") + def CFV_collide_coarse(level: int, f_0_fd: Any, f_1_fd: Any, bc_mask_fd: Any, missing_mask_fd: Any, omega: Any, timestep: int): + """Collision on CFV voxels only β€” skips both solid and SFV.""" + num_levels = f_0_fd.get_grid().num_levels + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + if level + 1 < f_0_fd.get_grid().num_levels: + f_0_pn = loader.get_mres_write_handle(f_0_fd, neon.Loader.Operation.stencil_up) + f_1_pn = loader.get_mres_write_handle(f_1_fd, neon.Loader.Operation.stencil_up) + else: + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if _boundary_id == wp.uint8(BC_SFV): + return + if not wp.neon_has_child(f_0_pn, index): + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + collide_bc_accum( + index, + timestep, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f0_thread, + omega, + num_levels, + level, + f_1_pn, + ) + else: + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(0)) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="stream_coarse_step_ABC") + def stream_coarse_step_ABC( + level: int, + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + coalescence_factor: Any, + timestep: int, + ): + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + coalescence_factor_pn = loader.get_mres_read_handle(coalescence_factor) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if wp.neon_has_child(f_0_pn, index): + return + + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = neon_stream_explode_coalesce(index, f_0_pn, coalescence_factor_pn) + + _f_post_stream = apply_bc_post_streaming( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_collision, _f_post_stream + ) + neon_apply_aux_recovery_bc(index, _boundary_id, _missing_mask, f_0_pn, f_1_pn) + + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_stream[l])) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="SFV_stream_coarse_step_ABC") + def SFV_stream_coarse_step_ABC( + level: int, + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + coalescence_factor: Any, + timestep: int, + ): + """Stream on CFV voxels only β€” skips SFV and solid.""" + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + coalescence_factor_pn = loader.get_mres_read_handle(coalescence_factor) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SFV): + return + if _boundary_id == wp.uint8(BC_SOLID): + return + if wp.neon_has_child(f_0_pn, index): + return + + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = neon_stream_explode_coalesce(index, f_0_pn, coalescence_factor_pn) + + _f_post_stream = apply_bc_post_streaming( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_collision, _f_post_stream + ) + neon_apply_aux_recovery_bc(index, _boundary_id, _missing_mask, f_0_pn, f_1_pn) + + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_stream[l])) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="SFV_reset_bc_mask") + def SFV_reset_bc_mask( + level: int, + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + ): + """ + Setting the BC type to BC_SFV + """ + + def ll_stream_coarse(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + + f_0_pn = loader.get_mres_read_handle(f_0_fd) + + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + _c = self.velocity_set.c + + @wp.func + def cl_stream_coarse(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if _boundary_id != 0: + return + + if wp.neon_has_child(f_0_pn, index): + # we are a halo cell so we just exit + return + + # do stream normally + _missing_mask = _missing_mask_vec() + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + + for l in range(self.velocity_set.q): + if l == lattice_central_index: + continue + + pull_direction = wp.neon_ngh_idx(wp.int8(-_c[0, l]), wp.int8(-_c[1, l]), wp.int8(-_c[2, l])) + + has_ngh_at_same_level = wp.bool(False) + wp.neon_read_ngh(f_0_pn, index, pull_direction, l, self.store_dtype(0), has_ngh_at_same_level) + + if not wp.neon_has_finer_ngh(f_0_pn, index, pull_direction): + if not has_ngh_at_same_level: + if wp.neon_has_parent(f_0_pn, index): + has_a_coarser_ngh = wp.bool(False) + wp.neon_lbm_read_coarser_ngh(f_0_pn, index, pull_direction, l, self.store_dtype(0), has_a_coarser_ngh) + if has_a_coarser_ngh: + # Explosion: not an SFV + return + else: + if has_ngh_at_same_level: + # Coalescence: not an SFV + return + + # Voxel is a pure fluid cell with no multi-resolution interactions β€” mark as SFV + wp.neon_write(bc_mask_pn, index, 0, wp.uint8(BC_SFV)) + + loader.declare_kernel(cl_stream_coarse) + + return ll_stream_coarse + + @neon.Container.factory(name="SFV_stream_coarse_step") + def SFV_stream_coarse_step(level: int, f_0_fd: Any, f_1_fd: Any, bc_mask_fd: Any, missing_mask_fd: Any): + def ll_stream_coarse(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + _c = self.velocity_set.c + + @wp.func + def cl_stream_coarse(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id != wp.uint8(BC_SFV): + return + # BC_SFV voxel type: + # - They are not BC voxels + # - They are not on a resolution jump -> they do not do coalescence or explosion + # - They are not mr halo cells + + _missing_mask = _missing_mask_vec() + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_stream[l])) + + loader.declare_kernel(cl_stream_coarse) + + return ll_stream_coarse + + @wp.func + def neon_stream_finest_with_explosion( + index: Any, + f_0_pn: Any, + explosion_src_pn: Any, + ): + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + + for l in range(self.velocity_set.q): + if l == lattice_central_index: + continue + + pull_direction = wp.neon_ngh_idx(wp.int8(-_c[0, l]), wp.int8(-_c[1, l]), wp.int8(-_c[2, l])) + + has_ngh_at_same_level = wp.bool(False) + wp.neon_read_ngh(f_0_pn, index, pull_direction, l, self.store_dtype(0), has_ngh_at_same_level) + + if not has_ngh_at_same_level: + # No same-level ngh β€” could we have a coarser-level ngh? + if wp.neon_has_parent(f_0_pn, index): + # Parent exists β€” try to read the exploded population from the coarser level + has_a_coarser_ngh = wp.bool(False) + exploded_pop = wp.neon_lbm_read_coarser_ngh( + explosion_src_pn, index, pull_direction, l, self.store_dtype(0), has_a_coarser_ngh + ) + if has_a_coarser_ngh: + # No finer ngh in pull direction, no same-level ngh, + # but a parent (ghost cell) exists with a coarser ngh + # -> Explosion: read the exploded population from the + # coarser level's halo. + _f_post_stream[l] = self.compute_dtype(exploded_pop) + + return _f_post_stream + + @neon.Container.factory(name="finest_fused_pull") + def finest_fused_pull( + level: int, + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + omega: Any, + timestep: Any, + is_f1_the_explosion_src_field: bool, + ): + if level != 0: + raise Exception("Only the finest level is supported for now") + num_levels = f_0_fd.get_grid().num_levels + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + if level + 1 < f_0_fd.get_grid().num_levels: + f_0_pn = loader.get_mres_write_handle(f_0_fd, neon.Loader.Operation.stencil_up) + f_1_pn = loader.get_mres_write_handle(f_1_fd, neon.Loader.Operation.stencil_up) + else: + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + explosion_src_pn = f_1_pn if is_f1_the_explosion_src_field else f_0_pn + accumulation_pn = f_1_pn if is_f1_the_explosion_src_field else f_0_pn + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if wp.neon_has_child(f_0_pn, index): + return + + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = neon_stream_finest_with_explosion(index, f_0_pn, explosion_src_pn) + + _f_post_stream = apply_bc_post_streaming( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_collision, _f_post_stream + ) + + collide_bc_accum( + index, + timestep, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f_post_stream, + omega, + num_levels, + level, + accumulation_pn, + ) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="CFV_finest_fused_pull") + def CFV_finest_fused_pull( + level: int, + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + omega: Any, + timestep: Any, + is_f1_the_explosion_src_field: bool, + ): + """Fused stream+collide on CFV voxels at the finest level β€” skips SFV and solid.""" + if level != 0: + raise Exception("Only the finest level is supported for now") + num_levels = f_0_fd.get_grid().num_levels + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + if level + 1 < f_0_fd.get_grid().num_levels: + f_0_pn = loader.get_mres_write_handle(f_0_fd, neon.Loader.Operation.stencil_up) + f_1_pn = loader.get_mres_write_handle(f_1_fd, neon.Loader.Operation.stencil_up) + else: + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + explosion_src_pn = f_1_pn if is_f1_the_explosion_src_field else f_0_pn + accumulation_pn = f_1_pn if is_f1_the_explosion_src_field else f_0_pn + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + if _boundary_id == wp.uint8(BC_SFV): + return + if wp.neon_has_child(f_0_pn, index): + return + + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + _f_post_stream = neon_stream_finest_with_explosion(index, f_0_pn, explosion_src_pn) + + _f_post_stream = apply_bc_post_streaming( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_collision, _f_post_stream + ) + + collide_bc_accum( + index, + timestep, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f_post_stream, + omega, + num_levels, + level, + accumulation_pn, + ) + + loader.declare_kernel(device) + + return ll + + @neon.Container.factory(name="SFV_finest_fused_pull") + def SFV_finest_fused_pull(level: int, f_0_fd: Any, f_1_fd: Any, bc_mask_fd: Any, missing_mask_fd: Any, omega: Any): + """Fused stream+collide on SFV voxels at the finest level β€” no BCs, no explosion.""" + if level != 0: + raise Exception("Only the finest level is supported for now") + + def ll(loader: neon.Loader): + loader.set_mres_grid(bc_mask_fd.get_grid(), level) + f_0_pn = loader.get_mres_read_handle(f_0_fd) + f_1_pn = loader.get_mres_write_handle(f_1_fd) + bc_mask_pn = loader.get_mres_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_mres_read_handle(missing_mask_fd) + + @wp.func + def device(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id != wp.uint8(BC_SFV): + return + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + collide_simple( + index, + 0, + _boundary_id, + _missing_mask, + f_0_pn, + f_1_pn, + _f_post_stream, + omega, + 0, + 0, + f_1_pn, + ) + + loader.declare_kernel(device) + + return ll + + return None, { + "collide_coarse": collide_coarse, + "stream_coarse_step_ABC": stream_coarse_step_ABC, + "finest_fused_pull": finest_fused_pull, + "CFV_finest_fused_pull": CFV_finest_fused_pull, + "SFV_finest_fused_pull": SFV_finest_fused_pull, + "SFV_reset_bc_mask": SFV_reset_bc_mask, + "CFV_collide_coarse": CFV_collide_coarse, + "SFV_collide_coarse": SFV_collide_coarse, + "SFV_stream_coarse_step_ABC": SFV_stream_coarse_step_ABC, + "SFV_stream_coarse_step": SFV_stream_coarse_step, + } + + def add_to_app(self, **kwargs): + """Append a container invocation to the Neon skeleton application list. + + Required keyword arguments are ``op_name`` (str) and ``app`` (list). + All remaining keyword arguments are forwarded to the container + factory for the given ``op_name``. Argument validation is performed + before the call, and a ``ValueError`` is raised on mismatch. + """ + import inspect + + def validate_kwargs_forward(func, kwargs): + """ + Check whether `func(**kwargs)` would be valid, + and return *all* the issues instead of raising on the first one. + + Returns a dict; empty dict means "everything is OK". + """ + sig = inspect.signature(func) + params = sig.parameters + + errors = {} + + # --- 1. Positional-only required params (cannot be given via kwargs) --- + pos_only_required = [name for name, p in params.items() if p.kind == inspect.Parameter.POSITIONAL_ONLY and p.default is inspect._empty] + if pos_only_required: + errors["positional_only_required"] = pos_only_required + + # --- 2. Unexpected kwargs (if no **kwargs in target) --- + has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + if not has_var_kw: + allowed_kw = { + name + for name, p in params.items() + if p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + unexpected = sorted(set(kwargs) - allowed_kw) + if unexpected: + errors["unexpected_kwargs"] = unexpected + + # --- 3. Missing required keyword-passable params --- + missing_required = [ + name + for name, p in params.items() + if p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + and p.default is inspect._empty # no default + and name not in kwargs # not provided + ] + if missing_required: + errors["missing_required"] = missing_required + + return errors + + container_generator = None + try: + op_name = kwargs.pop("op_name") + app = kwargs.pop("app") + except KeyError: + raise ValueError("op_name and app must be provided as keyword arguments") + + try: + container_generator = self.neon_container[op_name] + except KeyError: + raise ValueError(f"Operator {op_name} not found in neon container. Available operators: {list(self.neon_container.keys())}") + + errors = validate_kwargs_forward(container_generator, kwargs) + if errors: + raise ValueError(f"Cannot forward kwargs to target: {errors}") + + nvtx.push_range(f"New Container {op_name}", color="yellow") + app.append(container_generator(**kwargs)) + nvtx.pop_range() + + @Operator.register_backend(ComputeBackend.NEON) + def neon_launch(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + raise NotImplementedError("Use MultiresSimulationManager.step() instead of launching this stepper directly.") diff --git a/xlb/operator/stepper/nse_stepper.py b/xlb/operator/stepper/nse_stepper.py new file mode 100644 index 00000000..ee69750e --- /dev/null +++ b/xlb/operator/stepper/nse_stepper.py @@ -0,0 +1,663 @@ +""" +Single-resolution incompressible Navier-Stokes stepper. + +Implements the full LBM step (stream, collide, apply BCs) for a single- +resolution grid. Supports pull and push streaming schemes on JAX, a +pull-only fused kernel on Warp, and a pull-only Neon container. +""" + +from functools import partial + +from jax import jit +import warp as wp +from typing import Any + +from xlb import DefaultConfig +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import Precision +from xlb.operator import Operator +from xlb.operator.stream import Stream +from xlb.operator.collision import BGK, KBC, SmagorinskyLESBGK +from xlb.operator.equilibrium import QuadraticEquilibrium +from xlb.operator.macroscopic import Macroscopic +from xlb.operator.stepper import Stepper +from xlb.operator.boundary_condition.boundary_condition import ImplementationStep +from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry +from xlb.operator.collision import ForcedCollision +from xlb.operator.boundary_masker import ( + IndicesBoundaryMasker, + MeshVoxelizationMethod, + MeshMaskerAABB, + MeshMaskerRay, + MeshMaskerWinding, + MeshMaskerAABBClose, +) +from xlb.helper import check_bc_overlaps +from xlb.helper.nse_fields import create_nse_fields +from xlb.operator.boundary_condition.helper_functions_bc import EncodeAuxiliaryData +from xlb.cell_type import BC_SOLID + + +class IncompressibleNavierStokesStepper(Stepper): + """Single-resolution incompressible Navier-Stokes LBM stepper. + + Composes streaming, collision, equilibrium, macroscopic, and boundary- + condition operators into a complete timestep. + + Parameters + ---------- + grid : Grid + Computational grid. + boundary_conditions : list of BoundaryCondition + Boundary conditions to apply each step. + collision_type : str + ``"BGK"``, ``"KBC"``, or ``"SmagorinskyLESBGK"``. + streaming_scheme : str + ``"pull"`` (default) or ``"push"`` (JAX only). + forcing_scheme : str + Forcing scheme name (used when *force_vector* is given). + force_vector : array-like, optional + External body force vector. + backend_config : dict + Backend-specific options (e.g. Neon OCC configuration). + """ + + def __init__( + self, + grid, + boundary_conditions=[], + collision_type="BGK", + streaming_scheme="pull", + forcing_scheme="exact_difference", + force_vector=None, + backend_config={}, + ): + super().__init__(grid, boundary_conditions) + self.backend_config = backend_config + + # Construct the collision operator + if collision_type == "BGK": + self.collision = BGK(self.velocity_set, self.precision_policy, self.compute_backend) + elif collision_type == "KBC": + self.collision = KBC(self.velocity_set, self.precision_policy, self.compute_backend) + elif collision_type == "SmagorinskyLESBGK": + self.collision = SmagorinskyLESBGK(self.velocity_set, self.precision_policy, self.compute_backend) + + if force_vector is not None: + self.collision = ForcedCollision(collision_operator=self.collision, forcing_scheme=forcing_scheme, force_vector=force_vector) + + # Choose the implementation based on backend and streaming scheme + self.streaming_scheme = streaming_scheme + if self.compute_backend != ComputeBackend.JAX: + assert streaming_scheme == "pull", f"Unknown or unimplemented streaming scheme for backend: {self.compute_backend}" + + # Construct the operators + self.stream = Stream(self.velocity_set, self.precision_policy, self.compute_backend) + self.equilibrium = QuadraticEquilibrium(self.velocity_set, self.precision_policy, self.compute_backend) + self.macroscopic = Macroscopic(self.velocity_set, self.precision_policy, self.compute_backend) + + def prepare_fields(self, initializer=None): + """Prepare the fields required for the stepper. + + Args: + initializer: Optional operator to initialize the distribution functions. + If provided, it should be a callable that takes (grid, velocity_set, + precision_policy, compute_backend) as arguments and returns initialized f_0. + If None, default equilibrium initialization is used with rho=1 and u=0. + + Returns: + Tuple of (f_0, f_1, bc_mask, missing_mask): + - f_0: Initial distribution functions + - f_1: Copy of f_0 for double-buffering + - bc_mask: Boundary condition mask indicating which BC applies to each node + - missing_mask: Mask indicating which populations are missing at boundary nodes + """ + # Create fields using the helper function + _, f_0, f_1, missing_mask, bc_mask = create_nse_fields( + grid=self.grid, velocity_set=self.velocity_set, compute_backend=self.compute_backend, precision_policy=self.precision_policy + ) + + # Copy f_0 using backend-specific copy to f_1 + if self.compute_backend == ComputeBackend.JAX: + f_1 = f_0.copy() + if self.compute_backend == ComputeBackend.WARP: + wp.copy(f_1, f_0) + if self.compute_backend == ComputeBackend.NEON: + f_1.copy_from_run(f_0, 0) + + # Important note: XLB uses f_1 buffer (center index and missing directions) to store auxiliary data for boundary conditions. + # Process boundary conditions and update masks + f_1, bc_mask, missing_mask = self._process_boundary_conditions(self.boundary_conditions, f_1, bc_mask, missing_mask) + + # Initialize auxiliary data if needed + f_1 = self._initialize_auxiliary_data(self.boundary_conditions, f_1, bc_mask, missing_mask) + # bc_mask.update_host(0) + # missing_mask.update_host(0) + wp.synchronize() + # bc_mask.export_vti("bc_mask.vti", 'bc_mask') + # missing_mask.export_vti("missing_mask.vti", 'missing_mask') + + # Initialize distribution functions if initializer is provided + if initializer is not None: + f_0 = initializer(bc_mask, f_0) + else: + from xlb.helper.initializers import initialize_eq + + f_0 = initialize_eq(f_0, self.grid, self.velocity_set, self.precision_policy, self.compute_backend) + + return f_0, f_1, bc_mask, missing_mask + + def _process_boundary_conditions(self, boundary_conditions, f_1, bc_mask, missing_mask): + """Process boundary conditions and update boundary masks.""" + + # Check for boundary condition overlaps + check_bc_overlaps(boundary_conditions, DefaultConfig.velocity_set.d, DefaultConfig.default_backend) + + # Create boundary maskers + indices_masker = IndicesBoundaryMasker( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + grid=self.grid, + ) + + # Split boundary conditions by type + bc_with_vertices = [bc for bc in boundary_conditions if bc.mesh_vertices is not None] + bc_with_indices = [bc for bc in boundary_conditions if bc.indices is not None] + + # Process indices-based boundary conditions + if bc_with_indices: + bc_mask, missing_mask = indices_masker(bc_with_indices, bc_mask, missing_mask) + + # Process mesh-based boundary conditions for 3D + if DefaultConfig.velocity_set.d == 3 and bc_with_vertices: + for bc in bc_with_vertices: + if bc.voxelization_method.id is MeshVoxelizationMethod("AABB").id: + mesh_masker = MeshMaskerAABB( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + elif bc.voxelization_method.id is MeshVoxelizationMethod("RAY").id: + mesh_masker = MeshMaskerRay( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + elif bc.voxelization_method.id is MeshVoxelizationMethod("WINDING").id: + mesh_masker = MeshMaskerWinding( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + ) + elif bc.voxelization_method.id is MeshVoxelizationMethod("AABB_CLOSE").id: + mesh_masker = MeshMaskerAABBClose( + velocity_set=DefaultConfig.velocity_set, + precision_policy=DefaultConfig.default_precision_policy, + compute_backend=DefaultConfig.default_backend, + close_voxels=bc.voxelization_method.options.get("close_voxels"), + ) + else: + raise ValueError(f"Unsupported voxelization method: {bc.voxelization_method}") + # Apply the mesh masker to the boundary condition + f_1, bc_mask, missing_mask = mesh_masker(bc, f_1, bc_mask, missing_mask) + + return f_1, bc_mask, missing_mask + + @staticmethod + def _initialize_auxiliary_data(boundary_conditions, f_1, bc_mask, missing_mask): + """Initialize auxiliary data for boundary conditions that require it.""" + for bc in boundary_conditions: + if bc.needs_aux_init and not bc.is_initialized_with_aux_data: + # Create the encoder operator for storing the auxiliary data + encode_auxiliary_data = EncodeAuxiliaryData( + bc.id, + bc.num_of_aux_data, + bc.profile, + velocity_set=bc.velocity_set, + precision_policy=bc.precision_policy, + compute_backend=bc.compute_backend, + ) + + # Encode the auxiliary data in f_1 + f_1 = encode_auxiliary_data(f_1, bc_mask, missing_mask) + bc.is_initialized_with_aux_data = True + return f_1 + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0,)) + def jax_implementation(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + if self.streaming_scheme == "pull": + return self.jax_implementation_pull(f_0, f_1, bc_mask, missing_mask, omega, timestep) + elif self.streaming_scheme == "push": + return self.jax_implementation_push(f_0, f_1, bc_mask, missing_mask, omega, timestep) + else: + raise ValueError(f"Unknown streaming scheme: {self.streaming_scheme}") + + @partial(jit, static_argnums=(0,)) + def jax_implementation_pull(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + """ + Perform a single step of the lattice boltzmann method + """ + # Cast to compute precision + f_0 = self.precision_policy.cast_to_compute_jax(f_0) + f_1 = self.precision_policy.cast_to_compute_jax(f_1) + + # Apply streaming + f_post_stream = self.stream(f_0) + + # Apply boundary conditions + for bc in self.boundary_conditions: + if bc.implementation_step == ImplementationStep.STREAMING: + f_post_stream = bc( + f_0, + f_post_stream, + bc_mask, + missing_mask, + ) + + # Compute the macroscopic variables + rho, u = self.macroscopic(f_post_stream) + + # Compute equilibrium + feq = self.equilibrium(rho, u) + + # Apply collision + f_post_collision = self.collision(f_post_stream, feq, omega) + + # Apply collision type boundary conditions + for bc in self.boundary_conditions: + f_post_collision = bc.assemble_auxiliary_data(f_post_stream, f_post_collision, bc_mask, missing_mask) + if bc.implementation_step == ImplementationStep.COLLISION: + f_post_collision = bc( + f_post_stream, + f_post_collision, + bc_mask, + missing_mask, + ) + + # Copy back to store precision + f_1 = self.precision_policy.cast_to_store_jax(f_post_collision) + + return f_0, f_1 + + @partial(jit, static_argnums=(0,)) + def jax_implementation_push(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + """ + Perform a single step of the lattice boltzmann method + """ + # Cast to compute precision + f_0 = self.precision_policy.cast_to_compute_jax(f_0) + f_1 = self.precision_policy.cast_to_compute_jax(f_1) + + # Assign f_post_stream + f_post_stream = f_0 + + # Compute the macroscopic variables + rho, u = self.macroscopic(f_post_stream) + + # Compute equilibrium + feq = self.equilibrium(rho, u) + + # Apply collision + f_post_collision = self.collision(f_post_stream, feq, omega) + + # Apply collision type boundary conditions + for bc in self.boundary_conditions: + f_post_collision = bc.update_bc_auxiliary_data(f_post_stream, f_post_collision, bc_mask, missing_mask) + if bc.implementation_step == ImplementationStep.COLLISION: + f_post_collision = bc( + f_post_stream, + f_post_collision, + bc_mask, + missing_mask, + ) + + # Apply streaming + f_post_stream = self.stream(f_post_collision) + + # Apply boundary conditions + for bc in self.boundary_conditions: + if bc.implementation_step == ImplementationStep.STREAMING: + f_post_stream = bc( + f_post_collision, + f_post_stream, + bc_mask, + missing_mask, + ) + + # Copy back to store precision + f_0 = self.precision_policy.cast_to_store_jax(f_post_collision) + f_1 = self.precision_policy.cast_to_store_jax(f_post_stream) + + return f_0, f_1 + + def _construct_warp(self): + # Set local constants + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8) + _opp_indices = self.velocity_set.opp_indices + lattice_central_index = self.velocity_set.center_index + + # Read the list of bc_to_id created upon instantiation + bc_to_id = boundary_condition_registry.bc_to_id + + # Gather IDs of ExtrapolationOutflowBC boundary conditions + extrapolation_outflow_bc_ids = [] + for bc_name, bc_id in bc_to_id.items(): + if bc_name.startswith("ExtrapolationOutflowBC"): + extrapolation_outflow_bc_ids.append(bc_id) + + @wp.func + def apply_bc( + index: Any, + timestep: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + is_post_streaming: bool, + ): + f_result = f_post + + # Unroll the loop over boundary conditions + for i in range(wp.static(len(self.boundary_conditions))): + if is_post_streaming: + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.STREAMING): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].warp_functional)(index, timestep, _missing_mask, f_0, f_1, f_pre, f_post) + else: + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.COLLISION): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].warp_functional)(index, timestep, _missing_mask, f_0, f_1, f_pre, f_post) + if wp.static(self.boundary_conditions[i].id in extrapolation_outflow_bc_ids): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].assemble_auxiliary_data)( + index, timestep, _missing_mask, f_0, f_1, f_pre, f_post + ) + return f_result + + @wp.func + def get_thread_data( + f0_buffer: wp.array4d(dtype=Any), + missing_mask: wp.array4d(dtype=Any), + index: Any, + ): + # Read thread data for populations + _f0_thread = _f_vec() + _missing_mask = _missing_mask_vec() + for l in range(self.velocity_set.q): + # q-sized vector of pre-streaming populations + _f0_thread[l] = self.compute_dtype(f0_buffer[l, index[0], index[1], index[2]]) + _missing_mask[l] = missing_mask[l, index[0], index[1], index[2]] + + return _f0_thread, _missing_mask + + @wp.func + def apply_aux_recovery_bc( + index: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + ): + # Note: + # In XLB, the BC auxiliary data (e.g. prescribed values of pressure or normal velocity) are stored in (i) central index of f_1 and/or + # (ii) missing directions of f_1. Some BCs may or may not need all these available storage space. This function checks whether + # the BC needs recovery of auxiliary data and then recovers the information for the next iteration (due to buffer swapping) by + # writting the values of f_1 into f_0. + + # Unroll the loop over boundary conditions + for i in range(wp.static(len(self.boundary_conditions))): + if wp.static(self.boundary_conditions[i].needs_aux_recovery): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + for l in range(self.velocity_set.q): + # Perform the swapping of data + if l == lattice_central_index: + # (i) Recover the values stored in the central index of f_1 + f_0[l, index[0], index[1], index[2]] = self.store_dtype(f_1[l, index[0], index[1], index[2]]) + elif _missing_mask[l] == wp.uint8(1): + # (ii) Recover the values stored in the missing directions of f_1 + f_0[_opp_indices[l], index[0], index[1], index[2]] = self.store_dtype( + f_1[_opp_indices[l], index[0], index[1], index[2]] + ) + + @wp.kernel + def kernel( + f_0: wp.array4d(dtype=Any), + f_1: wp.array4d(dtype=Any), + bc_mask: wp.array4d(dtype=Any), + missing_mask: wp.array4d(dtype=Any), + omega: Any, + timestep: int, + ): + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + _boundary_id = bc_mask[0, index[0], index[1], index[2]] + if _boundary_id == wp.uint8(BC_SOLID): + return + + # Apply streaming + _f_post_stream = self.stream.warp_functional(f_0, index) + + _f0_thread, _missing_mask = get_thread_data(f_0, missing_mask, index) + _f_post_collision = _f0_thread + + # Apply post-streaming boundary conditions + _f_post_stream = apply_bc(index, timestep, _boundary_id, _missing_mask, f_0, f_1, _f_post_collision, _f_post_stream, True) + + _rho, _u = self.macroscopic.warp_functional(_f_post_stream) + _feq = self.equilibrium.warp_functional(_rho, _u) + _f_post_collision = self.collision.warp_functional(_f_post_stream, _feq, omega) + + # Apply post-collision boundary conditions + _f_post_collision = apply_bc(index, timestep, _boundary_id, _missing_mask, f_0, f_1, _f_post_stream, _f_post_collision, False) + + # Apply auxiliary recovery for boundary conditions (swapping) + apply_aux_recovery_bc(index, _boundary_id, _missing_mask, f_0, f_1) + + # Store the result in f_1 + for l in range(self.velocity_set.q): + f_1[l, index[0], index[1], index[2]] = self.store_dtype(_f_post_collision[l]) + + return None, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + wp.launch( + self.warp_kernel, + inputs=[f_0, f_1, bc_mask, missing_mask, omega, timestep], + dim=f_0.shape[1:], + device=f_0.device, + ) + return f_0, f_1 + + def _construct_neon(self): + import neon + + # Set local constants + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8) + _opp_indices = self.velocity_set.opp_indices + lattice_central_index = self.velocity_set.center_index + + # Read the list of bc_to_id created upon instantiation + bc_to_id = boundary_condition_registry.bc_to_id + + # Gather IDs of ExtrapolationOutflowBC boundary conditions + extrapolation_outflow_bc_ids = [] + for bc_name, bc_id in bc_to_id.items(): + if bc_name.startswith("ExtrapolationOutflowBC"): + extrapolation_outflow_bc_ids.append(bc_id) + + @wp.func + def apply_bc( + index: Any, + timestep: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0: Any, + f_1: Any, + f_pre: Any, + f_post: Any, + is_post_streaming: bool, + ): + f_result = f_post + + # Unroll the loop over boundary conditions + for i in range(wp.static(len(self.boundary_conditions))): + if is_post_streaming: + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.STREAMING): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].neon_functional)(index, timestep, _missing_mask, f_0, f_1, f_pre, f_post) + else: + if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.COLLISION): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].neon_functional)(index, timestep, _missing_mask, f_0, f_1, f_pre, f_post) + if wp.static(self.boundary_conditions[i].id in extrapolation_outflow_bc_ids): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + f_result = wp.static(self.boundary_conditions[i].assemble_auxiliary_data)( + index, timestep, _missing_mask, f_0, f_1, f_pre, f_post + ) + return f_result + + @wp.func + def neon_get_thread_data( + f0_pn: Any, + missing_mask_pn: Any, + index: Any, + ): + # Read thread data for populations + _f0_thread = _f_vec() + _missing_mask = _missing_mask_vec() + for l in range(self.velocity_set.q): + # q-sized vector of pre-streaming populations + _f0_thread[l] = self.compute_dtype(wp.neon_read(f0_pn, index, l)) + _missing_mask[l] = wp.neon_read(missing_mask_pn, index, l) + + return _f0_thread, _missing_mask + + @wp.func + def neon_apply_aux_recovery_bc( + index: Any, + _boundary_id: Any, + _missing_mask: Any, + f_0_pn: Any, + f_1_pn: Any, + ): + # Note: + # In XLB, the BC auxiliary data (e.g. prescribed values of pressure or normal velocity) are stored in (i) central index of f_1 and/or + # (ii) missing directions of f_1. Some BCs may or may not need all these available storage space. This function checks whether + # the BC needs recovery of auxiliary data and then recovers the information for the next iteration (due to buffer swapping) by + # writting the values of f_1 into f_0. + + # Unroll the loop over boundary conditions + for i in range(wp.static(len(self.boundary_conditions))): + if wp.static(self.boundary_conditions[i].needs_aux_recovery): + if _boundary_id == wp.static(self.boundary_conditions[i].id): + for l in range(self.velocity_set.q): + # Perform the swapping of data + if l == lattice_central_index: + # (i) Recover the values stored in the central index of f_1 + _f1_thread = wp.neon_read(f_1_pn, index, l) + wp.neon_write(f_0_pn, index, l, self.store_dtype(_f1_thread)) + elif _missing_mask[l] == wp.uint8(1): + # (ii) Recover the values stored in the missing directions of f_1 + _f1_thread = wp.neon_read(f_1_pn, index, _opp_indices[l]) + wp.neon_write(f_0_pn, index, _opp_indices[l], self.store_dtype(_f1_thread)) + + @neon.Container.factory(name="nse_stepper") + def container( + f_0_fd: Any, + f_1_fd: Any, + bc_mask_fd: Any, + missing_mask_fd: Any, + omega: Any, + timestep: int, + ): + def nse_stepper_ll(loader: neon.Loader): + loader.set_grid(bc_mask_fd.get_grid()) + + f_0_pn = loader.get_read_handle( + f_0_fd, + operation=neon.Loader.Operation.stencil, + discretization=neon.Loader.Discretization.lattice, + ) + bc_mask_pn = loader.get_read_handle(bc_mask_fd) + missing_mask_pn = loader.get_read_handle(missing_mask_fd) + + f_1_pn = loader.get_write_handle(f_1_fd) + + @wp.func + def nse_stepper_cl(index: Any): + _boundary_id = wp.neon_read(bc_mask_pn, index, 0) + if _boundary_id == wp.uint8(BC_SOLID): + return + # Apply streaming + _f_post_stream = self.stream.neon_functional(f_0_pn, index) + + _f0_thread, _missing_mask = neon_get_thread_data(f_0_pn, missing_mask_pn, index) + _f_post_collision = _f0_thread + + # Apply post-streaming boundary conditions + _f_post_stream = apply_bc(index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_collision, _f_post_stream, True) + + _rho, _u = self.macroscopic.neon_functional(_f_post_stream) + _feq = self.equilibrium.neon_functional(_rho, _u) + _f_post_collision = self.collision.neon_functional(_f_post_stream, _feq, omega) + + # Apply post-collision boundary conditions + _f_post_collision = apply_bc( + index, timestep, _boundary_id, _missing_mask, f_0_pn, f_1_pn, _f_post_stream, _f_post_collision, False + ) + + # Apply auxiliary recovery for boundary conditions (swapping) + neon_apply_aux_recovery_bc(index, _boundary_id, _missing_mask, f_0_pn, f_1_pn) + + # Store the result in f_1 + for l in range(self.velocity_set.q): + wp.neon_write(f_1_pn, index, l, self.store_dtype(_f_post_collision[l])) + + loader.declare_kernel(nse_stepper_cl) + + return nse_stepper_ll + + return None, container + + @Operator.register_backend(ComputeBackend.NEON) + def neon_launch(self, f_0, f_1, bc_mask, missing_mask, omega, timestep): + if timestep == 0: + self.prepare_skeleton(f_0, f_1, bc_mask, missing_mask, omega) + self.sk[self.sk_iter].run() + self.sk_iter = (self.sk_iter + 1) % 2 + return f_0, f_1 + + def prepare_skeleton(self, f_0, f_1, bc_mask, missing_mask, omega): + """Build the Neon odd/even skeletons for double-buffered time stepping.""" + import neon + + grid = f_0.get_grid() + bk = grid.backend + self.neon_skeleton = {"odd": {}, "even": {}} + self.neon_skeleton["odd"]["container"] = self.neon_container(f_0, f_1, bc_mask, missing_mask, omega, 0) + self.neon_skeleton["even"]["container"] = self.neon_container(f_1, f_0, bc_mask, missing_mask, omega, 1) + # check if 'occ' is a valid key + if "occ" not in self.backend_config: + occ = neon.SkeletonConfig.OCC.none() + else: + occ = self.backend_config["occ"] + # check that occ is of type neon.SkeletonConfig.OCC + if not isinstance(occ, neon.SkeletonConfig.OCC): + print(type(occ)) + raise ValueError("occ must be of type neon.SkeletonConfig.OCC") + + for key in self.neon_skeleton: + self.neon_skeleton[key]["app"] = [self.neon_skeleton[key]["container"]] + self.neon_skeleton[key]["skeleton"] = neon.Skeleton(backend=bk) + self.neon_skeleton[key]["skeleton"].sequence(name="mres_nse_stepper", containers=self.neon_skeleton[key]["app"], occ=occ) + + self.sk = [self.neon_skeleton["odd"]["skeleton"], self.neon_skeleton["even"]["skeleton"]] + self.sk_iter = 0 diff --git a/xlb/operator/stepper/stepper.py b/xlb/operator/stepper/stepper.py new file mode 100644 index 00000000..b2ed7418 --- /dev/null +++ b/xlb/operator/stepper/stepper.py @@ -0,0 +1,34 @@ +# Base class for all stepper operators +from xlb.operator import Operator +from xlb import DefaultConfig + + +class Stepper(Operator): + """ + Class that handles the construction of lattice boltzmann stepping operator + """ + + def __init__(self, grid, boundary_conditions): + self.grid = grid + self.boundary_conditions = boundary_conditions + # Get velocity set, precision policy, and compute backend + velocity_set = DefaultConfig.velocity_set + precision_policy = DefaultConfig.default_precision_policy + compute_backend = DefaultConfig.default_backend + + # Initialize operator + super().__init__(velocity_set, precision_policy, compute_backend) + + def prepare_fields(self, initializer=None): + """Initialize the fields required for the stepper. + + Args: + initializer: Optional operator to initialize the distribution functions. + If provided, it should be a callable that takes (grid, velocity_set, + precision_policy, compute_backend) as arguments and returns initialized f_0. + If None, default equilibrium initialization is used with rho=1 and u=0. + + Returns: + Tuple of (f_0, f_1, bc_mask, missing_mask) + """ + raise NotImplementedError("Subclasses must implement prepare_fields()") diff --git a/xlb/operator/stream/__init__.py b/xlb/operator/stream/__init__.py new file mode 100644 index 00000000..9093da71 --- /dev/null +++ b/xlb/operator/stream/__init__.py @@ -0,0 +1 @@ +from xlb.operator.stream.stream import Stream diff --git a/xlb/operator/stream/stream.py b/xlb/operator/stream/stream.py new file mode 100644 index 00000000..e337ff39 --- /dev/null +++ b/xlb/operator/stream/stream.py @@ -0,0 +1,154 @@ +""" +Streaming operator for the Lattice Boltzmann Method. + +Implements the pull-scheme propagation step: each voxel reads populations +from its lattice neighbours according to the velocity-set directions. +""" + +from functools import partial +import jax.numpy as jnp +from jax import jit, vmap +import warp as wp +from typing import Any + +from xlb.compute_backend import ComputeBackend +from xlb.operator.operator import Operator + + +class Stream(Operator): + """Pull-scheme streaming operator. + + Propagates distribution functions by reading each population from the + upstream neighbour along the corresponding lattice direction. Periodic + boundaries are applied automatically when a pull index falls outside + the domain (Warp backend only; JAX uses ``jnp.roll``). + + Supports JAX, Warp backends. + """ + + @Operator.register_backend(ComputeBackend.JAX) + @partial(jit, static_argnums=(0)) + def jax_implementation(self, f): + """ + JAX implementation of the streaming step. + + TODO: Make sure this works with pull scheme. + + Parameters + ---------- + f: jax.numpy.ndarray + The distribution function. + """ + + def _streaming_jax_i(f, c): + """ + Perform individual streaming operation in a direction. + + Parameters + ---------- + f: The distribution function. + c: The streaming direction vector. + + Returns + ------- + jax.numpy.ndarray + The updated distribution function after streaming. + """ + if self.velocity_set.d == 2: + return jnp.roll(f, (c[0], c[1]), axis=(0, 1)) + elif self.velocity_set.d == 3: + return jnp.roll(f, (c[0], c[1], c[2]), axis=(0, 1, 2)) + + return vmap(_streaming_jax_i, in_axes=(0, 0), out_axes=0)(f, jnp.array(self.velocity_set.c).T) + + def _construct_warp(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _c = self.velocity_set.c + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + # Construct the funcional to get streamed indices + @wp.func + def functional( + f: wp.array4d(dtype=Any), + index: Any, + ): + # Pull the distribution function + _f = _f_vec() + for l in range(self.velocity_set.q): + # Get pull index + pull_index = type(index)() + for d in range(self.velocity_set.d): + pull_index[d] = index[d] - _c[d, l] + + # impose periodicity for out of bound values + if pull_index[d] < 0: + pull_index[d] = f.shape[d + 1] - 1 + elif pull_index[d] >= f.shape[d + 1]: + pull_index[d] = 0 + + # Read the distribution function + # Unlike other functionals, we need to cast the type here since we read from the buffer + _f[l] = self.compute_dtype(f[l, pull_index[0], pull_index[1], pull_index[2]]) + + return _f + + # Construct the warp kernel + @wp.kernel + def kernel( + f_0: wp.array4d(dtype=Any), + f_1: wp.array4d(dtype=Any), + ): + # Get the global index + i, j, k = wp.tid() + index = wp.vec3i(i, j, k) + + # Set the output + _f = functional(f_0, index) + + # Write the output + for l in range(self.velocity_set.q): + f_1[l, index[0], index[1], index[2]] = self.store_dtype(_f[l]) + + return functional, kernel + + @Operator.register_backend(ComputeBackend.WARP) + def warp_implementation(self, f_0, f_1): + # Launch the warp kernel + wp.launch( + self.warp_kernel, + inputs=[ + f_0, + f_1, + ], + dim=f_0.shape[1:], + ) + return f_1 + + def _construct_neon(self): + # Set local constants TODO: This is a hack and should be fixed with warp update + _c = self.velocity_set.c + _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) + + # Construct the funcional to get streamed indices + @wp.func + def functional( + f: Any, + index: Any, + ): + # Pull the distribution function + _f = _f_vec() + for l in range(self.velocity_set.q): + # Get pull offset + ngh = wp.neon_ngh_idx(wp.int8(-_c[0, l]), wp.int8(-_c[1, l]), wp.int8(-_c[2, l])) + unused_is_valid = wp.bool(False) + + # Read the distribution function from the neighboring cell in the pull direction + _f[l] = self.compute_dtype(wp.neon_read_ngh(f, index, ngh, l, self.store_dtype(0), unused_is_valid)) + return _f + + return functional, None + + @Operator.register_backend(ComputeBackend.NEON) + def neon_implementation(self, f_0, f_1): + # raise exception as this feature is not implemented yet + raise NotImplementedError("This feature is not implemented in XLB with the NEON backend yet.") diff --git a/xlb/physics_type.py b/xlb/physics_type.py new file mode 100644 index 00000000..39841fcc --- /dev/null +++ b/xlb/physics_type.py @@ -0,0 +1,8 @@ +# Enum used to keep track of the physics types supported by different operators + +from enum import Enum, auto + + +class PhysicsType(Enum): + NSE = auto() # Navier-Stokes Equations + ADE = auto() # Advection-Diffusion Equations diff --git a/xlb/precision_policy.py b/xlb/precision_policy.py new file mode 100644 index 00000000..32a6d567 --- /dev/null +++ b/xlb/precision_policy.py @@ -0,0 +1,110 @@ +""" +Precision and precision-policy enumerations for XLB. + +:class:`Precision` maps symbolic precisions to Warp and JAX dtypes. +:class:`PrecisionPolicy` pairs a *compute* precision (used during +arithmetic) with a *store* precision (used in memory), enabling +mixed-precision simulations. +""" + +from enum import Enum, auto + + +class Precision(Enum): + """Scalar precision levels with Warp and JAX dtype accessors.""" + + FP64 = auto() + FP32 = auto() + FP16 = auto() + UINT8 = auto() + BOOL = auto() + + @property + def wp_dtype(self): + import warp as wp + + if self == Precision.FP64: + return wp.float64 + elif self == Precision.FP32: + return wp.float32 + elif self == Precision.FP16: + return wp.float16 + elif self == Precision.UINT8: + return wp.uint8 + elif self == Precision.BOOL: + return wp.bool + else: + raise ValueError("Invalid precision") + + @property + def jax_dtype(self): + import jax.numpy as jnp + + if self == Precision.FP64: + return jnp.float64 + elif self == Precision.FP32: + return jnp.float32 + elif self == Precision.FP16: + return jnp.float16 + elif self == Precision.UINT8: + return jnp.uint8 + elif self == Precision.BOOL: + return jnp.bool_ + else: + raise ValueError("Invalid precision") + + +class PrecisionPolicy(Enum): + """Mixed-precision policy pairing compute and store precisions. + + The naming convention is ````, e.g. ``FP32FP16`` + computes in FP32 and stores results in FP16. + """ + + FP64FP64 = auto() + FP64FP32 = auto() + FP64FP16 = auto() + FP32FP32 = auto() + FP32FP16 = auto() + + @property + def compute_precision(self): + if self == PrecisionPolicy.FP64FP64: + return Precision.FP64 + elif self == PrecisionPolicy.FP64FP32: + return Precision.FP64 + elif self == PrecisionPolicy.FP64FP16: + return Precision.FP64 + elif self == PrecisionPolicy.FP32FP32: + return Precision.FP32 + elif self == PrecisionPolicy.FP32FP16: + return Precision.FP32 + else: + raise ValueError("Invalid precision policy") + + @property + def store_precision(self): + if self == PrecisionPolicy.FP64FP64: + return Precision.FP64 + elif self == PrecisionPolicy.FP64FP32: + return Precision.FP32 + elif self == PrecisionPolicy.FP64FP16: + return Precision.FP16 + elif self == PrecisionPolicy.FP32FP32: + return Precision.FP32 + elif self == PrecisionPolicy.FP32FP16: + return Precision.FP16 + else: + raise ValueError("Invalid precision policy") + + def cast_to_compute_jax(self, array): + import jax.numpy as jnp + + compute_precision = self.compute_precision + return jnp.array(array, dtype=compute_precision.jax_dtype) + + def cast_to_store_jax(self, array): + import jax.numpy as jnp + + store_precision = self.store_precision + return jnp.array(array, dtype=store_precision.jax_dtype) diff --git a/xlb/precision_policy/precision_policy.py b/xlb/precision_policy/precision_policy.py new file mode 100644 index 00000000..6b400ee7 --- /dev/null +++ b/xlb/precision_policy/precision_policy.py @@ -0,0 +1,49 @@ +from xlb.compute_backend import ComputeBackend +from xlb import DefaultConfig +from xlb.precision_policy.jax_precision_policy import ( + JaxFp32Fp32, + JaxFp32Fp16, + JaxFp64Fp64, + JaxFp64Fp32, + JaxFp64Fp16, +) + + +class Fp64Fp64: + def __new__(cls): + if DefaultConfig.compute_backend == ComputeBackend.JAX: + return JaxFp64Fp64() + else: + raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}") + + +class Fp64Fp32: + def __new__(cls): + if DefaultConfig.compute_backend == ComputeBackend.JAX: + return JaxFp64Fp32() + else: + raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}") + + +class Fp32Fp32: + def __new__(cls): + if DefaultConfig.compute_backend == ComputeBackend.JAX: + return JaxFp32Fp32() + else: + raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}") + + +class Fp64Fp16: + def __new__(cls): + if DefaultConfig.compute_backend == ComputeBackend.JAX: + return JaxFp64Fp16() + else: + raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}") + + +class Fp32Fp16: + def __new__(cls): + if DefaultConfig.compute_backend == ComputeBackend.JAX: + return JaxFp32Fp16() + else: + raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}") diff --git a/xlb/utils/__init__.py b/xlb/utils/__init__.py new file mode 100644 index 00000000..07de394d --- /dev/null +++ b/xlb/utils/__init__.py @@ -0,0 +1,19 @@ +from .utils import ( + downsample_field, + warp_array_to_jax, + jax_has_gpu_devices, + save_image, + save_fields_vtk, + save_BCs_vtk, + rotate_geometry, + voxelize_stl, + axangle2mat, + ToJAX, + UnitConvertor, + save_usd_vorticity, + save_usd_q_criterion, + update_usd_lagrangian_parts, + plot_object_placement, + colorize_scalars, +) +from .mesher import make_cuboid_mesh, MultiresIO diff --git a/xlb/utils/mesher.py b/xlb/utils/mesher.py new file mode 100644 index 00000000..e8e9ca97 --- /dev/null +++ b/xlb/utils/mesher.py @@ -0,0 +1,943 @@ +""" +Multi-resolution mesh utilities. + +Provides geometry preparation and I/O for multi-resolution LBM simulations: + +* :func:`make_cuboid_mesh` β€” builds a strongly-balanced cuboid mesh hierarchy + from an STL file and a sequence of domain multipliers. +* :func:`prepare_sparsity_pattern` β€” converts level data into the sparsity + arrays required by :func:`multires_grid_factory`. +* :class:`MultiresIO` β€” exports multi-resolution Neon field data to HDF5 / + XDMF, 2-D slice images, and 1-D line profiles. +""" + +import numpy as np +import trimesh +from typing import Any, Optional + +import warp as wp +from xlb.utils.utils import UnitConvertor + + +def adjust_bbox(cuboid_max, cuboid_min, voxel_size_up): + """ + Adjust the bounding box to the nearest points of one level finer grid that encloses the desired region. + + Args: + cuboid_min (np.ndarray): Desired minimum coordinates of the bounding box. + cuboid_max (np.ndarray): Desired maximum coordinates of the bounding box. + voxel_size_up (float): Voxel size of one level higher (finer) grid. + + Returns: + tuple: (adjusted_min, adjusted_max) snapped to grid points of one level higher. + """ + adjusted_min = np.round(cuboid_min / voxel_size_up) * voxel_size_up + adjusted_max = np.round(cuboid_max / voxel_size_up) * voxel_size_up + return adjusted_min, adjusted_max + + +def prepare_sparsity_pattern(level_data): + """ + Prepare the sparsity pattern for the multiresolution grid based on the level data. "level_data" is expected to be formatted as in + the output of "make_cuboid_mesh". + """ + num_levels = len(level_data) + level_origins = [] + sparsity_pattern = [] + for lvl in range(num_levels): + # Get the level mask from the level data + level_mask = level_data[lvl][0] + + # Ensure level_0 is contiguous int32 + level_mask = np.ascontiguousarray(level_mask, dtype=np.int32) + + # Append the padded level mask to the sparsity pattern + sparsity_pattern.append(level_mask) + + # Get the origin for this level + level_origins.append(level_data[lvl][2]) + + return sparsity_pattern, level_origins + + +def make_cuboid_mesh(voxel_size, cuboids, stl_filename): + """ + Create a strongly-balanced multi-level cuboid mesh with a sequence of bounding boxes. + Outputs mask arrays that are set to True only in regions not covered by finer levels. + + Args: + voxel_size (float): Voxel size of the finest grid . + cuboids (list): List of multipliers defining each level's domain. + stl_name (str): Path to the STL file. + + Returns: + list: Level data with mask arrays, voxel sizes, origins, and levels. + """ + # Load the mesh and get its bounding box + mesh = trimesh.load_mesh(stl_filename, process=False) + assert not mesh.is_empty, "Loaded mesh is empty or invalid." + + mesh_vertices = mesh.vertices + min_bound = mesh_vertices.min(axis=0) + max_bound = mesh_vertices.max(axis=0) + partSize = max_bound - min_bound + + level_data = [] + adjusted_bboxes = [] + max_voxel_size = voxel_size * pow(2, (len(cuboids) - 1)) + # Step 1: Generate all levels and store their data + for level in range(len(cuboids)): + # Compute desired bounding box for this level + cuboid_min = np.array( + [ + min_bound[0] - cuboids[level][0] * partSize[0], + min_bound[1] - cuboids[level][2] * partSize[1], + min_bound[2] - cuboids[level][4] * partSize[2], + ], + dtype=float, + ) + + cuboid_max = np.array( + [ + max_bound[0] + cuboids[level][1] * partSize[0], + max_bound[1] + cuboids[level][3] * partSize[1], + max_bound[2] + cuboids[level][5] * partSize[2], + ], + dtype=float, + ) + + # Set voxel size for this level + voxel_size_level = max_voxel_size / pow(2, level) + + # Adjust bounding box to align with one level up (finer grid) + if level > 0: + voxel_level_up = max_voxel_size / pow(2, level - 1) + else: + voxel_level_up = voxel_size_level + adjusted_min, adjusted_max = adjust_bbox(cuboid_max, cuboid_min, voxel_level_up) + + xmin, ymin, zmin = adjusted_min + xmax, ymax, zmax = adjusted_max + + # Compute number of voxels based on level-specific voxel size + nx = int(np.round((xmax - xmin) / voxel_size_level)) + ny = int(np.round((ymax - ymin) / voxel_size_level)) + nz = int(np.round((zmax - zmin) / voxel_size_level)) + print(f"Domain {nx}, {ny}, {nz} Origin {adjusted_min} Voxel Size {voxel_size_level} Voxel Level Up {voxel_level_up}") + + voxel_matrix = np.ones((nx, ny, nz), dtype=bool) + + origin = adjusted_min + level_data.append((voxel_matrix, voxel_size_level, origin, level)) + adjusted_bboxes.append((adjusted_min, adjusted_max)) + + # Step 2: Adjust coarser levels to exclude regions covered by finer levels + for k in range(len(level_data) - 1): # Exclude the finest level + # Current level's data + voxel_matrix_k = level_data[k][0] + origin_k = level_data[k][2] + voxel_size_k = level_data[k][1] + nx, ny, nz = voxel_matrix_k.shape + + # Next finer level's bounding box + adjusted_min_k1, adjusted_max_k1 = adjusted_bboxes[k + 1] + + # Compute index ranges in level k that overlap with level k+1's bounding box + # Use epsilon (1e-10) to handle floating-point precision + i_start = max(0, int(np.ceil((adjusted_min_k1[0] - origin_k[0] - 1e-10) / voxel_size_k))) + i_end = min(nx, int(np.floor((adjusted_max_k1[0] - origin_k[0] + 1e-10) / voxel_size_k))) + j_start = max(0, int(np.ceil((adjusted_min_k1[1] - origin_k[1] - 1e-10) / voxel_size_k))) + j_end = min(ny, int(np.floor((adjusted_max_k1[1] - origin_k[1] + 1e-10) / voxel_size_k))) + k_start = max(0, int(np.ceil((adjusted_min_k1[2] - origin_k[2] - 1e-10) / voxel_size_k))) + k_end = min(nz, int(np.floor((adjusted_max_k1[2] - origin_k[2] + 1e-10) / voxel_size_k))) + + # Set overlapping region to zero + voxel_matrix_k[i_start:i_end, j_start:j_end, k_start:k_end] = 0 + + # Step 3 Convert to Indices from STL units + num_levels = len(level_data) + level_data = [(dr, int(v / voxel_size), np.round(dOrigin / v).astype(int), num_levels - 1 - l) for dr, v, dOrigin, l in level_data] + + return list(reversed(level_data)) + + +class MultiresIO(object): + """I/O helper for multi-resolution Neon field data. + + Converts hierarchical Neon ``mGrid`` fields into merged unstructured + hexahedral meshes and exports them as HDF5 + XDMF (for ParaView), + 2-D slice PNG images, or 1-D line CSV profiles. + + The constructor precomputes the merged geometry (coordinates, + connectivity, centroids) and allocates intermediate Warp fields so + that repeated exports only need to transfer data from the Neon fields. + """ + + def __init__( + self, + field_name_cardinality_dict, + levels_data, + unit_convertor: UnitConvertor = None, + offset: Optional[tuple] = (0.0, 0.0, 0.0), + store_precision=None, + ): + """ + Initialize the MultiresIO object. + + Parameters + ---------- + field_name_cardinality_dict : dict + A dictionary mapping field names to their cardinalities. + Example: {'velocity_x': 1, 'velocity_y': 1, 'velocity': 3, 'density': 1} + levels_data : list of tuples + Each tuple contains (data, voxel_size, origin, level). + unit_convertor : UnitConvertor + An instance of the UnitConvertor class for unit conversions. + offset : tuple, optional + Offset to be applied to the coordinates. + store_precision : str, optional + The precision policy for storing data. + """ + # Set the unit convertor object + self.unit_convertor = unit_convertor + + # Process the multires geometry and extract coordinates and connectivity in the coordinate system of the finest level + coordinates, connectivity, level_id_field, total_cells = self.process_geometry(levels_data) + + # Ensure that coordinates and connectivity are not empty + assert coordinates.size != 0, "Error: No valid data to process. Check the input levels_data." + + # Merge duplicate points + coordinates, connectivity = self._merge_duplicates(coordinates, connectivity, levels_data) + + # Transform coordinates to physical units and apply offset if provided + coordinates = self._transform_coordinates(coordinates, offset) + + # Assign to self + self.field_name_cardinality_dict = field_name_cardinality_dict + self.levels_data = levels_data + self.coordinates = coordinates + self.connectivity = connectivity + self.level_id_field = level_id_field + self.total_cells = total_cells + self.centroids = np.mean(coordinates[connectivity], axis=1) + + # Set the default precision policy if not provided + from xlb import DefaultConfig + + if store_precision is None: + self.store_precision = DefaultConfig.default_precision_policy.store_precision + self.store_dtype = DefaultConfig.default_precision_policy.store_precision.wp_dtype + + # Prepare and allocate the inputs for the NEON container + self.field_warp_dict, self.origin_list = self._prepare_container_inputs() + + # Construct the NEON container for exporting multi-resolution data + self.container = self._construct_neon_container() + + def process_geometry(self, levels_data): + """Build merged coordinates and connectivity from all levels. + + Returns + ------- + coordinates : np.ndarray, shape (N, 3) + Vertex positions (8 per active voxel, before deduplication). + connectivity : np.ndarray, shape (M, 8) + Hexahedral connectivity (one row per active voxel). + level_id_field : np.ndarray, shape (M,) + Grid level index for each cell. + total_cells : int + Total number of active voxels across all levels. + """ + num_voxels_per_level = [np.sum(data) for data, _, _, _ in levels_data] + num_points_per_level = [8 * nv for nv in num_voxels_per_level] + point_id_offsets = np.cumsum([0] + num_points_per_level[:-1]) + + all_corners = [] + all_connectivity = [] + level_id_field = [] + total_cells = 0 + + for level_idx, (data, voxel_size, origin, level) in enumerate(levels_data): + origin = origin * voxel_size + corners_list, conn_list = self._process_level(data, voxel_size, origin, point_id_offsets[level_idx]) + + if corners_list: + print(f"\tProcessing level {level}: Voxel size {voxel_size}, Origin {origin}, Shape {data.shape}") + all_corners.extend(corners_list) + all_connectivity.extend(conn_list) + num_cells = sum(c.shape[0] for c in conn_list) + level_id_field.extend([level] * num_cells) + total_cells += num_cells + else: + print(f"\tSkipping level {level} (no unique data)") + + # Stacking coordinates and connectivity + coordinates = np.concatenate(all_corners, axis=0).astype(np.float32) + connectivity = np.concatenate(all_connectivity, axis=0).astype(np.int32) + level_id_field = np.array(level_id_field, dtype=np.uint8) + + return coordinates, connectivity, level_id_field, total_cells + + def _process_level(self, data, voxel_size, origin, point_id_offset): + """ + Given a voxel grid, returns all corners and connectivity in NumPy for this resolution level. + """ + true_indices = np.argwhere(data) + if true_indices.size == 0: + return [], [] + + max_voxels_per_chunk = 268_435_450 + chunks = np.array_split(true_indices, max(1, (len(true_indices) + max_voxels_per_chunk - 1) // max_voxels_per_chunk)) + + all_corners = [] + all_connectivity = [] + pid_offset = point_id_offset + + for chunk in chunks: + if chunk.size == 0: + continue + corners, connectivity = self._process_voxel_chunk(chunk, np.asarray(origin, dtype=np.float32), voxel_size, pid_offset) + all_corners.append(corners) + all_connectivity.append(connectivity) + pid_offset += len(chunk) * 8 + + return all_corners, all_connectivity + + def _process_voxel_chunk(self, true_indices, origin, voxel_size, point_id_offset): + """ + Given a set of voxel indices, returns 8 corners and connectivity for each cube using NumPy. + """ + true_indices = np.asarray(true_indices, dtype=np.float32) + mins = origin + true_indices * voxel_size + offsets = np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ], + dtype=np.float32, + ) + + corners = (mins[:, None, :] + offsets[None, :, :] * voxel_size).reshape(-1, 3).astype(np.float32) + base_ids = point_id_offset + np.arange(len(true_indices), dtype=np.int32) * 8 + connectivity = (base_ids[:, None] + np.arange(8, dtype=np.int32)).astype(np.int32) + + return corners, connectivity + + def save_xdmf(self, h5_filename, xmf_filename, total_cells, num_points, fields={}): + """Write an XDMF descriptor that references the companion HDF5 file.""" + # Generate an XDMF file to accompany the HDF5 file + print(f"\tGenerating XDMF file: {xmf_filename}") + hdf5_rel_path = h5_filename.split("/")[-1] + with open(xmf_filename, "w") as xmf: + xmf.write(f''' + + + + + + + {hdf5_rel_path}:/Mesh/Connectivity + + + + + {hdf5_rel_path}:/Mesh/Points + + + + + {hdf5_rel_path}:/Mesh/Level + + + ''') + for field_name in fields.keys(): + xmf.write(f''' + + + {hdf5_rel_path}:/Fields/{field_name} + + + ''') + xmf.write(""" + + + + """) + print("\tXDMF file written successfully") + return + + def save_hdf5_file(self, filename, coordinates, connectivity, level_id_field, fields_data, compression="gzip", compression_opts=0): + """Write the processed mesh data to an HDF5 file. + Parameters + ---------- + filename : str + The name of the output HDF5 file. + coordinates : numpy.ndarray + An array of all coordinates. + connectivity : numpy.ndarray + An array of all connectivity data. + level_id_field : numpy.ndarray + An array of all level data. + fields_data : dict + A dictionary of all field data. + compression : str, optional + The compression method to use for the HDF5 file. + compression_opts : int, optional + The compression options to use for the HDF5 file. + """ + import h5py + + with h5py.File(filename + ".h5", "w") as f: + f.create_dataset("/Mesh/Points", data=coordinates, compression=compression, compression_opts=compression_opts, chunks=True) + f.create_dataset( + "/Mesh/Connectivity", + data=connectivity, + compression=compression, + compression_opts=compression_opts, + chunks=True, + ) + f.create_dataset("/Mesh/Level", data=level_id_field, compression=compression, compression_opts=compression_opts) + fg = f.create_group("/Fields") + for fname, fdata in fields_data.items(): + fg.create_dataset(fname, data=fdata.astype(np.float32), compression=compression, compression_opts=compression_opts, chunks=True) + + def _merge_duplicates(self, coordinates, connectivity, levels_data): + """Deduplicate vertices shared between adjacent voxels. + + Uses spatial hashing (grid-snapped coordinates) processed in + chunks to keep memory bounded. + """ + # Merging duplicate points + tolerance = 0.01 + chunk_size = 10_000_000 # Adjust based on GPU memory + num_points = coordinates.shape[0] + unique_points = [] + mapping = np.zeros(num_points, dtype=np.int32) + unique_idx = 0 + + # Get the grid shape of computational box at the finest level from the levels_data + num_levels = len(levels_data) + grid_shape_finest = np.array(levels_data[-1][0].shape) * 2 ** (num_levels - 1) + + for start in range(0, num_points, chunk_size): + end = min(start + chunk_size, num_points) + coords_chunk = coordinates[start:end] + + # Simple hashing: grid coordinates as tuple keys + grid_coords = np.round(coords_chunk / tolerance).astype(np.int64) + hash_keys = grid_coords[:, 0] + grid_coords[:, 1] * grid_shape_finest[0] + grid_coords[:, 2] * grid_shape_finest[0] * grid_shape_finest[1] + unique_hash, inverse = np.unique(hash_keys, return_inverse=True) + unique_hash, unique_indices, inverse = np.unique(hash_keys, return_index=True, return_inverse=True) + unique_chunk = coords_chunk[unique_indices] + + unique_points.append(unique_chunk) + mapping[start:end] = inverse + unique_idx + unique_idx += len(unique_hash) + + coordinates = np.concatenate(unique_points) + connectivity = mapping[connectivity] + return coordinates, connectivity + + def _transform_coordinates(self, coordinates, offset): + """Convert lattice coordinates to physical units and apply offset.""" + offset = np.array(offset, dtype=np.float32) + if self.unit_convertor is not None: + coordinates = self.unit_convertor.length_to_physical(coordinates) + return coordinates + offset + + def _prepare_container_inputs(self): + """Allocate dense Warp fields used as staging buffers for Neon-to-NumPy transfer.""" + # load necessary modules + from xlb.compute_backend import ComputeBackend + from xlb.grid import grid_factory + + # Get the number of levels from the levels_data + num_levels = len(self.levels_data) + + # Prepare lists to hold warp fields and origins allocated for each level + field_warp_dict = {} + origin_list = [] + for field_name, cardinality in self.field_name_cardinality_dict.items(): + field_warp_dict[field_name] = [] + for level in range(num_levels): + # get the shape of the grid at this level + box_shape = self.levels_data[level][0].shape + + # Use the warp backend to create dense fields to be written in multi-res NEON fields + grid_dense = grid_factory(box_shape, compute_backend=ComputeBackend.WARP) + field_warp_dict[field_name].append(grid_dense.create_field(cardinality=cardinality, dtype=self.store_precision)) + origin_list.append(wp.vec3i(*([int(x) for x in self.levels_data[level][2]]))) + + return field_warp_dict, origin_list + + def _construct_neon_container(self): + """ + Constructs a NEON container for exporting multi-resolution data to HDF5. + This container will be used to transfer multi-resolution NEON fields into stacked warp fields. + """ + import neon + + @neon.Container.factory(name="HDF5MultiresExporter") + def container( + field_neon: Any, + field_warp: Any, + origin: Any, + level: Any, + ): + def launcher(loader: neon.Loader): + loader.set_mres_grid(field_neon.get_grid(), level) + field_neon_hdl = loader.get_mres_read_handle(field_neon) + refinement = 2**level + + @wp.func + def kernel(index: Any): + cIdx = wp.neon_global_idx(field_neon_hdl, index) + # Get local indices by dividing the global indices (associated with the finest level) by 2^level + # Subtract the origin to get the local indices in the warp field + lx = wp.neon_get_x(cIdx) // refinement - origin[0] + ly = wp.neon_get_y(cIdx) // refinement - origin[1] + lz = wp.neon_get_z(cIdx) // refinement - origin[2] + + # write the values to the warp field + cardinality = field_warp.shape[0] + for card in range(cardinality): + field_warp[card, lx, ly, lz] = self.store_dtype(wp.neon_read(field_neon_hdl, index, card)) + + loader.declare_kernel(kernel) + + return launcher + + return container + + def get_fields_data(self, field_neon_dict): + """ + Extracts and prepares the fields data from the NEON fields for export. + """ + import neon + + # Check if the field_neon_dict is empty + if not field_neon_dict: + return {} + + # Ensure that this operator is called on multires grids + grid_mres = next(iter(field_neon_dict.values())).get_grid() + assert grid_mres.name == "mGrid", f"Operation {self.__class__.__name__} is only applicable to multi-resolution cases!" + + for field_name in field_neon_dict.keys(): + assert field_name in self.field_name_cardinality_dict.keys(), ( + f"Field {field_name} is not provided in the instantiation of the MultiresIO class!" + ) + + # number of levels + num_levels = grid_mres.num_levels + assert num_levels == len(self.levels_data), "Error: Inconsistent number of levels!" + + # Prepare the fields dictionary to be written by transfering multi-res NEON fields into stacked warp fields and then numpy arrays + fields_data = {} + for field_name, cardinality in self.field_name_cardinality_dict.items(): + if field_name not in field_neon_dict: + continue + for card in range(cardinality): + fields_data[f"{field_name}_{card}"] = [] + + # Iterate over each field and level to fill the dictionary with numpy fields + for field_name, cardinality in self.field_name_cardinality_dict.items(): + if field_name not in field_neon_dict: + continue + for level in range(num_levels): + # Create the container and run it to fill the warp fields + c = self.container(field_neon_dict[field_name], self.field_warp_dict[field_name][level], self.origin_list[level], level) + c.run(0, container_runtime=neon.Container.ContainerRuntime.neon) + + # Ensure all operations are complete before converting to JAX and Numpy arrays + wp.synchronize() + + # Convert the warp fields to numpy arrays and use level's mask to filter the data + mask = self.levels_data[level][0] + field_np = self.field_warp_dict[field_name][level].numpy() + for card in range(cardinality): + field_np_card = field_np[card][mask] + fields_data[f"{field_name}_{card}"].append(field_np_card) + + # Concatenate all field data + for field_name in fields_data.keys(): + fields_data[field_name] = np.concatenate(fields_data[field_name]) + assert fields_data[field_name].size == self.total_cells, f"Error: Field {field_name} size mismatch!" + + # Unit conversion if applicable + if self.unit_convertor is not None: + if "velocity" in field_name.lower(): + fields_data[field_name] = self.unit_convertor.velocity_to_physical(fields_data[field_name]) + elif "density" in field_name.lower(): + fields_data[field_name] = self.unit_convertor.density_to_physical(fields_data[field_name]) + elif "pressure" in field_name.lower(): + fields_data[field_name] = self.unit_convertor.pressure_to_physical(fields_data[field_name]) + # Add more physical quantities as needed + + return fields_data + + def to_hdf5(self, output_filename, field_neon_dict, compression="gzip", compression_opts=0): + """ + Export the multi-resolution mesh data to an HDF5 file. + Parameters + ---------- + output_filename : str + The name of the output HDF5 file (without extension). + field_neon_dict : a dictionary of neon mGrid Fields + Eg. The NEON fields containing velocity and density data as { "velocity": velocity_neon, "density": density_neon} + compression : str, optional + The compression method to use for the HDF5 file. + compression_opts : int, optional + The compression options to use for the HDF5 file. + """ + import time + + # Get the fields data from the NEON fields + fields_data = self.get_fields_data(field_neon_dict) + + # Save XDMF file + self.save_xdmf(output_filename + ".h5", output_filename + ".xmf", self.total_cells, len(self.coordinates), fields_data) + + # Writing HDF5 file + print("\tWriting HDF5 file") + tic_write = time.perf_counter() + self.save_hdf5_file(output_filename, self.coordinates, self.connectivity, self.level_id_field, fields_data, compression, compression_opts) + toc_write = time.perf_counter() + print(f"\tHDF5 file written in {toc_write - tic_write:0.1f} seconds") + + def to_slice_image( + self, + output_filename, + field_neon_dict, + plane_point, + plane_normal, + slice_thickness=1.0, + bounds=[0, 1, 0, 1], + grid_res=512, + cmap=None, + component=None, + show_axes=False, + show_colorbar=False, + **kwargs, + ): + """ + Export an arbitrary-plane slice from unstructured point data to PNG. + + Parameters + ---------- + output_filename : str + Output PNG filename (without extension). + field_neon_dict : dict + A dictionary of NEON fields containing the data to be plotted. + Example: {"velocity": velocity_neon, "density": density_neon} + plane_point : array_like + A point [x, y, z] on the plane. + plane_normal : array_like + Plane normal vector [nx, ny, nz]. + slice_thickness : float + How thick (in units of the coordinate system) the slice should be. + grid_resolution : tuple + Resolution of output image (pixels in plane u, v directions). + grid_size : tuple + Physical size of slice grid (width, height). + cmap : str + Matplotlib colormap. + """ + # Get the fields data from the NEON fields + assert len(field_neon_dict.keys()) == 1, "Error: This function is designed to plot a single field at a time." + fields_data = self.get_fields_data(field_neon_dict) + + # Check if the component is within the valid range + if component is None: + print("\tCreating slice image of the field magnitude!") + cell_data = list(fields_data.values()) + squared = [comp**2 for comp in cell_data] + cell_data = np.sqrt(sum(squared)) + field_name = list(fields_data.keys())[0].split("_")[0] + "_magnitude" + else: + assert component < max(self.field_name_cardinality_dict.values()), ( + f"Error: Component {component} is out of range for the provided fields." + ) + print(f"\tCreating slice image for component {component} of the input field!") + field_name = list(fields_data.keys())[component] + cell_data = fields_data[field_name] + + # Plot each field in the dictionary + self._to_slice_image_single_field( + f"{output_filename}_{field_name}", + cell_data, + plane_point, + plane_normal, + slice_thickness=slice_thickness, + bounds=bounds, + grid_res=grid_res, + cmap=cmap, + show_axes=show_axes, + show_colorbar=show_colorbar, + **kwargs, + ) + print(f"\tSlice image for field {field_name} saved as {output_filename}.png") + + def _to_slice_image_single_field( + self, + output_filename, + field_data, + plane_point, + plane_normal, + slice_thickness, + bounds, + grid_res, + cmap, + show_axes, + show_colorbar, + **kwargs, + ): + """ + Helper function to create a slice image for a single field. + """ + from matplotlib import cm + import numpy as np + import matplotlib.pyplot as plt + from scipy.spatial import cKDTree + + # field data are associated with the cells centers + cell_values = field_data + + # get the normalized plane normal + plane_normal = np.asarray(np.abs(plane_normal)) + n = plane_normal / np.linalg.norm(plane_normal) + + # Compute signed distances of each cell center to the plane + plane_point *= plane_normal + sdf = np.dot(self.centroids - plane_point, n) + + # Filter: cells with centroid near plane + mask = np.abs(sdf) <= slice_thickness / 2 + if not np.any(mask): + raise ValueError("No cells intersect the plane within thickness.") + + # Project centroids to plane + centroids_slice = self.centroids[mask] + sdf_slice = sdf[mask] + proj = centroids_slice - np.outer(sdf_slice, n) + + values = cell_values[mask] + + # Build in-plane basis + if np.allclose(n, [1, 0, 0]): + u1 = np.array([0, 1, 0]) + else: + u1 = np.array([1, 0, 0]) + u2 = np.abs(np.cross(n, u1)) + + local_x = np.dot(proj - plane_point, u1) + local_y = np.dot(proj - plane_point, u2) + + # Define extent of the plot + xmin, xmax, ymin, ymax = local_x.min(), local_x.max(), local_y.min(), local_y.max() + Lx = xmax - xmin + Ly = ymax - ymin + extent = np.array([xmin + bounds[0] * Lx, xmin + bounds[1] * Lx, ymin + bounds[2] * Ly, ymin + bounds[3] * Ly]) + mask_bounds = (extent[0] <= local_x) & (local_x <= extent[1]) & (extent[2] <= local_y) & (local_y <= extent[3]) + + if cmap is None: + cmap = cm.nipy_spectral + + # Adjust vertical resolution based on bounds + bounded_x_min = local_x[mask_bounds].min() + bounded_x_max = local_x[mask_bounds].max() + bounded_y_min = local_y[mask_bounds].min() + bounded_y_max = local_y[mask_bounds].max() + width_x = bounded_x_max - bounded_x_min + height_y = bounded_y_max - bounded_y_min + aspect_ratio = height_y / width_x + grid_resY = max(1, int(np.round(grid_res * aspect_ratio))) + + # Create grid + grid_x = np.linspace(bounded_x_min, bounded_x_max, grid_res) + grid_y = np.linspace(bounded_y_min, bounded_y_max, grid_resY) + xv, yv = np.meshgrid(grid_x, grid_y, indexing="xy") + + # Fast KDTree-based interpolation + points = np.column_stack((local_x[mask_bounds], local_y[mask_bounds])) + tree = cKDTree(points) + + # Query points + query_points = np.column_stack((xv.ravel(), yv.ravel())) + + # Find k nearest neighbors for smoother interpolation + k = min(4, len(points)) # Use 4 neighbors or less if not enough points + distances, indices = tree.query(query_points, k=k, workers=-1) # -1 uses all cores + + # Inverse distance weighting + epsilon = 1e-10 + weights = 1.0 / (distances + epsilon) + weights /= weights.sum(axis=1, keepdims=True) + + # Interpolate values + neighbor_values = values[mask_bounds][indices] + grid_field = (neighbor_values * weights).sum(axis=1).reshape(grid_resY, grid_res) + + # Plot + if show_colorbar or show_axes: + dpi = 300 + plt.imshow( + grid_field, + extent=[bounded_x_min, bounded_x_max, bounded_y_min, bounded_y_max], + cmap=cmap, + origin="lower", + aspect="equal", + **kwargs, + ) + if show_colorbar: + plt.colorbar() + if not show_axes: + plt.axis("off") + plt.savefig(output_filename + ".png", dpi=dpi, bbox_inches="tight", pad_inches=0) + plt.close() + else: + plt.imsave(output_filename + ".png", grid_field, cmap=cmap, origin="lower") + + def to_line( + self, + output_filename, + field_neon_dict, + start_point, + end_point, + resolution, + component=None, + radius=1.0, + **kwargs, + ): + """ + Extract field data along a line between start_point and end_point and save to a CSV file. + + This function performs two main steps: + 1. Extracts field data from field_neon_dict, handling components or computing magnitude. + 2. Interpolates the field values along a line defined by start_point and end_point, + then saves the results (coordinates and field values) to a CSV file. + + Parameters + ---------- + output_filename : str + The name of the output CSV file (without extension). Example: "velocity_profile". + field_neon_dict : dict + A dictionary containing the field data to extract, with a single key-value pair. + The key is the field name (e.g., "velocity"), and the value is the NEON data object + containing the field values. Example: {"velocity": velocity_neon}. + start_point : array_like + The starting point of the line in 3D space (e.g., [x0, y0, z0]). + Units must match the coordinate system used in the class (voxel units if untransformed, + or model units if scale/offset are applied). + end_point : array_like + The ending point of the line in 3D space (e.g., [x1, y1, z1]). + Units must match the coordinate system used in the class. + resolution : int + The number of points along the line where the field will be interpolated. + Example: 100 for 100 evenly spaced points. + component : int, optional + The specific component of the field to extract (e.g., 0 for x-component, 1 for y-component). + If None, the magnitude of the field is computed. Default is None. + radius : int + The specified distance (in units of the coordinate system) to prefilter and query for line plot + + Returns + ------- + None + The function writes the output to a CSV file and prints a confirmation message. + + Notes + ----- + - The output CSV file will contain columns: 'x', 'y', 'z', and the value of the field name (e.g., 'velocity_x' or 'velocity_magnitude'). + """ + + # Get the fields data from the NEON fields + assert len(field_neon_dict.keys()) == 1, "Error: This function is designed to plot a single field at a time." + fields_data = self.get_fields_data(field_neon_dict) + + # Check if the component is within the valid range + if component is None: + print("\tCreating csv plot of the field magnitude!") + cell_data = list(fields_data.values()) + squared = [comp**2 for comp in cell_data] + cell_data = np.sqrt(sum(squared)) + field_name = list(fields_data.keys())[0].split("_")[0] + "_magnitude" + + else: + assert component < max(self.field_name_cardinality_dict.values()), ( + f"Error: Component {component} is out of range for the provided fields." + ) + print(f"\tCreating csv plot for component {component} of the input field!") + field_name = list(fields_data.keys())[component] + cell_data = fields_data[field_name] + + # Plot each field in the dictionary + self._to_line_field( + f"{output_filename}_{field_name}", + cell_data, + start_point, + end_point, + resolution, + radius=radius, + **kwargs, + ) + print(f"\tLine Plot for field {field_name} saved as {output_filename}.csv") + + def _to_line_field( + self, + output_filename, + cell_data, + start_point, + end_point, + resolution, + radius, + **kwargs, + ): + """ + Helper function to create a line plot for a single field. + """ + import numpy as np + + # cell_points = self.coordinates[self.connectivity] # Shape: (M, K, 3), where M is num cells, K is nodes per cell + # centroids = np.mean(cell_points, axis=1) # Shape: (M, 3) + centroids = self.centroids + p0 = np.array(start_point, dtype=np.float32) + p1 = np.array(end_point, dtype=np.float32) + + # direction and parameter t for each centroid + d = p1 - p0 + L = np.linalg.norm(d) + d_unit = d / L + v = centroids - p0 + t = v.dot(d_unit) + closest = p0 + np.outer(t, d_unit) + perp_dist = np.linalg.norm(centroids - closest, axis=1) + + # optionally mask to [0,L] or a small perp-radius + mask = (t >= 0) & (t <= L) & (perp_dist <= radius) + t, data = t[mask], cell_data[mask] + + # sort by t + idx = np.argsort(t) + t_sorted = t[idx] + data_sorted = data[idx] + + # target samples + t_line = np.linspace(0, L, resolution) + + # 1D linear interpolation + vals_line = np.interp(t_line, t_sorted, data_sorted, left=np.nan, right=np.nan) + + # reconstruct (x,y,z) + line_xyz = p0[None, :] + t_line[:, None] * d_unit[None, :] + + # vectorized CSV dump + out = np.hstack([line_xyz, vals_line[:, None]]) + np.savetxt(output_filename + ".csv", out, delimiter=",", header="x,y,z,value", comments="") diff --git a/xlb/utils/utils.py b/xlb/utils/utils.py new file mode 100644 index 00000000..fe2e6b22 --- /dev/null +++ b/xlb/utils/utils.py @@ -0,0 +1,1088 @@ +""" +General-purpose utilities for XLB. + +Includes helpers for field downsampling, VTK/image/USD I/O, geometry +rotation, STL voxelization, Neon-to-JAX field transfer, and +physical-to-lattice unit conversion. +""" + +import numpy as np +import matplotlib.pylab as plt +from matplotlib import cm +from time import time +import pyvista as pv +from jax.image import resize +from jax import jit +import jax.numpy as jnp +from functools import partial +import trimesh +import warp as wp + +import os +import __main__ +import importlib +from contextlib import nullcontext + + +@partial(jit, static_argnums=(1, 2)) +def downsample_field(field, factor, method="bicubic"): + """ + Downsample a JAX array by a factor of `factor` along each axis. + + Parameters + ---------- + field : jax.numpy.ndarray + The input vector field to be downsampled. This should be a 3D or 4D JAX array where the last dimension is 2 or 3 (vector components). + factor : int + The factor by which to downsample the field. The dimensions of the field will be divided by this factor. + method : str, optional + The method to use for downsampling. Default is 'bicubic'. + + Returns + ------- + jax.numpy.ndarray + The downsampled field. + """ + if factor == 1: + return field + else: + new_shape = tuple(dim // factor for dim in field.shape[:-1]) + downsampled_components = [] + for i in range(field.shape[-1]): # Iterate over the last dimension (vector components) + resized = resize(field[..., i], new_shape, method=method) + downsampled_components.append(resized) + + return jnp.stack(downsampled_components, axis=-1) + + +def save_image(fld, timestep=None, prefix=None, **kwargs): + """ + Save an image of a field at a given timestep. + + Parameters + ---------- + timestep : int + The timestep at which the field is being saved. + fld : jax.numpy.ndarray + The field to be saved. This should be a 2D or 3D JAX array. If the field is 3D, the magnitude of the field will be calculated and saved. + prefix : str, optional + A prefix to be added to the filename. The filename will be the name of the main script file by default. + + Returns + ------- + None + + Notes + ----- + This function saves the field as an image in the PNG format. + The filename is based on the name of the main script file, the provided prefix, and the timestep number. + If the field is 3D, the magnitude of the field is calculated and saved. + The image is saved with the 'nipy_spectral' colormap and the origin set to 'lower'. + """ + if prefix is None: + fname = os.path.basename(__main__.__file__) + fname = os.path.splitext(fname)[0] + else: + fname = prefix + + if timestep is not None: + fname = fname + "_" + str(timestep).zfill(4) + + if len(fld.shape) > 3: + raise ValueError("The input field should be 2D!") + if len(fld.shape) == 3: + fld = np.sqrt(fld[0, ...] ** 2 + fld[1, ...] ** 2 + fld[2, ...] ** 2) + + plt.clf() + kwargs.pop("cmap", None) + plt.imsave(fname + ".png", fld.T, cmap=cm.nipy_spectral, origin="lower", **kwargs) + + +def save_fields_vtk(fields, timestep, output_dir=".", prefix="fields"): + """ + Save VTK fields to the specified directory. + + Parameters + ---------- + timestep (int): The timestep number to be associated with the saved fields. + fields (Dict[str, np.ndarray]): A dictionary of fields to be saved. Each field must be an array-like object + with dimensions (nx, ny) for 2D fields or (nx, ny, nz) for 3D fields, where: + - nx : int, number of grid points along the x-axis + - ny : int, number of grid points along the y-axis + - nz : int, number of grid points along the z-axis (for 3D fields only) + The key value for each field in the dictionary must be a string containing the name of the field. + output_dir (str, optional, default: '.'): The directory in which to save the VTK files. Defaults to the current directory. + prefix (str, optional, default: 'fields'): A prefix to be added to the filename. Defaults to 'fields'. + + Returns + ------- + None + + Notes + ----- + This function saves the VTK fields in the specified directory, with filenames based on the provided timestep number + and the filename. For example, if the timestep number is 10 and the file name is fields, the VTK file + will be saved as 'fields_0000010.vtk'in the specified directory. + + """ + # Assert that all fields have the same dimensions + for key, value in fields.items(): + if key == list(fields.keys())[0]: + dimensions = value.shape + else: + assert value.shape == dimensions, "All fields must have the same dimensions!" + + output_filename = os.path.join(output_dir, prefix + "_" + f"{timestep:07d}.vtk") + + # Add 1 to the dimensions tuple as we store cell values + dimensions = tuple([dim + 1 for dim in dimensions]) + + # Create a uniform grid + if value.ndim == 2: + dimensions = dimensions + (1,) + + grid = pv.ImageData(dimensions=dimensions) + + # Add the fields to the grid + for key, value in fields.items(): + grid[key] = value.flatten(order="F") + + # Save the grid to a VTK file + start = time() + grid.save(output_filename, binary=True) + print(f"Saved {output_filename} in {time() - start:.6f} seconds.") + + +def save_BCs_vtk(timestep, BCs, gridInfo, output_dir="."): + """ + Save boundary conditions as VTK format to the specified directory. + + Parameters + ---------- + timestep (int): The timestep number to be associated with the saved fields. + BCs (List[BC]): A list of boundary conditions to be saved. Each boundary condition must be an object of type BC. + + Returns + ------- + None + + Notes + ----- + This function saves the boundary conditions in the specified directory, with filenames based on the provided timestep number + and the filename. For example, if the timestep number is 10, the VTK file + will be saved as 'BCs_0000010.vtk'in the specified directory. + """ + + # Create a uniform grid + if gridInfo["nz"] == 0: + gridDimensions = (gridInfo["nx"] + 1, gridInfo["ny"] + 1, 1) + fieldDimensions = (gridInfo["nx"], gridInfo["ny"], 1) + else: + gridDimensions = (gridInfo["nx"] + 1, gridInfo["ny"] + 1, gridInfo["nz"] + 1) + fieldDimensions = (gridInfo["nx"], gridInfo["ny"], gridInfo["nz"]) + + grid = pv.ImageData(dimensions=gridDimensions) + + # Dictionary to keep track of encountered BC names + bcNamesCount = {} + + for bc in BCs: + bcName = bc.name + if bcName in bcNamesCount: + bcNamesCount[bcName] += 1 + else: + bcNamesCount[bcName] = 0 + bcName += f"_{bcNamesCount[bcName]}" + + if bc.isDynamic: + bcIndices, _ = bc.update_function(timestep) + else: + bcIndices = bc.indices + + # Convert indices to 1D indices + if gridInfo["dim"] == 2: + bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions[:-1], order="F") + else: + bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions, order="F") + + grid[bcName] = np.zeros(fieldDimensions, dtype=bool).flatten(order="F") + grid[bcName][bcIndices] = True + + # Save the grid to a VTK file + output_filename = os.path.join(output_dir, "BCs_" + f"{timestep:07d}.vtk") + + start = time() + grid.save(output_filename, binary=True) + print(f"Saved {output_filename} in {time() - start:.6f} seconds.") + + +def rotate_geometry(indices, origin, axis, angle): + """ + Rotates a voxelized mesh around a given axis. + + Parameters + ---------- + indices : array-like + The indices of the voxels in the mesh. + origin : array-like + The coordinates of the origin of the rotation axis. + axis : array-like + The direction vector of the rotation axis. This should be a 3-element sequence. + angle : float + The angle by which to rotate the mesh, in radians. + + Returns + ------- + tuple + The indices of the voxels in the rotated mesh. + + Notes + ----- + This function rotates the mesh by applying a rotation matrix to the voxel indices. The rotation matrix is calculated + using the axis-angle representation of rotations. The origin of the rotation axis is assumed to be at (0, 0, 0). + """ + indices_rotated = (jnp.array(indices).T - origin) @ axangle2mat(axis, angle) + origin + return tuple(jnp.rint(indices_rotated).astype("int32").T) + + +def voxelize_stl(stl_filename, length_lbm_unit=None, transformation_matrix=None, pitch=None): + """ + Converts an STL file to a voxelized mesh. + + Parameters + ---------- + stl_filename : str + The name of the STL file to be voxelized. + length_lbm_unit : float, optional + The unit length in LBM. Either this or 'pitch' must be provided. + transformation_matrix : array-like, optional + A transformation matrix to be applied to the mesh before voxelization. + pitch : float, optional + The pitch of the voxel grid. Either this or 'length_lbm_unit' must be provided. + + Returns + ------- + trimesh.VoxelGrid, float + The voxelized mesh and the pitch of the voxel grid. + + Notes + ----- + This function uses the trimesh library to load the STL file and voxelized the mesh. If a transformation matrix is + provided, it is applied to the mesh before voxelization. The pitch of the voxel grid is calculated based on the + maximum extent of the mesh and the provided lattice Boltzmann unit length, unless a pitch is provided directly. + """ + if length_lbm_unit is None and pitch is None: + raise ValueError("Either 'length_lbm_unit' or 'pitch' must be provided!") + mesh = trimesh.load_mesh(stl_filename, process=False) + length_phys_unit = mesh.extents.max() + if transformation_matrix is not None: + mesh.apply_transform(transformation_matrix) + if pitch is None: + pitch = length_phys_unit / length_lbm_unit + mesh_voxelized = mesh.voxelized(pitch=pitch) + return mesh_voxelized, pitch + + +def axangle2mat(axis, angle, is_normalized=False): + """Rotation matrix for rotation angle `angle` around `axis` + Parameters + ---------- + axis : 3 element sequence + vector specifying axis for rotation. + angle : scalar + angle of rotation in radians. + is_normalized : bool, optional + True if `axis` is already normalized (has norm of 1). Default False. + Returns + ------- + mat : array shape (3,3) + rotation matrix for specified rotation + Notes + ----- + From : https://github.com/matthew-brett/transforms3d + Ref : http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle + """ + x, y, z = axis + if not is_normalized: + n = jnp.sqrt(x * x + y * y + z * z) + x = x / n + y = y / n + z = z / n + c = jnp.cos(angle) + s = jnp.sin(angle) + C = 1 - c + xs = x * s + ys = y * s + zs = z * s + xC = x * C + yC = y * C + zC = z * C + xyC = x * yC + yzC = y * zC + zxC = z * xC + return jnp.array([ + [x * xC + c, xyC - zs, zxC + ys], + [xyC + zs, y * yC + c, yzC - xs], + [zxC - ys, yzC + xs, z * zC + c], + ]) + + +def jax_has_gpu_devices() -> bool: + """Return True if JAX can use at least one GPU (CUDA/ROCm) device.""" + import jax + + try: + return any(getattr(d, "platform", None) == "gpu" for d in jax.devices()) + except Exception: + return False + + +def warp_array_to_jax(warp_array): + """Convert a Warp array to a JAX array. + + ``warp.to_jax`` uses DLPack and expects JAX to support the same device + (e.g. CUDA). If Warp data is on GPU but only a **CPU** ``jaxlib`` is + installed, DLPack triggers ``RuntimeError: Unknown backend cuda``. In + that case we copy via the host with :meth:`warp.array.numpy`. + """ + dev = warp_array.device + if dev is not None and getattr(dev, "is_cuda", False) and not jax_has_gpu_devices(): + wp.synchronize() + return jnp.asarray(warp_array.numpy()) + return wp.to_jax(warp_array) + + +class ToJAX(object): + """Convert a Neon field to a JAX array via an intermediate Warp grid.""" + + def __init__(self, field_name, field_cardinality, grid_shape, store_precision=None): + """Initialise the Neon-to-JAX converter. + + Parameters + ---------- + field_name : str + The name of the field to be converted. + field_cardinality : int + The cardinality of the field to be converted. + grid_shape : tuple + The shape of the grid on which the field is defined. + store_precision : Precision, optional + Storage precision. Defaults to the global config value. + """ + from xlb.compute_backend import ComputeBackend + from xlb.grid import grid_factory + from xlb import DefaultConfig + + # Assign to self + self.field_name = field_name + self.field_cardinality = field_cardinality + self.grid_shape = grid_shape + self.compute_backend = DefaultConfig.default_backend + self.velocity_set = DefaultConfig.velocity_set + if store_precision is None: + self.store_precision = DefaultConfig.default_precision_policy.store_precision + self.store_dtype = DefaultConfig.default_precision_policy.store_precision.wp_dtype + + if self.compute_backend == ComputeBackend.NEON: + # Allocate warp fields for copying neon fields + # Use the warp backend to create dense fields for copying NEON dGrid fields + grid_dense = grid_factory(grid_shape, compute_backend=ComputeBackend.WARP) + self.warp_field = grid_dense.create_field(cardinality=self.field_cardinality, dtype=self.store_precision) + + def copy_neon_to_warp(self, neon_field): + """Convert a dense neon field to a warp field by copying.""" + import warp as wp + import neon + from typing import Any + + assert neon_field.get_grid().name == "dGrid", "to_warp only supports dense grids" + _d = self.velocity_set.d + + @neon.Container.factory("to_warp") + def container(src_field: Any, dst_field: Any, cardinality: wp.int32): + def loading_step(loader: neon.Loader): + loader.set_grid(src_field.get_grid()) + src_pn = loader.get_read_handle(src_field) + + @wp.func + def cloning(gridIdx: Any): + cIdx = wp.neon_global_idx(src_pn, gridIdx) + gx = wp.neon_get_x(cIdx) + gy = wp.neon_get_y(cIdx) + gz = wp.neon_get_z(cIdx) + + # XLB is flattening the z dimension in 3D, while neon uses the y dimension + if _d == 2: + gy, gz = gz, gy + + for card in range(cardinality): + value = wp.neon_read(src_pn, gridIdx, card) + dst_field[card, gx, gy, gz] = value + + loader.declare_kernel(cloning) + + return loading_step + + cardinality = neon_field.cardinality + c = container(neon_field, self.warp_field, cardinality) + c.run(0) + wp.synchronize() + return self.warp_field + + def __call__(self, field): + from xlb.compute_backend import ComputeBackend + import warp as wp + + if self.compute_backend == ComputeBackend.JAX: + return field + elif self.compute_backend == ComputeBackend.WARP: + return warp_array_to_jax(field) + elif self.compute_backend == ComputeBackend.NEON: + assert field.cardinality == self.field_cardinality, ( + f"Field cardinality mismatch! Expected {self.field_cardinality}, got {field.cardinality}!" + ) + return warp_array_to_jax(self.copy_neon_to_warp(field)) + + else: + raise ValueError("Unsupported compute backend!") + + +class UnitConvertor(object): + def __init__( + self, + velocity_lbm_unit: float, + velocity_physical_unit: float, + voxel_size_physical_unit: float, + density_physical_unit: float = 1.2041, + pressure_physical_unit: float = 1.101325e5, + ): + """ + Initialize the UnitConvertor object. + + Parameters + ---------- + velocity_lbm_unit : float + The reference velocity in lattice Boltzmann units. + velocity_physical_unit : float + The reference velocity in physical units (e.g., m/s). + voxel_size_physical_unit : float + The size of a voxel in physical units (e.g., meters). + density_physical_unit : float, optional + The reference density in physical units (e.g., kg/m^3). Default is 1.2041 (density of air at room temperature). + pressure_physical_unit : float, optional + The reference pressure in physical units (e.g., Pascals). Default is 1.101325e5 (atmospheric pressure at sea level). + """ + + self.voxel_size = voxel_size_physical_unit + self.velocity_lbm_unit = velocity_lbm_unit + self.velocity_phys_unit = velocity_physical_unit + + # Reference density and pressure in physical units + self.reference_density = density_physical_unit + self.referece_pressure = pressure_physical_unit + + @property + def time_step_physical(self): + return self.voxel_size * self.velocity_lbm_unit / self.velocity_phys_unit + + @property + def reference_length(self): + return self.voxel_size + + @property + def reference_time(self): + return self.time_step_physical + + @property + def reference_velocity(self): + return self.reference_length / self.reference_time + + def length_to_lbm(self, length_phys): + return length_phys / self.reference_length + + def length_to_physical(self, length_lbm): + return length_lbm * self.reference_length + + def time_to_lbm(self, time_phys): + return time_phys / self.reference_time + + def time_to_physical(self, time_lbm): + return time_lbm * self.reference_time + + def density_to_lbm(self, rho_phys): + return rho_phys / self.reference_density + + def density_to_physical(self, rho_lbm): + return rho_lbm * self.reference_density + + def velocity_to_lbm(self, velocity_phys): + return velocity_phys / self.reference_velocity + + def velocity_to_physical(self, velocity_lbm): + return velocity_lbm * self.reference_velocity + + def viscosity_to_lbm(self, viscosity_phys): + return viscosity_phys * (self.reference_time / (self.reference_length**2)) + + def viscosity_to_physical(self, viscosity_lbm): + return viscosity_lbm * (self.reference_length**2 / self.reference_time) + + def pressure_to_lbm(self, pressure_phys): + pressure_perturbation = pressure_phys - self.reference_pressure + return pressure_perturbation / self.reference_density / self.reference_velocity**2 + + def pressure_to_physical(self, pressure_lbm): + pressure_perturbation = pressure_lbm - 1.0 / 3.0 + return self.referece_pressure + pressure_perturbation * self.reference_density * (self.reference_velocity**2) + + +@wp.kernel +def get_color( + low: float, + high: float, + values: wp.array(dtype=float), + out_color: wp.array(dtype=wp.vec3), +): + """ + Colorize scalars using a rainbow color map. + + Parameters + ---------- + low : float + The lower bound of the color map. + high : float + The upper bound of the color map. + values : wp.array(dtype=float) + The values to colorize. + out_color : wp.array(dtype=wp.vec3) + The output colors. + + Returns + ------- + None + """ + tid = wp.tid() + v = values[tid] + r = 1.0 + g = 1.0 + b = 1.0 + if v < low: + v = low + if v > high: + v = high + dv = high - low + if v < (low + 0.25 * dv): + r = 0.0 + g = 4.0 * (v - low) / dv + elif v < (low + 0.5 * dv): + r = 0.0 + b = 1.0 + 4.0 * (low + 0.25 * dv - v) / dv + elif v < (low + 0.75 * dv): + r = 4.0 * (v - low - 0.5 * dv) / dv + b = 0.0 + else: + g = 1.0 + 4.0 * (low + 0.75 * dv - v) / dv + b = 0.0 + out_color[tid] = wp.vec3(r, g, b) + + +def colorize_scalars(scalars, device=None, value_range=None, percentiles=(5, 95), target=None): + """ + Colorize scalars using a rainbow color map. + + Parameters + ---------- + scalars : wp.array(dtype=float) + The scalars to colorize. + device : wp.Device, optional + The device to use for the colorization. + value_range : tuple, optional + The value range to use for the colorization. + percentiles : tuple, optional + The percentiles to use for the colorization. + target : wp.array(dtype=wp.vec3), optional + The target array to store the colors. + + Returns + ------- + wp.array(dtype=wp.vec3) + The colors. + tuple + The value range used for the colorization. + """ + if device is None: + device = scalars.device + colors = target if target is not None else wp.empty(scalars.shape[0], dtype=wp.vec3, device=device) + if value_range is None: + scalars_np = scalars.numpy() + low = float(np.percentile(scalars_np, percentiles[0])) + high = float(np.percentile(scalars_np, percentiles[1])) + else: + low = float(value_range[0]) + high = float(value_range[1]) + if abs(high - low) < 1e-6: + high = low + 1e-6 + wp.launch( + kernel=get_color, + dim=scalars.shape[0], + inputs=(low, high, scalars), + outputs=(colors,), + device=device, + ) + return colors, (low, high) + + +def _normalize_clip_values(values): + """ + Internal function to normalize clip values. + + Parameters + ---------- + values : tuple, optional + The clip values to normalize. + + Returns + ------- + tuple + The normalized clip values. + """ + if values is None: + return (0, 0, 0) + if isinstance(values, (int, float)): + v = int(values) + return (v, v, v) + seq = tuple(int(x) for x in values) + if len(seq) != 3: + raise ValueError("clip values must have length 3") + return seq + + +def _slice_velocity_field(field, clip_lower, clip_upper): + """ + Internal function to slice a velocity field. + + Parameters + ---------- + field : wp.array(dtype=float) + The velocity field to slice. + clip_lower : tuple, optional + The lower clip values. + clip_upper : tuple, optional + The upper clip values. + + Returns + ------- + wp.array(dtype=float) + The sliced velocity field. + """ + lower = _normalize_clip_values(clip_lower) + upper = _normalize_clip_values(clip_upper) + slices = [slice(None)] + for l, u in zip(lower, upper): + start = l if l > 0 else None + stop = -u if u > 0 else None + slices.append(slice(start, stop)) + return field[tuple(slices)] + + +def _clone_to_device(array, device): + """ + Internal function to clone an array to a device. + + Parameters + ---------- + array : wp.array(dtype=float) + The array to clone. + device : wp.Device + The device to clone the array to. + + Returns + ------- + wp.array(dtype=float) + The cloned array. + """ + if hasattr(array, "device") and array.device == device: + return array + return wp.clone(array, device=device) + + +def _get_usd_modules(): + """ + Internal function to get the USD modules. + + Returns + ------- + tuple + The USD modules (UsdGeom, Vt). + """ + UsdGeom = importlib.import_module("pxr.UsdGeom") + Vt = importlib.import_module("pxr.Vt") + return UsdGeom, Vt + + +def save_usd_vorticity( + timestep, + post_process_interval, + bc_mask, + f_current, + grid_shape, + usd_mesh, + vorticity_operator, + precision_policy, + vorticity_threshold, + usd_stage, + device=None, + clip_lower=None, + clip_upper=None, + color_percentiles=(5, 95), + color_range=None, +): + """ + Save the vorticity field to a USD mesh. + + Parameters + ---------- + timestep : int + The timestep. + post_process_interval : int + The post-process interval. + bc_mask : wp.array(dtype=bool) + The boundary mask. + f_current : wp.array(dtype=float) + The current field. + grid_shape : tuple + The shape of the grid. + usd_mesh : pxr.Usd.Mesh + The USD mesh to save the vorticity field to. + vorticity_operator : xlb.operator.vorticity.VorticityOperator + The vorticity operator. + precision_policy : xlb.precision_policy.PrecisionPolicy + The precision policy. + vorticity_threshold : float + The vorticity threshold. + usd_stage : pxr.Usd.Stage + The USD stage. + device : wp.Device, optional + The device to use for the computation. + clip_lower : tuple, optional + The lower clip values. + clip_upper : tuple, optional + The upper clip values. + color_percentiles : tuple, optional + The percentiles to use for the colorization. + color_range : tuple, optional + The value range to use for the colorization. + + Returns + ------- + None + """ + from xlb.compute_backend import ComputeBackend + from xlb.operator.macroscopic import Macroscopic + from xlb.operator.postprocess import GridToPoint + import xlb + + if device is None: + device = getattr(f_current, "device", "cpu") + clip_lower = _normalize_clip_values(clip_lower) + clip_upper = _normalize_clip_values(clip_upper) + f_current_dev = _clone_to_device(f_current, device) + bc_mask_dev = _clone_to_device(bc_mask, device) + with wp.ScopedDevice(device): + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.WARP) + macro_wp = Macroscopic(compute_backend=ComputeBackend.WARP, precision_policy=precision_policy, velocity_set=velocity_set) + rho = wp.zeros((1, *grid_shape), dtype=wp.float32, device=device) + u = wp.zeros((3, *grid_shape), dtype=wp.float32, device=device) + rho, u = macro_wp(f_current_dev, rho, u) + u = _slice_velocity_field(u, clip_lower, clip_upper) + vorticity = wp.zeros((3, *u.shape[1:]), dtype=wp.float32, device=device) + vorticity_magnitude = wp.zeros((1, *u.shape[1:]), dtype=wp.float32, device=device) + vorticity, vorticity_magnitude = vorticity_operator(u, bc_mask_dev, vorticity, vorticity_magnitude) + max_verts = grid_shape[0] * grid_shape[1] * grid_shape[2] * 5 + max_tris = grid_shape[0] * grid_shape[1] * grid_shape[2] * 3 + mc = wp.MarchingCubes(nx=u.shape[1], ny=u.shape[2], nz=u.shape[3], max_verts=max_verts, max_tris=max_tris, device=device) + mc.surface(vorticity_magnitude[0], vorticity_threshold) + if mc.verts.shape[0] == 0: + print(f"Warning: No vertices found for vorticity at timestep {timestep}.") + return + grid_to_point_op = GridToPoint(precision_policy=precision_policy, compute_backend=ComputeBackend.WARP) + scalars = wp.zeros(mc.verts.shape[0], dtype=wp.float32, device=device) + scalars = grid_to_point_op(vorticity_magnitude, mc.verts, scalars) + colors, value_range = colorize_scalars( + scalars, + device=device, + value_range=color_range, + percentiles=color_percentiles, + ) + vertices = mc.verts.numpy() + indices = mc.indices.numpy() + colors_np = colors.numpy() + tri_count = len(indices) // 3 + time_code = timestep // post_process_interval + UsdGeom, _ = _get_usd_modules() + usd_mesh.GetPointsAttr().Set(vertices.tolist(), time=time_code) + usd_mesh.GetFaceVertexCountsAttr().Set([3] * tri_count, time=time_code) + usd_mesh.GetFaceVertexIndicesAttr().Set(indices.tolist(), time=time_code) + usd_mesh.GetDisplayColorAttr().Set(colors_np.tolist(), time=time_code) + UsdGeom.Primvar(usd_mesh.GetDisplayColorAttr()).SetInterpolation("vertex") + print(f"Vorticity visualization at timestep {timestep}:") + print(f" Number of vertices: {len(vertices)}") + print(f" Number of triangles: {tri_count}") + print(f" Vorticity range: [{value_range[0]:.6f}, {value_range[1]:.6f}]") + + +def save_usd_q_criterion( + timestep, + post_process_interval, + bc_mask, + f_current, + grid_shape, + usd_mesh, + q_criterion_operator, + precision_policy, + q_threshold, + usd_stage, + device=None, + clip_lower=None, + clip_upper=None, + color_range=(0.0, 0.1), + color_percentiles=None, +): + """ + Save the Q-criterion field to a USD mesh. + + Parameters + ---------- + timestep : int + The timestep. + post_process_interval : int + The post-process interval. + bc_mask : wp.array(dtype=bool) + The boundary mask. + f_current : wp.array(dtype=float) + The current field. + grid_shape : tuple + The shape of the grid. + usd_mesh : pxr.Usd.Mesh + The USD mesh to save the Q-criterion field to. + q_criterion_operator : xlb.operator.q_criterion.QCriterionOperator + The Q-criterion operator. + precision_policy : xlb.precision_policy.PrecisionPolicy + The precision policy. + q_threshold : float + The Q-criterion threshold. + usd_stage : pxr.Usd.Stage + The USD stage. + device : wp.Device, optional + The device to use for the computation. + clip_lower : tuple, optional + The lower clip values. + clip_upper : tuple, optional + The upper clip values. + color_range : tuple, optional + The value range to use for the colorization. + color_percentiles : tuple, optional + The percentiles to use for the colorization. + + Returns + ------- + None + """ + from xlb.compute_backend import ComputeBackend + from xlb.operator.macroscopic import Macroscopic + from xlb.operator.postprocess import GridToPoint + import xlb + + if device is None: + device = getattr(f_current, "device", "cpu") + clip_lower = _normalize_clip_values(clip_lower) + clip_upper = _normalize_clip_values(clip_upper) + f_current_dev = _clone_to_device(f_current, device) + bc_mask_dev = _clone_to_device(bc_mask, device) + with wp.ScopedDevice(device): + velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, compute_backend=ComputeBackend.WARP) + macro_wp = Macroscopic(compute_backend=ComputeBackend.WARP, precision_policy=precision_policy, velocity_set=velocity_set) + rho = wp.zeros((1, *grid_shape), dtype=wp.float32, device=device) + u = wp.zeros((3, *grid_shape), dtype=wp.float32, device=device) + rho, u = macro_wp(f_current_dev, rho, u) + u = _slice_velocity_field(u, clip_lower, clip_upper) + norm_mu = wp.zeros((1, *u.shape[1:]), dtype=wp.float32, device=device) + q_field = wp.zeros((1, *u.shape[1:]), dtype=wp.float32, device=device) + norm_mu, q_field = q_criterion_operator(u, bc_mask_dev, norm_mu, q_field) + max_verts = grid_shape[0] * grid_shape[1] * grid_shape[2] * 5 + max_tris = grid_shape[0] * grid_shape[1] * grid_shape[2] * 3 + mc = wp.MarchingCubes(nx=u.shape[1], ny=u.shape[2], nz=u.shape[3], max_verts=max_verts, max_tris=max_tris, device=device) + mc.surface(q_field[0], q_threshold) + if mc.verts.shape[0] == 0: + print(f"Warning: No vertices found for Q-criterion at timestep {timestep}.") + return + grid_to_point_op = GridToPoint(precision_policy=precision_policy, compute_backend=ComputeBackend.WARP) + scalars = wp.zeros(mc.verts.shape[0], dtype=wp.float32, device=device) + scalars = grid_to_point_op(norm_mu, mc.verts, scalars) + if color_range is None: + percentiles = color_percentiles if color_percentiles is not None else (5, 95) + colors, used_range = colorize_scalars( + scalars, + device=device, + percentiles=percentiles, + ) + else: + colors, used_range = colorize_scalars( + scalars, + device=device, + value_range=color_range, + percentiles=color_percentiles if color_percentiles is not None else (5, 95), + ) + vertices = mc.verts.numpy() + indices = mc.indices.numpy() + colors_np = colors.numpy() + tri_count = len(indices) // 3 + time_code = timestep // post_process_interval + UsdGeom, _ = _get_usd_modules() + usd_mesh.GetPointsAttr().Set(vertices.tolist(), time=time_code) + usd_mesh.GetFaceVertexCountsAttr().Set([3] * tri_count, time=time_code) + usd_mesh.GetFaceVertexIndicesAttr().Set(indices.tolist(), time=time_code) + usd_mesh.GetDisplayColorAttr().Set(colors_np.tolist(), time=time_code) + UsdGeom.Primvar(usd_mesh.GetDisplayColorAttr()).SetInterpolation("vertex") + print(f"Q-criterion visualization at timestep {timestep}:") + print(f" Number of vertices: {len(vertices)}") + print(f" Number of triangles: {tri_count}") + print(f" Scalar range: [{used_range[0]:.6f}, {used_range[1]:.6f}]") + + +def update_usd_lagrangian_parts( + timestep, + post_process_interval, + vertices_wp, + parts, + vertex_offset=None, + lag_forces=None, + force_component=0, + device=None, +): + """ + Update the USD lagrangian parts. The lagrangian parts are updated with the vertices and faces. The forces are colorized using a rainbow color map. + The color map is a linear interpolation between the lower and upper bounds. + Parameters + ---------- + timestep : int + The timestep. + post_process_interval : int + The post-process interval. + vertices_wp : wp.array(dtype=float) + The vertices of the lagrangian parts. + parts : list + The parts of the lagrangian mesh. + vertex_offset : tuple, optional + The vertex offset. + lag_forces : wp.array(dtype=float), optional + The forces of the lagrangian parts. + force_component : int, optional + The component of the forces to colorize. + device : wp.Device, optional + The device to use for the computation. + + Returns + ------- + None + """ + vertices_np = vertices_wp.numpy() + if vertex_offset is not None: + offset_array = np.asarray(vertex_offset, dtype=np.float64) + if offset_array.ndim == 0: + offset_array = np.repeat(offset_array, 3) + vertices_np = vertices_np - offset_array + time_code = timestep // post_process_interval + lag_forces_np = lag_forces.numpy() if lag_forces is not None else None + if device is None: + device = getattr(vertices_wp, "device", None) + if device is None: + device = "cpu" + UsdGeom, Vt = _get_usd_modules() + for part in parts: + start = int(part.get("start", 0)) + end = int(part.get("end", vertices_np.shape[0])) + faces = np.asarray(part["faces"], dtype=np.int32) + if part.get("shift_indices", True): + faces = faces - start + usd_mesh = part["usd_mesh"] + part_vertices = vertices_np[start:end] + tri_count = len(faces) + usd_mesh.GetPointsAttr().Set(part_vertices.tolist(), time=time_code) + usd_mesh.GetFaceVertexCountsAttr().Set(Vt.IntArray([3] * tri_count), time=time_code) + usd_mesh.GetFaceVertexIndicesAttr().Set(Vt.IntArray(faces.flatten().tolist()), time=time_code) + if lag_forces_np is not None and part.get("colorize", False): + component = int(part.get("force_component", force_component)) + forces_np = lag_forces_np[start:end, component] + context = wp.ScopedDevice(device) if device is not None else nullcontext() + with context: + force_wp = wp.from_numpy(forces_np.astype(np.float32), dtype=wp.float32, device=device) + colors_wp = wp.zeros(part_vertices.shape[0], dtype=wp.vec3, device=device) + if part.get("color_percentiles") is not None: + colors_wp, used_range = colorize_scalars( + force_wp, + device=device, + percentiles=part["color_percentiles"], + target=colors_wp, + ) + else: + value_range = part.get("color_range") + if value_range is None: + low = float(np.min(forces_np)) + high = float(np.max(forces_np)) + if abs(high - low) < 1e-6: + high = low + 1e-6 + value_range = (low, high) + colors_wp, used_range = colorize_scalars( + force_wp, + device=device, + value_range=value_range, + target=colors_wp, + ) + colors_np = colors_wp.numpy() + usd_mesh.GetDisplayColorAttr().Set(colors_np.tolist(), time=time_code) + UsdGeom.Primvar(usd_mesh.GetDisplayColorAttr()).SetInterpolation("vertex") + print(f"Lagrangian meshes updated at timestep {timestep}") + + +def plot_object_placement(vertices_wp, grid_shape, filename, title, object_label="Object"): + """ + Plot the object placement. + The object placement is plotted as a polygon in the domain. The domain is the bounding box of the vertices. + The plot is saved as a PNG file. + + Parameters + ---------- + vertices_wp : wp.array(dtype=float) + The vertices of the object. + grid_shape : tuple + The shape of the grid. + filename : str + The filename to save the plot to. + title : str + The title of the plot. + object_label : str, optional + The label of the object. + + Returns + ------- + None + """ + verts = vertices_wp.numpy() + obj_min = verts.min(axis=0) + obj_max = verts.max(axis=0) + plt.figure(figsize=(10, 5)) + domain_x = [0, grid_shape[0], grid_shape[0], 0, 0] + domain_y = [0, 0, grid_shape[1], grid_shape[1], 0] + plt.plot(domain_x, domain_y, "k-", linewidth=1, label="Domain") + poly_x = [obj_min[0], obj_max[0], obj_max[0], obj_min[0], obj_min[0]] + poly_y = [obj_min[1], obj_min[1], obj_max[1], obj_max[1], obj_min[1]] + plt.plot(poly_x, poly_y, "r-", linewidth=2, label=object_label) + plt.title(title) + plt.xlabel("X") + plt.ylabel("Y") + plt.grid(True, linestyle="--", alpha=0.3) + plt.axis("equal") + plt.legend() + plt.savefig(filename, dpi=150, bbox_inches="tight") + plt.close() diff --git a/xlb/velocity_set/__init__.py b/xlb/velocity_set/__init__.py new file mode 100644 index 00000000..5b7b737f --- /dev/null +++ b/xlb/velocity_set/__init__.py @@ -0,0 +1,4 @@ +from xlb.velocity_set.velocity_set import VelocitySet +from xlb.velocity_set.d2q9 import D2Q9 +from xlb.velocity_set.d3q19 import D3Q19 +from xlb.velocity_set.d3q27 import D3Q27 diff --git a/xlb/velocity_set/d2q9.py b/xlb/velocity_set/d2q9.py new file mode 100644 index 00000000..69dad633 --- /dev/null +++ b/xlb/velocity_set/d2q9.py @@ -0,0 +1,24 @@ +# Description: Lattice class for 2D D2Q9 lattice. + +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet + + +class D2Q9(VelocitySet): + """ + Velocity Set for 2D D2Q9 lattice. + + D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the + Lattice Boltzmann Method for simulating fluid flows in two dimensions. + """ + + def __init__(self, precision_policy, compute_backend): + # Construct the velocity vectors and weights + cx = [0, 0, 0, 1, -1, 1, -1, 1, -1] + cy = [0, 1, -1, 0, 1, -1, 0, 1, -1] + c = np.array(tuple(zip(cx, cy))).T + w = np.array([4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 9, 1 / 36, 1 / 36]) + + # Call the parent constructor + super().__init__(2, 9, c, w, precision_policy=precision_policy, compute_backend=compute_backend) diff --git a/xlb/velocity_set/d3q19.py b/xlb/velocity_set/d3q19.py new file mode 100644 index 00000000..c2a9ab4c --- /dev/null +++ b/xlb/velocity_set/d3q19.py @@ -0,0 +1,30 @@ +# Description: Lattice class for 3D D3Q19 lattice. + +import itertools +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet + + +class D3Q19(VelocitySet): + """ + Velocity Set for 3D D3Q19 lattice. + + D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the + Lattice Boltzmann Method for simulating fluid flows in three dimensions. + """ + + def __init__(self, precision_policy, compute_backend): + # Construct the velocity vectors and weights + c = np.array([ci for ci in itertools.product([0, -1, 1], repeat=3) if np.sum(np.abs(ci)) <= 2]).T + w = np.zeros(19) + for i in range(19): + if np.sum(np.abs(c[:, i])) == 0: + w[i] = 1.0 / 3.0 + elif np.sum(np.abs(c[:, i])) == 1: + w[i] = 1.0 / 18.0 + elif np.sum(np.abs(c[:, i])) == 2: + w[i] = 1.0 / 36.0 + + # Initialize the lattice + super().__init__(3, 19, c, w, precision_policy=precision_policy, compute_backend=compute_backend) diff --git a/xlb/velocity_set/d3q27.py b/xlb/velocity_set/d3q27.py new file mode 100644 index 00000000..8110fd6c --- /dev/null +++ b/xlb/velocity_set/d3q27.py @@ -0,0 +1,32 @@ +# Description: Lattice class for 3D D3Q27 lattice. + +import itertools +import numpy as np + +from xlb.velocity_set.velocity_set import VelocitySet + + +class D3Q27(VelocitySet): + """ + Velocity Set for 3D D3Q27 lattice. + + D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the + Lattice Boltzmann Method for simulating fluid flows in three dimensions. + """ + + def __init__(self, precision_policy, compute_backend): + # Construct the velocity vectors and weights + c = np.array(list(itertools.product([0, -1, 1], repeat=3))).T + w = np.zeros(27) + for i in range(27): + if np.sum(np.abs(c[:, i])) == 0: + w[i] = 8.0 / 27.0 + elif np.sum(np.abs(c[:, i])) == 1: + w[i] = 2.0 / 27.0 + elif np.sum(np.abs(c[:, i])) == 2: + w[i] = 1.0 / 54.0 + elif np.sum(np.abs(c[:, i])) == 3: + w[i] = 1.0 / 216.0 + + # Initialize the Lattice + super().__init__(3, 27, c, w, precision_policy=precision_policy, compute_backend=compute_backend) diff --git a/xlb/velocity_set/velocity_set.py b/xlb/velocity_set/velocity_set.py new file mode 100644 index 00000000..6f9634b2 --- /dev/null +++ b/xlb/velocity_set/velocity_set.py @@ -0,0 +1,265 @@ +""" +Base velocity-set class for the Lattice Boltzmann Method. + +Defines lattice directions, weights, and derived properties (opposite +indices, moments, etc.) for any DdQq stencil. Backend-specific constants +(Warp vectors, JAX arrays, Neon lattice objects) are initialised lazily. +""" + +import math +import numpy as np +import warp as wp +import jax.numpy as jnp +import jax + +from xlb import DefaultConfig +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy + + +class VelocitySet(object): + """ + Base class for the velocity set of the Lattice Boltzmann Method (LBM), e.g. D2Q9, D3Q27, etc. + + Parameters + ---------- + d: int + The dimension of the lattice. + q: int + The number of velocities of the lattice. + c: numpy.ndarray + The velocity vectors of the lattice. Shape: (q, d) + w: numpy.ndarray + The weights of the lattice. Shape: (q,) + """ + + def __init__(self, d, q, c, w, precision_policy, compute_backend): + # Store the dimension and the number of velocities + self.d = d + self.q = q + self.precision_policy = precision_policy + self.compute_backend = compute_backend + + # Updating JAX config in case fp64 is requested + if compute_backend == ComputeBackend.JAX and (precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32): + jax.config.update("jax_enable_x64", True) + + # Create all properties in NumPy first + self._init_numpy_properties(c, w) + + # Convert properties to backend-specific format + if self.compute_backend == ComputeBackend.WARP: + self._init_warp_properties() + elif self.compute_backend == ComputeBackend.NEON: + self._init_neon_properties() + elif self.compute_backend == ComputeBackend.JAX: + self._init_jax_properties() + else: + raise ValueError(f"Unsupported compute backend: {self.compute_backend}") + + # Set up backend-specific constants + self._init_backend_constants() + + def _init_numpy_properties(self, c, w): + """ + Initialize all properties in NumPy first. + """ + self._c = np.array(c) + self._w = np.array(w) + self._opp_indices = self._construct_opposite_indices() + self._cc = self._construct_lattice_moment() + self._c_float = self._c.astype(np.float64) + self._qi = self._construct_qi() + + # Constants in NumPy + self.cs = np.float64(math.sqrt(3) / 3.0) + self.cs2 = np.float64(1.0 / 3.0) + self.inv_cs2 = np.float64(3.0) + + # Indices + self.main_indices = self._construct_main_indices() + self.right_indices = self._construct_right_indices() + self.left_indices = self._construct_left_indices() + self.center_index = self._get_center_index() + + def _init_warp_properties(self): + """ + Convert NumPy properties to Warp-specific properties. + """ + dtype = self.precision_policy.compute_precision.wp_dtype + self.c = wp.constant(wp.mat((self.d, self.q), dtype=wp.int32)(self._c)) + self.w = wp.constant(wp.vec(self.q, dtype=dtype)(self._w)) + self.opp_indices = wp.constant(wp.vec(self.q, dtype=wp.int32)(self._opp_indices)) + self.cc = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._cc)) + self.c_float = wp.constant(wp.mat((self.d, self.q), dtype=dtype)(self._c_float)) + self.qi = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._qi)) + + def _init_neon_properties(self): + """ + Convert NumPy properties to Neon-specific properties which are identical to Warp. + """ + self._init_warp_properties() + + def _init_jax_properties(self): + """ + Convert NumPy properties to JAX-specific properties. + """ + dtype = self.precision_policy.compute_precision.jax_dtype + self.c = jnp.array(self._c, dtype=jnp.int32) + self.w = jnp.array(self._w, dtype=dtype) + self.opp_indices = jnp.array(self._opp_indices, dtype=jnp.int32) + self.cc = jnp.array(self._cc, dtype=dtype) + self.c_float = jnp.array(self._c_float, dtype=dtype) + self.qi = jnp.array(self._qi, dtype=dtype) + + def _init_backend_constants(self): + """ + Initialize the constants for the backend. + """ + if self.compute_backend == ComputeBackend.WARP: + dtype = self.precision_policy.compute_precision.wp_dtype + self.cs = wp.constant(dtype(self.cs)) + self.cs2 = wp.constant(dtype(self.cs2)) + self.inv_cs2 = wp.constant(dtype(self.inv_cs2)) + elif self.compute_backend == ComputeBackend.JAX: + dtype = self.precision_policy.compute_precision.jax_dtype + self.cs = jnp.array(self.cs, dtype=dtype) + self.cs2 = jnp.array(self.cs2, dtype=dtype) + self.inv_cs2 = jnp.array(self.inv_cs2, dtype=dtype) + + def warp_lattice_vec(self, dtype): + return wp.vec(len(self.c), dtype=dtype) + + def warp_u_vec(self, dtype): + return wp.vec(self.d, dtype=dtype) + + def warp_stream_mat(self, dtype): + return wp.mat((self.q, self.d), dtype=dtype) + + def _construct_qi(self): + # Qi = cc - cs^2*I + dim = self.d + Qi = self._cc.copy() + if dim == 3: + diagonal, offdiagonal = (0, 3, 5), (1, 2, 4) + elif dim == 2: + diagonal, offdiagonal = (0, 2), (1,) + else: + raise ValueError(f"dim = {dim} not supported") + + # multiply off-diagonal elements by 2 because the Q tensor is symmetric + Qi[:, diagonal] += -1.0 / 3.0 + Qi[:, offdiagonal] *= 2.0 + return Qi + + def _construct_lattice_moment(self): + """ + This function constructs the moments of the lattice. + + The moments are the products of the velocity vectors, which are used in the computation of + the equilibrium distribution functions and the collision operator in the Lattice Boltzmann + Method (LBM). + + Returns + ------- + cc: numpy.ndarray + The moments of the lattice. + """ + c = self._c.T + # Counter for the loop + cntr = 0 + c = self._c.T + # nt: number of independent elements of a symmetric tensor + nt = self.d * (self.d + 1) // 2 + cc = np.zeros((self.q, nt)) + cntr = 0 + for a in range(self.d): + for b in range(a, self.d): + cc[:, cntr] = c[:, a] * c[:, b] + cntr += 1 + return cc + + def _construct_opposite_indices(self): + """ + This function constructs the indices of the opposite velocities for each velocity. + + The opposite velocity of a velocity is the velocity that has the same magnitude but the + opposite direction. + + Returns + ------- + opposite: numpy.ndarray + The indices of the opposite velocities. + """ + c = self._c.T + return np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)]) + + def _construct_main_indices(self): + """ + This function constructs the indices of the main velocities. + + The main velocities are the velocities that have a magnitude of 1 in lattice units. + + Returns + ------- + numpy.ndarray + The indices of the main velocities. + """ + c = self._c.T + if self.d == 2: + return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0] + elif self.d == 3: + return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0] + + def _construct_right_indices(self): + """ + This function constructs the indices of the velocities that point in the positive + x-direction. + + Returns + ------- + numpy.ndarray + The indices of the right velocities. + """ + return np.nonzero(self._c.T[:, 0] == 1)[0] + + def _construct_left_indices(self): + """ + This function constructs the indices of the velocities that point in the negative + x-direction. + + Returns + ------- + numpy.ndarray + The indices of the left velocities. + """ + return np.nonzero(self._c.T[:, 0] == -1)[0] + + def _get_center_index(self): + """ + This function returns the index of the center point in the lattice associated with (0,0,0) + + Returns + ------- + numpy.ndarray + The index of the zero lattice velocity. + """ + arr = self._c.T + if self.d == 2: + target = np.array([0, 0]) + else: + target = np.array([0, 0, 0]) + match = np.all(arr == target, axis=1) + return int(np.nonzero(match)[0][0]) + + def __str__(self): + """ + This function returns the name of the lattice in the format of DxQy. + """ + return self.__repr__() + + def __repr__(self): + """ + This function returns the name of the lattice in the format of DxQy. + """ + return "D{}Q{}".format(self.d, self.q)