現実世界を動かす技術をシェアしよう！～身近なものづくりから最先端テクノロジーまで大募集！～

LBM流体解析ツールXLBの活用【01:JAXcpu導入方法と動作確認】

Last updated at 2026-10-12Posted at 2026-09-22

この文書では、WSL2のUbuntu24.04上でLBM（格子ボルツマン法）流体解析ツールXLBを活用する方法を紹介します。

「XLB（Accelerated Lattice Boltzmann）」は、Autodeskが公開している、GPUによる高速計算と機械学習との連携を重視したPython製LBM（格子ボルツマン法）ライブラリです。
2次元・3次元の流体解析に対応し、標準的なBGKモデルや高レイノルズ数流れに適したKBCモデルなどを利用できます。
最大の特徴は、NVIDIA Warp、JAX、Neonなどの計算バックエンドを利用できる点です。特にWarpを用いることで、PythonコードからGPUを活用した高速なLBM計算を比較的容易に実装できます。
また、計算過程を微分可能（Differentiable） に設計しているため、流体解析と深層学習を組み合わせたPhysicsML（Physics-based Machine Learning）、形状最適化、逆問題などへの応用も期待されています。
さらに、大規模GPU計算、マルチGPU、格子細分化にも対応しており、従来の「流体を計算するだけのLBMソフトウェア」ではなく、高速CFDとAI・機械学習を統合する次世代型のLBM基盤として位置付けられます。
Pythonベースで拡張しやすいため、LBMの学習から先端的なPhysics AI研究まで幅広く活用できるライブラリです。

XLB公式サイト：https://www.research.autodesk.com/publications/xlb-differentiable-massively-parallel-lattice-boltzmann-library-python/
XLB公式GitHubリポジトリ

なお、XLBに関する技術的な詳細を記述した論文が、以下で公開されています。
「XLB: A differentiable massively parallel lattice Boltzmann library in Python」
https://arxiv.org/abs/2311.16080

この解説は、以下のQiita記事を、全面的に参考にしています。ありがとうございます。

XLBの重要機能の解説

「XLBには何ができるのか」 という機能分析の観点で整理すると、XLBは単なるLBM計算ライブラリではなく、JAXを中心として、CPU・GPU・マルチGPUまで同じPythonコードで扱える、拡張性の高いLBM計算基盤と捉えることができます。

主な特徴（Key Features）

複数のバックエンドに対応：JAX、NVIDIA Warp、Neonの3種類の計算バックエンドに対応しています。JAXではCPU・GPU・TPU、Warpでは主に単一GPU、NeonではマルチGPUやマルチ解像度計算を利用できます。同じXLBのLBM計算を、目的や計算環境に応じて異なるバックエンドで実行できます。
マルチ解像度格子：計算領域を一様な格子で計算するだけでなく、必要な領域を細かい格子に分割するGrid Refinementに対応しています。例えば、流体中の物体周辺だけを高解像度にし、それ以外の領域を粗い格子にすることで、解析精度を確保しながら全体の計算量を抑えることができます。
JAXエコシステムとの統合：JAXを利用しているため、Flax、Haiku、Optaxなどの機械学習ライブラリと組み合わせることができます。例えば、XLBで流体シミュレーションを実行し、その結果をニューラルネットワークの学習データとして利用するなど、LBMと機械学習を組み合わせた計算を構築できます。
微分可能LBM：LBMの計算過程を自動微分できるDifferentiable LBMに対応しています。これにより、シミュレーション結果からパラメータに対する勾配を計算し、形状最適化、境界条件の最適化、物性値の推定、機械学習モデルの訓練などに利用できます。物理シミュレーションをAIの学習過程に組み込めることが大きな特徴です。
大規模計算への拡張性：複数のGPUに計算を分散することで、大規模なLBMシミュレーションを実行できます。JAXやNeonの分散計算機能を利用することで、単一GPUでは扱いにくい大規模な計算領域にも対応し、数百GPU・数十億セル規模まで計算を拡張できる設計になっています。
多様なLBMモデルに対応：LBMに必要となるさまざまな境界条件や**衝突モデル（Collision Kernel）**を利用できます。例えば、固体壁の表現に用いるBounce Back系の境界条件などを組み合わせることで、流路、物体周りの流れ、キャビティ流れなど、さまざまな流体問題を構築できます。
使いやすいPythonインターフェース：XLBはPythonを中心に構築されているため、LBMの計算モデルを比較的容易に記述・変更できます。新しい境界条件や衝突モデル、計算ステップなどをPythonで追加できるため、研究者が独自の物理モデルやアルゴリズムを試す研究開発環境として利用しやすい設計になっています。
JAX ArrayとShardmapの活用：JAXの統一されたArray型やshardmapを利用し、NumPyに近い記述方法でLBM計算を構築できます。GPU間のデータ分割や並列処理などの低レベルな処理をユーザーが細かく記述する必要を減らし、物理計算の内容に集中してプログラムできることが特徴です。
幅広い計算環境に対応：同じXLBの計算コードを、マルチコアCPU、単一GPU、マルチGPU、TPUなど、異なる計算環境で実行できます。小規模な問題をCPUで開発・検証し、その後GPUやマルチGPU環境へ移行して大規模計算を行う、といった段階的な利用も可能です。
可視化機能：LBMの計算結果をファイルに保存して後から表示するだけでなく、PhantomGazeなどを利用したGPU上でのIn-situ可視化にも対応しています。大規模計算では大量のデータをCPUやストレージへ転送するコストが問題になりますが、計算中にGPU上で可視化することで、流体の状態を効率的に確認できます。

