text:
add_verbose_scope(:Automorphisms)

################################################################################
#
#  Automorphisms
#
################################################################################

function _automorphisms(K::NfAbsNS; isabelian::Bool = false)
  pols = fmpq_poly[isunivariate(x)[2] for x in K.pol]
  rt = Vector{Vector{NfAbsNSElem}}(undef, length(pols))
  for i = 1:length(pols)
    rt[i] = roots(pols[i], K)
  end
  auts = Vector{NfAbsNSToNfAbsNS}(undef, prod(length(x) for x in rt))
  ind = 1
  I = cartesian_product_iterator([1:length(x) for x in rt], inplace = true)
  for i in I
    auts[ind] = hom(K, K, elem_type(K)[rt[i[j]] for j = 1:length(pols)])
    ind += 1
  end
  return auts
end

function _automorphisms(K::AnticNumberField; isabelian::Bool = false)
  if degree(K) == 1
    return NfToNfMor[hom(K, K, one(K))]
  end
  if Nemo.iscyclo_type(K)
    f = get_special(K, :cyclo)::Int
    a = gen(K)
    A, mA = unit_group(ResidueRing(FlintZZ, f, cached = false))
    auts = NfToNfMor[ hom(K, K, a^lift(mA(g)), check = false) for g in A]
    return auts
  end
  if isabelian
    return _automorphisms_abelian(K)
  end
  c = get_special(K, :isabelian)
  if c !== nothing && c
    return _automorphisms_abelian(K)
  end
  f = K.pol
  ord_aut = _order_bound(K)
  if ord_aut == 1
    return NfToNfMor[id_hom(K)]
  end
  #=
  e, q = ispower(ord_aut)
  if e == 2 && isprime(q)
    return _automorphisms_center(K)
  end
  =#
  Kt, t = PolynomialRing(K, "t", cached = false)
  f1 = change_base_ring(K, f, parent = Kt)
  divpol = Kt(nf_elem[-gen(K), K(1)])
  f1 = divexact(f1, divpol)
  lr = roots(f1, max_roots = div(ord_aut, 2))
  Aut1 = Vector{NfToNfMor}(undef, length(lr)+1)
  for i = 1:length(lr)
    Aut1[i] = hom(K, K, lr[i], check = false)
  end
  Aut1[end] = id_hom(K)
  auts = closure(Aut1, degree(K))
  return auts
end

function _order_bound(K::AnticNumberField)
  p = 101
  i = 0
  ord = degree(K)
  isnormal = true
  while i < 15 && !isone(ord)
    p = next_prime(p)
    F = GF(p, cached = false)
    Fx, x = PolynomialRing(F, cached = false)
    fF = Fx(K.pol)
    if degree(fF) != degree(K) || iszero(discriminant(fF))
      continue
    end
    i += 1
    sh = factor_shape(fF)
    if length(sh) != 1
      isnormal = false
    end
    if haskey(sh, 1)
      ord = gcd(ord, sh[1])
    end
  end
  if ord == degree(K) && !isnormal
    lf = factor(ord)
    divs = collect(keys(lf))
    ord = div(ord, min(divs))
  end
  return ord
end

function _generator_automorphisms(K::AnticNumberField)
  if degree(K) == 1
    return NfToNfMor[]
  end
  if Nemo.iscyclo_type(K)
    f = get_special(K, :cyclo)::Int
    a = gen(K)
    A, mA = unit_group(ResidueRing(FlintZZ, f, cached = false))
    auts = NfToNfMor[ hom(K, K, a^lift(mA(g)), check = false) for g in gens(A)]
    return auts
  end
  f = K.pol
  Kt, t = PolynomialRing(K, "t", cached = false)
  f1 = change_base_ring(K, f, parent = Kt)
  divpol = Kt(nf_elem[-gen(K), K(1)])
  f1 = divexact(f1, divpol)
  lr = roots(f1, max_roots = div(degree(K), 2))
  Aut1 = Vector{NfToNfMor}(undef, length(lr))
  for i = 1:length(lr)
    Aut1[i] = hom(K, K, lr[i], check = false)
  end
  return small_generating_set(Aut1)
end

automorphism_type(::AnticNumberField) = NfToNfMor
automorphism_type(::NfAbsNS) = NfAbsNSToNfAbsNS

@doc Markdown.doc"""
    automorphisms(K::AnticNumberField) -> Vector{NfToNfMor}

Returns the set of automorphisms of K
"""
function automorphisms(K::NumField{fmpq}; copy::Bool = true, isabelian::Bool = false)
  T = automorphism_type(K)
  if isautomorphisms_known(K)
    Aut = get_automorphisms(K)
    if copy
      v = Vector{T}(undef, length(Aut))
      for i = 1:length(v)
        v[i] = Aut[i]
      end
      return v
    else
      return Aut::Vector{T}
    end
  end
  auts = _automorphisms(K, isabelian = isabelian)
  set_automorphisms(K, auts)
  if copy
    v = Vector{T}(undef, length(auts))
    for i = 1:length(v)
      v[i] = auts[i]
    end
    return v
  else
    return auts
  end
end

function isautomorphisms_known(K::AnticNumberField)
  return _get_automorphisms_nf(K, false) != nothing
end

function isautomorphisms_known(K::NfAbsNS)
  return get_special(K, :automorphisms) !== nothing
end

function get_automorphisms(K::AnticNumberField)
  return _get_automorphisms_nf(K)::Vector{NfToNfMor}
end

function get_automorphisms(K::NfAbsNS)
  return get_special(K, :automorphisms)::Vector{NfAbsNSToNfAbsNS}
end

function set_automorphisms(K::AnticNumberField, auts::Vector{NfToNfMor})
  _set_automorphisms_nf(K, auts)
  return nothing
end

function set_automorphisms(K::NfAbsNS, auts::Vector{NfAbsNSToNfAbsNS})
  set_special(K, :automorphisms => auts)
  return nothing
end

################################################################################
#
#  Automorphism Group
#
################################################################################

@doc Markdown.doc"""
    automorphism_group(K::AnticNumberField) -> GenGrp, GrpGenToNfMorSet

Given a number field $K$, this function returns a group $G$ and a map from $G$ to the automorphisms of $K$.
"""
function automorphism_group(K::AnticNumberField)
  if Nemo.iscyclo_type(K)
    return _automorphism_group_cyclo(K)
  else
    return _automorphism_group_generic(K)
  end
end

function _automorphism_group_cyclo(K)
  f = get_special(K, :cyclo)
  a = gen(K)
  A, mA = unit_group(ResidueRing(FlintZZ, f))
  G, AtoG, GtoA = generic_group(collect(A), +)
  aut = NfToNfMor[ hom(K, K, a^lift(mA(GtoA[g])), check = false) for g in G]
  _set_automorphisms_nf(K, aut)
  return G, GrpGenToNfMorSet(G, aut, K)
end

function _automorphism_group_generic(K)
  aut = automorphisms(K)
  n = degree(K)
  #First, find a good prime
  p = 11
  d = numerator(discriminant(K.pol))
  while mod(d, p) == 0
    p = next_prime(p)
  end
  R = GF(p, cached = false)
  Rx, x = PolynomialRing(R, "x", cached = false)
  fmod = Rx(K.pol)
  pols = gfp_poly[Rx(image_primitive_element(g)) for g in aut]
  Dcreation = Vector{Tuple{gfp_poly, Int}}(undef, length(pols))
  for i = 1:length(pols)
    Dcreation[i] = (pols[i], i)
  end
  D = Dict{gfp_poly, Int}(Dcreation)
  mult_table = Array{Int, 2}(undef, length(aut), length(aut))
  for s = 1:length(aut)
    for i = 1:length(aut)
      mult_table[s, i] = D[Hecke.compose_mod(pols[s], pols[i], fmod)]
    end
  end
  G = GrpGen(mult_table)
  return G, GrpGenToNfMorSet(G, aut, K)
end

###############################################################################
#
#  NfToNfMor closure
#
###############################################################################

function closure(S::Vector{NfToNfMor}, final_order::Int = -1)

  K = domain(S[1])
  d = numerator(discriminant(K.pol))
  p = 11
  while mod(d, p) == 0
    p = next_prime(p)
  end
  R = GF(p, cached = false)
  Rx, x = PolynomialRing(R, "x", cached = false)
  fmod = Rx(K.pol)

  t = length(S)
  order = 1
  elements = NfToNfMor[id_hom(K)]
  pols = gfp_poly[x]
  gpol = Rx(image_primitive_element(S[1]))
  if gpol != x
    push!(pols, gpol)
    push!(elements, S[1])
    order += 1

    gpol = compose_mod(gpol, pols[2], fmod)

    while gpol != x
      order = order +1
      push!(elements, S[1]*elements[end])
      push!(pols, gpol)
      gpol = compose_mod(gpol, pols[2], fmod)
    end
  end

  if order == final_order
    return elements
  end

  for i in 2:t
    if !(S[i] in elements)
      pi = Rx(image_primitive_element(S[i]))
      previous_order = order
      order = order + 1
      push!(elements, S[i])
      push!(pols, Rx(image_primitive_element(S[i])))
      for j in 2:previous_order
        order = order + 1
        push!(pols, compose_mod(pols[j], pi, fmod))
        push!(elements, elements[j]*S[i])
      end
      if order == final_order
        return elements
      end
      rep_pos = previous_order + 1
      while rep_pos <= order
        for k in 1:i
          s = S[k]
          po = Rx(image_primitive_element(s))
          att = compose_mod(pols[rep_pos], po, fmod)
          if !(att in pols)
            elt = elements[rep_pos]*s
            order = order + 1
            push!(elements, elt)
            push!(pols, att)
            for j in 2:previous_order
              order = order + 1
              push!(pols, compose_mod(pols[j], att, fmod))
              push!(elements, elements[j] *elt)
            end
            if order == final_order
              return elements
            end
          end
        end
        rep_pos = rep_pos + previous_order
      end
    end
  end
  return elements
end

function generic_group(G::Vector{NfToNfMor}, ::typeof(*), full::Bool = true)
  K = domain(G[1])
  n = length(G)
  #First, find a good prime
  p = 11
  d = numerator(discriminant(K.pol))
  while mod(d, p) == 0
    p = next_prime(p)
  end
  R = GF(p, cached = false)
  Rx, x = PolynomialRing(R, "x", cached = false)
  fmod = Rx(K.pol)
  pols = gfp_poly[Rx(image_primitive_element(g)) for g in G]
  Dcreation = Vector{Tuple{gfp_poly, Int}}(undef, length(pols))
  for i = 1:length(pols)
    Dcreation[i] = (pols[i], i)
  end
  D = Dict{gfp_poly, Int}(Dcreation)
  #full && @assert length(D) == degree(K)
  permutations = Array{Array{Int, 1},1}(undef, n)

  m_table = Array{Int, 2}(undef, n, n)

  for s = 1:n
    for i = 1:n
      m_table[s, i] =  D[Hecke.compose_mod(pols[s], pols[i], fmod)]
    end
  end

  Gen = GrpGen(m_table)
  GentoG = Dict{GrpGenElem, eltype(G)}(Gen[i] => G[i] for i in 1:length(G))
  GtoGen = Dict{eltype(G), GrpGenElem}(G[i] => Gen[i] for i in 1:length(G))
  return Gen, GtoGen, GentoG
end

################################################################################
#
#  Automorphisms abelian fields
#
################################################################################

function _automorphisms_abelian(K::AnticNumberField)

  #@assert isabelian(K)
  auts = NfToNfMor[id_hom(K)]
  p = 2
  dp = denominator(K.pol)
  while length(auts) != degree(K)
    p = next_prime(p)
    if divisible(dp, p)
      continue
    end
    F = GF(p, cached = false)
    Fx = PolynomialRing(F, cached = false)[1]
    fF = Fx(K.pol)
    if degree(fF) != degree(K) || iszero(discriminant(fF))
      continue
    end
    @vprint :Automorphisms 1 "Trying $p \n"
    @vtime :Automorphisms 1 isnew, h = _frobenius_at(K, p, auts)
    if !isnew
      @vprint :Automorphisms "Not new! \n"
      continue
    end
    @vprint :Automorphisms "New! Computing closure \n"
    push!(auts, h)
    auts = closure(auts)
  end
  return auts
end

function lift_root(K::AnticNumberField, b::nmod_poly, bound::Int)
  Fx = parent(b)
  fF = Fx(K.pol)
  Zx = PolynomialRing(FlintZZ, "x")[1]
  p = modulus(Fx)
  test = 2^10
  dfF = derivative(fF)
  dF = K(derivative(K.pol))
  evdfF = compose_mod(dfF, b, fF)
  w = invmod(evdfF, fF)
  b_0 = lift(Zx, b)
  w_0 = lift(Zx, w)
  #Now, the lifting
  r_old = one(K)
  modu = fmpz(p)^2
  R = ResidueRing(FlintZZ, modu, cached = false)
  Rx = PolynomialRing(R, "x", cached = false)[1]
  fR = map_coeffs(R, Zx(K.pol), parent = Rx)
  Rb_0 = Rx(b_0)
  Rw_0 = Rx(w_0)
  bi = compose_mod(fR, Rb_0, fR)
  bi = mulmod(Rw_0, bi, fR)
  bi = sub!(bi, Rb_0, bi)
  b_0 = lift(Zx, bi)
  r = K(parent(K.pol)(b_0))
  r_old = r + 1
  mul!(r, r, dF)
  mod_sym!(r, modu)
  i = 0
  while i < bound && r != r_old && !check_root(K, test, r)
    i += 1
    modu = modu^2
    R = ResidueRing(FlintZZ, modu, cached = false)
    Rx = PolynomialRing(R, "x", cached = false)[1]
    fR = Rx(K.pol)
    Rb_0 = Rx(b_0)
    Rw_0 = Rx(w_0)
    @vtime :Automorphisms 4 A, fRinv = precomp_compose_mod(Rb_0, fR)
    @vtime :Automorphisms 4 wi = compose_mod_precomp(derivative(fR), A, fR, fRinv)
    wi = wi * Rw_0
    wi = 2 - wi
    wi = mulmod(wi, Rw_0, fR)
    @vtime :Automorphisms 4 bi = my_compose_mod_precomp(fR, A, fR, fRinv)
    bi = mulmod(bi, wi, fR)
    bi = Rb_0 - bi
    b_0 = lift(Zx, bi)
    w_0 = lift(Zx, wi)
    r = K(parent(K.pol)(b_0))
    r = mul!(r, r, dF)
    r = mod_sym!(r, modu)
  end
  if i == bound
    if check_root(K, test, r)
      return true, divexact(r, dF)
    else
      return false, gen(K)
    end
  end
  return true, divexact(r, dF)
