text:
################################################################################
#
#   COLLECTION OF FUNCTIONS FOR
#   BRILLOUIN ZONES
#
#   FILE CONTAINS
#   - mathematical helper functions for planes, lines and points
#   - construction of BZ parts (corners, faces, etc.)
#   - construction of BZ objects from reciprocal unitcells
#
################################################################################

################################################################################
#
#   MATHEMATICAL HELPER FUNCTIONS
#
################################################################################

# function to find out if two vectors are parallel
function parallelVectors(
            v1 :: Vector{<:Real},
            v2 :: Vector{<:Real}
        ) :: Bool

    # check if they can be scaled into one another (i.e. the rank of the matrix)
    r = rank([v1'; v2'])
    if r == 1
        return true
    elseif r == 2
        return false
    else
        error("Got " * string(r) * " from calculating the rank of two vectors:\n" *
        "v1 = " * string(v1) * ", and v2 = " * string(v2))
    end
end

# Find out if a point point_p is on a plane given by (plane_p,plane_n)
function pointOnPlane(
            point_p :: Vector{<:Real},
            plane_p :: Vector{<:Real},
            plane_n :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # check if the plane normal is perpendicular to the connecting vector between plane basis point and point
    return dot(plane_p .- point_p, plane_n) < error_allowed
end

# Find out if two planes (p1,n2) and (p2,n2) are identical
function identicalPlanes(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # check if p2 is on plane 1 and the normals are parallel
    return pointOnPlane(p2, p1,n1, error_allowed) && parallelVectors(n1,n2)
end

# Find out if three planes (pi,ni) have one intersecting point
function haveSingleIntersection(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            p3 :: Vector{<:Real},
            n3 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # calculate the determinant
    determinant = dot(n3, cross(n1, n2))
    # check if they intersect in a point
    if abs(determinant) < error_allowed
        # there will not be a single point of intersection
        return false
    else
        # they will intersect in one point
        return true
    end
end

# Get the intersecting point of three planes
function getSingleIntersection(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            p3 :: Vector{<:Real},
            n3 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Vector{Float64}

    # calculate the main determinant
    determinant = dot(n3, cross(n1, n2))

    # calculate the other relevant determinants
    D = Float64[-dot(n1,p1), -dot(n2,p2), -dot(n3,p3)]
    det_x = det([
        D[1] n1[2] n1[3]
        D[2] n2[2] n2[3]
        D[3] n3[2] n3[3]
    ])
    det_y = det([
        n1[1] D[1] n1[3]
        n2[1] D[2] n2[3]
        n3[1] D[3] n3[3]
    ])
    det_z = det([
        n1[1] n1[2] D[1]
        n2[1] n2[2] D[2]
        n3[1] n3[2] D[3]
    ])
    # get the interesection point
    return [det_x, det_y, det_z] ./ determinant
end

################################################################################
#
#   COMPUTING BRILLOUIN ZONE OBJECTS
#   (GENERAL FUNCTIONS)
#
################################################################################

# obtain the Brillouin zone for a given reciprocal unitcell (L=2,N=2)
function getBrillouinZone(
            ::Type{BZ},
            reciprocal_unitcell :: RU
            ;
            max_r :: Real = 1.001
        ) :: BZ where {L,P<:AbstractReciprocalPoint{2},B<:AbstractBond{L,2},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    ##########
    # STEP 1 - Compute the reciprocal lattice points that are needed for the calculation
    ##########

    # build a test lattice
    reciprocal_lattice = getLatticePeriodic(reciprocal_unitcell, 2)

    # find the minimal distance to a the Gamma point
    max_distance = maximum([dot(k,k) for k in point.(sites(reciprocal_lattice))])
    min_distance = minimum([dot(k,k) for k in point.(sites(reciprocal_lattice))])

    # reconstruct the reciprocal lattice
    reciprocal_lattice = getLatticeInSphere(reciprocal_unitcell, sqrt(max_distance)*max_r)

    # list of mid points
    mid_points = Vector{Float64}[
        k.*0.5  for k in point.(sites(reciprocal_lattice)) if sum(k.*k) > 1e-5
    ]

    # list of normals
    normals = Vector{Float64}[
        [k[2], -k[1]] ./ (k[1]*k[1] + k[2]*k[2]) for k in mid_points
    ]

    ##########
    # STEP 2 - Compute the corners of the BZ
    ##########

    ############
    # STEP 2.1 - Compute insersection points
    ############

    # list of intersection points
    intersections = Vector{Float64}[]

    # find the intersections
    for i in 1:length(mid_points)
    for j in 1:length(mid_points)
        # continue if same index
        if i==j
            continue
        end
        # get data for point 1
        p1 = mid_points[i]
        d1 = normals[i]
        # get data for point 2
        p2 = mid_points[j]
        d2 = normals[j]
        # check the interesection point
        delta_1 =  d1[2]*p1[1] - d1[1]*p1[2]
        delta_2 =  d2[2]*p2[1] - d2[1]*p2[2]
        delta   = -d1[2]*d2[1] + d2[2]*d1[1]
        # push to the intersections (if not devided by zero)
        if abs(delta) > 1e-8
            push!(intersections,[
                (d1[1]*delta_2 - d2[1]*delta_1) / delta,
                (d1[2]*delta_2 - d2[2]*delta_1) / delta
            ])
        end
    end
    end

    ############
    # STEP 2.2 - Find closest intersection points --> points of BZ
    ############

    # list of distances
    intersection_distances = Float64[
        sum(isec.*isec) for isec in intersections
    ]

    # find the minimal intersection distance
    min_distance = minimum(intersection_distances)

    # list of points which are only the minimal intersection distance away (=corners)
    corners_BZ = Array{Float64, 1}[]

    # check all points
    for i in 1:length(intersections)
        # check the relative deviation to the calculated minimal distance
        if intersection_distances[i] < min_distance * (1.0 + 1e-8)
            # check if not inserted already
            inserted_already = false
            for p in corners_BZ
                if sum((p.-intersections[i]).*(p.-intersections[i])) < 1e-8
                    inserted_already = true
                end
            end
            # found a new point
            if !inserted_already
                push!(corners_BZ, intersections[i])
            end
        end
    end

    ##########
    # STEP 3 - Compute the faces of the BZ (based on corners)
    ##########

    # list of points if they are contained in the loop
    in_loop = Bool[false for p in corners_BZ]

    # the list of the loop (to bring the points in correct order)
    loop = Int64[]

    # start with the first point
    push!(loop, 1)
    in_loop[1] = true

    # look for next point (which is closest to the given point in all remaing ones)
    while length(loop) < length(corners_BZ)
        # check for the closest
        closest_distance = sum(corners_BZ[1].*corners_BZ[1]) * 1000
        closest_point = -1
        # check all other points
        for i in 1:length(corners_BZ)
            # if already in loop, continue
            if in_loop[i]
                continue
            end
            # compare to the nearest one
            if sum((corners_BZ[i].-corners_BZ[loop[end]]).*(corners_BZ[i].-corners_BZ[loop[end]])) < closest_distance
                closest_point    = i
                closest_distance = sum((corners_BZ[i].-corners_BZ[loop[end]]).*(corners_BZ[i].-corners_BZ[loop[end]]))
            end
        end
        # insert the closest point as the next in line
        push!(loop, closest_point)
        in_loop[closest_point] = true
    end

    # push the starting point into the loop to wrap up and make a closed line
    #push!(loop, loop[1])

    # list of edges (only contains this one loop)
    faces_BZ = Vector{Int64}[loop]

    ##########
    # STEP 4 - Construct a new object and return it
    ##########

    # return the finished BZ
    return newBrillouinZone(
        BZ,
        reciprocal_unitcell,
        corners_BZ,
        faces_BZ
    )
end

# obtain the Brillouin zone for a given reciprocal unitcell (L=3,N=3)
function getBrillouinZone(
            ::Type{BZ},
            reciprocal_unitcell :: RU
            ;
            max_r :: Real = 1.001
        ) :: BZ where {L,P<:AbstractReciprocalPoint{3},B<:AbstractBond{L,3},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    ##########
    # STEP 1 - Compute the reciprocal lattice points that are needed for the calculation
    ##########

    # build a test lattice
    reciprocal_lattice = getLatticePeriodic(reciprocal_unitcell, 2)

    # find the minimal distance to a the Gamma point
    max_distance = maximum([dot(k,k) for k in point.(sites(reciprocal_lattice))])
    min_distance = minimum([dot(k,k) for k in point.(sites(reciprocal_lattice))])

    # reconstruct the reciprocal lattice
    reciprocal_lattice = getLatticeInSphere(reciprocal_unitcell, sqrt(max_distance)*max_r)

    # k points
    k_points = point.(sites(reciprocal_lattice))

    # list of mid points
    mid_points = Vector{Float64}[
        k.*0.5  for k in point.(sites(reciprocal_lattice)) if sum(k.*k) > 1e-5
    ]

    # list of normals of planes
    normals = [m ./ dot(m,m) for m in mid_points]

    ##########
    # STEP 2 - Find all intersection points of 3 planes and check that they are closest to the Gamma point
    ##########

    # list of points which are only the minimal intersection distance to Gamma point
    points = Array{Float64, 1}[]

    # temp point
    point_tmp = zeros(Float64, 3)

    # lists for all planes that aided in the construction
    planes_n = Array{Float64,1}[]
    planes_p = Array{Float64,1}[]

    # find the intersections
    for i in 1:length(mid_points)
        # get data for plane 1
        p1 = mid_points[i]
        n1 = normals[i]
        # get a second plane
        for j in i+1:length(mid_points)
            # get data for plane 2
            p2 = mid_points[j]
            n2 = normals[j]
            # continue if they are parallel
            if dot(cross(n1,n2), cross(n1,n2)) < 1e-8
                continue
            end
            # get a third plane
            for l in j+1:length(mid_points)
                # get data for plane 3
                p3 = mid_points[l]
                n3 = normals[l]
                # calculate the determinant
                determinant = dot(n3, cross(n1, n2))
                # check if they intersect in a point
                if abs(determinant) < 1e-8
                    # there will not be one point of intersection
                    continue
                end
                # calculate the other relevant determinants
                #A = Float64[n1[1], n2[1], n3[1]]
                #B = Float64[n1[2], n2[2], n3[2]]
                #C = Float64[n1[3], n2[3], n3[3]]
                D = Float64[-dot(n1,p1), -dot(n2,p2), -dot(n3,p3)]
                #det_x = dot(D, cross(B, C))
                #det_y = dot(A, cross(D, C))
                #det_z = dot(A, cross(B, D))
                det_x = det([
                    D[1] n1[2] n1[3]
                    D[2] n2[2] n2[3]
                    D[3] n3[2] n3[3]
                ])
                det_y = det([
                    n1[1] D[1] n1[3]
                    n2[1] D[2] n2[3]
                    n3[1] D[3] n3[3]
                ])
                det_z = det([
                    n1[1] n1[2] D[1]
                    n2[1] n2[2] D[2]
                    n3[1] n3[2] D[3]
                ])
                # get the interesection point
                point_tmp[1] = det_x / determinant
                point_tmp[2] = det_y / determinant
                point_tmp[3] = det_z / determinant
                # check if the distance to any other k point in the lattice is shorter than to the Gamma point
                shorter_distance_found = false
                distance_to_Gamma = dot(point_tmp, point_tmp)
                # check all points
                for k in k_points
                    if dot(point_tmp .- k, point_tmp .- k)*(1 + 1e-8) < distance_to_Gamma
                        shorter_distance_found = true
                        break
                    end
                end
                # insert point into list if no shorter distance found
                if !shorter_distance_found
                    # check if not inserted already
                    inserted_already = false
                    # check all points
                    for p in points
                        if dot(p.-point_tmp, p.-point_tmp) < 1e-6
                            inserted_already = true
                            break
                        end
                    end
                    # if not inserted already, insert now
                    if !inserted_already
                        # insert the point into the list
                        push!(points, copy(point_tmp))
                        # see if the planes that cross here, have to be inserted as well

                        # check plane 1
                        insert_p1 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p1 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n1), cross(planes_n[index], n1)) < 1e-6
                                    # normals also match in direction
                                    insert_p1 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p1
                            push!(planes_p, p1)
                            push!(planes_n, n1)
                        end

                        # check plane 2
                        insert_p2 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p2 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n2), cross(planes_n[index], n2)) < 1e-6
                                    # normals also match in direction
                                    insert_p2 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p2
                            push!(planes_p, p2)
                            push!(planes_n, n2)
                        end

                        # check plane 3
                        insert_p3 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p3 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n3), cross(planes_n[index], n3)) < 1e-6
                                    # normals also match in direction
                                    insert_p3 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p3
                            push!(planes_p, p3)
                            push!(planes_n, n3)
                        end

                    end
                end
            end
        end
    end

    ##########
    # STEP 5 - Connect the points in planes of the BZ --> edges of BZ
    ##########

    # the list of all the loops (i.e. all edge loops of the BZ)
    edges        = Array{Int64,1}[]
    # the list of all the loops but sorted without double point for the start+end
    edges_sorted = Array{Int64,1}[]

    # find all loops of all points by constructing planes

    # iterate over all planes
    for i in 1:length(planes_n)

        # find the subset of points that lie within the plane
        in_plane = Bool[abs(dot(p .- planes_p[i], planes_n[i])) < 1e-6 for p in points]

        # use the 2D loop finding within the plane
        point_indices = collect(1:length(points))[in_plane]

        # check if there are enough points to form a loop
        if length(point_indices) < 3
            # not enough points found
            continue
        end

        # list of points if they are contained in the loop
        in_loop = [false for i in point_indices]

        # the list of the loop (to bring the points in correct order)
        # with relative indices (have to be translated back)
        loop = Int64[]

        # start with the first (relative) point
        push!(loop, 1)
        in_loop[1] = true

        # look for next point (which is closest to the given point in all remaing ones)
        while length(loop) < length(point_indices)
            # check for the closest
            closest_distance = 1e32
            closest_point = -1
            # check all other points
            for i in 1:length(point_indices)
                # if already in loop, continue
                if in_loop[i]
                    continue
                end
                # compare to the nearest one
                if dot(points[point_indices[i]].-points[point_indices[loop[end]]], points[point_indices[i]].-points[point_indices[loop[end]]]) < closest_distance
                    closest_point    = i
                    closest_distance = dot(points[point_indices[i]].-points[point_indices[loop[end]]], points[point_indices[i]].-points[point_indices[loop[end]]])
                end
            end
            # security check
            if closest_point == -1
                println()
                println(closest_distance)
                println(points[in_plane])
            end
            # insert the closest point as the next in line
            push!(loop, closest_point)
            in_loop[closest_point] = true
        end

        # get back the real indices and reconstruct loop
        loop = point_indices[loop]
        # get the sorted
        loop_sorted = sort(loop)
        # push the starting point into the loop
        #push!(loop, loop[1])

        # loop not found yet
        loop_found = false
        # check if the loop already exists
        for l in edges_sorted
            if l==loop_sorted
                loop_found = true
                break
            end
        end
        # if the loop is not found, insert the loop into the edges list
        if !loop_found
            push!(edges, loop)
        end

    end

    ##########
    # STEP 4 - Construct a new object and return it
    ##########

    # return the finished BZ
    return newBrillouinZone(
        BZ,
        reciprocal_unitcell,
        points,
        edges
    )