XLBのバックエンドとしてのJAX

このXLBは、格子ボルツマン法の高速数値計算のために、いくつかのバックエンドが利用可能です。「NVIDIA Warp・Google JAX・Autodesk Neon」が利用可能です。それぞれの概要と、３つの比較をまとめます。

【Gogle JAXバックエンド】

汎用性と機械学習連携に優れたバックエンドです。JAXはGoogleが開発したPython数値計算ライブラリで、NumPyに似た書き方をしながら、GPUやTPUによる高速計算を利用できます。

特徴としては、
CPU、GPU、TPUに対応　　自動微分が可能　　JITコンパイルによる高速化　　マルチGPU・分散計算に対応　　FlaxやOptaxなど機械学習ライブラリと連携可能

特に、流体解析の結果を利用してニューラルネットワークを学習する場合や、逆問題・形状最適化・Physics AIなどでは有力な選択肢です。

【NVIDIA Warpバックエンド】

NVIDIA GPU上で高速にLBM計算を実行することを重視したバックエンドです。NVIDIA Warpは、Python風のコードからGPU並列計算を記述できるフレームワークです。XLBでは、LBMのCollisionやStreamingなどの処理をGPUカーネルとして効率的に実行します。

特徴としては、
NVIDIA GPUに特化　　単一GPUで非常に高い性能　　PythonベースでGPUプログラムを記述可能　　物理シミュレーションとの親和性が高い　　CFD計算を高速に実行可能

特に、「まず高速なLBM流体解析をGPUで実行したい」という目的には、Warpバックエンドが最も扱いやすい選択肢です。XLB公式READMEでもWarpは高性能な単一GPU計算向けとして位置付けられています。

【Autodesk Neonバックエンド】

大規模計算や格子細分化（Grid Refinement）を目的とした新しいバックエンドです。Neonは、GPU上の格子ベース計算を効率的に実行するためのプログラミングモデルです。XLBでは特に、大規模なLBM計算への対応を目的に利用されています。

特徴としては、
マルチGPU計算に対応　　大規模格子に対応　　格子細分化に対応　　単一GPUでもマルチ解像度計算が可能　　Warpを基盤とした高速計算技術を利用

例えば、通常は計算領域全体を同じ格子幅で計算します。Neonでは、必要な場所だけ細かくするような計算が可能です。そのため、物体周辺や渦の発生部分だけ計算精度を高めるといった大規模CFDに適しています。

【３つのバックエンドの比較】

バックエンド	主な目的	GPU	マルチGPU	自動微分	特徴
JAX	AI・Physics ML	○	○	○	汎用性が高い
Warp	高速LBM	NVIDIA	主に単一GPU	用途に注意	高速な物理計算
Neon	大規模・高解像度計算	NVIDIA系	○	—	格子細分化

本検証【01:JAXcpu導入方法と動作確認】では、高速計算ではなく基本的なXLBの動作を、標準的なCPUのみの計算機環境を想定して、行うことを考えます。この時、XLBのCPU動作に関しては、以下の特徴があります。

