text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function create_cache(::Type{IndicatorHennemannGassner}, equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()
  alpha_tmp = similar(alpha)

  A = Array{real(basis), ndims(equations)}
  indicator_threaded  = [A(undef, nnodes(basis), nnodes(basis)) for _ in 1:Threads.nthreads()]
  modal_threaded      = [A(undef, nnodes(basis), nnodes(basis)) for _ in 1:Threads.nthreads()]
  modal_tmp1_threaded = [A(undef, nnodes(basis), nnodes(basis)) for _ in 1:Threads.nthreads()]

  return (; alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded)
end

# this method is used when the indicator is constructed as for AMR
function create_cache(typ::Type{IndicatorHennemannGassner}, mesh, equations::AbstractEquations{2}, dg::DGSEM, cache)
  create_cache(typ, equations, dg.basis)
end

function (indicator_hg::IndicatorHennemannGassner)(u::AbstractArray{<:Any,4},
                                                   mesh, equations, dg::DGSEM, cache;
                                                   kwargs...)
  @unpack alpha_max, alpha_min, alpha_smooth, variable = indicator_hg
  @unpack alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded = indicator_hg.cache
  # TODO: Taal refactor, when to `resize!` stuff changed possibly by AMR?
  #       Shall we implement `resize!(semi::AbstractSemidiscretization, new_size)`
  #       or just `resize!` whenever we call the relevant methods as we do now?
  resize!(alpha, nelements(dg, cache))
  if alpha_smooth
    resize!(alpha_tmp, nelements(dg, cache))
  end

  # magic parameters
  threshold = 0.5 * 10^(-1.8 * (nnodes(dg))^0.25)
  parameter_s = log((1 - 0.0001)/0.0001)

  @threaded for element in eachelement(dg, cache)
    indicator  = indicator_threaded[Threads.threadid()]
    modal      = modal_threaded[Threads.threadid()]
    modal_tmp1 = modal_tmp1_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = indicator_hg.variable(u_local, equations)
    end

    # Convert to modal representation
    multiply_scalar_dimensionwise!(modal, dg.basis.inverse_vandermonde_legendre, indicator, modal_tmp1)

    # Calculate total energies for all modes, without highest, without two highest
    total_energy = zero(eltype(modal))
    for j in 1:nnodes(dg), i in 1:nnodes(dg)
      total_energy += modal[i, j]^2
    end
    total_energy_clip1 = zero(eltype(modal))
    for j in 1:(nnodes(dg)-1), i in 1:(nnodes(dg)-1)
      total_energy_clip1 += modal[i, j]^2
    end
    total_energy_clip2 = zero(eltype(modal))
    for j in 1:(nnodes(dg)-2), i in 1:(nnodes(dg)-2)
      total_energy_clip2 += modal[i, j]^2
    end

    # Calculate energy in higher modes
    energy = max((total_energy - total_energy_clip1) / total_energy,
                 (total_energy_clip1 - total_energy_clip2) / total_energy_clip1)

    alpha_element = 1 / (1 + exp(-parameter_s / threshold * (energy - threshold)))

    # Take care of the case close to pure DG
    if alpha_element < alpha_min
      alpha_element = zero(alpha_element)
    end

    # Take care of the case close to pure FV
    if alpha_element > 1 - alpha_min
      alpha_element = one(alpha_element)
    end

    # Clip the maximum amount of FV allowed
    alpha[element] = min(alpha_max, alpha_element)
  end

  if alpha_smooth
    apply_smoothing_2d!(alpha, alpha_tmp, dg, cache)
  end

  return alpha
end

# Diffuse alpha values by setting each alpha to at least 50% of neighboring elements' alpha
function apply_smoothing_2d!(alpha, alpha_tmp, dg, cache)
  # Copy alpha values such that smoothing is indpedenent of the element access order
  alpha_tmp .= alpha

  # Loop over interfaces
  for interface in eachinterface(dg, cache)
    # Get neighboring element ids
    left  = cache.interfaces.neighbor_ids[1, interface]
    right = cache.interfaces.neighbor_ids[2, interface]

    # Apply smoothing
    alpha[left]  = max(alpha_tmp[left],  0.5 * alpha_tmp[right], alpha[left])
    alpha[right] = max(alpha_tmp[right], 0.5 * alpha_tmp[left],  alpha[right])
  end

  # Loop over L2 mortars
  for mortar in eachmortar(dg, cache)
    # Get neighboring element ids
    lower = cache.mortars.neighbor_ids[1, mortar]
    upper = cache.mortars.neighbor_ids[2, mortar]
    large = cache.mortars.neighbor_ids[3, mortar]

    # Apply smoothing
    alpha[lower] = max(alpha_tmp[lower], 0.5 * alpha_tmp[large], alpha[lower])
    alpha[upper] = max(alpha_tmp[upper], 0.5 * alpha_tmp[large], alpha[upper])
    alpha[large] = max(alpha_tmp[large], 0.5 * alpha_tmp[lower], alpha[large])
    alpha[large] = max(alpha_tmp[large], 0.5 * alpha_tmp[upper], alpha[large])
  end

  return alpha
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function create_cache(::Type{IndicatorLöhner}, equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()

  A = Array{real(basis), ndims(equations)}
  indicator_threaded = [A(undef, nnodes(basis), nnodes(basis)) for _ in 1:Threads.nthreads()]

  return (; alpha, indicator_threaded)
end

# this method is used when the indicator is constructed as for AMR
function create_cache(typ::Type{IndicatorLöhner}, mesh, equations::AbstractEquations{2}, dg::DGSEM, cache)
  create_cache(typ, equations, dg.basis)
end

function (löhner::IndicatorLöhner)(u::AbstractArray{<:Any,4},
                                   mesh, equations, dg::DGSEM, cache;
                                   kwargs...)
  @assert nnodes(dg) >= 3 "IndicatorLöhner only works for nnodes >= 3 (polydeg > 1)"
  @unpack alpha, indicator_threaded = löhner.cache
  resize!(alpha, nelements(dg, cache))

  @threaded for element in eachelement(dg, cache)
    indicator = indicator_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = löhner.variable(u_local, equations)
    end

    estimate = zero(real(dg))
    for j in eachnode(dg), i in 2:nnodes(dg)-1
      # x direction
      u0 = indicator[i,   j]
      up = indicator[i+1, j]
      um = indicator[i-1, j]
      estimate = max(estimate, local_löhner_estimate(um, u0, up, löhner))
    end

    for j in 2:nnodes(dg)-1, i in eachnode(dg)
      # y direction
      u0 = indicator[i, j  ]
      up = indicator[i, j+1]
      um = indicator[i, j-1]
      estimate = max(estimate, local_löhner_estimate(um, u0, up, löhner))
    end

    # use the maximum as DG element indicator
    alpha[element] = estimate
  end

  return alpha
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function create_cache(::Type{IndicatorMax}, equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()

  A = Array{real(basis), ndims(equations)}
  indicator_threaded = [A(undef, nnodes(basis), nnodes(basis)) for _ in 1:Threads.nthreads()]

  return (; alpha, indicator_threaded)
end

# this method is used when the indicator is constructed as for AMR
function create_cache(typ::Type{IndicatorMax}, mesh, equations::AbstractEquations{2}, dg::DGSEM, cache)
  cache = create_cache(typ, equations, dg.basis)
end

function (indicator_max::IndicatorMax)(u::AbstractArray{<:Any,4},
                                       mesh, equations, dg::DGSEM, cache;
                                       kwargs...)
  @unpack alpha, indicator_threaded = indicator_max.cache
  resize!(alpha, nelements(dg, cache))

  @threaded for element in eachelement(dg, cache)
    indicator = indicator_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = indicator_max.variable(u_local, equations)
    end

    alpha[element] = maximum(indicator)
  end

  return alpha
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
# empty cache is default
function create_cache(::Type{IndicatorNeuralNetwork},
                      equations::AbstractEquations{2}, basis::LobattoLegendreBasis)
  return NamedTuple()
end

# cache for NeuralNetworkPerssonPeraire-type indicator
function create_cache(::Type{IndicatorNeuralNetwork{NeuralNetworkPerssonPeraire}},
                      equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()
  alpha_tmp = similar(alpha)
  A = Array{real(basis), ndims(equations)}

  @assert nnodes(basis) >= 4 "Indicator only works for nnodes >= 4 (polydeg > 2)"

  prototype = A(undef, nnodes(basis), nnodes(basis))
  indicator_threaded  = [similar(prototype) for _ in 1:Threads.nthreads()]
  modal_threaded      = [similar(prototype) for _ in 1:Threads.nthreads()]
  modal_tmp1_threaded = [similar(prototype) for _ in 1:Threads.nthreads()]

  return (; alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded)
end

# cache for NeuralNetworkRayHesthaven-type indicator
function create_cache(::Type{IndicatorNeuralNetwork{NeuralNetworkRayHesthaven}},
                      equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()
  alpha_tmp = similar(alpha)
  A = Array{real(basis), ndims(equations)}

  prototype = A(undef, nnodes(basis), nnodes(basis))
  indicator_threaded  = [similar(prototype) for _ in 1:Threads.nthreads()]
  modal_threaded      = [similar(prototype) for _ in 1:Threads.nthreads()]
  modal_tmp1_threaded = [similar(prototype) for _ in 1:Threads.nthreads()]

  network_input = Vector{Float64}(undef, 15)
  neighbor_ids= Array{Int64}(undef, 8)
  neighbor_mean = Array{Float64}(undef, 4, 3)

  return (; alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded,
            network_input, neighbor_ids, neighbor_mean)
end

# cache for NeuralNetworkCNN-type indicator
function create_cache(::Type{IndicatorNeuralNetwork{NeuralNetworkCNN}},
                      equations::AbstractEquations{2}, basis::LobattoLegendreBasis)

  alpha = Vector{real(basis)}()
  alpha_tmp = similar(alpha)
  A = Array{real(basis), ndims(equations)}

  prototype = A(undef, nnodes(basis), nnodes(basis))
  indicator_threaded  = [similar(prototype) for _ in 1:Threads.nthreads()]
  n_cnn = 4
  nodes,_ = gauss_lobatto_nodes_weights(nnodes(basis))
  cnn_nodes,_= gauss_lobatto_nodes_weights(n_cnn)
  vandermonde = polynomial_interpolation_matrix(nodes, cnn_nodes)
  network_input = Array{Float32}(undef, n_cnn, n_cnn, 1, 1)

  return (; alpha, alpha_tmp, indicator_threaded, nodes, cnn_nodes, vandermonde, network_input)
end

# this method is used when the indicator is constructed as for AMR
function create_cache(typ::Type{<:IndicatorNeuralNetwork},
                      mesh, equations::AbstractEquations{2}, dg::DGSEM, cache)
  create_cache(typ, equations, dg.basis)
end

function (indicator_ann::IndicatorNeuralNetwork{NeuralNetworkPerssonPeraire})(
    u, mesh::TreeMesh{2}, equations, dg::DGSEM, cache; kwargs...)

  @unpack indicator_type, alpha_max, alpha_min, alpha_smooth, alpha_continuous, alpha_amr, variable, network = indicator_ann

  @unpack alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded = indicator_ann.cache
  # TODO: Taal refactor, when to `resize!` stuff changed possibly by AMR?
  #       Shall we implement `resize!(semi::AbstractSemidiscretization, new_size)`
  #       or just `resize!` whenever we call the relevant methods as we do now?
  resize!(alpha, nelements(dg, cache))
  if alpha_smooth
    resize!(alpha_tmp, nelements(dg, cache))
  end

  @threaded for element in eachelement(dg, cache)
    indicator  = indicator_threaded[Threads.threadid()]
    modal      = modal_threaded[Threads.threadid()]
    modal_tmp1 = modal_tmp1_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = indicator_ann.variable(u_local, equations)
    end

    # Convert to modal representation
    multiply_scalar_dimensionwise!(modal, dg.basis.inverse_vandermonde_legendre, indicator, modal_tmp1)

    # Calculate total energies for all modes, without highest, without two highest
    total_energy = zero(eltype(modal))
    for j in 1:nnodes(dg), i in 1:nnodes(dg)
      total_energy += modal[i, j]^2
    end
    total_energy_clip1 = zero(eltype(modal))
    for j in 1:(nnodes(dg)-1), i in 1:(nnodes(dg)-1)
      total_energy_clip1 += modal[i, j]^2
    end
    total_energy_clip2 = zero(eltype(modal))
    for j in 1:(nnodes(dg)-2), i in 1:(nnodes(dg)-2)
      total_energy_clip2 += modal[i, j]^2
    end
    total_energy_clip3 = zero(eltype(modal))
    for j in 1:(nnodes(dg)-3), i in 1:(nnodes(dg)-3)
      total_energy_clip3 += modal[i, j]^2
    end

    # Calculate energy in higher modes and polynomial degree for the network input
    X1 = (total_energy - total_energy_clip1)/total_energy
    X2 = (total_energy_clip1 - total_energy_clip2)/total_energy_clip1
    X3 = (total_energy_clip2 - total_energy_clip3)/total_energy_clip2
    X4 = nnodes(dg)
    network_input = SVector(X1, X2, X3, X4)

    # Scale input data
    network_input = network_input / max(maximum(abs, network_input), one(eltype(network_input)))
    probability_troubled_cell = network(network_input)[1]

    # Compute indicator value
    alpha[element] = probability_to_indicator(probability_troubled_cell, alpha_continuous,
                                              alpha_amr, alpha_min, alpha_max)
  end

  if alpha_smooth
    apply_smoothing_2d!(alpha, alpha_tmp, dg, cache)
  end

  return alpha
end

function (indicator_ann::IndicatorNeuralNetwork{NeuralNetworkRayHesthaven})(
    u, mesh::TreeMesh{2}, equations, dg::DGSEM, cache; kwargs...)

  @unpack indicator_type, alpha_max, alpha_min, alpha_smooth, alpha_continuous, alpha_amr, variable, network = indicator_ann

  @unpack alpha, alpha_tmp, indicator_threaded, modal_threaded, modal_tmp1_threaded, network_input, neighbor_ids, neighbor_mean = indicator_ann.cache #X, network_input
  # TODO: Taal refactor, when to `resize!` stuff changed possibly by AMR?
  #       Shall we implement `resize!(semi::AbstractSemidiscretization, new_size)`
  #       or just `resize!` whenever we call the relevant methods as we do now?
  resize!(alpha, nelements(dg, cache))
  if alpha_smooth
    resize!(alpha_tmp, nelements(dg, cache))
  end

  c2e = zeros(Int, length(mesh.tree))
  for element in eachelement(dg, cache)
    c2e[cache.elements.cell_ids[element]] = element
  end

  X = Array{Float64}(undef, 3, nelements(dg, cache))

  @threaded for element in eachelement(dg, cache)
    indicator  = indicator_threaded[Threads.threadid()]
    modal      = modal_threaded[Threads.threadid()]
    modal_tmp1 = modal_tmp1_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = indicator_ann.variable(u_local, equations)
    end

    # Convert to modal representation
    multiply_scalar_dimensionwise!(modal, dg.basis.inverse_vandermonde_legendre, indicator, modal_tmp1)
    # Save linear modal coefficients for the network input
    X[1,element] = modal[1,1]
    X[2,element] = modal[1,2]
    X[3,element] = modal[2,1]
  end

  @threaded for element in eachelement(dg, cache)
    cell_id = cache.elements.cell_ids[element]

    network_input[1] = X[1,element]
    network_input[2] = X[2,element]
    network_input[3] = X[3,element]

    for direction in eachdirection(mesh.tree)
      if direction == 1 # -x
          dir = 4
      elseif direction == 2 # +x
          dir = 1
      elseif direction == 3 # -y
          dir = 3
      elseif direction == 4 # +y
          dir = 2
      end

      # Of no neighbor exists and current cell is not small
      if !has_any_neighbor(mesh.tree, cell_id, direction)
        network_input[3*dir+1] = X[1, element]
        network_input[3*dir+2] = X[2, element]
        network_input[3*dir+3] = X[3, element]
        continue
      end

      # Get Input data from neighbors
      if has_neighbor(mesh.tree, cell_id, direction)
        neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
        if has_children(mesh.tree, neighbor_cell_id) # Cell has small neighbor
          # Mean over 4 neighbor cells
          neighbor_ids[1] = mesh.tree.child_ids[1, neighbor_cell_id]
          neighbor_ids[2] = mesh.tree.child_ids[2, neighbor_cell_id]
          neighbor_ids[3] = mesh.tree.child_ids[3, neighbor_cell_id]
          neighbor_ids[4] = mesh.tree.child_ids[4, neighbor_cell_id]

          for i in 1:4
            if has_children(mesh.tree, neighbor_ids[i])
              neighbor_ids5 = c2e[mesh.tree.child_ids[1, neighbor_ids[i]]]
              neighbor_ids6 = c2e[mesh.tree.child_ids[2, neighbor_ids[i]]]
              neighbor_ids7 = c2e[mesh.tree.child_ids[3, neighbor_ids[i]]]
              neighbor_ids8 = c2e[mesh.tree.child_ids[4, neighbor_ids[i]]]

              neighbor_mean[i,1] = (X[1,neighbor_ids5] + X[1,neighbor_ids6] + X[1,neighbor_ids7] + X[1,neighbor_ids8])/4
              neighbor_mean[i,2] = (X[2,neighbor_ids5] + X[2,neighbor_ids6] + X[2,neighbor_ids7] + X[2,neighbor_ids8])/4
              neighbor_mean[i,3] = (X[3,neighbor_ids5] + X[3,neighbor_ids6] + X[3,neighbor_ids7] + X[3,neighbor_ids8])/4
            else
              neighbor_id = c2e[neighbor_ids[i]]
              neighbor_mean[i,1] = X[1,neighbor_id]
              neighbor_mean[i,2] = X[2,neighbor_id]
              neighbor_mean[i,3] = X[3,neighbor_id]
            end
          end
          network_input[3*dir+1] = (neighbor_mean[1,1] + neighbor_mean[2,1] + neighbor_mean[3,1] + neighbor_mean[4,1])/4
          network_input[3*dir+2] = (neighbor_mean[1,2] + neighbor_mean[2,2] + neighbor_mean[3,2] + neighbor_mean[4,2])/4
          network_input[3*dir+3] = (neighbor_mean[1,3] + neighbor_mean[2,3] + neighbor_mean[3,3] + neighbor_mean[4,3])/4

        else # Cell has same refinement level neighbor
          neighbor_id = c2e[neighbor_cell_id]
          network_input[3*dir+1] = X[1,neighbor_id]
          network_input[3*dir+2] = X[2,neighbor_id]
          network_input[3*dir+3] = X[3,neighbor_id]
        end
      else # Cell is small and has large neighbor
        parent_id = mesh.tree.parent_ids[cell_id]
        neighbor_id = c2e[mesh.tree.neighbor_ids[direction, parent_id]]

        network_input[3*dir+1] = X[1,neighbor_id]
        network_input[3*dir+2] = X[2,neighbor_id]
        network_input[3*dir+3] = X[3,neighbor_id]
      end
    end

    # Scale input data
    network_input = network_input / max(maximum(abs, network_input), one(eltype(network_input)))
    probability_troubled_cell = network(network_input)[1]

    # Compute indicator value
    alpha[element] = probability_to_indicator(probability_troubled_cell, alpha_continuous,
                                              alpha_amr, alpha_min, alpha_max)
  end

  if alpha_smooth
    apply_smoothing_2d!(alpha, alpha_tmp, dg, cache)
  end

  return alpha
end

function (indicator_ann::IndicatorNeuralNetwork{NeuralNetworkCNN})(
    u, mesh::TreeMesh{2}, equations, dg::DGSEM, cache; kwargs...)
  @unpack indicator_type, alpha_max, alpha_min, alpha_smooth, alpha_continuous, alpha_amr, variable, network = indicator_ann

  @unpack alpha, alpha_tmp, indicator_threaded, nodes, cnn_nodes, vandermonde, network_input = indicator_ann.cache
  # TODO: Taal refactor, when to `resize!` stuff changed possibly by AMR?
  #       Shall we implement `resize!(semi::AbstractSemidiscretization, new_size)`
  #       or just `resize!` whenever we call the relevant methods as we do now?
  resize!(alpha, nelements(dg, cache))
  if alpha_smooth
    resize!(alpha_tmp, nelements(dg, cache))
  end

  @threaded for element in eachelement(dg, cache)
    indicator  = indicator_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    for j in eachnode(dg), i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, j, element)
      indicator[i, j] = indicator_ann.variable(u_local, equations)
    end

    # Interpolate nodal data to 4x4 LGL nodes
    for j in 1:4, i in 1:4
      acc = zero(eltype(indicator))
      for jj in eachnode(dg), ii in eachnode(dg)
        acc += vandermonde[i,ii] * indicator[ii,jj] * vandermonde[j,jj]
      end
      network_input[i,j,1,1] = acc
    end

    # Scale input data
    network_input = network_input / max(maximum(abs, network_input), one(eltype(network_input)))
    probability_troubled_cell = network(network_input)[1]

    # Compute indicator value
    alpha[element] = probability_to_indicator(probability_troubled_cell, alpha_continuous,
                                              alpha_amr, alpha_min, alpha_max)
  end

  if alpha_smooth
    apply_smoothing_2d!(alpha, alpha_tmp, dg, cache)
  end

  return alpha
end

end # @muladd


text:
function plankton_update!(plank, nuts, rnd, p, plk, diags_spcs, ΔT, t, arch::Architecture, mode::QuotaMode)
    plankton_growth!(plank, nuts, rnd, p, ΔT, t, arch)

    calc_consume!(plk.DIC.data, plk.DOC.data, plk.NH4.data, plk.NO3.data, plk.PO4.data,
                  plank, plank.ac, plank.xi, plank.yi, plank.zi, ΔT, arch)
    ##### diagnostics of processes for each species
    diags_spcs!(diags_spcs, plank, plank.ac, plank.xi, plank.yi, plank.zi, mode, arch)

    ##### check the probabilities every 10 time steps or 1 hour whichever is shorter
    if t%(ΔT*(min(10,3600÷ΔT))) == 0
        ##### grazing and its diagnostic
        diags_proc!(diags_spcs.graz, plank.graz, plank.ac, plank.xi, plank.yi, plank.zi, arch)

        grazing!(plank, arch, plk, p)

        ###### mortality and its diagnostic
        diags_proc!(diags_spcs.mort, plank.mort, plank.ac, plank.xi, plank.yi, plank.zi, arch)

        mortality!(plank, arch, plk, p)

        ###### cell division diagnostic
        diags_proc!(diags_spcs.dvid, plank.dvid, plank.ac, plank.xi, plank.yi, plank.zi, arch)

        ##### division
        ##### check if the number of individuals exceeded
        dvidnum = dot(plank.dvid, plank.ac)
        deactive_ind = findall(x -> x == 0.0, plank.ac)
        if dvidnum > length(deactive_ind)
            throw(ArgumentError("number of individual exceeds the capacity"))
        end
        ##### do not copy inactive individuals
        plank.dvid .*= plank.ac
        divide!(plank, deactive_ind, arch)
    end

    ##### diagnostic for individual distribution
    diags_proc!(diags_spcs.num, plank.ac, plank.ac, plank.xi, plank.yi, plank.zi, arch)
end

function plankton_update!(plank, nuts, rnd, p, plk, nothing::Nothing, ΔT, t, arch::Architecture, mode::QuotaMode)
    plankton_growth!(plank, nuts, rnd, p, ΔT, t, arch)

    calc_consume!(plk.DIC.data, plk.DOC.data, plk.NH4.data, plk.NO3.data, plk.PO4.data,
                  plank, plank.ac, plank.xi, plank.yi, plank.zi, ΔT, arch)

    ##### check the probabilities every 10 time steps or 1 hour whichever is shorter
    if t%(ΔT*(min(10,3600÷ΔT))) == 0
        ##### grazing
        grazing!(plank, arch, plk, p)

        ###### mortality
        mortality!(plank, arch, plk, p)

        ##### division
        ##### check if the number of individuals exceeded
        dvidnum = dot(plank.dvid, plank.ac)
        deactive_ind = findall(x -> x == 0.0, plank.ac)
        if dvidnum > length(deactive_ind)
            throw(ArgumentError("number of individual exceeds the capacity"))
        end
        ##### do not copy inactive individuals
        plank.dvid .*= plank.ac
        divide!(plank, deactive_ind, arch)
    end
end

ext: jl


text:
# lightweight poly for abstract cost estimation
abstract type AbstractPoly{D,T<:Number} end
Base.one(x::AbstractPoly) = one(typeof(x))
Base.zero(x::AbstractPoly) = zero(typeof(x))

function Base.show(io::IO, p::AbstractPoly{D,T}) where {D,T<:Real}
    N = degree(p)
    for i = N:-1:0
        if i > 0
            print(io, "$(abs(p[i]))*")
            print(io, "$D")
            i > 1 && print(io, "^$i")
            print(io, p[i-1] < 0 ? " - " : " + ")
        else
            print(io, "$(abs(p[i]))")
        end
    end
end
function Base.show(io::IO, p::AbstractPoly{D,T}) where {D,T<:Complex}
    N = degree(p)
    for i = N:-1:0
        if i > 0
            print(io, "($(p[i]))*")
            print(io, "$D")
            i > 1 && print(io, "^$i")
            print(io, " + ")
        else
            print(io, "($(p[i]))")
        end
    end
end
struct Power{D,T} <: AbstractPoly{D,T}
    coeff::T
    N::Int
end
degree(p::Power) = p.N
Base.getindex(p::Power{D,T}, i::Int) where {D,T} = (i == p.N ? p.coeff : zero(p.coeff))
Power{D}(coeff::T, N::Int = 0) where {D,T} = Power{D,T}(coeff, N)

Base.one(::Type{Power{D,T}}) where {D,T} = Power{D,T}(one(T), 0)
Base.zero(::Type{Power{D,T}}) where {D,T} = Power{D,T}(zero(T), 0)

Base.convert(::Type{Power{D}}, coeff::Number) where {D} = Power{D}(coeff, 0)
Base.convert(::Type{Power{D,T}}, coeff::Number) where {D,T} = Power{D,T}(coeff, 0)
Base.convert(::Type{Power{D,T}}, p::Power{D}) where {D,T} = Power{D,T}(p.coeff, p.N)

function Base.show(io::IO,p::Power{D,T}) where {D,T}
    if p.coeff == 1
    elseif p.coeff == -1
        print(io, "-")
    elseif isa(p.coeff, Complex)
        print(io, "($(p.coeff))")
    else
        print(io, "$(p.coeff)")
    end
    p.coeff == 1 || p.coeff == -1 || p.N == 0 || print(io, "*")
    p.N == 0 && (p.coeff == 1 || p.coeff ==- 1) && print(io, "1")
    p.N > 0 && print(io, "$D")
    p.N > 1 && print(io, "^$(p.N)")
end

Base.:*(p1::Power{D}, p2::Power{D}) where {D} = Power{D}(p1.coeff*p2.coeff, degree(p1)+degree(p2))
Base.:*(p::Power{D}, s::Number) where {D} = Power{D}(p.coeff*s, degree(p))
Base.:*(s::Number, p::Power) = *(p,s)
Base.:/(p::Power{D}, s::Number) where {D} = Power{D}(p.coeff/s, degree(p))
Base.:\(s::Number, p::Power) = /(p,s)
Base.:^(p::Power{D}, n::Int) where {D} = Power{D}(p.coeff^n, n*degree(p))

struct Poly{D,T} <: AbstractPoly{D,T}
    coeffs::Vector{T}
    function Poly{D,T}(coeffs::Vector{T}) where {D,T}
        if length(coeffs) == 0 || coeffs[end] == 0
            i = findlast(!iszero, coeffs)
            return i === nothing ? new{D,T}(T[0]) : new{D,T}(coeffs[1:i])
        else
            return new{D,T}(coeffs)
        end
    end
end
degree(p::Poly) = max(0,length(p.coeffs)-1)
Base.getindex(p::Poly{D,T}, i::Int) where {D,T} = (0<=i<=degree(p) ? p.coeffs[i+1] : zero(p[0]))
Poly{D}(coeffs::Vector{T}) where {D,T} = Poly{D,T}(coeffs)
Poly{D}(c0::T) where {D,T} = Poly{D,T}([c0])
Poly{D}(p::Power{D,T}) where {D,T} = Poly{D,T}(vcat(zeros(T, p.N), p.coeff))
Poly{D,T}(c0::Number) where {D,T} = Poly{D,T}([T(c0)])
Poly{D,T1}(p::Power{D,T2}) where {D,T1,T2} = Poly{D,T1}(vcat(zeros(T1,p.N),T1(p.coeff)))

Base.one(::Type{Poly{D,T}}) where {D,T} = Poly{D,T}([one(T)])
Base.zero(::Type{Poly{D,T}}) where {D,T} = Poly{D,T}([zero(T)])

Base.convert(::Type{Poly{D}}, x::Number) where {D} = Poly{D}([x])
Base.convert(::Type{Poly{D,T}}, x::Number) where {D,T} = Poly{D,T}(T[x])
Base.convert(::Type{Poly{D}}, p::Power{D}) where {D} = Poly{D}(vcat(fill(zero(p.coeff), p.N), p.coeff))
Base.convert(::Type{Poly{D,T}}, p::Power{D}) where {D,T} = Poly{D,T}(vcat(fill(zero(T), p.N), convert(T, p.coeff)))
Base.convert(::Type{Poly{D,T}}, p::Poly{D}) where {D,T} = Poly{D,T}(convert(Vector{T}, p.coeffs))

Base.:+(p::AbstractPoly{D},s::Number) where {D} = Poly{D}([p[i]+ifelse(i==0, s, zero(s)) for i=0:degree(p)])
Base.:+(s::Number,p::AbstractPoly) = +(p,s)
Base.:+(p1::AbstractPoly{D}, p2::AbstractPoly{D}) where {D} = Poly{D}([p1[i]+p2[i] for i=0:max(degree(p1),degree(p2))])

Base.:-(p::Poly{D}) where {D} = Poly{D}(-p.coeffs)
Base.:-(p::AbstractPoly{D},s::Number) where {D} = Poly{D}([p[i]-ifelse(i==0, s, zero(s)) for i=0:degree(p)])
Base.:-(s::Number,p::AbstractPoly{D}) where {D} = Poly{D}([-p[i]+ifelse(i==0, s, zero(s)) for i=0:degree(p)])
Base.:-(p1::AbstractPoly{D}, p2::AbstractPoly{D}) where {D} = Poly{D}([p1[i]-p2[i] for i=0:max(degree(p1),degree(p2))])

Base.:*(p1::Power{D}, p2::Poly{D}) where {D} = Poly{D}([p1.coeff*p2[n-degree(p1)] for n=0:degree(p1)+degree(p2)])
Base.:*(p1::Poly{D}, p2::Power{D}) where {D} = *(p2,p1)
Base.:*(p::Poly{D}, s::Number) where {D} = Poly{D}(s*p.coeffs)
Base.:*(s::Number, p::Poly) = *(p,s)
Base.:/(p::Poly{D}, s::Number) where {D} = Poly{D}(p.coeffs/s)
Base.:\(s::Number, p::Poly) = /(p,s)
function Base.:*(p1::Poly{D}, p2::Poly{D}) where {D}
    N = degree(p1)+degree(p2)
    s = p1[0]*p2[0]
    coeffs = zeros(typeof(s), N+1)
    for i = 0:degree(p1)
        for j = 0:degree(p2)
            coeffs[i+j+1] += p1[i]*p2[j]
        end
    end
    return Poly{D}(coeffs)
end

Base.promote_rule(::Type{Power{D,T1}}, ::Type{Power{D,T2}}) where {D,T1<:Number,T2<:Number} = Power{D,promote_type(T1,T2)}
Base.promote_rule(::Type{Power{D,T1}}, ::Type{T2}) where {D,T1<:Number,T2<:Number} = Power{D,promote_type(T1,T2)}
Base.promote_rule(::Type{Poly{D,T1}}, ::Type{Poly{D,T2}}) where {D,T1<:Number,T2<:Number} = Poly{D,promote_type(T1,T2)}
Base.promote_rule(::Type{Poly{D,T1}}, ::Type{Power{D,T2}}) where {D,T1<:Number,T2<:Number} = Poly{D,promote_type(T1,T2)}
Base.promote_rule(::Type{Poly{D,T1}}, ::Type{T2}) where {D,T1<:Number,T2<:Number} = Poly{D,promote_type(T1,T2)}

function Base.:(==)(p1::AbstractPoly{D}, p2::AbstractPoly{D}) where {D}
    for i = max(degree(p1),degree(p2)):-1:0
        p1[i] == p2[i] || return false
    end
    return true
end
Base.:(==)(p1::AbstractPoly, p2::Number) = degree(p1) == 0 && p1[0] == p2
Base.:(==)(p1::Number, p2::AbstractPoly) = degree(p2) == 0 && p2[0] == p1
function Base.:<(p1::AbstractPoly{D}, p2::AbstractPoly{D}) where {D}
    for i = max(degree(p1),degree(p2)):-1:0
        p1[i] < p2[i] && return true
        p1[i] > p2[i] && return false
    end
    return false
end
Base.isless(p1::AbstractPoly{D}, p2::AbstractPoly{D}) where {D} = p1 < p2


text:
### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ 5eebac68-891b-11eb-37ba-c1d9481c6134
begin
  using Pkg; Pkg.activate("MLJ_env", shared=true)
  using PlutoUI
	using Test
	using Random
	using RDatasets
	using DataFrames
	using Plots
	# using Gadfly

	push!(LOAD_PATH, "./src")
	using YaLinearAlgebra
end

# ╔═╡ 87e5e2ec-8915-11eb-362b-6bf13a36b8e4
md"""
## k-Nearest Neightbors

ref. from book **"Data Science from Scratch"**, Chap 12

$(html"
© Pascal, April 2021
")
"""

# ╔═╡ cb25501a-8915-11eb-3626-a1ae6d233f92
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ cb0838ae-8915-11eb-082f-3991192a101f
html"""

"""

# ╔═╡ caf03ada-8915-11eb-2935-01e110accba6
html"""

"""

# ╔═╡ cad66356-8915-11eb-0173-ada200678f63
md"""
### The model

Nearest makes no mathematical assumption and only requires two elements:
  1. A distance
  1. An assumption that points close to each other (given the distance) are similar.

The prediction for a given data point depends only on a few datapoints (k of them) "close" to it.

In this notebook the data points will be vector and the distance will be the Euclidean one.

Assuming we pick a number k, we then want to classify a new data point by finding the k nearest ones and let them vote on the new ouptut.
To achieve this we need a function that counts the votes.
"""

# ╔═╡ dc3b7280-8915-11eb-258a-1789476897b9
begin
	function raw_majority_vote(labels::AbstractVector{Symbol})::Symbol
		votes = count(labels)
		most_common(votes; k=1)[1]
	end

	function count(labels::AbstractVector{Symbol})::Dict{Symbol, Integer}
		hcnt = Dict{Symbol, Integer}()
		for label ∈ labels
			hcnt[label] = get(hcnt, label, 0) + 1
		end
		hcnt
	end

	function most_common(hcnt::Dict{Symbol, Integer};
			k::Integer=1)::Pair{Symbol,Integer}
		collect(hcnt) |> l -> sort(l, by=t -> t[2]; rev=true)[k]
	end
end

# ╔═╡ dc209b9a-8915-11eb-0f63-ab3b7c33793c
begin
	h = count([:b, :a, :b, :a, :c, :b, :d])
	@test h[:b] == 3
	@test h[:a] == 2
	@test h[:c] == h[:d] == 1

	@test sort(collect(h), by=t -> t[2]; rev=true)[1] == (:b => 3)

	@test raw_majority_vote([:b, :a, :b, :a, :c, :b, :d]) == :b
end

# ╔═╡ d7aade44-891b-11eb-00d3-abe7406c4b09
md"""
We need to take into account the possibility of ties. If this happens we will reduce k until we find a winner which will happen eventually. In the worse case we will go down until k equals 1.
"""

# ╔═╡ 9ae99efe-891c-11eb-3d1f-89182e3f9ff8
function majority_vote(labels::AbstractVector{Symbol})::Symbol
	"""
	Assuming labels is ordered from nearest to farthest
	"""
	vote_counts = count(labels)
	winner, winner_cnt = most_common(vote_counts)
	num_winners = length([
		cnt for cnt ∈ values(vote_counts) if cnt == winner_cnt
	])

	num_winners == 1 && (return winner)
	majority_vote(view(labels, 1:length(labels)-1))
end

# ╔═╡ ea1f4a2c-891d-11eb-3027-bd447dd9d1e1
begin
	@test majority_vote([:b, :a, :b, :a, :c, :b, :d]) == :b
	@test majority_vote([:b, :a, :b, :a, :c, :b, :d, :a]) == :b
	@test majority_vote([:b, :a, :b, :a, :c, :a, :d, :b]) == :a
end

# ╔═╡ d7931368-891b-11eb-2899-f5e5a3f4bd35
md"""
Let us now create our classifer applying knn.
"""

# ╔═╡ 516d69f0-891e-11eb-0f7a-3f05981cc1e1
begin
	const VF = AbstractVector{T} where T <: AbstractFloat

	struct LabeledPoint
		point::VF
		label::Symbol

		LabeledPoint(point::VF, sym::Symbol) = new(point, sym)
	end

	function knn(k::Integer, dataset::Vector{LabeledPoint}, qpoint::VF)::Symbol
		## order from nearest to farthest using euclidean distance
		bydist = sort(dataset, by=lp -> distance(lp.point, qpoint), rev=false)

		## find label from k closest
		k_near_labels = map(lp -> lp.label, view(bydist, 1:k))

		## vote and return decision
		majority_vote(k_near_labels)
	end
end

# ╔═╡ 50fe5114-891e-11eb-18ee-35d98ad32b7b
md"""
Let us test this on an existing dataset.
"""

# ╔═╡ dc0c5a5e-8915-11eb-2c62-5328b42add69
html"""

"""

# ╔═╡ dbefe77a-8915-11eb-39f1-5d0c84729739
md"""
### Example with this Iris Dataset

We will use the Iris dataset which is kind of the "hello wolrd" of machine learning. It contains a set of measurements for 150 flowers representing three species of iris. For each flower we have its petal length, petal width, sepal length, and sepal width,
as well as its species.
"""

# ╔═╡ dbd86924-8915-11eb-351f-8362f09ba984
begin
	iris_ds = dataset("datasets", "iris");
	DataFrames.describe(iris_ds, :eltype, :first => first)
end

# ╔═╡ d32b86e0-8928-11eb-193c-e3d7c85d23dd
measurements_mx = convert(Matrix, select(iris_ds, Not(:Species)))

# ╔═╡ d2d8c676-8928-11eb-2e4a-557f5f9fa1c3
begin
	# [row[1:end] for row ∈ eachrow(measurements)] # copy => Vector{Vector}... slow
	# Vector{Float64}[row for row ∈ eachrow(measurements)] # view
	measurements = [row[1:end] for row ∈ eachrow(measurements_mx)]
	labels = Symbol.(iris_ds[!, :Species])
	iris_data = [
		LabeledPoint(p, l) for (p, l) ∈ zip(measurements, labels)
	]
	N = length(iris_data);
end

# ╔═╡ 9eb1ef78-8953-11eb-3046-331119a0dc9b
@test N == size(iris_ds)[1]

# ╔═╡ 5809f59e-8916-11eb-1c09-6fa00c043e7a
md"""
let us start with some visual exploration
"""

# ╔═╡ 4ffc3b48-895a-11eb-1ed2-f78b522d0ae9
begin
	points_by_species = Dict{Symbol, Vector{VF}}()
	for iris ∈ iris_data
		ary = get(points_by_species, iris.label, VF[])
		push!(ary, iris.point)
		points_by_species[iris.label] = ary
	end
	points_by_species
end

# ╔═╡ 90f58d04-8958-11eb-170e-0f17904c9f2c
begin
	metrics = names(iris_ds)[1:end-1]
	pairs = [(i, j) for i ∈ 1:4 for j ∈ (i+1):4]
	marks = [:v, :o, :^] # we have 3 classes, so 3 markers
	ps = []
	for row ∈ 1:2, col ∈ 1:3
		i, j = pairs[3 * (row - 1) + col]
		k = 0
		for (mark, (species, points)) ∈ zip(marks, points_by_species)
			xs = [p[i] for p ∈ points]
			ys = [p[j] for p ∈ points]
			if k % 3 == 0
				push!(ps, scatter(xs, ys, marker=mark,
						label=string(species),
						legend=false,
						titlefontsize=8, title="$(metrics[i]) vs $(metrics[j])"))
			elseif row == 2 && col == 3
				ps[end] = scatter!(xs, ys, marker=mark,
					label=string(species),
					legendfontsize=6,
					legend=:bottomright)
			else
				ps[end] = scatter!(xs, ys, marker=mark, label=string(species))
			end
			k += 1
		end
	end
end

# ╔═╡ 90d83cae-8958-11eb-2d5c-bd298437f953
begin
	lg = grid(2, 3, widths=[0.33, 0.33, 0.33], heights=[0.5, 0.5, 0.5])
	plot(ps..., layout=lg)
end

# ╔═╡ ab9ac56e-8969-11eb-3340-337957fd81b7
md"""
The measurement seem to really cluster by species. Looking at sepal length vs sepal width alone (top left graph) does not allow for a clean spearation between versiclor and virgina, but adding petal length and width help in the separation. Our nearest neighbors should be able to predict the species. Let us see...
"""

# ╔═╡ 5f5e6d02-892a-11eb-32d8-dfc1a7f259fa
names(iris_ds)

# ╔═╡ 073f6442-892e-11eb-0fc2-bbeb1d12619e
function train_test_split(ds::Vector{LabeledPoint};
		split=0.8, seed=42, shuffled=true)
	Random.seed!(seed)
	nr = length(ds)
	row_ixes = shuffled ? shuffle(1:nr) : collect(1:nr)
	nrp = round(Int, length(row_ixes) * split)
	(ds[row_ixes[1:nrp]], ds[row_ixes[nrp+1:nr]])
end

# ╔═╡ 25fb8f9a-8952-11eb-35bb-b173ac466835
iris_train, iris_test = train_test_split(iris_data; split=0.70)

# ╔═╡ 65ec8554-8953-11eb-3973-b56039754312
begin
	@test length(iris_train) == round(Int, 0.7 * N)
	@test length(iris_test) == round(Int, (1 - 0.7) * N)
end

# ╔═╡ 0c8055b8-89b2-11eb-3046-331119a0dc9b
length(iris_test)

# ╔═╡ d728d526-8956-11eb-3c22-e788d025e8b4
function run_knn(iris_train, iris_test; k=5)
	num_correct = 0
	confusion_matrix = Dict{Tuple{Symbol, Symbol}, Integer}()
	for iris ∈ iris_test
		ŷ = knn(k, iris_train, iris.point)
		y = iris.label
		ŷ == y && (num_correct += 1)
		confusion_matrix[(ŷ, y)] = get(confusion_matrix, (ŷ, y), 0) + 1
	end
	perc_correct = num_correct / length(iris_test)
	(perc_correct, confusion_matrix)
end

# ╔═╡ ad388bba-8955-11eb-34ee-4db60fb0db8a
perc_correct, confusion_matrix = run_knn(iris_train, iris_test)

# ╔═╡ ad036de6-89b1-11eb-3973-b56039754312
md"""
Not bad, we have two mis-matches: versicolor/virginica (and conversely), otherwise the rest looks good.
"""

# ╔═╡ dbbf22a2-8915-11eb-00eb-4b0278c0283d
html"""

"""

# ╔═╡ dba976a2-8915-11eb-0ff0-b38952cbc38b
md"""
### The Curse of Dimensionality

The k-nearest neighbors algorithm has problems in high dimensions due to a phenomena called the “curse of dimensionality”.
In a nutshell high-dimensional spaces are large (tautology) and points in such high-dimensional spaces tend to be far to one another.

One way to see this is by randomly generating pairs of points in a n-dimensional “unit cube” varying the dimensions and calculating the distances between them.
"""

# ╔═╡ ca91a0d6-8915-11eb-36cc-a18fc8efaf40
random_point(n::Integer) = rand(Float64, n)

# ╔═╡ 194958a6-89b4-11eb-3340-337957fd81b7
function random_distances(dim::Integer, num_pairs::Integer)
	[distance(random_point(dim), random_point(dim)) for _ ∈ 1:num_pairs]
end

# ╔═╡ fa1fe202-89b7-11eb-34ee-4db60fb0db8a
begin
	const NS = 100
	const M = 10_000
	avg_dist, min_dist, min_avg_ratio = [], [], []
	Random.seed!(42)
	for dim ∈ 1:NS
		d = random_distances(dim, M)
		sum_d, min_d = sum(d) / M, minimum(d)
		push!(avg_dist, sum_d)
		push!(min_dist, min_d)
		push!(min_avg_ratio, min_d / sum_d);
	end
	# avg_dist, min_dist, min_avg_ratio;
end

# ╔═╡ 432cd594-89b9-11eb-2d4a-6f46f08a511d
begin
	plot(1:NS, avg_dist, legendfontsize=7, legend=:topleft,
		label="avg. dist.",
		title="$(M) random distances",
		titlefontsize=9,
		xlabel="# of dimensions"
	)
	plot!(1:NS, min_dist, label="min. dist.")
end

# ╔═╡ d0692f3c-89bb-11eb-2c0b-15edb9044d52
plot(1:NS, min_avg_ratio, legend=false, xlabel="# of dimensions",
	title="min. dist. / avg. dist.",
	titlefontsize=9)

# ╔═╡ 2ddeec6c-89bc-11eb-0796-f7dc9885b0de
md"""
In low-dimensional datasets, the nearest points tend to be much closer than
average.

The closeness of two points depends on how close thes points are in each dimension and note that every extra dimension is an opportunity for each point to be farther away from each other.

Thus when we have a lot of dimensions, the likelihood is that the closest points are not much closer than average (unless there is a lot of structure in the data)...
"""

# ╔═╡ Cell order:
# ╟─87e5e2ec-8915-11eb-362b-6bf13a36b8e4
# ╠═5eebac68-891b-11eb-37ba-c1d9481c6134
# ╟─cb25501a-8915-11eb-3626-a1ae6d233f92
# ╟─cb0838ae-8915-11eb-082f-3991192a101f
# ╟─caf03ada-8915-11eb-2935-01e110accba6
# ╟─cad66356-8915-11eb-0173-ada200678f63
# ╠═dc3b7280-8915-11eb-258a-1789476897b9
# ╠═dc209b9a-8915-11eb-0f63-ab3b7c33793c
# ╟─d7aade44-891b-11eb-00d3-abe7406c4b09
# ╠═9ae99efe-891c-11eb-3d1f-89182e3f9ff8
# ╠═ea1f4a2c-891d-11eb-3027-bd447dd9d1e1
# ╟─d7931368-891b-11eb-2899-f5e5a3f4bd35
# ╠═516d69f0-891e-11eb-0f7a-3f05981cc1e1
# ╟─50fe5114-891e-11eb-18ee-35d98ad32b7b
# ╟─dc0c5a5e-8915-11eb-2c62-5328b42add69
# ╟─dbefe77a-8915-11eb-39f1-5d0c84729739
# ╠═dbd86924-8915-11eb-351f-8362f09ba984
# ╠═d32b86e0-8928-11eb-193c-e3d7c85d23dd
# ╠═d2d8c676-8928-11eb-2e4a-557f5f9fa1c3
# ╠═9eb1ef78-8953-11eb-3046-331119a0dc9b
# ╟─5809f59e-8916-11eb-1c09-6fa00c043e7a
# ╠═4ffc3b48-895a-11eb-1ed2-f78b522d0ae9
# ╠═90f58d04-8958-11eb-170e-0f17904c9f2c
# ╠═90d83cae-8958-11eb-2d5c-bd298437f953
# ╟─ab9ac56e-8969-11eb-3340-337957fd81b7
# ╠═5f5e6d02-892a-11eb-32d8-dfc1a7f259fa
# ╠═073f6442-892e-11eb-0fc2-bbeb1d12619e
# ╠═25fb8f9a-8952-11eb-35bb-b173ac466835
# ╠═65ec8554-8953-11eb-3973-b56039754312
# ╠═0c8055b8-89b2-11eb-3046-331119a0dc9b
# ╠═d728d526-8956-11eb-3c22-e788d025e8b4
# ╠═ad388bba-8955-11eb-34ee-4db60fb0db8a
# ╟─ad036de6-89b1-11eb-3973-b56039754312
# ╟─dbbf22a2-8915-11eb-00eb-4b0278c0283d
# ╟─dba976a2-8915-11eb-0ff0-b38952cbc38b
# ╠═ca91a0d6-8915-11eb-36cc-a18fc8efaf40
# ╠═194958a6-89b4-11eb-3340-337957fd81b7
# ╠═fa1fe202-89b7-11eb-34ee-4db60fb0db8a
# ╠═432cd594-89b9-11eb-2d4a-6f46f08a511d
# ╠═d0692f3c-89bb-11eb-2c0b-15edb9044d52
# ╟─2ddeec6c-89bc-11eb-0796-f7dc9885b0de

max_stars_repo_path: Data Science From Scratch/12_kNN.jl

max_stars_repo_name: pascal-p/julia-notebooks

max_issues_repo_path: Data Science From Scratch/12_kNN.jl

max_issues_repo_name: pascal-p/julia-notebooks

max_forks_repo_path: Data Science From Scratch/12_kNN.jl

max_forks_repo_name: pascal-p/julia-notebooks


text:
########################### Submodular Minimization ############################
# greedy algorithm for submodular minimization, or maximization??
# WARNING: before using, check min / max
struct Submodular{T, F, X, D} <: Update{T}
    f::F # set-valued input function
    x::X # set to choose from
    δ::D # bound on improvement of adding a certain element
    # function Submodular(f, x::Union{AbstractVector, AbstractSet})
    #     δ = fill(-Inf, length(x))
    #     new{eltype(x), typeof(f), typeof(x), typeof(δ)}(f, x, δ)
    # end
end
function Submodular(f, x::Union{AbstractVector, AbstractSet}, n::Int)
    δ = fill(-Inf, length(x))
    Submodular(f, x, δ)
end

# TODO: parallelization, especially for first pass
function update!(S::Submodular, x::Union{AbstractVector, AbstractSet})
    # length(x) ≥ S.n ? (return x) : nothing # only add if we haven't hit maximum cardinality
    f, δ, grid = S.f, S.δ, S.x
    ind = sortperm(δ) # sort grid according to increasing δ (decreasing in magnitude)
    permute!(δ, ind)
    permute!(grid, ind)
    fx = f(x)
    min_i = 0. # keeps track of smallest (i.e. greatest magnitude of) δ
    for (i, xs) in enumerate(grid) # parallelizing here does not seem to benefit much
        if min_i > δ[i] # if previously recorded δ is smaller (larger in magnitude) than the current minimum
            push!(x, xs) # already minimized allocation for (x, xs)?
            δ[i] = f(x) - fx
            min_i = min(min_i, δ[i])
            pop!(x)
        end
    end
    push!(x, grid[argmin(δ)]) # choose next point greedily
end

# Old code:
# function submodular_minimization(x::AbstractArray{T}, grid::AbstractArray{T},
#                             objective::Function,
#                             num_next::Int = 1) where {T}
#     num_next < 1 && error("num_next < 1")
#
#     δ = fill(-Inf, length(grid))
#     base_objective = objective(x)
#     #
#     Threads.@threads for i in 1:length(grid)
#         δ[i] = objective(vcat(x, grid[i])) - base_objective
#     end
#     next = grid[argmin(δ)] # choose point which reduces the objective most
#
#     if num_next == 1
#         return [next], δ
#     else
#         recursive_next, δ = submodular_minimization(vcat(x, next), grid,
#                                                 objective, num_next-1, δ)
#         return vcat(next, recursive_next), δ # could instead return, and spawn task to complete the recursive step
#     end
# end
#
# # posterior uncertainty sampling for batch learning of optical data
# function submodular_minimization(x::AbstractArray{T}, grid::AbstractArray{T},
#                                 objective::Function, num_next::Int,
#                                 δ::Array{<:Real}) where {T}
#     length(δ) ≠ length(grid) && error("length(δ) ≠ length(grid)")
#     num_next < 1 && error("num_next < 1")
#     # sort grid according to increasing δ (decreasing in magnitude)
#     ind = sortperm(δ)
#     δ = δ[ind]
#     grid = grid[ind]
#
#     base_objective = objective(x)
#     min_i = 0 # keeps track of smallest (i.e. greatest magnitude of) δ
#     for (i, xs) in enumerate(grid) # parallelizing here does not seem to benefit much
#         if min_i > δ[i] # if previously recorded δ is smaller (larger in magnitude) than the current minimum
#             δ[i] = objective(vcat(x, xs)) - base_objective
#             min_i = min(min_i, δ[i])
#         end
#     end
#     next = grid[argmin(δ)] # choose the point which reduces posterior uncertainty most
#
#     if num_next == 1
#         return [next], δ
#     else
#         recursive_next, δ = submodular_minimization(vcat(x, next), grid,
#                                                     objective, num_next-1, δ)
#         return vcat(next, recursive_next), δ
#     end
# end


text:
using Test
using Quiqbox
using LinearAlgebra: norm

# function nnRepulsions

nuc = ["H", "H", "O"]
coords = [[-0.7,0.0,0.0], [0.6,0.0,0.0], [0.0, 0.0, 0.0]]

@test nnRepulsions(nuc, coords) == 1*1/norm(coords[1] - coords[2]) +
                                   1*8/norm(coords[1] - coords[3]) +
                                   1*8/norm(coords[2] - coords[3])


text:
#=
    This is a script to test the hole simulation chain with all standard detector types.
    The script also produces some output plots but only very basic ones, since it
    is only to test the core functionality of the package. (So not detector specific plots)
=#

outputdir = joinpath(ENV["HOME"], "tmp/test_solidstatedetectors.jl/")
mkpath(outputdir)
@info "Test output dir: $outputdir"

@info "Loading packages"
using Plots; pyplot()
using SolidStateDetectors

T = Float32
@info "Testing now for Float32:"

plot() # creates a plot so that the plots during the following loop pop up.

key = :CGD

for key in  [:InvertedCoax, :BEGe, :Coax, :CGD, :Spherical]
# for key in keys(SSD_examples)
    @info "Now test detector type: $key"

    det = SolidStateDetector{T}(SSD_examples[key])
    S = SSD.get_coordinate_system(det)

    simulation = Simulation(det);

    apply_initial_state!(simulation, ElectricPotential)
    p = if S == :cartesian
        plot(
            plot(simulation.electric_potential, y = 0.002),
            plot(simulation.ρ, y = 0.002),
            plot(simulation.point_types, y = 0.002),
            size = (1000, 600), layout= (1, 3)
        )
    else
        plot(
            plot(simulation.electric_potential),
            plot(simulation.ρ),
            plot(simulation.point_types),
            size = (1000, 600), layout= (1, 3)
        )
    end
    savefig(joinpath(outputdir, "$(key)_0_init_setup"))

    nrefs = if key == :InvertedCoax
        0:3
    elseif key == :Spherical
        0:3
    elseif key == :CGD
        0:3
    else
        0:1
    end
    for nref in nrefs
        update_till_convergence!(simulation, ElectricPotential)
        p = if S == :cartesian
            plot(
                plot(simulation.electric_potential, y = 0.002),
                plot(simulation.ρ, y = 0.002),
                plot(simulation.point_types, y = 0.002),
                size = (1000, 600), layout= (1, 3)
            )
        else
            plot(
                plot(simulation.electric_potential),
                plot(simulation.ρ),
                plot(simulation.point_types),
                size = (1000, 600), layout= (1, 3)
            )
        end
        savefig(joinpath(outputdir, "$(key)_1_Electric_Potential_$(nref)_refinements"))
        if nref != nrefs[end]
            refine!(simulation, ElectricPotential, (1e-5, 1e-5, 1e-5), (1e-5, 1e-5, 1e-5), update_other_fields = true)
        end
        @show size(simulation.electric_potential.grid)
    end

    for contact in simulation.detector.contacts
        calculate_weighting_potential!(simulation, contact.id, max_refinements = key == :Coax ? 0 : 1, verbose = true)
    end
    wp_plots = if S != :cartesian
        [ plot(simulation.weighting_potentials[contact.id]) for contact in simulation.detector.contacts ]
    else
        [ plot(simulation.weighting_potentials[contact.id], y = 0.002) for contact in simulation.detector.contacts ]
    end
    plot( wp_plots..., size = (1000, 1000))
    savefig(joinpath(outputdir, "$(key)_2_Weighting_Potentials"))

    calculate_electric_field!(simulation)

    plot( simulation.electric_field.grid[1], simulation.electric_field.grid[3], SSD.get_electric_field_strength(simulation.electric_field)[:, div(length(simulation.electric_field.grid[2].ticks), 2), :]',
          st=:heatmap, title = "Electric Field Streng [V / m]", xlabel = "x / m", ylabel = "x / m", aspect_ratio = 1, size = (900, 900))
    savefig(joinpath(outputdir, "$(key)_3_Electric_Field_strength"))

    if S == :cylindrical
        plot_electric_field(simulation, φ=deg2rad(0), spacing = 3.0)
        savefig(joinpath(outputdir, "$(key)_3_1_Electric_Field_Lines"))
    else
        plot_electric_field(simulation, y = 0, spacing = 3.0)
        savefig(joinpath(outputdir, "$(key)_3_1_Electric_Field_Lines"))
    end

    set_charge_drift_model!(simulation, ADLChargeDriftModel())

    apply_charge_drift_model!(simulation)

    pos = if key == :InvertedCoax
        CylindricalPoint{T}[ CylindricalPoint{T}( 0.02, deg2rad(10), 0.025 ) ]
    elseif key == :CGD
        CartesianPoint{T}[ CartesianPoint{T}( 0.006, 0.005, 0.005  ) ]
    elseif key == :BEGe
        CylindricalPoint{T}[ CylindricalPoint{T}( 0.016, deg2rad(10), 0.015  ) ]
    elseif key == :Coax
        CylindricalPoint{T}[ CylindricalPoint{T}( 0.016, deg2rad(10), 0.005  ) ]
    elseif key == :Spherical
        CylindricalPoint{T}[ CylindricalPoint{T}( 0.00, deg2rad(0), 0.0  ) ]
    end
    energy_depos = T[1460]
    @assert in(pos[1], simulation.detector) "Test point $(pos[1]) not inside the detector $(key)."

    event = Event(pos, energy_depos)
    simulate!(event, simulation)

    plot(simulation.detector)
    plot!(event.drift_paths)
    savefig(joinpath(outputdir, "$(key)_4_charge_drift"))

    # signals[:, 2] *= -1
    plot(event.signals, size = (1200, 600), lw = 1.5)
    savefig(joinpath(outputdir, "$(key)_5_induced_signals"))

end

@info "Finished testing."
@info "Test output saved in: $outputdir"


text:
# Using the ZigZagBoomerang with Turing with the BouncyParticle sampler
# (The approach taken here is retrieving the likelihood function from Turing and sampling
# directly with ZigZagBoomerang and not using Turings `AbstractMCMC` )

using Turing
using ZigZagBoomerang
const ZZB = ZigZagBoomerang
using LinearAlgebra
const ∅ = nothing
using DelimitedFiles

include("plot_chain.jl") # simple visualization of chains with GLMakie

# define Turing Logit regression model
# following https://github.com/TuringLang/Turing.jl/blob/master/benchmarks/nuts/lr.jl
@model lr_nuts(x, y, σ) = begin

    N,D = size(x)

    α ~ Normal(0, σ)
    β ~ MvNormal(zeros(D), ones(D)*σ)

    for n = 1:N
        y[n] ~ BinomialLogit(1, dot(x[n,:], β) + α)
    end
end

# read data
function readlrdata()
    fname = joinpath(dirname(@__FILE__), "lr_nuts.data")
    z = readdlm(fname)
    x = z[:,1:end-1]
    y = z[:,end] .- 1
    return x, y
end
x, y = readlrdata()

# define problem
model = lr_nuts(x, y, 100.0)

# sample First with Turing and Nuts

n_samples = 1_000 # Sampling parameter settings
nuts_chain = @time sample(model, NUTS(0.65), n_samples) # (a bit frickle, sometimes adapts wrong)
# sampling took 383 s

# plot NUTS
fig2 = plot_chain(1:n_samples, collect(eachrow(dropdims(nuts_chain[nuts_chain.name_map.parameters].value.data, dims=3)) ))
save("lrnuts.png", fig2)

# sample with ZigZagBoomerang

using ForwardDiff
using ForwardDiff: Dual, value
"""
    make_gradient_and_dhessian_logp(turingmodel) -> ∇nlogp!

Gradient of negative log-likelihood and second derivative in direction of movement

Following https://github.com/TuringLang/Turing.jl/blob/master/src/core/ad.jl
"""
function make_gradient_and_dhessian_neglogp(
    model::Turing.Model,
    sampler=Turing.SampleFromPrior(),
    ctx::Turing.DynamicPPL.AbstractContext = DynamicPPL.DefaultContext()
)
    vi = Turing.VarInfo(model)

    # define function to compute log joint.
    function ℓ(θ)
        new_vi = Turing.VarInfo(vi, sampler, θ)
        model(new_vi, sampler, ctx)
        logp = Turing.getlogp(new_vi)
        return logp
    end

    return function (y, t, x, θ, args...)
        x_ = x + Dual{:hSrkahPmmC}(0.0, 1.0)*θ
        y_ = ForwardDiff.gradient(x->-ℓ(x), x_)
        y .= value.(y_)
        y, dot(θ, y_).partials[]
    end
end

∇neglogp! = make_gradient_and_dhessian_neglogp(model)

d = 1 + 24 # number of parameters
t0 = 0.0
x0 = zeros(d) # starting point sampler
θ0 = randn(d) # starting direction sampler
T = 200. # end time (similar to number of samples in MCMC)
c = 50.0 # initial guess for the bound

# define BouncyParticle sampler (has two relevant parameters)
Z = BouncyParticle(∅, ∅, # ignored
    10.0, # momentum refreshment rate
    0.95, # momentum correlation / only gradually change momentum in refreshment/momentum update
    0.0 # ignored
)

trace, final, (acc, num), cs = @time pdmp(∇neglogp!, # problem
        t0, x0, θ0, T, # initial state and duration
        ZZB.LocalBound(c), # use Hessian information
        Z; # sampler
        adapt=true, # adapt bound c
        progress=true, # show progress bar
        subsample=true # keep only samples at refreshment times
)
# took 272 s

# obtain direction change times and points of piecewise linear trace
t, x = ZigZagBoomerang.sep(trace)

# plot bouncy particle sampler
fig3 = plot_chain(t, x, false)
save("lrbouncy.png", fig3)

# check visually
# lines(mean(trace))
# lines!(mean(nuts_chain).nt[:mean])

# show both in one plot
fig4 = plot_chain(1:n_samples, collect(eachrow(dropdims(nuts_chain[nuts_chain.name_map.parameters].value.data, dims=3)) ),
    color=:red, title="Green: Bouncy Particle. Red: NUTS.")
fig4 = plot_chain!(fig4, t*n_samples/T, x, false, color=:green)
save("lrboth.png", fig4)


text:
@testset "arrayOfEveryValue()" begin
    trailer = SeaPearl.Trailer()

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(3, 4, "y", trailer)

    @test SeaPearl.arrayOfEveryValue(SeaPearl.AbstractIntVar[x, y]) == [4, 2, 3]
end

@testset "cpVertexFromIndex()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    true_idToNode = [
        SeaPearl.ConstraintVertex(model.constraints[1]),
        SeaPearl.ConstraintVertex(model.constraints[2]),
        SeaPearl.VariableVertex(x),
        SeaPearl.VariableVertex(y),
        SeaPearl.ValueVertex(2),
        SeaPearl.ValueVertex(3)
    ]

    for i in 1:6
        @test SeaPearl.cpVertexFromIndex(g, i) == true_idToNode[i]
    end
end

@testset "indexFromCpVertex()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    true_idToNode = [
        SeaPearl.ConstraintVertex(model.constraints[1]),
        SeaPearl.ConstraintVertex(model.constraints[2]),
        SeaPearl.VariableVertex(x),
        SeaPearl.VariableVertex(y),
        SeaPearl.ValueVertex(2),
        SeaPearl.ValueVertex(3)
    ]

    for i in 1:6
        @test SeaPearl.indexFromCpVertex(g, true_idToNode[i]) == i
    end
end

@testset "caracteristics" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    g = SeaPearl.CPLayerGraph(model)

    @test eltype(g) == Int64
    @test LightGraphs.edgetype(g) == LightGraphs.SimpleEdge{Int64}
    @test !LightGraphs.is_directed(SeaPearl.CPLayerGraph)
end

@testset "has_vertex()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    for i in 1:6
        @test LightGraphs.has_vertex(g, i)
    end
    @test !LightGraphs.has_vertex(g, 0)
    @test !LightGraphs.has_vertex(g, 7)
end

@testset "has_edge()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    SeaPearl.assign!(x, 2)

    @test !has_edge(g, 1, 2)
    @test !has_edge(g, 2, 1)
    @test !has_edge(g, 3, 4)

    @test has_edge(g, 1, 3)

    @test has_edge(g, 3, 5)
    @test !has_edge(g, 3, 6)
    @test has_edge(g, 4, 6)
    @test has_edge(g, 4, 5)
end

@testset "edges() & ne() & nv()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    SeaPearl.assign!(x, 2)

    @test sort(LightGraphs.edges(g); by=(e -> (e.src, e.dst))) == sort([
        LightGraphs.SimpleEdge{Int64}(1, 4),
        LightGraphs.SimpleEdge{Int64}(2, 4),
        LightGraphs.SimpleEdge{Int64}(4, 5),
        LightGraphs.SimpleEdge{Int64}(3, 5),
        LightGraphs.SimpleEdge{Int64}(4, 6),
        LightGraphs.SimpleEdge{Int64}(1, 3),
        LightGraphs.SimpleEdge{Int64}(2, 3)
    ]; by=(e -> (e.src, e.dst)))

    @test LightGraphs.ne(g) == 7
end

@testset "inneighbors()/outneightbors()" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)

    SeaPearl.assign!(x, 2)

    @test LightGraphs.inneighbors(g, 3) == [1, 2, 5]
    @test LightGraphs.inneighbors(g, 4) == [1, 2, 5, 6]
    @test LightGraphs.inneighbors(g, 6) == [4]
    @test LightGraphs.inneighbors(g, 5) == [3, 4]

    @test LightGraphs.outneighbors(g, 3) == [1, 2, 5]
    @test LightGraphs.outneighbors(g, 4) == [1, 2, 5, 6]
    @test LightGraphs.outneighbors(g, 6) == [4]
    @test LightGraphs.outneighbors(g, 5) == [3, 4]
end

@testset "CPLayerGraph => Simplegraph features" begin
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)

    x = SeaPearl.IntVar(2, 3, "x", trailer)
    y = SeaPearl.IntVar(2, 3, "y", trailer)
    SeaPearl.addVariable!(model, x)
    SeaPearl.addVariable!(model, y)
    SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
    SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

    g = SeaPearl.CPLayerGraph(model)
    sg = SimpleGraph(g)

    @test nv(sg) == 6
    @test ne(sg) == 8

    z = SeaPearl.IntVar(2, 3, "Z", trailer)
    SeaPearl.addVariable!(model, z)   #add an isolated variable

    g = SeaPearl.CPLayerGraph(model)
    sg = SimpleGraph(g)

    @test nv(sg) == 7
    @test ne(sg) == 10

    @test adjacency_matrix(g)==adjacency_matrix(sg)
