Does PySR support a custom ODE loss (weak-form / integral-residual over grouped time-series)?

[yurbro]

Hi PySR team — thanks for the amazing library!

Use case (ODE / weak-form SR)

We model an autonomous ODE:
dQ/dt = f(Q, C)
where C = (C_pol, C_eth, C_pg)

Data: For each formulation (fixed C), we have a short time series
{(t_k, Q_k)}_{k=1..K} with K ≈ 7 (0.5–6 h).
Across 30 formulations we have 30 trajectories.

Instead of fitting pointwise y ≈ f(X), I want to learn an expression for f(Q, C) by minimizing a grouped, integral residual (weak form). For each trajectory g:

L_g = sum_{k=1..K_g} [ int_{t0}^{t_k} f(Q_obs(tau), C_g) d(tau)
                        - (Q_obs(t_k) - Q_obs(t0)) ]^2
Total loss = sum_g L_g

This avoids numerical differentiation of Q and is more robust for short, noisy trajectories.

Question 1 — custom loss API

Does PySR allow a custom objective that depends on X, y, and the candidate predictions yhat (here yhat = f(Q, C) evaluated at each sample), so that inside the loss I can:

• group samples by trajectory (e.g., via a group_id column in X);

• sort each group by time t;

• integrate yhat along t (e.g., trapezoid rule) to approximate int f(Q, C) dt;

• compare the cumulative integral to observed Q(t) - Q(t0) and sum squared residuals?

Concretely, is there a full_objective-style hook in PySR where the signature gives access to X, y, yhat (and optionally weights) so I can implement the above? If yes, what’s the recommended way to pass it from Python — e.g., a Julia function string, or a Python callback? A minimal example would be super helpful.

Question 2 — grouping / time columns

Any guidance on best practices for grouped losses? For example, if I pack columns as

X = [C_pol, C_eth, C_pg, Q_obs, t, group_id]

is it OK that the loss uses t and group_id only for bookkeeping, while the expression itself only “sees” [Q_obs, C_pol, C_eth, C_pg] as variables?

Question 3 — operators & safe numerics

For physical robustness we often use safe atoms like log1p, protected division, etc. Is it correct that we can provide these via unary_operators / binary_operators and map them with extra_sympy_mappings / extra_torch_mappings for export? Any caveats when using such custom operators together with a custom objective?

Minimal pseudo-code (intended)

from pysr import PySRRegressor

# X columns: [C_pol, C_eth, C_pg, Q_obs, t, group_id]
# y: Q_obs  (we won’t use y as a pointwise target; loss integrates yhat over t)

pysr = PySRRegressor(
    niterations=2000,
    # Expression should output dQ/dt given [Q_obs, C_pol, C_eth, C_pg]
    unary_operators=["log1p_abs(x)"],   # example safe atom
    binary_operators=["pdiv(x, y)"],    # protected division (example)
    # --- The key part I’m asking about: ---
    # Is there a way to pass a full objective that gets (X, y, yhat, w)?
    loss="""
    function loss(X, y, yhat, w)
        # X: N x D (includes t and group_id)
        # y: observed Q
        # yhat: model’s dQ/dt at each row (f(Q,C))

        # 1) group by group_id
        # 2) within each group, sort by t
        # 3) trapezoid integrate yhat over t to get Q_hat(t_k)-Q(t_0)
        # 4) compute sum of squared residuals vs (y - y0)
        # return scalar total loss
        return total_loss
    end
    """,
)

pysr.fit(X, y)

Notes: The above is pseudo-code to illustrate the desired API; I’m asking what the actual way is to implement such a grouped integral-residual objective in PySR.

Thanks a lot! Any pointers to documentation or examples for a trajectory-level (non-elementwise) objective in PySR would be greatly appreciated.

[MilesCranmer]

I think the right way to do this is via TemplateExpressionSpec and also loss_function

Take a look at, e.g., #1002 (comment) and other discussions about the template specs.

It is arbitrary Julia code, and doesn't need to be differentiable. So technically you could do anything you want there!