Custom loss with partial derivatives for 2 inputs and 2 outputs function.

[NiccoloAntonelliDziri]

Hi everyone,

I am currently trying to use PySR for finding an approximant for the inverse of a function. I apologize if there is something trivial that I missed but I haven't been able to find a discussion or an issue for something similar.

I want to find two functions, f1 and f2, such that:

d = f1(P, h)
T = f2(P, h)

I also have a custom condition on their partial derivatives that I want to enforce during the optimization process:

(df1/dh) / (df2/dh) = (df1/dP) / (df2/dP) + (2*(df1/dh)*(df1/dP)) / ((df2/dP)*(df1/dh) + (df2/dh)*(df1/dP))

My hope is that this condition will help steer the algorithm toward solutions that are physically meaningful.

I have X and y of shape (241127, 2), where:

X contains the inputs [P, h]
y contains the targets [d, T]

I used a custom loss function in PySR to combine the standard MSE loss with the derivative condition.

When I tried to implement the custom loss, I got a BoundsError:
BoundsError: attempt to access 2×241127 Matrix{Float32} at index [1:2, 1, 2]

This suggests that eval_grad_tree_array only returns a (2, N) matrix (2 variables × N data points), not the full Jacobian for both outputs.
My attempts to extract four derivatives (df1/dP, df1/dh, df2/dP, df2/dh) failed because PySR appears to call the loss function separately for each output (or at least it's my guess).

Is it possible to enforce such a condition in PySR with a single loss function for two outputs?
If not, is there a recommended way to achieve this ?

d = np.linspace(0.001, 5, 500)
T = np.linspace(200, 1000, 500)

D, T_grid = np.meshgrid(d, T)

p, h = ph_dt(D, T_grid)

pp = p[~np.isnan(p) & ~np.isnan(h)]
hh = h[~np.isnan(p) & ~np.isnan(h)]
DD = D[~np.isnan(p) & ~np.isnan(h)]
TT = T_grid[~np.isnan(p) & ~np.isnan(h)]

X = np.column_stack((pp.ravel(), hh.ravel()))
y = np.column_stack((DD.ravel(), TT.ravel()))

from pysr import jl, PySRRegressor
jl.seval("using Zygote") # For gradient calculations

extra_sympy_mappings={
    "pow_n": lambda x, n: sympy.Pow(x, n),
}

model = PySRRegressor(
    unary_operators=["log", "exp", "sqrt"],
    nested_constraints={
        "log": {"log": 0, "exp": 0},
        "exp": {"exp": 0, "log": 0},
        "sqrt": {"sqrt": 0},
    },

    maxsize=30,
    weight_optimize=0.001,
    batching=True,                 # Enable mini-batch training
    batch_size=1024,               # Recommended batch size (adjust as needed)
    niterations=100,               # Number of iterations (adjust for complexity)
    binary_operators=["+", "-", "*", "/", "pow_n(x::T, n::T) where(T) = x > 0 ? convert(T, x^n) : convert(T, NaN)"],  # Operators
    extra_sympy_mappings=extra_sympy_mappings,  # Custom unary operators
    constraints={"pow_n": (-1,1)}, # Allow the model to fit the exponent "a"
    populations=36,               # Number of populations (influences exploration)
    population_size=64,            # Number of individuals in each population
    ncycles_per_iteration=100,     # Number of total mutations to run, per 10 samples of the population, per iteration
    model_selection="best",        # Keep the best-performing model
    parsimony=1e-5,                # parsimony (times complexity) added to the loss
    maxdepth=7,                    # Max depth of the equation
    progress=True,                 # Display training progress
    verbosity=1,                   # Show intermediate output

    loss_function = """
function custom_derivative_loss(tree, dataset::Dataset{T,L}, options) where {T,L}
    # Evaluate the expression tree
    (prediction, gradients, completion) = eval_grad_tree_array(tree, dataset.X, options; variable=true)
    !completion && return L(Inf)

    # Compute MSE loss (standard)
    diffs = prediction .- dataset.y
    mse_loss = sum(abs2.(diffs)) / length(diffs)


    df1_dp = gradients[:, 1, 1]  # ∂f1/∂p for all N data points
    df1_dh = gradients[:, 1, 2]  # ∂f1/∂h for all N data points
    df2_dp = gradients[:, 2, 1]  # ∂f2/∂p for all N data points
    df2_dh = gradients[:, 2, 2]  # ∂f2/∂h for all N data points

    denominator = (df2_dp .* df1_dh .+ df2_dh .* df1_dp)

    if abs(denominator) < 1e-10 || isnan(denominator)
        derivative_condition_loss = L(Inf)
    else
        rhs = (df1_dp ./ df2_dp) .+ (2 .* df1_dh .* df1_dp) ./ denominator

        lhs = df1_dh ./ df2_dh

        derivative_condition_loss = sum(abs2.(lhs .- rhs)) / length(lhs)
    end

    total_loss = mse_loss + derivative_condition_loss

    return total_loss
end
"""
)
model.fit(X, y, variable_names = ["p", "h"])

[MilesCranmer]

I think you could try doing this with TemplateExpressionSpec and loss_function_expression and you could get faster symbolic derivatives?

---

I'll do some digging for relevant issues... (@MilesCranmerBot can you please read through discussions/issues over the past year and point to the most relevant ones? Ideally it would be nice if you could find a discussion that shows how to use loss_function_expression in tandem with template expressions. Although I think for this we might not even need to, and can simply have the template return the per-row loss?)

[NiccoloAntonelliDziri]

Thank you for your answer, I spent a few hours trying to use TemplateExpressionSpec :

template = TemplateExpressionSpec(
    expressions=["f1", "f2"],  # The two functions to learn
    variable_names=["p", "h"],  # Inputs
    combine="(f1(p,h),f2(p,h))",  # Output both functions
    # parameters={"p":2,"h":2},
)

but could not get it to work (with and without setting a custom loss). Maybe I haven't tried hard enough but setting combine="f1(p,h) + f2(p,h)" instantly seemed to work. However doing it with a single output and this sum will loose potentially crucial information for what I'm trying to do.
Is this an intended behavior of TemplateExpressionSpec or am I still missing something ?

[MilesCranmer]

Yes this will not work b/c combine has to return a single value rather than a tuple in PySR.

Sorry my bot didn't run, I ran out of LLM calls for the week... Will run it later today

[MilesCranmerBot]

I dug through the most relevant recent threads, and I think the clearest pointers are:

• PySR discussion Define Multiple Outputs with TemplateExpressionSpec #1002, “Define Multiple Outputs with TemplateExpressionSpec”: Define Multiple Outputs with TemplateExpressionSpec #1002
  • This is the closest match to the current question. The recommended pattern is to pass y_0, y_1, etc. as features, have the template return the per-row loss, and use elementwise_loss="(pred, target) -> pred".
• PySR discussion BoundsError when using TemplateExpressionSpec with custom loss function #1011, “BoundsError when using TemplateExpressionSpec with custom loss function”: BoundsError when using TemplateExpressionSpec with custom loss function #1011
  • Confirms that with templates you generally want loss_function_expression, not loss_function.
  • Also suggests that for derivative-based objectives it is often cleaner to put the derivative terms directly in the template, e.g. with D(...), and again return the per-row residual.
• PySR discussion [BUG]: TemplateExpressionSpec and loss_function_expression not compatible with multiprocessing #939 / issue [BUG]: TemplateExpressionSpec and loss_function_expression not compatible with multiprocessing #938: [BUG]: TemplateExpressionSpec and loss_function_expression not compatible with multiprocessing #939 and [BUG]: TemplateExpressionSpec and loss_function_expression not compatible with multiprocessing #938
  • Example of TemplateExpressionSpec working with loss_function_expression; the bug there was just the custom function name/signature.
• SymbolicRegression.jl issue How to compute the complexity of a function inside the custom objective function? #380: [BUG]: when using loss_function, Template Expression cannot have operators outside of binary_operators astroautomata/SymbolicRegression.jl#380
  • This is the backend thread that motivated loss_function_expression for TemplateExpression.
  • It also explicitly suggests the workaround of using the template itself as the loss, by passing the true targets in X and returning the residual from the template.

So I think the answer to the second part is yes: for this use case, the template can directly return the per-row loss, and that is probably the cleanest PySR route.

The important limitation is that in PySR the template combine still needs to resolve to a single per-row value (internally a ValidVector), not a tuple of outputs. So something like (f1(p,h), f2(p,h)) will not work directly, but something like

(y1 - f1(p,h))^2 + (y2 - f2(p,h))^2

will.

There is now also a docs example mirroring this pattern: “Vector-valued expressions” in the PySR examples: https://ai.damtp.cam.ac.uk/pysr/examples/#13-vector-valued-expressions

[NiccoloAntonelliDziri]

I'll look into that on monday and come back to you ! Thanks!