end

function _frobenius_at(K::AnticNumberField, p::Int, auts::Vector{NfToNfMor} = NfToNfMor[]; bound::Int = 100)

  Zx = FlintZZ["x"][1]
  F = ResidueRing(FlintZZ, p, cached = false)
  Fx, gFx = PolynomialRing(F, "x", cached = false)
  fF = map_coeffs(F, Zx(K.pol), parent = Fx)
  b = powmod(gFx, p, fF)
  if b in nmod_poly[Fx(image_primitive_element(x)) for x in auts]
    return false, id_hom(K)
  end
  fl, rt = lift_root(K, b, bound)
  if fl
    return true, hom(K, K, rt, check = false)
  else
    return false, id_hom(K)
  end
end

function _coefficients_bound(K::AnticNumberField)
  r1, r2 = signature(K)
  bound_root = Vector{arb}(undef, r1 + r2)
  a = gen(K)
  dfa = K(derivative(K.pol))
  dfa_conjs = conjugates_arb(dfa, 32)
  RR = ArbField(64, cached = false)
  RRt, t = PolynomialRing(RR, "t", cached = false)
  ub_f = roots_upper_bound(RRt(K.pol))
  for i in 1:r1+r2
    bound_root[i] = ub_f * abs(dfa_conjs[i])
  end
  E = EquationOrder(K)
  c1, c2 = norm_change_const(E)

  #First, t2 norm
  R = parent(bound_root[1])
  bd = zero(R)
  for i in 1:r1
    bd += bound_root[i]^2
  end

  for i=1:r2
    bd += 2*bound_root[i+r1]^2
  end
  boundt2 = max(bd, one(R))
  return upper_bound(sqrt(R(c2)*boundt2), fmpz)
end

function check_root(K::AnticNumberField, p::Int, el::nf_elem)
  isroot = true
  cnt = 0
  q = p
  while cnt < 10
    q = next_prime(q)
    F = GF(q, cached = false)
    Fx = PolynomialRing(F, cached = false)[1]
    fF = Fx(K.pol)
    if degree(fF) != degree(K) || iszero(discriminant(fF))
      continue
    end
    cnt += 1
    img_el = Fx(el)*invmod(derivative(fF), fF)
    if !iszero(compose_mod(fF, img_el, fF))
      isroot = false
      break
    end
  end
  if !isroot
    @vprint :Automorphisms 4 "Not yet a root!\n"
  end
  return isroot
end

function _automorphisms_center(K::AnticNumberField)
  auts = NfToNfMor[id_hom(K)]
  p = 2
  dp = denominator(K.pol)
  coeffs_bound = 2*_coefficients_bound(K)
  cnt = 0
  ord = _order_bound(K)
  threshold = max(60, ord^2)
  while length(auts) < ord && cnt < threshold
    cnt += 1
    p = next_prime(p)
    if divisible(dp, p)
      continue
    end
    F = GF(p, cached = false)
    Fx = PolynomialRing(F, cached = false)[1]
    fF = Fx(K.pol)
    if degree(fF) != degree(K) || iszero(discriminant(fF))
      continue
    end
    lf = factor_distinct_deg(fF)
    if length(lf) != 1
      continue
    end
    it_bound = clog(fmpz(clog(coeffs_bound, p)), 2)
    @vprint :Automorphisms "Trying $p \n"
    isnew, h = _frobenius_at(K, p, auts, bound = it_bound)
    if !isnew
      @vprint :Automorphisms "Not new! \n"
      continue
    end
    cnt = 0
    @vprint :Automorphisms "New! Computing closure \n"
    push!(auts, h)
    auts = closure(auts)
  end
  return auts
end

function isabelian2(K::AnticNumberField)
  if isautomorphisms_known(K)
    return isabelian(automorphism_group(K)[1])
  end
  auts = NfToNfMor[id_hom(K)]
  p = 2
  dp = denominator(K.pol)
  coeffs_bound = 2*_coefficients_bound(K)
  while length(auts) != degree(K)
    p = next_prime(p)
    if divisible(dp, p)
      continue
    end
    F = GF(p, cached = false)
    Fx, gFx = PolynomialRing(F, cached = false)
    fF = Fx(K.pol)
    if degree(fF) != degree(K) || iszero(discriminant(fF))
      continue
    end
    lf = factor_distinct_deg(fF)
    if length(lf) != 1
      return false
    end
    it_bound = clog(fmpz(clog(coeffs_bound, p)), 2)
    @vprint :Automorphisms 1 "Trying $p \n"
    b = powmod(gFx, p, fF)
    if b in nmod_poly[Fx(x(gen(K))) for x in auts]
      continue
    end
    fl, rt = lift_root(K, b, it_bound)
    if !fl
      return false
    end
    cnt = 0
    @vprint :Automorphisms 1 "New! Computing closure \n"
    push!(auts, hom(K, K, rt, check = false))
    auts = closure(auts)
  end
  Hecke._set_automorphisms_nf(K, auts)
  return true
end


text:
struct VegaBumpCluster
  factorBegin::Int
  factorEnd::Int
  rateBegin::Int
  rateEnd::Int
  stepBegin::Int
  stepEnd::Int
end

struct VegaBumpCollection{M <: AbstractMarketModel}
  allBumps::Vector{VegaBumpCluster}
  associatedVolStructure::M
  checked::Bool
  nonOverlapped::Bool
  isFull::Bool
end

function VegaBumpCollection(volStructure::AbstractMarketModel, factorwiseBumping::Bool)
  steps = volStructure.numberOfSteps
  rates = volStructure.numberOfRates
  factors = volStructure.numberOfFactors

  allBumps = VegaBumpCluster[]

  @inbounds @simd for s = 1:steps
    for r = volStructure.evolution.firstAliveRate[s]:rates
      if factorwiseBumping
        for f = 1:factors
          push!(allBumps, VegaBumpCluster(f, f, r, r, s, s))
        end
      else
        push!(allBumps, VegaBumpCluster(1, factors, r, r, s, s))
      end
    end
  end

  return VegaBumpCollection(allBumps, volStructure, true, true, true)
end

number_of_bumps(vegaCol::VegaBumpCollection) = length(vegaCol.allBumps)
associated_model(vegaCol::VegaBumpCollection) = vegaCol.associatedVolStructure

ext: jl


text:
module MRForecaster

using HDF5
using Distributions, StatsBase
using Printf

# Constants
# Model range
const MODEL_MIN_MASS = 3e-4
const MODEL_MIN_LOGMASS = -3.522
const MODEL_MAX_MASS = 3e5
const MODEL_MAX_LOGMASS = 5.477
const MODEL_MIN_RADIUS = 1e-1
const MODEL_MAX_RADIUS = 1e2
const N_POPS = 4

# Conversion factors
const MEARTH_MJUP = 317.828
const MEARTH_MSUN = 333060.4
const REARTH_RJUP = 11.21
const REARTH_RSUN = 109.2

# Load hyperparameters
datapath = joinpath(dirname(pathof(MRForecaster)), "..", "data")
hyper_file = joinpath(datapath, "fitting_parameters.h5")
all_hyper = h5read(hyper_file, "hyper_posterior")

include("utils.jl")

"""
    mass_to_radius(mass::AbstractVector, unit, classify)

Forecast a radius distribution given a mass distribution.
"""
function mass_to_radius(mass::Union{AbstractVector{<:AbstractFloat}, AbstractFloat};
                        unit="Earth", classify=false)
    # Convert mass units
    if unit == "Jupiter"
        mass *= MEARTH_MJUP
    elseif unit == "Sun" || unit == "Solar"
        mass*= MEARTH_MSUN
    elseif unit != "Earth"
        throw(ArgumentError(unit, "Mass unit must be Jupiter, Sun, or Earth."))
    end

    # Ensure we are within model mass range
    if (minimum(mass) < MODEL_MIN_MASS) || (maximum(mass) > MODEL_MAX_MASS)
        throw(ArgumentError(mass, "Mass outside model range"))
    end

    logm = log10.(mass)
    sample_size = length(mass)
    logr = Vector{Float64}(undef, sample_size)

    # Draw random realizations of the posterior
    hyper_params = all_hyper[:, rand(1:size(all_hyper)[2], sample_size)]

    if classify
        classification(logm, hyper_params[end-2:end,:])
    end

    # Add random dispersion around mean
    prob = rand(sample_size)

    for i = 1:sample_size
        logr[i] = piece_linear(hyper_params[:,i], logm[i], prob[i])[1]
    end

    radius = 10 .^ logr

    # Convert back to right unit
    if unit == "Jupiter"
        radius /= REARTH_RJUP
    elseif unit == "Sun" || unit == "Solar"
        radius /= REARTH_RSUN
    end

    return radius

end # mass_to_radius

"""
    Mstat2R(mass_mean::AbstractFloat, mass_std::AbstractFloat)

Forecast a radius mean and upper/lower uncertainties given a mass
mean and standard deviation.

Assumes a normal distribution with the mean and standard deviation,
truncated at the mass range limit of the model.
"""
function Mstat2R(mass_mean::AbstractFloat,
                 mass_std::AbstractFloat,
                 unit="Earth", sample_size::Int=1000, classify::Bool=false)
    # Convert mass units
    if unit == "Jupiter"
        mass_mean *= mearth2mjup
        mass_std *= mearth2mjup
    elseif unit == "Sun" || unit == "Solar"
        mass_mean *= mearth2msun
        mass_std *= mearth2msun
    elseif unit != "Earth"
        throw(ArgumentError(unit, "Mass unit must be Jupiter, Sun, or Earth."))
    end

    # Draw random samples
    mass_dist = TruncatedNormal(mass_mean, mass_std, MODEL_MIN_MASS, MODEL_MAX_MASS)
    mass_rvs = rand(mass_dist, sample_size)

    radius = mass_to_radius(mass_rvs, unit="Earth", classify=classify)

    # Convert to original units
    if unit == "Jupiter"
        radius /= REARTH_RJUP
    elseif unit == "Sun" || unit == "Solar"
        radius /= REARTH_RSUN
    end

    r_stat = quantile(radius, [0.5, 0.841, 0.159])

    return r_stat[1], r_stat[2:3] .- r_stat[1]
end # Mstat2R

"""
    radius_to_mass(radius::AbstractVector, unit, grid_size, classify)

Forecast a mass given a radius or radius distribution.
"""
function radius_to_mass(radius::Union{AbstractVector{<:AbstractFloat}, AbstractFloat};
                        unit="Earth", grid_size::Int=1000, classify::Bool=false)

    # Convert radius units
    if unit == "Jupiter"
        radius *= REARTH_RJUP
    elseif unit == "Sun" || unit == "Solar"
        radius *= REARTH_RSUN
    elseif unit != "Earth"
        throw(ArgumentError(unit, "Radius unit must be Jupiter, Sun, or Earth."))
    end

    # Ensure we are within model radius range
    if (minimum(radius) < MODEL_MIN_RADIUS) || (maximum(radius) > MODEL_MAX_RADIUS)
        throw(ArgumentError(string(radius), "Radius outside model range"))
    end
    # Generate a sample grid
    if grid_size < 10
        throw(ArgumentError(string(grid_size), "Grid size must be at least 10."))
    end

    logr = log10.(radius)
    sample_size = length(radius)
    logm = Vector{Float64}(undef, sample_size)

    # Draw random realizations of the posterior
    hyper_params = all_hyper[:, rand(1:size(all_hyper)[2], sample_size)]

    logm_grid = LinRange(MODEL_MIN_LOGMASS, MODEL_MAX_LOGMASS, grid_size)

    for i = 1:sample_size
        prob = prob_R_given_M(logr[i], logm_grid, hyper_params[:,i])
        logm[i] = sample(logm_grid, Weights(prob))
    end

    mass = 10 .^ logm

    # Convert back to right unit
    if unit == "Jupiter"
        mass /= MEARTH_MJUP
    elseif unit == "Sun" || unit == "Solar"
        mass /= MEARTH_MSUN
    end

    return mass

end # radius_to_mass

