API reference

Auto-generated from the source docstrings.

Estimator

jax_ibart.ibart.IBart

Iterative BART descriptor discovery.

Parameters:

Name	Type	Description	Default
backend	('bart', 'diff')	Selection backend; see :mod:jax_ibart.selection.	"bart"
schedule	tuple of {"u", "b"}	Feature-engineering steps to interleave with selection.	DEFAULT_SCHEDULE
max_dim	int	Maximum size of the final descriptor (l0 best-subset).	3
n_keep	int	Candidates passed from LASSO screening into the l0 search.	10
score	('bic', 'cv')	Selection score for the l0 best-subset.	"bic"
seed	int	PRNG seed for the (functional) selection randomness.	0
unary_ops	list of :class:`~jax_ibart.operators.Operator`	Operator sets for the "u" and "b" steps. Default to :func:~jax_ibart.operators.default_unary and :func:~jax_ibart.operators.default_binary (the catalysis-study sets).	None
binary_ops	list of :class:`~jax_ibart.operators.Operator`	Operator sets for the "u" and "b" steps. Default to :func:~jax_ibart.operators.default_unary and :func:~jax_ibart.operators.default_binary (the catalysis-study sets).	None
carry_selected	bool	If True, each screen's survivors are unioned with everything kept by earlier screens, and that union feeds the next engineering step and the final LASSO/l0 stage. This is what the R iBART package does (X_selected accumulates across iterations). If False (default), a binary step replaces the pool, so a descriptor kept in one round can drop out of the final pool.	False

Attributes (after fit)

descriptors_ : list[sympy.Expr] The selected descriptors (symbolic, interpretable). coef_ : np.ndarray OLS coefficients [intercept, *slopes] aligned with descriptors_. history_ : list[dict] One record per pipeline stage, {"stage", "generated", "kept"}: the pool size after feature engineering and after screening. Shows the grow-then-prune rhythm that keeps iBART tractable. pool_ : engineer.Pool The final screened descriptor pool that the LASSO/l0 stage chose from.

Source code in src/jax_ibart/ibart.py

class IBart:
    """Iterative BART descriptor discovery.

    Parameters
    ----------
    backend : {"bart", "diff"}
        Selection backend; see :mod:`jax_ibart.selection`.
    schedule : tuple of {"u", "b"}
        Feature-engineering steps to interleave with selection.
    max_dim : int
        Maximum size of the final descriptor (l0 best-subset).
    n_keep : int
        Candidates passed from LASSO screening into the l0 search.
    score : {"bic", "cv"}
        Selection score for the l0 best-subset.
    seed : int
        PRNG seed for the (functional) selection randomness.
    unary_ops, binary_ops : list of :class:`~jax_ibart.operators.Operator`, optional
        Operator sets for the ``"u"`` and ``"b"`` steps. Default to
        :func:`~jax_ibart.operators.default_unary` and
        :func:`~jax_ibart.operators.default_binary` (the catalysis-study sets).
    carry_selected : bool
        If True, each screen's survivors are unioned with everything kept by
        earlier screens, and that union feeds the next engineering step and the
        final LASSO/l0 stage. This is what the R iBART package does
        (``X_selected`` accumulates across iterations). If False (default), a
        binary step replaces the pool, so a descriptor kept in one round can
        drop out of the final pool.

    Attributes (after ``fit``)
    --------------------------
    descriptors_ : list[sympy.Expr]
        The selected descriptors (symbolic, interpretable).
    coef_ : np.ndarray
        OLS coefficients ``[intercept, *slopes]`` aligned with ``descriptors_``.
    history_ : list[dict]
        One record per pipeline stage, ``{"stage", "generated", "kept"}``: the
        pool size after feature engineering and after screening. Shows the
        grow-then-prune rhythm that keeps iBART tractable.
    pool_ : engineer.Pool
        The final screened descriptor pool that the LASSO/l0 stage chose from.
    """

    def __init__(
        self,
        backend: str = "bart",
        schedule=DEFAULT_SCHEDULE,
        max_dim: int = 3,
        n_keep: int = 10,
        score: str = "bic",
        screen_primaries: bool = False,
        seed: int = 0,
        unary_ops=None,
        binary_ops=None,
        carry_selected: bool = False,
        **selector_kwargs,
    ):
        self.backend = backend
        self.schedule = schedule
        self.max_dim = max_dim
        self.n_keep = n_keep
        self.score = score
        # Leading BART screen on the primary features before the first
        # engineering step (the "BART ->" prefix of the paper's schedule). The R
        # iBART package selects primaries first (e.g. 59 -> 12) and only then
        # expands. This is more faithful and much faster, but at small MCMC
        # budgets the leading screen can commit a PAN false negative and drop an
        # important primary before it is ever combined (observed: it dropped Oxv
        # on catalysis). Default off; recommended with a larger MCMC budget.
        # (Found via the R-iBART oracle cross-check.)
        self.screen_primaries = screen_primaries
        self.seed = seed
        self.unary_ops = unary_ops
        self.binary_ops = binary_ops
        self.carry_selected = carry_selected
        self._selector_kwargs = dict(selector_kwargs)
        self.selector = make_selector(backend, **selector_kwargs)
        self.descriptors_ = None
        self.coef_ = None
        self.history_ = None
        self.pool_ = None
        self._feature_names = None

    def get_params(self, deep: bool = True) -> dict:
        """Constructor kwargs, for cloning a fresh unfitted estimator."""
        return dict(
            backend=self.backend,
            schedule=self.schedule,
            max_dim=self.max_dim,
            n_keep=self.n_keep,
            score=self.score,
            screen_primaries=self.screen_primaries,
            seed=self.seed,
            unary_ops=self.unary_ops,
            binary_ops=self.binary_ops,
            carry_selected=self.carry_selected,
            **self._selector_kwargs,
        )

    def set_params(self, **params) -> "IBart":
        """Update constructor kwargs (scikit-learn style); returns ``self``."""
        own = set(self.get_params()) - set(self._selector_kwargs)
        for k, v in params.items():
            if k in own:
                setattr(self, k, v)
            else:
                self._selector_kwargs[k] = v
        self.selector = make_selector(self.backend, **self._selector_kwargs)
        return self

    def clone(self) -> "IBart":
        """A fresh, unfitted copy with the same configuration."""
        return IBart(**self.get_params())

    def fit(self, X, y, feature_names: list[str] | None = None, units=None):
        """Fit iBART.

        ``feature_names`` defaults to the columns of a pandas DataFrame ``X``,
        or ``x0, x1, ...``. Names become the symbols in the discovered
        expressions.

        ``units`` (optional): an (n_features, n_dims) integer unit-exponent
        matrix. When given, descriptors that are dimensionally inconsistent
        (e.g. summing two different units) are never generated.

        The fit always runs JAX in float32, even if 64-bit mode is enabled
        globally (some libraries, e.g. discopt, enable it on import), so results
        do not depend on what else is imported. The final regression is float64.
        """
        if feature_names is None:
            cols = getattr(X, "columns", None)
            p = np.shape(X)[1] if np.ndim(X) == 2 else 0
            feature_names = (
                [str(c) for c in cols]
                if cols is not None
                else [f"x{i}" for i in range(p)]
            )
        X, y, feature_names, units = _check_inputs(X, y, feature_names, units)
        with _x64_off():
            return self._fit(X, y, feature_names, units)

    def _fit(self, X, y, feature_names, units):
        if self.score not in ("bic", "cv"):
            raise ValueError(f"score must be 'bic' or 'cv', got {self.score!r}")
        # Built here so changes to ``backend`` after construction take effect.
        self.selector = make_selector(self.backend, **self._selector_kwargs)
        key = jax.random.PRNGKey(self.seed)
        unary = default_unary() if self.unary_ops is None else list(self.unary_ops)
        binary = default_binary() if self.binary_ops is None else list(self.binary_ops)
        pool = engineer.from_primary(X, feature_names, units=units)
        n = pool.n_descriptors
        history = [{"stage": "primary", "generated": n, "kept": n}]

        # Union of everything kept so far (used when carry_selected).
        selected = None

        # Leading screen on the primary features (matches R iBART).
        if self.screen_primaries and pool.n_descriptors > 1:
            key, sk = jax.random.split(key)
            pool = self._screen(pool, y, sk)
            history[0]["kept"] = pool.n_descriptors
            selected = pool

        for step in self.schedule:
            if step not in ("u", "b"):
                raise ValueError(f"schedule steps must be 'u' or 'b', got {step!r}")
            grown = (
                engineer.apply_unary(pool, unary)
                if step == "u"
                else engineer.apply_binary(pool, binary)
            )
            generated = grown.n_descriptors
            if generated == 0:
                # Nothing valid to engineer (e.g. one feature, or every value
                # outside the operators' domains): stop growing, keep the pool.
                history.append(
                    {"stage": step, "generated": 0, "kept": pool.n_descriptors}
                )
                break
            key, sk = jax.random.split(key)
            pool = self._screen(grown, y, sk)
            if self.carry_selected:
                selected = pool if selected is None else _union(selected, pool)
                pool = selected
            history.append(
                {"stage": step, "generated": generated, "kept": pool.n_descriptors}
            )

        # Final parametric stage: LASSO screen -> l0 best subset (f64 host-side).
        Xpool = np.asarray(pool.X, dtype=np.float64)
        cand = sparsify.lasso_screen(Xpool, y, n_keep=self.n_keep)
        local, coef = sparsify.l0_best_subset(
            Xpool[:, cand], y, max_dim=self.max_dim, score=self.score
        )
        history.append({"stage": "l0", "generated": len(cand), "kept": len(local)})

        # Set fitted state only once the fit has succeeded.
        self._feature_names = feature_names
        self._X_train = X
        self.pool_ = pool
        self.descriptors_ = [pool.exprs[cand[i]] for i in local]
        self.coef_ = coef
        self.history_ = history
        self._check_consistency(X, [cand[i] for i in local])
        return self

    def _check_consistency(self, X, cols):
        """Warn if the symbolic descriptors don't reproduce their numeric columns.

        That happens when a custom operator's ``fn`` and ``sym`` disagree, and
        would make ``predict`` silently differ from the fitted model.
        """
        num = np.asarray(self.pool_.X, dtype=np.float64)[:, cols]
        with np.errstate(all="ignore"):
            sym = self.transform(X)
        scale = np.abs(num).max(axis=0) + 1e-12
        bad = [
            str(self.descriptors_[k])
            for k in range(len(cols))
            if not np.allclose(sym[:, k] / scale[k], num[:, k] / scale[k], atol=1e-3)
        ]
        if bad:
            warnings.warn(
                "symbolic and numeric forms disagree on the training data for "
                f"{bad}; check custom operators' fn and sym",
                stacklevel=3,
            )

    def _screen(self, pool, y, key):
        """Run the selection backend and reduce the pool, preserving units.

        Safeguard: if screening would empty the pool, keep it unchanged.
        """
        keep = np.asarray(self.selector.select(pool.X, y, key=key))
        if not keep.any():
            return pool
        return engineer.Pool(
            X=pool.X[:, keep],
            exprs=[e for e, k in zip(pool.exprs, keep) if bool(k)],
            units=None if pool.units is None else pool.units[keep],
        )

    def _check_fitted(self):
        if self.coef_ is None:
            raise RuntimeError("this IBart is not fitted yet; call fit first")

    def transform(self, X):
        """Evaluate the selected descriptors on raw primary features ``X``.

        Returns the ``(n, k)`` design matrix the final linear model uses.
        """
        self._check_fitted()
        syms = [sp.Symbol(n) for n in self._feature_names]
        X = np.asarray(X, dtype=np.float64)
        if X.ndim != 2 or X.shape[1] != len(syms):
            raise ValueError(
                f"X must have shape (n, {len(syms)}), got {tuple(X.shape)}"
            )
        cols = []
        for expr in self.descriptors_:
            f = sp.lambdify(syms, expr, modules="numpy")
            cols.append(
                np.asarray(f(*[X[:, i] for i in range(X.shape[1])]), dtype=np.float64)
            )
        return np.column_stack(cols)

    _design = transform  # backward-compatible alias

    def predict(self, X):
        """Predict ``y`` for raw primary features ``X``.

        The symbolic descriptors are evaluated as written, so operators such as
        ``log`` and ``sqrt`` give NaN outside their real domain (the training
        pipeline evaluates them on ``|x|``).
        """
        with np.errstate(all="ignore"):
            D = self.transform(X)
        pred = self.coef_[0] + D @ self.coef_[1:]
        if not np.isfinite(pred).all():
            warnings.warn(
                f"{int((~np.isfinite(pred)).sum())} predictions are NaN or inf; "
                "inputs are probably outside a descriptor's domain (e.g. log of a "
                "non-positive value)",
                stacklevel=2,
            )
        return pred

    def expression(self, digits: int = 4) -> sp.Expr:
        """The fitted model as a single symbolic expression.

        Coefficients are rounded to ``digits`` significant figures for display
        (pass ``digits=None`` for full precision).
        """
        self._check_fitted()
        expr = sp.Float(self.coef_[0])
        for c, d in zip(self.coef_[1:], self.descriptors_):
            expr = expr + sp.Float(c) * d
        return sp.N(expr, digits) if digits else expr

    def to_discopt(self, m, x=None, **kwargs):
        """Add this model to a discopt optimization model.

        Returns ``(x, y)``: the input variables by feature name and a variable
        constrained to equal the prediction. Keyword arguments are passed on;
        see :func:`jax_ibart.export.to_discopt`.
        """
        from .export import to_discopt

        return to_discopt(self, m, x, **kwargs)

    def latex(self, lhs: str = "y", digits: int = 4) -> str:
        """The fitted model as a LaTeX string, e.g. ``y = -0.099 BVS Hfo + ...``."""
        return f"{lhs} = {sp.latex(self.expression(digits))}"

    def to_html(self, lhs: str = "y", digits: int = 4) -> str:
        """An HTML snippet rendering the model (MathJax/KaTeX delimiters)."""
        return f"<p>\\({self.latex(lhs, digits)}\\)</p>"

    def _repr_latex_(self) -> str:
        """Rich LaTeX display in Jupyter/notebooks."""
        if self.coef_ is None:
            return "IBart(unfitted)"
        return f"${self.latex()}$"