end


text:
# customs names should take care to set nodedata & edgedata to appropriate values.
# nodedata:
#   nodedata[1] should be the objective value of the selected solution
#   nodedata[2] should be the 'accept' param (0 or 1), 1 if the selected solution will be accepted as the new incumbent solution
#   all other elements in this vector are optional
# edgedata:
#   all elements in this vector are optional
function metagraph(st::SearchTrace; nodedata::Vector{Symbol}=[:v_s, :accept, :it, :output], edgedata::Vector{Symbol}=[:operator, :params])
    mg = MetaDiGraph(SimpleDiGraph())
    origin = 1
    add_vertex!(mg)
    set_prop!(mg, 1, :v, st.v_init)
    set_prop!(mg, 1, :it, 0)
    objval = nodedata[1]
    accept = nodedata[2]
    for i in 1:size(st.df, 1)
        # add a node for all iterations, this will be vertex(mg, i+1)
        add_vertex!(mg)
        # every node should have all its metadata
        set_prop!(mg, i+1, :v, st.df[i, objval])
        for j in 3:length(nodedata)
            set_prop!(mg, i+1, nodedata[j], st.df[i, nodedata[j]])
        end
        # add the "parent" node in the connected graph
        add_edge!(mg, origin, i+1)
        for j in 1:length(edgedata)
                set_prop!(mg, origin, i+1, edgedata[j], st.df[i, edgedata[j]])
        end
        # if this node was accepted, it becomes the new parent!
        if st.df[i, accept] == 1
            origin = i+1
        end
    end
    set_prop!(mg, :lastaccept, st.df[size(st.df, 1), accept])
    mg
