Metadata-Version: 2.4
Name: uncpy
Version: 0.2.0
Summary: Linear uncertainty propagation powered by JAX autodiff: real/complex values, N-D arrays, linear algebra, FFT.
Author-email: Ziad Hatab <zi.hatab@gmail.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/ZiadHatab/uncpy
Project-URL: Repository, https://github.com/ZiadHatab/uncpy
Project-URL: Issues, https://github.com/ZiadHatab/uncpy/issues
Keywords: uncertainty,metrology,error propagation,covariance,jax,measurement
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.11
Requires-Dist: numpy>=1.26
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Dynamic: license-file

<p align="center">
  <img src="https://raw.githubusercontent.com/ZiadHatab/uncpy/main/assets/logo-errorbar.png" alt="UncPy" width="440">
</p>

Linear uncertainty propagation library powered by [JAX](https://github.com/jax-ml/jax)
autodiff. Given uncertain input parameters, uncpy computes the uncertainty
of any function of them. JAX computes the Jacobian of the function through
automatic differentiation, and uncpy propagates the input covariance to
the output covariance:

$$
\Sigma_{\mathrm{out}} = J \Sigma_{\mathrm{in}} J^{\top},
\qquad J = \frac{\partial f}{\partial x}
$$

`uncpy` reads like NumPy: one array type and one set of familiar function
names. There are no hand-written propagation rules; anything JAX can
differentiate propagates uncertainty correctly.

- real and complex values, scalars and N-D arrays
- full correlation tracking through arbitrary computations
- numpy-style math (`up.sin`, `up.matmul`, `up.einsum`, ...), linear algebra and FFT
- native NumPy interoperability: `np.sin(x)`, `np.linalg.solve(A, b)` work on uncertain arrays
- lift *any* JAX-traceable function with `up.wrap`
- uncertainty budgets, per-input sensitivities, effective degrees of freedom

## Installation

From PyPI:

```bash
pip install uncpy
```

From this repository (latest development version):

```bash
pip install git+https://github.com/ZiadHatab/uncpy.git
```

Or, for a local editable clone:

```bash
git clone https://github.com/ZiadHatab/uncpy.git
cd uncpy
pip install -e .
```

Requires Python ≥ 3.11, `jax ≥ 0.11`, `numpy`.

## Quick example

```python
import uncpy as up

x = up.ufloat(1.0, 0.1, name="length")
y = up.ufloat(2.0, 0.2, name="width")
z = x * up.sin(y)

print(z)               # 0.909297 ± 0.123
z.value                # 0.9092974268256817; plain floats and numpy arrays,
z.stdunc               # 0.1232706...        never wrapped Array types

print(up.budget(z))    # input   stdunc contribution
                       # ---------------------------
                       # length  0.09093
                       # width   0.08323
```

## Uncertain inputs

One universal constructor, `uarray`: shape and dtype follow the nominal
value, and the uncertainty (elementwise `stdunc` or a full `covariance`)
is validated against them.

```python
import numpy as np
import uncpy as up

a = up.uarray([1.0, 2.0, 3.0], stdunc=0.1)                            # independent
c = up.uarray([1.0, 2.0], covariance=[[0.01, 0.002], [0.002, 0.04]])  # correlated
S = up.uarray(np.array([0.1 + 0.2j, 0.9 - 0.1j]), stdunc=0.01)        # complex

# independent elements, each with its own covariance: pass a stack.
# Here: B independent complex S21 measurements, one 2x2 (Re, Im)
# covariance per measurement point
S21 = up.uarray(s21_values,                # shape (B,), complex
                covariance=s21_covs)       # shape (B, 2, 2), one per point

x = up.ufloat(1.0, 0.1, name="x", dof=9)   # scalar convenience
z = up.ucomplex(50.0, 0.1)                 # force complex: u(Re) = u(Im) = 0.1
w = up.ucomplex(1 + 2j, 0.1 + 0.3j)        # u(Re) = 0.1, u(Im) = 0.3

# like np.array, uarray also assembles existing uncertain values into an
# array; their uncertainty and correlations are preserved
m = up.uarray([x, 2.0 * x, x + 1.0])
```

Complex quantities are backed by one (Re, Im) pair of real random
variables per element: complex `stdunc` entries give each part its own
uncertainty, and `covariance` is `(2m, 2m)` over the flattened elements
with interleaved `(Re v0, Im v0, Re v1, ...)` ordering, the layout
`up.covariance` returns. Quantities created in *separate* calls are
independent; components created *together* may be correlated.

You do not have to pass the stack form to get its efficiency: a joint
`covariance` that is block-diagonal over the leading axis is recognized
as one, and a rank-deficient covariance (a shared systematic, or a
sample covariance from fewer repeats than quantities) is stored with one
base variable per unit of rank rather than per element.

## Correlation is tracked automatically

```python
x = up.ufloat(3.0, 0.1)
up.stdunc(x * x)       # 0.6; knows x*x is x**2, not a product of two inputs
up.stdunc(x - x)       # 0.0; exact cancellation
```

## Arrays, linear algebra, FFT

```python
A = up.uarray(np.eye(3) + 0.1, stdunc=0.01)
b = up.uarray([1.0, 2.0, 3.0], stdunc=0.1)

up.linalg.solve(A, b)      # solve, inv, pinv, det, cholesky, eig, eigh, svd, qr, ...
up.fft.fft(b)              # linear transform -> propagation is exact

np.mean(b)                 # numpy functions dispatch to uncpy natively
np.linalg.solve(A, b)
```

An `UncArray` never turns into plain numbers by accident: `float(x)`
raises, and `np.array` applied to uncertain values keeps the objects
(an object-dtype array) rather than stripping them.
When you want just the nominal value (for plotting, saving, printing),
ask for it explicitly:

```python
float(b[0])              # TypeError: this would discard the uncertainty
b[0].value               # 1.0; the explicit way to leave uncertainty behind
plt.plot(b.value)        # deliberate, and visible in the code
```

## Eigendecomposition

Non-symmetric `eig` propagates uncertainty of eigenvalues *and*
eigenvectors (JAX ≥ 0.11; assumes distinct eigenvalues, LAPACK
normalization gauge):

```python
A = up.uarray([[2.0, 1.0], [0.5, 3.0]], stdunc=0.01)
w, v = up.linalg.eig(A)

up.stdunc(w)               # eigenvalue uncertainty, u(Re) + 1j*u(Im)
up.jacobian(w[0].real, A)  # d(lambda)/dA; matches the analytic
                           # (w^T dA v) / (w^T v) left/right-eigenvector formula
```

See [Jacobian of eigendecomposition](https://ziadhatab.github.io/posts/jacobian-of-eigendecomposition/)
for the underlying math; the test suite verifies uncpy against it.

## Any JAX function

```python
import jax
import jax.numpy as jnp

def s2t(S):             # scattering -> transfer parameters (2x2 complex)
    return jnp.array([[-jnp.linalg.det(S), S[0, 0]],
                      [-S[1, 1], 1.0]]) / S[1, 0]

S = up.uarray(np.array([[0.01 + 0.02j, 0.9 - 0.3j],
                        [0.9 - 0.3j, 0.05 + 0.01j]]),
              stdunc=0.005)               # uncertain 2x2 matrix

s2t_unc = up.wrap(s2t)
T = s2t_unc(S)     # uncertainty propagates through det, indexing, ...
```

## Performance

Each uncpy operation asks JAX to linearize a function, so the cost is per
operation, not per arithmetic flop. Two habits follow from that.

**Lift whole models, not single operations.** Write the model in plain
`jax.numpy` and lift it once:

```python
def model(a, b, c):                     # plain JAX, no uncpy inside
    z = a
    for _ in range(40):
        z = jnp.sin(z * b) + c / (1.0 + z * z)
    return z

fast = up.wrap(jax.jit(model))          # lift once, compile once
```

| how the model is written | time |
| --- | --- |
| the same model with `up.*` ops in a Python loop | 333 ms |
| `up.wrap(model)` | 87 ms |
| `up.wrap(jax.jit(model))` | **14 ms** |

> Do not put `jax.jit` *over* uncertain values. It works, but each input's
> identity is part of the pytree structure, so every fresh set of inputs
> recompiles. Jit the plain-JAX function and lift the result, as above.

**Work on arrays, not elements.** An array created with `stdunc=` has
independent elements, so uncpy stores O(m) sensitivities instead of a
dense m×m Jacobian:

```python
M = up.uarray(np.tile(np.eye(2) * 2.0, (50_000, 1, 1)), stdunc=0.01)

up.stdunc(up.linalg.inv(M))   # 50k independent 2x2 inversions: 40 ms, 6 MB
```

The same 50 000 elements in a Python loop take about 8 minutes.

This compact storage survives elementwise math, and any function that
numpy batches over leading axes: `inv`, `solve`, `det`, `eig`, `eigh`,
`svd`, `qr`, `cholesky`, `@`, and the 1-D FFTs. You get it by writing
normal numpy; there is nothing to switch on.

When an operation really does mix elements, uncpy switches to the dense
m×m Jacobian. That happens for `up.sum(x)` and other reductions, for an
FFT of a single signal, and for inverting a single matrix. The result is
exact at any size, but the memory grows quadratically.

A function you build from `up.*` operations keeps the compact storage on
its own, because uncpy sees each step. `up.wrap` is different: it hands
uncpy one opaque function, so the dense path is the only safe choice.
When that function handles a single element, `up.vmap` applies it across
the array and keeps the O(m) storage:

```python
s2t_unc = up.wrap(s2t)                  # s2t from above: one 2x2 matrix
Sweep = up.uarray(meas, stdunc=0.005)   # a stack of them, (50_000, 2, 2)

T = up.vmap(s2t_unc)(Sweep)             # 1.3 s, 26 MB
```

This is the same reason `jax.vmap` exists: `s2t` is written for one
matrix, so it has to be mapped to run over many. See the `up.vmap`
docstring for mapping several arguments at once.

## Reading out results

```python
up.value(z)             # nominal value                  (also z.value)
up.stdunc(z)            # standard uncertainty           (also z.stdunc)
up.stdunc(z, wrt=x)     # uncertainty w.r.t. input x     (or wrt="name")
up.covariance([z, w])   # joint covariance (off-diagonal block = cross-cov)
up.covariance(z, wrt=x) # covariance w.r.t. input x only
up.correlation(z)       # correlation matrix
up.jacobian(z, x)       # dz/dx; up.jacobian(z) for base-input sensitivities
up.dof(z, wrt=x)        # effective dof (Welch-Satterthwaite), optionally wrt
up.budget(z)            # per-input contributions, printable table
up.coverage_interval(z) # (lower, upper) at 95%; Student-t at the effective
                        # dof, e.g. up.coverage_interval(z, 0.99)
```

Everything returns plain Python scalars (0-d) or numpy arrays, and every
statistic is also available inline on the object: `z.stdunc`, `z.dof`,
`z.jacobian([x, y])`, `z.covariance()`, `z.correlation()`, `z.budget()`,
`z.inputs()`.

## Precision

JAX defaults to 32-bit floats. On import, uncpy switches JAX to 64-bit,
the double precision NumPy users expect. This switch is global to the
Python process: any other JAX code in the same program becomes 64-bit
as well.

If you embed uncpy in a JAX application that must stay 32-bit, set the
environment variable before the first import of uncpy:

```python
import os
os.environ["UNCPY_DISABLE_X64"] = "1"   # must run before "import uncpy"

import uncpy as up   # JAX stays at its 32-bit defaults
```

You can also flip the JAX switch directly at any point in a running
program:

```python
import jax
jax.config.update("jax_enable_x64", False)   # back to JAX defaults
jax.config.update("jax_enable_x64", True)    # back to double precision
```

## Notes

- Propagation is first order. Strongly nonlinear models with large
  uncertainties should be cross-checked with Monte Carlo (the test suite
  shows how).
- Comparisons (`<`, `==`, ...) act on nominal values and return plain arrays.
- `linalg.svd` propagates only the thin form (the default); `eigh` assumes
  non-degenerate eigenvalues; `eig` additionally assumes no ties in
  eigenvector components and prefers complex-typed input for real matrices
  with complex-conjugate eigenvalue pairs.
- Correlated elementary inputs must be created together (one `covariance=`);
  there is no post-hoc correlation between existing quantities.
- Uncertain quantities pickle, and input identity survives: what was
  correlated stays correlated, what was independent stays independent,
  including across processes (`multiprocessing`, saved results).
- `stdunc` and `budget` stay O(m) on large independent inputs, but a full
  `covariance` is inherently m×m; for a 100 000-point sweep that is 80 GB,
  and uncpy raises rather than trying.

## About this library

I have always liked working with NumPy, and JAX's autodiff capabilities impressed me. This is a project I have wanted for years but was never able to build myself, so I vibe coded it 😉. I will keep updating it as I learn more about autodiff and how to do it better. If you find bugs or run into trouble using it, please open an issue.

## License

[![Apache License 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://choosealicense.com/licenses/apache-2.0/)
