Metadata-Version: 2.3
Name: dinox
Version: 0.5
Summary: Derivative Informed Neural Operators in JAX and Equinox
Author: Joshua Chen, Michael Brennan, Lianghao Cao, Thomas O'Leary-Roseberry
Author-email: Joshua Chen <joshuawchen@icloud.com>, Michael Brennan <mcbrenn@mit.edu>, Lianghao Cao <lianghao@caltech.edu>, Thomas O'Leary-Roseberry <tom.olearyroseberry@utexas.edu>
License: MIT License
         
         Copyright (c) 2025 Joshua Chen
         
         Permission is hereby granted, free of charge, to any person obtaining a copy
         of this software and associated documentation files (the "Software"), to deal
         in the Software without restriction, including without limitation the rights
         to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
         copies of the Software, and to permit persons to whom the Software is
         furnished to do so, subject to the following conditions:
         
         The above copyright notice and this permission notice shall be included in all
         copies or substantial portions of the Software.
         
         THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
         IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
         FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
         AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
         LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
Requires-Dist: jax>=0.4.30
Requires-Dist: jaxlib>=0.4.30
Requires-Dist: jaxtyping
Requires-Dist: optax>=0.2.3
Requires-Dist: bayesflux[hippylib]>=0.8
Requires-Dist: equinox
Requires-Dist: hickle>=5.0.3
Requires-Dist: pydantic
Requires-Dist: cupy-cuda12x ; extra == 'cupy'
Requires-Dist: pytest ; extra == 'dev'
Requires-Dist: black ; extra == 'dev'
Requires-Dist: isort ; extra == 'dev'
Requires-Dist: flake8 ; extra == 'dev'
Requires-Dist: flake8-pyproject ; extra == 'dev'
Requires-Python: >=3.9
Project-URL: Homepage, https://github.com/dinoSciML/dinox
Project-URL: Repository, https://github.com/dinoSciML/dinox
Provides-Extra: cupy
Provides-Extra: dev
Description-Content-Type: text/markdown

# dinox

`dinox` is a `JAX` implementation of Reduced Basis Derivative Informed Neural Operators, built for high performance in single-GPU environments where all training data fits in GPU memory.

The library is designed primarily for PDE learning workflows based on:

- `FEniCS 2019.1`
- Jacobians computed via `hippylib`
- Subspace methods provided by `bayesflux`

---

## Overview

`dinox` provides:

- Reduced basis neural operator architectures
- Derivative-informed training using PDE Jacobians
- GPU-accelerated implementations in `JAX` and `Equinox`
- Integration with FEniCS-discretized PDEs

---

## Important: FEniCS & Hippylib Environment Required

`dinox` depends on `bayesflux[hippylib]`, which requires:

- hippylib
- FEniCS 2019.1

`FEniCS` has system-level dependencies and cannot be installed via pip alone.

You must first create a conda environment with FEniCS 2019.1 before installing dinox.

---

## Installation

### Prerequisites

- NVIDIA driver >= 525 (check with `nvidia-smi`)
- conda or mamba

> **Note on CUDA libraries:** You do **not** need to install CUDA Toolkit, cuDNN, or cuSPARSE via conda or your system package manager. The pip wheels for JAX and CuPy bundle their own CUDA 12 runtime libraries. Installing system CUDA alongside pip-bundled CUDA is the most common source of GPU detection failures.

### Step 1 — Create a FEniCS 2019.1 environment

```bash
conda create -n fenics-2019.1_env -c conda-forge fenics==2019.1.0 python=3.11
conda activate fenics-2019.1_env
```

### Step 2 — Fix `LD_LIBRARY_PATH` (critical for GPU)

A system-level or conda-set `LD_LIBRARY_PATH` pointing to a CUDA installation will conflict with the CUDA libraries bundled in the JAX and CuPy pip wheels, causing errors like `Unable to load cuSPARSE`. This fixes this issue

```bash 
unset LD_LIBRARY_PATH
```

### Step 3 — Install GPU-enabled JAX

```bash
pip install "jax[cuda12]" cupy-cuda12x
```

> JAX versions 0.6.1 and 0.6.2 have a [known cuSPARSE loading bug](https://github.com/jax-ml/jax/issues/30050). Version 0.6.0 is the recommended stable release for CUDA 12.

### Step 4 — Install dinox

```bash
# With CuPy GPU support (recommended)
pip install dinox[cupy]

# Without CuPy
pip install dinox
```

### Step 5 — Verify GPU

```bash
python -c "import jax; print('JAX devices:', jax.devices())"
python -c "import cupy; print('CuPy GPU count:', cupy.cuda.runtime.getDeviceCount())"
```

You should see your NVIDIA GPU listed. If JAX shows only `CpuDevice`, check that `LD_LIBRARY_PATH` is unset (see Step 2).

---

## GPU Support

- Designed for single-GPU workflows where all data fits in GPU memory
- Requires CUDA 12-enabled JAX (`pip install "jax[cuda12]"`) — the pip wheel bundles its own CUDA runtime
- Optional CuPy arrays for GPU operations via `dinox[cupy]`
- Without GPU, CPU fallback is automatic

---

## Development

```bash
conda create -n fenics-2019.1_env -c conda-forge fenics==2019.1.0 python=3.11
conda activate fenics-2019.1_env
unset LD_LIBRARY_PATH  # or use the permanent conda hook above

pip install "jax[cuda12]==0.6.0"
pip install -e ".[dev]"
```

This installs:

- dinox (editable)
- development tools (pytest, black, flake8, isort)
- bayesflux[hippylib]
- hippylib
- all required JAX dependencies

---

## Requirements

- Python >= 3.9
- FEniCS 2019.1 (via conda)
- JAX >= 0.4.30 (for GPU: `pip install "jax[cuda12]==0.6.0"`)
- NVIDIA driver >= 525 (for GPU)
- Optional: CuPy for GPU array operations (`pip install dinox[cupy]`)

---

## Troubleshooting

| Problem | Solution |
|---|---|
| `Unable to load cuSPARSE` | `unset LD_LIBRARY_PATH` before running Python |
| JAX shows only `CpuDevice` | Ensure `jax[cuda12]` was installed (not just `jax`) and `LD_LIBRARY_PATH` is unset |
| `nvidia-smi` not found | Install or update NVIDIA driver (>= 525) |
| JAX/CuPy CUDA version conflict | Do **not** `conda install cudatoolkit` — let pip wheels provide CUDA |

---

## Repository

- Homepage: https://github.com/dinoSciML/dinox
- Repository: https://github.com/dinoSciML/dinox