"""
    Rstat2M(radius_mean, radius_std)

Forecast a mass mean and upper/lower uncertainties given a radius mean
and standard deviation.

Assumes a normal distribution with the mean and standard deviation,
truncated at the radius range limit of the model.
"""
function Rstat2M(radius_mean::AbstractFloat,
                 radius_std::AbstractFloat,
                 unit="Earth", sample_size::Int=1000, grid_size::Int=1000,
                 classify::Bool=false)
    # Convert radius units
    if unit == "Jupiter"
        radius_mean *= REARTH_RJUP
        radius_std *= REARTH_RJUP
    elseif unit == "Sun" || unit == "Solar"
        radius_mean *= REARTH_RSUN
        radius_std *= REARTH_RSUN
    elseif unit != "Earth"
        throw(ArgumentError(unit, "Radius unit must be Jupiter, Sun, or Earth."))
    end

    # Draw random samplse
    radius_dist = TruncatedNormal(radius_mean, radius_std, MODEL_MIN_RADIUS, MODEL_MAX_RADIUS)
    radius_rvs = rand(radius_dist, sample_size)

    mass = radius_to_mass(radius_rvs, unit="Earth", classify=classify)

    # Convert back to original units
    if unit == "Jupiter"
        mass /= MEARTH_MJUP
    elseif unit == "Sun" || unit == "Solar"
        mass /= MEARTH_MSUN
    end

    m_stat = quantile(mass, [0.5, 0.841, 0.159])

    return m_stat[1], m_stat[2:3] .- m_stat[1]
end # Rstat2M

"""
    classification(logm, trans)

Classify the planet according to its mass. Units should be Earth units.
"""
function classification(logm::Union{AbstractVector{<:AbstractFloat}, AbstractFloat},
                        trans::AbstractArray{<:AbstractFloat})
    println("Classify")
    count = Vector{Float64}(undef, N_POPS)
    sample_size = length(logm)

    for iclass in 1:N_POPS
        for isample in 1:sample_size
            ind = indicate(logm[isample], trans[:,isample], iclass)
            count[iclass] += ind
        end
    end

    prob = count / sum(count) * 100.
    @printf("Terran %.1f%%, Neptunian %.1f%%, Jovian %.1f%%, Star %.1f%%\n",
            prob[1], prob[2], prob[3], prob[4])
    return
end

end # module


text:
module Libfive

using libfive_jll

include("wrapper.jl") # clang.jl autogenerated

end


text:
function circlecenter(coords1, coords2, radius)
    # two circle intersection -> https://math.stackexchange.com/a/1367732
    dist = norm(coords1-coords2)
    midpoint = (coords1+coords2)/2
    normaldist = sqrt(radius^2 - dist^2/4)
    normalvec = reverse((coords1 - coords2)) .* [1, -1]
    normalvec /= norm(normalvec)
    center = midpoint .+ sign(radius) .* normalvec .* normaldist

    return center
end

function circlesection(coords1, coords2, center, u)
    leg1 = coords1 .- center
    leg2 = coords2 .- center

    γ = angle(vec(leg1), vec(leg2))

    x(au) = leg1[1] * cos(au*γ) - leg1[2] * sin(au*γ) + center[1]
    y(au) = leg1[1] * sin(au*γ) + leg1[2] * cos(au*γ) + center[2]

    return (x.(u), y.(u))
end

function boundaryfun(boundary::StraightBoundary, vertices, inverse_direction=false)
    coords1 = vertices[boundary.vertice_start]
    coords2 = vertices[boundary.vertice_end]

    connection = coords2 - coords1

    if inverse_direction
        return coords2, u -> coords2 - connection .* u
    else
        return coords1, u -> coords1 + connection .* u
    end
end

function boundaryfun(boundary::CircularBoundary, vertices, inverse_direction=false)
    coords1 = vertices[boundary.vertice_start]
    coords2 = vertices[boundary.vertice_end]

    center = circlecenter(coords1, coords2, boundary.radius)

    leg1 = coords1 .- center
    leg2 = coords2 .- center

    γ = angle(vec(leg1), vec(leg2))

    T(u) = [cos(u*γ) -sin(u*γ); sin(u*γ) cos(u*γ)]

    if inverse_direction
        return coords2, u -> T(1-u) * leg1 + center
    else
        return coords1, u -> T(u) * leg1 + center
    end
end


text:
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using InteractiveDynamics, Random
using DynamicalBilliards
import GLMakie
using GLMakie: to_color, RGBf0, RGBAf0
using DynamicalSystems

# Set style to book colors
InteractiveDynamics.obcolor(::Antidot) = to_color(COLORS[3])
InteractiveDynamics.obcolor(::Obstacle) = to_color(COLORS[3])
InteractiveDynamics.obfill(::Antidot) = RGBAf0(0,0,0,0)
InteractiveDynamics.obls(::Antidot) = nothing

bd = Billiard(Antidot(Float32[0, 0], 1.0f0, false))

# make particles at the rim
js = [(0.2, sind(30)), (1.5, sind(30 + sqrt(2)/10))]

colors = [to_color(COLORS[i]) for i in 1:length(js)]

ps = Particle{Float32}[]
for (ξ, s) in js
    x, y = from_bcoords(Float32(ξ), Float32(s), bd)
    push!(ps, Particle(x..., atan(y[2], y[1])))
end

fig, bmapax = billiard_bmap_plot(bd, ps;
colors = colors, tail = 100000, steps = 100003, backgroundcolor = RGBf0(1,1,1),
ms = 10, vr = 0.1, tailwidth = 2.5, arrowsize = 15, particle_size = 1.5)
GLMakie.ylims!(bmapax, 0.2, 0.9)

# %% Chaotic billiard (sinai)
InteractiveDynamics.obcolor(::Antidot) = to_color(COLORS[1])
InteractiveDynamics.obfill(::Antidot) = RGBAf0(0,0,0,0)
InteractiveDynamics.obls(::Antidot) = nothing

colors = [to_color(COLORS[i]) for i in (5, 1,2)]

bd = billiard_sinai(0.25, 1, 1)

ps = particlebeam(0.2, 0.75, π + π/4 + 0.414235, 100, 0.002)
fig, bmapax = billiard_bmap_plot(bd, ps; colors = colors,
tail = 3100, steps = 3400, backgroundcolor = RGBf0(1,1,1),
ms = 10, vr = 0.05)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
"""
	mutable struct MultiEntry <: JSONEntry
		childs::Vector{Any}
	end

support for JSON which does not adhere to a fixed type.
Container for multiple types of entry which are observed on the same place in JSON.
"""
mutable struct MultiEntry <: JSONEntry
	childs::Vector{JSONEntry}
	updated::Int
end

MultiEntry(s::MultiEntry) = s
MultiEntry(s::Vector{T}) where {T<:MultiEntry} = only(s)

Base.getindex(s::MultiEntry, k::Int) = s.childs[k]
Base.setindex!(s::MultiEntry, i, k::Int) = s.childs[k] = i
Base.isempty(e::MultiEntry) = isempty(e.childs)
Base.keys(e::MultiEntry) = collect(1:length(e.childs))

function update!(s::MultiEntry, d; path = "")
	s.updated += 1
	for (i, c) in enumerate(s.childs)
		update!(c, d, path="$path[$i]") && return true
	end
	new_child = safe_update!(newentry(d), d, path=path)
	push!(s.childs, new_child)
	return true
end

findfirst_with_el(testf::Function, A) = (idx = findfirst(testf, A); (idx, A[idx]))
tryparse_entry_keys(e::Entry, T::Type) = Entry(Dict(tryparse(T, k)=>v for (k, v) in e.counts), e.updated)
function merge_entries_with_cast(e::MultiEntry, concrete_type, abstract_type)
	leaf_entries = filter(x->x isa Entry, e.childs)
	if any(is_numeric_entry.(leaf_entries, abstract_type)) && any(is_castable.(leaf_entries, concrete_type))
		e_new = MultiEntry(copy(e.childs), e.updated)
		to_merge_idx, to_merge = findfirst_with_el(x->is_castable(x, concrete_type), e_new.childs)
		to_merge_into_idx, to_merge_into = findfirst_with_el(x->is_numeric_entry(x, abstract_type), e_new.childs)
		new_entry = tryparse_entry_keys(to_merge, concrete_type)
		merged_entry = merge(to_merge_into, new_entry)
		deleteat!(e_new.childs, sort([to_merge_idx, to_merge_into_idx]))
		push!(e_new.childs, merged_entry)
		return e_new
	end
	e
end

# todo: benchmark on large cuckoo schemas, and optimize if needed
function suggestextractor(e::MultiEntry, settings = NamedTuple(); path = "", child_less_than_parent = false)
	# consolidation of types, type wrangling of numeric strings takes place here
	# trying to unify types and create new child entries for them. Merging string + numbers
	e = merge_entries_with_cast(e, Int32, Real)
	e = merge_entries_with_cast(e, FloatType, Real)
	# we need to filter out empty things in multientry too, same manner as dict
	ks = filter(k->!isempty(e.childs[k]), keys(e.childs))
	# child extractors of multi representation will always gene incompatible type which is treated as missing
	# otherwise there would not be the need for MultiRepresentation at all
	# that's why we enforce true here
	MultipleRepresentation(map(k -> suggestextractor(e.childs[k], settings,
			path = path,
			# in case of numberic strings and numbers I may have full samples so it should not be uniontype in such case
			child_less_than_parent = child_less_than_parent || length(e.childs) > 1
		),ks))
end

function merge(es::MultiEntry...)
	updates_merged = sum(updated.(es))
	entries = childs.(es) |> Iterators.flatten |> collect
	entry_types = entries .|> typeof |> unique
	merged_childs = [merge(filter(x->x isa t, entries)...) for t in entry_types]
	MultiEntry(merged_childs, updates_merged)
end

# this is merging of different types, merging them
function merge(es::JSONEntry...)
	multi_entries = filter(e->e isa MultiEntry, es)
	multi_entry = if isempty(multi_entries)
		MultiEntry([], 0)
	else
		merge(multi_entries...)
	end

	other_entries = filter(e->!(e isa MultiEntry), es)
	updates_merged = sum(updated.(other_entries))
	multi_entry.updated += updates_merged
	entry_types = other_entries .|> typeof |> unique
	merged_entries = [merge(filter(x->x isa t, other_entries)...) for t in entry_types]
	multi_entry_types = map(typeof, multi_entry.childs)
	for e in merged_entries
		idx = findfirst(x->e isa x, multi_entry_types)
		if isnothing(idx)
			push!(multi_entry.childs, e)
		else
			multi_entry.childs[idx] = merge(multi_entry.childs[idx], e)
		end
	end
	multi_entry
end

childs(s::T) where {T<:MultiEntry} = s.childs
Base.hash(e::MultiEntry, h::UInt) = hash((e.childs, e.updated), h)
Base.:(==)(e1::MultiEntry, e2::MultiEntry) = e1.updated === e2.updated && e1.childs == e2.childs
# because we have all the structure inside types, I don't actually need to extract values for all children
sample_synthetic(e::MultiEntry) = sample_synthetic(first(e.childs))


text:
import JSON

function make_json(PARAMDAT)
    fff=[]
    open(PARAMDAT, "r") do f
        fff=read(f,String)
    end

    fff=replace(fff,r"#(.+)\n" => "\n")

    read_PARAMDAT="read_"*PARAMDAT
    open(read_PARAMDAT, "w") do f
        write(f,fff)
    end
    return read_PARAMDAT
end

function read_json(read_PARAMDAT)
    dict=1
    open(read_PARAMDAT, "r") do f
        dicttxt = read(f,String)  # file information to string
        dict=JSON.parse(dicttxt)  # parse and transform data
    end
    return dict
end

function read_para(dict)
    # 出力情報
    every_outnum=Int(parse(Float64,dict["every_outnum"]))      # 出力毎タイムステップ
    out_file_front=dict["out_file_front"]
    out_ext=dict["out_ext"]

    # リスタート
    restart_file = dict["Restart"]
    restart_num  = dict["restartnum"]
    restart_ext  = dict["in_ext"]
    restart_file = restart_file*restart_num*restart_ext
    restart_num  = Int(parse(Float64,restart_num))

    # -- time step --
    nt=Int(parse(Float64,dict["nt"]))                    # 時間ステップ数
    dt=parse(Float64,dict["dt"])                    # 時間刻み幅

    # 物性値
    lam1 = parse(Float64,dict["lam1"])
    rho1 = parse(Float64,dict["rho1"])
    c1 = parse(Float64,dict["c1"])

    # 初期値
    init_T=parse(Float64,dict["init_val"])

    # 境界条件
    #=
    bdcon[i][j]
    i:境界条件番号(1,2,3..)
    j=1:境界条件番号（0or1）
    j=2:rho
    j=3:u
    j=4:v
    j=5:p
    =#
    bdcon = Any[]
    k=1
    while true
        try
            bd=[parse(Int,dict["bd"*string(k)*"_con"]),parse(Float64,dict["bd"*string(k)*"_val"])]
            push!(bdcon,bd)
        catch
            break
        end
        k+=1
    end

    return nt,dt,init_T,lam1,rho1,c1,bdcon,every_outnum,out_file_front,out_ext,restart_file,restart_num
end

function input_para(PARAMDAT)
   read_PARAMDAT=make_json(PARAMDAT)
   dict=read_json(read_PARAMDAT)
   nt,dt,init_T,lam1,rho1,c1,bdcon,every_outnum,out_file_front,out_ext,restart_file,restart_num=read_para(dict)
   println("fin read para")
   return nt,dt,init_T,lam1,rho1,c1,bdcon,every_outnum,out_file_front,out_ext,restart_file,restart_num
end


text:
if basename(pwd()) == "aoc"
    cd("2016/5")
end

using MD5

function part1(input::AbstractString)
    result = []
    for i in Iterators.countfrom(0)
        hash = bytes2hex(md5(input * string(i)))
        if startswith(hash, "00000")
            push!(result, hash[6])
            if length(result) == 8
                return join(result)
            end
        end
    end
end