end

# for non-default DataFrame column names, please supply a dfsymbols::Vector{Symbol} list where
#   dfsymbols[1] = "accept" column name, (shows whether this solution was accepted or not during the search)
#   dfsymbols[2] = name of column with the objective values of the selected solution
#   dfsymbols[3] = name of column with the objective values of the incumbent solution
#   dfsymbols::Vector{Symbol}= [:accept, :v_s, :v_i]
# all other indices of dfsymbols will be ignored.

# accept column name can be specified using the accept param
# current solution objval column name can be specified using the curr_obj param
# incumbent solution value column name can be specified using the  inc_obj param

function to_df(mg::MetaDiGraph; accept::Symbol = :accept, curr_obj::Symbol = :v_s, inc_obj::Symbol = :v_i)
    # extract all column names edge(1, 2) into empty df to set types
    edgepairs = collect(props(mg, 1, 2))
    df_cons_edge = map(prop -> (prop[1] => Vector{typeof(prop[2])}()), edgepairs)

    # same for nodes, but collect datatype of (:v) and filter this line out to know type of inc_obj
    v_example = filter(prop -> prop[1] == :v, collect(props(mg, 2)))[1][2]

    nodepairs = copy(collect(props(mg, 2)))
    push!(nodepairs, (inc_obj => v_example))
    push!(nodepairs, (accept => get_prop(mg, :lastaccept)))

    df_cons_node = map(prop -> (prop[1] => Vector{typeof(prop[2])}()), nodepairs)

    df = DataFrame(union(df_cons_edge, df_cons_node))
    # now we should populate this.
    for vertex in outneighbors(mg, 1)
        fill!(df, mg, inc_obj, accept; cv=vertex)
    end
    rename!(df, :v => curr_obj)
    df[size(df, 1), accept] = get_prop(mg, :lastaccept)
    df