【XLBのCPU動作の比較】

バックエンド	単体CPU動作	マルチコアCPU	CPU検討環境	主目的
JAX	◎	◎	◎	CPU/GPU汎用
Warp	〇	〇	〇	GPU高速計算
Neon	△	△	×	マルチGPU

以上を踏まえて、XLBをCPUで動作させるには、「Google Jax」をバックエンドとして用いることが、適切な選択になります。なお、JaxはCPUだけでなく、GPU,TPUなどの高速演算デバイスが利用できます。

Google JAXの概要

JAXは、Googleが開発した、高性能な数値計算と機械学習研究向けのPythonオープンソースライブラリです。

JAX公式解説サイト：https://docs.jax.dev/en/latest/
JAX公式GitHubリポジトリ

主な特徴

NumPy互換のAPI: NumPyと非常に似た操作感を持ちながら、CPUやGPU、TPUなどのアクセラレータ上で高速に動作します。
強力な関数型変換: 自動微分（grad）、自動ベクター化（vmap）、JITコンパイル（jit）、並列化（pmap）などの機能を組み合わせて利用できます。
XLAによる最適化: OpenXLAのXLA（Accelerated Linear Algebra）コンパイラを使用し、計算処理を効率化します。

詳細なドキュメントやチュートリアルはJAX公式ページから確認できます。
JAXについて、インストール方法や具体的なコードの書き方（自動微分やJITコンパイルなど）について知りたいですか？

XLBのバックエンド

Autodeskが開発したXLB（Accelerated Lattice Boltzmann Method library）は、JAXの強力な機能（自動微分・JITコンパイル・並列化）を流体力学シミュレーション（LBM）に応用した、Pythonライブラリによるツールです。

物理シミュレーションと機械学習を融合させる「SciML（科学的機械学習）」の文脈で非常に注目されています。JAXをバックエンドにしたXLBが、どのようにその利点を引き出しているかを詳しく解説します。

【XLBにおけるJAXの「3大変換機能」の活用】
XLBは、JAXのコア機能をフルに活用して流体計算の高速化と最適化を実現しています。

JITコンパイル（jax.jit）による超高速化
格子ボルツマン法（LBM）は、空間のグリッド（格子）ごとに衝突や移動の計算を繰り返すため、計算量が非常に多くなります。
XLBは、Pythonで書かれた流体計算の数式を、JAXを介してOpenXLAコンパイラへと引き渡します。コンパイラがカーネルの融合（複数の計算の結合）やメモリレイアウトの最適化を行い、GPU/TPUに最適化されたマシンコードを動的に生成するため、C++で書かれたネイティブコードに匹敵する速度をPythonだけで叩き出せます。
自動微分（jax.grad）による「逆問題・最適化」の容易さ
従来のシミュレータは、結果の出力を算出する（順方向）だけでした。しかし、XLBはJAXを基盤にしているため、シミュレーション全体が「微分可能（Differentiable）」です。
これにより、例えば「車の空気抵抗を最小にする最適な形状」や「特定の流速を作るためのパイプの設計」といった逆問題を、勾配降下法（Gradient Descent）などの機械学習の手法を使ってシミュレータ内で直接最適化（Adjoint-based optimization）できます。
並列化（jax.vmap / shard_map）による大規模分散計算
XLBは、JAXの高度な並列処理アルゴリズムを利用することで、単一のGPUだけでなく、複数のGPUや大規模なクラスター環境へシミュレーションをスケール（大規模化）させることができます。
大規模な格子データ（数億〜数十億セル）を複数のGPUメモリへ自動的に効率よく分割配置（シャード）して計算させることが可能です。

XLB-JAXcpuの動作環境

まず初めに、ここで取り組むXLBの動作環境の概要を示しておきます。

Windows 11
 └─ WSL2 (Ubuntu 24.04LTS: Ubuntu-24.04-XLBjaxCPU)
    └─ Python 仮想環境 (XLBjaxCPU)
       ├─ Google Jax (CPU動作)
       └─ XLB (Lattice Boltzmann Method)

XLB公式GitHubリポジトリの説明を見ると、Installationは、Pythonのpipを用いて簡単にできるようです。指定した環境というより、適切なPythonが動作していれば良いと思いました。むしろ、バックエンドに何を利用するかで、動作環境が選定されます。