@assert part1("abc") == "18f47a30"
part1("uqwqemis")

function part2(input::AbstractString)
    result = repeat([' '], 8)
    for i in Iterators.countfrom(0)
        hash = bytes2hex(md5(input * string(i)))
        if startswith(hash, "00000")
            i = parse(Int, hash[6], base = 16)
            if i < 8 && isspace(result[i+1])
                result[i+1] = hash[7]
                if all(!isspace, result)
                    return join(result)
                end
            end
        end
    end
end
@assert part2("abc") == "05ace8e3"
part2("uqwqemis")

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
function gradient_wrt_input(model, input, output_indices)
    return only(gradient((in) -> model(in)[output_indices], input))
end

function gradients_wrt_batch(model, input::AbstractArray{T,N}, output_indices) where {T,N}
    # To avoid computing a sparse jacobian, we compute individual gradients
    # by calling `gradient_wrt_input` on slices of the input along the batch dimension.
    out = similar(input)
    inds_before_N = ntuple(Returns(:), N - 1)
    for (i, ax) in enumerate(axes(input, N))
        view(out, inds_before_N..., ax, :) .= gradient_wrt_input(
            model, view(input, inds_before_N..., ax, :), drop_batch_index(output_indices[i])
        )
    end
    return out
end

"""
    Gradient(model)

Analyze model by calculating the gradient of a neuron activation with respect to the input.
"""
struct Gradient{C<:Chain} <: AbstractXAIMethod
    model::C
    Gradient(model::Chain) = new{typeof(model)}(Flux.testmode!(check_output_softmax(model)))
end
function (analyzer::Gradient)(input, ns::AbstractNeuronSelector)
    output = analyzer.model(input)
    output_indices = ns(output)
    grad = gradients_wrt_batch(analyzer.model, input, output_indices)
    return Explanation(grad, output, output_indices, :Gradient, Nothing)
end

"""
    InputTimesGradient(model)

Analyze model by calculating the gradient of a neuron activation with respect to the input.
This gradient is then multiplied element-wise with the input.
"""
struct InputTimesGradient{C<:Chain} <: AbstractXAIMethod
    model::C
    function InputTimesGradient(model::Chain)
        return new{typeof(model)}(Flux.testmode!(check_output_softmax(model)))
    end
end
function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
    output = analyzer.model(input)
    output_indices = ns(output)
    attr = input .* gradients_wrt_batch(analyzer.model, input, output_indices)
    return Explanation(attr, output, output_indices, :InputTimesGradient, Nothing)
end

"""
    SmoothGrad(analyzer, [n=50, std=0.1, rng=GLOBAL_RNG])
    SmoothGrad(analyzer, [n=50, distribution=Normal(0, σ²=0.01), rng=GLOBAL_RNG])

Analyze model by calculating a smoothed sensitivity map.
This is done by averaging sensitivity maps of a `Gradient` analyzer over random samples
in a neighborhood of the input, typically by adding Gaussian noise with mean 0.

# References
[1] Smilkov et al., SmoothGrad: removing noise by adding noise
"""
SmoothGrad(model, n=50, args...) = NoiseAugmentation(Gradient(model), n, args...)

"""
    IntegratedGradients(analyzer, [n=50])
    IntegratedGradients(analyzer, [n=50])

Analyze model by using the Integrated Gradients method.

# References
[1] Sundararajan et al., Axiomatic Attribution for Deep Networks
"""
IntegratedGradients(model, n=50) = InterpolationAugmentation(Gradient(model), n)


text:
module ENN

export TimedAutomata, TimePetriNets, Neurons

include(joinpath("TimedAutomata/","TimedAutomata.jl"))
include(joinpath("TimePetriNets/","TimePetriNets.jl"))
include(joinpath("Neurons/","Neurons.jl"))

end


text:
using WaterLily
using LinearAlgebra: norm2

"""
    circle(n, m; Re=250)

This function is copy from [WaterLily](https://github.com/weymouth/WaterLily.jl)
"""
function circle(n, m; Re=250)
    # Set physical parameters
    U, R, center = 1., m/8., [m/2, m/2]
    ν = U * R / Re

    body = AutoBody((x,t) -> norm2(x .- center) - R)
    Simulation((n+2, m+2), [U, 0.], R; ν, body)
end

function gen_data(ts::AbstractRange)
    n, m = 3(2^5), 2^6
    circ = circle(n, m)

    𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
    for (i, t) in enumerate(ts)
        sim_step!(circ, t)
        𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
    end

    return 𝐩s
end

function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Float64=0.95, batchsize=100)
    data = gen_data(ts)

    n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))

    𝐱_train, 𝐲_train = data[:, :, :, 1:(n_train-1)], data[:, :, :, 2:n_train]
    loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)

    𝐱_test, 𝐲_test = data[:, :, :, (end-n_test+1):(end-1)], data[:, :, :, (end-n_test+2):end]
    loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)

    return loader_train, loader_test
end


text:
export DiscontinuePipelineStage

"""
    DiscontinuePipelineStage <: AbstractPipelineStage

Applies the `discontinue()` operator from `Rocket.jl` library to the given pipeline

See also: [`AbstractPipelineStage`](@ref), [`apply_pipeline_stage`](@ref), [`EmptyPipelineStage`](@ref), [`CompositePipelineStage`](@ref)
"""
struct DiscontinuePipelineStage <: AbstractPipelineStage end

apply_pipeline_stage(::DiscontinuePipelineStage, factornode, tag, stream) = stream |> discontinue()


text:
import Interact

"""
    manipulate(f::Function; parameters, config=())

Create an interactive plot updated by callback `f`. Only works in Jupyter Notebook.

# Arguments
- `f::Function`: callback for generating a plot; interactively updated configuration `c` is provided.
- `parameters`: parameters adjustable with interactive widgets; value should be an iterable.
- `config=()`: a baseline configuration.
"""
manipulate(f::Function; parameters, config=()) = begin
    P = configure(parameters)
    C = configure(config)
    W = []
    L = []
    for (s, Q) in P
        n = Interact.node(:div, string(s))
        #HACK: use similar style/color (:light_magenta) to Config
        l = Interact.style(n, "font-family" => "monospace", "color" => :darkorchid)
        push!(L, l)
        for (k, V) in Q
            label = string(k)
            v = option(C, s, k)
            kw = ismissing(v) ? (; label) : (; label, value=v)
            #TODO: deunitfy() based on known units from parameters()
            w = Interact.widget(V; kw...)
            #HACK: use similar style/color (:light_blue) to Config
            d = w.layout(w).children[1].dom
            d.props[:style] = Dict("font-family" => "monospace", "width" => "25%")
            d.children[1].children[1].props[:style]["color"] = :royalblue
            push!(W, w)
            push!(L, w)
        end
    end
    K = parameterkeys(P)
    O = [Interact.onchange(w) for w in W]
    c = map(O...) do (W...)
        V = getindex.(W)
        configure(config, parameterzip(K, V))
    end
    if isempty(Interact.WebIO.providers_initialised)
        @warn "interactive plot only works with a WebIO provider loaded"
        return f(c[])
    end
    output = Interact.@map f(&c)
    z = Interact.Widget{:Cropbox}(Dict(zip(K, W)); output)
    Interact.@layout!(z, Interact.vbox(L..., output))
end

"""
    manipulate(args...; parameters, kwargs...)

Create an interactive plot by calling `manipulate` with `visualize` as a callback.

See also: [`visualize`](@ref)

# Arguments
- `args`: positional arguments for `visualize`.
- `parameters`: parameters for `manipulate`.
- `kwargs`: keyword arguments for `visualize`.
"""
manipulate(args...; parameters, config=(), kwargs...) = manipulate(function (c)
    visualize(args...; config=c, kwargs...)
end; parameters, config)

export manipulate


text:
# Converts the weigths of a PyTorch model to a Flux model from Metalhead

# PyTorch need to be installed

# Tested on ResNet and VGG models

using Flux
import Metalhead
using DataStructures
using Statistics
using BSON
using PyCall
using Images
using Test

torchvision = pyimport("torchvision")
torch = pyimport("torch")

modellib = [
    ("vgg11", () -> Metalhead.VGG(11), torchvision.models.vgg11),
    ("vgg13", () -> Metalhead.VGG(13), torchvision.models.vgg13),
    ("vgg16", () -> Metalhead.VGG(16), torchvision.models.vgg16),
    ("vgg19", () -> Metalhead.VGG(19), torchvision.models.vgg19),
    ("resnet18", () -> Metalhead.ResNet(18), torchvision.models.resnet18),
    ("resnet34", () -> Metalhead.ResNet(34), torchvision.models.resnet34),
    ("resnet50", () -> Metalhead.ResNet(50), torchvision.models.resnet50),
    ("resnet101",() -> Metalhead.ResNet(101),torchvision.models.resnet101),
    ("resnet152",() -> Metalhead.ResNet(152),torchvision.models.resnet152),
]

function _list_state(node::Flux.BatchNorm,channel,prefix)
    # use the same order of parameters than PyTorch
    put!(channel, (prefix * ".γ", node.γ)) # weigth (learnable)
    put!(channel, (prefix * ".β", node.β)) # bias (learnable)
    put!(channel, (prefix * ".μ", node.μ))  # running mean
    put!(channel, (prefix * ".σ²", node.σ²)) # running variance
end

function _list_state(node::Union{Flux.Conv,Flux.Dense},channel,prefix)
    put!(channel, (prefix * ".weight", node.weight))

    if node.bias !== Flux.Zeros()
        put!(channel, (prefix * ".bias", node.bias))
    end
end

_list_state(node,channel,prefix) = nothing

function _list_state(node::Union{Flux.Chain,Flux.Parallel},channel,prefix)
    for (i,n) in enumerate(node.layers)
        _list_state(n,channel,prefix * ".layers[$i]")
    end
end

function list_state(node; prefix = "model")
    Channel() do channel
        _list_state(node,channel,prefix)
    end
end

for (modelname,jlmodel,pymodel) in modellib

    model = jlmodel()
    pytorchmodel = pymodel(pretrained=true)

    state = OrderedDict(list_state(model.layers))

    # pytorchmodel.state_dict() looses the order
    state_dict = OrderedDict(pycall(pytorchmodel.state_dict,PyObject).items())
    pytorch_pp = OrderedDict((k,v.numpy()) for (k,v) in state_dict if !occursin("num_batches_tracked",k))

    # loop over all parameters
    for ((flux_key,flux_param),(pytorch_key,pytorch_param)) in zip(state,pytorch_pp)
        if size(flux_param) == size(pytorch_param)
            # Dense weight and vectors
            flux_param .= pytorch_param
        elseif size(flux_param) == reverse(size(pytorch_param))
            tmp = pytorch_param
            tmp = permutedims(tmp,ndims(tmp):-1:1)

            if ndims(flux_param)  == 4
                # convolutional weights
                flux_param .= reverse(tmp,dims=(1,2))
            else
                flux_param .= tmp
            end
        else
            @debug begin
                @show size(flux_param), size(pytorch_param)
            end
            error("incompatible shape $flux_key $pytorch_key")
        end
    end

    @info "saving model $modelname"
    BSON.@save joinpath(@__DIR__,"weights","$(modelname).bson") model
end

max_stars_repo_name: Alexander-Barth/MetalheadWeights

max_issues_repo_name: Alexander-Barth/MetalheadWeights

max_forks_repo_name: Alexander-Barth/MetalheadWeights


text:
# TODO -- temporary overwriting of this function, and should be removed from here when IIF 0.4.0 is tagged.
import IncrementalInference: compare

function compare(a::IncrementalInference.GenericFunctionNodeData{T1,S},b::IncrementalInference.GenericFunctionNodeData{T2,S}) where {T1, T2, S}
  # TODO -- beef up this comparison to include the gwp
  TP = true
  TP = TP && a.fncargvID == b.fncargvID
  TP = TP && a.eliminated == b.eliminated
  TP = TP && a.potentialused == b.potentialused
  TP = TP && a.edgeIDs == b.edgeIDs
  TP = TP && a.frommodule == b.frommodule
  # TP = TP && typeof(a.fnc) == typeof(b.fnc)
  return TP
end

"""
$(TYPEDEF)
"""
mutable struct Pose2Pose2{T} <: IncrementalInference.FunctorPairwise where {T <: IIF.SamplableBelief}
  z::T
  Pose2Pose2{T}() where {T <: IIF.SamplableBelief} = new{T}()
  Pose2Pose2{T}(z1::T) where {T <: IIF.SamplableBelief} = new{T}(z1)
end
Pose2Pose2(z::T) where {T <: IIF.SamplableBelief} = Pose2Pose2{T}(z)
function Pose2Pose2(mean::Array{Float64,1}, cov::Array{Float64,2})
  @warn "Pose2Pose2(mu,cov) is deprecated in favor of Pose2Pose2(T(...)) -- use for example Pose2Pose2(MvNormal(mu, cov))"
  Pose2Pose2(MvNormal(mean, cov))
end
function Pose2Pose2(mean::Array{Float64,1}, cov::Array{Float64,2}, w::Vector{Float64})
  @warn "Pose2Pose2(mu,cov,w) is deprecated in favor of Pose2Pose2(T(...)) -- use for example Pose2Pose2(MvNormal(mu, cov))"
  Pose2Pose2(MvNormal(mean, cov))
