# JNLR Jax-based non-linear reconciliation and learning 

<p align="center">
  <img src="docs/logo.webp" alt="JNLR Logo" width="300">
</p>

**J-NLR** is a Python library for non-linear reconciliation, learning, and geometric analysis on constraint manifolds. Built on [JAX](https://github.com/google/jax), it leverages automatic differentiation and GPU/TPU acceleration to efficiently project predicted values onto surfaces defined by implicit constraints $f(z) = 0$.

📚 **[Documentation](https://supsi-dacd-isaac.github.io/JNLR/)** - Full API reference, examples, and interactive notebooks

## Key Features

- **Non-linear Reconciliation**: Multiple solvers (Augmented Lagrangian, curvature-aware Newton, vanilla projections) for projecting forecasts onto constraint manifolds
- **SHOULD Analysis**: Curvature-based methods to determine *when* reconciliation is beneficial—verify if RMSE is guaranteed to reduce before applying corrections
- **Manifold Sampling**: Sample from explicit (graph) or implicit manifolds using volume-weighted sampling, Latin hypercube, or Langevin dynamics on the constraint surface
- **Mesh Generation**: Create triangulated meshes from explicit parameterizations for visualization and geodesic computation
- **Geodesics**: Compute geodesic distances and shortest paths on manifolds via exact MMP algorithm or fast graph-based approximations; includes probabilistic scores like pointcloud geodesic distance
- **Visualization**: Interactive 3D rendering of manifolds, projections, and geodesic paths with Plotly
- **JAX-native**: Fully JIT-compiled and vectorized (`vmap`) for high-performance batch processing


## Running the notebooks
If you want to run the notebooks, the suggested way is to install `uv` packet manager, cloning the repo and, from a terminal:
`uv pip install -e .`


## Citation

If you use **JNLR** in academic work, please cite the associated paper:

Lorenzo Nespoli, Anubhab Biswas, Roberto Rocchetta, and Vasco Medici.  
"Nonlinear reconciliation: Error reduction theorems."  
Transactions on Machine Learning Research (TMLR), 2026.  
OpenReview: https://openreview.net/forum?id=dXRWuogm3J

### BibTeX

```bibtex
@article{nespoli2026nonlinear_reconciliation,
  title   = {Nonlinear reconciliation: Error reduction theorems},
  author  = {Nespoli, Lorenzo and Biswas, Anubhab and Rocchetta, Roberto and Medici, Vasco},
  journal = {Transactions on Machine Learning Research},
  year    = {2026},
  url     = {https://openreview.net/forum?id=dXRWuogm3J},
  note    = {Accepted by TMLR}
}
```

## Acknowledgements

This work has been funded by the Swiss State Secretariat for Education, Research and Innovation (SERI) under the Swiss contribution to the Horizon Europe projects DR-RISE (Horizon Europe, Grant Agreement No. 101104154) and REEFLEX (Horizon Europe, Grant Agreement No. 101096192).
