acropole baseline

.github/workflows/docs.yml

@@ -1,28 +1,20 @@
-name: build docs
+name: docs
 
 on:
   workflow_dispatch:
+  pull_request:
   push:
-    tags:
-      - "*"
 
 permissions:
   contents: read
-  pages: write      # for deploying to GitHub Pages
-  id-token: write   # for authenticating deployment
 
 jobs:
-  deploy:
+  build:
+    name: Build Docs
     runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
     steps:
       - uses: actions/checkout@v6
 
-      - name: Configure Pages
-        uses: actions/configure-pages@v6
-
       - name: Install uv
         uses: astral-sh/setup-uv@v8.2.0
         with:
@@ -33,17 +25,31 @@ jobs:
         with:
           python-version: "3.14"
 
-      - name: Install the project
-        run: uv sync --all-extras --all-groups
+      - name: Install docs dependencies
+        run: uv sync --no-default-groups --group docs --group typing
 
       - name: Build with Zensical
-        run: uv run zensical build --clean  # no cache for zensical v0.0.x
+        run: uv run --no-default-groups --group docs --group typing zensical build --clean
 
-      - name: Upload artifact
+      - name: Upload Pages artifact
         uses: actions/upload-pages-artifact@v5
         with:
           path: site
 
+  deploy:
+    name: Deploy GitHub Pages
+    needs: build
+    if: >-
+      github.event_name == 'workflow_dispatch' ||
+      startsWith(github.ref, 'refs/tags/')
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write      # for deploying to GitHub Pages
+      id-token: write   # for authenticating deployment
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v5
+        uses: actions/deploy-pages@v5

.github/workflows/tests.yml

@@ -30,7 +30,7 @@ jobs:
           python-version: "3.10"
 
       - name: install the project
-        run: uv sync --all-extras --all-groups
+        run: uv sync --all-groups --extra all --extra torch_cpu
 
       - name: style checking
         run: |
@@ -67,7 +67,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: install the project
-        run: uv sync --all-extras --all-groups
+        run: uv sync --all-groups --extra all --extra torch_cpu
 
       - name: run non-network tests
         run: uv run pytest -q -m "not network"
@@ -97,7 +97,7 @@ jobs:
         run: sudo apt-get update && sudo apt-get install -y zstd
 
       - name: install the project
-        run: uv sync --all-extras --all-groups
+        run: uv sync --all-groups --extra all --extra torch_cpu
 
       - name: run network tests
         env:

README.md

