Metadata-Version: 2.4
Name: unit-jit
Version: 0.1.0
Summary: JIT unit-stripping decorator: write Pint-annotated code, run at float speed
Author: Matthias Függer
License: Apache-2.0
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: libcst>=1.4
Requires-Dist: pint>=0.24
Provides-Extra: dev
Requires-Dist: numpy>=1.26; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Description-Content-Type: text/markdown

# unit-jit

JIT unit-stripping decorator for [Pint](https://pint.readthedocs.io)-annotated Python. Write clean, unit-safe code; pay no Pint overhead in hot loops.

```python
from unit_jit import unit_jit, ureg

@unit_jit
def simulate(n: int) -> np.ndarray:
    mrna = 10.0 * ureg.mol / ureg.L
    dt   =  0.1 * ureg.s
    delta = 0.01 / ureg.s
    out = np.empty(n)
    for i in range(n):
        mrna = mrna - delta * mrna * dt
        out[i] = mrna.to_base_units().magnitude
    return out
```

First call runs the original Pint function (warm-up). All subsequent calls run a rewritten, pure-float version — **~10× faster** for tight loops.

## How it works

1. **Module-level compilation** — on first call, all `@unit_jit` functions in the same module are rewritten together: `.magnitude`, `.to_base_units()`, and `cast("Quantity", x)` are stripped from the source.
2. **Eager snapshot** — Quantity attributes on objects (e.g. `self.params.alpha`) are pre-converted to SI floats once at boundary entry. Attribute access inside the loop is a plain dict lookup.
3. **Fast zone** — a thread-local flag marks the outermost `@unit_jit` frame. Inner `@unit_jit` calls skip boundary conversion entirely.
4. **Return wrapping** — the SI unit of the return value is inferred from the first call and used to wrap subsequent results back into `Quantity`.
5. **Dimension guard** — argument dimensions are cached from the first call; any later call with a different dimension raises `TypeError` immediately.

The right entry point is the **outermost function that owns the hot loop** — not the leaf functions it calls.

## Installation

```bash
# uv (recommended)
uv add unit-jit

# pip
pip install unit-jit
```

From source:

```bash
git clone ...
cd unit-jit

# uv
uv sync --extra dev

# pip
pip install -e ".[dev]"
```

## Usage

### Simple function

```python
from unit_jit import unit_jit, ureg

@unit_jit
def velocity(d, t):
    return d / t

velocity(10 * ureg.m, 2 * ureg.s)   # warm-up (runs Pint)
velocity(10 * ureg.m, 2 * ureg.s)   # fast (pure float internally)
velocity(10 * ureg.cm, 2 * ureg.s)  # fine — same dimension, different unit
velocity(10 * ureg.m, 2 * ureg.m)   # TypeError — wrong dimension for arg 1
```

### Class with Quantity attributes

```python
from dataclasses import dataclass
from unit_jit import unit_jit, ureg

@dataclass
class Params:
    alpha: Quantity   # [mol/L/s]
    delta: Quantity   # [1/s]

class Model:
    def __init__(self, params):
        self.params = params

    @unit_jit
    def rate(self, mrna):
        return self.params.alpha - self.params.delta * mrna

    @unit_jit                          # ← entry point: owns the hot loop
    def simulate(self, n):
        mrna = self.params.alpha / self.params.delta
        out = np.empty(n)
        for i in range(n):
            mrna = mrna + self.rate(mrna) * (0.1 * ureg.s)   # rate() in fast zone
            out[i] = mrna.to_base_units().magnitude
        return out
```

`self.params.alpha` and all other Quantity attributes are converted to SI floats once when `simulate` is first called fast; `self.rate()` is called from inside the fast zone, so it skips boundary conversion entirely.

## Running tests

```bash
pytest
```

## License

Apache-2.0
