epgx 0.1.dev0__py3-none-any.whl

epgx/stats.py

@@ -0,0 +1,117 @@
+"""Cramer-Rao lower bound and confidence intervals
+
+Differentiable (jnp): the gradient of the CRLB w.r.t. sequence parameters is a
+`jax.grad` (no separate gradient function).
+
+Layout (that of `simulate` and `diff.derivatives`):
+    J: (nobs, *batch, nparam)          Jacobian, observations first
+    H: (nobs, *batch, nparam, nparam)  Hessian
+    obs, pred: (nobs, *batch)
+Results: (*batch, ...). Singular information matrices (condition number above
+1 / machine epsilon) give NaN.
+"""
+
+import jax
+import jax.numpy as jnp
+
+
+def crlb(J, *, W=None, sigma2=1, log=False):
+    """Cramer-Rao lower bound: sum of the weighted variance bounds
+
+    Args:
+        J: Jacobian (nobs, *batch, nparam)
+        W: weights of the parameters (nparam,), default: 1
+        sigma2: noise variance
+        log: return log10 of the bound
+
+    Returns:
+        (*batch,) bound
+    """
+    bounds = crlb_split(J, W=W, sigma2=sigma2)
+    cost = jnp.sum(bounds, axis=-1)
+    return jnp.log10(cost) if log else cost
+
+
+def crlb_split(J, *, W=None, sigma2=1, log=False):
+    """Cramer-Rao lower bound of each parameter
+
+    Returns:
+        (*batch, nparam) weighted variance bounds
+    """
+    lb = _inverse(_information(J) / sigma2)
+    bounds = jnp.diagonal(lb, axis1=-2, axis2=-1)
+    if W is not None:
+        bounds = bounds * jnp.asarray(W)
+    return jnp.log10(bounds) if log else bounds
+
+
+def confint(obs, pred, J, H=None, *, conflevel=0.95):
+    """confidence intervals of the parameters and band of the prediction (delta method)
+
+    Args:
+        obs, pred: observed and predicted signals (nobs, *batch)
+        J: Jacobian of the prediction (nobs, *batch, nparam)
+        H: Hessian of the prediction (nobs, *batch, nparam, nparam), optional:
+            covariance from the Hessian of the least-squares cost instead of
+            the Gauss-Newton approximation
+        conflevel: confidence level
+
+    Returns:
+        intervals (*batch, nparam), band (nobs, *batch): half-widths
+    """
+    J = jnp.asarray(J)
+    nobs, nparam = J.shape[0], J.shape[-1]
+    dof = nobs - nparam
+    res = jnp.asarray(obs) - jnp.asarray(pred)
+    sse = jnp.sum(jnp.abs(res) ** 2, axis=0)
+
+    # Hessian of the cost sum |res|^2 / 2 w.r.t. the parameters
+    info = _information(J)
+    if H is not None:
+        info = info - jnp.einsum("n...,n...pq->...pq", res.conj(), H).real
+    cov = _inverse(info) * (sse / dof)[..., None, None]
+
+    tval = t_quantile((1 + conflevel) / 2, dof)
+    intervals = tval * jnp.sqrt(jnp.diagonal(cov, axis1=-2, axis2=-1))
+    band = tval * jnp.sqrt(jnp.einsum("n...p,...pq,n...q->n...", J.conj(), cov, J).real)
+    return intervals, band
+
+
+def t_quantile(p, dof):
+    """quantile of Student's t distribution (p > 1/2), without scipy
+
+    CDF from the regularized incomplete beta function, inverted by bisection.
+    """
+    p, dof = jnp.asarray(p, dtype=float), jnp.asarray(dof, dtype=float)
+
+    def cdf(t):
+        return 1 - 0.5 * jax.scipy.special.betainc(dof / 2, 0.5, dof / (dof + t**2))
+
+    def bisect(_, bounds):
+        low, high = bounds
+        mid = (low + high) / 2
+        above = cdf(mid) > p
+        return jnp.where(above, low, mid), jnp.where(above, mid, high)
+
+    low, high = jax.lax.fori_loop(
+        0, 100, bisect, (jnp.zeros_like(p), jnp.full_like(p, 1e6))
+    )
+    return (low + high) / 2
+
+
+#
+# private
+
+
+def _information(J):
+    """Fisher information J^H J, real part (*batch, nparam, nparam)"""
+    J = jnp.asarray(J)
+    return jnp.einsum("n...p,n...q->...pq", J.conj(), J).real
+
+
+def _inverse(matrix):
+    """inverse, NaN for singular matrices"""
+    cond = jnp.linalg.cond(matrix)
+    singular = ~(cond < 1 / jnp.finfo(matrix.dtype).eps)  # also NaN/inf
+    inverse = jnp.linalg.inv(matrix)
+    return jnp.where(singular[..., None, None], jnp.nan, inverse)