end

# cv = currentvertex
function fill!(df::DataFrame, mg::MetaDiGraph, inc_obj::Symbol, accept::Symbol; cv::Int=2)
    parent = inneighbors(mg, cv)[1]
    edgedict = copy(props(mg, parent, cv))
    nodedict = copy(props(mg, cv))
    push!(nodedict, inc_obj => props(mg, parent)[:v])
    haschildren = length(outneighbors(mg, cv)) > 0
    push!(nodedict, accept => haschildren)
    push!(df, Dict(union(edgedict, nodedict)))
    for nb in outneighbors(mg, cv)
        fill!(df, mg, inc_obj, accept; cv = nb)
    end
end

function isaccepted(mg::MetaDiGraph, vertex::Int)
    if vertex == 1
        return true
    elseif vertex == nv(mg)
        return get_prop(mg, :lastaccept)
    end
    length(outneighbors(mg, vertex)) > 0
end

max_stars_repo_name: Michiel-VL/PolyaViz

max_issues_repo_name: Michiel-VL/PolyaViz

max_forks_repo_name: Michiel-VL/PolyaViz


text:
@testset "test BasisMatrices.lookup" begin
    table1 = [1.0, 4.0]
    table2 = [1.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0]

    x = [0.5, 1.0, 1.5, 4.0, 5.5]
    x2 = [0.5, 2.0]
    @test BasisMatrices.lookup(table1, x, 0) == [0, 1, 1, 2, 2]
    @test BasisMatrices.lookup(table1, x, 1) == [1, 1, 1, 2, 2]
    @test BasisMatrices.lookup(table1, x, 2) == [0, 1, 1, 1, 1]
    @test BasisMatrices.lookup(table1, x, 3) == [1, 1, 1, 1, 1]

    @test BasisMatrices.lookup(table2, x, 0) == [0, 3, 3, 12, 12]
    @test BasisMatrices.lookup(table2, x, 1) == [3, 3, 3, 12, 12]
    @test BasisMatrices.lookup(table2, x, 2) == [0, 3, 3, 8, 8]
    @test BasisMatrices.lookup(table2, x, 3) == [3, 3, 3, 8, 8]

    @test BasisMatrices.lookup([1.0], x2, 0) == [0, 1]
    @test BasisMatrices.lookup([1.0], x2, 1) == [1, 1]
    @test BasisMatrices.lookup([1.0], x2, 2) == [0, 1]
    @test BasisMatrices.lookup([1.0], x2, 3) == [1, 1]

    # test scalar version of lookup
    x2 = collect(linspace(-2.0, 4.0, 10))
    @test [BasisMatrices.lookup(x2, -3.0, i) for i=0:3] == [0, 1, 0, 1]
    @test [BasisMatrices.lookup(x2, 5.0, i) for i=0:3] == [10, 10, 9, 9]
    @test [BasisMatrices.lookup(x2, i, 0) for i=x2] == collect(0:length(x2)-1)
    @test [BasisMatrices.lookup(x2, i, 1) for i=x2] == [1; 1:length(x2)-1]