ここでは、Windowsで手軽に利用できるUbuntu環境として、「WSL2上のUbuntu24.04LTS」を用いることにして、ここに、XLB+JAX用のPython環境を作ることにします。

バックエンドjaxの条件

JAXをCPU環境のみで利用する場合、特別なハードウェア（GPUやTPU）は不要ですが、JAXが提供する自動微分やXLAによるJITコンパイル（高速化）の恩恵を最大限に受けるためには、いくつかのOSやCPUアーキテクチャ、ソフトウェア面での条件があります。

JAX公式ドキュメントおよび動作要件に基づく計算機環境の条件は以下の通りです。
https://docs.jax.dev/en/latest/installation.html#installation

【対応オペレーティングシステム（OS）とアーキテクチャ】
JAXの動作に必要なコアバイナリ（jaxlib）は、以下の環境向けにビルド（Pre-built Wheels）が公式提供されています。

Linux: x86_64（一般的なIntel/AMD）および aarch64（ARM64）
macOS: aarch64（Apple Silicon: M1/M2/M3/M4などのMac）
インテル製プロセッサ（x86_64）のMac用jaxlibは公式配布が終了しているため、ソースからのビルドが必要です。
Windows: x86_64 （実験的サポート）
Windows環境で安定して動作させたい、あるいは大規模に動かす場合は、Windows Subsystem for Linux (WSL2)上のLinux環境で利用することが推奨されています。

【推奨されるCPUハードウェアスペック】
JAXを快適、あるいは実用的に動かすための一般的なシステム構成の目安です。

基本利用 (プロトタイプ・小規模計算): Intel または AMD の 6コア以上
高度な利用 (大規模シミュレーションなど): 8コア〜12コア以上
メモリ (RAM): 16GB以上（シミュレーションの格子サイズや配列の大きさに応じて、32GB〜64GB以上が望ましいケースもあります）。

【ソフトウェア環境の条件】

Pythonバージョン: JAXがサポートしている現行のPythonバージョン（通常は最新から3つほどのマイナーバージョン。例：Python 3.10〜3.12など）が必要です。
C++ランタイム（Windowsの場合）: Windowsで直接動かす場合、Microsoft Visual Studio 2019 Redistributable（再頒布可能パッケージ）のインストールが必要です。

【CPU環境でJAX（XLBなど）を扱う際の注意点とTips】

インストール方法
CPUのみの環境であれば、以下の非常にシンプルなコマンドでインストールが完了します。

pip install -U jax

マルチコアの擬似シミュレート
JAXはデフォルトでは1つのプロセス（1つのCPUデバイス）として扱われますが、環境変数 XLA_FLAGS を設定することで、1台のCPU内に複数のデバイスが存在するようにシミュレートさせることができます（並列処理 pmap や shard_map のテストに便利です）。

# 8つのマルチコア/デバイスとしてシミュレートする場合
export XLA_FLAGS="--xla_force_host_platform_device_count=8"

データ型（精度）のデフォルト値
JAXは計算の高速化のため、デフォルトで32ビット精度（float32）で計算を行います。流体シミュレーションなどでより精密な64ビット精度（float64）が必要な場合は、コードの冒頭で明示的に設定を有効化する必要があります。

import jax
jax.config.update("jax_enable_x64", True)

XBL+JAXのバージョン対応

今回の「XLB+JAX」の組合わせて、今回想定するバージョンの組み合わせは以下の通りです。

ソフト	推奨バージョン	備考
Python	3.12	今回のvenvに適している
XLB	0.3.2	現在の最新リリース
JAX	0.10～0.11系を優先	XLB 0.3.2のJAXバックエンド用
jaxlib	JAXと互換する同系統	CPU版
NumPy	JAXが要求する版	pipに任せる
backend	JAX	XLB側
device	CPU	`JAX_PLATFORMS=cpu`

上記の、XLB+JAXは最新版なのですが、今回は特別にバージョンを指定せずに、JAXは単純にpipインストールしますが、XLBは例題との対応を考えて、リポジトリの内容からpip導入します。

動作環境Ubuntu-24.04-XLBjaxCPUの準備

ここでは、「WSL2上のUbuntu24.04LTS」を用いることにして、XLB+JAXcpu用のPython環境を作ることにします。初めて取り組む方は、以下の手順でWSL2上にUbuntu24.04を、導入してください。