end
getSample(s::Pose2Pose2{<:IIF.SamplableBelief}, N::Int=1) = (rand(s.z,N), )
function (s::Pose2Pose2{<:IIF.SamplableBelief})(res::Vector{Float64},
            userdata,
            idx::Int,
            meas::Tuple,
            wxi::Array{Float64,2},
            wxj::Array{Float64,2}  )
  #
  wXjhat = SE2(wxi[1:3,idx])*SE2(meas[1][1:3,idx])
  jXjhat = SE2(wxj[1:3,idx]) \ wXjhat
  se2vee!(res, jXjhat)
  nothing
end

"""
$(TYPEDEF)

Introduce direct observations on all dimensions of a Pose2 variable:

Example:
--------
```julia
PriorPose2( MvNormal([10; 10; pi/6.0], Matrix(Diagonal([0.1;0.1;0.05].^2))) )
```
"""
mutable struct PriorPose2{T} <: IncrementalInference.FunctorSingleton  where {T <: IncrementalInference.SamplableBelief}
    Z::T
    PriorPose2{T}() where T = new{T}()
    PriorPose2{T}(x::T) where {T <: IncrementalInference.SamplableBelief}  = new{T}(x)
end
PriorPose2(x::T) where {T <: IncrementalInference.SamplableBelief} = PriorPose2{T}(x)
function PriorPose2(mu::Array{Float64}, cov::Array{Float64,2}, W::Vector{Float64})
  @warn "PriorPose2(mu,cov,W) is deprecated in favor of PriorPose2(T(...)) -- use for example PriorPose2(MvNormal(mu, cov))"
  PriorPose2(MvNormal(mu[:], cov))
end
function getSample(p2::PriorPose2, N::Int=1)
  return (rand(p2.Z,N), )
end

function compare(a::PriorPose2,b::PriorPose2; tol::Float64=1e-10)
  TP = true
  TP = TP && norm(a.Z.μ-b.Z.μ) < tol
  TP = TP && norm(a.Z.Σ.mat-b.Z.Σ.mat) < tol
  return TP
end
function compare(a::Pose2Pose2,b::Pose2Pose2; tol::Float64=1e-10)
  TP = true
  TP = TP && norm(a.z.μ-b.z.μ) < (tol + 1e-5)
  TP = TP && norm(a.z.Σ.mat-b.z.Σ.mat) < tol
  return TP
end

# Project all particles (columns) Xval with Z, that is for all  SE3(Xval[:,i])*Z
function projectParticles(Xval::Array{Float64,2}, Z::Array{Float64,2}, Cov::Array{Float64,2})
  # TODO optimize convert SE2 to a type
  @warn "projectParticles is an old function, rather standardize on approxConv instead."

  r,c = size(Xval)
  RES = zeros(r,c) #*cz

  # ent, x = SE3(0), SE3(0)
  j=1
  # for j in 1:cz
  ENT = rand( MvNormal(Z[:,1], Cov), c )
    for i in 1:c
      x = SE2(Xval[1:3,i])
      dx = SE2(ENT[1:3,i])
      RES[1:r,i*j] = se2vee(x*dx)
    end
  # end
  #
  return RES
end

⊕(Xpts::Array{Float64,2}, z::Pose2Pose2) = projectParticles(Xpts, z.Zij, z.Cov)
⊕(Xvert::Graphs.ExVertex, z::Pose2Pose2) = ⊕(getVal(Xvert), z)

"""
$(TYPEDEF)
"""
mutable struct PartialPriorYawPose2{T} <: IncrementalInference.FunctorSingleton  where {T <: IIF.SamplableBelief}
    Z::T
    partial::Tuple{Int}
    PartialPriorYawPose2{T}() where T = new{T}()
    PartialPriorYawPose2{T}(x::T) where {T <: IIF.SamplableBelief}  = new{T}(x, (3,))
end
PartialPriorYawPose2(x::T) where {T <: IIF.SamplableBelief} = PartialPriorYawPose2{T}(x)

function getSample(p2::PartialPriorYawPose2, N::Int=1)
  return (reshape(rand(p2.Z,N),1,N), )
end

"""
$(TYPEDEF)
"""
mutable struct PackedPartialPriorYawPose2 <: IncrementalInference.PackedInferenceType
    Z::String
    PackedPartialPriorYawPose2() = new()
    PackedPartialPriorYawPose2(x::T) where {T <: AbstractString}  = new(x)
end

function convert(::Type{PackedPartialPriorYawPose2}, d::PartialPriorYawPose2)
  PackedPartialPriorYawPose2(string(d.Z))
end
function convert(::Type{PartialPriorYawPose2}, d::PackedPartialPriorYawPose2)
  PartialPriorYawPose2(extractdistribution(d.Z))
end

# NOTE, for database support -- will be reduced to macro in future
# ------------------------------------

"""
$(TYPEDEF)
"""
mutable struct PackedPriorPose2  <: IncrementalInference.PackedInferenceType
    # vecZij::Array{Float64,1} # 0rotations, 1translation in each column
    # dimz::Int
    # vecCov::Array{Float64,1}
    # dimc::Int
    # W::Array{Float64,1}
    str::String
    PackedPriorPose2() = new()
    PackedPriorPose2(x::String) = new(x)
end
function convert(::Type{PackedPriorPose2}, d::PriorPose2)
  # v1 = d.Zi[:];
  # v2 = d.Cov[:];
  return PackedPriorPose2(string(d.Z))
end
function convert(::Type{PriorPose2}, d::PackedPriorPose2)
  # Zi = reshapeVec2Mat(d.vecZij,d.dimz)
  # Cov = reshapeVec2Mat(d.vecCov, d.dimc)
  distr = extractdistribution(d.str)
  return PriorPose2{typeof(distr)}(distr)
end

# --------------------------------------------

"""
$(TYPEDEF)
"""
mutable struct PackedPose2Pose2  <: IncrementalInference.PackedInferenceType
  datastr::String
  PackedPose2Pose2() = new()
  PackedPose2Pose2(x::String) = new(x)
end
function convert(::Type{Pose2Pose2}, d::PackedPose2Pose2)
  return Pose2Pose2(extractdistribution(d.datastr))
end
function convert(::Type{PackedPose2Pose2}, d::Pose2Pose2)
  return PackedPose2Pose2(string(d.z))
end

# --------------------------------------------


text:
function pcnet(X::AbstractMatrix{T}, p::Int=3;
              scalein::Bool=true, scaleout::Bool=false,
              symmout::Bool=false) where T<:Real
    if scalein
        σ=std(X,dims=1)
        σ[σ.==0].=1.0
        X=(X.-mean(X,dims=1))./σ
    end
    ℊ=size(X,2)
    A=1.0 .-Matrix(I,ℊ,ℊ)
    Threads.@threads for k in 1:ℊ
        y=X[:,k]
        𝒳=X[:,1:end.≠k]
        ϕ=TSVD.tsvd(𝒳,p)[3]
        s=𝒳*ϕ
        s ./= (vecnorm(s).^2)'
        b=sum(y.*s,dims=1)
        𝒷=ϕ*b'
        @inbounds A[k,A[k,:].==1.0]=𝒷
    end
    if symmout
        A=0.5*(A+A')
    end
    if scaleout
        A=A./maximum(abs.(A))
    end
    return convert(Matrix{Float16},A)
end

max_stars_repo_name: jamesjcai/ScTenifold

max_issues_repo_name: jamesjcai/ScTenifold

max_forks_repo_name: jamesjcai/ScTenifold


text:
export HWBS, distr, rand_rw, zb

using Distributions
using DataFrames

"""
The 1+1 Hull-White Black-Scholes model for both
real world and risk-neutral projections. The time line is the
interval [0,T], discretized into nᵀ steps of the same length.
"""
mutable struct HWBS
  "Number of steps per year"
  nʸ::Int
  "HW: initial interest rate"
  r₀::Real
  "HW: real world expected interest rate"
  r∞::Real
  "HW:convergence parameter"
  a::Real
  "HW: volatility parameter"
  σ::Real
  "BS: initial index price"
  S₀::Real
  "BS: expected real world yield from index"
  μ∞::Real
  "BS: volatility parameter"
  η::Real
  "BS/HW: correlation"
  ρ::Real
  "HW: spot rates etc. (risk neutral data from prices)"
  prices::DataFrame
end

"""
The input `df_spot` has at least two columns:
   - `:t`: maturities for spot rates.
   - `:spot`: spot rates at time zero for corresponding  maturities in `df_spot[:t]`
"""
function HWBS(p::Dict{Symbol, Real}, df_spot::DataFrame)
  prices = DataFrame()
  prices[!, :t] = df_spot[!, :t]
  prices[!, :spot] = df_spot[!, :spot]
  prices[!, :zb] = exp.(-prices[!, :t] .* prices[!, :spot])
  # Instantaneous forward rates
  # Approximate last forward with previous forward
  prices[!, :forw] =
    vcat( prices[1:(end-1), :t] .*
            diff(prices[!, :spot]) ./ diff(prices[!, :t]),
          prices[end, :t] *
            diff(prices[!, :spot])[end] /
            diff(prices[!, :t])[end])
  ## the differences are not long enough and the last entry
  ## is not correct
  ## Approximate correction: repeat the last (correct) value twice
  tmp =  diff(prices[!, :forw]) ./ diff(prices[!, :t])
  prices[!, :dforw] =
    vcat(tmp[1:end-1], tmp[end-1], tmp[end-1])

  HWBS( p[:nʸ], df_spot[1,:spot], p[:r∞], p[:a], p[:σ],
        p[:S₀], p[:μ∞], p[:η], p[:ρ], prices)
end

function ev_rw(cm::HWBS, t::Real)
  [(cm.r₀-cm.r∞) * exp(-cm.a * t) + cm.r∞,
  log(cm.S₀) + (cm.μ∞ - cm.η * cm.η /2)* t]
end

function cov_rw(cm::HWBS, t::Real)
  cov11 = cm.σ * cm.σ / (2cm.a) * (1 - exp(-2cm.a * t))
  cov12 = cm.ρ * cm.σ * cm.η / cm.a * (1 - exp(-cm.a * t))
  cov22 = cm.η * cm.η * t
  [cov11 cov12; cov12 cov22]
end

"""
Real world distribution of the HWBS model
* first component: ``r_t``
* second component: ``\\ln S_t``  **(ln may not be expected)**
"""
function distr_rw(cm::HWBS, t::Real)
  MvNormal(ev_rw(cm, t), cov_rw(cm, t))
end

"""
Real world random values for the HWBS model
* first component: ``r_t``
* second component: ``S_t``
"""
function rand_rw(cm::HWBS, t::Real, n)
  tmp = rand(distr_rw(cm, t), n)
  tmp[2,:] = exp.(tmp[2,:])
  tmp
end

"""
Price of zero bond which matures at time t₂, based on information at time 0.
"""
function zb(cm::HWBS, t₂::Real)
    n₂ = convert(Int, round(t₂ * cm.nʸ, digits = 0))
  if n₂ > nrow(cm.prices)
    error("t₂ too large")
  else
    return cm.prices[n₂, :zb]
  end
end

"""
Price of zero bond which matures at time ``t₂``, as estimated  at time ``t₁``, under the condition that ``r_{t_1} = r_1``.
"""
function zb(cm::HWBS, r₁::Real, t₁::Real, t₂::Real)
  Δt = t₂-t₁
  n₁ = convert(Int, round(t₁ * cm.nʸ, digits = 0))
  n₂ = convert(Int, round(t₂ * cm.nʸ, digits = 0))
  if n₂ > nrow(cm.prices)
    error("t₂ too large")
  elseif Δt < 0
    error("t₂ < t₁")
  elseif n₁ == 0
    return zb(cm, t₂)
  else
    B = (1 - exp(-cm.a * Δt)) / cm.a
    A = log(cm.prices[n₂, :zb]) - log(cm.prices[n₁, :zb]) +
        cm.prices[n₁, :forw] * B -
        cm.σ^2 / (4 * cm.a) * B^2 * (1 - exp(-2 * cm.a * Δt))
    return exp(A) * exp(-B * r₁)
  end
end


text:
struct InterfaceOperators
    operators::Any
    celltooperator::Any
    ncells::Any
    numoperators::Any
    function InterfaceOperators(operators, celltooperator)
        ncells = length(celltooperator)
        numcouples, numoperators = size(operators)

        @assert numcouples == 4
        @assert all(celltooperator .>= 0)
        @assert all(celltooperator .<= numoperators)
        new(operators, celltooperator, ncells, numoperators)
    end
end

function Base.getindex(interfaceoperator::InterfaceOperators, s1, s2, cellid)
    row = cell_couple_sign_to_row(s1, s2)
    idx = interfaceoperator.celltooperator[cellid]
    idx > 0 || error("Cell $cellid does not have an operator")
    return interfaceoperator.operators[row, idx]
end

function Base.show(io::IO, interfaceoperator::InterfaceOperators)
    ncells = number_of_cells(interfaceoperator)
    nops = number_of_operators(interfaceoperator)
    str = "InterfaceOperator\n\tNum. Cells: $ncells\n\tNum. Operators: $nops"
    print(io, str)
end

function number_of_cells(interfaceoperator::InterfaceOperators)
    return interfaceoperator.ncells
end

function number_of_operators(interfaceoperator::InterfaceOperators)
    return interfaceoperator.numoperators
end

function interface_mass_operators(basis, interfacequads, cellmap, cellsign, penalty)

    ncells = length(cellsign)
    hasinterface = cellsign .== 0
    numinterfaces = count(hasinterface)
    operators = Matrix{Any}(undef, 4, numinterfaces)
    celltooperator = zeros(Int, ncells)
    dim = dimension(basis)

    cellids = findall(hasinterface)

    for (idx, cellid) in enumerate(cellids)
        normal = interface_normals(interfacequads, cellid)
        facescale = scale_area(cellmap, normal)
        for s1 in [+1, -1]
            quad1 = interfacequads[s1, cellid]
            for s2 in [+1, -1]
                row = cell_couple_sign_to_row(s1, s2)
                quad2 = interfacequads[s2, cellid]
                mass = penalty * interface_mass_matrix(basis, quad1, quad2, facescale)
                operators[row, idx] = mass
            end
        end
        celltooperator[cellid] = idx
    end
    return InterfaceOperators(operators, celltooperator)
end

function interface_mass_operators(basis, interfacequads, cutmesh, penalty)
    cellmap = cell_map(cutmesh, 1)
    cellsign = cell_sign(cutmesh)
    return interface_mass_operators(basis, interfacequads, cellmap, cellsign, penalty)
end

function interface_incoherent_mass_operators(
    basis,
    interfacequads,
    cellmap,
    cellsign,
    penalty,
)

    ncells = length(cellsign)
    hasinterface = cellsign .== 0
    numinterfaces = count(hasinterface)
    operators = Matrix{Any}(undef, 4, numinterfaces)
    celltooperator = zeros(Int, ncells)
    dim = dimension(basis)

    cellids = findall(hasinterface)

    for (idx, cellid) in enumerate(cellids)
        normal = interface_normals(interfacequads, cellid)
        components = normal
        facescale = scale_area(cellmap, normal)
        for s1 in [+1, -1]
            quad1 = interfacequads[s1, cellid]
            for s2 in [+1, -1]
                row = cell_couple_sign_to_row(s1, s2)
                quad2 = interfacequads[s2, cellid]
                mass =
                    penalty * interface_component_mass_matrix(
                        basis,
                        quad1,
                        quad2,
                        components,
                        facescale,
                    )
                operators[row, idx] = mass
            end
        end
        celltooperator[cellid] = idx
    end
    return InterfaceOperators(operators, celltooperator)
end

function interface_incoherent_mass_operators(basis, interfacequads, cutmesh, penalty)
    cellmap = cell_map(cutmesh, 1)
    cellsign = cell_sign(cutmesh)
    return interface_incoherent_mass_operators(
        basis,
        interfacequads,
        cellmap,
        cellsign,
        penalty,
    )
end

function interface_traction_operators(basis, interfacequads, stiffness, cellmap, cellsign)

    ncells = length(cellsign)
    hasinterface = cellsign .== 0
    numinterfaces = count(hasinterface)
    operators = Matrix{Any}(undef, 4, numinterfaces)
    celltooperator = zeros(Int, ncells)

    cellids = findall(hasinterface)

    for (idx, cellid) in enumerate(cellids)
        normal = interface_normals(interfacequads, cellid)
        for s1 in [+1, -1]
            quad1 = interfacequads[s1, cellid]
            for s2 in [+1, -1]
                row = cell_couple_sign_to_row(s1, s2)
                quad2 = interfacequads[s2, cellid]
                top = coherent_traction_operator(
                    basis,
                    quad1,
                    quad2,
                    normal,
                    stiffness[s2],
                    cellmap,
                )
                operators[row, idx] = top
            end
        end
        celltooperator[cellid] = idx
    end
    return InterfaceOperators(operators, celltooperator)
end

function interface_traction_operators(basis, interfacequads, stiffness, cutmesh)
    cellmap = cell_map(cutmesh, 1)
    cellsign = cell_sign(cutmesh)
    return interface_traction_operators(basis, interfacequads, stiffness, cellmap, cellsign)
end

function interface_incoherent_traction_operators(basis, interfacequads, stiffness, cutmesh)
    cellsign = cell_sign(cutmesh)
    cellmap = cell_map(cutmesh, 1)

    ncells = length(cellsign)
    hasinterface = cellsign .== 0
    numinterfaces = count(hasinterface)
    operators = Matrix{Any}(undef, 4, numinterfaces)
    celltooperator = zeros(Int, ncells)

    cellids = findall(hasinterface)

    for (idx, cellid) in enumerate(cellids)
        normals = interface_normals(interfacequads, cellid)
        for s1 in [+1, -1]
            quad1 = interfacequads[s1, cellid]
            for s2 in [+1, -1]
                row = cell_couple_sign_to_row(s1, s2)
                quad2 = interfacequads[s2, cellid]
                top = incoherent_traction_operator(
                    basis,
                    quad1,
                    quad2,
                    normals,
                    stiffness[s2],
                    cellmap,
                )
                operators[row, idx] = top
            end
        end
        celltooperator[cellid] = idx
    end
    return InterfaceOperators(operators, celltooperator)
end

struct InterfaceCondition
    tractionoperator::Any
    massoperator::Any
    penalty::Any
    ncells::Any
    function InterfaceCondition(
        tractionoperator::InterfaceOperators,
        massoperator::InterfaceOperators,
        penalty,
    )

        ncells = number_of_cells(massoperator)
        @assert number_of_cells(tractionoperator) == ncells
        new(tractionoperator, massoperator, penalty, ncells)
    end
end

function coherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)
    tractionoperator =
        interface_traction_operators(basis, interfacequads, stiffness, cutmesh)
    massoperator = interface_mass_operators(basis, interfacequads, cutmesh, penalty)
    return InterfaceCondition(tractionoperator, massoperator, penalty)