end

@testset "test BasisMatrices._check_order" begin
    # check (::Int, ::Int) method
    @test BasisMatrices._check_order(10, 0) == fill(0, 1, 10)
    @test BasisMatrices._check_order(1, 0) == fill(0, 1, 1)

    # check (::Int, ::Vector) method
    ov = [0, 0]
    @test BasisMatrices._check_order(2, ov) == reshape(ov, 1, 2)
    @test_throws DimensionMismatch BasisMatrices._check_order(3, ov)

    # check (::Int, ::Matrix) method
    om = [0 0]
    @test BasisMatrices._check_order(2, om) == om
    @test BasisMatrices._check_order(1, om) == om'
    @test_throws DimensionMismatch BasisMatrices._check_order(3, ov)
end

@testset "test BasisMatrices.gridmake" begin
    # TODO: these tests are very incomplete. They just test that gridmake
    #       preserves input type

    for T in [Float16, Float32, Float64, Int128, Int16, Int32, Int64, Int8]
        @test eltype(BasisMatrices.gridmake(rand(T, 2), rand(T, 2))) == T
    end
end

@testset "test BasisMatrices.ckronx" begin
    # will test by constructing an interpoland, then evaluating at the nodes
    # and verifying that we get back our original function
    basis = BasisMatrices.Basis(BasisMatrices.Basis(BasisMatrices.Spline(), [-1.0, 1.0], 13, 3),
                           BasisMatrices.Basis(BasisMatrices.Spline(), [-5.0, 3.0], 18, 3))
    X, x12 = BasisMatrices.nodes(basis);

    # make up a funciton and evaluate at the nodes
    f(x1, x2) = cos(x1) ./ exp(x2)
    f(X::Matrix) = f(X[:, 1], X[:, 2])
    y = f(X)

    # fit the interpoland in Tensor form (tensor b/c using x12)
    c, bs = BasisMatrices.funfitxy(basis, x12, y);

    # verify that we are actually interpolating -- all heavy lifting in funeval
    # is done by ckronx so this is effectively testing that we wrote that
    # function properly
    @test maxabs(funeval(c, bs, [0 0]) - y) <= 1e-13
