contrax 0.1.0__tar.gz

contrax-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Givani Boekestijn
+
+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.

contrax-0.1.0/PKG-INFO

@@ -0,0 +1,172 @@
+Metadata-Version: 2.4
+Name: contrax
+Version: 0.1.0
+Summary: JAX-native control theory library
+Author-email: Givani Boekestijn <givaniboek@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/givani30/Contrax
+Project-URL: Repository, https://github.com/givani30/Contrax
+Project-URL: Bug Tracker, https://github.com/givani30/Contrax/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: diffrax>=0.7.2
+Requires-Dist: equinox>=0.13.6
+Requires-Dist: jax>=0.9.2
+Requires-Dist: lineax>=0.0.8
+Requires-Dist: optimistix>=0.1.0
+Provides-Extra: dev
+Requires-Dist: build>=1.2.2; extra == "dev"
+Requires-Dist: matplotlib>=3.10.8; extra == "dev"
+Requires-Dist: mkdocstrings-python>=2.0.3; extra == "dev"
+Requires-Dist: pillow>=12.2.0; extra == "dev"
+Requires-Dist: pre-commit>=4.0.1; extra == "dev"
+Requires-Dist: pytest>=9.0.3; extra == "dev"
+Requires-Dist: pytest-xdist>=3.8.0; extra == "dev"
+Requires-Dist: ruff>=0.11.13; extra == "dev"
+Requires-Dist: zensical>=0.0.32; extra == "dev"
+Dynamic: license-file
+
+# Contrax
+
+**Differentiable control theory primitives for JAX.**
+
+[![CI](https://github.com/givani30/Contrax/actions/workflows/ci.yml/badge.svg)](https://github.com/givani30/Contrax/actions/workflows/ci.yml)
+[![Docs](https://img.shields.io/badge/docs-online-teal)](https://givani30.github.io/Contrax/)
+[![PyPI](https://img.shields.io/pypi/v/contrax)](https://pypi.org/project/contrax/)
+[![Python](https://img.shields.io/pypi/pyversions/contrax)](https://pypi.org/project/contrax/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+Contrax is a JAX-native systems, estimation, and control toolbox.
+MATLAB-familiar names at the API surface — `ss`, `lqr`, `kalman`, `place` —
+with `jit`, `vmap`, and `grad` behavior underneath.
+
+**[Documentation](https://givani30.github.io/Contrax/)** · [Getting started](https://givani30.github.io/Contrax/getting-started/) · [API reference](https://givani30.github.io/Contrax/api/systems/)
+
+---
+
+## Install
+
+```bash
+pip install contrax
+```
+
+Requires Python 3.11+ and JAX 0.4+. For GPU support, install JAX separately
+per the [JAX installation guide](https://jax.readthedocs.io/en/latest/installation.html)
+before installing Contrax.
+
+## Why Contrax
+
+A controller design step can sit inside an ordinary differentiable objective
+instead of being a separate offline calculation:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp
+import contrax as cx
+
+A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
+B = jnp.array([[0.0], [0.05]])
+SYS = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)
+X0 = jnp.array([1.0, 0.0])
+
+
+def closed_loop_cost(log_q_diag, log_r):
+    K = cx.lqr(SYS, jnp.diag(jnp.exp(log_q_diag)), jnp.exp(log_r)[None, None]).K
+    _, xs, _ = cx.simulate(SYS, X0, lambda t, x: -K @ x, num_steps=80)
+    return jnp.sum(xs**2) + 1e-2 * jnp.sum((xs[:-1] @ K.T) ** 2)
+
+
+cost, (dq, dr) = jax.jit(jax.value_and_grad(closed_loop_cost, argnums=(0, 1)))(
+    jnp.zeros(2), jnp.array(0.0)
+)
+```
+
+This is the central Contrax idea: control primitives that behave like normal
+JAX building blocks.
+
+## Scope
+
+**Systems**
+`ss`, `dss`, `c2d`, `nonlinear_system`, `phs_system`, `canonical_J`,
+`schedule_phs`, `linearize`, `linearize_ss`, `phs_to_ss`, `series`, `parallel`
+
+**Simulation**
+`lsim`, `simulate`, `rollout`, `step_response`, `impulse_response`,
+`initial_response`, `sample_system`, `foh_inputs`
+
+**Control**
+`lqr`, `lqi`, `dare`, `care`, `place`, `feedback`, `state_feedback`,
+`augment_integrator`
+
+**Estimation**
+`kalman`, `ekf`, `ukf`, `rts`, `uks`
+`kalman_predict`, `kalman_update`, `kalman_step`
+`ekf_predict`, `ekf_update`, `ekf_step`
+`kalman_gain`, `mhe_objective`, `mhe`, `mhe_warm_start`
+
+**Analysis**
+`ctrb`, `obsv`, `poles`, `evalfr`, `freqresp`, `dcgain`,
+`ctrb_gramian`, `obsv_gramian`, `lyap`, `dlyap`, `zeros`
+
+**Diagnostics**
+`innovation_diagnostics`, `likelihood_diagnostics`, `ukf_diagnostics`,
+`smoother_diagnostics`, `phs_diagnostics`, `innovation_rms`
+
+**Parameterization**
+`positive_exp`, `positive_softplus`, `spd_from_cholesky_raw`, `diagonal_spd`,
+`lower_triangular`
+
+**Interoperability**
+`contrax.compat.python_control` — optional bidirectional conversion with
+[python-control](https://python-control.readthedocs.io/) (`pip install control`)
+
+## Solver Status
+
+- **`dare`** — structured-doubling forward solve with implicit-differentiation VJP. Most mature path.
+- **`care`** — Hamiltonian stable-subspace solver with implicit backward pass. Validated, less benchmarked than `dare`.
+- **`place`** — JAX-native KNV0/YT-style robust pole placement; Ackermann retained as SISO fallback only.
+- **`ekf` / `ukf` / `rts` / `uks`** — differentiable nonlinear filtering and smoothing with full JAX transform support.
+- **`mhe`** — LBFGS-backed fixed-window MHE. Useful optimization-based estimation primitive; not a full NLP framework.
+
+LTI workflows are the most mature slice. Nonlinear models, PHS support, and
+fixed-window MHE are real public capabilities, but they should be read with
+more explicit solver-maturity caution than the core discrete design path.
+
+## Development
+
+```bash
+git clone https://github.com/givani30/Contrax.git
+cd Contrax
+uv sync --group dev
+uv run pre-commit install
+uv run pytest tests/ -q
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines.
+
+## Acknowledgements
+
+Contrax is built on top of several excellent JAX-ecosystem libraries:
+
+- **[JAX](https://github.com/google/jax)** — the foundation: JIT compilation, autodiff, and vectorization.
+- **[Equinox](https://github.com/patrick-kidger/equinox)** — pytree-compatible modules used for all result bundles and system types.
+- **[Diffrax](https://github.com/patrick-kidger/diffrax)** — ODE solvers powering `simulate()` on continuous-time models.
+- **[Optimistix](https://github.com/patrick-kidger/optimistix)** — LBFGS solver backing `mhe()`.
+- **[Lineax](https://github.com/patrick-kidger/lineax)** — linear solvers used in the estimation and Riccati paths.
+
+The DARE structured-doubling custom VJP is adapted from
+[trajax](https://github.com/google/trajax) (Google) and the DiLQR approach
+(ICML 2025). The UKF sigma-point rules follow Wan & van der Merwe (2000).

contrax-0.1.0/README.md

@@ -0,0 +1,133 @@
+# Contrax
+
+**Differentiable control theory primitives for JAX.**
+
+[![CI](https://github.com/givani30/Contrax/actions/workflows/ci.yml/badge.svg)](https://github.com/givani30/Contrax/actions/workflows/ci.yml)
+[![Docs](https://img.shields.io/badge/docs-online-teal)](https://givani30.github.io/Contrax/)
+[![PyPI](https://img.shields.io/pypi/v/contrax)](https://pypi.org/project/contrax/)
+[![Python](https://img.shields.io/pypi/pyversions/contrax)](https://pypi.org/project/contrax/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+Contrax is a JAX-native systems, estimation, and control toolbox.
+MATLAB-familiar names at the API surface — `ss`, `lqr`, `kalman`, `place` —
+with `jit`, `vmap`, and `grad` behavior underneath.
+
+**[Documentation](https://givani30.github.io/Contrax/)** · [Getting started](https://givani30.github.io/Contrax/getting-started/) · [API reference](https://givani30.github.io/Contrax/api/systems/)
+
+---
+
+## Install
+
+```bash
+pip install contrax
+```
+
+Requires Python 3.11+ and JAX 0.4+. For GPU support, install JAX separately
+per the [JAX installation guide](https://jax.readthedocs.io/en/latest/installation.html)
+before installing Contrax.
+
+## Why Contrax
+
+A controller design step can sit inside an ordinary differentiable objective
+instead of being a separate offline calculation:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp
+import contrax as cx
+
+A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
+B = jnp.array([[0.0], [0.05]])
+SYS = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)
+X0 = jnp.array([1.0, 0.0])
+
+
+def closed_loop_cost(log_q_diag, log_r):
+    K = cx.lqr(SYS, jnp.diag(jnp.exp(log_q_diag)), jnp.exp(log_r)[None, None]).K
+    _, xs, _ = cx.simulate(SYS, X0, lambda t, x: -K @ x, num_steps=80)
+    return jnp.sum(xs**2) + 1e-2 * jnp.sum((xs[:-1] @ K.T) ** 2)
+
+
+cost, (dq, dr) = jax.jit(jax.value_and_grad(closed_loop_cost, argnums=(0, 1)))(
+    jnp.zeros(2), jnp.array(0.0)
+)
+```
+
+This is the central Contrax idea: control primitives that behave like normal
+JAX building blocks.
+
+## Scope
+
+**Systems**
+`ss`, `dss`, `c2d`, `nonlinear_system`, `phs_system`, `canonical_J`,
+`schedule_phs`, `linearize`, `linearize_ss`, `phs_to_ss`, `series`, `parallel`
+
+**Simulation**
+`lsim`, `simulate`, `rollout`, `step_response`, `impulse_response`,
+`initial_response`, `sample_system`, `foh_inputs`
+
+**Control**
+`lqr`, `lqi`, `dare`, `care`, `place`, `feedback`, `state_feedback`,
+`augment_integrator`
+
+**Estimation**
+`kalman`, `ekf`, `ukf`, `rts`, `uks`
+`kalman_predict`, `kalman_update`, `kalman_step`
+`ekf_predict`, `ekf_update`, `ekf_step`
+`kalman_gain`, `mhe_objective`, `mhe`, `mhe_warm_start`
+
+**Analysis**
+`ctrb`, `obsv`, `poles`, `evalfr`, `freqresp`, `dcgain`,
+`ctrb_gramian`, `obsv_gramian`, `lyap`, `dlyap`, `zeros`
+
+**Diagnostics**
+`innovation_diagnostics`, `likelihood_diagnostics`, `ukf_diagnostics`,
+`smoother_diagnostics`, `phs_diagnostics`, `innovation_rms`
+
+**Parameterization**
+`positive_exp`, `positive_softplus`, `spd_from_cholesky_raw`, `diagonal_spd`,
+`lower_triangular`
+
+**Interoperability**
+`contrax.compat.python_control` — optional bidirectional conversion with
+[python-control](https://python-control.readthedocs.io/) (`pip install control`)
+
+## Solver Status
+
+- **`dare`** — structured-doubling forward solve with implicit-differentiation VJP. Most mature path.
+- **`care`** — Hamiltonian stable-subspace solver with implicit backward pass. Validated, less benchmarked than `dare`.
+- **`place`** — JAX-native KNV0/YT-style robust pole placement; Ackermann retained as SISO fallback only.
+- **`ekf` / `ukf` / `rts` / `uks`** — differentiable nonlinear filtering and smoothing with full JAX transform support.
+- **`mhe`** — LBFGS-backed fixed-window MHE. Useful optimization-based estimation primitive; not a full NLP framework.
+
+LTI workflows are the most mature slice. Nonlinear models, PHS support, and
+fixed-window MHE are real public capabilities, but they should be read with
+more explicit solver-maturity caution than the core discrete design path.
+
+## Development
+
+```bash
+git clone https://github.com/givani30/Contrax.git
+cd Contrax
+uv sync --group dev
+uv run pre-commit install
+uv run pytest tests/ -q
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines.
+
+## Acknowledgements
+
+Contrax is built on top of several excellent JAX-ecosystem libraries:
+
+- **[JAX](https://github.com/google/jax)** — the foundation: JIT compilation, autodiff, and vectorization.
+- **[Equinox](https://github.com/patrick-kidger/equinox)** — pytree-compatible modules used for all result bundles and system types.
+- **[Diffrax](https://github.com/patrick-kidger/diffrax)** — ODE solvers powering `simulate()` on continuous-time models.
+- **[Optimistix](https://github.com/patrick-kidger/optimistix)** — LBFGS solver backing `mhe()`.
+- **[Lineax](https://github.com/patrick-kidger/lineax)** — linear solvers used in the estimation and Riccati paths.
+
+The DARE structured-doubling custom VJP is adapted from
+[trajax](https://github.com/google/trajax) (Google) and the DiLQR approach
+(ICML 2025). The UKF sigma-point rules follow Wan & van der Merwe (2000).

contrax-0.1.0/contrax/__init__.py

@@ -0,0 +1,197 @@
+from contrax.analysis import (
+    ctrb,
+    ctrb_gramian,
+    dcgain,
+    dlyap,
+    evalfr,
+    freqresp,
+    lyap,
+    obsv,
+    obsv_gramian,
+    poles,
+    zeros,
+)
+from contrax.control import (
+    augment_integrator,
+    care,
+    dare,
+    feedback,
+    lqi,
+    lqr,
+    place,
+    state_feedback,
+)
+from contrax.estimation import (
+    ekf,
+    ekf_predict,
+    ekf_step,
+    ekf_update,
+    innovation_diagnostics,
+    innovation_rms,
+    kalman,
+    kalman_gain,
+    kalman_predict,
+    kalman_step,
+    kalman_update,
+    likelihood_diagnostics,
+    mhe,
+    mhe_objective,
+    mhe_warm_start,
+    rts,
+    smoother_diagnostics,
+    soft_quadratic_penalty,
+    ukf,
+    ukf_diagnostics,
+    uks,
+)
+from contrax.parameterization import (
+    diagonal_spd,
+    lower_triangular,
+    positive_exp,
+    positive_softplus,
+    spd_from_cholesky_raw,
+)
+from contrax.simulation import (
+    as_ode_term,
+    foh_inputs,
+    impulse_response,
+    initial_response,
+    lsim,
+    rollout,
+    sample_system,
+    simulate,
+    step_response,
+)
+from contrax.systems import (
+    ContLTI,
+    DiscLTI,
+    NonlinearSystem,
+    PHSSystem,
+    block_matrix,
+    block_observation,
+    c2d,
+    canonical_J,
+    dss,
+    linearize,
+    linearize_ss,
+    nonlinear_system,
+    parallel,
+    partition_state,
+    phs_diagnostics,
+    phs_system,
+    phs_to_ss,
+    project_psd,
+    schedule_phs,
+    series,
+    ss,
+    symmetrize_matrix,
+)
+from contrax.types import (
+    InnovationDiagnostics,
+    KalmanGainResult,
+    KalmanResult,
+    LikelihoodDiagnostics,
+    LQRResult,
+    MHEResult,
+    PHSStructureDiagnostics,
+    RTSResult,
+    SmootherDiagnostics,
+    UKFResult,
+)
+
+__all__ = [
+    # system types
+    "ContLTI",
+    "DiscLTI",
+    "NonlinearSystem",
+    "PHSSystem",
+    "partition_state",
+    "block_observation",
+    "block_matrix",
+    "symmetrize_matrix",
+    "project_psd",
+    "phs_diagnostics",
+    # constructors
+    "ss",
+    "dss",
+    "c2d",
+    "nonlinear_system",
+    "phs_system",
+    "phs_to_ss",
+    "schedule_phs",
+    "canonical_J",
+    "linearize",
+    "linearize_ss",
+    "series",
+    "parallel",
+    # control
+    "lqr",
+    "lqi",
+    "dare",
+    "care",
+    "place",
+    "state_feedback",
+    "feedback",
+    "augment_integrator",
+    # simulation
+    "rollout",
+    "foh_inputs",
+    "sample_system",
+    "lsim",
+    "simulate",
+    "step_response",
+    "impulse_response",
+    "initial_response",
+    "as_ode_term",
+    # estimation
+    "kalman",
+    "kalman_gain",
+    "kalman_predict",
+    "kalman_update",
+    "kalman_step",
+    "innovation_diagnostics",
+    "innovation_rms",
+    "likelihood_diagnostics",
+    "ekf",
+    "ekf_predict",
+    "ekf_update",
+    "ekf_step",
+    "mhe_objective",
+    "mhe_warm_start",
+    "mhe",
+    "soft_quadratic_penalty",
+    "smoother_diagnostics",
+    "ukf",
+    "ukf_diagnostics",
+    "uks",
+    "rts",
+    # analysis
+    "ctrb",
+    "obsv",
+    "poles",
+    "zeros",
+    "evalfr",
+    "freqresp",
+    "dcgain",
+    "ctrb_gramian",
+    "obsv_gramian",
+    "lyap",
+    "dlyap",
+    # parameterization
+    "positive_exp",
+    "positive_softplus",
+    "lower_triangular",
+    "spd_from_cholesky_raw",
+    "diagonal_spd",
+    # result types
+    "LQRResult",
+    "InnovationDiagnostics",
+    "SmootherDiagnostics",
+    "KalmanGainResult",
+    "KalmanResult",
+    "LikelihoodDiagnostics",
+    "UKFResult",
+    "RTSResult",
+    "MHEResult",
+    "PHSStructureDiagnostics",
+]