end

function incoherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)
    tractionoperator =
        interface_incoherent_traction_operators(basis, interfacequads, stiffness, cutmesh)
    massoperator =
        interface_incoherent_mass_operators(basis, interfacequads, cutmesh, penalty)
    return InterfaceCondition(tractionoperator, massoperator, penalty)
end

function traction_operator(interfacecondition::InterfaceCondition, s1, s2, cellid)
    return interfacecondition.tractionoperator[s1, s2, cellid]
end

function mass_operator(interfacecondition::InterfaceCondition, s1, s2, cellid)
    return interfacecondition.massoperator[s1, s2, cellid]
end

function Base.show(io::IO, interfacecondition::InterfaceCondition)
    ncells = interfacecondition.ncells
    penalty = interfacecondition.penalty
    str = "InterfaceCondition\n\tNum. Cells: $ncells\n\tDisplacement Penalty: $penalty"
    print(io, str)
end

function coherent_traction_operator(basis, quad1, quad2, normals, stiffness, cellmap)
    numqp = length(quad1)
    @assert length(quad2) == size(normals)[2] == numqp
    dim = dimension(basis)
    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    matrix = zeros(ndofs, ndofs)
    vectosymmconverter = vector_to_symmetric_matrix_converter()

    jac = jacobian(cellmap)
    scalearea = scale_area(cellmap, normals)

    for qpidx = 1:numqp
        p1, w1 = quad1[qpidx]
        p2, w2 = quad2[qpidx]
        @assert w1 ≈ w2

        vals = basis(p1)
        grad = transform_gradient(gradient(basis, p2), jac)
        normal = normals[:, qpidx]
        NK = sum([make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim])
        N = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])
        NI = interpolation_matrix(vals, dim)

        matrix .+= NI' * N * stiffness * NK * scalearea[qpidx] * w1
    end
    return matrix
end

function component_traction_operator(
    basis,
    quad1,
    quad2,
    components,
    normals,
    stiffness,
    cellmap,
)
    numqp = length(quad1)
    @assert length(quad2) == size(normals)[2] == size(components)[2] == numqp
    dim = dimension(basis)
    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    matrix = zeros(ndofs, ndofs)
    vectosymmconverter = vector_to_symmetric_matrix_converter()

    jac = jacobian(cellmap)
    scalearea = scale_area(cellmap, normals)

    for qpidx = 1:numqp
        p1, w1 = quad1[qpidx]
        p2, w2 = quad2[qpidx]
        @assert w1 ≈ w2

        vals = basis(p1)
        grad = transform_gradient(gradient(basis, p2), jac)

        normal = normals[:, qpidx]
        component = components[:, qpidx]
        projector = component * component'

        NK = sum([make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim])
        N = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])
        NI = interpolation_matrix(vals, dim)

        matrix .+= NI' * projector * N * stiffness * NK * scalearea[qpidx] * w1
    end
    return matrix
end

function incoherent_traction_operator(basis, quad1, quad2, normals, stiffness, cellmap)

    components = normals
    return component_traction_operator(
        basis,
        quad1,
        quad2,
        components,
        normals,
        stiffness,
        cellmap,
    )
end

function interface_mass_matrix(basis, quad1, quad2, scale)
    numqp = length(quad1)
    @assert length(quad2) == length(scale) == numqp
    nf = number_of_basis_functions(basis)
    dim = dimension(basis)
    totaldofs = dim * nf
    matrix = zeros(totaldofs, totaldofs)
    for qpidx = 1:numqp
        p1, w1 = quad1[qpidx]
        p2, w2 = quad2[qpidx]
        @assert w1 ≈ w2

        vals1 = basis(p1)
        vals2 = basis(p2)

        NI1 = interpolation_matrix(vals1, dim)
        NI2 = interpolation_matrix(vals2, dim)

        matrix .+= NI1' * NI2 * scale[qpidx] * w1
    end
    return matrix
end

function interface_component_mass_matrix(basis, quad1, quad2, components, scale)
    numqp = length(quad1)
    @assert length(quad2) == length(scale) == size(components)[2] == numqp
    nf = number_of_basis_functions(basis)
    dim = dimension(basis)
    totaldofs = dim * nf
    matrix = zeros(totaldofs, totaldofs)
    for qpidx = 1:numqp
        p1, w1 = quad1[qpidx]
        p2, w2 = quad2[qpidx]
        @assert w1 ≈ w2

        component = components[:, qpidx]
        projector = component * component'

        vals1 = basis(p1)
        vals2 = basis(p2)

        NI1 = interpolation_matrix(vals1, dim)
        NI2 = make_row_matrix(projector, vals2)

        matrix .+= NI1' * NI2 * scale[qpidx] * w1
    end
    return matrix
end


text:
### A Hamiltonian is composed of blocks (kpoints), which have a list of RealFourierOperator
# corresponding to each term
# This is the "high-level" interface, provided for convenience

struct HamiltonianBlock
    basis::PlaneWaveBasis
    kpoint::Kpoint

    # The operators are vectors of RealFourierOperator,
    # not typed because of type invariance issues.
    operators::Vector  # the original list of RealFourierOperator
                       # (as many as there are terms), kept for easier exploration
    optimized_operators::Vector  # the optimized list of RealFourierOperator, to be used for applying
    scratch  # Pre-allocated scratch arrays for fast application
end
function HamiltonianBlock(basis, kpt, operators, scratch)
    HamiltonianBlock(basis, kpt, operators, optimize_operators_(operators), scratch)
end
Base.eltype(block::HamiltonianBlock) = complex(eltype(block.basis))
Base.size(block::HamiltonianBlock, i::Integer) = i < 3 ? size(block)[i] : 1
function Base.size(block::HamiltonianBlock)
    n_G = length(G_vectors(block.basis, block.kpoint))
    (n_G, n_G)
end

struct Hamiltonian
    basis::PlaneWaveBasis
    blocks::Vector{HamiltonianBlock}
end

# Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
# This is a fallback function that works for all hamiltonians;
# for "usual cases", there is a faster implementation below
@views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray,
                                                                        H::HamiltonianBlock,
                                                                        ψ::AbstractArray)
    # Special-case of DFT Hamiltonian: go to a fast path
    fh = fast_hblock(H)
    fh !== nothing && return fast_hblock_mul!(Hψ, fh, ψ)

    basis = H.basis
    T = eltype(basis)
    kpt = H.kpoint
    nband = size(ψ, 2)

    Hψ_fourier = similar(Hψ[:, 1])
    ψ_real = zeros(complex(T), basis.fft_size...)
    Hψ_real = zeros(complex(T), basis.fft_size...)

    # take ψi, IFFT it to ψ_real, apply each term to Hψ_fourier and Hψ_real, and add it to Hψ
    for iband = 1:nband
        Hψ_real .= 0
        Hψ_fourier .= 0
        G_to_r!(ψ_real, basis, kpt, ψ[:, iband])
        for op in H.optimized_operators
            apply!((fourier=Hψ_fourier, real=Hψ_real),
                   op,
                   (fourier=ψ[:, iband], real=ψ_real))
        end
        Hψ[:, iband] .= Hψ_fourier
        r_to_G!(Hψ_fourier, basis, kpt, Hψ_real)
        Hψ[:, iband] .+= Hψ_fourier
    end

    Hψ
end
Base.:*(H::HamiltonianBlock, ψ) = mul!(similar(ψ), H, ψ)

# Returns a fast path hamiltonian if eligible, nothing if not
function fast_hblock(H::HamiltonianBlock)
    length(H.optimized_operators) == 3 || return nothing
    fourier_ops = filter(o -> o isa FourierMultiplication, H.optimized_operators)
    real_ops = filter(o -> o isa RealSpaceMultiplication, H.optimized_operators)
    nonlocal_ops = filter(o -> o isa NonlocalOperator, H.optimized_operators)
    (length(fourier_ops) == length(real_ops) == length(nonlocal_ops) == 1) || return nothing
    (fourier_op=only(fourier_ops), real_op=only(real_ops), nonlocal_op=only(nonlocal_ops), H=H)
end

