Is it possible to account for the covariance matrix in PySR?

[zouzaxd]

Hello,

I’m working with data that includes a covariance matrix, and I was wondering if it’s possible to incorporate or consider this covariance matrix directly in PySR. If yes, could you provide some guidance or examples on how to implement this?

Thank you for your help!

[MilesCranmer]

Do you mean covariance between rows of y? If this is what you are looking for:

import numpy as np
from pysr import PySRRegressor
from pysr import jl

X = np.random.randn(100, 2)  # Your input data
y = np.random.randn(100)     # Your target values
cov_matrix = # Your covariance matrix

# First compute and register inverse of covariance matrix
inv_cov = np.linalg.inv(cov_matrix)  # Compute inverse in Python

# Important we make it `const` here for performance reasons
jl.seval("x -> @eval const INV_COV_MATRIX = $x")(inv_cov)

loss_function = """
function custom_loss(tree, dataset::Dataset{T,L}, options) where {T,L}
    # Compute predictions
    prediction, completed = eval_tree_array(tree, dataset.X, options)
    
    if !completed
        return L(Inf)
    end
    
    # Get residuals
    diffs = prediction .- dataset.y
    n = length(diffs)
    
    # Compute (diffs' * INV_COV_MATRIX * diffs) / n
    loss = 0.0
    for i in 1:n
        weighted_sum = 0.0
        for j in 1:n
            weighted_sum += INV_COV_MATRIX[i,j] * diffs[j]
        end
        loss += diffs[i] * weighted_sum
    end
    
    return L(loss / n)
end
"""

model = PySRRegressor(
    loss_function=loss_function,
    # ... other parameters as needed ...
)

[zouzaxd]

Thank you so much for your quick reply! It worked, I appreciate your help and the great work you've done with PySR.