Main optimization step of direct SNR optimization

SNROptConvex.jl

using Convex
using COSMO
using LinearAlgebra 



"""

INPUT:
 - Matrix containing the flattened subtraction region of all the images 
 - Matrix containing the flattened optimization region of all the images 
 - Vector for forward modelled photometry of the planet PSF centered in the S region for each image.
 - c0 initial coefficients. Should be all zeros with ones at the locations of the target image(s)

Given a subtraction region, optimization region, photemtry vector, and target
index, return the coefficients for optimally combining the subtraction regions
to achieve a good subtraction. Also return some SNR stats.
"""
function optimize_subtraction(
    S::AbstractArray{ValType,2},
    O::AbstractArray{ValType,2},
    phot,
    c0,
) where {ValType}

    SNR = let phot = phot
        function (images, c)
            return phot ⋅ c / norm(images * c) * sqrt(size(images, 1))
        end
    end

    pixel_mask = reshape(all(isfinite, O, dims=2), :)
    O = O[pixel_mask,:]

    # Initial SNR in the subtraction and optimization regions
    SNR_Si = SNR(S, c0)
    SNR_Oi = SNR(O, c0)


    # Set up variables
    c = Variable(length(phot))
    noise = norm(O*c)
    regularization = sumsquares(c)

    # Best coefficients will go here
    cf = copy(c0)
    cf_temp = similar(c0)

    # Regularization parameter to prevent overfitting.
    # We will step through optimizing the SNR in the optimization region
    # and tesing the SNR in the subtraction region until the subtraction
    # region stops improving.
    # of the 
    αs = (
        0.0,
        1e-4,
        1e-3,
        1e-2,
        1e-1,
        1e0,
        1e1,
        1e2,
        1e3,
        1e4,
        1e5,
        1e6
    )
    SNR_S_previous = SNR_S = SNR_Si
    best_α = first(αs)
    i=0
    for α in αs
        if α != 0.0
            cost = noise + α*regularization
        else
            cost = noise
        end

        problem = minimize(cost)
        problem.constraints += (
            dot(phot, c) == 1
        )

        # Solve!
        solve!(problem, () -> COSMO.Optimizer(verbose=false), verbose=false, warmstart=i>0)
        # solve!(problem, () -> SCS.Optimizer(verbose=0), verbose=false, warmstart=i>0)
        
        # Extract best parameter values
        if isnothing(c.value)
            println(problem)
            error("Solution did not converge")
        end
        if length(c.value) > 1
            cf_temp .= view(c.value, :)
        else
            cf_temp .= c.value
        end

        SNR_S = SNR(S, cf_temp)

        if SNR_S ≤ SNR_S_previous
            # Reset to the last values since that was better
            SNR_S = SNR_S_previous
            break
        end
        best_α = α
        SNR_S_previous = SNR_S
        # else, this was an improvement.
        # Update the best coefficients
        cf .= cf_temp
        i+=1
    end

    SNR_O = SNR(O, cf)

    return (; cf, SNR_Si, SNR_Oi, SNR_S, SNR_O)
end