WSL2上に複数のUbuntu24.04を利用する場合には、基盤となるUbuntuを導入して、これをexportしてから、目的の環境としてimportします。具体的な以下の手順で行います。環境名としては「Ubuntu-24.04-XLBjaxCPU」とします。

PS C:\Users\ryos> wsl --import Ubuntu-24.04-XLBjaxCPU .\WSL2VM\Ubuntu-24.04-XLBjaxCPU .\Ubuntu-24.04-Base.tar
[========================  41.6%                     ]
↓
この操作を正しく終了しました。

PS C:\Users\ryos> wsl -l -v
  NAME                         STATE           VERSION
* Ubuntu-24.04                 Stopped         2
  Ubuntu-24.04-XLBjaxCPU       Stopped         2

この導入で、作成したUbuntu-24.04-XLBjaxCPUを起動したとき、プロンプトが＃でroot環境の時には、先のQiitaに示した手順で、一般ユーザー（ここではryos）にしておきます。

Ubuntu環境を起動して、最新「Ubuntu 24.04.5 LTS」に更新しておきます。

PS C:\Users\ryos> wsl -d Ubuntu-24.04-XLBjaxCPU
ryos@dalab2509dt01:/mnt/c/Users/ryos$ cd
ryos@dalab2509dt01:~$ sudo apt update
[sudo] password for ryos:
：
ryos@dalab2509dt01:~$ sudo apt upgrade
：

ryos@dalab2509dt01:~$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 24.04.5 LTS
Release:        24.04
Codename:       noble

Python環境の確認と準備

Pythonバージョンの条件

Autodesk/XLB および google/jax のGitHubにおける仕様やソースコードから、必要なPythonバージョンを確認します。

【XLB 側の根拠（GitHub ソースコードおよび公式解説より）】
ソースコード（setup.py / メタデータ）の仕様として、
XLBの最新版のパッケージ要件（python_requires）では、>=3.11（Python 3.11以上）が明示的に指定されています。

【JAX 側の根拠（GitHub Discussions / 過去の Issue より）】
JAXの最新仕様: JAXのリポジトリ（setup.py）自体は、現在のバージョンで Python 3.9 以降を広くサポートしています。

以上により、「version 3.11以上」なら良さそうです。

Pythonバージョンの確認とツールの導入

ここで用いるUbunut24.04では、以下の通り「3.12.3」で対応しています。
現在のPythonは、ほぼバージョン３を利用しますが、この状態だとコマンド「python3」と入力します。これを「python」だけでも起動できる設定にしておきます。
さらに、Python3.12で仮想環境を作るツール「python3-venv」と、Python3のインストールツールpip「python3-pip」も導入します。

ryos@dalab2509dt01:~$ python3 --version
Python 3.12.3
ryos@dalab2509dt01:~$ sudo apt install python-is-python3
：
ryos@dalab2509dt01:~$ python --version
Python 3.12.3

ryos@dalab2509dt01:~$ sudo apt install python3.12-venv

ryos@dalab2509dt01:~$ sudo apt install python3-pip
：
ryos@dalab2509dt01:~$ pip --version
pip 24.0 from /usr/lib/python3/dist-packages/pip (python 3.12)

作業ディレクトリとPython仮想環境の準備

準備として、まずJAXcpuを用いたXLBのファイルを保存する作業ディレクトリ「~/XLBjaxcpuWork」を作ります。
次に、Pythonの利用において、目的に応じて仮想環境「XLBjaxCPU」を作って行うことにします。以下のように準備し、sourceで有効にすると、プロンプトの表示が変わります。環境変数VIRTUAL_ENVに名称が設定されています。

ryos@dalab2509dt01:~$ cd
ryos@dalab2509dt01:~$ ls
ryos@dalab2509dt01:~$ mkdir XLBjaxcpuWork
ryos@dalab2509dt01:~$ ls
XLBjaxcpuWork

ryos@dalab2509dt01:~$ python -m venv XLBjaxCPU
ryos@dalab2509dt01:~$ ls
XLBjaxCPU  XLBjaxcpuWork

