Implementing Mahalanobis Distance Loss

[katcheng]

Hello,

First, thank you for making such a helpful package.

I am currently trying to code a custom loss function calculates the Mahalanobis distance (https://en.wikipedia.org/wiki/Mahalanobis_distance ) as the loss. The below code is what I have so far, along with a test dataset to use:

#create the dataset columns
X = abs(np.random.randn(12, 3))
manif = [1, 1, 1, 3, 3, 3,1, 1, 1, 3, 3, 3]
y= [50, 150, 100, 75, 80 , 250 ,50, 150, 100, 75, 80 , 250 ]
weights3 = [1,1,1,1,1,1, 1,1,1,1,1,1]
#create the whole dataset
Xwithmanif = pd.concat([pd.DataFrame(X), pd.Series(manif)], axis = 1)
Xwithmanif.columns = ['x1', 'x2', 'x3', 'manif']
#initialize Statistics package in Julia
jl.seval('using Statistics')

#create the mahal loss
mahalloss2 = """
function mahalloss2(tree, dataset::Dataset{T,L}, options) where {T,L}

    X = copy(dataset.X)
    y = copy(dataset.y)

    weights2 = dataset.weights 
    temploss = Vector{Float64}()
    
    X_without_manifs = vcat(
        X[1:3, :],  # The actual data
        transpose(X[4, :] .* 0 )) # Pass zeroed version
#since the manifs are used to just separate out the different cov matrices and not for actual prediction
    
    prediction, flag = eval_tree_array(tree, X_without_manifs, options)
    if !flag
        return L(Inf)
    end

    Xnomanifs = X[1:3, :]
    
    Xtotal = hcat(transpose(X), y)

    Xwithpred = hcat(transpose(Xnomanifs), prediction)

    #break out the cov matrices
    Xclass12 = Xtotal[(Xtotal[:,4] .== 1) ,:]
    Xclass3 = Xtotal[(Xtotal[:,4] .== 3) ,:]
    
    Xclass12nomanif = Xclass12[:, 1:end .!= 4]
    Xclass3nomanif = Xclass3[:, 1:end .!= 4]

    invcov12 = inv(Statistics.cov(Xclass12nomanif))
    invcov3 = inv(Statistics.cov(Xclass3nomanif))

    #get the mean of each X feature for each class
    class12mean = Statistics.mean.(eachcol(Xclass12nomanif))
    class3mean = Statistics.mean.(eachcol(Xclass3nomanif))

    #calculate Mahanal
    for i in 1:length(y)
        if prediction[i] < 0 #prediction can't be less than 0
            #append!(temploss, 999)
            return L(Inf)
        elseif y[i] < 200 #if class 1/2
            x_minus_mu = Xwithpred[i, :] - class12mean
            tempmahal = transpose(x_minus_mu) * (invcov12) * x_minus_mu
            append!(temploss, tempmahal)
        else #if class 3
            x_minus_mu = Xwithpred[i, :] - class3mean
            tempmahal = transpose(x_minus_mu) * (invcov3) * x_minus_mu
            append!(temploss, tempmahal)
        end

    end
    return sum(temploss.* weights2)

end 
"""
#run through pysr
model = PySRRegressor(
    niterations=10,
    binary_operators=["*", "+", "-", "greater"],
    loss_function=mahalloss2,
)
model.fit(Xwithmanif, y, weights = weights3 )
model.get_best().equation

The issue I am having is that currently it will run near infinitely (I haven't been able to get it to stop and have had to restart kernel multiple times to interrupt it). I have tried doing print statement debugging and just watch the code constantly repeat and never reaching the final return of the loss function. Some specific questions I have:
-Should I return L(Inf) or just append a very high loss for if predictions are negative?
-How come the code seems to be running many more loops than the niterations I set?
-Why is it not returning the final sum, which should be the loss?

Please let me know if I can help supply any more information that might be helpful to answer these questions.

[MilesCranmer]

Thanks for the question, glad to hear it's useful!

For Mahalanobis I would first try out Distances.jl instead of writing one from scratch as they likely have a fairly optimised version of it: https://github.com/JuliaStats/Distances.jl, and it will plug in just fine. Optimising the metric might be the issue in all honesty.

See the Primes.jl example in the PySR documentation for how to install and then interface with an external package.

Let me know if this works!

[katcheng]

Hi Miles,

Thank you for the quick reply! I have tried the Mahalanobis from Distances.jl in a small non-loss based test case before, but unfortunately I wasn't able to completely understand how to implement it especially after reading through this thread: https://discourse.julialang.org/t/understanding-mahalanobis-in-distances-jl/59656/3 and also this thread: https://discourse.julialang.org/t/efficient-calculation-of-mahalanobis-distance/90121/4 where it seems like the calculation isn't using the inverted covariance matrix. I think the mismatch of use cases might be me needing to calculate a Mahalanobis distance of a "new" point that isn't in the preexisting dataframe used to create the covariance matrix, so a pre-calculation like what I think the source code in Julia is doing wouldn't work. Please let me know if I have made an incorrect assumption, though, and your help is greatly appreciated!