epgx-0.1.dev0.dist-info/METADATA

@@ -0,0 +1,64 @@
+Metadata-Version: 2.4
+Name: epgx
+Version: 0.1.dev0
+Summary: Extended Phase Graph (EPG) simulations in JAX
+Author: Pierre-Yves Baudin
+License-Expression: BSD-3-Clause
+Project-URL: Repository, https://github.com/py-baudin/epgx
+Keywords: MRI,extended phase graph,EPG,simulation,JAX,differentiable
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: numpy
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Provides-Extra: plot
+Requires-Dist: matplotlib; extra == "plot"
+Requires-Dist: cmocean; extra == "plot"
+Provides-Extra: xarray
+Requires-Dist: xarray; extra == "xarray"
+Provides-Extra: dev
+Requires-Dist: epgx[plot,test,xarray]; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+# epgx
+
+A JAX implementation of the Extended Phase Graph (EPG) formalism, successor of
+[epgpy](https://github.com/py-baudin/epgpy).
+
+Work in progress: the API may change between versions.
+
+## Installation
+
+``` shell
+pip install epgx            # core (jax, numpy)
+pip install epgx[plot]      # with EPG diagrams (matplotlib, cmocean)
+```
+
+Development:
+
+``` shell
+conda create -n epgx python=3.12
+conda activate epgx
+pip install -e .[dev]
+pip install --group parity  # epgpy, for the parity tests (tests/parity)
+pytest
+```
+
+## Precision
+
+`epgx` follows JAX's precision setting (float32 by default) and never changes it.
+For float64, enable it at the start of your program:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+```

epgx-0.1.dev0.dist-info/RECORD

@@ -0,0 +1,38 @@
+epgx/__init__.py,sha256=3aDyN8qXa_0Zd-P8X5XDNyDrynI8lx1f7zRJHOidVvs,588
+epgx/axes.py,sha256=U_VfGBsJCSuQu6mEbO1TR83AmcGlaI5nqovgNUEfQMY,10248
+epgx/common.py,sha256=34x4UaU5edDfkBOA41yoSPLt4RbNAu5OJpQ-umXzNbo,2055
+epgx/compiler.py,sha256=-M-8ILKbpmlH5qjANbTajQxgR6lamVsGRp9k7ZbcWnc,11986
+epgx/diff.py,sha256=BFACor42z9ibVrE4jc_zON3yOItCgYoGsS3t_l0Nwkk,8855
+epgx/examples.py,sha256=ZGauA9QKsd7uyzJEhGshyC1kPx6DXmCl5Ye8X1qPmVo,4888
+epgx/functions.py,sha256=-UkbcBPLUNW09TA-GInhlt-KUhKZqsDrlxW2ROerWdA,22506
+epgx/lowering.py,sha256=oKTkbCL40J6-hnWeoVoOJ26txMZZXzyszHsbdg2l2BA,3107
+epgx/merging.py,sha256=gsoiEYeXbm6v_S_5bf8V2WFVdMhi-V-somJMqu0t5gI,4560
+epgx/methods.py,sha256=XTwZWctRGAduaVHi80rHEcEcAIIJbQj4tskx_vtRhIA,12740
+epgx/plotting.py,sha256=Ww9LAuCZvHmfY0c90tprZVHEqJymjdRLrsRLR9gWAMI,8219
+epgx/progress.py,sha256=Nv1L-CaEC2-Y6bIwYvci8dky7YWCPN-u6xO90XJznjA,2752
+epgx/sample.py,sha256=Pb_lVea2-1MBiPNXvdKQDhg8IiwqaI2nHEnVFihVFdQ,4506
+epgx/sequences.py,sha256=_IzV4_DzwLWAm4AzXASmGT4pe2_btZb-A6jS1xgzj7M,32483
+epgx/stats.py,sha256=nmVUEPrLep7vju1hORdmLVOUUIlxFhDIql3xK2KbDhw,3792
+epgx/kernels/__init__.py,sha256=gfuTvWuy8Z_byah2rY6wv4_nmYcinf7zll2HKN3LYJs,356
+epgx/kernels/affine.py,sha256=tzlMSwKONd1uzfOxjuDbDTroFNuDsdktWTEoM0hjyy8,6168
+epgx/kernels/evolution.py,sha256=0MrFtXVUHQKQxOCsTRStvYipwCpp4XR70eBqSWh0Ox8,1449
+epgx/kernels/operator.py,sha256=t4HDsFdmQzkvlWHL7z4_eiE985nNguSWK2YsLVhdch4,1139
+epgx/kernels/probe.py,sha256=0lwX_k7sFatbVc7z5xew-59lqUl7ccUPqP9If-4wrnI,3865
+epgx/kernels/shift.py,sha256=_vYxLjnOkSGLCy3FIOSmL0cKr6uKVQ9ITqFgOgkqpfw,3464
+epgx/kernels/state.py,sha256=yuVfaa-_K1R1A6dZA7lObchRrbM1qoKPNoBVg0A3Eiw,16289
+epgx/kernels/transition.py,sha256=cuQ3xfM9NUs89nkKfSMK3xTIs_SGvlMkV45nL--rJuU,1279
+epgx/library/__init__.py,sha256=KVvf8PLA1_0fV5Hh3nn7OnOtheAJxEq5RYB9MuOHwlI,109
+epgx/library/sequences.py,sha256=njOKEmyHMG3omb0y-f6PRbUBe3nqOxwRntvGOyuqpv4,7461
+epgx/ops/__init__.py,sha256=Z-SMFeGhEt4no3gcWTwqtPUelpJ3iEzRM2YEAOPI9cY,593
+epgx/ops/dephasing.py,sha256=A9dhsGbXIDsnI_SN4CsQcN_wON1QX5G46YD0H71e-j0,2752
+epgx/ops/evolution.py,sha256=JkPnFhP6g4Iywt4N9YcsbWqeATKoX5RH0E9Uve36vM4,4434
+epgx/ops/functions.py,sha256=rqmLt9Kj3XFHz4UT-BK17jrhWmp4d51n9wVJG38XPaw,5291
+epgx/ops/operator.py,sha256=N5vg9S2qYegcM9z-anAxdyS6O-qp2_XOfJNMNdd4RCg,7290
+epgx/ops/probe.py,sha256=X5W_281dALFRDRfHqaJD2qm8sZ7VhlM49B-D51R00UE,3912
+epgx/ops/shift.py,sha256=UcNKgygVN3BjUaElUafxne-7KRhwNjYjyqnHKxGh_Lk,4303
+epgx/ops/transition.py,sha256=M00yhd-9GRR1CCuk4fhcgCJOCwDBS8gUdyapHjxKB8c,1491
+epgx-0.1.dev0.dist-info/licenses/LICENSE,sha256=TLhi4uBEksu2oRv12uW3Ar9k9OPMF_cqyVANPiP6uwA,1501
+epgx-0.1.dev0.dist-info/METADATA,sha256=fWwymCNLH1o6Ca-dOTc_62NVC2as9mdpyfK8cABXcEo,1838
+epgx-0.1.dev0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+epgx-0.1.dev0.dist-info/top_level.txt,sha256=mvERoxaeTHo7FgZKTjnE3kn5MKwvIVYiIyKZisyqVMU,5
+epgx-0.1.dev0.dist-info/RECORD,,

epgx-0.1.dev0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

epgx-0.1.dev0.dist-info/licenses/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023-2026, py-baudin
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

epgx-0.1.dev0.dist-info/top_level.txt

@@ -0,0 +1 @@
+epgx