Anri: GPU-accelerated Diffraction Microstructure Imaging analysis#

repo docs GitHub License tests codecov style platform jax

What is Anri?#

Anri is a Python package for the GPU-accelerated analysis of Diffraction Microstructure Imaging data, such as (Scanning) Three-Dimensional X-Ray Diffraction [(s)3DXRD].

How does it work?#

The core concept of Anri is as follows:

  • JAX-native code

(almost) all of Anri is implemented in JAX, a Python library for high-performance (e.g. GPU-accelerated) numerical computing. The benefit is that Anri will work on any major x86 (and some ARM) CPUs and any recent CUDA-compatible GPU.

  • Intensity-aware forward model

At the core of Anri is a high-performance forward model that goes from sample space (e.g. a grid of UBI matrices) to detector space (e.g. [slow, fast] coordinates). This can be used to investigate the performance of existing (s)3DXRD analysis packages such as ImageD11 by comparing forward-projected data to the raw data that you measured. Intensities are computed using the structure factors thanks to Dan’s Diffraction and accumulate in detector pixels.

  • Differentiability

Great effort has been undertaken to ensure that JAX-native Anri functions are differentiable, using the powerful auto-diff capabilities of JAX. This has two obvious use-cases:

  • Peak shapes

By expressing instrumental parameters such as incident beam divergence and energy spread as Gaussian distributions, Anri can use the per-peak Jacobians produced by JAX to propagate these parameters into detector space as a covariance matrix in output space, thereby rendering fairly realistic peak shapes that are not just simple detector point spread functions. Therefore, a spread in beam energy (for example) manifests as a radial distribution on the detector.

  • Gradient-aware optimisation (in progress)

Anri will take advantage of the differentiable, intensity-aware forward model to perform iterative refinement of grain maps produced by ImageD11 (and perhaps other programs in the future) to yield refined maps of orientation gradients and strains.

AI usage disclaimer#

Some parts of the code are developed with the assistance of LLMs, such as Claude and Gemini.

Installation#

From Conda#

Coming soon!

From source (for developers)#

Anri may (eventually) rely on packages from both conda and pip. For ease of installation, it is recommended to use unidep which can install packages from both sources.

Clone the repository#

git clone git@github.com:jadball/anri.git anri
cd anri

Set up a Conda environment#

conda create -n <env-name>
conda activate <env-name>

Ensure pip is running from the Conda environment#

which pip  # should yield something inside the environment <env-name>

Install build dependencies#

pip install --upgrade pip unidep

Install conda, then pip deps, then the package itself (with dev optional deps) as editable:#

unidep install .[dev] -e