ryos@dalab2509dt01:~$ source ~/XLBjaxCPU/bin/activate
(XLBjaxCPU) ryos@dalab2509dt01:~$ printenv | grep VIRTUAL
VIRTUAL_ENV=/home/ryos/XLBjaxCPU

さらに導入で用いるpipを更新しておきます

(XLBjaxCPU) ryos@dalab2509dt01:~$ pip install --upgrade pip
：
Successfully installed pip-26.2.1
(XLBjaxCPU) ryos@dalab2509dt01:~$ pip --version
pip 26.2.1 from /home/ryos/XLBjaxCPU/lib/python3.12/site-packages/pip (python 3.12)

Google JAXの導入と確認

導入手順は、以下の公式資料を見ると「pip install -U jax」とありますが、ここでは「CPU」を指定して導入します。
https://github.com/jax-ml/jax#installation

(XLBjaxCPU) ryos@dalab2509dt01:~$ pip install --upgrade "jax[cpu]"
：
Installing collected packages: opt_einsum, numpy, scipy, ml_dtypes, jaxlib, jax
Successfully installed jax-0.11.2 jaxlib-0.11.2 ml_dtypes-0.6.0 numpy-2.5.3 opt_einsum-3.4.0 scipy-1.18.1

JAXの動作をCPUを用いて行うように、以下のように環境変数「export JAX_PLATFORMS=cpu」を「~/XLBjaxCPU/bin/activate」に追加して有効にしておきます。

(XLBjaxCPU) ryos@dalab2509dt01:~$ cd
(XLBjaxCPU) ryos@dalab2509dt01:~$ vi ~/XLBjaxCPU/bin/activate
(XLBjaxCPU) ryos@dalab2509dt01:~$ source ~/XLBjaxCPU/bin/activate
(XLBjaxCPU) ryos@dalab2509dt01:~$ printenv | grep JAX
JAX_PLATFORMS=cpu

以下のPythonスクリプトをプロンプトに入力して、JAXの導入と状態を確認します。

python - <<'PY'
import jax
print("JAX version:", jax.__version__)
print("Devices:", jax.devices())
print("Backend:", jax.default_backend())
PY

動作結果は、以下のようになり、JAXのバージョン0.11.2が導入され、デバイスとしてCPUが設定されています。

JAX version: 0.11.2
Devices: [CpuDevice(id=0)]
Backend: cpu

Autodesk XLBの導入と確認

導入手順は、以下の公式資料では、「CPU-only」の条件で「pip install xlb」で導入しています。
https://github.com/Autodesk/XLB#basic-installation-cpu-only

しかし、この方法では、以下の２つの問題があります。

CPU用として導入したJAXですが、計算には利用しないのに、バックエンドとしてNVIDIA Warpを利用する部分があり、ライブラリを読み込む部分でエラーが生じます
動作検証において、公式githubで公開されている例題を実行しますが、この部分でpip導入したXLBと、例題が想定しているバージョンの不一致でエラーが生じます

そこで、XLBの公式Githubから、以下のリポジトリを入手して、その中にある情報を用いて、XLBを準備します。
https://github.com/Autodesk/XLB.git

まずgitコマンドを導入しておきます。

(XLBjaxCPU) ryos@dalab2509dt01:~$ sudo apt install git

XLBjaxCPUの作業用ディレクトリ「XLBjaxcpuWork」の中に、リポジトリ「XLB-git」を作成するために、以下のようにgitコマンドを実行します。

(XLBjaxCPU) ryos@dalab2509dt01:~$ cd ~/XLBjaxcpuWork/
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ git clone https://github.com/Autodesk/XLB.git XLB-git
Cloning into 'XLB-git'...
：
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ ls
XLB-git
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ cd XLB-git/
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/XLB-git$ ls
AUTHORS       CODE_OF_CONDUCT.md  LICENSE    assets  examples    requirements.txt  setup.py  xlb
CHANGELOG.md  CONTRIBUTING.md     README.md  docs    mkdocs.yml  ruff.toml         tests

次に、ツールや例題のバージョンを固定するために、以下の手順で、Gitでmainのcheckpoint（コミット）を固定しておくことにします。
そのうえで、XLBをリポジトリから導入します。

