mosquito_code/MosquitoPackage.jl

module MosquitoPackage

using DifferentialEquations

include("base.jl")
include("misc.jl")
include("model.jl")
include("resim.jl")
include("resim_new.jl")
function f1(x, y)
    return 3x + y
end

CENTER_COORDS = [0.0 0.0]
SCALING_FACTOR = 1.0

# function distance_to_center!(v, r; center_coords=CENTER_COORDS)
#     for k in 1:length(r)
#         v[k] = r[k] - center_coords[k]
#     end
#     dd = mynorm(v)
#     v ./= dd
#     return dd
# end

function distance_to_center!(v, r, center_coords)
    for k in 1:length(r)
        v[k] = r[k] - center_coords[k]
    end
    dd = mynorm(v)
    v ./= dd
    return dd
end


function define_model(; ndim=3, f1=LaguerreFunction{Float64}(), n1=collect(0:5), p1=ones(1), c1=zeros(length(n1)), bfield1=nullfunc!,
                        f2a=LaguerreFunction{Float64}(), f2b=LaguerreFunction{Float64}(), f2ab=LegendreP{Float64}(),
                        n2a=collect(0:5), n2b=collect(0:5), n2ab=collect(0:5), c2=zeros(length(n2a)*length(n2b)*length(n2ab)),
                        p2=ones(2), bfield2=nullfunc!, D=0.0, len_scale=1.0, target=:visualco2)

    model=MosquitoModel(; ndim=ndim, f1=f1, f2a=f2a, f2b=f2b, f2ab=f2ab);

    model.islearned[:visual] = true

    if target == :visualco2
        model.funcs[:visual] = [Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 3, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 3, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 3, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 3, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 0, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 3, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 2, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 3, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 3, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 3, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 2, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 3, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 2, 3, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 0, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 2, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 2, 0, :bhatorth2)]
        model.coeffs[:visual] = vcat(
        [-26.58508768656102, -10.025540504214353, -19.281931331711576, 16.977296218016942, 3.5274605792090674, -17.538921168146157, -0.8978633007550898, 0.5757674869632454, -0.34720605748622313, 27.285528208991952, -1.6629950732453829, -1.0093785167806952, -6.6835266268737215, -10.976471308197574, 87.25051545826713, 6.2020878634158185, 0.3504622511215739, 56.91679044163722, -60.54927762496129, 1.3669974738018196, 3.3570117102041386, -58.61752756159205, 1.8720678662134222],
        [-968.4160609423942, 2.938765454316324, 21.388261050998572, 3358.4406597434454, -34.40532797779392, -3931.911838118647, 4.436187969207309, 1549.2916146950502, -1.2742586401918503, 25.447437000300674, 14.92461449802172, -11.169183496972131, -3.548778046602168, -1.9539308398909119, -3.3637899732746517, -9.122537038444928]
        )
        model.params[:visual] = [2.0, len_scale]
    end

    if target == :visual
        model.funcs[:visual] = [Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 5, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 5, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 2, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 1, 4, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 0, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 1, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 2, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 2, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(5, 1, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(5, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(5, 2, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 2, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 0, 2, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(4, 2, 2, :bhatorth2)]    
        model.coeffs[:visual] = vcat(
        [-23.049553930068313, 1.2805363220377013, 15.634000112901866, -1.6178711813269901, -1.4481040246501162, -9.628047542097333, -1.4859135403284893, 2.4996786477168658, 2.0661356872364505, 29.29125803278182, -1.6741239900815652, -2.291264005934315, -27.03018666229796, -1.1292483619710232, 1.3418076524051603, -2.2882006697246338, 23.078303770442247, 6.851874052669045, -3.2337825983991344, -5.2918543547805506, -1.8571939716476507],
        [19.393682461565486, 3.7691775260410947, -55.05475208276624, 42.11202847465953, -4.30170800452535, 10.455570562893309, -1.423554049042593, -22.845514925029146, 6.7854873442575485, 1.9311392844973874, -3.406814112103596]
        )
        model.params[:visual] = [1.0, len_scale]
    end

    if target == :co2
        model.funcs[:visual] = [Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 0, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 1, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 2, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 1, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 3, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 0, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 1, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 2, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 3, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(2, 3, 1, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 0, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 0, 2, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 3, 0, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 3, 3, :ahat), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(0, 3, 3, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 0, 1, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(1, 1, 0, :bhatorth2), Ψ2{LaguerreFunction{Float64}(), LaguerreFunction{Float64}(), LegendreP{Float64}()}(3, 0, 0, :bhatorth2)]
        model.coeffs[:visual] = vcat(
        [-39.563505410353486, -9.311926053131815, -1.0251077241944562, -1.4051674625111896, 11.839589512069683, -3.775510590106278, -2.305645376591694, -7.486152180561698, 120.53205753070903, 6.36056677489857, 3.2375029382577227, -5.827400073489639, -122.51054970551223, 9.463663710553451, -14.758220179371005, 0.5255283106744485, 3.3970311310452312, 10.791621285634644, 42.03259704739897, -3.337070191539953, 0.7980748466933036, 3.7103808472126425],
        [1.2229449926283504, 0.9822827479198946, 0.8021727670846446, -2.4670708959778227, 1.1549935390711488]
        )
        model.params[:visual] = [1.0, len_scale]
    end

    if target == :none
        model.islearned[:visual] = false
        model.islearned[:velpot] = true
        model.funcs[:velpot] = [Ψ1(f1, n, :a) for n in [0, 4, 5]];
        model.coeffs[:velpot]= [2.26539, -6.20361, 18.9581];
        model.params[:velpot] = [2.0]
    end

    model.bfields[:visual] = bfield2
    model.diffD = D

    return model

end

f_basis = LaguerreFunction{Float64}();
g_basis = LegendreP{Float64}();
none_model = define_model(; ndim=2, f2a=f_basis, f2b=f_basis, f2ab=g_basis, # bfield1 = distance_to_center!, 
        bfield2 = distance_to_center!, D=0.3, len_scale=SCALING_FACTOR, target=:none);
co2_model = define_model(; ndim=2, f2a=f_basis, f2b=f_basis, f2ab=g_basis, # bfield1 = distance_to_center!, 
        bfield2 = distance_to_center!, D=0.3, len_scale=SCALING_FACTOR, target=:co2);
visual_model = define_model(; ndim=2, f2a=f_basis, f2b=f_basis, f2ab=g_basis, # bfield1 = distance_to_center!, 
        bfield2 = distance_to_center!, D=0.3, len_scale=SCALING_FACTOR, target=:visual);
visualco2_model = define_model(; ndim=2, f2a=f_basis, f2b=f_basis, f2ab=g_basis, # bfield1 = distance_to_center!, 
        bfield2 = distance_to_center!, D=0.3, len_scale=SCALING_FACTOR, target=:visualco2);

get_deriv_none(u, center_coords) = get_rhs(u, 0.0, center_coords, none_model)
get_deriv_co2(u, center_coords) = get_rhs(u, 0.0, center_coords, co2_model)
get_deriv_visual(u, center_coords) = get_rhs(u, 0.0, center_coords, visual_model)
get_deriv_visualco2(u, center_coords) = get_rhs(u, 0.0, center_coords, visualco2_model)
get_w_none(u) = get_w(u, 0.0, none_model)
get_w_co2(u) = get_w(u, 0.0, co2_model)
get_w_visual(u) = get_w(u, 0.0, visual_model)
get_w_visualco2(u) = get_w(u, 0.0, visualco2_model)

function condition(u, t, integrator)
    # Check if outside box boundaries
    x, y = u[1], u[2]
    half_x = SCALING_FACTOR/2
    half_y = SCALING_FACTOR/2
    
    return abs(x) > half_x || abs(y) > half_y
end

function affect!(integrator)
    # Get current position and velocity
    u = integrator.u
    x, y = u[1], u[2]
    vx, vy = u[3], u[4]
    
    half_x = SCALING_FACTOR
    half_y = SCALING_FACTOR
    
    # Handle x boundaries
    if abs(x) > half_x
        # Move to boundary
        integrator.u[1] = sign(x) * half_x
        # Reflect velocity in x direction
        integrator.u[3] = -vx
    end
    
    # Handle y boundaries
    if abs(y) > half_y
        # Move to boundary
        integrator.u[2] = sign(y) * half_y
        # Reflect velocity in y direction
        integrator.u[4] = -vy
    end
    
    # Notify the integrator that we modified u
    u_modified!(integrator, true)
end

cb = DiscreteCallback(condition, affect!);

simulate_none(u0, tspan, tsave, center_coords, dts) = sde_sim(none_model, u0, tspan, tsave, center_coords, dts=dts, callback=cb)
simulate_co2(u0, tspan, tsave, center_coords, dts) = sde_sim(co2_model, u0, tspan, tsave, center_coords, dts=dts, callback=cb)
simulate_visual(u0, tspan, tsave, center_coords, dts) = sde_sim(visual_model, u0, tspan, tsave, center_coords, dts=dts, callback=cb)
simulate_visualco2(u0, tspan, tsave, center_coords, dts) = sde_sim(visualco2_model, u0, tspan, tsave, center_coords, dts=dts, callback=cb)

function sample_init(; r₀ = 0.5, σᵣ = 0.1, v₀ = 0.5, σᵥ = 0.2 )
    ρ =  r₀ + σᵣ * randn()
    ϕ =  rand()*2π
    θ =  rand()*2π
    v =  v₀ + σᵥ * randn()
    return [ρ*cos(ϕ), ρ*sin(ϕ), v*cos(θ), v*sin(θ)]
end

export CENTER_COORDS, SCALING_FACTOR, f1, get_deriv_none, get_deriv_co2, get_deriv_visual, get_deriv_visualco2, get_w_none, get_w_co2, get_w_visual, get_w_visualco2, sample_init, simulate_none, simulate_co2, simulate_visual, simulate_visualco2

end