# Fast version, specialized on DFT models with just 3 operators: real, fourier and nonlocal
# Minimizes the number of allocations
@views function fast_hblock_mul!(Hψ::AbstractArray,
                                 fast_hblock::NamedTuple,
                                 ψ::AbstractArray)
    H = fast_hblock.H
    basis = H.basis
    kpt = H.kpoint
    nband = size(ψ, 2)

    potential = fast_hblock.real_op.potential
    potential /= prod(basis.fft_size)  # because we use unnormalized plans
    @timing "kinetic+local" begin
        Threads.@threads for iband = 1:nband
            tid = Threads.threadid()
            ψ_real = H.scratch.ψ_reals[tid]

            G_to_r!(ψ_real, basis, kpt, ψ[:, iband]; normalize=false)
            ψ_real .*= potential
            r_to_G!(Hψ[:, iband], basis, kpt, ψ_real; normalize=false)  # overwrites ψ_real
            Hψ[:, iband] .+= fast_hblock.fourier_op.multiplier .* ψ[:, iband]
        end
    end

    # Apply the nonlocal operator
    @timing "nonlocal" begin
        apply!((fourier=Hψ, real=nothing), fast_hblock.nonlocal_op, (fourier=ψ, real=nothing))
    end

    Hψ
end

function LinearAlgebra.mul!(Hψ, H::Hamiltonian, ψ)
    for ik = 1:length(H.basis.kpoints)
        mul!(Hψ[ik], H.blocks[ik], ψ[ik])
    end
    Hψ
end
# need `deepcopy` here to copy the elements of the array of arrays ψ (not just pointers)
Base.:*(H::Hamiltonian, ψ) = mul!(deepcopy(ψ), H, ψ)

# Get energies and Hamiltonian
# kwargs is additional info that might be useful for the energy terms to precompute
# (eg the density ρ)
@timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occ; kwargs...)
    # it: index into terms, ik: index into kpoints
    @timing "ene_ops" ene_ops_arr = [ene_ops(term, basis, ψ, occ; kwargs...)
                                     for term in basis.terms]
    energies  = [eh.E for eh in ene_ops_arr]
    operators = [eh.ops for eh in ene_ops_arr]         # operators[it][ik]

    # flatten the inner arrays in case a term returns more than one operator
    function flatten(arr)
        ret = []
        for a in arr
            if a isa RealFourierOperator
                push!(ret, a)
            else
                push!(ret, a...)
            end
        end
        ret
    end
    hks_per_k   = [flatten([blocks[ik] for blocks in operators])
                   for ik = 1:length(basis.kpoints)]      # hks_per_k[ik][it]

    # Preallocated scratch arrays
    T = eltype(basis)
    scratch = (
        ψ_reals=[zeros(complex(T), basis.fft_size...) for tid = 1:Threads.nthreads()],
    )

    H = Hamiltonian(basis, [HamiltonianBlock(basis, kpt, hks, scratch)
                            for (hks, kpt) in zip(hks_per_k, basis.kpoints)])
    E = Energies(basis.model.term_types, energies)
    (E=E, H=H)
end
function Hamiltonian(basis::PlaneWaveBasis; ψ=nothing, occ=nothing, kwargs...)
    _, H = energy_hamiltonian(basis, ψ, occ; kwargs...)
    H
end

import Base: Matrix, Array
function Matrix(block::HamiltonianBlock)
    sum(Matrix(h) for h in block.optimized_operators)
end
Array(block::HamiltonianBlock) = Matrix(block)

"""
Get the total local potential of the given Hamiltonian, in real space
in the spin components.
"""
function total_local_potential(ham::Hamiltonian)
    n_spin = ham.basis.model.n_spin_components
    pots = map(1:n_spin) do σ
        # Get the first Hamiltonian block of this spin component
        # (works since all local potentials are the same)
        block = ham.blocks[first(krange_spin(ham.basis, σ))]
        rs = [o for o in block.optimized_operators if o isa RealSpaceMultiplication]
        only(rs).potential
    end
    cat(pots..., dims=4)
end

"""
Returns a new Hamiltonian with local potential replaced by the given one
"""
function hamiltonian_with_total_potential(ham::Hamiltonian, V)
    @assert size(V, 4) == ham.basis.model.n_spin_components
    newblocks = [hamiltonian_with_total_potential(Hk, V[:, :, :, Hk.kpoint.spin])
                 for Hk in ham.blocks]
    Hamiltonian(ham.basis, newblocks)
end
function hamiltonian_with_total_potential(Hk::HamiltonianBlock, V)
    operators = [op for op in Hk.operators if !(op isa RealSpaceMultiplication)]
    push!(operators, RealSpaceMultiplication(Hk.basis, Hk.kpoint, V))
    HamiltonianBlock(Hk.basis, Hk.kpoint, operators, Hk.scratch)
end


text:
#
#      MatrixManifold -- a matrix manifold,
#       i.e. the values/points on the manifold are matrices
#		this is here implemented as traits, such that every manifold can
#		“gain” this feature by setting @traitimlp IsMatrixM{}
#		for the manifold, its point and the tangent vector
#
# Manopt.jl, R. Bergmann, 2018-06-26
export IsEmbeddedM, IsEmbeddedP, IsEmbeddedV
"""
    IsEmbeddedM{X}
An abstract [`Manifold`](@ref) that is embedded in some Euclidean space.
These manifolds may have projections and converters for gradient and Hessian.
"""
@traitdef IsEmbeddedM{X}
"""
    IsEmbeddedP{X}
An abstract [`MPoint`](@ref) belonging to an embedded manifold.
"""
@traitdef IsEmbeddedP{X}
"""
    IsEmbeddedV{X}
An abstract [`TVector`](@ref) belonging to an embedded manifold.
"""
@traitdef IsEmbeddedV{X}


text:
#=
SimulationTest:
- Julia version:
- Author: inaba
- Date: 2021-11-20
=#
module SimulationTest
using Test
using LightGraphs
using Statistics: mean

include("../src/Simulation.jl")
const sim = Simulation  # alias

@testset "Agent" begin
    @testset "Deterministic test" begin
        agent = sim.Agent(1)
        @test agent.id == 1
        @test agent.next_is_cooperator == false
        @test agent.next_is_punisher == false
        @test agent.payoff == 0.0
    end

    # is_cooperator と is_punisher が50%程度の確率で設定されていることを確認。
    @testset "Probabilistic test" begin
        agents = [sim.Agent(id) for id in 1:1000]
        average_cooperator_rate = mean([agent.is_cooperator for agent in agents])
        average_punisher_rate = mean([agent.is_punisher for agent in agents])
        @test 0.45 <= average_cooperator_rate <= 0.55
        @test 0.45 <= average_punisher_rate <= 0.55
    end
end

@testset "Model" begin
    @testset "neighbours" begin
        agents = [sim.Agent(id) for id in 1:5]
        graph = cycle_graph(length(agents))
        model = sim.Model(agents, graph, g = 0.0, p = 1.0, a = 0.0)
        # 1の隣人は 4, 5, 2, 3 であることを確認
        @test sort(model.neighbours[1]) == sort([4, 5, 2, 3])
        # 2の隣人は 5, 1, 3, 4 であることを確認
        @test sort(model.neighbours[2]) == sort([5, 1, 3, 4])
        # 3の隣人は 1, 2, 4, 5 であることを確認
        @test sort(model.neighbours[3]) == sort([1, 2, 4, 5])
        # 4の隣人は 2, 3, 5, 1 であることを確認
        @test sort(model.neighbours[4]) == sort([2, 3, 5, 1])
        # 5の隣人は 3, 4, 1, 2 であることを確認
        @test sort(model.neighbours[5]) == sort([3, 4, 1, 2])
    end

    @testset "calc_payoffs!" begin
        agents = [sim.Agent(id) for id in 1:10]

        # agentの戦略を固定する。(戦略を確率論的にではなく決定論的に決めるのが目的なので、どの戦略にするかは恣意的に決めて良い)
        for agent in agents
            agent.is_cooperator = (agent.id % 3 == 2)
            agent.is_punisher = (agent.id % 3 == 1)
        end

        graph = cycle_graph(length(agents))
        model = sim.Model(agents, graph, g = 0.0, p = 1.0, a = 0.0)
        sim.calc_payoffs!(model)
    end
end
end


text:
using Documenter
using Pkg
using NahaJuliaLib

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

makedocs(;
    modules=[NahaJuliaLib],
    authors="MarkNahabedian <naha@mit.edu> and contributors",
    repo="https://github.com/MarkNahabedian/NahaJuliaLib.jl/blob/{commit}{path}#{line}",
    sitename="NahaJuliaLib.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://MarkNahabedian.github.io/NahaJuliaLib.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/MarkNahabedian/NahaJuliaLib.jl",
    devbranch="master",
)


text:
"""
    check_map(ref, cur)
---
Check the SNP in map `cur` have the same chromosome and base pair position.
"""
function check_map(ref, cur)
    title("Check maps if same")
    if !isfile(ref) || !isfile(cur)
        warning("File $ref or $cur doesn't exist")
        return
    end
    dic = Dict()
    for line in eachline(ref)
        snp, chr, bp = split(line)
        push!(dic, snp=>(chr, bp))
    end
    for line in eachline(cur)
        snp, chr, bp = split(line)
        if haskey(dic, snp)
            if dic[snp] != (chr, bp)
                warning("The two maps are different")
                return false
            end
        end
    end

    return true
end

"""
    update_map(cur, ref)
---
Update SNP map locate in `cur` with info from `ref`.  If a SNP is only in `cur`, keep its info.  Results are written to `out`.
"""
function update_map(cur, ref, out)
    title("Update map $cur with $ref")
    if !isfile(ref) || !isfile(cur)
        warning("File $ref or $cur doesn't exist")
        return
    end

    dic = Dict()
    for line in eachline(ref)
        snp, chr, bp = split(line)
        push!(dic, snp=>(chr, bp))
    end

    tmp = mktempdir(".")
    open("$tmp/t.map", "w") do io
        for line in eachline(cur)
            snp, chr, bp = split(line)
            if haskey(dic, snp)
                chr, bp = dic[snp]
            end
            println(io, join([snp, chr, bp], '\t'))
        end
    end
    run(pipeline("$tmp/t.map", `sort -nk2 -nk3`, out))
    done()
end


text:
using C2Julia
using CSyntax
using Clang
using Test

@testset "switch" begin
    tu = parse_header(joinpath(@__DIR__, "c", "switch.h"), includes=[LLVM_INCLUDE])
    GC.@preserve tu begin
        root_cursor = getcursor(tu)
        switch_func = search(root_cursor, "switch_func")[1]
        body = children(switch_func)[1]
        switch = children(body)[3]
        x = translate(switch)
        expr = Expr(:macrocall, Symbol("@switch"), nothing, :str, Expr(:block,
                    Expr(:macrocall, Symbol("@case"), nothing, 'a'),
                    :(x = 1),
                    :(x = 0),
                    Expr(:break),
                    Expr(:macrocall, Symbol("@case"), nothing, 'b'),
                    Expr(:if, :(x == 2), :(x = 7)),
                    Expr(:macrocall, Symbol("@case"), nothing, 'c'),
                    Expr(:macrocall, Symbol("@default"), nothing),
                    :(x = 10)
                    ))
        y = Base.remove_linenums!(expr)
        @test string(x) == string(y)
   end
end


text:
module Utils
    function dtot32(d)
        return map(x->UInt32(mod(round(Int64,mod(x,1.0)*2.0^32),Int64(2)^32)),d)
    end
    function gaussian32(μ,α)
        return dtot32(α*randn(length(μ)) .+ μ)
    end
end


text:
using XDiff
import ReverseDiffSource

function ann(w1, w2, w3, x1)
    x2 = w1 * x1
    x3 = w2 * x2
    x4 = sum(w3 * x3)
    return x4
end

# ANN input parameter types
types = (Matrix{Float64}, Matrix{Float64}, Matrix{Float64}, Vector{Float64})

# generate example input
w1, w2, w3 = randn(10,10), randn(10,10), randn(1,10)
x1 = randn(10)

dann = ReverseDiffSource.rdiff(ann, types)
_, rds_dw1_val, _, _, _ = dann(w1, w2, w3, x1)

dw1, _, _, _ = fdiff(ann, types)
dw1_val = dw1(w1, w2, w3, x1)

if isapprox(dw1_val, rds_dw1_val)
    println("Equal!")
else
    error("Not equal!")
end


text:
export SIM, SimPosterior, Distances

struct SIM{T<:Metric}
    mode::Vector{Path{Int}} # Mode
    γ::Real # Precision
    dist::T # Distance metric
    V::UnitRange{Int} # Vertex Set
    K_inner::DimensionRange # Maximum interaction sequence size
    K_outer::DimensionRange # Maximum path (interaction) length
end

SIM(
    mode::InteractionSequence{Int},
    γ::Real,
    dist::Metric,
    V::UnitRange
) = SIM(
    mode,
    γ,
    dist, V,
    DimensionRange(1,Inf),
    DimensionRange(1,Inf)
)

SIM(
    mode::InteractionSequence{Int},
    γ::Real,
    dist::Metric,
    V::UnitRange,
    K_inner::Real, K_outer::Real
) = SIM(
    mode,
    γ,
    dist, V,
    DimensionRange(1,K_inner),
    DimensionRange(1,K_outer)
)

function Base.show(
    io::IO, model::SIM
    )

    title = "$(typeof(model))"
    n = length(title)
    println(io, title)
    println(io, "-"^n)
    println(io, "Parameters:")
    for par in fieldnames(typeof(model))
        println(io, par, " = $(getfield(model, par))")
    end
end

function Base.similar(
    model::SIM{T},
    mode::Vector{Vector{Int}},
    γ::Float64
    ) where {T}
    return SIM(mode, γ, model.dist, model.V, model.K_inner, model.K_outer)
end