end

################################################################################
#
#   COMPUTING BRILLOUIN ZONE OBJECTS
#   (Wrapper)
#
################################################################################

# wrapper for all reciprocal unitcells
function getBrillouinZone(
            reciprocal_unitcell :: RU
            ;
            kwargs...
        ) :: BrillouinZone{RU} where {L,N,D,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B}}

    # call the respective function
    return getBrillouinZone(BrillouinZone{RU}, reciprocal_unitcell; kwargs...)
end

# wrapper for all real space unitcells
function getBrillouinZone(
            unitcell :: U
            ;
            kwargs...
        ) :: BrillouinZone{ReciprocalUnitcell{ReciprocalPoint{D},B}} where {L,D,S<:AbstractSite{L,D},B,U<:AbstractUnitcell{S,B}}

    # call the respective function
    return getBrillouinZone(BrillouinZone{ReciprocalUnitcell{ReciprocalPoint{D},B}}, getReciprocalUnitcell(unitcell); kwargs...)
end

# Export the function
export getBrillouinZone

################################################################################
#
#   FUNCTIONS FOR HIGH SYMMETRY MESH
#
################################################################################

# contructor for symmetric mesh of given brillouin zone
function getPointsHighSymmetryMesh(
            bz              :: BZ,
            line_resolution :: Integer
            ;
            include_bz_corners :: Bool           = true,
            include_face_centers :: Bool         = true,
            include_bz_edges :: Bool             = true,  # edges of faces (i.e. lines connecting the corners)
            include_lines_in_faces :: Bool       = true,
            include_lines_to_face_center :: Bool = true,
            include_lines_to_gamma :: Bool       = true,
            include_random_face_points :: Bool   = true,
            face_resolution :: Integer = 100
        ) :: Vector{Vector{Float64}} where {D,B, P<:AbstractReciprocalPoint{D}, R<:AbstractReciprocalUnitcell{P,B}, BZ<:AbstractBrillouinZone{R}}

    # get the gamma point
    gamma = point(gammaPoint(reciprocalUnitcell(bz)))
    # the list of points that form the mesh
    mesh = Vector{Float64}[]

    # iterate over all faces of the BZ
    for f in faces(bz)

        # obtain the relevant corners
        k_corners = unique([corners(bz)[i] for i in f])
        # points that connect to gamma
        points_to_gamma = Vector{Float64}[]

        # INCLUDE CORNER POINTS IN MESH
        if include_bz_corners
            for c in k_corners
                # put the corner into the mesh
                push!(mesh, c)
                # put the corner into the list of points that connect to the gamma point
                push!(points_to_gamma, c)
            end
        end

        # INCLUDE FACE CENTERS IN MESH
        if include_face_centers
            # push the center into the mesh
            push!(mesh, getPointCenter(k_corners))
            # push the center into the list of connecting points
            push!(points_to_gamma, getPointCenter(k_corners))
        end

        # INCLUDE EDGES IN MESH
        if include_bz_edges
            # the edge that wraps in the point list
            line = getPointsOnLine(k_corners[end], k_corners[1], line_resolution)
            for p in line
                # put the point into the mesh
                push!(mesh, p)
            end
            # put the start, end and center of line into the gamma point connecting point list
            push!(points_to_gamma, k_corners[end])
            push!(points_to_gamma, k_corners[1])
            push!(points_to_gamma, getPointCenter(k_corners[1], k_corners[end]))
            # all other edges
            for i in 1:length(k_corners)-1
                # construct line
                line = getPointsOnLine(k_corners[i], k_corners[i+1], line_resolution)
                for p in line
                    # put the point into the mesh
                    push!(mesh, p)
                end
                # put the start, end and center of line into the gamma point connecting point list
                push!(points_to_gamma, k_corners[i])
                push!(points_to_gamma, k_corners[i+1])
                push!(points_to_gamma, getPointCenter(k_corners[i], k_corners[i+1]))
            end
        end

        # INCLUDE LINES IN THE FACE
        if include_lines_in_faces
            # iterate over all combination of points
            for i in 1:length(k_corners)
            for j in i+1:length(k_corners)
                # build the line from i to j
                line = getPointsOnLine(k_corners[i], k_corners[j], line_resolution)
                for p in line
                    push!(mesh, p)
                end
                # put the start, end and center of line into the gamma point connecting point list
                push!(points_to_gamma, k_corners[i])
                push!(points_to_gamma, k_corners[j])
                push!(points_to_gamma, getPointCenter(k_corners[i], k_corners[j]))
            end
            end
        end

        # INCLUDE LINES TO THE FACE CENTER
        if include_lines_to_face_center
            # build the face center
            center = getPointCenter(k_corners)
            # iterate over all points
            for i in 1:length(k_corners)
                # build the line from i to j
                line = getPointsOnLine(k_corners[i], center, line_resolution)
                for p in line
                    push!(mesh, p)
                end
            end
        end

        # INCLUDE LINES CONNECTING TO GAMMA POINT
        if include_lines_to_gamma
            # get a unique list of connecting points
            unique!(points_to_gamma)
            # build lines from these points to the gamma point (if they are not the gamma point)
            for ptg in points_to_gamma
                if dot(ptg.-gamma, ptg.-gamma) < 1e-8
                    continue
                end
                line = getPointsOnLine(ptg, gamma, line_resolution)
                for p in line
                    push!(mesh, p)
                end
            end
        end

        # INCLUDE FACE POINTS RANDOMLY SAMPLED
        if include_random_face_points
            points = getPointsRandomInConvexHull(face_resolution, k_corners...)
            for p in points
                push!(mesh, p)
            end
        end

    end

    # remove redundant points
    unique!(mesh)

    # return the mesh
    return mesh
end

export getPointsHighSymmetryMesh


text:
using Flux
using Flux: Chain, Dense, params, train!, glorot_uniform
using Flux.Tracker: Params, gradient, update!

f = x -> (1/(0.82+x))^12 -2 * (1/(.82+x))^6
#f = x -> sin.(x.*7.0)

nsamples = 200
XS = rand(nsamples)
YS = f.(XS)
#samples = zip(XS,YS)
#samples = [(x, f(x)) for x in XS]
#samples= [(XS, YS)]

xerr = range(0, 1, length=1000)
eval_nn(model, x) = model(x).data
maxfit_nn(model) = maximum(abs, f.(XS) - model(XS))

function train_nn(nnodes, η, samples, XX, YY)
    model = Chain(Dense(nnodes, nnodes, Flux.tanh),   # first hidden layer
                  Dense(nnodes, 1))      # output layer

    loss(x, y) = Flux.mse(model(x), y)

    θ = Flux.params(model)
    #println(θ)

    for i=1:1000
        #Flux.train!(loss, θ, samples, Descent(η))
        Flux.train!(loss, θ, repeated((samples,), 2_000), ADAM(0.05))
        if i%100==1
            @show loss(XX, YY)
            @show maxfit_nn(model)
        end
    end
    return model
end

η = 0.00001 # Learning rate (should be smaller corresponding to the complex regression times the # of nodes)
NN_nodes = 3:5
errs_nn = Float64[]

for N in NN_nodes
    XX = repeat(XS,1,N)
    XX = XX'
    samples = [([x for i=1:N], f(x)) for x in XS]
    @show N
    nn = train_nn(N,η,samples,XX,YS)
    #println(Flux.params(nn))
    #@show e = err_nn(nn)
    #push!(errs_nn, e)
end
@show errs_nn

ext: jl

max_stars_repo_name: berkonat/ann-examples

max_issues_repo_name: berkonat/ann-examples

max_forks_repo_name: berkonat/ann-examples


text:
using Flux
using FluxGAN
using BSON: @save
using MLDatasets

const n = parse(Int, ARGS[1])
const m = parse(Int, ARGS[2])
const skip = parse(Int, ARGS[3])

# dataset
train_tensor, train_labels = CIFAR10.traindata(Float32)

animal_classnames = [
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse"
]

classnames = CIFAR10.classnames()

# removing any images that aren't animals
animal_indices = findall(lab -> classnames[lab + 1] in animal_classnames, train_labels)
animal_tensor = train_tensor[:,:,:,animal_indices]

const img_size = size(animal_tensor)[1:end-1]
const img_num  = size(animal_tensor)[end]
const img_dim  = *(img_size...)

# centering pixels
animal_tensor = @. 2f0 * animal_tensor - 1f0

const d = 100

generator = Chain(
    Dense(d, 8000, relu),
    Dense(8000, 8000, tanh_fast),
    x -> reshape(x, 10, 10, 80, size(x)[end]),
    ConvTranspose((5, 5), 80 => 3, tanh_fast; stride=(3,3))
)

