Issues with Regressing Circuit Using PySR: Positive Parameters and Custom Loss Functions

[frank-ccc]

Dear Developers,

Thank you for your excellent work on the PySR repository. It is truly fantastic. I have been working on regressing a ZARC circuit using the code below, but I have encountered some issues:

1. The code seems unable to fit the model because it uses complex parameters (not only complex inputs) with some semi-default loss function (see below).
2. When I attempt to use a custom loss function, the regression still does not work, to handle positivity as discussed in How to constrain all the constants to be positive #449

Could there be a workaround for this? Specifically, I am looking for a solution where the symbolic regression parameters are positive, but the function input and output remain complex.

Thank you for your kind help.

P.S. This question is related but somewhat different to #338

CODE:

import numpy as np
import matplotlib.pyplot as plt
from pysr import PySRRegressor
from sympy import latex, re

# File name
file_name = './results/ZARC.csv'

# Parameters for ZARC model
R_inf, R_ct, phi, tau_0 = 10, 50, 0.8, 0.01

# Compute impedance of ZARC model
def Z_ZARC(s, R_inf, R_ct, phi, tau_0):
    return R_inf + R_ct / (1 + (s * tau_0)**phi)

# Frequency range
N_freqs = 71
freq_vec = np.logspace(-1, 6, N_freqs)

# Exact impedance
Z_exact = Z_ZARC(1j * 2 * np.pi * freq_vec, R_inf, R_ct, phi, tau_0)

# Synthetic data with noise
np.random.seed(1225)
sigma_n_exp = 0.2
Z_exp = Z_exact + sigma_n_exp * (np.random.normal(0, 1, N_freqs) + 1j * np.random.normal(0, 1, N_freqs))

# Normalize data
Z_abs_mean, Z_abs_std = np.mean(np.abs(Z_exp)), np.std(np.abs(Z_exp))
Z_normalized = (Z_exp - Z_abs_mean) / Z_abs_std


# PySRRegressor model setup
model = PySRRegressor(
    equation_file=file_name,
    binary_operators=["+", "-", "*", "/", "mypow(x,y)=(x)^(y.re)"],
    constraints={"/": (-1, 9)},
    extra_sympy_mappings ={"mypow": lambda x,y: x**re(y)},
    niterations=100,
    maxsize=30,
    maxdepth=10,
    adaptive_parsimony_scaling=100.0,
    precision=64,
    verbosity=1,
    elementwise_loss="f(x, y) = abs2(x - y)/abs2(x)"
)

# Fit the model
model.fit((1j * 2 * np.pi * freq_vec).reshape(-1, 1), Z_normalized)

# Predict using the model
Z_pred_normalized = model.predict((1j * 2 * np.pi * freq_vec).reshape(-1, 1))
Z_pred = Z_abs_mean + Z_abs_std * Z_pred_normalized

# Plot results
plt.figure(figsize=(10, 5))
plt.plot(np.real(Z_exact), -np.imag(Z_exact), '--', linewidth=2, color='blue', label='Exact')
plt.plot(np.real(Z_exp), -np.imag(Z_exp), 'o', markersize=7, color='red', label='Synthetic Exp')
plt.plot(np.real(Z_pred), -np.imag(Z_pred), 'x', markersize=7, color='green', label='Symbolic Regression')
plt.legend(frameon=False, fontsize=15)
plt.axis('scaled')
plt.gca().set_aspect('equal', adjustable='box')
plt.xlabel(r'$Z_{\mathrm{re}}/\Omega$', fontsize=20)
plt.ylabel(r'$-Z_{\mathrm{im}}/\Omega$', fontsize=20)
plt.tight_layout()
plt.show()

# Print best equation
best_equation = model.sympy()
print('*********')
print('Best equation found:', best_equation)
print('Best equation in LaTeX:', latex(best_equation))

[MilesCranmer]

Thanks!

Specifically, I am looking for a solution where the symbolic regression parameters are positive, but the function input and output remain complex.

By this, do you mean parameters which have no imaginary component and whose real component is positive? (Since “positive” doesn’t apply to complex numbers)

Or do you mean positive imaginary and positive real component?

[frank-ccc]

I am sorry that was not clear; by " parameters are positive," I mean that the "parameters are positive real numbers."

[MilesCranmer]

Thanks, will write an answer with that in mind

[MilesCranmer]

So, right now, there are various assumptions that the expressions have has the same numeric type as the dataset. i.e., expressions will have complex-valued constants.

However, perhaps you could try a soft constraint that the parameters are close to the positive real line? See https://astroautomata.com/PySR/examples/#9-custom-objectives for an example. There are many other examples in the discussions page.

Here's how you could do it:

function my_objective(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    (prediction, completion) = eval_tree_array(tree, dataset.X, options)
    y = dataset.y
    if !completion
        return L(Inf)
    end
    
    function node_penalty(node)
        is_constant_node = node.degree == 0 && node.constant
        if is_constant_node
            val = node.val
            return L(abs(abs(val) - val))  # How far the constant is from its projection to the positive real line
        else
            return L(0)
        end
    end
    
    # Aggregate the node penalty over every node in the tree:
    total_node_penalty = sum(node_penalty, tree)
    total_node_penalty *= 1000.0  # Upweight penalty (likely need to tune this?)
    
    # Regular relative MSE loss:
    prediction_loss = sum(i -> abs2(prediction[i] - y[i])/abs2(y[i]), eachindex(y)) / length(y)

    return L(total_node_penalty + prediction_loss)
end

then pass that as a string to loss_function.

Let me know how that works!

Note that constants will still have some small complex component. But hopefully this will eliminate most of them...

[MilesCranmer]

P.S., your initial loss was provided as elementwise_loss

    elementwise_loss="f(prediction, target) = abs2(prediction - target)/abs2(prediction)"

however, you don't want to do this, as the genetic algorithm will "game the system" and send prediction to infinity! You probably want to use

    elementwise_loss="f(prediction, target) = abs2(prediction - target)/abs2(target)"

instead.

But use the loss_function parameter instead with the full loss function in #651 (comment)

[MilesCranmer]

Also make sure none of the targets are zero.