get_params

get_params(deep=True)

Constructor kwargs, for cloning a fresh unfitted estimator.

Source code in src/jax_ibart/ibart.py

def get_params(self, deep: bool = True) -> dict:
    """Constructor kwargs, for cloning a fresh unfitted estimator."""
    return dict(
        backend=self.backend,
        schedule=self.schedule,
        max_dim=self.max_dim,
        n_keep=self.n_keep,
        score=self.score,
        screen_primaries=self.screen_primaries,
        seed=self.seed,
        unary_ops=self.unary_ops,
        binary_ops=self.binary_ops,
        carry_selected=self.carry_selected,
        **self._selector_kwargs,
    )

set_params

set_params(**params)

Update constructor kwargs (scikit-learn style); returns self.

Source code in src/jax_ibart/ibart.py

def set_params(self, **params) -> "IBart":
    """Update constructor kwargs (scikit-learn style); returns ``self``."""
    own = set(self.get_params()) - set(self._selector_kwargs)
    for k, v in params.items():
        if k in own:
            setattr(self, k, v)
        else:
            self._selector_kwargs[k] = v
    self.selector = make_selector(self.backend, **self._selector_kwargs)
    return self

clone

clone()

A fresh, unfitted copy with the same configuration.

Source code in src/jax_ibart/ibart.py

def clone(self) -> "IBart":
    """A fresh, unfitted copy with the same configuration."""
    return IBart(**self.get_params())

fit

fit(X, y, feature_names=None, units=None)

Fit iBART.

feature_names defaults to the columns of a pandas DataFrame X, or x0, x1, .... Names become the symbols in the discovered expressions.

units (optional): an (n_features, n_dims) integer unit-exponent matrix. When given, descriptors that are dimensionally inconsistent (e.g. summing two different units) are never generated.

The fit always runs JAX in float32, even if 64-bit mode is enabled globally (some libraries, e.g. discopt, enable it on import), so results do not depend on what else is imported. The final regression is float64.

Source code in src/jax_ibart/ibart.py

def fit(self, X, y, feature_names: list[str] | None = None, units=None):
    """Fit iBART.

    ``feature_names`` defaults to the columns of a pandas DataFrame ``X``,
    or ``x0, x1, ...``. Names become the symbols in the discovered
    expressions.

    ``units`` (optional): an (n_features, n_dims) integer unit-exponent
    matrix. When given, descriptors that are dimensionally inconsistent
    (e.g. summing two different units) are never generated.

    The fit always runs JAX in float32, even if 64-bit mode is enabled
    globally (some libraries, e.g. discopt, enable it on import), so results
    do not depend on what else is imported. The final regression is float64.
    """
    if feature_names is None:
        cols = getattr(X, "columns", None)
        p = np.shape(X)[1] if np.ndim(X) == 2 else 0
        feature_names = (
            [str(c) for c in cols]
            if cols is not None
            else [f"x{i}" for i in range(p)]
        )
    X, y, feature_names, units = _check_inputs(X, y, feature_names, units)
    with _x64_off():
        return self._fit(X, y, feature_names, units)

transform

transform(X)

Evaluate the selected descriptors on raw primary features X.

Returns the (n, k) design matrix the final linear model uses.

Source code in src/jax_ibart/ibart.py

def transform(self, X):
    """Evaluate the selected descriptors on raw primary features ``X``.

    Returns the ``(n, k)`` design matrix the final linear model uses.
    """
    self._check_fitted()
    syms = [sp.Symbol(n) for n in self._feature_names]
    X = np.asarray(X, dtype=np.float64)
    if X.ndim != 2 or X.shape[1] != len(syms):
        raise ValueError(
            f"X must have shape (n, {len(syms)}), got {tuple(X.shape)}"
        )
    cols = []
    for expr in self.descriptors_:
        f = sp.lambdify(syms, expr, modules="numpy")
        cols.append(
            np.asarray(f(*[X[:, i] for i in range(X.shape[1])]), dtype=np.float64)
        )
    return np.column_stack(cols)

predict

predict(X)

Predict y for raw primary features X.

The symbolic descriptors are evaluated as written, so operators such as log and sqrt give NaN outside their real domain (the training pipeline evaluates them on |x|).

Source code in src/jax_ibart/ibart.py

def predict(self, X):
    """Predict ``y`` for raw primary features ``X``.

    The symbolic descriptors are evaluated as written, so operators such as
    ``log`` and ``sqrt`` give NaN outside their real domain (the training
    pipeline evaluates them on ``|x|``).
    """
    with np.errstate(all="ignore"):
        D = self.transform(X)
    pred = self.coef_[0] + D @ self.coef_[1:]
    if not np.isfinite(pred).all():
        warnings.warn(
            f"{int((~np.isfinite(pred)).sum())} predictions are NaN or inf; "
            "inputs are probably outside a descriptor's domain (e.g. log of a "
            "non-positive value)",
            stacklevel=2,
        )
    return pred

expression

expression(digits=4)

The fitted model as a single symbolic expression.

Coefficients are rounded to digits significant figures for display (pass digits=None for full precision).

Source code in src/jax_ibart/ibart.py