discriminator = Chain(
    Maxout(() -> Conv((8, 8), 3 => 32; pad=4), 2),
    MeanPool((4, 4); stride=(2, 2)),
    Maxout(() -> Conv((8, 8), 32 => 32, pad=3), 2),
    MeanPool((4, 4); stride=(2, 2)),
    Maxout(() -> Conv((5, 5), 32 => 192, pad=3), 2),
    MeanPool((2, 2); stride=(2, 2)),
    x -> reshape(x, 3072, size(x)[end]),
    Maxout(() -> Dense(3072, 500), 5),
    Dense(500, 1)
)

# learning rate
const η = 0.0001f0

model = GAN(generator, discriminator;
    η_gen=η,
    η_dscr=η,
    minibatch=m,
    latent_dim=d,
    img_size=img_size
)

println("beginning training...\n")

train!(model, animal_tensor, skip=skip, iterations=n,
    gif=true,
    gif_filename="CIFAR10/conv",
    gif_fps=5
)

println("\nfinished!\n")

max_stars_repo_name: aarontrowbridge/GAN

max_issues_repo_name: aarontrowbridge/GAN

max_forks_repo_name: aarontrowbridge/GAN


text:
struct BlockSparseDiagonal{T,Diag<:Diagonal{T},Blocks} <: AbstractMatrix{T}
    diag::Diag
    blocks::Blocks
    block_starts::Vector{Int}
end

Base.size(A::BlockSparseDiagonal, args...) = size(A.diag, args...)

function find_block(A::BlockSparseDiagonal, i, j)
    k = findlast(≤(i), A.block_starts)
    isnothing(k) && return nothing
    bs = A.block_starts[k]
    interval = bs:bs+size(A.blocks[k],1)-1
    i ∈ interval && j ∈ interval || return nothing
    k
end

function Base.getindex(A::BlockSparseDiagonal{T}, i, j) where T
    v = zero(T)
    b = find_block(A, i, j)
    if !isnothing(b)
        bs = A.block_starts[b]
        v += A.blocks[b][i-bs+1,j-bs+1]
    end
    if i == j
        v += A.diag.diag[i]
    end
    v
end

Base.replace_in_print_matrix(A::BlockSparseDiagonal,i::Integer,j::Integer,s::AbstractString) =
    (i == j || !isnothing(find_block(A, i, j))
     ? s
     : Base.replace_with_centered_mark(s))

Base.:(*)(α::Number, A::BlockSparseDiagonal) =
    BlockSparseDiagonal(α*A.diag, α .* A.blocks, A.block_starts)

Base.:(*)(A::BlockSparseDiagonal, α::Number) =
    BlockSparseDiagonal(α*A.diag, α .* A.blocks, A.block_starts)

Base.:(/)(A::BlockSparseDiagonal, α::Number) = A * inv(α)

function LinearAlgebra.mul!(y::AbstractMatrix, A::BlockSparseDiagonal, x::AbstractMatrix,
                            α::Number=true, β::Number=false)
    mul!(y, A.diag, x, α, β)
    for (b,bs) in zip(A.blocks, A.block_starts)
        interval = bs:bs+size(b,2)-1
        mul!(view(y, interval, :), b, view(x, interval, :), α, true)
    end
    y
end

function LinearAlgebra.mul!(y::AbstractVector, A::BlockSparseDiagonal, x::AbstractVector,
                            α::Number=true, β::Number=false)
    mul!(y, A.diag, x, α, β)
    for (b,bs) in zip(A.blocks, A.block_starts)
        interval = bs:bs+size(b,2)-1
        mul!(view(y, interval), b, view(x, interval), α, true)
    end
    y
end

export BlockSparseDiagonal


text:
using Test
import ArchGDAL as AG

@testset "test_gdal_tutorials.jl" begin
    @testset "Raster Tutorial" begin
        driver = AG.getdriver("GTiff")

        AG.read("data/utmsmall.tif") do dataset
            @test AG.shortname(driver) == "GTiff"
            @test AG.longname(driver) == "GeoTIFF"
            @test AG.width(dataset) == 100
            @test AG.height(dataset) == 100
            @test AG.nraster(dataset) == 1

            nad27_prefix = "PROJCS[\"NAD27 / UTM zone 11N\",GEOGCS[\"NAD27\",DATUM[\"North_American_Datum_1927\","
            @test startswith(AG.getproj(dataset), nad27_prefix) == true
            @test AG.getgeotransform(dataset) ≈
                  [440720.0, 60.0, 0.0, 3.75132e6, 0.0, -60.0]

            band = AG.getband(dataset, 1)
            @test AG.blocksize(band) ≈ [100, 81]
            @test AG.pixeltype(band) == UInt8
            @test AG.getname(AG.getcolorinterp(band)) == "Gray"

            @test AG.minimum(band) ≈ 0.0
            @test AG.maximum(band) ≈ 255.0

            @test AG.noverview(band) == 0

            # Reading Raster Data
            @test AG.width(band) == 100
            data = map(Float32, AG.read(dataset, 1))
            #! format: off
            @test data[:, 1] ≈ Float32[107.0f0, 123.0f0, 132.0f0, 115.0f0, 132.0f0, 132.0f0, 140.0f0, 132.0f0, 132.0f0, 132.0f0, 107.0f0, 132.0f0, 107.0f0, 132.0f0, 132.0f0, 107.0f0, 123.0f0, 115.0f0, 156.0f0, 148.0f0, 107.0f0, 132.0f0, 107.0f0, 115.0f0, 99.0f0, 123.0f0, 99.0f0, 74.0f0, 115.0f0, 82.0f0, 115.0f0, 115.0f0, 107.0f0, 123.0f0, 123.0f0, 99.0f0, 123.0f0, 123.0f0, 115.0f0, 115.0f0, 107.0f0, 90.0f0, 99.0f0, 107.0f0, 107.0f0, 99.0f0, 123.0f0, 107.0f0, 140.0f0, 123.0f0, 123.0f0, 115.0f0, 99.0f0, 132.0f0, 123.0f0, 115.0f0, 115.0f0, 123.0f0, 132.0f0, 115.0f0, 123.0f0, 132.0f0, 214.0f0, 156.0f0, 165.0f0, 148.0f0, 115.0f0, 148.0f0, 156.0f0, 148.0f0, 140.0f0, 165.0f0, 156.0f0, 197.0f0, 156.0f0, 197.0f0, 140.0f0, 173.0f0, 156.0f0, 165.0f0, 148.0f0, 156.0f0, 206.0f0, 214.0f0, 181.0f0, 206.0f0, 173.0f0, 222.0f0, 206.0f0, 255.0f0, 214.0f0, 173.0f0, 214.0f0, 255.0f0, 214.0f0, 247.0f0, 255.0f0, 230.0f0, 206.0f0, 197.0f0]
            #! format: on

            @test AG.metadatadomainlist(dataset) ==
                  ["IMAGE_STRUCTURE", "", "DERIVED_SUBDATASETS"]
            @test AG.metadata(dataset) == ["AREA_OR_POINT=Area"]
            @test AG.metadataitem(dataset, "AREA_OR_POINT") == "Area"
            @test AG.metadata(dataset, domain = "IMAGE_STRUCTURE") ==
                  ["INTERLEAVE=BAND"]
            @test AG.metadata(dataset, domain = "") == ["AREA_OR_POINT=Area"]
            @test AG.metadata(dataset, domain = "DERIVED_SUBDATASETS") == [
                "DERIVED_SUBDATASET_1_NAME=DERIVED_SUBDATASET:LOGAMPLITUDE:data/utmsmall.tif",
                "DERIVED_SUBDATASET_1_DESC=log10 of amplitude of input bands from data/utmsmall.tif",
            ]
        end

        # Get metadata from a RasterDataset
        AG.readraster("data/utmsmall.tif") do dataset
            # interestingly the list order below is different from the order above
            @test AG.metadatadomainlist(dataset) ==
                  ["IMAGE_STRUCTURE", "DERIVED_SUBDATASETS", ""]
            @test AG.metadata(dataset) == ["AREA_OR_POINT=Area"]
            @test AG.metadataitem(dataset, "AREA_OR_POINT") == "Area"
        end

        # Techniques for Creating Files
        @test AG.metadataitem(driver, "DCAP_CREATE", domain = "") == "YES"
        @test AG.metadataitem(driver, "DCAP_CREATECOPY", domain = "") == "YES"

        AG.read("data/utmsmall.tif") do ds_src
            AG.write(ds_src, "/vsimem/utmsmall.tif")
            AG.read("/vsimem/utmsmall.tif") do ds_copy
                @test AG.read(ds_src) == AG.read(ds_copy)
            end
        end
    end

    @testset "Vector Tutorial" begin
        AG.read("data/point.geojson") do dataset
            @test AG.nlayer(dataset) == 1
            layer = AG.getlayer(dataset, 0)
            @test (AG.getname(layer) in ["point", "OGRGeoJSON"]) == true
            layerbyname = AG.getlayer(dataset, "point")
            @test layerbyname.ptr == layer.ptr
            AG.resetreading!(layer)

            featuredefn = AG.layerdefn(layer)
            @test AG.nfield(featuredefn) == 2
            fielddefn = AG.getfielddefn(featuredefn, 0)
            @test AG.gettype(fielddefn) == AG.OFTReal
            fielddefn = AG.getfielddefn(featuredefn, 1)
            @test AG.gettype(fielddefn) == AG.OFTString

            AG.nextfeature(layer) do feature
                @test AG.asdouble(feature, 0) ≈ 2.0
                @test AG.asstring(feature, 1) == "point-a"
            end
            AG.nextfeature(layer) do feature # second feature
                @test AG.asdouble(feature, 0) ≈ 3.0
                @test AG.asstring(feature, 1) == "point-b"

                geometry = AG.getgeom(feature)
                @test AG.geomname(geometry) == "POINT"
                @test AG.getgeomtype(geometry) == AG.wkbPoint
                @test AG.nfield(featuredefn) == 2
                @test AG.getx(geometry, 0) ≈ 100.2785
                @test AG.gety(geometry, 0) ≈ 0.0893
                @test AG.getpoint(geometry, 0) == (100.2785, 0.0893, 0.0)
            end
        end

        AG.create(AG.getdriver("MEMORY")) do dataset
            layer = AG.createlayer(
                name = "point_out",
                dataset = dataset,
                geom = AG.wkbPoint,
            )
            AG.addfielddefn!(layer, "Name", AG.OFTString, nwidth = 32)
            featuredefn = AG.layerdefn(layer)
            @test AG.getname(featuredefn) == "point_out"
            @test AG.nfeature(layer) == 0
            AG.createfeature(layer) do feature
                AG.setfield!(
                    feature,
                    AG.findfieldindex(feature, "Name"),
                    "myname",
                )
                AG.setgeom!(feature, AG.createpoint(100.123, 0.123))
                return nothing
            end
            @test AG.nfeature(layer) == 1
        end
    end
end


text:
##  Copyright 2021 Myrl Marmarelis
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.

using StatsBase
using StaticArrays # another package to watch is TupleVectors
using PDMats
using Distributions
using LinearAlgebra
using Random: rand!

draw_cohort(mixture::VoterMixture) =
  mixture.cohorts[ Categorical(mixture.shares) |> rand ]

function draw_marginal(mixture::VoterMixture{N,M,T}, candidate::Int)::T where {N,M,T}
  @assert 1 <= candidate <= N
  cohort = draw_cohort(mixture)
  means = mean(cohort)
  covariances = cov(cohort)
  location = means[candidate]
  variance = covariances[candidate, candidate]
  sqrt(variance)*randn() + location
end

