raytrace.jl

struct ModSub
    n::Vector{Float64}
    x::Vector{Float64}
    y::Vector{Float64}
    v::Vector{Float64}
end

struct ModelConst
    u::ModSub
    A::ModSub
    t::ModSub
    error::Vector{Float64}
    BIC::Vector{Float64}
    fit::Vector{Float64}
    nfit::Vector{Float64}
    llh::Vector{Float64}
end

struct RayConst
    L::Vector{Float64}
    u::Vector{Float64}
    A::Vector{Float64}
    theta::Vector{Float64}
    phi::Vector{Float64}
    t::Vector{Float64}
    x::Vector{Float64}
    y::Vector{Float64}
    sta::Float64
    evt::Float64
    pick::Vector{Float64}
end


function raytracer!(ray::RayConst, model::ModelConst )

    np = length(ray.x)

    u_mat = Matrix{Float64}(undef, (length(model.u.x), length(model.u.y)))
    A_mat = Matrix{Float64}(undef, (length(model.A.x), length(model.A.y)))
    t_mat = Matrix{Float64}(undef, (length(model.t.x), length(model.t.y)))
    for k = 1:np
        v_dist!(u_mat, ray.x[k], ray.y[k], model.u)
        u = sum(model.u.v ./ u_mat) / sum(1 ./ u_mat)

        v_dist!(A_mat, ray.x[k], ray.y[k], model.A)
        A = sum(model.A.v ./ A_mat) / sum(1 ./ A_mat)

        v_dist!(t_mat, ray.x[k], ray.y[k], model.t)
        t = sum(model.t.v ./ t_mat) / sum(1 ./ t_mat)
        if k != np
            @inbounds ray.u[k] = u
            @inbounds ray.A[k] = A
            @inbounds ray.theta[k] = t
        end
        if k != 1
            @inbounds ray.u[k - 1] += u
            @inbounds ray.u[k - 1] /= 2
            @inbounds ray.A[k - 1] += A
            @inbounds ray.A[k - 1] /= 2
            @inbounds ray.theta[k - 1] += t
            @inbounds ray.theta[k - 1] /= 2
        end
    end

    dx      = diff(ray.x);
    dy      = diff(ray.y);

    @. ray.L   = sqrt(dx^2 + dy^2)
    @. ray.phi = atan(dy/dx)

    ray.t .= sum(@. ray.L * (ray.u * (1 + ray.A * cos(2*(ray.phi - ray.theta)))))

    return ray
end


model = ModelConst( ModSub(8*ones(1), 100*rand(Float64, 8), 100*rand(Float64, 8), rand(Float64, 8)),
    ModSub(8*ones(1), 100*rand(Float64, 8), 100*rand(Float64, 8), 0.1*rand(Float64, 8)),
    ModSub(8*ones(1), 100*rand(Float64, 8), 100*rand(Float64, 8), pi*(rand(Float64, 8).-1)),
    0.01*ones(1), ones(1), ones(1), ones(1), ones(1))

r = 1:20

ray = RayConst(ones(length(r) - 1), ones(length(r) - 1), ones(length(r) - 1), ones(length(r) - 1),
    ones(length(r) - 1), ones(1), range(0, 100, length=20), range(0, 100, length=20), 1, 1, ones(1))


function Base.isequal(a::RayConst, b::RayConst)
    for field in fieldnames(RayConst)
        if !isequal(getproperty(a,field), getproperty(b,field))
            @info field
            return false
        end
    end
    return true
end

## original - 25.494 μs (203 allocations: 115.08 KiB)

function raytracer!(ray::RayConst, model::ModelConst )
    
    np = length(ray.x)
    
    u = zeros(np)
    A = zeros(np)
    t = zeros(np)
    
    for k = 1:np
    
        ru = [ ((ray.x[k] - mx).^2) .+ ((ray.y[k] - my).^2) for mx in model.u.x, my in model.u.y]
        rA = [ ((ray.x[k] - mx).^2) .+ ((ray.y[k] - my).^2) for mx in model.A.x, my in model.A.y]
        rt = [ ((ray.x[k] - mx).^2) .+ ((ray.y[k] - my).^2) for mx in model.t.x, my in model.t.y]
    
        u[k] = sum((model.u.v ./ ru))./sum(1 ./ ru)
        A[k] = sum((model.A.v ./ rA))./sum(1 ./ rA)
        t[k] = sum((model.t.v ./ rt))./sum(1 ./ rt)
    
    end
    
    dx      = diff(ray.x);
    dy      = diff(ray.y);
    
    ray.L   .= sqrt.(dx.^2 .+ dy.^2)
    ray.phi .= atan.(dy./dx)
    
    ray.u     .= 0.5*(2*u[1:(length(u)-1)] .+ diff(u));
    ray.A     .= 0.5*(2*A[1:(length(A)-1)] .+ diff(A));
    ray.theta .= 0.5*(2*t[1:(length(t)-1)] .+ diff(t));
    
    ray.t .= sum(ray.L.*(ray.u.*(1 .+ ray.A.*cos.(2*(ray.phi .- ray.theta)))))
    
    return ray
    
end

rayt = deepcopy(ray); @btime raytracer!(rayt, model); truth = deepcopy(rayt);

# Option 1 - 20.179 μs (135 allocations: 77.77 KiB)
function v_dist!(M::Matrix, x::Float64, y::Float64, m::ModSub)
    nr,nc = size(M)
    @inbounds for i in 1:nr
        x_val = (m.x[i] - x)^2
        for j in 1:nc
            M[i,j] = x_val + (m.y[j] - y)^2
        end
    end
    M