def expression(self, digits: int = 4) -> sp.Expr:
    """The fitted model as a single symbolic expression.

    Coefficients are rounded to ``digits`` significant figures for display
    (pass ``digits=None`` for full precision).
    """
    self._check_fitted()
    expr = sp.Float(self.coef_[0])
    for c, d in zip(self.coef_[1:], self.descriptors_):
        expr = expr + sp.Float(c) * d
    return sp.N(expr, digits) if digits else expr

to_discopt

to_discopt(m, x=None, **kwargs)

Add this model to a discopt optimization model.

Returns (x, y): the input variables by feature name and a variable constrained to equal the prediction. Keyword arguments are passed on; see :func:jax_ibart.export.to_discopt.

Source code in src/jax_ibart/ibart.py

def to_discopt(self, m, x=None, **kwargs):
    """Add this model to a discopt optimization model.

    Returns ``(x, y)``: the input variables by feature name and a variable
    constrained to equal the prediction. Keyword arguments are passed on;
    see :func:`jax_ibart.export.to_discopt`.
    """
    from .export import to_discopt

    return to_discopt(self, m, x, **kwargs)

latex

latex(lhs='y', digits=4)

The fitted model as a LaTeX string, e.g. y = -0.099 BVS Hfo + ....

Source code in src/jax_ibart/ibart.py

def latex(self, lhs: str = "y", digits: int = 4) -> str:
    """The fitted model as a LaTeX string, e.g. ``y = -0.099 BVS Hfo + ...``."""
    return f"{lhs} = {sp.latex(self.expression(digits))}"

to_html

to_html(lhs='y', digits=4)

An HTML snippet rendering the model (MathJax/KaTeX delimiters).

Source code in src/jax_ibart/ibart.py

def to_html(self, lhs: str = "y", digits: int = 4) -> str:
    """An HTML snippet rendering the model (MathJax/KaTeX delimiters)."""
    return f"<p>\\({self.latex(lhs, digits)}\\)</p>"

Selection backends

jax_ibart.selection.make_selector

make_selector(backend='bart', **kwargs)

Factory: "bart" -> BART-G.SE (faithful), "diff" -> differentiable.

Source code in src/jax_ibart/selection.py

def make_selector(backend: str = "bart", **kwargs) -> Selector:
    """Factory: ``"bart"`` -> BART-G.SE (faithful), ``"diff"`` -> differentiable."""
    if backend == "bart":
        return BartGSESelector(**kwargs)
    if backend == "diff":
        return DiffSelector(**kwargs)
    raise ValueError(f"unknown selection backend: {backend!r}")

jax_ibart.selection.BartGSESelector

Faithful BART-G.SE screening backend (default).

Parameters:

Name	Type	Description	Default
n_perm	int	Number of label permutations for the null distribution of inclusion proportions.	50
alpha	float	Simultaneous coverage level for the global standard-error (G.SE) threshold. Larger alpha lowers the threshold and keeps more descriptors. This is the key knob: the PAN criterion wants intermediate screening to tolerate false positives but never drop a true building block (no false negatives), so a permissive (larger) alpha is correct for iBART's interleaved screening. An over-conservative alpha over-prunes; on catalysis (the R-iBART oracle cross-check) alpha=0.05 screened out the Hfo/Oxv coupling, while alpha=0.25 selects descriptors built on the published abs(Hfo/Oxv). Default 0.10; use ~0.2-0.25 for the most faithful reproduction. (Clean, well-separated problems are insensitive to alpha; it matters on collinear, high-dimensional screens.)	0.1
n_tree	int	BART ensemble size and MCMC schedule (passed through to bartz). Note: a larger MCMC budget does not narrow the catalysis gap; the screening threshold (alpha) does. See docs/oracle-crosscheck.md.	50
n_draw	int	BART ensemble size and MCMC schedule (passed through to bartz). Note: a larger MCMC budget does not narrow the catalysis gap; the screening threshold (alpha) does. See docs/oracle-crosscheck.md.	50
n_burn	int	BART ensemble size and MCMC schedule (passed through to bartz). Note: a larger MCMC budget does not narrow the catalysis gap; the screening threshold (alpha) does. See docs/oracle-crosscheck.md.	50

Source code in src/jax_ibart/selection.py

class BartGSESelector:
    """Faithful BART-G.SE screening backend (default).

    Parameters
    ----------
    n_perm : int
        Number of label permutations for the null distribution of inclusion
        proportions.
    alpha : float
        Simultaneous coverage level for the global standard-error (G.SE)
        threshold. Larger alpha lowers the threshold and keeps *more*
        descriptors. This is the key knob: the PAN criterion wants intermediate
        screening to tolerate false positives but never drop a true building
        block (no false negatives), so a permissive (larger) alpha is correct
        for iBART's interleaved screening. An over-conservative alpha
        over-prunes; on catalysis (the R-iBART oracle cross-check) alpha=0.05
        screened out the Hfo/Oxv coupling, while alpha=0.25 selects descriptors
        built on the published abs(Hfo/Oxv). Default 0.10; use ~0.2-0.25 for the most
        faithful reproduction. (Clean, well-separated problems are insensitive
        to alpha; it matters on collinear, high-dimensional screens.)
    n_tree, n_draw, n_burn : int
        BART ensemble size and MCMC schedule (passed through to ``bartz``).
        Note: a *larger* MCMC budget does not narrow the catalysis gap; the
        screening threshold (alpha) does. See docs/oracle-crosscheck.md.
    """

    def __init__(
        self,
        n_perm: int = 50,
        alpha: float = 0.10,
        n_tree: int = 50,
        n_draw: int = 1000,
        n_burn: int = 1000,
    ):
        self.n_perm = n_perm
        self.alpha = alpha
        self.n_tree = n_tree
        self.n_draw = n_draw
        self.n_burn = n_burn

    def select(self, X, y, *, key):
        """BART-G.SE screening (Bleich et al. 2014, Global SE rule).

        Returns a boolean keep-mask over the columns of ``X`` (n, d). A column
        is kept if its observed inclusion proportion exceeds a permutation-null
        threshold ``mu_j + g * sd_j`` with a single global multiplier ``g``
        calibrated for simultaneous (1-alpha) coverage.
        """
        X32 = np.asarray(X, dtype=np.float32)
        y32 = np.asarray(y, dtype=np.float32)
        n, d = X32.shape
        kw = dict(n_tree=self.n_tree, n_draw=self.n_draw, n_burn=self.n_burn)

        # Derive reproducible int seeds from the functional key.
        base = int(jax.random.randint(key, (), 0, 2**31 - 1))
        rng = np.random.default_rng(base)

        # 1. Observed inclusion proportions.
        q_obs = _bart_varprop(X32, y32, base, **kw)
        assert q_obs.shape == (d,), (q_obs.shape, d)

        # 2. Permutation null: refit on permuted y.
        null = np.empty((self.n_perm, d), dtype=np.float64)
        for b in range(self.n_perm):
            yb = y32[rng.permutation(n)]
            null[b] = _bart_varprop(X32, yb, base + 1 + b, **kw)

        # 3. Per-variable null mean/sd, then the global SE multiplier: the
        #    (1-alpha) quantile of the per-permutation max standardized score.
        mu = null.mean(axis=0)
        sd = null.std(axis=0) + _EPS
        z_null = (null - mu) / sd  # (n_perm, d)
        g = float(np.quantile(z_null.max(axis=1), 1.0 - self.alpha))

        # 4. Select where the observed score clears the global threshold.
        keep = (q_obs - mu) / sd > g
        # Diagnostics from the last call, for plotting: a column is kept where
        # q_obs > threshold, i.e. its usage beats the permutation null.
        self.diagnostics_ = {"q_obs": q_obs, "null": null, "threshold": mu + g * sd}
        return jax.numpy.asarray(keep)

select

select(X, y, *, key)

BART-G.SE screening (Bleich et al. 2014, Global SE rule).

Returns a boolean keep-mask over the columns of X (n, d). A column is kept if its observed inclusion proportion exceeds a permutation-null threshold mu_j + g * sd_j with a single global multiplier g calibrated for simultaneous (1-alpha) coverage.

Source code in src/jax_ibart/selection.py

def select(self, X, y, *, key):
    """BART-G.SE screening (Bleich et al. 2014, Global SE rule).

    Returns a boolean keep-mask over the columns of ``X`` (n, d). A column
    is kept if its observed inclusion proportion exceeds a permutation-null
    threshold ``mu_j + g * sd_j`` with a single global multiplier ``g``
    calibrated for simultaneous (1-alpha) coverage.
    """
    X32 = np.asarray(X, dtype=np.float32)
    y32 = np.asarray(y, dtype=np.float32)
    n, d = X32.shape
    kw = dict(n_tree=self.n_tree, n_draw=self.n_draw, n_burn=self.n_burn)

    # Derive reproducible int seeds from the functional key.
    base = int(jax.random.randint(key, (), 0, 2**31 - 1))
    rng = np.random.default_rng(base)

    # 1. Observed inclusion proportions.
    q_obs = _bart_varprop(X32, y32, base, **kw)
    assert q_obs.shape == (d,), (q_obs.shape, d)

    # 2. Permutation null: refit on permuted y.
    null = np.empty((self.n_perm, d), dtype=np.float64)
    for b in range(self.n_perm):
        yb = y32[rng.permutation(n)]
        null[b] = _bart_varprop(X32, yb, base + 1 + b, **kw)

    # 3. Per-variable null mean/sd, then the global SE multiplier: the
    #    (1-alpha) quantile of the per-permutation max standardized score.
    mu = null.mean(axis=0)
    sd = null.std(axis=0) + _EPS
    z_null = (null - mu) / sd  # (n_perm, d)
    g = float(np.quantile(z_null.max(axis=1), 1.0 - self.alpha))

    # 4. Select where the observed score clears the global threshold.
    keep = (q_obs - mu) / sd > g
    # Diagnostics from the last call, for plotting: a column is kept where
    # q_obs > threshold, i.e. its usage beats the permutation null.
    self.diagnostics_ = {"q_obs": q_obs, "null": null, "threshold": mu + g * sd}
    return jax.numpy.asarray(keep)