end

@testset "test row_kron" begin
    h = ["a" "b"; "c" "d"]
    z = ["1" "2" "3"; "4" "5" "6"]
    want = ["a1" "a2" "a3" "b1" "b2" "b3"; "c4" "c5" "c6" "d4" "d5" "d6"]
    @test row_kron(h, z) == want

    # now test on some bigger matrices
    a = randn(400, 3)
    b = randn(400, 5)
    out = row_kron(a, b)
    @test size(out) == (400, 15)

    rows_good = true
    for row=1:400
        rows_good &= out[row, :] == kron(a[row, :], b[row, :])
    end
    @test rows_good == true

    for i in 1:100
        A = sprandn(30, 5, 0.3)
        B = sprandn(30, 3, 0.3)

        want = row_kron(A, B)

        @test maxabs(want - row_kron(full(A), B)) == 0
        @test maxabs(want - row_kron(A, full(B))) == 0
        @test maxabs(want - row_kron(full(A), full(B))) == 0
    end
end

@testset "cdprodx" begin
    for nrow in 3:3:100
        for nb in 2:5
            b = [sprand(nrow, rand(5:13), 0.3) for _ in 1:nb]
            c = rand(prod([size(_, 2) for _ in b]))

            full_b = reduce(row_kron, b)
            want = full_b * c
            have = BasisMatrices.cdprodx(b, c)

            @test maxabs(want - have) < 1e-12

            # test RowKron object
            rk = RowKron(b...)
            @test maxabs(want - rk*c) < 1e-12

            # now test transpose
            c2 = rand(size(rk, 1))
            want2 = full_b'c2
            @test maxabs(want2 - Base.At_mul_B(rk, c2)) < 1e-12
        end
    end
end


text:
module CUSPARSE

using ..APIUtils

using ..CUDA
using ..CUDA: CUstream, cuComplex, cuDoubleComplex, libraryPropertyType, cudaDataType
using ..CUDA: libcusparse, unsafe_free!, @retry_reclaim

using CEnum

const SparseChar = Char

# core library
include("libcusparse_common.jl")
include("error.jl")
include("libcusparse.jl")
include("libcusparse_old.jl")

# low-level wrappers
include("array.jl")
include("util.jl")
include("types.jl")
include("helpers.jl")
include("wrappers.jl")

# high-level integrations
include("interfaces.jl")

# thread cache for task-local library handles
const thread_handles = Vector{Union{Nothing,cusparseHandle_t}}()

function handle()
    tid = Threads.threadid()
    if @inbounds thread_handles[tid] === nothing
        ctx = context()
        thread_handles[tid] = get!(task_local_storage(), (:CUSPARSE, ctx)) do
            handle = cusparseCreate()
            finalizer(current_task()) do task
                CUDA.isvalid(ctx) || return
                context!(ctx) do
                    cusparseDestroy(handle)
                end
            end

            handle
        end
    end
    @inbounds thread_handles[tid]
end

function __init__()
    resize!(thread_handles, Threads.nthreads())
    fill!(thread_handles, nothing)

    CUDA.atdeviceswitch() do
        tid = Threads.threadid()
        thread_handles[tid] = nothing
    end

    CUDA.attaskswitch() do
        tid = Threads.threadid()
        thread_handles[tid] = nothing
    end
end

end


text:
using CUDA, Interpolations,Plots
src = Float32[sin(i/5)+i/5 for i in 1:5000]

plot([src],seriestype=:scatter)

plot([gpu_dst],seriestype=:scatter)

idx = collect(1:0.01:41.9599)
int = interpolate(src, BSpline(Cubic(Line(OnGrid()))))
dst = similar(src, size(idx))
dst .= int.(idx)