struct SimPosterior
    data::InteractionSequenceSample{Int}
    S_prior::SIM
    γ_prior::ContinuousUnivariateDistribution
    dist::Metric
    V::AbstractArray{Int}
    K_inner::DimensionRange
    K_outer::DimensionRange
    sample_size::Int
    function SimPosterior(
        data::InteractionSequenceSample{Int},
        S_prior::SIM,
        γ_prior::ContinuousUnivariateDistribution
        )

        dist = S_prior.dist
        V = S_prior.V
        K_inner = S_prior.K_inner
        K_outer = S_prior.K_outer
        sample_size = length(data)
        new(data, S_prior, γ_prior, dist, V, K_inner, K_outer, sample_size)
    end
end

max_stars_repo_path: src/Metric Models/SIM.jl

max_issues_repo_path: src/Metric Models/SIM.jl

max_forks_repo_path: src/Metric Models/SIM.jl


text:
using MPI

using CLIMA.Topologies
using CLIMA.Grids
using CLIMA.AtmosDycore.VanillaAtmosDiscretizations
using CLIMA.MPIStateArrays
using CLIMA.ODESolvers
using CLIMA.LowStorageRungeKuttaMethod
using CLIMA.GenericCallbacks
using CLIMA.AtmosDycore
using CLIMA.MoistThermodynamics
using LinearAlgebra
using Logging, Printf, Dates

using CLIMA.ParametersType
using CLIMA.PlanetParameters: R_d, cp_d, grav, cv_d, MSLP

# FIXME: Will these keywords args be OK?
function rising_thermal_bubble(x...; ntrace=0, nmoist=0, dim=3)
  DFloat = eltype(x)

  p0::DFloat      = MSLP
  R_gas::DFloat   = R_d
  c_p::DFloat     = cp_d
  c_v::DFloat     = cv_d
  gravity::DFloat = grav
  q_tot::DFloat   = 0

  r = sqrt((x[1] - 500)^2 + (x[dim] - 350)^2)
  rc::DFloat = 250
  θ_ref::DFloat = 300
  θ_c::DFloat = 0.5
  Δθ::DFloat = 0.0
  if r <= rc
    Δθ = θ_c * (1 + cos(π * r / rc)) / 2
  end
  θ = θ_ref + Δθ
  π_k = 1 - gravity / (c_p * θ) * x[dim]

  ρ = p0 / (R_gas * θ) * (π_k)^ (c_v / R_gas)
  u = zero(DFloat)
  v = zero(DFloat)
  w = zero(DFloat)
  U = ρ * u
  V = ρ * v
  W = ρ * w
  P = p0 * (R_gas * (ρ * θ) / p0)^(c_p / c_v)
  T = P / (ρ * R_gas)
  # Calculation of energy per unit mass
  e_kin = (u^2 + v^2 + w^2) / 2
  e_pot = gravity * x[dim]
  e_int = MoistThermodynamics.internal_energy(T)
  # Total energy
  E = ρ * MoistThermodynamics.total_energy(e_kin, e_pot, T)
  (ρ=ρ, U=U, V=V, W=W, E=E, Qmoist=(ρ * q_tot,))
end

function main(mpicomm, DFloat, ArrayType, brickrange, nmoist, ntrace, N,
              timeend, bricktopo; gravity=true, viscosity=0, dt=nothing,
              exact_timeend=true)
  dim = length(brickrange)
  topl = bricktopo(# MPI communicator to connect elements/partition
                   mpicomm,
                   # tuple of point element edges in each dimension
                   # (dim is inferred from this)
                   brickrange,
                   periodicity=(true, ntuple(j->false, dim-1)...))

  grid = DiscontinuousSpectralElementGrid(topl,
                                          # Compute floating point type
                                          FloatType = DFloat,
                                          # This is the array type to store
                                          # data: CuArray = GPU, Array = CPU
                                          DeviceArray = ArrayType,
                                          # polynomial order for LGL grid
                                          polynomialorder = N,
                                          # how to skew the mesh degrees of
                                          # freedom (for instance spherical
                                          # or topography maps)
                                          # warp = warpgridfun
                                         )

  # spacedisc = data needed for evaluating the right-hand side function
  spacedisc = VanillaAtmosDiscretization(grid,
                                        gravity=gravity,
                                        viscosity=viscosity,
                                        ntrace=ntrace,
                                        nmoist=nmoist)

  # This is a actual state/function that lives on the grid
  initialcondition(x...) = rising_thermal_bubble(x...;
                                               ntrace=ntrace,
                                               nmoist=nmoist,
                                               dim=dim)
  Q = MPIStateArray(spacedisc, initialcondition)

  # Determine the time step
  (dt == nothing) && (dt = VanillaAtmosDiscretizations.estimatedt(spacedisc, Q))
  if exact_timeend
    nsteps = ceil(Int64, timeend / dt)
    dt = timeend / nsteps
  end

  # Initialize the Method (extra needed buffers created here)
  # Could also add an init here for instance if the ODE solver has some
  # state and reading from a restart file

  # TODO: Should we use get property to get the rhs function?
  lsrk = LowStorageRungeKutta(getrhsfunction(spacedisc), Q; dt = dt, t0 = 0)

  # Get the initial energy
  eng0 = norm(Q)
  @info @sprintf """Starting
  norm(Q₀) = %.16e""" eng0

  # Set up the information callback
  starttime = Ref(now())
  cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s=false)
    if s
      starttime[] = now()
    else
      energy = norm(Q)
      @info @sprintf """Update
  simtime = %.16e
  runtime = %s
  norm(Q) = %.16e""" ODESolvers.gettime(lsrk) Dates.format(convert(Dates.DateTime, Dates.now()-starttime[]), Dates.dateformat"HH:MM:SS") energy
    end
  end

  #= Paraview calculators:
  P = (0.4) * (E  - (U^2 + V^2 + W^2) / (2*ρ) - 9.81 * ρ * coordsZ)
  theta = (100000/287.0024093890231) * (P / 100000)^(1/1.4) / ρ
  =#
   step = [0]
  mkpath("vtk")
  cbvtk = GenericCallbacks.EveryXSimulationSteps(100) do (init=false)
    outprefix = @sprintf("vtk/IS_%dD_rank_%d_of_%d_step%04d", dim,
                         MPI.Comm_rank(mpicomm)+1, MPI.Comm_size(mpicomm),
                         step[1])
    @debug "doing VTK output" outprefix
    step[1] == 0 &&
      VanillaAtmosDiscretizations.writevtk(outprefix, Q, spacedisc)
    step[1] += 1
    nothing
  end

  solve!(Q, lsrk; timeend=timeend, callbacks=(cbinfo, cbvtk))

  # Print some end of the simulation information
  engf = norm(Q)
  @info @sprintf """Finished
  norm(Q)            = %.16e
  norm(Q) / norm(Q₀) = %.16e
  norm(Q) - norm(Q₀) = %.16e""" engf engf/eng0 engf-eng0

  engf
end

using Test
let
  MPI.Initialized() || MPI.Init()

  Sys.iswindows() || (isinteractive() && MPI.finalize_atexit())
  mpicomm = MPI.COMM_WORLD
  if MPI.Comm_rank(mpicomm) == 0
    ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
    loglevel = ll == "DEBUG" ? Logging.Debug :
               ll == "WARN"  ? Logging.Warn  :
               ll == "ERROR" ? Logging.Error : Logging.Info
    global_logger(ConsoleLogger(stderr, loglevel))
  else
    global_logger(NullLogger())
  end

  nmoist = 1
  ntrace = 0
  Ne = (10, 10, 1)
  N = 3
  dt = 1e-3
  timeend = 2dt
  for DFloat in (Float64, )
    for ArrayType in (Array,)
      for bricktopo in (BrickTopology, StackedBrickTopology)
        for dim in 2:3
          brickrange = ntuple(j->range(DFloat(0); length=Ne[j]+1, stop=1000),
                              dim)
          main(mpicomm, DFloat, ArrayType, brickrange, nmoist, ntrace, N,
               timeend, bricktopo; dt=dt)
          @test true
        end
      end
    end
  end
end

isinteractive() || MPI.Finalize()

max_stars_repo_name: christophernhill-dev/CLIMA

max_issues_repo_name: christophernhill-dev/CLIMA

max_forks_repo_name: christophernhill-dev/CLIMA


text:
function splice_splines(left_spline::Schumaker, right_spline::Schumaker, splice_point::Real)
    end_in_left_spline = searchsortedlast(left_spline.IntStarts_, splice_point)
    left_starts        = left_spline.IntStarts_[1:end_in_left_spline]
    left_coefficients  = left_spline.coefficient_matrix_[1:end_in_left_spline,:]

    start_in_right     = searchsortedlast(right_spline.IntStarts_, splice_point)
    right_starts       = right_spline.IntStarts_[start_in_right:length(right_spline.IntStarts_)]
    right_coefficients = right_spline.coefficient_matrix_[start_in_right:length(right_spline.IntStarts_),:]

    # As the splice_point is probably not an interval start in the right_spline we need to adjust.
    G                       = splice_point - right_starts[1]
    a                       = right_coefficients[1,1]
    b                       = right_coefficients[1,2]
    c                       = right_coefficients[1,3]
    right_starts[1]         = splice_point
    right_coefficients[1,2] = 2*G*a + b
    right_coefficients[1,3] = a*(G^2) + b*G + c

    IntStarts = vcat(left_starts, right_starts)
    Coeffs    = vcat(left_coefficients, right_coefficients)
    return Schumaker(IntStarts, Coeffs)
end


text:
#AnalogMTK: Diodes
#-------------------------------------------------------------------------------

const kb = 1.380649e-23 #[J/K] Boltzman constant
const qe = -1.602176634e-19 #[C] Electron charge

#Defaults
const DFLT_T = 300 #[K] temperature
const DFLT_VT = kb*DFLT_T/abs(qe) #Diode thermal voltage

struct Shockley
end

#==
===============================================================================#

#Has convergeance issues?
function Diode(::Ideal; name)
@warn("Has convergeance issues???")
	gmin = 1e-12

	@named(p = Pin()); @named(n = Pin())
	@variables v(t) i(t)
	eqs = [
		v ~ p.v - n.v
#		0 ~ p.i + n.i - gmin*(p.v-n.v) #in = -out (w/gmin)
#		i ~ p.i - gmin*p.v #Positive current flows into p terminal
		0 ~ p.i + n.i #in = -out
		i ~ p.i - gmin*v #Positive current flows into p terminal
		0 ~ ifelse(v<0, i, v)
	]
	ODESystem(eqs, t, [v, i], [], systems=[p, n], defaults=Dict(), name=name)
end

function Diode(::Shockley; name::Symbol, IS=1e-15, VT=DFLT_VT, η=1.0)
	ηVT = η*VT

	#Rename before re-using variables:
	_IS = IS; _VT=VT; _η=η
	@parameters IS, VT, η
	defaults = Dict(IS=>_IS, VT=>_VT, η=>_η)

	#Rename before re-using pin name:
#	_n=n
	@named(p = Pin()); @named(n = Pin())
	@variables v(t) i(t)
	eqs = [
		v ~ p.v - n.v
		0 ~ p.i + n.i #in = -out
		i ~ p.i #Positive current flows into p terminal
		i ~ IS*(exp(v/ηVT)-1)
	]

	ODESystem(eqs, t, [v, i], keys(defaults), systems=[p, n]; defaults, name)
end

#Last line


text:
mutable struct FwdsNode
  x::AbstractArray{Float64, 1}
  u::AbstractArray{Float64, 1}
  r::Float64

  FwdsNode(x) = new(x, Float64[], 0.0)
  FwdsNode(x, u) = new(x, u, 0.0)
end

function plan_fwds(x::AbstractArray{Float64, 1}, agent::Agent, world::World,
                   reward::Function, ctrl_d::Discretization, depth::Int)
  root = Tree(FwdsNode(x))

  rs = select_action_fwds(root, agent, world, reward, ctrl_d, depth)
  for node in root.next
    if node.value.r >= rs
      return node.value.u
    end
  end
  error("No best action? FWDS has a bug")

  return Float64[]
end

function select_action_fwds(node::Tree, agent::Agent, world::World,
                            reward::Function, ctrl_d::Discretization,
                            depth::Int)
  if depth <= 0
    return 0.0
  end

  max_r = -Inf
  us = Float64[]

  node.value.x[1] = mod(node.value.x[1], world.road.path.S[end])

  la_len = ctrl_d.thr[end] * ctrl_d.pt[end] # number of actions to survey
  node.next = Array{Tree, 1}(undef, la_len)
  for la in 0:(la_len - 1)
    u = dis.ls2x(ctrl_d, la)
    agent.custom = u
    nx = copy(node.value.x)
    sim.advance!(sim.default_dynamics!, nx, Pair(agent, world), 0.0, olm_dt,
                 olm_h)
    nx[1] = mod(nx[1], world.road.path.S[end])
    r = reward(node.value.x, u, nx, agent, world)
    value = FwdsNode(nx, u)
    node.next[la + 1] = Tree(value)
    next_r = select_action_fwds(node.next[la + 1], agent, world, reward,
                                ctrl_d, depth - 1)
    r = reward(node.value.x, u, nx, agent, world) + olm_gamma * next_r
    node.next[la + 1].value.r = r

    #println("x = $(node.value.x), u = $(u), nx = $(nx), r = $(r)")

    if r > max_r
      max_r = r
    end
  end

  return max_r
end

function controller_fwds!(u::AbstractArray{Float64},
                          x::AbstractArray{Float64},
                          dx::AbstractArray{Float64},
                          agent_world::Pair{Agent, World}, t::Float64)
  agent = agent_world.first
  u[1] = agent.custom[1]
  u[2] = agent.custom[2]
end