function sample_mixture_shares(voters::Vector{VoterRealization{N,M,T}},
    dirichlet_weights::AbstractVector{T}, voter_weights::Vector{T}=T[]
    )::SVector{M,T} where {N,M,T}
  n_voters = length(voters)
  @assert length(voter_weights) in (0, n_voters)
  @assert length(dirichlet_weights) == M
  counts = @MVector zeros(T, M)
  for (voter_index, voter) in enumerate(voters)
    cohort = voter.membership
    weight = T(1)
    if length(voter_weights) > 0
      weight = voter_weights[voter_index]
    end
    counts[cohort] += weight
  end
  #count_map = countmap(voter.membership for voter in voters)
  #counts = SVector{M,T}(get(count_map, cohort, 0) for cohort in 1:M)
  posterior_weights = counts + dirichlet_weights
  Dirichlet(posterior_weights) |> rand
end

function weighted_mean(f::Function, iterand, weights)
  sum(zip(iterand, weights)) do (value, weight)
    weight * f(value)
  end / sum(weights)
end

# eventually, the best way to improve performance would be to avoid the gradual
# memory fragmentation due to many small matrix allocations. the "purest" approach
# would be to get rid of the crutch that is our dependency on `MvNormal` algotegher.
# retrofit the library procedures entirely in terms of static matrices.
# first step: relieve most of the memory burden/footprint via the persistent objects.
function sample_mixture_posteriors(voters::Vector{VoterRealization{N,M,T}};
    precision_scale::AbstractPDMat{T}, precision_dof::T, mean_loc::AbstractVector{T},
    mean_scale::T, voter_weights::Vector{T}=T[])::SVector{M} where {N,M,T} # VC<:VoterCohort{N,T}
  n_voters = length(voters) # weight for recency (online prediction) or inverse propensity scoring to adjust for bias
  @assert length(voter_weights) in (0, n_voters)
  SVector{M}( begin
    cohort_voter_indices = Iterators.filter(1:n_voters) do voter_index
      voters[voter_index].membership == cohort
    end
    cohort_voters = (voters[voter_index] for voter_index in cohort_voter_indices)
    cohort_weights = Iterators.map(cohort_voter_indices) do voter_index
      length(voter_weights) > 0 ? voter_weights[voter_index] : T(1)
    end
    cohort_utilities = (voter.utility for voter in cohort_voters)
    n_samples = sum(cohort_weights, init=T(0))
    utility_covariance = zero(SMatrix{N,N,T,N*N})
    utility_means = SVector{N,T}(mean_loc)
    #error("what to do with an empty cohort?")
    #make it draw from the prior, basically
    if n_samples > 0
      utility_covariance = weighted_mean(
          cohort_utilities, cohort_weights) do utility
        utility * utility'
      end
      utility_means = weighted_mean(identity,
        cohort_utilities, cohort_weights)
    end
    # to be proper, I would actually invert `precision_scale`
    precision_prior_scale = precision_scale + n_samples*(
      utility_covariance + mean_scale/(mean_scale+n_samples)
        * (utility_means-mean_loc) * (utility_means-mean_loc)' )
    # alongside that Bayesreg.pdf, see https://en.wikipedia.org/wiki/Normal-Wishart_distribution
    # it's rather difficult to find reliable sources on this..
    precision_prior_dof = precision_dof + n_samples
    precision_dist = Wishart(precision_prior_dof,
      inv(precision_prior_scale) |> Symmetric |> cholesky) # not the most resourceful sequence of operations?
    precision_sample = rand(precision_dist) .|> T
    covariance_sample = inv(precision_sample)
    mean_means = (mean_loc*mean_scale + n_samples*utility_means) / (n_samples+mean_scale)
    mean_covariance = covariance_sample / (n_samples+mean_scale)
    mean_dist = MvNormal(mean_means, Symmetric(mean_covariance)) # Hermitian simply to counter numerical instability post-inversion
    mean_sample = rand(mean_dist) |> SVector{N,T}
    covariance_decomp = cholesky(covariance_sample |> SMatrix{N,N,T,N*N} |> Symmetric)
    MvNormal(mean_sample, PDMat(covariance_decomp))
  end for cohort in 1:M )
end

# allow zeros to demarcate absent choices (abstained). first element MUST be nonzero.
# question. say there are five candidates and I fill out
#   3, _, _, _, 2.
# does that actually imply that #2 yields lower utility than the remaining unnamed candidates?
# this special case does not lead to any actionable generalization; it falls apart immediately
# say if R<N or I fill in the fourth entry rather than the fifth.
function obeys_ranking(utility::Utility{N,T}, # AbstractVector for ranking?
    ranking::RankedChoice{R}, indif::Bool,
    interval::UnitRange{Int}=1:N)::Bool where {N,R,T}
  @assert ranking[1] > 0
  # don't take the time to check that all elements are within [1,N]
  seen = @MVector fill(false, N)
  # passing in a view from the get-go is costlier because the size is dynamic
  offset = interval.start - 1
  last_index = 1
  seen[ranking[1] + offset] = true # first item is always seen
  @inbounds for rank_index in 2:R # pairwise comparisons down the line
    if ranking[rank_index] == 0
      continue
    end
    first = ranking[last_index] + offset
    second = ranking[rank_index] + offset
    #ref_index = min(rank_index-1, R)
    if utility[second] >= utility[first]
      return false
    end
    seen[first] = seen[second] = true # logically separate these housekeeping lines from the above
    last_index = rank_index
  end
  # valid so far
  if indif
    return true
  end
  last_ranking = ranking[last_index]
  last_ranked = utility[last_ranking]
  @inbounds for candidate in interval
    if seen[candidate]
      continue
    end
    if utility[candidate] >= last_ranked
      return false
    end
  end
  return true
end # TODO use LoopVectorization.jl here?

function obeys_ranking(utility::Utility{N,T},
    voter_index::Int, vote::MultiIssueVote, indif::Bool)::Bool where {N,T}
  offset = 0
  all(vote) do issue
    interval = (offset+1) : (offset+issue.n_candidates)
    offset += issue.n_candidates
    obeys_ranking(utility, issue.choices[voter_index], indif, interval)
  end
end

function check_issues(vote::MultiIssueVote, n_voters::Int, N::Int)
  @assert all( length(issue.choices) == n_voters for issue in vote )
  @assert sum( issue.n_candidates for issue in vote ) == N
end

using Base.Threads

# straight into informed sampling. rejection sampler with variable runtime. sample_ordered_gaussian_mixture()
function rejection_sample_utilities!(voters::Vector{VoterRealization{N,M,T}},
    vote::MultiIssueVote, mixture::VoterMixture{N,M,T},
    n_sample_attempts::Int, indif::Bool=false)::Int where {N,M,T}
  n_voters = length(voters)
  n_issues = length(vote)
  @assert isapprox(sum(mixture.shares), 1)
  check_issues(vote, n_voters, N)
  categorical = Categorical(mixture.shares) # is it wiser to store `shares` as a `Categorical`?
  ## the following is what I would do if I were to Gibbs-sample from each
  ## realization `sub_interval` at a time, conditioning on the rest
  #outside_interval = vcat(
  #  1:(sub_interval.start-1), (sub_interval.stop+1):N )
  #conditional_mean_scales = [
  #  ( covariance[sub_interval, outside_interval]
  #    * inv(covariance)[outside_interval, outside_interval] )
  #  for covariance in cov.(mixture.cohorts) ]
  #conditional_covariances = [
  #  invcov(cohort)[sub_interval, sub_interval] |> inv
  #  for cohort in mixture.cohorts ]
  #centered_conditional_cohorts = [
  #  VoterCohort{T}(zeros(T, length(sub_interval)), covariance)
  #  for covariance in conditional_covariances ]
  n_threads = nthreads()
  sample = zeros(T, N, n_threads) # allocate an array to act as an intermediary in which to store results before assigning to a Vector of SVectors
  #utility = zeros(Utility{N,T}, n_threads) why did I have this? be careful when removing it (ref: when you see a fence with an unknown purpose...)
  n_failures = zeros(Int, n_threads)
  @threads for voter_index in 1:n_voters
    thread_id = threadid()
    this_sample = @view sample[:, thread_id]
    cohort_index = 0
    n_failures[thread_id] += 1 # this shall be deemed a failure until proven otherwise
    utility = zero(Utility{N,T})
    @inbounds for attempt in 1:n_sample_attempts # beyond this we simply give up and accept a suboptimal fate
      cohort_index = rand(categorical)
      cohort = mixture.cohorts[cohort_index]
      rand!(cohort, this_sample) # per the docs, the global RNG is thread-safe as of Julia 1.3 (I've been following the development since at least sophomore year!)
      utility = Utility{N,T}(this_sample)
      offset = 0
      failed = false
      for issue in vote # better praxis would be to wrap this into a smaller function..
        ranking = issue.choices[voter_index]
        interval = (offset+1):(offset+issue.n_candidates)
        offset += issue.n_candidates # non-overlapping
        satisfaction = obeys_ranking(utility, ranking, indif, interval)
        if !satisfaction # && break
          failed = true
          break
        end # idea: settle on the one with the FEWEST violations? also try to Gibbs sampler on constrained marginals...
      end
      if !failed
        n_failures[thread_id] -= 1
        break
      end
    end
    realization = VoterRealization{N,M,T}(cohort_index, utility)
    voters[voter_index] = realization
  end
  sum(n_failures)
end


text:
export Beta

"""
Description:

    Beta factor node

    Real scalars
    a > 0
    b > 0

    f(out, a, b) = Beta(out|a, b) = Γ(a + b)/(Γ(a) Γ(b)) out^{a - 1} (1 - out)^{b - 1}

Interfaces:

    1. out
    2. a
    3. b

Construction:

    Beta(id=:some_id)
"""
mutable struct Beta <: SoftFactor
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    function Beta(out, a, b; id=generateId(Beta))
        @ensureVariables(out, a, b)
        self = new(id, Array{Interface}(undef, 3), Dict{Symbol,Interface}())
        addNode!(currentGraph(), self)
        self.i[:out] = self.interfaces[1] = associate!(Interface(self), out)
        self.i[:a] = self.interfaces[2] = associate!(Interface(self), a)
        self.i[:b] = self.interfaces[3] = associate!(Interface(self), b)

        return self
    end
end

slug(::Type{Beta}) = "Beta"

format(dist::ProbabilityDistribution{Univariate, Beta}) = "$(slug(Beta))(a=$(format(dist.params[:a])), b=$(format(dist.params[:b])))"

ProbabilityDistribution(::Type{Univariate}, ::Type{Beta}; a=1.0, b=1.0) = ProbabilityDistribution{Univariate, Beta}(Dict(:a=>a, :b=>b))
ProbabilityDistribution(::Type{Beta}; a=1.0, b=1.0) = ProbabilityDistribution{Univariate, Beta}(Dict(:a=>a, :b=>b))

dims(dist::ProbabilityDistribution{Univariate, Beta}) = 1

vague(::Type{Beta}) = ProbabilityDistribution(Univariate, Beta, a=1.0, b=1.0)

isProper(dist::ProbabilityDistribution{Univariate, Beta}) = (dist.params[:a] > 0.0) && (dist.params[:b] > 0.0)

unsafeMean(dist::ProbabilityDistribution{Univariate, Beta}) = dist.params[:a]/(dist.params[:a] + dist.params[:b])

unsafeLogMean(dist::ProbabilityDistribution{Univariate, Beta}) = digamma(dist.params[:a]) - digamma(dist.params[:a] + dist.params[:b]) # E[log(X)]

unsafeMirroredLogMean(dist::ProbabilityDistribution{Univariate, Beta}) = digamma(dist.params[:b]) - digamma(dist.params[:a] + dist.params[:b]) # E[log(1 - X)]

unsafeVar(dist::ProbabilityDistribution{Univariate, Beta}) = dist.params[:a]*dist.params[:b]/((dist.params[:a] + dist.params[:b])^2*(dist.params[:a] + dist.params[:b] + 1.0))

logPdf(dist::ProbabilityDistribution{Univariate, Beta}, x) = (dist.params[:a]-1)*log(x) + (dist.params[:b]-1)*log(1.0-x) - loggamma(dist.params[:a]) - loggamma(dist.params[:b]) + loggamma(dist.params[:a]+dist.params[:b])

function prod!( x::ProbabilityDistribution{Univariate, Beta},
                y::ProbabilityDistribution{Univariate, Beta},
                z::ProbabilityDistribution{Univariate, Beta}=ProbabilityDistribution(Univariate, Beta, a=1.0, b=1.0))

    z.params[:a] = x.params[:a] + y.params[:a] - 1.0
    z.params[:b] = x.params[:b] + y.params[:b] - 1.0

    return z
end

@symmetrical function prod!(x::ProbabilityDistribution{Univariate, Beta},
                            y::ProbabilityDistribution{Univariate, PointMass},
                            z::ProbabilityDistribution{Univariate, PointMass}=ProbabilityDistribution(Univariate, PointMass, m=0.0))

    (0.0 <= y.params[:m] <= 1.0) || error("PointMass location $(y.params[:m]) should be between 0 and 1")
    z.params[:m] = y.params[:m]

    return z
end

sample(dist::ProbabilityDistribution{Univariate, Beta}) = betainvcdf(dist.params[:a], dist.params[:b], rand())

# Entropy functional
function differentialEntropy(dist::ProbabilityDistribution{Univariate, Beta})
    labsbeta(dist.params[:a], dist.params[:b]) -
    (dist.params[:a] - 1.0)*digamma(dist.params[:a]) -
    (dist.params[:b] - 1.0)*digamma(dist.params[:b]) +
    (dist.params[:a] + dist.params[:b] - 2.0)*digamma(dist.params[:a] + dist.params[:b])
end

# Average energy functional
function averageEnergy(::Type{Beta}, marg_out::ProbabilityDistribution{Univariate}, marg_a::ProbabilityDistribution{Univariate, PointMass}, marg_b::ProbabilityDistribution{Univariate, PointMass})
    labsbeta(marg_a.params[:m], marg_b.params[:m]) -
    (marg_a.params[:m] - 1.0)*unsafeLogMean(marg_out) -
    (marg_b.params[:m] - 1.0)*unsafeMirroredLogMean(marg_out)
end

# By Stirling's approximation and Monte Carlo summation
function averageEnergy(::Type{Beta}, marg_out::ProbabilityDistribution{Univariate}, marg_a::ProbabilityDistribution{Univariate}, marg_b::ProbabilityDistribution{Univariate})
    unsafeMeanLogMean(marg_a) - unsafeMean(marg_a) + 0.5*(log(2*pi)-unsafeLogMean(marg_a)) +
    unsafeMeanLogMean(marg_b) - unsafeMean(marg_b) + 0.5*(log(2*pi)-unsafeLogMean(marg_b)) -
    (unsafeMean(marg_a)-1)*unsafeLogMean(marg_out) -
    (unsafeMean(marg_b)-1)*unsafeMirroredLogMean(marg_out) -
    sum(loggamma.(sample(marg_a,default_n_samples).+sample(marg_b,default_n_samples)))/default_n_samples
end


text:
using XLA, NNlib, Test

@test @code_xla(2+3) isa XLA.IR

double = xla(x -> x + x)

@test double(21) == 42
@test double(4.5) == 9.0

xadd = xla(+)

@test xadd(2, 2.0) == 4.0

xadd2 = xla(xadd)

@test xadd2(2, 3) == 5

@test xla(() -> 2+2)() == 4

@test xla(x -> 3x^(1+1) + (2x + 1))(5) == 86

relu = x -> x > 0 ? x : zero(x)
xrelu = xla(relu)

@test xrelu(5) == 5
@test xrelu(-5) == 0
@test xrelu(5.0) == 5.0

let x = rand(3), y = rand(3)
  @test xadd(x, y) isa XLA.XArray
  @test collect(xadd(x, y)) == x+y
end

function updatezero!(env)
  if env[:x] < 0
    env[:x] = 0
  end
end

relu = xla() do x
  env = Dict()
  env[:x] = x
  updatezero!(env)
  return env[:x]
end

@test relu(5) == 5
@test relu(-5) == 0

poly(x) = x^2 + 1

poly(xs::AbstractArray) = poly.(xs)

xpoly = xla(poly)

@test xpoly(3) == 10

@test collect(xpoly([1, 2, 3])) == [2, 5, 10]

xsum = xla(xs -> reduce((a, b) -> a+b, xs))

@test xsum([1, 2, 3]) == 6

xexp = xla(x -> exp.(x))

@test xexp(2) == exp(2)

@test collect(xexp([1, 2, 3])) == exp.([1, 2, 3])

xsum1 = xla(xs -> sum(xs, dims = 1))

@test collect(xsum1([1, 2, 3, 4])) == [10]

@test collect(xsum1([1 2; 3 4])) == [4 6]

xsoftmax = xla(softmax)

@test collect(xsoftmax([1, 2, 3])) ≈ softmax([1, 2, 3])
@test collect(xsoftmax([1 2; 3 4])) == softmax([1 2; 3 4])

xsquare = xla(x -> x^2)

@test xsquare(1+2im) == -3+4im

xsquare = xla(x -> @show x^2)

function stdout_string(f)
  _stdout = stdout
  try
    rd, = redirect_stdout()
    f()
    return String(readavailable(rd))
  finally
    redirect_stdout(_stdout)
  end
end

output = stdout_string() do
  @test xsquare(1+2im) == -3+4im
end

@test output == "x ^ 2 = -3 + 4im\n"

@test isxla() == false
@test xla(isxla)() == true

timestable(x) = x .* x'

@test collect(xla(timestable)([1, 2, 3])) == timestable([1, 2, 3])

function sumloop(xs)
  s = zero(eltype(xs))
  for i = 1:length(xs)
    s += xs[i]
  end
  return s
end

sumloop([1, 2, 3])

@test xla(sumloop)([1, 2, 3]) == 6

let
  f(x) = reshape(x, length(x))
  @test collect(xla(f)([1 2; 3 4])) == [1, 2, 3, 4]
end


text:
cd(@__DIR__)

# if isfile("Project.toml") && isfile("Manifest.toml")
#     using Pkg
#     Pkg.activate(".")
# else
#     printstyled("You are missing environment files!\n"; color = :red, bold = true)
# end

using StaticArrays
using LaserTypes
using LinearAlgebra

using Plots
fmt = "pdf"

using LaTeXStrings
using RecursiveArrayTools
using Serialization
using OrdinaryDiffEq
using ThreadsX
using Distributions
using QuadGK
using UnPack
using LaserTypes: fundamental_constants

function folder_setup(path)
    path = replace(path,"/" => "")

    if basename(pwd()) == basename(path) # split(pwd(),"/")[end] == path
        println("already in $path")
        return
    end

    if !isdir(path)
        mkdir(path)
        println("making folder $path")
    end

    cd(path)
    println("moved in $path")
end

const N = 100_000 # number of particles
const c = 137.036 # speed of light
const ω = 0.057; const T₀ = 2π/ω # angular frequency, period
const k = ω/c # wavenumber
const λ = c*T₀ # wavelength
const w₀ = 75 * λ # is the waist radius
const a₀ = 1e-5; const AA = a₀*c # (qA₀)/(mc) = a₀
const w = 10.; const τ = w/ω; # temporal pulse decay time given in terms of number of oscilations and frequency
const n = 5 # number of periods to integrate before and after pulse collides with particles
const π₀ = a₀*c # in atomic units, a₀mc has units of linear momentum and sets the scale for linear momentum transfered to the particle, not that m = 1 for our particles
const maxR = 3.25w₀ # maximum radius for distributing praticles
const ϕ₀ = 0.

folder_setup("output_a$(a₀)_w$(round(w₀/λ, digits = 1))_phi$(round(ϕ₀, digits = 3))")

profile = GaussProfile(τ = w/ω, z₀ = 0.)

function init_lasers(λ, a₀, ϕ₀, w₀, p_max)
    Iterators.flatten([[setup_laser(LaguerreGaussLaser, :atomic; λ, a₀, ϕ₀, w₀, p, m, ξx, ξy, profile)
        for m in -p:p]
        for p in 0:p_max,
            (ξx, ξy) in [(1., 0.),
                         ((1. + 0im)/√2, (1im)/√2),
                         ((1. + 0im)/√2, (-1im)/√2)]
    ])
end

# p_max = 2;
lasers = init_lasers(λ, a₀, ϕ₀, w₀, 2)
lasers = [setup_laser(LaguerreGaussLaser, :atomic; λ, a₀, ϕ₀, w₀, p = 2, m = -2, ξx = 1., ξy = 0., profile)]

# roots don't change for m = ± |m|; adimensional ρ/w₀
rootsLUT = Base.ImmutableDict(
    (p = 0, m = 1) => [0., 1.954], # unphysical

    (p = 0, m = 0) => [0., 1.518],

    (p = 1, m = 0) => [0., 0.707, 2.323],
    (p = 1, m = 1) => [0., 1., 2.566],

    (p = 2, m = 0) => [0., 0.541, 1.306, 2.807],
    (p = 2, m = 1) => [0., 0.797, 1.538, 3.005],
    (p = 2, m = 2) => [0., 1.000, 1.731, 3.176]
)

rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
circ = Shape(Plots.partialcircle(0, 2π))

function add_info!(plt, type, isscatter = false)
    #there is either a bug or confusing behaviour with ylims and aspect_ratio
    xl, yl = (isscatter ? ((-5.658368025329326, 5.667117711980884), (-3.662004592945377, 3.613120699860083)) : (xlims(plt), ylims(plt)))
    frame_width = xl[2] - xl[1]
    frame_height =yl[2] - yl[1]

    ξy = type.polarization.ξy

    box_ratio = (w = 0.14 + 0.05 * abs(ξy*√2), h = 0.31) # hardcoded, I don't know how to get display length

    # lower left corner coordinates
    box = (x = (0.99 - box_ratio.w)*frame_width + xl[1],
           y = (0.99 - box_ratio.h)*frame_height + yl[1],
           w = box_ratio.w*frame_width,
           h = box_ratio.h*frame_height)

    # compute text positions
    text_x = box.x + box.w/2 # always centered on the box
    text_y = box.y + box.h # starting from the top, mind the minus
    δy = box.h/8

    # find the exact strings to print, related to polarization
    sgn = (imag(ξy) < 0 ? "-" : "")
    xi = (ξy == 0 ? "\$\\xi_y=0\$" : "\$\\xi_y= $sgn\\mathrm{i}/\\sqrt{2}\$")

    annotate!(plt, [(text_x, text_y - 1δy, Plots.text("\$a_0=$(type.a₀)\$", :center, 10, "computer modern")),
                    (text_x, text_y - 3δy, Plots.text(xi,                   :center, 10, "computer modern")),
                    (text_x, text_y - 5δy, Plots.text("\$p=$(type.p)\$",    :center, 10, "computer modern")),
                    (text_x, text_y - 7δy, Plots.text("\$m=$(type.m)\$",    :center, 10, "computer modern"))])

    # add a rectangle without changing the frame size
    plot!(plt, rectangle(box.w, box.h, box.x, box.y);
                        xlims = xl, ylims = yl,
                        linewidth = 1.5,
                        fillcolor = :white, label = false)
end

function add_rect!(plt, X, Y, roots; kwargs...)
    h = ylims(plt)[2]-ylims(plt)[1]
    y₀ = ylims(plt)[1]
    underlay = plot([],[], label = false)
    for i in 1:(length(roots) - 1)
        a = roots[i]
        b = roots[i+1]
        plot!(underlay,rectangle(b-a, h, a, y₀),
        color_palette = :Set1_4, seriescolor = i,
        fillalpha = .4, linecolor = nothing, label = false)
    end
    plot!(underlay, X, Y;
        xlims = xlims(plt), ylims = ylims(plt), kwargs...)
end

const τi = - n * τ
const τf =   n * τ
const time_samples = 2000
const Δτ = range(0, stop = 2*n*τ/T₀, length = time_samples+1)

function polarization(type)
    if imag(type.polarization.ξy)*√2 == 1.
        printstyled("\n Right Circular Polarization, p = $(type.p), m = $(type.m)\n"; color = :green, bold = true)
        return "RHCP"
    end
    if imag(type.polarization.ξy)*√2 == -1.
        printstyled("\n Left Circular Polarization, p = $(type.p), m = $(type.m)\n"; color = :green, bold = true)
        return "LHCP"
    end
    if imag(type.polarization.ξy)*√2 == 0.
        printstyled("\n Linear Polarization, p = $(type.p), m = $(type.m)\n"; color = :green, bold = true)
        return "LinP"
    end
end

io = open("output_log.txt","w")
write(io,"""Simulation Paramters

Number of Particles per Ring,           N = $N
Speed of Light [Atomic Units],          c = $c
Angular Frequency [Atomic Units],       ω = $ω
Oscillation Period [Atomic Units],      T₀ = $T₀
Wavenumber [Atomic Units],              k = $k
Wavelength [Atomic Units],              λ = $λ
Beam Waist Radius [Atomic Units],       w₀ = $w₀
Reduced Vector Potential,               a₀ = $a₀
Initial Phase                           φ₀ = $ϕ₀
Number of Oscillations in the Pulse,    w = $w
Envelope Decay Time,                    τ = $τ
Maximum Linear Momentum Transfered,     π₀ = $π₀
Disc Radius for Particle Distribution,  Rₘₐₓ = $maxR\n
""") # save out important parameters

# include("field-plots.jl")
include("electrodynamic-sims2.jl")

close(io)
cd("..") # return to file folder

max_stars_repo_name: toma-vlad/electrodynamic-simulations

max_issues_repo_name: toma-vlad/electrodynamic-simulations

max_forks_repo_name: toma-vlad/electrodynamic-simulations


text:
import DataStructures: counter

mutable struct Basin <: AgComponent
    name::String
    zones::Tuple
    climate::Climate
    managers::Tuple
    current_ts::NamedTuple{(:idx, :dt)}

    function Basin(; name, zone_spec, climate_data, managers)
        basin = new()
        basin.name = name

        climate::Climate = load_climate(climate_data)
        basin.climate = climate
        basin.current_ts = (idx=1, dt=climate.time_steps[1])

        basin.managers = Tuple(mg for (mg, _) in managers)

        basin.zones = Tuple(create(v, climate) for (k, v) in zone_spec)
        assign_managers!(managers, basin.zones)

        return basin
    end
end

"""
# Example

```julia
# Load zone specifications
zone_specs = load_spec("test/data/zones")

# Define managers of interest
Manager_A = BaseManager("optimizing")
Manager_B = RigidManager("rigid", 0.05)

# Create tuple-based relation between managers and zones (by name)
manage_zones = ((Manager_A, ("Zone_1", )), (Manager_B, ("Zone_2", "Zone_3")))

# Attach managers to their associated zones
assign_managers!(manage_zones, values(zone_specs))
```

Raises `ArgumentError` if duplicate zone names are found.
"""
function assign_managers!(rel::Tuple, zones)
    all_names = []
    for (mg, zone_names) in rel
        for tgt_name in zone_names
            for zone in zones
                if string(tgt_name) == zone.name
                    push!(all_names, tgt_name)
                    zone.manager = mg
                end
            end
        end
    end

    dups::Array{String} = String[k for (k, v) in counter(all_names) if v > 1]
    if length(dups) > 0
        throw(ArgumentError("Cannot co-manage a zone!\nDuplicates found: $(dups)"))
    end

    if length(all_names) != length(zones)
        @info "Names vs zones" all_names zones
        throw(ArgumentError("One or more zones are missing a manager"))
    end
end

function reset!(b::Basin)::Nothing
    for z in b.zones
        reset!(z)
    end

    b.current_ts = (idx=1, dt=b.climate.time_steps[1])

    return nothing
end

"""Collect model run results for a Basin"""
function collect_results(basin::Basin; last=false)::Dict
    results = Dict()
    sizehint!(results, length(basin.zones))

    for z in basin.zones
        results[z.name] = collect_results(z)
    end

    return results
end

max_stars_repo_name: ConnectedSystems/agtor-jl

max_issues_repo_name: ConnectedSystems/agtor-jl

max_forks_repo_name: ConnectedSystems/agtor-jl


text:
function sum_vector(x::Array{Float64, 1})
s = 0.0
for i = 1:length(x)
s = s + x[i]
end
return s
end
function sum_cols_matrix(x::Array{Float64, 2})
num_cols = size(x, 2)
s = zeros(num_cols)
for i = 1:num_cols
s[i] = sum_vector(x[:, i])
end
return s
end

max_stars_repo_path: Module 2/Chapter06/array_views.jl

max_issues_repo_path: Module 2/Chapter06/array_views.jl

max_forks_repo_path: Module 2/Chapter06/array_views.jl


text:
# set knob to normal
# KN
export psp_setknobnormal

const KNOB_NORMAL = [0x4b, 0x4E, 0x0d]

"""
    psp_setknobnormal(io)

Set knob to normal.
"""
function psp_setknobnormal(io_psp::IO)
  write(io_psp, KNOB_NORMAL)
  return nothing
end


text:
module TableIO

export read_table, write_table!, read_sql

using Tables, Requires, Suppressor
using CSV, DataFrames # required for multiple file types, therefore currently not optional

## definition of file formats and extensions

abstract type AbstractFormat end

struct CSVFormat <: AbstractFormat end
struct ZippedFormat <: AbstractFormat end
struct JDFFormat <: AbstractFormat end
struct ParquetFormat <: AbstractFormat end
struct ExcelFormat <: AbstractFormat end
struct SQLiteFormat <: AbstractFormat end
struct StataFormat <: AbstractFormat end
struct SPSSFormat <: AbstractFormat end
struct SASFormat <: AbstractFormat end
struct JSONFormat <: AbstractFormat end
struct ArrowFormat <: AbstractFormat end

# specify if a reader accepts an io buffer as input or if creation of a temp file is required
supports_io_input(::AbstractFormat) = false

supports_io_input(::CSVFormat) = true
supports_io_input(::JSONFormat) = true
supports_io_input(::ArrowFormat) = true

const FILE_EXTENSIONS = Dict(
    "zip" => ZippedFormat,
    "csv" => CSVFormat,
    "jdf" => JDFFormat,
    "parquet" => ParquetFormat,
    "xlsx" => ExcelFormat,
    "db" => SQLiteFormat,
    "sqlite" => SQLiteFormat,
    "sqlite3" => SQLiteFormat,
    "dta" => StataFormat,
    "sav" => SPSSFormat,
    "sas7bdat" => SASFormat,
    "json" => JSONFormat,
    "arrow" => ArrowFormat,
)

const IMPORT_PACKAGES = Dict(
    ZippedFormat => :ZipFile,
    JDFFormat => :JDF,
    ParquetFormat => :Parquet,
    ExcelFormat => :XLSX,
    SQLiteFormat => :SQLite,
    StataFormat => :StatFiles,
    SPSSFormat => :StatFiles,
    SASFormat => :StatFiles,
    JSONFormat => :JSONTables,
    ArrowFormat => :Arrow,
)

## Dispatching on file extensions

"""
    read_table(filename:: AbstractString; kwargs...)

`filename`: path and filename of the input file
`kwargs...`: keyword arguments passed to the underlying file reading function (e.g. `CSV.File`)

Returns a Tables.jl interface compatible object.

Example:

    df = DataFrame(read_table("my_data.csv"); copycols=false)

"""
function read_table(filename:: AbstractString, args...; kwargs...)
    data_type = _get_file_type(filename)()
    try
        # to speed up the standard case (format specific package is imported), this is tried first
        return read_table(data_type, filename, args...; kwargs...)
    catch ex
        if ex isa MethodError
            # import format specific package and invoke latest version of the function to avoid world age issues
            _import_package(data_type)
            return Base.invokelatest(read_table, data_type, filename, args...; kwargs...)
        else
            rethrow()
        end
    end

end

"""
    read_table(file_picker:: Dict, args...; kwargs...)

Reading tabular data from a PlutoUI.jl FilePicker.

Usage (in a Pluto.jl notebook):

    using PlutoUI, TableIO, DataFrames
    using XLSX # import the packages required for the uploaded file formats
    @bind f PlutoUI.FilePicker()
    df = DataFrame(read_table(f); copycols=false)

"""
function read_table(file_picker:: Dict, args...; kwargs...)
    filename, data = _get_file_picker_data(file_picker)
    data_type = _get_file_type(filename)()
    data_buffer = IOBuffer(data)

    _import_package(data_type)

    if supports_io_input(data_type)
        data_object = data_buffer # if it is supported by the corresponding package, creation of a temporary file is avoided and the IOBuffer is used directly
    else
        tmp_file = joinpath(mktempdir(), filename)
        write(tmp_file, data_buffer)
        data_object = tmp_file
    end

    try
        # to speed up the standard case (format specific package is imported), this is tried first
        read_table(data_type, data_object, args...; kwargs...)
    catch ex
        if ex isa MethodError
            # import format specific package and invoke latest version of the function to avoid world age issues
            _import_package(data_type)
            return Base.invokelatest(read_table, data_type, data_object, args...; kwargs...)
        else
            rethrow()
        end
    end
end

"""
    write_table!(filename:: AbstractString, table; kwargs...):: AbstractString

`filename`: path and filename of the output file
`table`: a Tables.jl compatible object (e.g. a DataFrame) for storage
`kwargs...`: keyword arguments passed to the underlying file writing function (e.g. `CSV.write`)

Example:

    write_table!("my_output.csv", df)

"""
function write_table!(filename:: AbstractString, table, args...; kwargs...)
    data_type = _get_file_type(filename)()
    try
        # to speed up the standard case (format specific package is imported), this is tried first
        write_table!(data_type, filename, table, args...; kwargs...)
    catch ex
        if ex isa MethodError
            # import format specific package and invoke latest version of the function to avoid world age issues
            _import_package(data_type)
            return Base.invokelatest(write_table!, data_type, filename, table, args...; kwargs...)
        else
            rethrow()
        end
    end
    _import_package(data_type)
    nothing
end

"""
    read_sql(db, sql:: AbstractString)

Returns the result of the SQL query as a Tables.jl compatible object.
"""
function read_sql end

## CSV - always supported because CSV.jl is required for multiple other file formats, too

read_table(::CSVFormat, filename:: AbstractString; kwargs...) = CSV.File(filename; kwargs...)
read_table(::CSVFormat, io:: IO; kwargs...) = CSV.File(read(io); kwargs...)

function write_table!(::CSVFormat, output:: Union{AbstractString, IO}, table; kwargs...)
    _checktable(table)
    table |> CSV.write(output; kwargs...)
    nothing
end

## conditional dependencies

function __init__()
    @require ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" include("zip.jl")
    @require JDF = "babc3d20-cd49-4f60-a736-a8f9c08892d3" include("jdf.jl")
    @require Parquet = "626c502c-15b0-58ad-a749-f091afb673ae" include("parquet.jl")
    @require XLSX = "fdbf4ff8-1666-58a4-91e7-1b58723a45e0" include("xlsx.jl")
    @require StatFiles = "1463e38c-9381-5320-bcd4-4134955f093a" include("stat_files.jl")
    @require SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9" include("sqlite.jl")
    @require LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1" include("postgresql.jl")
    @require JSONTables = "b9914132-a727-11e9-1322-f18e41205b0b" include("json.jl")
    @require Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45" include("arrow.jl")
end

## Utilities

_get_file_extension(filename) = lowercase(splitext(filename)[2][2:end])
_get_file_type(filename) = FILE_EXTENSIONS[_get_file_extension(filename)]

function _import_package(::T) where {T <: AbstractFormat}
    if T ∉ keys(IMPORT_PACKAGES)
        return
    end
    import_package = IMPORT_PACKAGES[T]

    # A warning is raised if a package is imported which is not a dependency of TableIO. This warning is suppressed.
    # If the package is not installed, an error message is raised.
    try
        @suppress @eval import $import_package
    catch ex
        # If the package is not installed, the error message is swallowed by @suppress, but the warning message for a missing TableIO dependeny is raised.
        # To get back the more helpful error message for a not installed package, it is regenerated below.
        if ex isa ArgumentError
            throw(ArgumentError("""
                ERROR: ArgumentError: Package $import_package not found in current path:
                - Run `import Pkg; Pkg.add("$import_package")` to install the $import_package package.
                """))
        else
            rethrow()
        end
    end

    # note that it is required to use Base.invokelatest for calling any functionality depending on the imported package, unless one returns to global scope before.
end

_checktable(table) = Tables.istable(typeof(table)) || error("table has no Tables.jl compatible interface")

# poor man's approach to prevent SQL injections / garbage inputs
_checktablename(tablename) = match(r"^[a-zA-Z0-9_]*$", tablename) === nothing && error("tablename must only contain alphanumeric characters and underscores")

function _get_file_picker_data(file_picker:: Dict)
    data = file_picker["data"]:: Vector{UInt8} # brings back type stability
    length(data) == 0 && error("no file selected yet")
    filename = file_picker["name"]:: String
    return filename, data
end

end


text:
using NetworkEpidemics
using LightGraphs
using Statistics
using Plots
using ColorSchemes
using Random

CorrectedMetapopulation(χ::Real, mp::Metapopulation{SI}) = Metapopulation(mp.h, mp.D, SI(χ*mp.dynamics.β))
CorrectedMetapopulation(χ::Real, mp::Metapopulation{SIS}) = Metapopulation(mp.h, mp.D, SIS(χ*mp.dynamics.β, mp.dynamics.γ))
CorrectedMetapopulation(χ::Real, mp::Metapopulation{SIR}) = Metapopulation(mp.h, mp.D, SIR(χ*mp.dynamics.β, mp.dynamics.δ))

N = 1000
M = 10

Random.seed!(2020) # for reproducibility

h = erdos_renyi(M, 0.3)
g = static_scale_free(N, Int(20*N), 2)

k = degree(g)

χ₁ = mean(closeness_centrality(g))
χ₂ = mean(k.^2)/(mean(k)*N)
χ₃ = mean(k)/N #ne(g)/(n*(n-1)/2)
χ₄ = global_clustering_coefficient(g)

β = 0.1
γ = 0.3
D = 0.1

mp = Metapopulation(h, D, SIR(β,γ))
mp_1 = CorrectedMetapopulation(χ₁, mp)
mp_2 = CorrectedMetapopulation(χ₂, mp)
mp_3 = CorrectedMetapopulation(χ₃, mp)
mp_4 = CorrectedMetapopulation(χ₄, mp)

mpx = Metaplex(g, h, D, SIR(β,γ))

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = [div(i-1,div(N,M))+1 for i in 1:N]

x0_mp = Array{Int,2}(undef, M,2)
x0_mp[2:M,:] .= [100 0]
x0_mp[1, :] .= [90, 10]

# infected seed randomly dispersed over meta-nodes
x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = rand(1:M, N)
x0_mp = zeros(Int, M, 3)
for i in 1:N
    x0_mp[x0_μ[i], x0_i[i]] += 1
end

tmax = 60.0
nmax = 100000
nsims = 1000
nbins = 200

ts, output = gillespie(mp, x0_mp, tmax=tmax, nmax=nmax)
ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax)