gpu_idx = CuArray(idx)
gpu_dst = CuArray{Float32}(undef, size(idx))
gpu_src = CuArray(src)
gpu_tex = CuTexture(CuTextureArray(gpu_src); interpolation=CUDA.CubicInterpolation());

broadcast!(f1,gpu_dst, gpu_idx, Ref(gpu_tex))

gpu_dst .= f1.(gpu_idx,Ref(gpu_tex))

f1(x,y)=y[x]

f1.(gpu_idx,Ref(gpu_tex))

gpu_dst[1344]

dst[1345]

idx[1345]|>f2

f2(i)=sin(i/5)+i/5

CUDA.allowscalar(true)

max_stars_repo_name: Yangyang-Tan/FRG-Code-Collection

max_issues_repo_name: Yangyang-Tan/FRG-Code-Collection

max_forks_repo_name: Yangyang-Tan/FRG-Code-Collection


text:
using Jags, StatsPlots, Random, Distributions

ProjDir = @__DIR__
Random.seed!(3431)

y = rand(Normal(0,1),50)

Model = "
model {
      for (i in 1:length(y)) {
            y[i] ~ dnorm(mu,sigma);
      }
      mu  ~ dnorm(0, 1/sqrt(10));
      sigma  ~ dt(0,1,1) T(0, );
  }
"

monitors = Dict(
  "mu" => true,
  "sigma" => true,
  )

jagsmodel = Jagsmodel(
  name="Gaussian",
  model=Model ,
  monitor=monitors,
  ncommands=2,
  nchains=2,
  nsamples=2000,
  #deviance=true, dic=true, popt=true,
  pdir=ProjDir
  )

data = Dict{String, Any}(
  "y" => y,
)

inits = [
  Dict("mu" => 0.0,"sigma" => 1.0,
    ".RNG.name" => "base::Mersenne-Twister",
    ".RNG.seed" => 314159,
  )
]

sim = jags(jagsmodel, data, inits, ProjDir)
show(sim)
println()

#plot(sim)


text:
using Documenter
using WignerD

DocMeta.setdocmeta!(WignerD, :DocTestSetup, :(using WignerD); recursive=true)

makedocs(;
    modules=[WignerD],
    authors="Jishnu Bhattacharya",
    repo="https://github.com/jishnub/WignerD.jl/blob/{commit}{path}#L{line}",
    sitename="WignerD.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://jishnub.github.io/WignerD.jl",
        assets=String[],
    ),
    pages=[
        "WignerD" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/jishnub/WignerD.jl",
)


text:
export gradient_descent!

function gradient_descent!(x, # the initial value of the variable
	objective::Function,
	grad_objective::Function,
	params::OptParams = SimpleParams(ConstantStepSize(1), 100);
	verbose::Bool = false)

	grad = grad_objective(x)
	objval = objective(x)
	if verbose
		@printf("%10s%12s\n", "iter", "obj")
		@printf("%10d%12.4e\n", 0, objval)
	end

	for iter = 1:params.maxiters
		curstep = step(params.stepsizerule, objective, x, grad; objval = objval)
		x -= curstep*grad

		grad = grad_objective(x)
		objval, oldobjval = objective(x), objval
		verbose && iter%1==0 && @printf("%10d%12.4e\n", iter, objval)

		if stop(params, objval, oldobjval, grad)
			verbose && println("stopping criterion satisfied")
			break
		end
	end
	return x
end

max_stars_repo_name: madeleineudell/FirstOrderOptimization


text:
include("line.jl")
include("functions.jl")
include("voronoi_utils.jl")

# Push system closer to LTE, makes convergence easier.
# This is instead of developing operator splitting...
const BOOST = 2.0e9

"""
    calculate_C(atmos::Atmosphere,
                LTE_pops::Array{<:NumberDensity, 4})

Calculate collisional rates on the regular grid
"""
function calculate_C(atmos::Atmosphere,
                     LTE_pops::Array{<:NumberDensity, 4})

    nz, nx, ny, nl = size(LTE_pops)
    n_levels = nl - 1

    C = Array{Float64, 5}(undef, (n_levels+1, n_levels+1, nz, nx, ny))u"s^-1"

    # ionization
    for level = 1:n_levels
        C[level,n_levels+1,:,:,:] = Cij(level, n_levels+1, atmos.electron_density*1.0, atmos.temperature*1.0, LTE_pops)
        C[n_levels+1,level,:,:,:] = Cij(n_levels+1, level, atmos.electron_density*1.0, atmos.temperature*1.0, LTE_pops)
    end

    # bb transition
    # for l=1:n_levels-1
        # for u=(l+1):n_levels

    l = 1
    u = 2

    C[l,u,:,:,:] = Cij(l, u, atmos.electron_density*1.0, atmos.temperature*1.0, LTE_pops)
    C[u,l,:,:,:] = Cij(u, l, atmos.electron_density*1.0, atmos.temperature*1.0, LTE_pops)

        # end
    # end

    # Fill diagonal with zeros, because #undef does not like arithmetics
    for l=1:n_levels+1
        C[l,l,:,:,:] .= 0u"s^-1"
    end

    return C # ε ~ 0.1 at the upper parts, or on average, minimum should be 1e-2
end

"""
    calculate_C(sites::VoronoiSites,
                LTE_pops::Matrix{<:NumberDensity})

Calculate collisional rates on the regular grid
"""
function calculate_C(sites::VoronoiSites,
                     LTE_pops::Matrix{<:NumberDensity})

    n, nl = size(LTE_pops)
    n_levels = nl - 1

    C = Array{Float64, 3}(undef, (n_levels+1, n_levels+1, n))u"s^-1"

    # ionization
    for level = 1:n_levels
        C[level,n_levels+1,:] = Cij(level, n_levels+1, sites.electron_density*1.0, sites.temperature*1.0, LTE_pops)
        C[n_levels+1,level,:] = Cij(n_levels+1, level, sites.electron_density*1.0, sites.temperature*1.0, LTE_pops)
    end

    # bb transition
    # for l=1:n_levels-1
        # for u=(l+1):n_levels

    l = 1
    u = 2

    C[l,u,:] = Cij(l, u, sites.electron_density*1.0, sites.temperature*1.0, LTE_pops)
    C[u,l,:] = Cij(u, l, sites.electron_density*1.0, sites.temperature*1.0, LTE_pops)

        # end
    # end

    # Fill diagonal with zeros, because #undef does not like arithmetics
    for l=1:n_levels+1
        C[l,l,:] .= 0u"s^-1"
    end

    return C
end

"""
    calculate_R(atmos::Atmosphere,
                     line::HydrogenicLine,
                     J_λ::Array{<:UnitsIntensity_λ,4},
                     damping_λ::Array{<:Float64, 4},
                     LTE_pops::Array{<:NumberDensity, 4})

Calculcates the radiative rates on the regular grid
"""
function calculate_R(atmos::Atmosphere,
                     line::HydrogenicLine,
                     J_λ::Array{<:UnitsIntensity_λ,4},
                     damping_λ::Array{<:Float64, 4},
                     LTE_pops::Array{<:NumberDensity, 4})

    nz, nx, ny, nl = size(LTE_pops)
    n_levels = nl - 1

    R = Array{Float64, 5}(undef, (n_levels+1, n_levels+1, nz, nx, ny))u"s^-1"

    # ionization
    for level = 1:n_levels
        start = line.λidx[level+1]+1
        stop = line.λidx[level+2]
        σ = σic(level, line, line.λ[start:stop])
        G = Gij(level, n_levels+1, line.λ[start:stop], atmos.temperature*1.0, LTE_pops)

        R[level,n_levels+1,:,:,:] = Rij(J_λ[start:stop,:,:,:], σ, line.λ[start:stop])
        R[n_levels+1,level,:,:,:] = Rji(J_λ[start:stop,:,:,:], σ, G, line.λ[start:stop])
    end

    # bb transition
    # for l=1:n_levels-1
        # for u=(l+1):n_levels

    start = line.λidx[1]+1
    stop = line.λidx[2]

    l = 1
    u = 2

    σ = σij(l, u, line, line.λ[start:stop], damping_λ[start:stop,:,:,:])
    G = Gij(l, u, line.λ[start:stop], atmos.temperature*1.0, LTE_pops)

    R[l,u,:,:,:] = Rij(J_λ[start:stop,:,:,:], σ, line.λ[start:stop])
    R[u,l,:,:,:] = Rji(J_λ[start:stop,:,:,:], σ, G, line.λ[start:stop])

        # end
    # end

    # Fill diagonal with zeros, because #undef does not like arithmetics
    for l=1:n_levels+1
        R[l,l,:,:,:] .= 0u"s^-1"
    end

    return R
end

"""
    calculate_R(sites::VoronoiSites,
                line::HydrogenicLine,
                J_λ::Matrix{<:UnitsIntensity_λ},
                damping_λ::Matrix{<:Float64},
                LTE_pops::Matrix{<:NumberDensity})

Calculcates the radiative rates on the irregular grid
"""
function calculate_R(sites::VoronoiSites,
                     line::HydrogenicLine,
                     J_λ::Matrix{<:UnitsIntensity_λ},
                     damping_λ::Matrix{<:Float64},
                     LTE_pops::Matrix{<:NumberDensity})

    n, nl = size(LTE_pops)
    n_levels = nl - 1

    R = Array{Float64, 3}(undef, (n_levels+1, n_levels+1, n))u"s^-1"

    # ionization
    for level = 1:n_levels
        start = line.λidx[level+1]+1
        stop = line.λidx[level+2]
        σ = σic(level, line, line.λ[start:stop])
        G = Gij(level, n_levels+1, line.λ[start:stop], sites.temperature*1.0, LTE_pops)

        R[level,n_levels+1,:] = Rij(J_λ[start:stop,:], σ, line.λ[start:stop])
        R[n_levels+1,level,:] = Rji(J_λ[start:stop,:], σ, G, line.λ[start:stop])
    end

    # bb transition
    # for l=1:n_levels-1
        # for u=(l+1):n_levels

    start = line.λidx[1]+1
    stop = line.λidx[2]

    l = 1
    u = 2

    σ = σij(l, u, line, line.λ[start:stop], damping_λ[start:stop, :])
    G = Gij(l, u, line.λ[start:stop], sites.temperature*1.0, LTE_pops)

    R[l,u,:] = Rij(J_λ[start:stop,:], σ, line.λ[start:stop])
    R[u,l,:] = Rji(J_λ[start:stop,:], σ, G, line.λ[start:stop])

        # end
    # end

    # Fill diagonal with zeros, because #undef does not like arithmetics
    for l=1:n_levels+1
        R[l,l,:] .= 0u"s^-1"
    end

    return R
end

"""
    Rij(J::Array{<:UnitsIntensity_λ, 4},
        σij::Array{<:Unitful.Area, 4},
        λ::Vector{<:Unitful.Length})

Raditive rate for excitation transitions.
"""
function Rij(J::Array{<:UnitsIntensity_λ, 4},
             σij::Array{<:Unitful.Area, 4},
             λ::Vector{<:Unitful.Length})

    nλ, nz, nx, ny = size(J)
    R = Array{Float64, 3}(undef, (nz, nx, ny))u"s^-1"
    fill!(R,0.0u"s^-1")

    for l=1:(nλ-1)
        R += @. 2π/hc*((λ[l]*σij[l,:,:,:]*J[l,:,:,:] +
                     λ[l+1]*σij[l+1,:,:,:]*J[l+1,:,:,:])*(λ[l+1] - λ[l]))/1000
    end

    return R
end
function Rij(J::Matrix{<:UnitsIntensity_λ},
             σij::Matrix{<:Unitful.Area},
             λ::Vector{<:Unitful.Length})

    nλ, n = size(J)
    R = Vector{Float64}(undef, n)u"s^-1"
    fill!(R,0.0u"s^-1")

    for l=1:(nλ-1)
        R += @. 2π/hc*((λ[l]*σij[l,:]*J[l,:] +
                     λ[l+1]*σij[l+1,:]*J[l+1,:])*(λ[l+1] - λ[l]))/1000
    end

    return R
end

"""
    Rij(J::Array{<:UnitsIntensity_λ, 4},
        σij::Array{<:Unitful.Area, 1},
        λ::Vector{<:Unitful.Length})

Radiative rate for ionisation transitions.
"""
function Rij(J::Array{<:UnitsIntensity_λ, 4},
             σij::Vector{<:Unitful.Area},
             λ::Vector{<:Unitful.Length})

    nλ, nz, nx, ny = size(J)
    R = Array{Float64, 3}(undef, (nz, nx, ny))u"s^-1"
    fill!(R,0.0u"s^-1")

    for l=1:(nλ-1)
        R += @. 2π/hc*(λ[l]*σij[l]*J[l,:,:,:] +
                     λ[l+1]*σij[l+1]*J[l+1,:,:,:])*(λ[l+1] - λ[l])/1000
    end

    return R
end
function Rij(J::Matrix{<:UnitsIntensity_λ},
             σij::Vector{<:Unitful.Area},
             λ::Vector{<:Unitful.Length})

    nλ, n = size(J)
    R = Vector{Float64}(undef, n)u"s^-1"
    fill!(R,0.0u"s^-1")

    for l=1:(nλ-1)
        R += @. 2π/hc*(λ[l]*σij[l]*J[l,:] +
                     λ[l+1]*σij[l+1]*J[l+1,:])*(λ[l+1] - λ[l])/1000
    end

    return R
end