jax_ibart.selection.DiffSelector

Differentiable end-to-end screening backend.

Replaces the discrete BART importance with a sparse-gated neural importance model (hard-concrete L0 gates, Louizos et al. 2018; see :mod:jax_ibart._diffscreen) so the whole selection step admits gradients. Acceptance gate: keep the true drivers on the synthetic fixtures and agree with :class:BartGSESelector (PAN: no false negatives).

Parameters:

Name	Type	Description	Default
steps	training hyperparameters for the gated MLP.		1500
lr	training hyperparameters for the gated MLP.		1500
l0	training hyperparameters for the gated MLP.		1500
hidden	training hyperparameters for the gated MLP.		1500
threshold	float	Keep features whose learned gate-open probability exceeds this. Kept low to over-select (PAN-safe: tolerate false positives, avoid false negatives).	0.5
min_keep	int	Always keep at least this many top-importance features (never empties the pool).	2

Source code in src/jax_ibart/selection.py

class DiffSelector:
    """Differentiable end-to-end screening backend.

    Replaces the discrete BART importance with a sparse-gated neural importance
    model (hard-concrete L0 gates, Louizos et al. 2018; see
    :mod:`jax_ibart._diffscreen`) so the whole selection step admits gradients.
    Acceptance gate: keep the true drivers on the synthetic fixtures and agree
    with :class:`BartGSESelector` (PAN: no false negatives).

    Parameters
    ----------
    steps, lr, l0, hidden : training hyperparameters for the gated MLP.
    threshold : float
        Keep features whose learned gate-open probability exceeds this. Kept low
        to over-select (PAN-safe: tolerate false positives, avoid false
        negatives).
    min_keep : int
        Always keep at least this many top-importance features (never empties
        the pool).
    """

    def __init__(
        self,
        steps: int = 1500,
        lr: float = 1e-2,
        l0: float = 0.03,
        hidden: int = 16,
        threshold: float = 0.5,
        min_keep: int = 2,
    ):
        self.steps = steps
        self.lr = lr
        self.l0 = l0
        self.hidden = hidden
        self.threshold = threshold
        self.min_keep = min_keep

    def select(self, X, y, *, key):
        try:
            import optax  # noqa: F401
        except ImportError as e:
            raise ImportError(
                "the diff backend needs optax: pip install 'jax-ibart[diff]'"
            ) from e
        from ._diffscreen import screen_importance

        imp = np.asarray(
            screen_importance(
                X, y, key, steps=self.steps, lr=self.lr, l0=self.l0, hidden=self.hidden
            )
        )
        self.importance_ = imp  # last call's gate-open probabilities
        keep = imp > self.threshold
        # PAN safeguard: never drop below min_keep (keep top-importance features).
        if keep.sum() < min(self.min_keep, imp.size):
            top = np.argsort(imp)[::-1][: min(self.min_keep, imp.size)]
            keep = np.zeros_like(keep)
            keep[top] = True
        return jax.numpy.asarray(keep)

Feature engineering

jax_ibart.engineer.Pool dataclass

A descriptor pool: numeric columns paired with symbolic provenance.

units is optional. When present it has shape (d, n_dims): row j is the integer unit-exponent vector of column j. When None, dimensional analysis is disabled and no unit filtering is done.

Source code in src/jax_ibart/engineer.py

@dataclass
class Pool:
    """A descriptor pool: numeric columns paired with symbolic provenance.

    ``units`` is optional. When present it has shape (d, n_dims): row j is the
    integer unit-exponent vector of column j. When None, dimensional analysis is
    disabled and no unit filtering is done.
    """

    X: jnp.ndarray  # (n, d)
    exprs: list[sp.Expr]  # length d
    units: np.ndarray | None = None  # (d, n_dims) or None

    @property
    def n_descriptors(self) -> int:
        return len(self.exprs)

jax_ibart.engineer.from_primary

from_primary(X, names, units=None)

Build the initial pool from primary features and their symbols.