plot(ts, output[2], label="")
plot(ts_mpx, output_mpx[2], label="")

Random.seed!(2020)

#ts_av, u_av = average(mp, x0_mp, tmax=tmax, nmax=nmax, nbins=nbins, nsims=nsims)
ts_av_mpx, u_av_mpx = average(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax, nbins=nbins, nsims=nsims)

ts_mf, u_mf = meanfield(mp, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_1, u_mf_1 = meanfield(mp_1, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_2, u_mf_2 = meanfield(mp_2, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_3, u_mf_3 = meanfield(mp_3, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_4, u_mf_4 = meanfield(mp_4, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_mpx, u_mf_mpx = meanfield(mpx, [x0_i,x0_μ], tmax=tmax, saveat=ts_av_mpx)

colors = ColorSchemes.tab10;

plot(ts_av_mpx, sum(u_av_mpx[2], dims=2)/N + sum(u_av_mpx[3], dims=2)/N,
    label="Metaplex",
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    linestyle=:dash,
    legend=:bottomright,
    color = colors[2]
    )
    plot!(ts_mf, sum(u_mf[2], dims=2)/N + sum(u_mf[3], dims=2)/N,
    #label="Meanfield (No correction)",
    label = "Metapopulation",
#    xlabel="Time",
#    ylabel="Fraction of infected individuals",
    #legend=:bottomright,
    color = colors[1]
    )
    plot!(ts_mf_mpx, sum(u_mf_mpx[2], dims=2)/N + sum(u_mf_mpx[3], dims=2)/N,
#    label="Meanfield (Metaplex)",
    label = "Metaplex IBMF",
    color = colors[2]
    )
    plot!(ts_mf_1, sum(u_mf_1[2], dims=2)/N + sum(u_mf_1[3], dims=2)/N,
#    label="Meanfield (Closeness Centrality)",
    label = "Closeness",
    color = colors[3]
    )
    plot!(ts_mf_2, sum(u_mf_2[2], dims=2)/N + sum(u_mf_2[3], dims=2)/N,
#    label="Meanfield (Second Moment)",
    label = "Second Moment",
    color = colors[4]
    )
    plot!(ts_mf_3, sum(u_mf_3[2], dims=2)/N + sum(u_mf_3[3], dims=2)/N,
#    label="Meanfield (Mean degree)",
    label = "Mean degree",
    color = colors[5]
    )
    plot!(ts_mf_4, sum(u_mf_4[2], dims=2)/N + sum(u_mf_4[3], dims=2)/N,
#    label="Meanfield (Clustering Coefficient)",
    label = "Clustering",
    color = colors[6]
    )

cols = map(x->[x.r, x.g, x.b], colors.colors)

using MATLAB;

mat"""
figure
title('SI')
xlabel('Time')
ylabel('Fraction of Infected Individuals')
hold on
plot($ts_av_mpx, $(sum(u_av_mpx[2], dims=2)/N + sum(u_av_mpx[3], dims=2)/N), 'LineStyle', '--', 'Color', $(cols[2]))

plot($ts_mf, $(sum(u_mf[2], dims=2)/N + sum(u_mf[3], dims=2)/N), 'Color', $(cols[1]))
plot($ts_mf_mpx, $(sum(u_mf_mpx[2], dims=2)/N + sum(u_mf_mpx[3], dims=2)/N), 'Color', $(cols[2]))
plot($ts_mf_1, $(sum(u_mf_1[2], dims=2)/N + sum(u_mf_1[3], dims=2)/N), 'Color', $(cols[3]))
plot($ts_mf_2, $(sum(u_mf_2[2], dims=2)/N + sum(u_mf_2[3], dims=2)/N), 'Color', $(cols[4]))
plot($ts_mf_3, $(sum(u_mf_3[2], dims=2)/N + sum(u_mf_3[3], dims=2)/N), 'Color', $(cols[5]))
plot($ts_mf_4, $(sum(u_mf_4[2], dims=2)/N + sum(u_mf_4[3], dims=2)/N), 'Color', $(cols[6]))
hold off
ylim([0.0,1.0])
legend({'Metaplex','Metapopulation', 'Metaplex IBMF', 'Closeness Centrality', '\$\\langle k^2\\rangle / \\langle k \\rangle\$', '\$\\langle k \\rangle \$', 'Clustering Coefficient'}, 'Interpreter', 'latex', 'Location', 'southeast')
""";

max_stars_repo_name: csimal/metaplex-networks

max_issues_repo_name: csimal/metaplex-networks

max_forks_repo_name: csimal/metaplex-networks


text:
struct ExperimentalConfig
    diffusion_alpha::Union{Float64, String}
    chemotaxis_mu::Union{Float64, String}
    est::String
    est_strength::Union{Float64, String}
end

function ExperimentalConfig(config::Dict{String,Any})

    valid_user_spec = Dict("diffusion_alpha" => [Float64, String],
                            "chemotaxis_mu" => [Float64, String],
                            "external_signal_type" => String,
                            "external_signal_receptor_strength" =>  [Float64, String])

    for (key, value) in  valid_user_spec
        if ! haskey(config,key)
            throw(DomainError(key, "Please specify value for " * key ))
        end
    end

    for (key, value) in config
        if haskey(valid_user_spec,key)

            if key == "external_signal_type"

                if ! (value in ["location","type"])
                    throw(DomainError((key,value), "Incorrect value type - please specify either 'location' or 'type'"))
                end

            elseif (typeof(value) in valid_user_spec[key])
                if typeof(value) == String
                    if value != "evolve"
                        throw(DomainError((key,value), "Incorrect value type - only possible string is 'evolve'"))
                    end
                end
            else
                throw(DomainError((key,value), "Incorrect value type"))
            end

        else
            throw(DomainError(key, "Unrecognized parameter type"))
        end
    end

    return ExperimentalConfig(config["diffusion_alpha"],
                                config["chemotaxis_mu"],
                                config["external_signal_type" ],
                                config["external_signal_receptor_strength"])
end

struct FGRNConfig

    protein_size::Int
    max_iter::Int
    mandel_thresh::Int

    substrate_protein::Array{Int,2}

    Ct::Float64
    Cs::Float64
    Cw::Float64
    Cp::Float64
    Ci::Float64

end

function FGRNConfig(config::Dict{String,Any})

    valid_user_spec = Dict("protein_size" => [Int],
                            "max_iter" => [Int],
                            "mandel_thresh" => [Int],
                            "Ct" =>  [Float64],
                            "Cs" =>  [Float64],
                            "Cw" =>  [Float64],
                            "Cp" =>  [Float64],
                            "Ci" =>  [Float64])

    for (key, value) in  valid_user_spec
        if ! haskey(config,key)
            throw(DomainError(key, "Please specify value for " * key ))
        end
    end

    for (key, value) in  config
        if !(typeof(value) in valid_user_spec[key])
            throw(DomainError((key,value), "Please specify value with correct type for " * key ))
        end
    end

    substrate_protein = zeros(Int,config["protein_size"],config["protein_size"]);

    return FGRNConfig(config["protein_size"],
                      config["max_iter"],
                      config["mandel_thresh"],
                      substrate_protein,
                      config["Ct"],
                      config["Cs"],
                        config["Cw"],
                        config["Cp"],
                        config["Ci"])
end

struct CellularConfig
    cell_size_upper_limit::Float64
    max_split::Float64
    area_lambda::Float64
    temperature::Float64
    chemotaxis_max_strength::Float64
    grid_size::Int
    Mx::Tridiagonal{Float64,Array{Float64,1}}
    My::Tridiagonal{Float64,Array{Float64,1}}
end

function CellularConfig(config::Dict{String,Any})

    valid_user_spec = Dict("cell_size_upper_limit" => [Float64],
                            "max_split" => [Float64],
                            "area_lambda" => [Float64],
                            "temperature" =>  [Float64],
                            "chemotaxis_max_strength" =>  [Float64],
                            "grid_size" =>  [Int])

    for (key, value) in  valid_user_spec
        if ! haskey(config,key)
            throw(DomainError(key, "Please specify value for " * key ))
        end
    end

    for (key, value) in  config
        if !(typeof(value) in valid_user_spec[key])
            throw(DomainError((key,value), "Please specify value with correct type for " * key ))
        end
    end

    Mx = Tridiagonal([1.0 for i in 1:config["grid_size"]-1],[-2.0 for i in 1:config["grid_size"]],[1.0 for i in 1:config["grid_size"]-1])
    My = copy(Mx)

    Mx[2,1] = 2.0
    Mx[end-1,end] = 2.0
    My[1,2] = 2.0
    My[end,end-1] = 2.0

    return CellularConfig(config["cell_size_upper_limit"],
                      config["max_split"],
                      config["area_lambda"],
                      config["temperature"],
                      config["chemotaxis_max_strength"],
                        config["grid_size"],Mx,My)
end

struct POOLConfig
    n_total_genes::Int64
    n_evaluation::Int64
    time_steps::Int64
    adult_pop_size::Int64
    child_pop_size::Int64
    n_pools::Int64

    elitism_percentage::Float64 #0.4
    elitism_migration_percentage::Float64
    emmigration_rate::Float64
    immigration_rate::Float64
    crossover::Float64
    flag::Matrix{Float64}
    genome_start::Vector{Int64}
    genome_preserve::Int64
end

function POOLConfig(config::Dict{String,Any})

    valid_user_spec = Dict("n_total_genes" => [Int64],
                            "n_evaluation" => [Int64],
                            "time_steps" => [Int64],
                            "adult_pop_size" =>  [Int64],
                            "child_pop_size" =>  [Int64],
                            "n_pools" =>  [Int64],
                            "elitism_percentage" =>  [Float64],
                            "elitism_migration_percentage" =>  [Float64],
                            "emmigration_rate" =>  [Float64],
                            "immigration_rate" =>  [Float64],
                            "crossover" => [Float64],
                            "flag" =>  [Matrix{Float64}],
                            "genome_start" => [Vector{Int64}])

    for (key, value) in  valid_user_spec
        if ! haskey(config,key)
            throw(DomainError(key, "Please specify value for " * key ))
        end
    end

    for (key, value) in config
        if !(typeof(value) in valid_user_spec[key])
            throw(DomainError((key,value), "Please specify value with correct type for " * key ))
        end
    end

    return POOLConfig(config["n_total_genes"],
                      config["n_evaluation"],
                      config["time_steps"],
                      config["adult_pop_size"],
                      config["child_pop_size"],
                        config["n_pools"],
                        config["elitism_percentage"],
                        config["elitism_migration_percentage"],
                        config["emmigration_rate"],
                        config["immigration_rate"],
                        config["crossover"],
                        config["flag"],
                        config["genome_start"],
                        size(config["genome_start"],1))
end

struct EMAPConfig
    n_total_genes::Int64
    n_evaluation::Int64
    time_steps::Int64
    adult_niche_size::Int64
    child_map_size::Int64
    initial_pop_size::Int64
    families::Vector{Int64}
    elitism_percentage::Float64 #0.4
    inter_species_crossover::Float64 #0.2
    normal_crossover::Float64 #0.8
    non_elite_crossover::Float64  #0.2
    genome_start::Vector{Int64}
    genome_preserve::Int64
end

function EMAPConfig(config::Dict{String,Any})

    valid_user_spec = Dict("n_total_genes" => [Int64],
                            "n_evaluation" => [Int64],
                            "time_steps" => [Int64],
                            "adult_niche_size" =>  [Int64],
                            "child_map_size" =>  [Int64],
                            "initial_pop_size" =>  [Int64],
                            "families" =>  [Vector{Int64}],
                            "elitism_percentage" =>  [Float64],
                            "inter_species_crossover" =>  [Float64],
                            "normal_crossover" =>  [Float64],
                            "non_elite_crossover" =>  [Float64],
                            "genome_start" => [Vector{Int64}])

    for (key, value) in  valid_user_spec
        if ! haskey(config,key)
            throw(DomainError(key, "Please specify value for " * key ))
        end
    end

    for (key, value) in  config
        if !(typeof(value) in valid_user_spec[key])
            throw(DomainError((key,value), "Please specify value with correct type for " * key ))
        end
    end

    n_rec = 0
    n_beh = 0
    n_reg = 0
    n_env = 0
    error = 0

    for i in config["genome_start"]
        if i == 1
            n_rec+=1
        elseif i == 2
            n_beh+=1
        elseif i == 3
            n_reg+=1
        elseif i == 4
            n_env+=1
        else
            error+=1
        end
    end

    @assert n_rec >= 4 "This organism requires at least 4 receptor proteins"
    @assert n_beh >= 3 "This organism requires at least 3 behaviorial proteins"
    @assert n_reg >= 1 "This organism requires at least 1 regulatory proteins"
    @assert n_env >= 1 "This organism requires at least 1 environmental proteins"
    @assert error == 0 "You have specified an undefined protein type. Please use values [1-4] and see reference."

    return EMAPConfig(config["n_total_genes"],
                      config["n_evaluation"],
                      config["time_steps"],
                      config["adult_niche_size"],
                      config["child_map_size"],
                        config["initial_pop_size"],
                        config["families"],
                        config["elitism_percentage"],
                        config["inter_species_crossover"],
                        config["normal_crossover"],
                        config["non_elite_crossover"],
                        config["genome_start"],
                        size(config["genome_start"],1))
end

struct ESTSetup
    affinity_1::Float64
    affinity_2::Float64
    affinity_3::Float64
    affinity_4::Float64

    left_segment::Int
    middle_segment::Int
end

max_stars_repo_name: harrybooth/FractalDevelopment

max_issues_repo_name: harrybooth/FractalDevelopment

max_forks_repo_name: harrybooth/FractalDevelopment


text:
using BioDataStructures
using Base.Test

include("test_IntRangeSet.jl")
include("test_IntRangeDict.jl")


text:
ENV["PYTHON"]=""
using Pkg
Pkg.add("Conda")
using Conda
Conda.update()
Conda.add("matplotlib")
Conda.add_channel("conda-forge")
Pkg.add("PyCall")
Pkg.build("PyCall")
Pkg.add("PyPlot")
Pkg.add("SpecialFunctions")
Pkg.add("ForwardDiff")
Pkg.add("DiffResults")
Pkg.add("Optim")
Pkg.add("Documenter")
Pkg.add("GSL")
Pkg.add("DelimitedFiles")
Pkg.add("IterativeSolvers")
Pkg.add("JLD2")
Pkg.add("MCMCDiagnostics")
Pkg.clone("https://github.com/rodluger/Limbdark.jl.git")


text:
@testset "Volume Image" begin
    fname = joinpath(dirname(@__FILE__), "../../data", "test-3d.dat")
    t = read_VolumeImage(fname)
    t2 = read_VolumeImage(fname)

    @testset "Reading" begin
        t = read_VolumeImage(fname)
        @test isa(t, EEG.VolumeImage) ==  true
    end

    @testset "Create" begin
        # Array
        n = VolumeImage(t.data, t.units, t.x, t.y, t.z, t.t, t.method, t.info, t.coord_system)
        @test isa(n, EEG.VolumeImage) == true
        @test isequal(n.data, t.data) == true

        # Vector
        d = zeros(length(t.data))
        x = zeros(length(t.data))
        y = zeros(length(t.data))
        z = zeros(length(t.data))
        s = zeros(length(t.data))  # t is already taken
        i = 1
        for xi in 1:length(t.x)
            for yi in 1:length(t.y)
                for zi in 1:length(t.z)
                    for ti in 1:length(t.t)
                        d[i] = ustrip(t.data[xi, yi, zi, ti])
                        x[i] = ustrip(t.x[xi])
                        y[i] = ustrip(t.y[yi])
                        z[i] = ustrip(t.z[zi])
                        s[i] = ustrip(t.t[ti])
                        i += 1
                    end
                end
            end
        end
        n2 = VolumeImage(d, t.units, x, y, z, s, t.method, t.info, t.coord_system)
        @test isa(n2, EEG.VolumeImage) == true
        @test isequal(n2.data, t.data) == true
        @test isequal(n2.x, t.x) == true
        @test isequal(n2.y, t.y) == true
        @test isequal(n2.z, t.z) == true
        @test isequal(n2.t, t.t) == true

    end

    @testset "Maths" begin

        @test isequal(t, t2) == true

        @test t + t ==  t * 2
        t4 = t * 4
        @test t4 / 2 == t + t
        @test t4 / 2 == t4 - t - t

        t2 = read_VolumeImage(fname)
        t2.units = "A/m^3"
        # @test_throws t + t2
    end

    @testset "Min/Max" begin

        @test maximum(t) == 33.2692985535
        @test minimum(t) == -7.5189352036

        b = t * 2

        @test b.data[1, 1, 1] == 2 * b.data[1, 1, 1]
        @test b.data[2, 2, 2] == 2 * b.data[2, 2, 2]

        c = mean([b, t])

        @test size(c.data) == size(b.data)
        @test c.data[1, 1, 1] == (b.data[1, 1, 1] + t.data[1, 1, 1]) / 2

        @test maximum([b, t]) == maximum(t) * 2
        @test minimum([b, t]) == minimum(b)

        n = normalise([b, t])

        @test maximum(n) <= 1.0
        @test minimum(n) >= -1.0

    end

    @testset "Dimension checks" begin
        t2 = deepcopy(t); t2.x = t2.x[1:3]
        @test_throws KeyError EEG.dimensions_equal(t, t2)
        t2 = deepcopy(t); t2.y = t2.y[1:3]
        @test_throws KeyError EEG.dimensions_equal(t, t2)
        t2 = deepcopy(t); t2.z = t2.z[1:3]
        @test_throws KeyError EEG.dimensions_equal(t, t2)
        t2 = read_VolumeImage(joinpath(dirname(@__FILE__), "../../data", "test-4d.dat"))
        @test_throws KeyError EEG.dimensions_equal(t, t2)
    end

    @testset "Average" begin
        s = mean(t)
    end

    @testset "Find dipoles" begin

        t = read_VolumeImage(fname)
        dips = find_dipoles(mean(t))
        @test size(dips) == (16,)

        #dips = EEG.new_dipole_method(mean(t))
        #@test size(dips) == (9,)

        fname = joinpath(dirname(@__FILE__), "../../data", "test-4d.dat")
        t2 = read_VolumeImage(fname)
        dips = find_dipoles(mean(t2))
        @test size(dips) == (3,)

        # Test dipoles are returned in order of size
        @test issorted([dip.size for dip in dips], rev = true) == true
    end

    @testset "Best Dipole" begin

        # dips = find_dipoles(mean(t))
        # # Plots.pyplot(size=(1400, 400))
        # # p = plot(t, c = :inferno)
        # # p = plot(p, Talairach(-0.04, 0.01, 0.02))
        # bd = best_dipole(Talairach(-0.05, 0, 0.01), dips)
        # #p = plot(p, Talairach(-0.05, 0, 0.01), c = :red)
        #
        # @test float(bd.x) == roughly(-0.0525, atol =  0.001)
        # @test float(bd.y) == roughly(-0.00378, atol =  0.001)
        # @test float(bd.z) == roughly(0.0168099, atol =  0.001)
        #
        # # Take closest
        # bd = best_dipole(Talairach(-0.05, 0, 0.01), dips, maxdist = 0.0015)
        # @test float(bd.x) == roughly(-0.0525, atol =  0.001)
        # @test float(bd.y) == roughly(-0.00378, atol =  0.001)
        # @test float(bd.z) == roughly(0.0168099, atol =  0.001)
        #
        # # Take only valid dipole
        # dists = [euclidean(Talairach(-0.05, 0, 0.01), dip) for dip=dips]
        # bd = best_dipole(Talairach(-0.05, 0, 0.01), dips, maxdist = 0.015)
        # @test float(bd.x) == roughly(-0.0525, atol =  0.001)
        # @test float(bd.y) == roughly(-0.00378, atol =  0.001)
        # @test float(bd.z) == roughly(0.0168099, atol =  0.001)
        #
        # bd = best_dipole(Talairach(-0.05, 0, 0.01), Dipole[])
        # @test isnan(bd) == true

    end

    @testset "Show" begin
        show(t)
        show(normalise(t))
        t.info["Regularisation"] = 1.2
        show(t)
    end

    @testset "Plotting" begin
        EEG.plot(mean(t))
        EEG.plot(mean(t), min_val = 0, max_val = 50)
        EEG.plot(mean(t), elp = joinpath(dirname(@__FILE__), "../../data", "test.elp"))
        p = EEG.plot(mean(t), threshold = 24 )

        @testset "Overlay dipole" begin

            dips = find_dipoles(mean(t))
            p = plot(p, dips)
            p = plot(p, dips[1], c = :red)

        end
    end
end