end
function raytracer!(ray::RayConst, model::ModelConst )

    np = length(ray.x)

    u = Vector{Float64}(undef,np)
    A = Vector{Float64}(undef,np)
    t = Vector{Float64}(undef,np)

    u_mat = Matrix{Float64}(undef, (length(model.u.x), length(model.u.y)))
    A_mat = Matrix{Float64}(undef, (length(model.A.x), length(model.A.y)))
    t_mat = Matrix{Float64}(undef, (length(model.t.x), length(model.t.y)))
    @inbounds for k = 1:np
        v_dist!(u_mat, ray.x[k], ray.y[k], model.u)
        u[k] = sum(model.u.v ./ u_mat) / sum(1 ./ u_mat)

        v_dist!(A_mat, ray.x[k], ray.y[k], model.A)
        A[k] = sum(model.A.v ./ A_mat) / sum(1 ./ A_mat)

        v_dist!(t_mat, ray.x[k], ray.y[k], model.t)
        t[k] = sum(model.t.v ./ t_mat) / sum(1 ./ t_mat)
    end

    dx      = diff(ray.x);
    dy      = diff(ray.y);

    @. ray.L   = sqrt(dx^2 + dy^2)
    @. ray.phi = atan(dy/dx)

    du = diff(u)
    @. ray.u     = u[1:(end-1)] + du/2;
    dA = diff(A)
    @. ray.A     = A[1:(end-1)] + dA/2;
    dt = diff(t)
    @. ray.theta = t[1:(end-1)] + dt/2;

    ray.t .= sum(@. ray.L * (ray.u * (1 + ray.A * cos(2*(ray.phi - ray.theta)))))

    return ray
end

rayt = deepcopy(ray); @btime raytracer!(rayt, model); isequal(truth, rayt)

# Option 2 - 19.806 μs (126 allocations: 75.66 KiB)
function raytracer!(ray::RayConst, model::ModelConst )

    np = length(ray.x)

    u_mat = Matrix{Float64}(undef, (length(model.u.x), length(model.u.y)))
    A_mat = Matrix{Float64}(undef, (length(model.A.x), length(model.A.y)))
    t_mat = Matrix{Float64}(undef, (length(model.t.x), length(model.t.y)))
    @inbounds for k = 1:np
        v_dist!(u_mat, ray.x[k], ray.y[k], model.u)
        u = sum(model.u.v ./ u_mat) / sum(1 ./ u_mat)

        v_dist!(A_mat, ray.x[k], ray.y[k], model.A)
        A = sum(model.A.v ./ A_mat) / sum(1 ./ A_mat)

        v_dist!(t_mat, ray.x[k], ray.y[k], model.t)
        t = sum(model.t.v ./ t_mat) / sum(1 ./ t_mat)
        if k != np
            ray.u[k] = u
            ray.A[k] = A
            ray.theta[k] = t
        end
        if k != 1
            ray.u[k - 1] += u
            ray.u[k - 1] /= 2
            ray.A[k - 1] += A
            ray.A[k - 1] /= 2
            ray.theta[k - 1] += t
            ray.theta[k - 1] /= 2
        end
    end

    dx      = diff(ray.x);
    dy      = diff(ray.y);

    @. ray.L   = sqrt(dx^2 + dy^2)
    @. ray.phi = atan(dy/dx)

    ray.t .= sum(@. ray.L * (ray.u * (1 + ray.A * cos(2*(ray.phi - ray.theta)))))

    return ray
end

rayt = deepcopy(ray); @btime raytracer!(rayt, model); isequal(truth, rayt)

# Option 3 - 5.209 μs (3 allocations: 720 bytes)
# This version does slightly different math that makes strict vector equality 
# to the previous solutions no longer hold, but they are the same within +/- 1e-15
function v_dist(x::Float64, y::Float64, m::ModSub)
    v_sum = zero(Float64)
    inv_sum = zero(Float64)
    @inbounds for i in 1:length(m.x)
        x_val = (m.x[i] - x)^2
        for j in 1:length(m.y)
            distance = x_val + (m.y[j] - y)^2

            v_sum += m.v[i] / distance
            inv_sum += 1 / distance
        end
    end
    v_sum / inv_sum
end
function raytracer!(ray::RayConst, model::ModelConst )

    np = length(ray.x)

    @inbounds for k = 1:np
        u = v_dist(ray.x[k], ray.y[k], model.u)

        A = v_dist(ray.x[k], ray.y[k], model.A)

        t = v_dist(ray.x[k], ray.y[k], model.t)
        if k != np
            ray.u[k] = u
            ray.A[k] = A
            ray.theta[k] = t
        end
        if k != 1
            ray.u[k - 1] += u
            ray.u[k - 1] /= 2
            ray.A[k - 1] += A
            ray.A[k - 1] /= 2
            ray.theta[k - 1] += t
            ray.theta[k - 1] /= 2
        end
    end

    dx      = diff(ray.x);
    dy      = diff(ray.y);

    @. ray.L   = sqrt(dx^2 + dy^2)
    @. ray.phi = atan(dy/dx)

    ray.t .= sum(@. ray.L * (ray.u * (1 + ray.A * cos(2*(ray.phi - ray.theta)))))

    return ray
end

rayt = deepcopy(ray); @btime raytracer!(rayt, model); isequal(truth, rayt) # fails, but differences are tiny and I'm not sure which version is "more correct"