Numeric columns are float32: this is the dtype bartz consumes, and global jax_enable_x64 must stay off (it breaks bartz's MCMC). Higher-precision work (final LASSO/OLS) happens host-side in numpy, not in JAX traces.

units (optional) is (d, n_dims): the unit-exponent vector per feature.

Source code in src/jax_ibart/engineer.py

def from_primary(X, names: list[str], units=None) -> Pool:
    """Build the initial pool from primary features and their symbols.

    Numeric columns are float32: this is the dtype bartz consumes, and global
    ``jax_enable_x64`` must stay off (it breaks bartz's MCMC). Higher-precision
    work (final LASSO/OLS) happens host-side in numpy, not in JAX traces.

    ``units`` (optional) is (d, n_dims): the unit-exponent vector per feature.
    """
    X = jnp.asarray(X, dtype=jnp.float32)
    u = None if units is None else np.asarray(units, dtype=np.int64)
    return Pool(X=X, exprs=[sp.Symbol(n) for n in names], units=u)

jax_ibart.engineer.apply_unary

apply_unary(pool, ops, corr_max_d=10000)

Apply unary operators to every column.

Vectorized: each operator is evaluated on all valid columns in one batched JAX call (op.fn broadcasts over the (n, k) slice), and unit vectors are transformed with a single array op. The symbolic provenance is built in a host-side loop (SymPy is not vectorizable). corr_max_d is forwarded to :func:clean (set 0 to keep collinear columns, e.g. for implicit discovery).

Source code in src/jax_ibart/engineer.py

def apply_unary(pool: Pool, ops: list[Operator], corr_max_d: int = 10000) -> Pool:
    """Apply unary operators to every column.

    Vectorized: each operator is evaluated on all valid columns in one batched
    JAX call (``op.fn`` broadcasts over the (n, k) slice), and unit vectors are
    transformed with a single array op. The symbolic provenance is built in a
    host-side loop (SymPy is not vectorizable). ``corr_max_d`` is forwarded to
    :func:`clean` (set 0 to keep collinear columns, e.g. for implicit discovery).
    """
    has_units = pool.units is not None
    blocks, exprs, units = [], [], ([] if has_units else None)
    for op in ops:
        idx = np.where(np.asarray(op.valid(pool.X)))[0]
        if idx.size == 0:
            continue
        blocks.append(op.fn(pool.X[:, idx]))  # (n, k) batched
        exprs.extend(op.sym(pool.exprs[j]) for j in idx)
        if has_units:
            units.append(op.unit(pool.units[idx]))  # vectorized unit rule
    return _concat(blocks, exprs, units, corr_max_d)

jax_ibart.engineer.apply_binary

apply_binary(pool, ops, corr_max_d=10000)

Apply binary operators to all column pairs.

Vectorized over pairs: index arrays (I, J) select the operands, op.fn evaluates all pairs at once via broadcasting (the O(d^2) blowup that is the prime JAX/GPU target), and the unit filter is a vectorized numpy mask.

Source code in src/jax_ibart/engineer.py

def apply_binary(pool: Pool, ops: list[Operator], corr_max_d: int = 10000) -> Pool:
    """Apply binary operators to all column pairs.

    Vectorized over pairs: index arrays (I, J) select the operands, ``op.fn``
    evaluates all pairs at once via broadcasting (the O(d^2) blowup that is the
    prime JAX/GPU target), and the unit filter is a vectorized numpy mask.
    """
    d = pool.n_descriptors
    has_units = pool.units is not None
    U = pool.units
    blocks, exprs, units = [], [], ([] if has_units else None)

    for op in ops:
        if op.commutative:
            I, J = np.triu_indices(d, k=1)
        else:
            ii, jj = np.meshgrid(np.arange(d), np.arange(d), indexing="ij")
            off = ii != jj
            I, J = ii[off], jj[off]
        if I.size == 0:
            continue

        new_u = None
        if has_units:
            if op.unit is u_same:  # additive: keep if equal
                ok = np.all(U[I] == U[J], axis=1)
                I, J = I[ok], J[ok]
                new_u = U[I]
            elif op.unit is u_add:  # multiply: add exponents
                new_u = U[I] + U[J]
            elif op.unit is u_sub:  # divide: subtract
                new_u = U[I] - U[J]
            else:  # generic fallback
                pairs = [(op.unit(U[a], U[b]), a, b) for a, b in zip(I, J)]
                keep = [(u, a, b) for (u, a, b) in pairs if u is not None]
                if not keep:
                    continue
                new_u = np.array([u for u, _, _ in keep])
                I = np.array([a for _, a, _ in keep])
                J = np.array([b for _, _, b in keep])
            if I.size == 0:
                continue

        ok = np.asarray(op.valid(pool.X[:, I], pool.X[:, J]), dtype=bool)
        if ok.ndim == 0:
            if not ok:
                continue
        else:
            I, J = I[ok], J[ok]
            if has_units:
                new_u = new_u[ok]
            if I.size == 0:
                continue

        blocks.append(op.fn(pool.X[:, I], pool.X[:, J]))  # (n, npairs) batched
        exprs.extend(op.sym(pool.exprs[a], pool.exprs[b]) for a, b in zip(I, J))
        if has_units:
            units.append(new_u)
    return _concat(blocks, exprs, units, corr_max_d)

jax_ibart.engineer.clean

clean(pool, corr_thresh=1 - 1e-08, corr_max_d=10000)

Drop invalid (NaN/inf), near-constant, and duplicate/collinear columns.

All numeric decisions are made in float64 (host-side numpy), even though the columns are stored f32 for the bartz pipeline: collinearity/rank judgements on small-n, highly correlated engineered descriptors are exactly where f32's ~7 digits are not enough. The returned pool keeps the original f32 columns; only the keep/drop choices use f64.

Source code in src/jax_ibart/engineer.py

def clean(pool: Pool, corr_thresh: float = 1 - 1e-8, corr_max_d: int = 10000) -> Pool:
    """Drop invalid (NaN/inf), near-constant, and duplicate/collinear columns.

    All numeric *decisions* are made in float64 (host-side numpy), even though
    the columns are stored f32 for the bartz pipeline: collinearity/rank
    judgements on small-n, highly correlated engineered descriptors are exactly
    where f32's ~7 digits are not enough. The returned pool keeps the original f32 columns; only the
    keep/drop choices use f64.
    """
    X64 = np.asarray(pool.X, dtype=np.float64)
    n, d = X64.shape
    keep = np.ones(d, dtype=bool)

    # 1. Finite and non-constant (f64).
    finite = np.all(np.isfinite(X64), axis=0)
    std = np.where(finite, X64.std(axis=0), 0.0)
    keep &= finite & (std > 1e-12)

    # 2. Symbolic-duplicate pruning (structural, keep first occurrence).
    #    SymPy auto-canonicalizes commutative args on construction, so srepr
    #    catches reorderings (a+b == b+a) directly. Deeper algebraic identities
    #    (e.g. sqrt(x**2) == abs(x)) are left to the numeric collinearity pass
    #    below -- using sp.simplify here instead costs ~100x more (it dominated
    #    feature-engineering time) for no gain over the numeric dedup.
    seen: set = set()
    for j, e in enumerate(pool.exprs):
        if not keep[j]:
            continue
        key = sp.srepr(e)
        if key in seen:
            keep[j] = False
        else:
            seen.add(key)

    # 3. Numeric collinearity dedup (f64): greedily drop columns that are
    #    near-perfectly correlated with an already-kept column. Catches
    #    numeric duplicates symbolic canonicalization misses (e.g. sqrt(x**2)
    #    vs abs(x)) and ill-conditioning before it reaches the linear solve.
    #    This greedy pass is O(d^2 n); skip it above corr_max_d (pre-selection,
    #    exact duplicates there are harmless and BART screening follows anyway).
    surv = np.where(keep)[0]
    if 1 < surv.size <= corr_max_d:
        Z = X64[:, surv]
        Z = (Z - Z.mean(0)) / (Z.std(0) + 1e-300)
        kept_cols: list[int] = []
        for col, j in enumerate(surv):
            if kept_cols:
                corr = np.abs(Z[:, kept_cols].T @ Z[:, col] / n)
                if corr.max() > corr_thresh:
                    keep[j] = False
                    continue
            kept_cols.append(col)

    idx = np.where(keep)[0]
    u = None if pool.units is None else pool.units[idx]
    return Pool(X=pool.X[:, idx], exprs=[pool.exprs[j] for j in idx], units=u)

Operators

jax_ibart.operators.Operator dataclass

Source code in src/jax_ibart/operators.py

@dataclass(frozen=True)
class Operator:
    name: str
    arity: int  # 1 (unary) or 2 (binary)
    fn: Callable  # numeric: (col,) or (a, b) -> col   [jnp arrays]
    sym: Callable  # symbolic: (expr,) or (e1, e2) -> sympy expr
    valid: Callable  # domain predicate -> bool mask (jnp)
    unit: Callable  # unit rule: (u,) or (u_a, u_b) -> vector | None (None=reject)
    commutative: bool = False  # binary ops: skip duplicate (j, i) pairs

jax_ibart.operators.default_unary

default_unary()

The O1 set used in the catalysis study (Liu et al. 2022).

Source code in src/jax_ibart/operators.py

def default_unary() -> list[Operator]:
    """The O1 set used in the catalysis study (Liu et al. 2022)."""
    return [
        UNARY[k] for k in ("id", "neg", "sq", "sqrt", "log", "exp", "abs", "sin", "cos")
    ]

jax_ibart.operators.default_binary

default_binary()

The O2 set used in the catalysis study (Liu et al. 2022).

Source code in src/jax_ibart/operators.py

def default_binary() -> list[Operator]:
    """The O2 set used in the catalysis study (Liu et al. 2022)."""
    return [BINARY[k] for k in ("add", "sub", "mul", "div", "absdiff")]

Final parametric stage

jax_ibart.sparsify.ols

ols(X, y)

Ordinary least squares with an intercept, in float64.

Returns (coef, rss) where coef is [intercept, *slopes] and rss is the residual sum of squares.

Source code in src/jax_ibart/sparsify.py

def ols(X, y):
    """Ordinary least squares with an intercept, in float64.

    Returns ``(coef, rss)`` where ``coef`` is ``[intercept, *slopes]`` and
    ``rss`` is the residual sum of squares.
    """
    X, y = _as_f64(X, y)
    A = np.column_stack([np.ones(len(y)), X])
    beta, *_ = np.linalg.lstsq(A, y, rcond=None)
    resid = y - A @ beta
    return beta, float(resid @ resid)

jax_ibart.sparsify.lasso_screen

lasso_screen(X, y, n_keep=10)

Return indices of up to n_keep descriptors ranked by LASSO entry.

Features are standardized for ranking only. Variables are ordered by absolute standardized coefficient at the cross-validated LASSO solution; any shortfall is filled by univariate correlation so exactly min(n_keep, d) indices come back.

Source code in src/jax_ibart/sparsify.py

def lasso_screen(X, y, n_keep: int = 10) -> np.ndarray:
    """Return indices of up to ``n_keep`` descriptors ranked by LASSO entry.

    Features are standardized for ranking only. Variables are ordered by
    absolute standardized coefficient at the cross-validated LASSO solution;
    any shortfall is filled by univariate correlation so exactly
    ``min(n_keep, d)`` indices come back.
    """
    X, y = _as_f64(X, y)
    n, d = X.shape
    if d <= n_keep:
        return np.arange(d)

    mu, sd = X.mean(0), X.std(0) + 1e-300
    Xs = (X - mu) / sd
    yc = y - y.mean()

    if n < 2 * _MIN_FOLD:  # too few samples to cross-validate the LASSO path
        chosen = []
    else:
        cv = min(5, n // _MIN_FOLD)
        coef = LassoLarsCV(cv=cv, max_iter=2000).fit(Xs, yc).coef_
        ranked = list(np.argsort(np.abs(coef))[::-1])
        chosen = [j for j in ranked if coef[j] != 0.0][:n_keep]

    if len(chosen) < n_keep:  # fill from univariate correlation
        corr = np.abs(Xs.T @ yc) / n
        for j in np.argsort(corr)[::-1]:
            if j not in chosen:
                chosen.append(int(j))
            if len(chosen) == n_keep:
                break
    return np.asarray(chosen[:n_keep], dtype=int)

jax_ibart.sparsify.l0_best_subset

l0_best_subset(X, y, max_dim=3, score='bic')

Exhaustive l0: best OLS subset up to max_dim columns.

Returns (indices, coef) where indices are column positions into X and coef is the OLS coefficient vector [intercept, *slopes] for that subset (f64).

Source code in src/jax_ibart/sparsify.py

def l0_best_subset(X, y, max_dim: int = 3, score: str = "bic"):
    """Exhaustive l0: best OLS subset up to ``max_dim`` columns.

    Returns ``(indices, coef)`` where ``indices`` are column positions into
    ``X`` and ``coef`` is the OLS coefficient vector ``[intercept, *slopes]``
    for that subset (f64).
    """
    if score not in ("bic", "cv"):
        raise ValueError(f"score must be 'bic' or 'cv', got {score!r}")
    X, y = _as_f64(X, y)
    n, d = X.shape
    if d == 0:
        raise ValueError("no candidate descriptors to choose from")
    max_dim = min(max_dim, d)

    def score_subset(cols):
        Xs = X[:, list(cols)]
        if score == "cv":
            return _cv_mse(Xs, y)
        _, rss = ols(Xs, y)
        return _bic(rss, n, k=len(cols) + 1)

    best_cols, best_score = None, np.inf
    for k in range(1, max_dim + 1):
        for cols in combinations(range(d), k):
            s = score_subset(cols)
            if s < best_score:
                best_score, best_cols = s, cols

    coef, _ = ols(X[:, list(best_cols)], y)
    return np.asarray(best_cols, dtype=int), coef

Uncertainty quantification

jax_ibart.uncertainty.conformal_intervals

conformal_intervals(factory, X, y, feature_names, units=None, alpha=0.1, cal_frac=0.3, seed=0)

Split-conformal prediction intervals around an iBART fit.

The estimator (selection + final model) is fit on a proper-training split and scored on a disjoint calibration split, so the residual quantile is an honest, post-selection-valid calibration. Returns a :class:ConformalModel.

Source code in src/jax_ibart/uncertainty.py

def conformal_intervals(
    factory: Callable,
    X,
    y,
    feature_names,
    units=None,
    alpha: float = 0.10,
    cal_frac: float = 0.3,
    seed: int = 0,
):
    """Split-conformal prediction intervals around an iBART fit.

    The estimator (selection + final model) is fit on a proper-training split
    and scored on a disjoint calibration split, so the residual quantile is an
    honest, post-selection-valid calibration. Returns a :class:`ConformalModel`.
    """
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    n = len(y)
    rng = np.random.default_rng(seed)
    idx = rng.permutation(n)
    n_cal = max(int(round(cal_frac * n)), 1)
    cal, tr = idx[:n_cal], idx[n_cal:]

    model = factory().fit(X[tr], y[tr], feature_names, units=units)
    resid = np.abs(y[cal] - np.asarray(model.predict(X[cal])))

    # Finite-sample conformal quantile. If ceil((n_cal + 1)(1 - alpha)) > n_cal
    # the calibration set is too small for the guarantee and the interval is
    # unbounded.
    rank = int(np.ceil((n_cal + 1) * (1 - alpha)))
    if rank > n_cal:
        warnings.warn(
            f"{n_cal} calibration points are too few for alpha={alpha}; "
            "the conformal interval is infinite",
            stacklevel=2,
        )
        qhat = float("inf")
    else:
        qhat = float(np.quantile(resid, rank / n_cal, method="higher"))
    return ConformalModel(model=model, half_width=qhat, alpha=alpha)

jax_ibart.uncertainty.empirical_coverage

empirical_coverage(conformal, X_test, y_test)

Fraction of test points whose true y falls inside the conformal interval.

Source code in src/jax_ibart/uncertainty.py

def empirical_coverage(conformal: ConformalModel, X_test, y_test):
    """Fraction of test points whose true y falls inside the conformal interval."""
    lo, _, hi = conformal.interval(X_test)
    y_test = np.asarray(y_test, dtype=np.float64)
    return float(np.mean((y_test >= lo) & (y_test <= hi)))

jax_ibart.uncertainty.ols_intervals

ols_intervals(model, X_train, y_train, X_new=None, alpha=0.1)

OLS coefficient SEs/CIs and Gaussian prediction intervals.

OPTIMISTIC: conditions on model.descriptors_ and ignores that they were selected from data, so coverage is below nominal. Use :func:conformal_intervals for honest intervals.

Source code in src/jax_ibart/uncertainty.py

def ols_intervals(model, X_train, y_train, X_new=None, alpha: float = 0.10):
    """OLS coefficient SEs/CIs and Gaussian prediction intervals.

    OPTIMISTIC: conditions on ``model.descriptors_`` and ignores that they were
    selected from data, so coverage is below nominal. Use
    :func:`conformal_intervals` for honest intervals.
    """
    y = np.asarray(y_train, dtype=np.float64)
    D = model.transform(X_train)  # (n, k) selected descriptors
    A = np.column_stack([np.ones(len(y)), D])  # design with intercept
    n, p = A.shape
    beta = model.coef_
    resid = y - A @ beta
    dof = max(n - p, 1)
    sigma2 = float(resid @ resid) / dof
    XtX_inv = np.linalg.inv(A.T @ A)
    se = np.sqrt(np.diag(sigma2 * XtX_inv))
    tcrit = stats.t.ppf(1 - alpha / 2, dof)

    names = ["intercept"] + [str(e) for e in model.descriptors_]
    coefs = [
        {
            "term": nm,
            "coef": float(b),
            "se": float(s),
            "ci_low": float(b - tcrit * s),
            "ci_high": float(b + tcrit * s),
        }
        for nm, b, s in zip(names, beta, se)
    ]
    out = {"coefficients": coefs, "sigma": float(np.sqrt(sigma2)), "dof": dof}

    if X_new is not None:
        Dn = model.transform(X_new)
        An = np.column_stack([np.ones(len(Dn)), Dn])
        pred = An @ beta
        # prediction interval: includes both estimation and noise variance
        lev = np.einsum("ij,jk,ik->i", An, XtX_inv, An)
        half = tcrit * np.sqrt(sigma2 * (1.0 + lev))
        out["prediction"] = {"point": pred, "lower": pred - half, "upper": pred + half}
    return out

jax_ibart.uncertainty.selection_stability

selection_stability(factory, X, y, feature_names, units=None, n_boot=25, seed=0)

Bootstrap inclusion frequencies for descriptors and primary features.

Refits the whole pipeline on n_boot bootstrap resamples and records how often each descriptor (by canonical form) and each primary feature is selected. Frequencies near 1 are stable signals; low frequencies are selection noise.

Source code in src/jax_ibart/uncertainty.py

def selection_stability(
    factory: Callable, X, y, feature_names, units=None, n_boot: int = 25, seed: int = 0
):
    """Bootstrap inclusion frequencies for descriptors and primary features.

    Refits the whole pipeline on ``n_boot`` bootstrap resamples and records how
    often each descriptor (by canonical form) and each primary feature is
    selected. Frequencies near 1 are stable signals; low frequencies are
    selection noise.
    """
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    n = len(y)
    rng = np.random.default_rng(seed)

    desc_counts: dict[str, int] = {}
    feat_counts: dict[str, int] = {}
    for _ in range(n_boot):
        bs = rng.integers(0, n, n)
        m = factory().fit(X[bs], y[bs], feature_names, units=units)
        seen_feats = set()
        for e in m.descriptors_:
            key = str(e)
            desc_counts[key] = desc_counts.get(key, 0) + 1
            seen_feats |= {str(s) for s in e.free_symbols}
        for f in seen_feats:
            feat_counts[f] = feat_counts.get(f, 0) + 1

    def freq(d):
        return {k: v / n_boot for k, v in sorted(d.items(), key=lambda kv: -kv[1])}

    return {
        "n_boot": n_boot,
        "descriptor_freq": freq(desc_counts),
        "feature_freq": freq(feat_counts),
    }

jax_ibart.uncertainty.repeated_split_rmse

repeated_split_rmse(factory, X, y, feature_names, units=None, n_splits=20, test_frac=0.25, seed=0)

Out-of-sample RMSE over n_splits random train/test partitions.

Source code in src/jax_ibart/uncertainty.py

def repeated_split_rmse(
    factory: Callable,
    X,
    y,
    feature_names,
    units=None,
    n_splits: int = 20,
    test_frac: float = 0.25,
    seed: int = 0,
):
    """Out-of-sample RMSE over ``n_splits`` random train/test partitions."""
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    n = len(y)
    rng = np.random.default_rng(seed)
    n_test = max(int(round(test_frac * n)), 1)

    rmses = []
    for _ in range(n_splits):
        idx = rng.permutation(n)
        te, tr = idx[:n_test], idx[n_test:]
        m = factory().fit(X[tr], y[tr], feature_names, units=units)
        err = y[te] - np.asarray(m.predict(X[te]))
        rmses.append(float(np.sqrt(np.mean(err**2))))
    rmses = np.array(rmses)
    return {
        "rmse": rmses,
        "mean": float(rmses.mean()),
        "std": float(rmses.std()),
        "q05": float(np.quantile(rmses, 0.05)),
        "q95": float(np.quantile(rmses, 0.95)),
    }

jax_ibart.uncertainty.bart_posterior_intervals

bart_posterior_intervals(model, X_train, y_train, X_new, alpha=0.1, n_tree=50, n_draw=500, n_burn=500, seed=0)

Posterior predictive intervals from a BART fit on the selected descriptors.

Propagates BART's function-value posterior and the error-scale posterior (sigma) to a predictive interval at X_new. Conditional on the selected descriptors (it does not propagate selection uncertainty -- use :func:selection_stability for that). Requires bartz.

Source code in src/jax_ibart/uncertainty.py

def bart_posterior_intervals(
    model,
    X_train,
    y_train,
    X_new,
    alpha: float = 0.10,
    n_tree: int = 50,
    n_draw: int = 500,
    n_burn: int = 500,
    seed: int = 0,
):
    """Posterior predictive intervals from a BART fit on the selected descriptors.

    Propagates BART's function-value posterior and the error-scale posterior
    (sigma) to a predictive interval at ``X_new``. Conditional on the selected
    descriptors (it does not propagate selection uncertainty -- use
    :func:`selection_stability` for that). Requires ``bartz``.
    """
    from .selection import bartz_layout, fit_gbart

    D_tr = np.asarray(model.transform(X_train), dtype=np.float32)
    D_new = np.asarray(model.transform(X_new), dtype=np.float32)
    y = np.asarray(y_train, dtype=np.float32)

    m = fit_gbart(
        D_tr,
        y,
        ntree=n_tree,
        ndpost=n_draw,
        nskip=n_burn,
        seed=int(seed),
        printevery=None,
    )
    fhat = np.asarray(
        m.predict(bartz_layout(D_new)), dtype=np.float64
    )  # (ndpost, n_new)
    sigma = np.asarray(m.sigma, dtype=np.float64)[-n_draw:]  # post-burn draws

    rng = np.random.default_rng(seed)
    noise = rng.standard_normal(fhat.shape) * sigma[:, None]
    ppd = fhat + noise  # posterior predictive
    lo = np.quantile(ppd, alpha / 2, axis=0)
    med = np.quantile(ppd, 0.5, axis=0)
    hi = np.quantile(ppd, 1 - alpha / 2, axis=0)
    return {"lower": lo, "median": med, "upper": hi}

Optimization (discopt export)

jax_ibart.export.to_discopt

to_discopt(model, m, x=None, *, bounds='data', domain=True, X_ref=None, name='y', y_bounds='auto')

Add a fitted iBART model to a discopt model.

Parameters:

Name	Type	Description	Default
model	IBart	A fitted model.	required
m	Model	The optimization model to add variables and constraints to.	required
x	dict or sequence	The inputs. A dict maps feature names to discopt expressions; a sequence (list, array, or a vector variable) is indexed in feature_names order. If None, one continuous variable per feature the model uses is created, named after the feature.	None
bounds	"data", tuple of arrays, or None	Bounds for created variables: "data" uses the range of X_ref, (lb, ub) gives arrays over all features. Ignored when x is given.	'data'
domain	bool	Add the domain constraints described in the module docstring.	True
X_ref	array	Reference data for bounds and domain constraints; defaults to the training X.	None
name	str	Name of the prediction variable (must be unique in m).	'y'
y_bounds	"auto", (lb, ub), or None	Bounds for the prediction variable. "auto" derives them by interval arithmetic over the input box when the inputs are created here (they are conservative, so they never exclude a reachable prediction, and they help the solver); with user-supplied x the box is not known and y is left unbounded, so pass (lb, ub) to bound it yourself. None leaves y unbounded.	'auto'

Returns:

Name	Type	Description
x	dict[str, Expression]	The inputs the model uses, by feature name.
y	Variable	The model prediction, constrained to equal coef_[0] + sum(coef_[k] * descriptor_k).

Examples:

>>> import discopt.modeling as dm
>>> m = dm.Model("design")
>>> x, y = to_discopt(fitted, m)
>>> m.maximize(y)
>>> res = m.solve()
>>> {k: float(res.value(v)) for k, v in x.items()}

Source code in src/jax_ibart/export.py

def to_discopt(
    model,
    m,
    x=None,
    *,
    bounds="data",
    domain=True,
    X_ref=None,
    name="y",
    y_bounds="auto",
):
    """Add a fitted iBART model to a discopt model.

    Parameters
    ----------
    model : IBart
        A fitted model.
    m : discopt.modeling.Model
        The optimization model to add variables and constraints to.
    x : dict or sequence, optional
        The inputs. A dict maps feature names to discopt expressions; a
        sequence (list, array, or a vector variable) is indexed in
        ``feature_names`` order. If None, one continuous variable per feature
        the model uses is created, named after the feature.
    bounds : "data", tuple of arrays, or None
        Bounds for created variables: ``"data"`` uses the range of ``X_ref``,
        ``(lb, ub)`` gives arrays over all features. Ignored when ``x`` is
        given.
    domain : bool
        Add the domain constraints described in the module docstring.
    X_ref : array, optional
        Reference data for bounds and domain constraints; defaults to the
        training ``X``.
    name : str
        Name of the prediction variable (must be unique in ``m``).
    y_bounds : "auto", (lb, ub), or None
        Bounds for the prediction variable. ``"auto"`` derives them by interval
        arithmetic over the input box when the inputs are created here (they
        are conservative, so they never exclude a reachable prediction, and
        they help the solver); with user-supplied ``x`` the box is not known
        and ``y`` is left unbounded, so pass ``(lb, ub)`` to bound it yourself.
        ``None`` leaves ``y`` unbounded.

    Returns
    -------
    x : dict[str, Expression]
        The inputs the model uses, by feature name.
    y : Variable
        The model prediction, constrained to equal
        ``coef_[0] + sum(coef_[k] * descriptor_k)``.

    Examples
    --------
    >>> import discopt.modeling as dm
    >>> m = dm.Model("design")
    >>> x, y = to_discopt(fitted, m)
    >>> m.maximize(y)
    >>> res = m.solve()
    >>> {k: float(res.value(v)) for k, v in x.items()}
    """
    try:
        import discopt.modeling as dm
    except ImportError as e:  # pragma: no cover - exercised without discopt
        raise ImportError(
            "to_discopt needs discopt (Python >= 3.12): "
            "pip install 'jax-ibart[discopt]'"
        ) from e

    model._check_fitted()
    names = list(model._feature_names)
    X_ref = model._X_train if X_ref is None else np.asarray(X_ref, dtype=np.float64)
    if X_ref.ndim != 2 or X_ref.shape[1] != len(names):
        raise ValueError(f"X_ref must have shape (n, {len(names)})")

    used = sorted(
        {str(s) for d in model.descriptors_ for s in d.free_symbols},
        key=names.index,
    )
    if x is None:
        lb, ub = _bounds(bounds, X_ref, len(names))
        inputs = {
            n: m.continuous(n, lb=lb[names.index(n)], ub=ub[names.index(n)])
            for n in used
        }
    elif isinstance(x, dict):
        missing = [n for n in used if n not in x]
        if missing:
            raise ValueError(f"x is missing inputs the model uses: {missing}")
        inputs = {n: x[n] for n in used}
    else:
        inputs = {n: x[names.index(n)] for n in used}

    walker = _Walker(dm, m, inputs, names, X_ref, domain)
    expr = float(model.coef_[0])
    for c, d in zip(model.coef_[1:], model.descriptors_):
        expr = expr + float(c) * walker(d)
    if isinstance(y_bounds, str) and y_bounds != "auto":
        raise ValueError(f"y_bounds must be 'auto', (lb, ub) or None, not {y_bounds!r}")
    y_lb, y_ub = -np.inf, np.inf
    if y_bounds is not None and not isinstance(y_bounds, str):
        y_lb, y_ub = (float(v) for v in y_bounds)
    elif y_bounds == "auto" and x is None:
        box = {n: (float(lb[names.index(n)]), float(ub[names.index(n)])) for n in used}
        y_lb, y_ub = _interval(model, box)
    kw = {}
    if np.isfinite(y_lb):
        kw["lb"] = y_lb
    if np.isfinite(y_ub):
        kw["ub"] = y_ub
    y = m.continuous(name, **kw)
    m.subject_to(y == expr, name=f"{name}_ibart")
    return inputs, y

Datasets

jax_ibart.datasets.load_metal_support

load_metal_support()

Load the M-S single-atom catalysis dataset.

Returns (X, y, feature_names, units):

• X -- 91 x 59 primary-feature matrix.
• y -- 91 metal/oxide binding energies (eV).
• feature_names -- the 59 feature names.
• units -- 59 x 12 integer unit-exponent matrix (row per feature, column per base dimension), for the dimensional-analysis filter. Pass it to IBart.fit(..., units=units) to drop unit-inconsistent descriptors.

This is the exact catalysis dataset shipped with the R iBART package (converted from the R package's data/catalysis.rda), so results are directly comparable. Notable features for the published descriptor abs(Hfo/Oxv): Hfo (oxide formation enthalpy) and Oxv (oxygen vacancy energy); both carry unit eV.

Source code in src/jax_ibart/datasets.py

def load_metal_support():
    """Load the M-S single-atom catalysis dataset.

    Returns ``(X, y, feature_names, units)``:

    * ``X``     -- 91 x 59 primary-feature matrix.
    * ``y``     -- 91 metal/oxide binding energies (eV).
    * ``feature_names`` -- the 59 feature names.
    * ``units`` -- 59 x 12 integer unit-exponent matrix (row per feature, column
      per base dimension), for the dimensional-analysis filter. Pass it to
      ``IBart.fit(..., units=units)`` to drop unit-inconsistent descriptors.

    This is the exact ``catalysis`` dataset shipped with the R ``iBART`` package
    (converted from the R package's ``data/catalysis.rda``), so results are directly comparable.
    Notable features for the published descriptor abs(Hfo/Oxv):
    ``Hfo`` (oxide formation enthalpy) and ``Oxv`` (oxygen vacancy energy); both
    carry unit ``eV``.
    """
    Xdf = pd.read_csv(DATA_DIR / "catalysis_X.csv")
    y = pd.read_csv(DATA_DIR / "catalysis_y.csv")["y"].to_numpy(dtype=np.float64)
    X = Xdf.to_numpy(dtype=np.float64)
    names = list(Xdf.columns)

    # unit CSV: rows = base dimensions, cols = features (+ a row-name column).
    udf = pd.read_csv(DATA_DIR / "catalysis_unit.csv", index_col=0)
    units = udf[names].to_numpy(dtype=np.int64).T  # -> (n_features, n_dims)
    return X, y, names, units

jax_ibart.datasets.load_reaction

load_reaction()

Load the example batch-reactor kinetics dataset.

Returns (C, dCdt, species, batch):

• C -- (n, 4) measured concentrations [CA, CB, CC, CD].
• dCdt -- (n, 4) time derivatives, estimated by central finite differences within each batch (the mole-balance RHS targets).
• species -- ["CA", "CB", "CC", "CD"].
• batch -- (n,) batch index of each row.

Synthetic data from a liquid-phase batch reactor with A+B->C (bimolecular) and C->D (first order); see scripts/make_reaction_data.py. Use it to rediscover the mole-balance RHS (e.g. with jax_ibart.dynamics.discover_ode(C, t=None, species, dXdt=dCdt, screen="bart")).

Source code in src/jax_ibart/datasets.py

def load_reaction():
    """Load the example batch-reactor kinetics dataset.

    Returns ``(C, dCdt, species, batch)``:

    * ``C``       -- (n, 4) measured concentrations [CA, CB, CC, CD].
    * ``dCdt``    -- (n, 4) time derivatives, estimated by central finite
      differences *within each batch* (the mole-balance RHS targets).
    * ``species`` -- ["CA", "CB", "CC", "CD"].
    * ``batch``   -- (n,) batch index of each row.

    Synthetic data from a liquid-phase batch reactor with A+B->C (bimolecular)
    and C->D (first order); see ``scripts/make_reaction_data.py``. Use it to
    rediscover the mole-balance RHS (e.g. with
    ``jax_ibart.dynamics.discover_ode(C, t=None, species, dXdt=dCdt,
    screen="bart")``).
    """
    from .dynamics import finite_difference

    df = pd.read_csv(DATA_DIR / "reaction_kinetics.csv")
    species = ["CA", "CB", "CC", "CD"]
    C, dCdt, batch = [], [], []
    for b, g in df.groupby("batch"):
        g = g.sort_values("t")
        Cb = g[species].to_numpy(dtype=np.float64)
        C.append(Cb)
        dCdt.append(finite_difference(Cb, g["t"].to_numpy(dtype=np.float64)))
        batch.append(np.full(len(g), b))
    return (np.vstack(C), np.vstack(dCdt), species, np.concatenate(batch))

Dynamical systems

jax_ibart.dynamics.discover_ode

discover_ode(X_t, t, feature_names, degree=2, max_terms=3, units=None, dXdt=None, coef_threshold=0.001, screen=None, seed=0)

Discover dx_i/dt = f_i(x) by sparse regression on a polynomial library.

Parameters:

Name	Type	Description	Default
X_t	(n_time, n_state) trajectory samples.		required
t	(n_time,) time stamps.		required
feature_names	names of the state variables.		required
degree	polynomial library degree (<=2 here).		2
max_terms	max non-zero terms per equation (l0 best subset).		3
units	optional (n_state, n_dims) unit-exponent matrix.		None
dXdt	optional precomputed derivatives; if None, central finite differences.	(For multi-trajectory data, compute derivatives per trajectory and pass the stacked array here.)	None
screen	None or "bart". With "bart", a BART-G.SE nonparametric screen prunes	the library against each derivative before the l0 search, i.e. BART finds which terms enter each mole-balance RHS. Needs bartz.	None

Returns:

Type	Description
DiscoveredSystem

Source code in src/jax_ibart/dynamics.py

def discover_ode(
    X_t,
    t,
    feature_names,
    degree: int = 2,
    max_terms: int = 3,
    units=None,
    dXdt=None,
    coef_threshold: float = 1e-3,
    screen: str | None = None,
    seed: int = 0,
):
    """Discover dx_i/dt = f_i(x) by sparse regression on a polynomial library.

    Parameters
    ----------
    X_t : (n_time, n_state) trajectory samples.
    t : (n_time,) time stamps.
    feature_names : names of the state variables.
    degree : polynomial library degree (<=2 here).
    max_terms : max non-zero terms per equation (l0 best subset).
    units : optional (n_state, n_dims) unit-exponent matrix.
    dXdt : optional precomputed derivatives; if None, central finite differences.
        (For multi-trajectory data, compute derivatives per trajectory and pass
        the stacked array here.)
    screen : None or "bart". With "bart", a BART-G.SE nonparametric screen prunes
        the library against each derivative before the l0 search, i.e. BART finds
        which terms enter each mole-balance RHS. Needs ``bartz``.

    Returns
    -------
    DiscoveredSystem
    """
    import jax

    if screen not in (None, "bart"):
        raise ValueError(f"screen must be None or 'bart', got {screen!r}")
    X_t = np.asarray(X_t, dtype=np.float64)
    dX = finite_difference(X_t, t) if dXdt is None else np.asarray(dXdt, np.float64)

    lib = polynomial_library(X_t, feature_names, degree=degree, units=units)
    L = np.asarray(lib.X, dtype=np.float64)

    bart = None
    if screen == "bart":
        from .selection import BartGSESelector

        bart = BartGSESelector(n_perm=30, alpha=0.20, n_tree=40, n_draw=400, n_burn=400)

    rhs, r2 = [], []
    for i in range(len(feature_names)):
        yi = dX[:, i]
        if bart is not None:
            key = jax.random.PRNGKey(seed + i)
            keep = np.asarray(bart.select(L, yi, key=key))
            cand = np.where(keep)[0] if keep.any() else np.arange(L.shape[1])
        else:
            cand = sparsify.lasso_screen(
                L, yi, n_keep=min(2 * max_terms + 2, L.shape[1])
            )
        idx, coef = sparsify.l0_best_subset(
            L[:, cand], yi, max_dim=max_terms, score="bic"
        )
        # STLSQ-style cleanup: drop terms whose coefficient is negligible
        # relative to the largest, then refit OLS on the survivors.
        slopes = coef[1:]
        big = np.abs(slopes) > coef_threshold * (np.abs(slopes).max() + 1e-30)
        idx = [j for j, b in zip(idx, big) if b]
        cols = [L[:, cand[j]] for j in idx]
        A = (
            np.column_stack([np.ones(len(yi))] + cols)
            if cols
            else np.ones((len(yi), 1))
        )
        coef, _ = (
            sparsify.ols(np.column_stack(cols), yi)
            if cols
            else (np.array([yi.mean()]), 0.0)
        )

        expr = sp.Float(coef[0])
        for c, j in zip(coef[1:], idx):
            expr = expr + sp.Float(c) * lib.exprs[cand[j]]
        pred = A @ coef
        ss = 1 - ((yi - pred) ** 2).sum() / (((yi - yi.mean()) ** 2).sum() + 1e-30)
        rhs.append(expr)
        r2.append(float(ss))
    return DiscoveredSystem(list(feature_names), rhs, r2)

jax_ibart.dynamics.implicit_constraints

implicit_constraints(X, feature_names, degree=2, units=None, tol=0.01)

PROTOTYPE: find algebraic constraints g(x) ~ 0 (the DAE algebraic part).

A differential-algebraic system couples dynamics dx/dt = f with algebraic constraints 0 = g(x, z). Discovering g is an implicit problem: find a sparse coefficient vector xi != 0 with Theta(x) @ xi ~ 0, i.e. a vector in (or near) the null space of the library matrix. This is NOT the regress-on-a- target problem iBART/SINDy solve, which is why it needs separate machinery.

This prototype: 1. builds the iBART library Theta (with the unit filter), 2. standardizes columns and takes the SVD, 3. returns the smallest right-singular vector as the leading candidate constraint, with its normalized residual ||Theta xi|| / ||Theta||.

A production DAE module (the "something special to develop") would add: - sparse null-space selection (only a few terms active), e.g. SINDy-PI's per-term implicit regressions or sparse-null-space algorithms, instead of the dense smallest singular vector; - a weak / integral formulation to handle noise and avoid derivatives (Messenger & Bortz 2021); - separation of differential vs algebraic relations and DAE index/ consistency handling.

See docs/references.md (Mangan et al. 2016; Kaheman et al. 2020; Messenger & Bortz 2021).

Returns a dict with the candidate constraint expression, its coefficient vector over the library terms, and the residual (smaller = closer to an exact constraint).

Source code in src/jax_ibart/dynamics.py

def implicit_constraints(
    X, feature_names, degree: int = 2, units=None, tol: float = 1e-2
):
    """PROTOTYPE: find algebraic constraints g(x) ~ 0 (the DAE algebraic part).

    A differential-algebraic system couples dynamics dx/dt = f with algebraic
    constraints 0 = g(x, z). Discovering g is an *implicit* problem: find a
    sparse coefficient vector xi != 0 with Theta(x) @ xi ~ 0, i.e. a vector in
    (or near) the null space of the library matrix. This is NOT the regress-on-a-
    target problem iBART/SINDy solve, which is why it needs separate machinery.

    This prototype:
      1. builds the iBART library Theta (with the unit filter),
      2. standardizes columns and takes the SVD,
      3. returns the smallest right-singular vector as the leading candidate
         constraint, with its normalized residual ||Theta xi|| / ||Theta||.

    A production DAE module (the "something special to develop") would add:
      - *sparse* null-space selection (only a few terms active), e.g. SINDy-PI's
        per-term implicit regressions or sparse-null-space algorithms, instead of
        the dense smallest singular vector;
      - a *weak / integral* formulation to handle noise and avoid derivatives
        (Messenger & Bortz 2021);
      - separation of differential vs algebraic relations and DAE index/
        consistency handling.

    See docs/references.md (Mangan et al. 2016; Kaheman et al. 2020; Messenger &
    Bortz 2021).

    Returns a dict with the candidate constraint expression, its coefficient
    vector over the library terms, and the residual (smaller = closer to an
    exact constraint).
    """
    lib = polynomial_library(
        X, feature_names, degree=degree, units=units, dedup_collinear=False
    )
    Theta = np.asarray(lib.X, dtype=np.float64)
    exprs = list(lib.exprs)
    # augment with a constant column so constraints like x^2 + y^2 - 1 = 0 are
    # representable (the engineer drops constant columns, so add it here).
    Theta = np.column_stack([Theta, np.ones(len(Theta))])
    exprs = exprs + [sp.Integer(1)]

    # scale columns to unit norm (not standardize: that breaks the constant
    # column), then the smallest right-singular vector is the best near-null
    # combination.
    norms = np.linalg.norm(Theta, axis=0) + 1e-30
    _, S, Vt = np.linalg.svd(Theta / norms, full_matrices=False)
    xi = Vt[-1] / norms
    residual = float(S[-1] / (S[0] + 1e-30))

    # normalize so the largest coefficient is 1, then threshold small terms
    xi = xi / (np.abs(xi).max() + 1e-30)
    expr = sp.Float(0.0)
    for c, e in zip(xi, exprs):
        if abs(c) > tol:
            expr = expr + sp.Float(round(c, 4)) * e
    return {
        "constraint": expr,
        "coef": xi,
        "terms": [str(e) for e in exprs],
        "residual": residual,
        "note": "PROTOTYPE: dense SVD null space; production needs sparse "
        "null-space selection + weak form (see docstring).",
    }

jax_ibart.dynamics.finite_difference

finite_difference(X_t, t)

Central finite-difference time derivatives of a trajectory.

X_t is (n_time, n_state); t is (n_time,). Returns dX/dt with the same shape (forward/backward differences at the endpoints). For noisy data prefer a smoothing or weak/integral formulation (see docs/references.md); this is the clean-data baseline.

Source code in src/jax_ibart/dynamics.py

def finite_difference(X_t, t):
    """Central finite-difference time derivatives of a trajectory.

    ``X_t`` is (n_time, n_state); ``t`` is (n_time,). Returns dX/dt with the same
    shape (forward/backward differences at the endpoints). For noisy data prefer
    a smoothing or weak/integral formulation (see docs/references.md); this is
    the clean-data baseline.
    """
    X_t = np.asarray(X_t, dtype=np.float64)
    t = np.asarray(t, dtype=np.float64)
    if X_t.ndim != 2 or X_t.shape[0] < 2:
        raise ValueError("X_t must be (n_time, n_state) with at least 2 time points")
    if t.shape != (X_t.shape[0],):
        raise ValueError(f"t must have shape ({X_t.shape[0]},), got {t.shape}")
    dX = np.empty_like(X_t)
    dX[1:-1] = (X_t[2:] - X_t[:-2]) / (t[2:] - t[:-2])[:, None]
    dX[0] = (X_t[1] - X_t[0]) / (t[1] - t[0])
    dX[-1] = (X_t[-1] - X_t[-2]) / (t[-1] - t[-2])
    return dX