"""
    Rji(J::Array{<:UnitsIntensity_λ, 4},
        σij::Array{<:Unitful.Area, 4},
        Gij::Array{Float64, 4},
        λ::Vector{<:Unitful.Length})

Radiative rate for de-excitation transitions.
"""
function Rji(J::Array{<:UnitsIntensity_λ, 4},
             σij::Array{<:Unitful.Area, 4},
             Gij::Array{Float64, 4},
             λ::Vector{<:Unitful.Length})

    nλ, nz, nx, ny = size(J)
    R = Array{Float64, 3}(undef, (nz, nx, ny))u"s^-1"
    fill!(R,0.0u"s^-1")

    # Trapezoid rule
    for l=1:(nλ-1)
        R += @. 2π/hc*(σij[l,:,:,:]*Gij[l,:,:,:]*λ[l]*(2*h*c_0^2/λ[l]^5 + J[l,:,:,:] ) +
                    σij[l+1,:,:,:]*Gij[l+1,:,:,:]*λ[l+1]*(2*h*c_0^2 / λ[l+1]^5 + J[l+1,:,:,:] ))*(λ[l+1] - λ[l])
    end

    return R
end
function Rji(J::Matrix{<:UnitsIntensity_λ},
             σij::Matrix{<:Unitful.Area},
             Gij::Matrix{Float64},
             λ::Vector{<:Unitful.Length})

    nλ, n = size(J)
    R = Vector{Float64}(undef, n)u"s^-1"
    fill!(R,0.0u"s^-1")

    # Trapezoid rule
    for l=1:(nλ-1)
        R += @. 2π/hc*(σij[l,:]*Gij[l,:]*λ[l]*(2*h*c_0^2/λ[l]^5 + J[l,:] ) +
                    σij[l+1,:]*Gij[l+1,:]*λ[l+1]*(2*h*c_0^2 / λ[l+1]^5 + J[l+1,:] ))*(λ[l+1] - λ[l])
    end

    return R
end

"""
    Rji(J::Array{<:UnitsIntensity_λ, 4},
        σij::Array{<:Unitful.Area, 1},
        Gij::Array{Float64, 4},
        λ::Vector{<:Unitful.Length})

Radiative rate for recombination transitions.
"""
function Rji(J::Array{<:UnitsIntensity_λ, 4},
             σij::Vector{<:Unitful.Area},
             Gij::Array{Float64, 4},
             λ::Vector{<:Unitful.Length})

    nλ, nz, nx, ny = size(J)
    R = Array{Float64, 3}(undef, (nz, nx, ny))u"s^-1"
    fill!(R,0.0u"s^-1")

    # Trapezoid rule
    for l=1:(nλ-1)
        R += @. 2π/hc * (σij[l]*Gij[l,:,:,:]*λ[l]*(2*hc*c_0/λ[l]^5 + J[l,:,:,:]) +
                      σij[l+1]*Gij[l+1,:,:,:]*λ[l+1]*(2*hc*c_0/λ[l+1]^5 + J[l+1,:,:,:]))*(λ[l+1] - λ[l])
    end

    return R
end
function Rji(J::Matrix{<:UnitsIntensity_λ},
             σij::Vector{<:Unitful.Area},
             Gij::Matrix{Float64},
             λ::Vector{<:Unitful.Length})

    nλ, n = size(J)
    R = Vector{Float64}(undef, n)u"s^-1"
    fill!(R,0.0u"s^-1")

    # Trapezoid rule
    for l=1:(nλ-1)
        R += @. 2π/hc * (σij[l]*Gij[l,:]*λ[l]*(2*hc*c_0/λ[l]^5 + J[l,:]) +
                      σij[l+1]*Gij[l+1,:]*λ[l+1]*(2*hc*c_0/λ[l+1]^5 + J[l+1,:]))*(λ[l+1] - λ[l])
    end

    return R
end

"""
    σij(i::Integer,
        j::Integer,
        atom::Atom,
        λ::Vector{<:Unitful.Length})

Calculates the bound-bound cross section.
"""
function σij(i::Integer,
             j::Integer,
             line::HydrogenicLine,
             λ::Vector{<:Unitful.Length},
             damping_λ::Array{<:Float64, 4})

    λ0 = line.λ0
    σ_constant = h*c_0/(4π*λ0) * line.Bij
    nλ = length(λ)
    nz, nx, ny = size(line.ΔD)
    σ = Array{Unitful.Area, 4}(undef, nλ, nz, nx, ny)

    for l=1:nλ
        v = (λ[l] - λ0) ./ line.ΔD
        profile = voigt_profile.(damping_λ[l, :, :, :], ustrip(v), line.ΔD)
        σ[l,:,:,:] = σ_constant .* profile
    end

    return σ
end
function σij(i::Integer,
             j::Integer,
             line::HydrogenicLine,
             λ::Vector{<:Unitful.Length},
             damping_λ::Matrix{<:Float64})

    λ0 = line.λ0
    σ_constant = h*c_0/(4π*λ0) * line.Bij
    nλ = length(λ)
    n = length(line.ΔD)
    σ = Matrix{Float64}(undef, (nλ, n))u"m^2"

    for l=1:nλ
        v = (λ[l] - λ0) ./ line.ΔD
        profile = voigt_profile.(damping_λ[l, :], ustrip(v), line.ΔD)
        σ[l, :] = σ_constant .* profile
    end

    return σ
end

"""
    σic(i::Integer,
        atom::Atom,
        λ::Vector{<:Unitful.Length})

Calculates the bound-free cross-section.
"""
function σic(i::Integer,
             line::HydrogenicLine,
             λ::Vector{<:Unitful.Length},
             λ_edge=nothing)

    if λ_edge == nothing
        λ_edge = λ[end]
    end
    λ3_ratio = (λ ./ λ_edge).^3
    n_eff = sqrt(E_∞ / (line.χj - line.χi)) |>u"J/J" # should be χu - χl
    charge = line.Z
    σ_constant = (4 * e^2 / (3 * π * sqrt(3) * ε_0 * m_e * c_0^2 * R_∞)) |> u"m^2"

    σ = (σ_constant * charge^4 * n_eff * λ3_ratio .* gaunt_bf.(λ, charge, n_eff))

    return σ
end

"""
    Gij(i::Integer,
        j::Integer,
        λ::Vector{<:Unitful.Length},
        temperature::Array{<:Unitful.Temperature, 3},
        LTE_populations::Array{<:NumberDensity, 4})

Factor to go into the de-excitation and recombination expressions.
"""
function Gij(i::Integer,
             j::Integer,
             λ::Vector{<:Unitful.Length},
             temperature::Array{<:Unitful.Temperature, 3},
             LTE_populations::Array{<:NumberDensity, 4})

    nλ = length(λ)
    nz, nx, ny = size(temperature)
    G = Array{Float64, 4}(undef, nλ, nz, nx, ny)

    n_ratio = LTE_populations[:,:,:,i] ./LTE_populations[:,:,:,j]

    for l=1:nλ
        G[l,:,:,:] = @. n_ratio * exp(- hc / (k_B*λ[l]*temperature))
    end

    return G
end
function Gij(i::Integer,
             j::Integer,
             λ::Vector{<:Unitful.Length},
             temperature::Vector{<:Unitful.Temperature},
             LTE_populations::Matrix{<:NumberDensity})

    nλ = length(λ)
    n = length(temperature)
    G = Matrix{Float64}(undef, (nλ, n))

    n_ratio = LTE_populations[:,i] ./LTE_populations[:,j]

    for l=1:nλ
        G[l,:] = @. n_ratio * exp(- hc / (k_B*λ[l]*temperature))
    end

    return G
end

"""
    Cij(i::Integer,
        j::Integer,
        electron_density::Array{<:NumberDensity,3},
        temperature::Array{<:Unitful.Temperature,3},
        LTE_populations::Array{<:NumberDensity,4})

Calculates the collisional rates for all possible
two-level atom transitions.
"""
function Cij(i::Integer,
             j::Integer,
             electron_density::Array{<:NumberDensity,3},
             temperature::Array{<:Unitful.Temperature,3},
             LTE_populations::Array{<:NumberDensity,4})

    ionisation_level = size(LTE_populations)[end]

    # If UP
    if i < j
        if j < ionisation_level
            C = coll_exc_hydrogen_johnson.(i, j, electron_density, temperature)
        elseif j == ionisation_level
            C = coll_ion_hydrogen_johnson.(i, electron_density, temperature)
        end

    # If DOWN
    elseif i > j
        if i < ionisation_level
            C = coll_exc_hydrogen_johnson.(j, i, electron_density, temperature)
        elseif i == ionisation_level
            C = coll_ion_hydrogen_johnson.(j, electron_density, temperature)
        end
        C = @. C * ( LTE_populations[:,:,:,j] / LTE_populations[:,:,:,i] )
    end

    return C.*BOOST
end
function Cij(i::Integer,
             j::Integer,
             electron_density::Array{<:NumberDensity,1},
             temperature::Array{<:Unitful.Temperature,1},
             LTE_populations::Array{<:NumberDensity,2})

    ionisation_level = size(LTE_populations)[end]

    # If UP
    if i < j
        if j < ionisation_level
            C = coll_exc_hydrogen_johnson.(i, j, electron_density, temperature)
        elseif j == ionisation_level
            C = coll_ion_hydrogen_johnson.(i, electron_density, temperature)
        end

    # If DOWN
    elseif i > j
        if i < ionisation_level
            C = coll_exc_hydrogen_johnson.(j, i, electron_density, temperature)
        elseif i == ionisation_level
            C = coll_ion_hydrogen_johnson.(j, electron_density, temperature)
        end
        C = @. C * ( LTE_populations[:,j] / LTE_populations[:,i] )
    end

    return C.*BOOST
end

"""
    gaunt_bf(charge::Int, n_eff::Number, λ::Unitful.Length)::Float64

Compute bound-free Gaunt factor for a given charge, effective principal
quantum number and wavelength λ. Taken from RH. Formula from
[Seaton (1960), Rep. Prog. Phys. 23, 313](https://ui.adsabs.harvard.edu/abs/1960RPPh...23..313S/abstract),
page 316. Recipes from Seaton. (This is copied from github.com/tiagopereira/Transparency.jl)
"""
function gaunt_bf(λ::Unitful.Length,
                  charge::Real,
                  n_eff::Real)::Float64

    x = ustrip(1 / (λ * R_∞ * charge^2) |> u"m/m")
    x3 = x^(1/3)
    nsqx = 1 / (n_eff^2 * x)
    g_bf = 1 + 0.1728 * x3 * (1 - 2 * nsqx) - 0.0496 * x3^2 * (1 - (1 - nsqx) * 0.66666667 * nsqx)
    @assert g_bf >= 0 "gaunt_bf negative, calculation will not be reliable"
    return g_bf
end

max_stars_repo_name: meudnaes/VoronoiRT

max_issues_repo_name: meudnaes/VoronoiRT

max_forks_repo_name: meudnaes/VoronoiRT


text:
module Solver

using JuMP
using Gurobi
using LinearAlgebra

export randomsubsetsum

const GRB_ENV = Gurobi.Env()

function randomsubsetsum(z; ε::Real, n::Real, minsecs::Real = Inf, verbose::Bool = false)
    @assert 0 < ε
    @assert n > 0
    @assert minsecs ≥ 0

    n = ceil(Int, n)
    # The random subset
    A = 2rand(1, n) .- 1
    if verbose
        @show ε n z
        flush(stdout)
    end

    model = Model(() -> Gurobi.Optimizer(GRB_ENV))
    set_optimizer_attribute(model, "Seed", 4231)
    verbose || set_optimizer_attribute(model, "OutputFlag", 0)

    @variable(model, x[1:n], Bin)
    @variable(model, t)
    @objective(
        model,
        Min,
        t
    )
    @constraint(model, c1, (z .- (A * x)) .≤ t)
    @constraint(model, c2, -(z .- (A * x)) .≤ t)
    @constraint(model, c3, t ≤ ε)
    # print(model)

    # Search for `minsecs` seconds (keep going even if a solution is found)
    if minsecs > 0
        set_optimizer_attribute(model, "TimeLimit", minsecs)
        optimize!(model)
    end
    # If no solution has been found, keep searching for as long as necessary
    if !has_values(model)
        set_optimizer_attribute(model, "TimeLimit", Inf)
        set_optimizer_attribute(model, "SolutionLimit", 1)
        optimize!(model)
    end
    # If still no solution has been found, restart without optimizations that could have discarded feasable solutions
    if termination_status(model) == MOI.INFEASIBLE_OR_UNBOUNDED
        @info "Needs restart without `DualReductions`"
        set_optimizer_attribute(model, "DualReductions", 0)
        optimize!(model)
    end

    # If a solution is possible, it should be available by now
    if has_values(model)
        result = sum(value.(x) .* eachcol(A))
        verbose && @show(sum(value.(x)), maximum(abs, z .- result), z, result)
        return convert(eltype(z), first(result))
    end
    # At this point we can trust the infeasability
    if termination_status(model) == MOI.INFEASIBLE
        @warn "There is no solution for this instance. Termination status: $(raw_status(model))"
        @show ε n z A
    else
        @error "Something went wrong with the solver. Termination status: $(raw_status(model))"
    end
    @info "Returning zeros"
    flush(stdout); flush(stderr)
    return zero(z)
end

end

max_stars_repo_name: ArthurWalraven/cnnslth

max_issues_repo_name: ArthurWalraven/cnnslth

max_forks_repo_name: ArthurWalraven/cnnslth