@@ -25,7 +25,7 @@ Depending on your use case, you can pick the optional dependencies you need:
 - `polars`: support for [polars](https://github.com/pola-rs/polars) and its Array API shim (postprocessing third party data)
 - `httpx`: support for [httpx](https://github.com/encode/httpx) (downloading data from external sources)
 - `xarray`: support for [xarray](https://github.com/pydata/xarray) (ARCO-ERA5 weather grids, working with NetCDF)
-- `jax`: support for [JAX](https://github.com/jax-ml/jax) (automatic differentiation support)
+- `jax`: support for [JAX](https://github.com/jax-ml/jax) (automatic differentiation/GPU acceleration support)
 - `matplotlib`: plotting
 - `platformdirs`: reading/writing cache/config files
 - `cli`: command line scripts
@@ -51,13 +51,16 @@ uv run aerocore --help
 git clone https://github.com/abc8747/aerocore --depth=1
 cd aerocore
 uv venv
-uv sync --all-extras --all-groups
+# standard
+uv sync --all-groups --extra all
+# if you want to benchmark with GPU support 
+uv sync --all-groups --all-extras --no-extra 'torch-cpu'
 ```
 
 To run scripts:
 
 ```sh
-uv run examples/autodiff.py
+uv run examples/capabilities_jax.py
 ```
 
 For documentation:

docs/acropole.md

@@ -0,0 +1,105 @@
+# Fuel Estimation (Acropole)
+
+Acropole is a neural network that estimates the fuel flow given an aircraft state vector and its parameters (see [Jarry et al. (2024)](https://doi.org/10.13140/RG.2.2.23229.27360)).
+
+Unlike the [upstream repository](https://github.com/DGAC/Acropole) which uses ONNX, this implementation is intentionally backend agnostic. You can use pure numpy, or optionally use JAX/Torch for GPU acceleration.
+
+## Download assets
+
+The model weights and aircraft database are distributed separately (not bundled in the library) because it is licensed under the GNU Affero General Public License v3.0. Download them manually:
+
+```sh
+uv add "aerocore[cli,httpx]"
+# by default, this is downloaded to the cache directory (xdg cache home on Linux/Mac)
+uv run aerocore data-acropole-sync
+```
+
+## Numpy Backend
+
+To run with the default numpy backend:
+
+```python
+--8<-- "examples/acropole_numpy.py:input0"
+```
+
+```text
+--8<-- "examples/acropole_numpy.py:output0"
+```
+
+!!! note "Inputs"
+
+    Trajectory arrays must have identical shapes and store time along the final axis. Aerocore does not broadcast, resample, smooth or interpolate trajectory inputs.
+
+    The paper reports using one-second QAR samples and says vertical speed and true airspeed were smoothed with a Savitzky-Golay filter, but does not specify a filter window. The upstream implementation later recommended approximately four-second sampling. Callers are responsible for any resampling, interpolation, and smoothing.
+
+Under the hood, it duplicates the parameters in the given [`AcropoleAircraft`][aerocore.acropole.AcropoleAircraft] across the batch, resulting in an array of shape `(12, T)`. But in cases where you want heterogenous aircraft types, you should construct the input arrays yourself and use lower level functions:
+
+```python
+--8<-- "examples/acropole_mixed_aircraft.py:input0"
+```
+
+```text
+--8<-- "examples/acropole_mixed_aircraft.py:output0"
+```
+
+## Other numerical backends
+
+Here is a benchmark of Acropole. On my machine `jax.jit` achieves 2-4x speedup compared to `onnxruntime-gpu` depending on the batch size. Note that input arrays are assumed to be on the device the model is located, so GPU benchmarks do not include the CPU-GPU synchronisation latency.
+
+![Acropole backend benchmark](assets/img/acropole-benchmark.png)
+
+To use JAX JIT on GPU, use [`tree_map()`][aerocore.utils.tree_map] to move the model to another device:
+
+```sh
+uv add "aerocore[jax_gpu,matplotlib]"
+```
+
+```python
+--8<-- "examples/acropole_jax_gpu.py:input0"
+```
+
+```text
+--8<-- "examples/acropole_jax_gpu.py:output0"
+```
+
+A toy example of using `jax.value_and_grad` to show the direction of the steepest increase in fuel flow over a grid of altitude and groundspeed, for a steady, level, unaccelerated A320 flight at 65 tonnes:
+
+<details>
+    <summary>code</summary>
+
+    ```python
+    --8<-- "examples/acropole_jax_gradient.py:input0"
+    ```
+</details>
+
+![Acropole fuel-flow gradient field](assets/img/acropole-fuel-flow-gradient.png)
+
+It supports PyTorch CUDA too:
+
+```sh
+uv add "aerocore[torch_gpu]"
+```
+
+```python
+--8<-- "examples/acropole_torch_gpu.py:input0"
+```
+
+```text
+--8<-- "examples/acropole_torch_gpu.py:output0"
+```
+
+Any Array API compatible interface is also supported.
+
+## Development
+
+```sh
+./scripts/acropole_prepare_model.py build \
+  --output-dir dist/acropole-v0.1.0
+
+# test the published assets
+uv run pytest -q tests/test_acropole.py
+
+./scripts/acropole_benchmark.py \
+  --output /tmp/acropole-benchmark.jsonl \
+  --plot docs/assets/img/acropole-benchmark.png
+```

docs/api/acropole.md

@@ -0,0 +1 @@
+::: aerocore.acropole

docs/capabilities.md

@@ -0,0 +1,31 @@
+# Capabilities
+
+Most aerocore functions can be used with multiple computational backends by passing an Array API-compatible namespace to the `xp` keyword argument.
+
+This allows us to do some cool tricks like automatically differentiating and vmapping a function via JAX:
+
+```python
+--8<-- "examples/capabilities_jax.py:input0"
+```
+
+```text
+--8<-- "examples/capabilities_jax.py:output0"
+```
+
+And even pipe an entire dataframe column through an aerocore function with Polars Lazy evaluation:
+
+```python
+--8<-- "examples/capabilities_polars.py:input0"
+```
+
+```text
+--8<-- "examples/capabilities_polars.py:output0"
+```
+
+```python
+--8<-- "examples/capabilities_polars.py:input1"
+```
+
+```text
+--8<-- "examples/capabilities_polars.py:output1"
+```

examples/acropole_jax_gpu.py

@@ -0,0 +1,43 @@
+# --8<-- [start:input0]
+import jax
+import jax.numpy as jnp
+
+import numpy as np
+from aerocore.acropole import fuel_flow, load_aircraft_database, load_model
+from aerocore.utils import tree_map
+
+device = jax.devices("gpu")[0]
+model = tree_map(lambda array: jnp.asarray(array, device=device), load_model())
+aircraft = load_aircraft_database()["A320"]
+
+
+def predict(
+    groundspeed: jax.Array, altitude: jax.Array, vertical_rate: jax.Array
+) -> jax.Array:
+    return fuel_flow(  # type: ignore[no-any-return]
+        model=model,
+        aircraft=aircraft,
+        groundspeed=groundspeed,
+        altitude=altitude,
+        vertical_rate=vertical_rate,
+        xp=jnp,
+    )
+
+
+predict = jax.jit(predict, device=device)
+print(
+    np.asarray(
+        predict(
+            groundspeed=jnp.asarray([180.0, 450.0, 250.0], device=device),
+            altitude=jnp.asarray([0.0, 30_000.0, 40_000.0], device=device),
+            vertical_rate=jnp.asarray([3_000.0, 0.0, -2_000.0], device=device),
+        )
+    )
+)
+# --8<-- [end:input0]
+
+"""
+--8<-- [start:output0]
+[1.9125862  0.75758696 0.09448687]
+--8<-- [end:output0]
+"""