(XLBjaxCPU) ryos@dalab2312nb01:~$ cd ~/XLBjaxcpuWork/XLB-git/
(XLBjaxCPU) ryos@dalab2312nb01:~/XLBjaxcpuWork/XLB-git$ git checkout main
Already on 'main'
Your branch is up to date with 'origin/main'.
(XLBjaxCPU) ryos@dalab2312nb01:~/XLBjaxcpuWork/XLB-git$ git pull
Already up to date.
(XLBjaxCPU) ryos@dalab2312nb01:~/XLBjaxcpuWork/XLB-git$ pip install -e .
Obtaining file:///home/dalab/XLBjaxcpuWork/XLB-git
：
Successfully built xlb
Installing collected packages: warp-lang, urllib3, typing-extensions, trimesh, six, scooby, ruff, pyparsing, pygments, platformdirs, pillow, packaging, nvtx, mdurl, kiwisolver, idna, fonttools, docstring-parser, cycler, contourpy, charset_normalizer, certifi, attrs, annotated-types, typing-inspection, requests, pyvista-validation, python-utils, python-dateutil, pydantic-core, markdown-it-py, rich, pydantic, pooch, numpy-stl, matplotlib, vtk, rich-rst, cyclopts, pyvista, xlb
Successfully installed annotated-types-0.8.0 attrs-26.1.0 certifi-2026.7.22 charset_normalizer-3.5.1 contourpy-1.4.0 cycler-0.12.1 cyclopts-4.25.3 docstring-parser-0.18.0 fonttools-4.65.0 idna-3.20 kiwisolver-1.5.1 markdown-it-py-4.2.0 matplotlib-3.11.2 mdurl-0.1.2 numpy-stl-4.0.1 nvtx-0.2.16 packaging-26.3 pillow-12.3.0 platformdirs-4.11.11 pooch-1.9.0 pydantic-2.13.5 pydantic-core-2.46.5 pygments-2.21.0 pyparsing-3.3.2 python-dateutil-2.9.0.post0 python-utils-4.0.1 pyvista-0.49.0 pyvista-validation-0.2.2 requests-2.34.2 rich-15.0.0 rich-rst-2.1.0 ruff-0.16.8 scooby-0.12.0 six-1.17.0 trimesh-5.1.0 typing-extensions-4.16.0 typing-inspection-0.4.4 urllib3-2.8.0 vtk-9.7.0 warp-lang-1.17.0 xlb-0.3.2

以上の準備ができたら、JAX+XLBが導入された状態を、改めて以下のスクリプトで確認します。

python - <<'PY'
import jax
import xlb
from xlb.compute_backend import ComputeBackend
print("JAX:", jax.__version__)
print("JAX devices:", jax.devices())
print("JAX backend:", jax.default_backend())
print("XLB:", xlb.__version__)
print("XLB JAX backend:", ComputeBackend.JAX)
PY

動作結果は、以下のように、「JAX：バージョン0.11.2」「XLB：バージョン0.3.2」として、「XLB→JAX→CPU」の関係で処理が行われることになります。

JAX: 0.11.2
JAX devices: [CpuDevice(id=0)]
JAX backend: cpu
XLB: 0.3.2
XLB JAX backend: ComputeBackend.JAX

XLBの公式例題の準備

作成したリポジトリ「XLB-git」の状態で、公式例題の中でJAXを使用しているものを検索します。

(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/XLB-git$ grep -RIl "ComputeBackend.JAX" examples --include="*.py"
examples/cfd/rotating_sphere_3d.py
examples/cfd/differentiable_lbm.py
examples/cfd/lid_driven_cavity_2d_distributed.py
examples/cfd/windtunnel_3d.py
examples/cfd/lid_driven_cavity_2d.py
examples/cfd/flow_past_sphere_3d.py
examples/cfd/turbulent_channel_3d.py
examples/performance/mlups_3d.py
examples/ibm/windtunnel_ibm.py
examples/ibm/airfoil_ibm.py
examples/ibm/wind_turbine_ibm.py
examples/ibm/sphere_ibm.py

この中で、流体解析の例題として標準的でよく使われるもので、２次元矩形空間のキャビティ流れ「lid_driven_cavity_2d.py」を用いて、動作検証を行います。

lid_driven_cavity_2d.pyによるXLBの動作確認

この例題用のディレクトリ「LDC2D」を、XLB作業ディレクトリ「XLBjaxcpuWork」に作り、この中に例題ファイル「lid_driven_cavity_2d.py」をコピーします。

(XLBjaxCPU) ryos@dalab2509dt01:~$ cd ~/XLBjaxcpuWork/
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ mkdir LDC2D
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ ls
LDC2D  XLB-git
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork$ cd LDC2D/
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/LDC2D$ cp ~/XLBjaxcpuWork/XLB-git/examples/cfd/lid_driven_cavity_2d.py .
(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/LDC2D$ ls
lid_driven_cavity_2d.py

以下のように、pythonを用いて、この例題スクリプトを実行します。
なお、この計算環境ではNVIDIA GPUがないので、WARNINGが表示されていますが、問題なくCPUで計算は進みます。

(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/LDC2D$ python lid_driven_cavity_2d.py
No GPU support is available; CPU fallback will be used.
registered bc EquilibriumBC_7914089725062 with id 1
registered bc HalfwayBounceBackBC_7914089725059 with id 2
WARNING! Assuming no-slip condition for BC type = HalfwayBounceBackBC!
Warp CUDA warning: Could not find or load the NVIDIA CUDA driver. GPU execution will not be available.
Warp 1.17.0 initialized:
   CUDA Toolkit 12.9, CUDA driver not available (NVRTC compilation available)
   Devices:
     "cpu"      : "x86_64"
   Kernel cache:
     /home/ryos/.cache/warp/1.17.0
Saved ./lid_driven_cavity_0000000.vtk in 0.029958 seconds.
：
Saved ./lid_driven_cavity_0049999.vtk in 0.015085 seconds.

(XLBjaxCPU) ryos@dalab2509dt01:~/XLBjaxcpuWork/LDC2D$ ls
lid_driven_cavity_0000.png     lid_driven_cavity_0025000.vtk  lid_driven_cavity_1000.png   lid_driven_cavity_32000.png
：
lid_driven_cavity_0024000.vtk  lid_driven_cavity_0049999.vtk  lid_driven_cavity_31000.png

実行すると、可視化用vtkファイルと結果画像pngファイルが、1000ステップごとに、49999ステップまでの５０段階で保存されます。ファイルのタイムスタンプを見ると解析に１１分間必要でした。
なお、例題を実行したPCのCPUは、コマンドlscpuで確認すると、「AMD Ryzen 7 5825U：８コア」で、Windowsで確かめると2.00GHz動作です。

ここで、JAXがマルチコアCPUを利用する状態を調べると、「JAXのCPUバックエンドが複数スレッドを利用し、OSを介してCPUの物理コア上でLBM計算を並列実行する」とのことです。
特別に指定しない場合には、単純に言うと物理コア全てを用いて、マルチコアでの並列処理を行っています。上記の場合には８コアと言えます。

lid_driven_cavity_2d.pyの結果確認

Windowsのファイルブラウザの左下に、Linuxの項目があり、ここにファイルシステム「Ubuntu-24.04-XLBjaxCPU」があります。このリンクをデスクトップに作っておきます。
例題ディレクトリ「\home\ryos\XLBjaxcpuWork\LDC2D」を開きます。最後のステップの画像ファイル「lid_driven_cavity_49999.png」は、以下のように解析最後の速度分布を表しています。

可視化用ファイルvtkを確認するには、可視化ツールParaViewを利用します。以下のリンクからWindows用のParaViewを、ファイルから導入して利用します。
https://www.paraview.org/download/

WindowsのParaViewを起動して、例題ディレクトリを開き、可視化用vtkファイル群を選択して、最後の状態の速度の大きさを可視化すると、以下のような状態です。

現在のまとめ

実は、XLBの公式ウエブの解説通りpipで単純にインストールすると、例題ファイルの実行では、色々と工夫が必要みたいで動作しません。

この主な原因は、XLB・JAX・NVIDIA Warpのバージョンと、各例題が想定するAPIの世代が一致していないことです。このように、例題ごとに利用するXLBの機能が異なるため、同じ環境でも異なる場所で互換性エラーが発生します。
したがって、XLB・JAX・Warpの対応するバージョンを一つの組合せとして固定するため、githubからリポジトリを入手して、checkpointで状態を固定して、XLBをリポジトリから導入し、その例題を用いると、上記のように、問題なく動作します。

それでは、ここで確認した例題「lid_driven_cavity_2d.py」について、以下では詳しく調べてみます。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up