text:
# General Julia modules
using ArgParse
using Dierckx
using Distributions
using DocStringExtensions
using LinearAlgebra
using NCDatasets
using Pkg.Artifacts
using Printf
using Random
using StaticArrays
using Test
using UnPack

# ClimateMachine Modules
using ClimateMachine

using ClimateMachine.ArtifactWrappers
using ClimateMachine.Atmos
using ClimateMachine.ConfigTypes
using ClimateMachine.GenericCallbacks
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.Diagnostics
using ClimateMachine.Mesh.Filters
using ClimateMachine.Mesh.Grids
using ClimateMachine.Orientations
using ClimateMachine.Thermodynamics
using ClimateMachine.BalanceLaws
using ClimateMachine.TurbulenceClosures
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.Writers

# Planet parameters
using CLIMAParameters
using CLIMAParameters.Planet: e_int_v0, grav, day
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()
# Physics specific imports
using ClimateMachine.Atmos: altitude, recover_thermo_state
import ClimateMachine.BalanceLaws: source, eq_tends, prognostic_vars

# Citation for problem setup
## CMIP6 Test Dataset - cfsites
## [Webb2017](@cite)

"""
    GCMRelaxation{FT} <: TendencyDef{Source}

"""
struct GCMRelaxation{FT} <: TendencyDef{Source}
    τ_relax::FT
end
prognostic_vars(::GCMRelaxation) = (Mass(), TotalMoisture())

function source(::Mass, s::GCMRelaxation, m, args)
    # TODO: write correct tendency
    return 0
end

function source(::TotalMoisture, s::GCMRelaxation, m, args)
    # TODO: write correct tendency
    return 0
end

"""
    LargeScaleProcess <: TendencyDef{Source}

# Energy tendency ∂_t ρe

Temperature tendency for the LES configuration based on quantities
from a GCM. Quantities are included in standard CMIP naming format.
Tendencies included here are

    tntha = temperature tendency due to horizontal advection
    tntva = temperature tendency due to vertical advection
    tntr = temperature tendency due to radiation fluxes
    ∂T∂z = temperature vertical gradient from GCM values

# Moisture tendency ∂_t ρq_tot

Moisture tendency for the LES configuration based on quantities
from a GCM. Quantities are included in standard CMIP naming format.
Tendencies included here are

    tnhusha = moisture tendency due to horizontal advection
    tnhusva = moisture tendency due to vertical advection
    ∂qt∂z = moisture vertical gradient from GCM values
"""
struct LargeScaleProcess <: TendencyDef{Source} end

prognostic_vars(::LargeScaleProcess) = (Mass(), Energy(), TotalMoisture())

function source(::Energy, s::LargeScaleProcess, m, args)
    @unpack state, aux, diffusive = args
    @unpack ts = args.precomputed
    # Establish problem float-type
    FT = eltype(state)
    # Establish vertical orientation
    k̂ = vertical_unit_vector(m, aux)
    param_set = parameter_set(m)
    _e_int_v0 = e_int_v0(param_set)
    # Unpack vertical gradients
    ∂qt∂z = diffusive.lsforcing.∇ᵥhus
    ∂T∂z = diffusive.lsforcing.∇ᵥta
    w_s = aux.lsforcing.w_s
    cvm = cv_m(ts)
    # Compute tendency terms
    # Temperature contribution
    T_tendency = aux.lsforcing.Σtemp_tendency + ∂T∂z * w_s
    # Moisture contribution
    q_tot_tendency = compute_q_tot_tend(m, args)

    return cvm * state.ρ * T_tendency + _e_int_v0 * state.ρ * q_tot_tendency
end

function compute_q_tot_tend(m, args)
    @unpack aux, diffusive = args
    # Establish problem float-type
    FT = eltype(aux)
    k̂ = vertical_unit_vector(m, aux)
    # Establish vertical orientation
    ∂qt∂z = diffusive.lsforcing.∇ᵥhus
    w_s = aux.lsforcing.w_s
    # Compute tendency terms
    return aux.lsforcing.Σqt_tendency + ∂qt∂z * w_s
end

function source(::Mass, s::LargeScaleProcess, m, args)
    @unpack state = args
    q_tot_tendency = compute_q_tot_tend(m, args)
    return state.ρ * q_tot_tendency
end

function source(::TotalMoisture, s::LargeScaleProcess, m, args)
    @unpack state, aux = args
    q_tot_tendency = compute_q_tot_tend(m, args)
    return state.ρ * q_tot_tendency
end

"""
    LargeScaleSubsidence <: TendencyDef{Source}

Large-scale subsidence tendency, given a vertical velocity
at the large scale, obtained from the GCM data.

```
wap = GCM vertical velocity [Pa s⁻¹]. Note the conversion required
```
"""
struct LargeScaleSubsidence <: TendencyDef{Source} end

prognostic_vars(::LargeScaleSubsidence) = (Mass(), Energy(), TotalMoisture())

function source(::Mass, s::LargeScaleSubsidence, m, args)
    @unpack state, aux, diffusive = args
    # Establish vertical orientation
    k̂ = vertical_unit_vector(m, aux)
    # Establish subsidence velocity
    w_s = aux.lsforcing.w_s
    return -state.ρ * w_s * dot(k̂, diffusive.moisture.∇q_tot)
end
function source(::Energy, s::LargeScaleSubsidence, m, args)
    @unpack state, aux, diffusive = args
    # Establish vertical orientation
    k̂ = vertical_unit_vector(m, aux)
    # Establish subsidence velocity
    w_s = aux.lsforcing.w_s
    return -state.ρ * w_s * dot(k̂, diffusive.energy.∇h_tot)
end
function source(::TotalMoisture, s::LargeScaleSubsidence, m, args)
    @unpack state, aux, diffusive = args
    # Establish vertical orientation
    k̂ = vertical_unit_vector(m, aux)
    # Establish subsidence velocity
    w_s = aux.lsforcing.w_s
    return -state.ρ * w_s * dot(k̂, diffusive.moisture.∇q_tot)
end

# Sponge relaxation
"""
    LinearSponge{FT} <: TendencyDef{Source}

Two parameter sponge (α_max, γ) for velocity relaxation to a reference
state.
    α_max = Sponge strength (can be interpreted as timescale)
    γ = Sponge exponent
    z_max = Domain height
    z_sponge = Altitude at which sponge layer starts
"""
struct LinearSponge{FT} <: TendencyDef{Source}
    "Maximum domain altitude (m)"
    z_max::FT
    "Altitude at with sponge starts (m)"
    z_sponge::FT
    "Sponge Strength 0 ⩽ α_max ⩽ 1"
    α_max::FT
    "Sponge exponent"
    γ::FT
end

prognostic_vars(::LinearSponge) = (Momentum(),)

function source(::Momentum, s::LinearSponge, m, args)
    @unpack state, aux = args
    #Unpack sponge parameters
    FT = eltype(state)
    @unpack z_max, z_sponge, α_max, γ = s
    # Establish sponge relaxation velocity
    u_geo = SVector(aux.lsforcing.ua, aux.lsforcing.va, 0)
    z = altitude(m, aux)
    # Accumulate sources
    if z_sponge <= z
        r = (z - z_sponge) / (z_max - z_sponge)
        #ZS: different sponge formulation?
        β_sponge = α_max .* sinpi(r / FT(2)) .^ γ
        return -β_sponge * (state.ρu .- state.ρ * u_geo)
    else
        return SVector{3, FT}(0, 0, 0)
    end
end

# We first specify the NetCDF file from which we wish to read our
# GCM values.
# Utility function to read and store variables directly from the
# NetCDF file
"""
    str2var(str::String, var::Any)

Helper function allowing variables read in from the GCM file
to be made available to the LES simulation.
"""
function str2var(str::String, var::Any)
    str = Symbol(str)
    @eval(($str) = ($var))
end

# Define the get_gcm_info function
const forcingfile = "HadGEM2-A_amip.2004-2008.07"
"""
    get_gcm_info(group_id)

For a specific global site, establish and store the GCM state
for each available vertical level. `group_id` refers to the integer
index of the specific global site that we are interested in.
"""
function get_gcm_info(group_id)

    @printf("--------------------------------------------------\n")
    @info @sprintf("""\n
     Experiment: GCM(HadGEM2-A) driven LES(ClimateMachine)
     """)

    @printf("\n")
    @printf("HadGEM2-A_LES = %s\n", group_id)
    @printf("--------------------------------------------------\n")

    lsforcing_dataset = ArtifactWrapper(
        @__DIR__,
        isempty(get(ENV, "CI", "")),
        "lsforcing",
        ArtifactFile[ArtifactFile(
            url = "https://caltech.box.com/shared/static/dszfbqzwgc9a55vhxd43yenvebcb6bcj.nc",
            filename = forcingfile,
        ),],
    )
    lsforcing_dataset_path = get_data_folder(lsforcing_dataset)
    data_file = joinpath(lsforcing_dataset_path, forcingfile)

    req_varnames = (
        "zg",
        "ta",
        "hus",
        "ua",
        "va",
        "pfull",
        "tntha",
        "tntva",
        "tntr",
        "tnhusha",
        "tnhusva",
        "wap",
        "hfls",
        "hfss",
        "ts",
        "alpha",
    )
    # Load NETCDF dataset (HadGEM2-A information)
    # Load the NCDataset (currently we assume all time-stamps are
    # in the same NCData file). We store this information in `data`.
    data = NCDataset(data_file)
    # To assist the user / inform them of the data processing step
    # we print out some useful information, such as groupnames
    # and a list of available variables
    @printf("Storing information for group %s ...", group_id)
    for (varname, var) in data.group[group_id]
        for reqvar in req_varnames
            reqvar ≠ varname && continue
            # Get average over time dimension
            var = mean(var, dims = 2)
            if any(varname == y for y in ("hfls", "hfss", "ts"))
                var = mean(var, dims = 1)[1]
            end
            # Assign the variable values to the appropriate converted string
            str2var(varname, var)
        end
        # Store key variables
    end
    @printf("Complete\n")
    @printf("--------------------------------------------------\n")
    @printf("Group data storage complete\n")
    return (
        zg = zg,
        ta = ta,
        hus = hus,
        ua = ua,
        va = va,
        pfull = pfull,
        tntha = tntha,
        tntva = tntva,
        tntr = tntr,
        tnhusha = tnhusha,
        tnhusva = tnhusva,
        wap = wap,
        hfls = hfls,
        hfss = hfss,
        ts = ts,
        alpha = alpha,
    )

end

# Initialise the CFSite experiment :D!
const seed = MersenneTwister(0)
function init_cfsites!(problem, bl, state, aux, localgeo, t, spl)
    FT = eltype(state)
    param_set = parameter_set(bl)
    (x, y, z) = localgeo.coord
    _grav = FT(grav(param_set))

    # Unpack splines, interpolate to z coordinate at
    # present grid index. (Functions are all pointwise)
    ta = FT(spl.spl_ta(z))
    hus = FT(spl.spl_hus(z))
    ua = FT(spl.spl_ua(z))
    va = FT(spl.spl_va(z))
    pfull = FT(spl.spl_pfull(z))
    tntha = FT(spl.spl_tntha(z))
    tntva = FT(spl.spl_tntva(z))
    tntr = FT(spl.spl_tntr(z))
    tnhusha = FT(spl.spl_tnhusha(z))
    tnhusva = FT(spl.spl_tnhusva(z))
    wap = FT(spl.spl_wap(z))
    ρ_gcm = FT(1 / spl.spl_alpha(z))

    param_set = parameter_set(bl)
    # Compute field properties based on interpolated data
    ρ = air_density(param_set, ta, pfull, PhasePartition(hus))
    e_int = internal_energy(param_set, ta, PhasePartition(hus))
    e_kin = (ua^2 + va^2) / 2
    e_pot = _grav * z
    # Assignment of state variables
    state.ρ = ρ
    state.ρu = ρ * SVector(ua, va, 0)
    state.energy.ρe = ρ * (e_kin + e_pot + e_int)
    state.moisture.ρq_tot = ρ * hus
    if z <= FT(400)
        state.energy.ρe += rand(seed) * FT(1 / 100) * (state.energy.ρe)
        state.moisture.ρq_tot +=
            rand(seed) * FT(1 / 100) * (state.moisture.ρq_tot)
    end

    # Assign and store the ref variable for sources
    aux.lsforcing.ta = ta
    aux.lsforcing.hus = hus
    aux.lsforcing.Σtemp_tendency = tntha + tntva + tntr
    aux.lsforcing.ua = ua
    aux.lsforcing.va = va
    aux.lsforcing.Σqt_tendency = tnhusha + tnhusva
    aux.lsforcing.w_s = -wap / ρ_gcm / _grav
    return nothing
end

function config_cfsites(
    ::Type{FT},
    N,
    resolution,
    xmax,
    ymax,
    zmax,
    hfls,
    hfss,
    ts,
    group_id,
) where {FT}
    # Boundary Conditions
    u_star = FT(0.28)

    physics = AtmosPhysics{FT}(
        param_set;
        turbulence = Vreman{FT}(0.23),
        moisture = EquilMoist(; maxiter = 5, tolerance = FT(2)),
        lsforcing = HadGEMVertical(),
    )

    problem = AtmosProblem(
        boundaryconditions = (
            AtmosBC(
                physics;
                momentum = Impenetrable(DragLaw(
                    (state, aux, t, normPu_int) -> (u_star / normPu_int)^2,
                )),
                energy = PrescribedEnergyFlux((state, aux, t) -> hfls + hfss),
                moisture = PrescribedMoistureFlux(
                    (state, aux, t) ->
                        hfls / latent_heat_vapor(param_set, ts),
                ),
            ),
            AtmosBC(physics;),
        ),
        init_state_prognostic = init_cfsites!,
    )

    model = AtmosModel{FT}(
        AtmosLESConfigType,
        physics;
        problem = problem,
        source = (
            Gravity(),
            LinearSponge{FT}(zmax, zmax * 0.85, 1, 4),
            LargeScaleProcess(),
            LargeScaleSubsidence(),
        ),
    )

    # Timestepper options

    # Explicit Solver
    ex_solver = ClimateMachine.ExplicitSolverType()

    # Multirate Explicit Solver
    mrrk_solver = ClimateMachine.MultirateSolverType(
        fast_model = AtmosAcousticGravityLinearModel,
        slow_method = LSRK144NiegemannDiehlBusch,
        fast_method = LSRK144NiegemannDiehlBusch,
        timestep_ratio = 10,
    )

    # IMEX Solver Type
    imex_solver = ClimateMachine.IMEXSolverType()

    # Configuration
    config = ClimateMachine.AtmosLESConfiguration(
        forcingfile * "_$group_id",
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        param_set,
        init_cfsites!,
        #ZS: multi-rate?
        solver_type = imex_solver,
        model = model,
    )
    return config
end

# Define the diagnostics configuration (Atmos-Default)
function config_diagnostics(driver_config)
    default_dgngrp = setup_atmos_default_diagnostics(
        AtmosLESConfigType(),
        "2500steps",
        driver_config.name,
    )
    core_dgngrp = setup_atmos_core_diagnostics(
        AtmosLESConfigType(),
        "2500steps",
        driver_config.name,
    )
    return ClimateMachine.DiagnosticsConfiguration([
        default_dgngrp,
        core_dgngrp,
    ])
end

function main()

    # Provision for custom command line arguments
    # Convenience args for slurm-array launches
    cfsite_args = ArgParseSettings(autofix_names = true)
    add_arg_group!(cfsite_args, "HadGEM2-A_SiteInfo")
    @add_arg_table! cfsite_args begin
        "--group-id"
        help = "Specify cfSite-ID for forcing data"
        metavar = "site<number>"
        arg_type = String
        default = "site17"
    end
    cl_args = ClimateMachine.init(
        parse_clargs = true,
        custom_clargs = cfsite_args,
        fix_rng_seed = true,
    )
    group_id = cl_args["group_id"]

    # Working precision
    FT = Float64
    # DG polynomial order
    N = 4
    # Domain resolution and size
    Δh = FT(75)
    Δv = FT(20)
    resolution = (Δh, Δh, Δv)
    # Domain extents
    xmax = FT(1800)
    ymax = FT(1800)
    zmax = FT(4000)
    # Simulation time
    t0 = FT(0)
    timeend = FT(600)
    #timeend = FT(3600 * 6)
    # Courant number
    CFL = FT(0.2)

    # Execute the get_gcm_info function
    ls_forcing = get_gcm_info(group_id)

    # Drop dimensions for compatibility with Dierckx
    z = dropdims(ls_forcing.zg; dims = 2)
    # Create spline objects and pass them into a named tuple
    splines = (
        spl_ta = Spline1D(z, view(ls_forcing.ta, :, 1)),
        spl_pfull = Spline1D(z, view(ls_forcing.pfull, :, 1)),
        spl_ua = Spline1D(z, view(ls_forcing.ua, :, 1)),
        spl_va = Spline1D(z, view(ls_forcing.va, :, 1)),
        spl_hus = Spline1D(z, view(ls_forcing.hus, :, 1)),
        spl_tntha = Spline1D(z, view(ls_forcing.tntha, :, 1)),
        spl_tntva = Spline1D(z, view(ls_forcing.tntva, :, 1)),
        spl_tntr = Spline1D(z, view(ls_forcing.tntr, :, 1)),
        spl_tnhusha = Spline1D(z, view(ls_forcing.tnhusha, :, 1)),
        spl_tnhusva = Spline1D(z, view(ls_forcing.tnhusva, :, 1)),
        spl_wap = Spline1D(z, view(ls_forcing.wap, :, 1)),
        spl_alpha = Spline1D(z, view(ls_forcing.alpha, :, 1)),
    )

    # Set up driver configuration
    driver_config = config_cfsites(
        FT,
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        hfls,
        hfss,
        ts,
        group_id,
    )
    # Set up solver configuration
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        splines;
        init_on_cpu = true,
        Courant_number = CFL,
        CFL_direction = HorizontalDirection(),
    )
    # Set up diagnostic configuration
    dgn_config = config_diagnostics(driver_config)

    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            ("moisture.ρq_tot",),
            solver_config.dg.grid,
            TMARFilter(),
        )
        nothing
    end

    #ZS: cutoff filter?
    filterorder = 2 * N
    filter = BoydVandevenFilter(solver_config.dg.grid, 0, filterorder)
    cbfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            AtmosFilterPerturbations(driver_config.bl),
            solver_config.dg.grid,
            filter,
        )
        nothing
    end

    cutoff_filter = CutoffFilter(solver_config.dg.grid, N - 1)
    cbcutoff = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(solver_config.Q, 1:6, solver_config.dg.grid, cutoff_filter)
        nothing
    end

    # Invoke solver (calls solve! function for time-integrator)
    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        #ZS: only tmar?
        user_callbacks = (cbtmarfilter,),
        check_euclidean_distance = true,
    )

end
main()


text:
using TOML
TOML.parsefile("ColorMacros.toml")

ext: jl

max_stars_repo_name: periodicaidan/oryx-kbd-colors

max_issues_repo_name: periodicaidan/oryx-kbd-colors

max_forks_repo_name: periodicaidan/oryx-kbd-colors


text:
using Test
import ShuffleProofs: decode, encode, Tree, unmarshal, marshal_publickey, unmarshal_publickey, unmarshal_privatekey
import CryptoGroups: ElGamal, PGroup, Dec

CIPHERTEXT_FILE = "$(@__DIR__)/../../ref/demo/ciphertexts"
CIPHERTEXTOUT_FILE = "$(@__DIR__)/../../ref/demo/ciphertextsout"
PUBLIC_KEY = "$(@__DIR__)/../../ref/demo/publicKey"
PRIVATE_KEY = "$(@__DIR__)/../../ref/privateKey"

y, g = let
    bytes = read(PUBLIC_KEY)
    tree = decode(bytes)
    unmarshal_publickey(tree)
end

sk, g′ = let
    bytes = read(PRIVATE_KEY)
    tree = decode(bytes)
    unmarshal_privatekey(tree)
end

@test g′ == g
@test y == g^sk # This way we can be sure that correct inputs have been choosen.

# It would be cool to also to input my own ciphertexts!

G = typeof(g)

#𝓖 = group(g)

𝐞 = let
    bytes = read(CIPHERTEXT_FILE)
    tree = decode(bytes)
    𝐚, 𝐛 = convert(Tuple{Vector{BigInt}, Vector{BigInt}}, tree)
    ElGamal{G}(𝐚, 𝐛)
end

𝐞′ = let
    bytes = read(CIPHERTEXTOUT_FILE)
    tree = decode(bytes)
    convert(ElGamal{G}, tree)
end

dec = Dec(sk)

𝐦 = g .^ (2:11)
@test dec(𝐞) == 𝐦

𝐦′ = dec(𝐞′)
@test sort(𝐦) == sort(𝐦′)


text:
module Docker

using HTTP
using JSON
using Compat

struct DockerError
    status::Int
    msg::Compat.String
end

const headers = Dict("Content-Type" => "application/json")

docker_uri(host) = HTTP.URI("http://$host/v1.21")
docker_uri(host,endpoint) = HTTP.URI("http://$host/v1.21/$endpoint")
parse_body(data) = JSON.parse(join(map(Char,data)))

function create_container(
        host, image;
        cmd::Cmd     = ``,
        entryPoint   = "",
        tty_stream   = true,
        attachStdin  = false,
        openStdin    = false,
        attachStdout = true,
        attachStderr = true,
        memory       = 0,
        cpuSets      = "",
        volumeDriver = "",
        portBindings = ["",""], # [ContainerPort,HostPort]
        ports        = [],
        workingDir   = ""
    )

    url = docker_uri(host)

    params = Dict(
        "Image" => image,
        "Tty" => tty_stream,
        "AttachStdin"   => attachStdin,
        "OpenStdin"     => openStdin,
        "AttachStdout"  => attachStdout,
        "AttachStderr"  => attachStderr,
        "ExposedPorts"  => Dict([Pair(string(dec(p),"/tcp"), Dict()) for p in ports]),
        "HostConfig"    => Dict(
            "Memory"       => memory,
            "CpusetCpus"   => cpuSets,
            "VolumeDriver" => volumeDriver,
            "PortBindings" => Dict(
                string(portBindings[1],"/tcp") => [
                    Dict("HostPort" => string(portBindings[2]))
                ]
            )
        )
    )

    if !isempty(entryPoint)
        params["Entrypoint"] = entryPoint
    end

    if !isempty(cmd.exec)
        params["Cmd"] = collect(cmd.exec)
    end

    if !isempty(workingDir)
        params["WorkingDir"] = workingDir
    end

    resp = HTTP.request("POST",HTTP.URI("$url/containers/create"),[],JSON.json(params),headers=headers)
    if resp.status != 201
        throw(DockerError(resp.status,resp.body))
    end
    parse_body(resp.body)
end

function inspect_container(host,id)
    resp = HTTP.request("GET",(docker_uri(host,"containers/$id/json")))
    if resp.status != 200
        throw(DockerError(resp.status,resp.body))
    end
    parse_body(resp.body)
end

function start_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/start")))
    print(resp)
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    id
end

function restart_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/restart")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    resp
end

function stop_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/stop")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    resp
end

function pause_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/pause")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    resp
end

function unpause_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/unpause")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    id
end

function kill_container(host, id)
    resp = HTTP.request("POST",(docker_uri(host,"containers/$id/kill")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    resp
end

function remove_container(host, id)
    resp = HTTP.request("DELETE",(docker_uri(host,"containers/$id?force=1")))
    if resp.status != 204
        throw(DockerError(resp.status,resp.body))
    end
    resp
end

function processes_container(host, id)
    resp = HTTP.request("GET",(docker_uri(host,"containers/$id/top")))
    if resp.status != 204
        println(resp.status)
    end
    if resp.status != 200
        throw(DockerError(resp.status,resp.body))
    end
    parse_body(resp.body)
end

function list_containers(host)
    resp = HTTP.request("GET",(docker_uri(host,"containers/json")))
    print(resp)
    if resp.status != 200
        throw(DockerError(resp.status,resp.body))
    end
    parse_body(resp.body)
end

function stats_container(host,id)
    resp = HTTP.request("GET",(docker_uri(host,"containers/$id/stats?stream=0")))
    if resp.status != 200
        throw(DockerError(resp.status,resp.body))
    end
    parse_body(resp.body)
end

function open_logs_stream(host, id; history=false)
    path = "containers/$id/attach?logs&follow=1&stdout=1"
    if history
        path *=  "&logs=1"
    end
    url = docker_uri(host,path)
    HTTP.open(url,"GET",[Dict("Content-Type"=>"plain/text")],"","POST")
end

function cleanse!(host)
    resp = HTTP.request("GET",(docker_uri(host,"containers/json?all=true")))
    if resp.status != 200
        throw(DockerError(resp.status,resp.body))
    end
    data = parse_body(resp.body)
    for c in data
        remove_container(host,c["Id"])
        kill_container(host,c["Id"])
    end
    nothing
end

end # module Docker


text:
"""
    read_fasta_alignment(filename::AbstractString, max_gap_fraction::Real)

Return a `L × M` matrix of integers (`L` is the sequence length, and `M` is the number of sequences) of the
multiple sequence alignment contained in the fasta file `filename` including all sequences with a fraction of gaps (`-`) ≤
`max_gap_fraction`.
"""
function read_fasta_alignment(filename::AbstractString, max_gap_fraction::Real)

    f = FastaReader(filename)

    max_gap_fraction = Float64(max_gap_fraction)

    # pass 1

    seqs = Int[]
    inds = Int[]
    fseqlen = 0

    for (name, seq) in f
        ngaps = 0
        if f.num_parsed == 1
            ls = length(seq)
            resize!(inds, ls)
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    fseqlen += 1
                    inds[fseqlen] = i
                    c == '-' && (ngaps += 1)
                end
            end
        else
            ls = length(seq)
            ls == length(inds) || error("inputs are not aligned")
            tstfseqlen = 0
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    tstfseqlen += 1
                    inds[tstfseqlen] == i || error("inconsistent inputs")
                    c == '-' && (ngaps += 1)
                end
            end
            tstfseqlen == fseqlen || error("inconsistent inputs")
        end
        ngaps / fseqlen <= max_gap_fraction && push!(seqs, f.num_parsed)
    end

    length(seqs) > 0 || error("Out of $(f.num_parsed) sequences, none passed the filter (max_gap_fraction=$max_gap_fraction)")

    # pass 2
    Z = Array{Int8}(undef, fseqlen, length(seqs))

    seqid = 1
    for (name, seq) in f
        seqs[end] < f.num_parsed && break
        seqs[seqid] == f.num_parsed || continue
        for i = 1:fseqlen
            c = seq[inds[i]]
            Z[i, seqid] = letter2num(c)
        end
        seqid += 1
    end
    @assert seqid == length(seqs) + 1

    close(f)

    return Z
end

let alphabet = [ 1,21, 2, 3, 4, 5, 6, 7, 8,21, 9,10,11,12,21,13,14,15,16,17,21,18,19,21,20]
               # A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R, S, T, U, V, W, X, Y
    global letter2num
    function letter2num(c::Union{Char,UInt8})
        i = UInt8(c) - 0x40
        1 <= i <= 25 && return alphabet[i]
        return 21
    end
end

"""
    remove_duplicate_sequences(Z::Matrix{Ti}) where Ti<:Integer
Remove duplicate sequences (columns) in the alignment matrix `Z`

# Examples
```
julia> Z = [1 2 3 1;
            1 3 2 1;]
2×4 Array{Int64,2}:
 1  2  3  1
 1  3  2  1

julia> remove_duplicate_sequences(Z)
removing duplicate sequences... done: 4 -> 3
([1 2 3; 1 3 2], [1, 2, 3])
```
"""
function remove_duplicate_sequences(Z::Matrix{Ti}) where Ti<:Integer

    N, M = size(Z)
    hZ = Array{UInt}(undef, M)
    @inbounds for i = 1:M
        hZ[i] = hash(Z[:,i])
    end
    print("removing duplicate sequences... ")
    ref_seq_ind = Array{Int}(undef, M)
    ref_seq = Dict{UInt,Int}()
    @inbounds for i = 1:M
        ref_seq_ind[i] = get!(ref_seq, hZ[i], i)
    end
    uniqueseqs = collect(values(ref_seq))
    # Check for collisions
    collided = falses(M)
    @inbounds for i in 1:M
        k = ref_seq_ind[i]
        k == i && continue
        for j = 1:N
            Z[j,i] != Z[j,k] && (collided[i] = true; break)
        end
    end
     if any(collided)
        nowcollided = BitArray(undef, M)
        while any(collided)
            # Collect index of first row for each collided hash
            empty!(ref_seq)
            @inbounds for i = 1:M
                collided[i] || continue
                ref_seq_ind[i] = get!(ref_seq, hZ[i], i)
            end
            for v in values(ref_seq)
                push!(uniqueseqs, v)
            end

            # Check for collisions
            fill!(nowcollided, false)
            @inbounds for i = 1:M
                k = ref_seq_ind[i]
                (!collided[i] || k == i) && continue
                for j = 1:N
                    Z[j,i] != Z[j,k] && (nowcollided[i] = true; break)
                end
            end
            collided, nowcollided = nowcollided, collided
        end
    end

    sort!(uniqueseqs)

    newM = length(uniqueseqs)
    newZ = Z[:,uniqueseqs]
    println("done: $M -> $newM")
    return newZ,uniqueseqs
end

max_stars_repo_name: pagnani/CorrDCA

max_issues_repo_name: pagnani/CorrDCA

max_forks_repo_name: pagnani/CorrDCA


text:
using ModelingToolkit, StaticArrays, LinearAlgebra
using StochasticDiffEq, SparseArrays
using Random,Test

# Define some variables
@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

noiseeqs = [0.1*x,
            0.1*y,
            0.1*z]

# ODESystem -> SDESystem shorthand constructor
@named sys = ODESystem(eqs,t,[x,y,z],[σ,ρ,β])
@test SDESystem(sys, noiseeqs, name=:foo) isa SDESystem

@named de = SDESystem(eqs,noiseeqs,t,[x,y,z],[σ,ρ,β])
f = eval(generate_diffusion_function(de)[1])
@test f(ones(3),rand(3),nothing) == 0.1ones(3)

f = SDEFunction(de)
prob = SDEProblem(SDEFunction(de),f.g,[1.0,0.0,0.0],(0.0,100.0),(10.0,26.0,2.33))
sol = solve(prob,SRIW1(),seed=1)

probexpr = SDEProblem(SDEFunction(de),f.g,[1.0,0.0,0.0],(0.0,100.0),(10.0,26.0,2.33))
solexpr = solve(eval(probexpr),SRIW1(),seed=1)

@test all(x->x==0,Array(sol-solexpr))

# Test no error
@test_nowarn SDEProblem(de,nothing,(0, 10.0))

noiseeqs_nd = [0.01*x 0.01*x*y 0.02*x*z
               σ      0.01*y   0.02*x*z
               ρ      β        0.01*z  ]
@named de = SDESystem(eqs,noiseeqs_nd,t,[x,y,z],[σ,ρ,β])
f = eval(generate_diffusion_function(de)[1])
@test f([1,2,3.0],[0.1,0.2,0.3],nothing) == [0.01*1   0.01*1*2   0.02*1*3
                                             0.1      0.01*2     0.02*1*3
                                             0.2      0.3        0.01*3  ]

f = eval(generate_diffusion_function(de)[2])
du = ones(3,3)
f(du,[1,2,3.0],[0.1,0.2,0.3],nothing)
@test du == [0.01*1   0.01*1*2   0.02*1*3
             0.1      0.01*2     0.02*1*3
             0.2      0.3        0.01*3  ]

f = SDEFunction(de)
prob = SDEProblem(SDEFunction(de),f.g,[1.0,0.0,0.0],(0.0,100.0),(10.0,26.0,2.33),
                  noise_rate_prototype = zeros(3,3))
sol = solve(prob,EM(),dt=0.001)

u0map = [
    x => 1.0,
    y => 0.0,
    z => 0.0
]

parammap = [
    σ => 10.0,
    β => 26.0,
    ρ => 2.33
]

prob = SDEProblem(de,u0map,(0.0,100.0),parammap)
@test size(prob.noise_rate_prototype) == (3,3)
@test prob.noise_rate_prototype isa Matrix
sol = solve(prob,EM(),dt=0.001)

prob = SDEProblem(de,u0map,(0.0,100.0),parammap,sparsenoise=true)
@test size(prob.noise_rate_prototype) == (3,3)
@test prob.noise_rate_prototype isa SparseMatrixCSC
sol = solve(prob,EM(),dt=0.001)

# Test eval_expression=false
function test_SDEFunction_no_eval()
    # Need to test within a function scope to trigger world age issues
    f = SDEFunction(de, eval_expression=false)
    @test f([1.0,0.0,0.0], (10.0,26.0,2.33), (0.0,100.0)) ≈ [-10.0, 26.0, 0.0]
end
test_SDEFunction_no_eval()

# modelingtoolkitize and Ito <-> Stratonovich sense
seed = 10
Random.seed!(seed)

# simple 2D diagonal noise
u0 = rand(2)
t =  randn()
trange = (0.0,100.0)
p = [1.01,0.87]
f1!(du,u,p,t) = (du .= p[1]*u)
σ1!(du,u,p,t) = (du .= p[2]*u)

prob = SDEProblem(f1!,σ1!,u0,trange,p)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) == p[1]*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == p[1]*u0 - 1//2*p[2]^2*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0 - 1//2*p[2]^2*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == p[1]*u0 + 1//2*p[2]^2*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0 + 1//2*p[2]^2*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# somewhat complicated 1D without explicit parameters but with explicit time-dependence
f2!(du,u,p,t) = (du[1] = sin(t) + cos(u[1]))
σ2!(du,u,p,t) = (du[1] = pi + atan(u[1]))

u0 = rand(1)
prob = SDEProblem(f2!,σ2!,u0,trange)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) ==  @. sin(t) + cos(u0)
@test fdif(u0,p,t) ==  pi .+ atan.(u0)
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == @. sin(t) + cos(u0)
fdif!(du,u0,p,t)
@test du == pi .+ atan.(u0)

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) ==  @. sin(t) + cos(u0) - 1//2*1/(1 + u0^2)*(pi + atan(u0))
@test fdif(u0,p,t) == pi .+ atan.(u0)
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ==  @. sin(t) + cos(u0) - 1//2*1/(1 + u0^2)*(pi + atan(u0))
fdif!(du,u0,p,t)
@test du == pi .+ atan.(u0)

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) ≈  @. sin(t) + cos(u0) + 1//2*1/(1 + u0^2)*(pi + atan(u0))
@test fdif(u0,p,t) == pi .+ atan.(u0)
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ≈  @. sin(t) + cos(u0) + 1//2*1/(1 + u0^2)*(pi + atan(u0))
fdif!(du,u0,p,t)
@test du == pi .+ atan.(u0)

# 2D diagonal noise with mixing terms (no parameters, no time-dependence)
u0 = rand(2)
t =  randn()
function f3!(du,u,p,t)
  du[1] = u[1]/2
  du[2] = u[2]/2
  return nothing
end
function σ3!(du,u,p,t)
  du[1] = u[2]
  du[2] = u[1]
  return nothing
end

prob = SDEProblem(f3!,σ3!,u0,trange,p)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) == u0/2
@test fdif(u0,p,t) == reverse(u0)
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == u0/2
fdif!(du,u0,p,t)
@test du ==  reverse(u0)

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == u0*0
@test fdif(u0,p,t) == reverse(u0)
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == u0*0
fdif!(du,u0,p,t)
@test du == reverse(u0)

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == u0
@test fdif(u0,p,t) == reverse(u0)
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == u0
fdif!(du,u0,p,t)
@test du == reverse(u0)

# simple 2D diagonal noise oop
u0 = rand(2)
t =  randn()
p = [1.01,0.87]
f1(u,p,t) = p[1]*u
σ1(u,p,t) = p[2]*u

prob = SDEProblem(f1,σ1,u0,trange,p)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) == p[1]*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == p[1]*u0 - 1//2*p[2]^2*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0 - 1//2*p[2]^2*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == p[1]*u0 + 1//2*p[2]^2*u0
@test fdif(u0,p,t) == p[2]*u0
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0 + 1//2*p[2]^2*u0
fdif!(du,u0,p,t)
@test du == p[2]*u0

# non-diagonal noise
u0 = rand(2)
t =  randn()
p = [1.01,0.3,0.6,1.2,0.2]
f4!(du,u,p,t) = du .= p[1]*u
function g4!(du,u,p,t)
  du[1,1] = p[2]*u[1]
  du[1,2] = p[3]*u[1]
  du[2,1] = p[4]*u[1]
  du[2,2] = p[5]*u[2]
  return nothing
end

prob = SDEProblem(f4!,g4!,u0,trange,noise_rate_prototype=zeros(2,2),p)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) == p[1]*u0
@test fdif(u0,p,t) == [p[2]*u0[1]   p[3]*u0[1]
                      p[4]*u0[1]     p[5]*u0[2] ]
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == p[1]*u0
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [p[2]*u0[1]   p[3]*u0[1]
             p[4]*u0[1]   p[5]*u0[2]  ]

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) ≈ [p[1]*u0[1] - 1//2*(p[2]^2*u0[1]+p[3]^2*u0[1]), p[1]*u0[2] - 1//2*(p[2]*p[4]*u0[1]+p[5]^2*u0[2])]
@test fdif(u0,p,t) == [p[2]*u0[1]   p[3]*u0[1]
                      p[4]*u0[1]     p[5]*u0[2] ]
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ≈ [p[1]*u0[1] - 1//2*(p[2]^2*u0[1]+p[3]^2*u0[1]), p[1]*u0[2] - 1//2*(p[2]*p[4]*u0[1]+p[5]^2*u0[2])]
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [p[2]*u0[1]   p[3]*u0[1]
            p[4]*u0[1]     p[5]*u0[2] ]

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) ≈ [p[1]*u0[1] + 1//2*(p[2]^2*u0[1]+p[3]^2*u0[1]), p[1]*u0[2] + 1//2*(p[2]*p[4]*u0[1]+p[5]^2*u0[2])]
@test fdif(u0,p,t) == [p[2]*u0[1]   p[3]*u0[1]
                      p[4]*u0[1]     p[5]*u0[2] ]
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ≈ [p[1]*u0[1] + 1//2*(p[2]^2*u0[1]+p[3]^2*u0[1]), p[1]*u0[2] + 1//2*(p[2]*p[4]*u0[1]+p[5]^2*u0[2])]
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [p[2]*u0[1]   p[3]*u0[1]
            p[4]*u0[1]     p[5]*u0[2] ]

# non-diagonal noise: Torus -- Strat and Ito are identical
u0 = rand(2)
t =  randn()
p =  rand(1)
f5!(du,u,p,t) = du .= false
function g5!(du,u,p,t)
  du[1,1] = cos(p[1])*sin(u[1])
  du[1,2] = cos(p[1])*cos(u[1])
  du[1,3] = -sin(p[1])*sin(u[2])
  du[1,4] = -sin(p[1])*cos(u[2])
  du[2,1] = sin(p[1])*sin(u[1])
  du[2,2] = sin(p[1])*cos(u[1])
  du[2,3] = cos(p[1])*sin(u[2])
  du[2,4] = cos(p[1])*cos(u[2])
  return nothing
end

prob = SDEProblem(f5!,g5!,u0,trange,noise_rate_prototype=zeros(2,4),p)
# no correction
sys = modelingtoolkitize(prob)
fdrift = eval(generate_function(sys)[1])
fdif = eval(generate_diffusion_function(sys)[1])
@test fdrift(u0,p,t) == 0*u0
@test fdif(u0,p,t) == [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
                        sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]
fdrift! = eval(generate_function(sys)[2])
fdif! = eval(generate_diffusion_function(sys)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du == 0*u0
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
                        sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]

# Ito -> Strat
sys2 = stochastic_integral_transform(sys,-1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) == 0*u0
@test fdif(u0,p,t) == [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
                        sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ==  0*u0
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
              sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]

# Strat -> Ito
sys2 = stochastic_integral_transform(sys,1//2)
fdrift = eval(generate_function(sys2)[1])
fdif = eval(generate_diffusion_function(sys2)[1])
@test fdrift(u0,p,t) ==  0*u0
@test fdif(u0,p,t) ==  [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
                         sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]
fdrift! = eval(generate_function(sys2)[2])
fdif! = eval(generate_diffusion_function(sys2)[2])
du = similar(u0)
fdrift!(du,u0,p,t)
@test  du ==  0*u0
du = similar(u0, size(prob.noise_rate_prototype))
fdif!(du,u0,p,t)
@test du == [ cos(p[1])*sin(u0[1])   cos(p[1])*cos(u0[1])   -sin(p[1])*sin(u0[2])   -sin(p[1])*cos(u0[2])
              sin(p[1])*sin(u0[1])   sin(p[1])*cos(u0[1])    cos(p[1])*sin(u0[2])    cos(p[1])*cos(u0[2])]

# issue #819
@testset "Combined system name collisions" begin
    @variables t
    eqs_short = [D(x) ~ σ*(y-x),
                D(y) ~ x*(ρ-z)-y,
                ]
    sys1 = SDESystem(eqs_short, noiseeqs, t, [x, y, z], [σ, ρ, β], name = :sys1)
    sys2 = SDESystem(eqs_short, noiseeqs, t, [x, y, z], [σ, ρ, β], name = :sys1)
    @test_throws ArgumentError SDESystem([sys2.y ~ sys1.z], t, [], [], [], systems = [sys1, sys2], name=:foo)
end


text:
"""
    CODE_TO_LANG::Dict{String,String}

    A constant dictionary maps 51 language codes to their names
"""
const CODE_TO_LANG =
Dict{String,String}(
    "sk" => "Slovak",
    "id" => "Indonesian",
    "lv" => "Latvian",
    "ca" => "Catalan",
    "sl" => "Slovenian",
    "hi" => "Hindi",
    "sa" => "Sanskrit",
    "el" => "Greek",
    "nl" => "Dutch",
    "hr" => "Croatian",
    "ug" => "Uyghur",
    "gl" => "Galician",
    "ru" => "Russian",
    "ur" => "Urdu",
    "en" => "English",
    "bxr" => "Buryat",
    "pl" => "Polish",
    "cop" => "Coptic",
    "grc" => "Ancient_Greek",
    "ta" => "Tamil",
    "eu" => "Basque",
    "fa" => "Persian",
    "da" => "Danish",
    "tr" => "Turkish",
    "it" => "Italian",
    "got" => "Gothic",
    "ko" => "Korean",
    "fi" => "Finnish",
    "fo" => "Faroese",
    "zh" => "Chinese",
    "pt" => "Portuguese",
    "ro" => "Romanian",
    "ar" => "Arabic",
    "de" => "German",
    "bg" => "Bulgarian",
    "la" => "Latin",
    "uk" => "Ukrainian",
    "et" => "Estonian",
    "cs" => "Czech",
    "ja" => "Japanese",
    "sv" => "Swedish",
    "no" => "Norwegian",
    "fr" => "French",
    "ga" => "Irish",
    "es" => "Spanish",
    "vi" => "Vietnamese",
    "he" => "Hebrew",
    "hu" => "Hungarian",
    "kk" => "Kazakh",
    "am" => "Amharic",
    "cu" => "Old_Church_Slavonic"
);


text:
using MacroTools: postwalk, splitdef, combinedef, @capture

macro model(expr)

    # Check it's a function and deconstruct.
    foo = splitdef(expr)

    # Transform provided expression to use a GPC with each GP.
    foo[:body] = Expr(:block,
        :(gpc = GPC()),
        postwalk(x->@capture(x, GP(xs__)) ? :(GP($(xs...), gpc)) : x, foo[:body]),
    )

    # Recombine into function expression.
    return esc(combinedef(foo))
end


text:
"""
    Both

Type whose singleton instance `both` represents both `false` and `true`.
"""
struct Both end

"""
    both

Singleton value of type `Both` representing both `false` and `true`.
"""
const both = Both()

Base.copy(::Both) = both
Base.show(io::IO, ::Both) = print(io, "both")

"""
    BooleanAbs

Belnap logic abstraction for Boolean values.
"""
const BooleanAbs = Union{Missing,Bool,Both}

BooleanAbs(x::BooleanAbs) = x

lub(a::Both, b::Both) = both
lub(a::Both, b::Bool) = both
lub(a::Bool, b::Bool) = a === b ? a : both
lub(a::Both, b::Missing) = both
lub(a::Bool, b::Missing) = a
lub(a::BooleanAbs, b::BooleanAbs) = lub(b, a)

glb(a::Both, b::Both) = both
glb(a::Both, b::Bool) = b
glb(a::Bool, b::Bool) = a === b ? a : missing
glb(a::Both, b::Missing) = missing
glb(a::Bool, b::Missing) = missing
glb(a::BooleanAbs, b::BooleanAbs) = glb(b, a)

widen(a::BooleanAbs, b::BooleanAbs) = lub(a, b)

equiv(a::Both, b::Bool) = true
equiv(a::Bool, b::Both) = true

Base.:(!)(::Both) = both

Base.:(&)(a::Both, b::Both) = both
Base.:(&)(a::Both, b::Missing) = false
Base.:(&)(a::Both, b::Bool) = b ? both : false
Base.:(&)(a::BooleanAbs, b::Both) = b & a

Base.:(|)(a::Both, b::Both) = both
Base.:(|)(a::Both, b::Missing) = true
Base.:(|)(a::Both, b::Bool) = b ? true : both
Base.:(|)(a::BooleanAbs, b::Both) = b | a

# Teach GenericState the way of true contradictions
function set_fluent!(state::GenericState, val::Both, term::Const)
    push!(state.facts, term, negate(term))
    return val
end

function set_fluent!(state::GenericState, val::Both, term::Compound)
    push!(state.facts, term, negate(term))
    return val
end


text:
using SIAN

println("Setting up the problem")

ode = @ODEmodel(
  xA'(t) = -k1 * xA(t),
  xB'(t) = k1 * xA(t) - k2 * xB(t),
  xC'(t) = k2 * xB(t),
  eA'(t) = 0,
  eC'(t) = 0,
  y1(t) = xC(t),
  y2(t) = eA(t) * xA(t) + eB * xB(t) + eC(t) * xC(t),
  y3(t) = eA(t),
  y4(t) = eC(t)
)

res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=0, nthrds=1)

println(res)

# The following computation shows that one can identify more from two experiments
# SIAN. generate_replica(ode, 2) generates a system consiting of two copies of sigma
# with the same parameters in ODEs but different inputs and initial conditions

r = 2

res = identifiability_ode(SIAN.generate_replica(ode, r), get_parameters(ode; initial_conditions=false); p=0.99, p_mod=0, nthrds=1)

println(res)

max_stars_repo_name: iliailmer/SIAN-Julia

max_issues_repo_name: iliailmer/SIAN-Julia

max_forks_repo_name: iliailmer/SIAN-Julia


text:
include("../src/TreeLas.jl")

using Test
import GraphIdx.Grid: num_nodes, num_edges, GridGraph
import GraphIdx: neighbors_lambda
import .TreeLas.Dual

@testset "GapLas   : 2x3                 " begin
    g = GridGraph(2, 3)
    m = 2*2 + 3
    n = 6
    @test num_nodes(g) == n
    @test num_edges(g) == m

    x = zeros(Float64, 2, 3)
    gam = zeros(Float64, m)
    alpha = zeros(Float64, m)
    gam .= -1
    @test all(gam .< 0)
    Dual.gap_vec!(gam, x, alpha, g)
    @test all(gam .≈ 0)

    x[1] = 13
    gam .= -1
    @test all(gam .< 0)
    Dual.gap_vec!(gam, x, alpha, g)
    @test gam[1] ≈ 13
    @test gam[4] ≈ 13
    for e = 1:m
        if e ∉ [1, 4]
            @test gam[e] ≈ 0
        end
    end

    x .= 0
    x[1, 1] = 13
    x[1, 2] = 13
    gam .= -1
    @test all(gam .< 0)
    Dual.gap_vec!(gam, x, alpha, g)
    @test gam[1] ≈ +13
    @test gam[2] ≈ +13
    @test gam[3] ≈ 0
    @test gam[4] ≈ 0
    @test gam[5] ≈ 13
    @test gam[6] ≈ 0
    @test gam[7] ≈ 0
end

@testset "IncidenceIndex(GridGraph(2, 3)  " begin
    g = GridGraph(2, 3)
    idx, lam = neighbors_lambda(g)
    @test all(lam .≈ 1)
    @test Set(idx[1]) == Set([(2, 1), (3, 4)])
end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
using RoME, Distributions
# using RoMEPlotting

import IncrementalInference: getSample

mutable struct NorthSouthPartial{T} <: FunctorSingleton
  Z::T
  partial::Tuple{Int}
  NorthSouthPartial() = new()
  NorthSouthPartial(Z::D) where {D <: Distribution} = new{D}(Z, (2,))
end
getSample(ns::NorthSouthPartial, N=1) = (reshape(rand(ns.Z, N),1,N),)

# Start with an empty graph
fg = initfg(sessionname="MULTIMODAL_2D_TUTORIAL")

# Add landmarks with Bearing range measurements
addVariable!(fg, :l1, Point2, labels=["LANDMARK"])
addVariable!(fg, :l2, Point2, labels=["LANDMARK"])

addFactor!(fg, [:l1], Prior(MvNormal([10.0;0.0], Matrix(Diagonal([1.0;1.0].^2)))) )
addFactor!(fg, [:l2], Prior(MvNormal([30.0;0.0], Matrix(Diagonal([1.0;1.0].^2)))) )

setVal!(getVert(fg, :l2), zeros(2,100))

addVariable!(fg, :x0, Pose2)
# addFactor!(fg, [:x0], Prior(MvNormal([0.0;0.0;0], Matrix(Diagonal([1.0;1.0;0.01].^2)))) )

p2br = Pose2Point2BearingRange(Normal(0,0.1),Normal(20.0,1.0))
addFactor!(fg, [:x0; :l1; :l2], p2br, multihypo=[1.0; 0.5; 0.5])

addFactor!(fg, [:x0;], NorthSouthPartial(Normal(0,1.0)))

# writeGraphPdf(fg)

ensureAllInitialized!(fg)

tree = wipeBuildNewTree!(fg)

[inferOverTreeR!(fg, tree, N=100) for i in 1:4]

# L1 = getVal(fg, :l1)
# L2 = getVal(fg, :l2)
X0 = getVertKDE(fg, :x0)

# plot(X0, dims=[1;2])

X0pts = getPoints(X0)

@test sum([0 < sum(-20 .< X0pts[1,:] .< 0);
           0 < sum(  0 .< X0pts[1,:] .< 20);
           0 < sum( 20 .< X0pts[1,:] .< 40);
           0 < sum( 40 .< X0pts[1,:] .< 60)]) > 2

# # Drive around in a hexagon
# for i in 0:5
#   psym = Symbol("x$i")
#   nsym = Symbol("x$(i+1)")
#   addVariable!(fg, nsym, Pose2, labels=["POSE"])
#   addFactor!(fg, [psym;nsym], Pose2Pose2(MvNormal([10.0;0;pi/3], Matrix(Diagonal([0.1;0.1;0.1].^2)))))
# end


text:
function moi_randsdp(optimizer, seed, n, m; verbose = false, test = false, atol = 1e-2, scalar = false)

    MOI.empty!(optimizer)
    if test
        @test MOI.is_empty(optimizer)
    end

    A, b, C = randsdp_data(seed, m, n)

    nvars = sympackedlen(n)

    X = MOI.add_variables(optimizer, nvars)

    Xsq = Matrix{MOI.VariableIndex}(undef,n,n)
    ivech!(Xsq, X)
    Xsq = Matrix(Symmetric(Xsq,:U))

    for k in 1:m
        ctr_k = vec([MOI.ScalarAffineTerm(A[k][i,j], Xsq[i,j]) for j in 1:n, i in 1:n])
        if scalar
            MOI.add_constraint(optimizer, MOI.ScalarAffineFunction(ctr_k, 0.0), MOI.EqualTo(b[k]))
        else
            MOI.add_constraint(optimizer,
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1], ctr_k), [-b[k]]), MOI.Zeros(1))
        end
    end

    vov = MOI.VectorOfVariables(X)
    cX = MOI.add_constraint(optimizer, vov, MOI.PositiveSemidefiniteConeTriangle(n))

    objf_t = vec([MOI.ScalarAffineTerm(C[i,j], Xsq[i,j]) for j in 1:n, i in 1:n])
    MOI.set(optimizer, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(objf_t, 0.0))

    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    MOI.optimize!(optimizer)

    objval = MOI.get(optimizer, MOI.ObjectiveValue())

    stime = -1.0
    try
        stime = MOI.get(optimizer, MOI.SolveTime())
    catch
        println("could not query time")
    end

    Xsq_s = MOI.get.(optimizer, MOI.VariablePrimal(), Xsq)

    minus_rank = length([eig for eig in eigen(Xsq_s).values if eig < -1e-4])
    if test
        @test minus_rank == 0
    end
    # rank = length([eig for eig in eigen(XX).values if eig > 1e-10])
    # @show rank
    if test
        @test tr(C * Xsq_s) - objval < atol
        for i in 1:m
            @test abs(tr(A[i] * Xsq_s)-b[i])/(1+abs(b[i])) < atol
        end
    end
    verbose && randsdp_eval(A,b,C,n,m,Xsq_s)

    rank = -1
    status = 0
    if MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
        status = 1
    end
    return (objval, stime, rank, status)
end


text:
function calc_x̂_minus_x(rr::RadauIntegrator{T_object, N, n_rule_max}, table::RadauTable{n_stage}, x0::Vector{Float64}) where {n_stage, n_rule_max, N, T_object}
    # Implements Eq. 8.17

    h = rr.step.h
    x̂_minus_x = zeros(Float64, length(x0))
    LinearAlgebra.BLAS.axpy!(table.b̂_0 * h, rr.cv.xx_0, x̂_minus_x)  # x̂_minus_x = (b̂_0 * h) * ẋ_0
    for k = 1:n_stage
        b̂_minus_b = table.b̂[k] - table.A[n_stage, k]
        LinearAlgebra.BLAS.axpy!(b̂_minus_b * h, rr.ct.F_X_stage[k], x̂_minus_x)  # x̂_minus_x += b̂_minus_b * F_X_stage[k]
    end
    return x̂_minus_x
end

function calc_x_err_norm(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}, x0::Vector{Float64},
    x_err::Vector{Float64}) where {n_stage, n_rule_max, N, TO}
    # Implements Eq 8.21 in Solving Ordinary Differential Equations II (Stiff and Differential-Algebraic Problems)

    x_final = get_X_final(rr, table)
    term_sigma = 0.0
    for k = 1:N
        sc_k = rr.step.tol_a + max(abs(x_final[k]), abs(x0[k])) * rr.step.tol_r
        term_sigma += (x_err[k] / sc_k)^2
    end
    return sqrt(term_sigma / N)
end

function update_x_err_norm!(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}, x0::Vector{Float64}) where {n_stage, n_rule_max, N, TO}
    rr.step.x_err_norm = rr.step.x_err_normᵏ⁺¹
    x̂_minus_x = calc_x̂_minus_x(rr, table, x0)
    x_err = rr.ct.inv_C_stage[1] * x̂_minus_x  # Eq. 8.19
    x_err = real.(x_err)
    rr.step.x_err_normᵏ⁺¹ = calc_x_err_norm(rr, table, x0, x_err)  # see definition in middle of pg. 124
end

function calc_h_new_estimate_1(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}) where {n_stage, n_rule_max, N, TO}
    #  Implements first half of part of 8.25
    k_iter = rr.rule.k_iter
    two_k_max_iter = 2 * rr.rule.k_iter_max
    fac = 0.9 * (two_k_max_iter + 1) / (two_k_max_iter + k_iter)
    term_1 = (1 / rr.step.x_err_normᵏ⁺¹)^get_exponent(table)
    return fac * rr.step.h * term_1
end

function calc_and_update_h!(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}, x0::Vector{Float64}) where {n_stage, n_rule_max, N, TO}
    if is_exit_flag_success(rr)
        h_new = calc_h_new_estimate_1(rr, table)
    else
        h_new = get_h(rr) * 0.1
    end
    h_new = min(rr.step.h_max, 2 * get_h(rr), h_new)
    update_h!(rr, h_new)
end

function update_h!(rr::RadauIntegrator{TO, N, n_rule_max}, h_new::Float64) where {n_rule_max, N, TO}
    (0.0 < h_new < Inf) || error("unacceptable h: $h_new")
    rr.step.hᵏ⁻¹ = rr.step.h
    rr.step.h    = h_new
    rr.step.h⁻¹  = 1 / h_new
end

function update_rule!(rr::RadauIntegrator{T_object, N, n_rule_max}) where {n_rule_max, N, T_object}
    # Implements stategy on page 14 of "Stiff differential equations solved by Radau methods (Hairer)"

    cooldown_reset = 10
    s = rr.rule.s
    if is_exit_flag_success(rr)
        rr.rule.n_increase_cooldown -= 1
        cool = rr.rule.n_increase_cooldown
        Ψ_k = rr.rule.Ψ_k
        if cool < 1  # has the cooldown expired
            if Ψ_k < 0.1  # convergence isn't terrible
                rr.rule.s = min(rr.rule.s + 1, rr.rule.max_rule)  # up rules never going above max_rule
            end
        end
    else
        rr.rule.n_increase_cooldown = cooldown_reset
        rr.rule.s = max(rr.rule.s - 1, 1)  # decrease rules number, but never below 1
    end
    (1 <= rr.rule.s ) || error("something is wrong")
    return nothing
end

# function predictive_correction(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}) where {n_stage, n_rule_max, N, TO}
#     # Implements second half of part of 8.25
#
#     # an x_err_norm can be 0.0 if the equation is integratable EXACTLY
#     if 0.0 == rr.step.x_err_normᵏ⁺¹  # things went well let's keep it that way
#         return 1.0
#     elseif rr.step.x_err_norm == 0.0  # things got much much worse let's drastically reduce step size in anticipation
#         return 0.1
#     else
#         ratio_h = rr.step.h / rr.step.hᵏ⁻¹
#         ratio_x_err_norm = rr.step.x_err_norm / rr.step.x_err_normᵏ⁺¹
#         pred_corr = ratio_h * (ratio_x_err_norm)^get_exponent(table)
#         return clamp(pred_corr, 0.1, 1.0)  # otherwise prediction will be WAY too pessimistic
#     end
# end

# function calc_h_new(rr::RadauIntegrator{TO, N, n_rule_max}, table::RadauTable{n_stage}, x0::Vector{Float64},
#     is_converge::Bool) where {n_stage, n_rule_max, N, TO}
#
#     if is_converge
#
#         h_est¹ = calc_h_new_estimate_1(rr, table)
#         h_new = h_est¹
#
#         # if rr.step.is_has_prev_step
#         #     h_new = h_est¹ * min(1.0, predictive_correction(rr, table))  # conservative == good
#         # else
#         #     rr.step.is_has_prev_step = true
#         #     h_new = h_est¹
#         # end
#
#         # return min(h_new, 2 * rr.step.h)  # otherwise it will try to use numbers O(10.0-50.0) next time
#     else  # failure
#         rr.step.is_has_prev_step = false
#         # return rr.step.h * 0.1  # drastically decrease time-step and return
#         h_new = rr.step.h * 0.1
#     end
#     # don't want to exceed maximum timestep or to increase timestep too quickly
#     return min(rr.step.h_max, 2 * rr.step.h, h_new)
# end


text:
"""
Code get from github:
https://github.com/JuliaGPU/CUDA.jl/issues/177#issuecomment-838311843
"""

import Zygote

import Base: _RepeatInnerOuter

@kernel function repeat_inner_kernel!(a::AbstractArray{<:Any, N}, inner::NTuple{N}, out) where {N}
    inds = @index(Global, NTuple)
    inds_a = ntuple(i -> (inds[i] - 1) ÷ inner[i] + 1, N)

    @inbounds out[inds...] = a[inds_a...]
end

function repeat_inner(a::TV, inner) where {TV<:AbstractArray}
    out = TV(undef, inner .* size(a))

    kernel! = if out isa CuArray
        repeat_inner_kernel!(CUDADevice(), 512)
    else
        repeat_inner_kernel!(CPU(), Threads.nthreads())
    end

    ev = kernel!(a, inner, out, ndrange=size(out))
    wait(ev)
    return out
end
# Non-cached coherent loads
@kernel function repeat_outer_kernel!(a::AbstractArray{<:Any, N}, sa::NTuple{N}, outer::NTuple{N}, out) where {N}
    inds = @index(Global, NTuple)
    inds_a = ntuple(i -> (inds[i] - 1) % sa[i] + 1, N)

    @inbounds out[inds...] = a[inds_a...]
end

function repeat_outer(a::TV, outer) where {TV<:AbstractArray}
    out = TV(undef, outer .* size(a))

    kernel! = if out isa CuArray
        repeat_outer_kernel!(CUDADevice(), 512)
    else
        repeat_outer_kernel!(CPU(), Threads.nthreads())
    end

    ev = kernel!(a, size(a), outer, out, ndrange=size(out))
    wait(ev)
    return out
end

# Overload methods used by `Base.repeat`.
# No need to implement `repeat_inner_outer` since this is implemented in `Base` as
# `repeat_outer(repeat_inner(arr, inner), outer)`.
function _RepeatInnerOuter.repeat_inner(a::CuArray{<:Any, N}, dims::NTuple{N}) where {N}
    return repeat_inner(a, dims)
end

function _RepeatInnerOuter.repeat_outer(a::CuArray{<:Any, N}, dims::NTuple{N}) where {N}
    return repeat_outer(a, dims)
end

function _RepeatInnerOuter.repeat_outer(a::CuVector, dims::Tuple{Any})
    return repeat_outer(a, dims)
end

function _RepeatInnerOuter.repeat_outer(a::CuMatrix, dims::NTuple{2, Any})
    return repeat_outer(a, dims)
end

### Adjoint implementation for `repeat`
@kernel function repeat_adjoint_gpu_kernel!(
    Δ::AbstractArray{T},
    inner::NTuple,
    outer::NTuple,
    out::AbstractArray{T},
    outsize::NTuple{N}
) where {N, T<:Real}
    dest_inds = @index(Global, NTuple)
    src_inds = ntuple(i -> mod1((dest_inds[i] - 1) ÷ inner[i] + 1, outsize[i]), N)
    CUDA.@atomic out[src_inds...] += Δ[dest_inds...]
end;

# FIXME: not threadsafe atm!!! And therefore not used anywhere (we only overload) the
# adjoint for `CuArray`. But if we have  something like `@atomic_addindex!`
# from https://github.com/JuliaLang/julia/pull/37683, we'd be golden.
@kernel function repeat_adjoint_cpu_kernel!(
    Δ::AbstractArray,
    inner::NTuple,
    outer::NTuple,
    out::AbstractArray,
    outsize::NTuple{N}
) where {N}
    dest_inds = @index(Global, NTuple)
    src_inds = ntuple(i -> mod1((dest_inds[i] - 1) ÷ inner[i] + 1, outsize[i]), N)
    # FIXME: make threadsafe
    out[src_inds...] += Δ[dest_inds...]
end;

function repeat_adjoint(
    x,
    Δ::AbstractArray{<:Any, N},
    inner::NTuple{N},
    outer::NTuple{N}
) where {N}
    out = zero(x)

    kernel! = if out isa CuArray
        repeat_adjoint_gpu_kernel!(CUDADevice(), 512)
    else
        repeat_adjoint_cpu_kernel!(CPU(), Threads.nthreads())
    end

    ev = kernel!(Δ, inner, outer, out, size(out), ndrange=size(Δ))
    wait(ev)

    return out
end;

Zygote.@adjoint function _RepeatInnerOuter.repeat_inner_outer(xs::AbstractArray, inner::Nothing, outer::Nothing)
    return xs, Δ -> (Δ, )
end
Zygote.@adjoint function _RepeatInnerOuter.repeat_inner_outer(
    xs::AbstractArray,
    inner::Nothing,
    outer::NTuple{N}
) where {N}
    inner_new = ntuple(_ -> 1, N)
    return (
        _RepeatInnerOuter.repeat_outer(xs, outer),
        Δ -> (repeat_adjoint(xs, Δ, inner_new, outer), )
    )
end
Zygote.@adjoint function _RepeatInnerOuter.repeat_inner_outer(
    xs::AbstractArray,
    inner::NTuple{N},
    outer::Nothing
) where {N}
    outer_new = ntuple(_ -> 1, N)
    return (
        _RepeatInnerOuter.repeat_inner(xs, inner),
        Δ -> (repeat_adjoint(xs, Δ, inner, outer_new), )
    )
end
Zygote.@adjoint function _RepeatInnerOuter.repeat_inner_outer(
    xs::AbstractArray,
    inner::NTuple{N},
    outer::NTuple{N}
) where {N}
    return (
        _RepeatInnerOuter.repeat_outer(_RepeatInnerOuter.repeat_inner(xs, inner), outer),
        Δ -> (repeat_adjoint(xs, Δ, inner, outer), )
    )
end

# We need to stop Zygote from using the rule implemented in
# https://github.com/FluxML/Zygote.jl/blob/d5be4d5ca80e79278d714eaac15ca71904a262e3/src/lib/array.jl#L149-L163
# We stop this adjoint from triggering, and then call the underlying `repeat_inner_outer`.
# Unfortunately we need to replicate the body of `_RepeatInnerOuter.repeat` since we can't do `Zygote.pullback`
# for kwargs.
Zygote.@adjoint function Base.repeat(arr::CuArray; inner = nothing, outer = nothing)
    _RepeatInnerOuter.check(arr, inner, outer)
    arr, inner, outer = _RepeatInnerOuter.resolve(arr, inner, outer)
    return Zygote.pullback(_RepeatInnerOuter.repeat_inner_outer, arr, inner, outer)
end


text:
# compute an ensemble of ϕ(u) scalar observable values
function run_obs_ensemble(u0data, Tmax, ϕ)

    n = length(u0data);
    samples = zeros(n);

    Threads.@threads for i in 1:n
        soln = solve_ode(u0data[i], Tmax);
        samples[i] = ϕ(soln)
    end

    return samples
end

# compute an ensemble of the paths at the indicated values of
function run_path_ensemble(u0data, Tmax, t_vals)

    n = length(u0data);
    u1_samples = zeros(n, length(t_vals));
    u2_samples = zeros(n, length(t_vals));

    Threads.@threads for i in 1:n
        soln = solve_ode(u0data[i], Tmax);
        soln_at_t = soln(t_vals);
        u1_samples[i,:] .= [u[1] for u in soln_at_t];
        u2_samples[i,:] .= [u[2] for u in soln_at_t];
    end

    return u1_samples,u2_samples
end


text:
using Definitions
using Results
using Parameters
using Household
using Calculations
module Driver

include( "driver.jl" )

end


text:
using DSGE, DataFrames, JLD
using Base.Test

path = dirname(@__FILE__)

# Set up arguments
m = AnSchorfheide(testing = true)
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))

df, system, z0, P0 = jldopen("$path/../reference/forecast_args.jld", "r") do file
    read(file, "df"), read(file, "system"), read(file, "z0"), read(file, "P0")
end

# Read expected output
exp_kal = jldopen("$path/../reference/filter_out.jld", "r") do file
    read(file, "exp_kal")
end

# Without providing z0 and P0
@testset "Check Kalman filter outputs without initializing state/state-covariance" begin
    kal = DSGE.filter(m, df, system)
    for out in fieldnames(kal)
        expect = exp_kal[out]
        actual = kal[out]

        if ndims(expect) == 0
            @test expect ≈ actual
        else
            @test @test_matrix_approx_eq(expect, actual)
        end
    end
end

# Providing z0 and P0
@testset "Check Kalman filter outputs initializing state/state-covariance" begin
    kal = DSGE.filter(m, df, system, z0, P0)
    for out in fieldnames(kal)
        expect = exp_kal[out]
        actual = kal[out]

        if ndims(expect) == 0
            @test expect ≈ actual
        else
            @test @test_matrix_approx_eq(expect, actual)
        end
    end
end

nothing


text:
Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:AbstractAxis,Y<:Axis} = promote_rule(Y, X)
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:Axis,Y<:AbstractAxis}
    return Axis{
        promote_type(keytype(X), keytype(Y)),
        promote_type(valtype(X), valtype(Y)),
        promote_type(_keys_type(X), _keys_type(Y)),
        promote_type(parent_type(X), parent_type(Y))
    }
end
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:Axis,Y<:Axis}
    return Axis{
        promote_type(keytype(X), keytype(Y)),
        promote_type(valtype(X), valtype(Y)),
        promote_type(_keys_type(X), _keys_type(Y)),
        promote_type(parent_type(X), parent_type(Y))
    }
end

Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:SimpleAxis,Y<:Axis} = promote_rule(Y, X)
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:Axis,Y<:SimpleAxis}
    return Axis{
        promote_type(keytype(X), keytype(Y)),
        promote_type(valtype(X), valtype(Y)),
        _keys_type(X),
        promote_type(parent_type(X), parent_type(Y))
    }
end

Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:AbstractAxis,Y<:SimpleAxis} = promote_rule(Y, X)
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:SimpleAxis,Y<:AbstractAxis}
    return SimpleAxis{
        promote_type(valtype(X),valtype(Y)),
        promote_type(parent_type(X), parent_type(Y))
    }
end
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:SimpleAxis,Y<:SimpleAxis}
    return SimpleAxis{
        promote_type(valtype(X),valtype(Y)),
        promote_type(parent_type(X),parent_type(Y))
    }
end

function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:AbstractUnitRange,Y<:SimpleAxis}
    return promote_rule(Y, X)
end
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:SimpleAxis,Y<:AbstractUnitRange}
    return SimpleAxis{
        promote_type(valtype(X),valtype(Y)),
        promote_type(parent_type(X), parent_type(Y))
    }
end

Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:AbstractUnitRange,Y<:Axis} = promote_rule(Y, X)
function Base.promote_rule(::Type{X}, ::Type{Y}) where {X<:Axis,Y<:AbstractUnitRange}
    return Axis{
        keytype(X),
        promote_type(valtype(X),valtype(Y)),
        _keys_type(X),
        promote_type(parent_type(X), parent_type(Y))
    }
end

# unfortunately, we need these to avoid ambiguities
function Base.promote_rule(::Type{UnitRange{T1}}, ::Type{Y}) where {T1,Y<:SimpleAxis}
    return promote_rule(Y,UnitRange{T1})
end
function Base.promote_rule(::Type{UnitRange{T1}}, ::Type{Y}) where {T1,Y<:Axis}
    return promote_rule(Y,UnitRange{T1})
end

# TODO delete this?
function StaticRanges.same_type(::Type{X}, ::Type{Y}) where {X<:AbstractAxis,Y<:AbstractAxis}
    return (X.name === Y.name)  & # TODO there should be a better way of doing this
       same_type(_keys_type(X), _keys_type(Y)) &
       same_type(parent_type(X), indices_type(Y))
end


text:
@inline isexpr(@nospecialize(stmt), head::Symbol) = isa(stmt, Expr) && stmt.head === head
@eval Core.UpsilonNode() = $(Expr(:new, Core.UpsilonNode))
Core.PhiNode() = Core.PhiNode(Any[], Any[])

struct Argument
    n::Int
end

struct GotoIfNot
    cond
    dest::Int
    GotoIfNot(@nospecialize(cond), dest::Int) = new(cond, dest)
end

struct ReturnNode
    val
    ReturnNode(@nospecialize(val)) = new(val)
    # unassigned val indicates unreachable
    ReturnNode() = new()
end

"""
Like UnitRange{Int}, but can handle the `last` field, being temporarily
< first (this can happen during compacting)
"""
struct StmtRange <: AbstractUnitRange{Int}
    first::Int
    last::Int
end
first(r::StmtRange) = r.first
last(r::StmtRange) = r.last
start(r::StmtRange) = 0
done(r::StmtRange, state) = r.last - r.first < state
next(r::StmtRange, state) = (r.first + state, state + 1)

StmtRange(range::UnitRange{Int}) = StmtRange(first(range), last(range))

struct BasicBlock
    stmts::StmtRange
    #error_handler::Bool
    preds::Vector{Int}
    succs::Vector{Int}
end
function BasicBlock(stmts::StmtRange)
    BasicBlock(stmts, Int[], Int[])
end
function BasicBlock(old_bb, stmts)
    BasicBlock(stmts, #= old_bb.error_handler, =# old_bb.preds, old_bb.succs)
end
copy(bb::BasicBlock) = BasicBlock(bb.stmts, #= bb.error_handler, =# copy(bb.preds), copy(bb.succs))

struct CFG
    blocks::Vector{BasicBlock}
    index::Vector{Int}
end

function block_for_inst(index, inst)
    searchsortedfirst(index, inst, lt=(<=))
end
block_for_inst(cfg::CFG, inst) = block_for_inst(cfg.index, inst)

function compute_basic_blocks(stmts::Vector{Any})
    jump_dests = IdSet{Int}(1)
    # First go through and compute jump destinations
    for (idx, stmt) in pairs(stmts)
        # Terminators
        if isa(stmt, Union{GotoIfNot, GotoNode, ReturnNode})
            if isa(stmt, GotoIfNot)
                push!(jump_dests, idx+1)
                push!(jump_dests, stmt.dest)
            else
                # This is a fake dest to force the next stmt to start a bb
                idx < length(stmts) && push!(jump_dests, idx+1)
                if isa(stmt, GotoNode)
                    push!(jump_dests, stmt.label)
                end
            end
        elseif isa(stmt, Expr) && stmt.head === :leave
            # :leave terminates a BB
            push!(jump_dests, idx+1)
        elseif isa(stmt, Expr) && stmt.head == :enter
            # :enter starts/ends a BB
            push!(jump_dests, idx)
            push!(jump_dests, idx+1)
        end
    end
    bb_starts = sort(collect(jump_dests))
    for i = length(stmts):-1:1
        if stmts[i] != nothing
            push!(bb_starts, i+1)
            break
        end
    end
    # Compute ranges
    basic_block_index = Int[]
    blocks = BasicBlock[]
    sizehint!(blocks, length(bb_starts)-1)
    for (first, last) in Iterators.zip(bb_starts, Iterators.drop(bb_starts, 1))
        push!(basic_block_index, first)
        push!(blocks, BasicBlock(StmtRange(first, last-1)))
    end
    popfirst!(basic_block_index)
    # Compute successors/predecessors
    for (num, b) in pairs(blocks)
        terminator = stmts[last(b.stmts)]
        # Conditional Branch
        if isa(terminator, GotoIfNot)
            block′ = block_for_inst(basic_block_index, terminator.dest)
            push!(blocks[block′].preds, num)
            push!(b.succs, block′)
        end
        if isa(terminator, GotoNode)
            block′ = block_for_inst(basic_block_index, terminator.label)
            push!(blocks[block′].preds, num)
            push!(b.succs, block′)
        elseif !isa(terminator, ReturnNode)
            if isa(terminator, Expr) && terminator.head == :enter
                # :enter gets a virtual edge to the exception handler and
                # the exception handler gets a virtual edge from outside
                # the function.
                # See the devdocs on exception handling in SSA form (or
                # bug Keno to write them, if you're reading this and they
                # don't exist)
                block′ = block_for_inst(basic_block_index, terminator.args[1])
                push!(blocks[block′].preds, num)
                push!(blocks[block′].preds, 0)
                push!(b.succs, block′)
            end
            if num + 1 <= length(blocks)
                push!(blocks[num+1].preds, num)
                push!(b.succs, num+1)
            end
        end
    end
    CFG(blocks, basic_block_index)
end

function first_insert_for_bb(code, cfg, block)
    for idx in cfg.blocks[block].stmts
        stmt = code[idx]
        if !isa(stmt, LabelNode) && !isa(stmt, PhiNode)
            return idx
        end
    end
end

const NewNode = Tuple{Int, #= reverse affinity =#Bool, Any, Any, #=LineNumber=#Int}

struct IRCode
    stmts::Vector{Any}
    types::Vector{Any}
    lines::Vector{Int}
    argtypes::Vector{Any}
    cfg::CFG
    new_nodes::Vector{NewNode}
    mod::Module
    meta::Vector{Any}

    function IRCode(stmts::Vector{Any}, types::Vector{Any}, lines::Vector{Int}, cfg::CFG, argtypes::Vector{Any}, mod::Module, meta::Vector{Any})
        return new(stmts, types, lines, argtypes, cfg, NewNode[], mod, meta)
    end
    function IRCode(ir::IRCode, stmts::Vector{Any}, types::Vector{Any}, lines::Vector{Int}, cfg::CFG, new_nodes::Vector{NewNode})
        return new(stmts, types, lines, ir.argtypes, cfg, new_nodes, ir.mod, ir.meta)
    end
end

function getindex(x::IRCode, s::SSAValue)
    if s.id <= length(x.stmts)
        return x.stmts[s.id]
    else
        return x.new_nodes[s.id - length(x.stmts)][4]
    end
end

function setindex!(x::IRCode, repl, s::SSAValue)
    @assert s.id <= length(x.stmts)
    x.stmts[s.id] = repl
    nothing
end

struct OldSSAValue
    id::Int
end

struct NewSSAValue
    id::Int
end

mutable struct UseRefIterator
    stmt::Any
end
getindex(it::UseRefIterator) = it.stmt

struct UseRef
    urs::UseRefIterator
    use::Int
end

struct OOBToken
end

struct UndefToken
end

function getindex(x::UseRef)
    stmt = x.urs.stmt
    if isa(stmt, Expr) && stmt.head === :(=)
        rhs = stmt.args[2]
        if isa(rhs, Expr) && is_relevant_expr(rhs)
            x.use > length(rhs.args) && return OOBToken()
            return rhs.args[x.use]
        end
        x.use == 1 || return OOBToken()
        return rhs
    elseif isa(stmt, Expr) && is_relevant_expr(stmt)
        x.use > length(stmt.args) && return OOBToken()
        return stmt.args[x.use]
    elseif isa(stmt, GotoIfNot)
        x.use == 1 || return OOBToken()
        return stmt.cond
    elseif isa(stmt, ReturnNode) || isa(stmt, PiNode) || isa(stmt, UpsilonNode)
        isdefined(stmt, :val) || return OOBToken()
        x.use == 1 || return OOBToken()
        return stmt.val
    elseif isa(stmt, PhiNode) || isa(stmt, PhiCNode)
        x.use > length(stmt.values) && return OOBToken()
        isassigned(stmt.values, x.use) || return UndefToken()
        return stmt.values[x.use]
    else
        return OOBToken()
    end
end

function is_relevant_expr(e::Expr)
    return e.head in (:call, :invoke, :new, :(=), :(&),
                      :gc_preserve_begin, :gc_preserve_end,
                      :foreigncall, :isdefined, :copyast,
                      :undefcheck, :throw_undef_if_not)
end

function setindex!(x::UseRef, @nospecialize(v))
    stmt = x.urs.stmt
    if isa(stmt, Expr) && stmt.head === :(=)
        rhs = stmt.args[2]
        if isa(rhs, Expr) && is_relevant_expr(rhs)
            x.use > length(rhs.args) && throw(BoundsError())
            rhs.args[x.use] = v
        else
            x.use == 1 || throw(BoundsError())
            stmt.args[2] = v
        end
    elseif isa(stmt, Expr) && is_relevant_expr(stmt)
        x.use > length(stmt.args) && throw(BoundsError())
        stmt.args[x.use] = v
    elseif isa(stmt, GotoIfNot)
        x.use == 1 || throw(BoundsError())
        x.urs.stmt = GotoIfNot(v, stmt.dest)
    elseif isa(stmt, ReturnNode) || isa(stmt, UpsilonNode)
        x.use == 1 || throw(BoundsError())
        x.urs.stmt = typeof(stmt)(v)
    elseif isa(stmt, PiNode)
        x.use == 1 || throw(BoundsError())
        x.urs.stmt = typeof(stmt)(v, stmt.typ)
    elseif isa(stmt, PhiNode) || isa(stmt, PhiCNode)
        x.use > length(stmt.values) && throw(BoundsError())
        isassigned(stmt.values, x.use) || throw(BoundsError())
        stmt.values[x.use] = v
    else
        throw(BoundsError())
    end
    return x
end

function userefs(@nospecialize(x))
    if (isa(x, Expr) && is_relevant_expr(x)) ||
        isa(x, Union{GotoIfNot, ReturnNode, PiNode, PhiNode, PhiCNode, UpsilonNode})
        UseRefIterator(x)
    else
        ()
    end
end

start(it::UseRefIterator) = 1
function next(it::UseRefIterator, use)
    x = UseRef(it, use)
    v = x[]
    v === UndefToken() && return next(it, use + 1)
    x, use + 1
end
function done(it::UseRefIterator, use)
    x, _ = next(it, use)
    v = x[]
    v === OOBToken() && return true
    false
end

function scan_ssa_use!(used, @nospecialize(stmt))
    if isa(stmt, SSAValue)
        push!(used, stmt.id)
    end
    for useref in userefs(stmt)
        val = useref[]
        if isa(val, SSAValue)
            push!(used, val.id)
        end
    end
end

function ssamap(f, @nospecialize(stmt))
    urs = userefs(stmt)
    urs === () && return stmt
    for op in urs
        val = op[]
        if isa(val, SSAValue)
            op[] = f(val)
        end
    end
    urs[]
end

function foreachssa(f, @nospecialize(stmt))
    for op in userefs(stmt)
        val = op[]
        if isa(val, SSAValue)
            f(val)
        end
    end
end

function insert_node!(ir::IRCode, pos::Int, @nospecialize(typ), @nospecialize(val), reverse_affinity::Bool=false)
    line = ir.lines[pos]
    push!(ir.new_nodes, (pos, reverse_affinity, typ, val, line))
    return SSAValue(length(ir.stmts) + length(ir.new_nodes))
end

# For bootstrapping
function my_sortperm(v)
    p = Vector{Int}(undef, length(v))
    for i = 1:length(v)
        p[i] = i
    end
    sort!(p, Sort.DEFAULT_UNSTABLE, Order.Perm(Sort.Forward,v))
    p
end

mutable struct IncrementalCompact
    ir::IRCode
    result::Vector{Any}
    result_types::Vector{Any}
    result_lines::Vector{Int}
    result_bbs::Vector{BasicBlock}
    ssa_rename::Vector{Any}
    used_ssas::Vector{Int}
    late_fixup::Vector{Int}
    # This could be Stateful, but bootstrapping doesn't like that
    perm::Vector{Int}
    new_nodes_idx::Int
    idx::Int
    result_idx::Int
    active_result_bb::Int
    function IncrementalCompact(code::IRCode)
        # Sort by position with reverse affinity nodes before regular ones
        perm = my_sortperm(Int[(code.new_nodes[i][1]*2 + Int(!code.new_nodes[i][2])) for i in 1:length(code.new_nodes)])
        new_len = length(code.stmts) + length(code.new_nodes)
        result = Array{Any}(undef, new_len)
        result_types = Array{Any}(undef, new_len)
        result_lines = Array{Int}(undef, new_len)
        used_ssas = fill(0, new_len)
        ssa_rename = Any[SSAValue(i) for i = 1:new_len]
        late_fixup = Vector{Int}()
        return new(code, result, result_types, result_lines, code.cfg.blocks, ssa_rename, used_ssas, late_fixup, perm, 1, 1, 1, 1)
    end

    # For inlining
    function IncrementalCompact(parent::IncrementalCompact, code::IRCode, result_offset)
        perm = my_sortperm(Int[code.new_nodes[i][1] for i in 1:length(code.new_nodes)])
        new_len = length(code.stmts) + length(code.new_nodes)
        ssa_rename = Any[SSAValue(i) for i = 1:new_len]
        used_ssas = fill(0, new_len)
        late_fixup = Vector{Int}()
        return new(code, parent.result, parent.result_types, parent.result_lines, parent.result_bbs, ssa_rename, parent.used_ssas,
            late_fixup, perm, 1, 1, result_offset, parent.active_result_bb)
    end
end

struct TypesView
    ir::Union{IRCode, IncrementalCompact}
end
types(ir::Union{IRCode, IncrementalCompact}) = TypesView(ir)

function getindex(compact::IncrementalCompact, idx)
    if idx < compact.result_idx
        return compact.result[idx]
    else
        return compact.ir.stmts[idx]
    end
end

function count_added_node!(compact, v)
    if isa(v, SSAValue)
        compact.used_ssas[v.id] += 1
    else
        for ops in userefs(v)
            val = ops[]
            if isa(val, SSAValue)
                compact.used_ssas[val.id] += 1
            end
        end
    end
end

function insert_node_here!(compact::IncrementalCompact, @nospecialize(val), @nospecialize(typ), ltable_idx::Int)
    if compact.result_idx > length(compact.result)
        @assert compact.result_idx == length(compact.result) + 1
        resize!(compact, compact.result_idx)
    end
    compact.result[compact.result_idx] = val
    compact.result_types[compact.result_idx] = typ
    compact.result_lines[compact.result_idx] = ltable_idx
    count_added_node!(compact, val)
    ret = SSAValue(compact.result_idx)
    compact.result_idx += 1
    ret
end

function getindex(view::TypesView, v::OldSSAValue)
    return view.ir.ir.types[v.id]
end

function setindex!(compact::IncrementalCompact, v, idx)
    if idx < compact.result_idx
        (compact.result[idx] == v) && return
        # Kill count for current uses
        for ops in userefs(compact.result[idx])
            val = ops[]
            if isa(val, SSAValue)
                @assert compact.used_ssas[val.id] >= 1
                compact.used_ssas[val.id] -= 1
            end
        end
        compact.result[idx] = v
        # Add count for new use
        count_added_node!(compact, v)
    else
        compact.ir.stmts[idx] = v
    end
    return nothing
end

function getindex(view::TypesView, idx)
    isa(idx, SSAValue) && (idx = idx.id)
    if isa(view.ir, IncrementalCompact) && idx < view.ir.result_idx
        return view.ir.result_types[idx]
    else
        ir = isa(view.ir, IncrementalCompact) ? view.ir.ir : view.ir
        if idx <= length(ir.types)
            return ir.types[idx]
        else
            return ir.new_nodes[idx - length(ir.types)][3]
        end
    end
end

# maybe use expr_type?
function value_typ(ir::IRCode, value)
    isa(value, SSAValue) && return ir.types[value.id]
    isa(value, GlobalRef) && return abstract_eval_global(value.mod, value.name)
    isa(value, Argument) && return ir.argtypes[value.n]
    # TODO: isa QuoteNode, etc.
    return typeof(value)
end

function value_typ(ir::IncrementalCompact, value)
    isa(value, SSAValue) && return types(ir)[value.id]
    isa(value, GlobalRef) && return abstract_eval_global(value.mod, value.name)
    isa(value, Argument) && return ir.ir.argtypes[value.n]
    # TODO: isa QuoteNode, etc.
    return typeof(value)
end

start(compact::IncrementalCompact) = (compact.idx, 1)
function done(compact::IncrementalCompact, (idx, _a)::Tuple{Int, Int})
    return idx > length(compact.ir.stmts) && (compact.new_nodes_idx > length(compact.perm))
end

function process_phinode_values(old_values, late_fixup, processed_idx, result_idx, ssa_rename, used_ssas)
    values = Vector{Any}(undef, length(old_values))
    for i = 1:length(old_values)
        isassigned(old_values, i) || continue
        val = old_values[i]
        if isa(val, SSAValue)
            if val.id > processed_idx
                push!(late_fixup, result_idx)
                val = OldSSAValue(val.id)
            else
                val = renumber_ssa!(val, ssa_rename, true, used_ssas)
            end
        end
        values[i] = val
    end
    values
end

function process_node!(result::Vector{Any}, result_idx::Int, ssa_rename::Vector{Any},
        late_fixup::Vector{Int}, used_ssas::Vector{Int}, @nospecialize(stmt),
        idx::Int, processed_idx::Int)
    ssa_rename[idx] = SSAValue(result_idx)
    if stmt === nothing
        ssa_rename[idx] = stmt
    elseif isa(stmt, GotoNode) || isa(stmt, GlobalRef)
        result[result_idx] = stmt
        result_idx += 1
    elseif isa(stmt, Expr) || isa(stmt, PiNode) || isa(stmt, GotoIfNot) || isa(stmt, ReturnNode) || isa(stmt, UpsilonNode)
        result[result_idx] = renumber_ssa!(stmt, ssa_rename, true, used_ssas)
        result_idx += 1
    elseif isa(stmt, PhiNode)
        result[result_idx] = PhiNode(stmt.edges, process_phinode_values(stmt.values, late_fixup, processed_idx, result_idx, ssa_rename, used_ssas))
        result_idx += 1
    elseif isa(stmt, PhiCNode)
        result[result_idx] = PhiCNode(process_phinode_values(stmt.values, late_fixup, processed_idx, result_idx, ssa_rename, used_ssas))
        result_idx += 1
    elseif isa(stmt, SSAValue)
        # identity assign, replace uses of this ssa value with its result
        stmt = ssa_rename[stmt.id]
        ssa_rename[idx] = stmt
    else
        # Constant assign, replace uses of this ssa value with its result
        ssa_rename[idx] = stmt
    end
    return result_idx
end
function process_node!(compact::IncrementalCompact, result_idx::Int, @nospecialize(stmt), idx::Int, processed_idx::Int)
    return process_node!(compact.result, result_idx, compact.ssa_rename,
        compact.late_fixup, compact.used_ssas, stmt, idx, processed_idx)
end

function resize!(compact::IncrementalCompact, nnewnodes)
    old_length = length(compact.result)
    resize!(compact.result, nnewnodes)
    resize!(compact.result_types, nnewnodes)
    resize!(compact.result_lines, nnewnodes)
    resize!(compact.used_ssas, nnewnodes)
    compact.used_ssas[(old_length+1):nnewnodes] = 0
end

function finish_current_bb!(compact, old_result_idx=compact.result_idx)
    bb = compact.result_bbs[compact.active_result_bb]
    # If this was the last statement in the BB and we decided to skip it, insert a
    # dummy `nothing` node, to prevent changing the structure of the CFG
    if compact.result_idx == first(bb.stmts)
        length(compact.result) < old_result_idx && resize!(compact, old_result_idx)
        compact.result[old_result_idx] = nothing
        compact.result_types[old_result_idx] = Nothing
        compact.result_lines[old_result_idx] = 0
        compact.result_idx = old_result_idx + 1
    end
    compact.result_bbs[compact.active_result_bb] = BasicBlock(bb, StmtRange(first(bb.stmts), compact.result_idx-1))
    compact.active_result_bb += 1
    if compact.active_result_bb <= length(compact.result_bbs)
        new_bb = compact.result_bbs[compact.active_result_bb]
        compact.result_bbs[compact.active_result_bb] = BasicBlock(new_bb,
            StmtRange(compact.result_idx, last(new_bb.stmts)))
    end
end

function reverse_affinity_stmt_after(compact::IncrementalCompact, idx::Int)
    compact.new_nodes_idx > length(compact.perm) && return false
    entry = compact.ir.new_nodes[compact.perm[compact.new_nodes_idx]]
    entry[1] == idx + 1 && entry[2]
end

function next(compact::IncrementalCompact, (idx, active_bb)::Tuple{Int, Int})
    old_result_idx = compact.result_idx
    if length(compact.result) < old_result_idx
        resize!(compact, old_result_idx)
    end
    bb = compact.ir.cfg.blocks[active_bb]
    if compact.new_nodes_idx <= length(compact.perm) && compact.ir.new_nodes[compact.perm[compact.new_nodes_idx]][1] == idx
        new_idx = compact.perm[compact.new_nodes_idx]
        compact.new_nodes_idx += 1
        _, reverse_affinity, typ, new_node, new_line = compact.ir.new_nodes[new_idx]
        new_idx += length(compact.ir.stmts)
        compact.result_types[old_result_idx] = typ
        compact.result_lines[old_result_idx] = new_line
        result_idx = process_node!(compact, old_result_idx, new_node, new_idx, idx)
        compact.result_idx = result_idx
        # If this instruction has reverse affinity and we were at the end of a basic block,
        # finish it now.
        if reverse_affinity && idx == last(bb.stmts)+1 && !reverse_affinity_stmt_after(compact, idx-1)
            active_bb += 1
            finish_current_bb!(compact, old_result_idx)
        end
        (old_result_idx == result_idx) && return next(compact, (idx, active_bb))
        return (old_result_idx, compact.result[old_result_idx]), (compact.idx, active_bb)
    end
    # This will get overwritten in future iterations if
    # result_idx is not, incremented, but that's ok and expected
    compact.result_types[old_result_idx] = compact.ir.types[idx]
    compact.result_lines[old_result_idx] = compact.ir.lines[idx]
    result_idx = process_node!(compact, old_result_idx, compact.ir.stmts[idx], idx, idx)
    stmt_if_any = old_result_idx == result_idx ? nothing : compact.result[old_result_idx]
    compact.result_idx = result_idx
    if idx == last(bb.stmts) && !reverse_affinity_stmt_after(compact, idx)
        active_bb += 1
        finish_current_bb!(compact, old_result_idx)
    end
    (old_result_idx == compact.result_idx) && return next(compact, (idx + 1, active_bb))
    compact.idx = idx + 1
    if !isassigned(compact.result, old_result_idx)
        @assert false
    end
    return (old_result_idx, compact.result[old_result_idx]), (compact.idx, active_bb)
end

function maybe_erase_unused!(extra_worklist, compact, idx)
    effect_free = stmt_effect_free(compact.result[idx], compact, compact.ir.mod)
    if effect_free
        for ops in userefs(compact.result[idx])
            val = ops[]
            if isa(val, SSAValue)
                if compact.used_ssas[val.id] == 1
                    if val.id < idx
                        push!(extra_worklist, val.id)
                    end
                end
                compact.used_ssas[val.id] -= 1
            end
        end
        compact.result[idx] = nothing
    end
end

function fixup_phinode_values!(compact, old_values)
    values = Vector{Any}(undef, length(old_values))
    for i = 1:length(old_values)
        isassigned(old_values, i) || continue
        val = old_values[i]
        if isa(val, OldSSAValue)
            val = compact.ssa_rename[val.id]
            if isa(val, SSAValue)
                compact.used_ssas[val.id] += 1
            end
        end
        values[i] = val
    end
    values
end

function just_fixup!(compact)
    for idx in compact.late_fixup
        stmt = compact.result[idx]::Union{PhiNode, PhiCNode}
        if isa(stmt, PhiNode)
            compact.result[idx] = PhiNode(stmt.edges, fixup_phinode_values!(compact, stmt.values))
        else
            compact.result[idx] = PhiCNode(fixup_phinode_values!(compact, stmt.values))
        end
    end
end

function finish(compact::IncrementalCompact)
    just_fixup!(compact)
    # Record this somewhere?
    result_idx = compact.result_idx
    resize!(compact.result, result_idx-1)
    resize!(compact.result_types, result_idx-1)
    resize!(compact.result_lines, result_idx-1)
    bb = compact.result_bbs[end]
    compact.result_bbs[end] = BasicBlock(bb,
                StmtRange(first(bb.stmts), result_idx-1))
    # Perform simple DCE for unused values
    extra_worklist = Int[]
    for (idx, nused) in Iterators.enumerate(compact.used_ssas)
        idx >= result_idx && break
        nused == 0 || continue
        maybe_erase_unused!(extra_worklist, compact, idx)
    end
    while !isempty(extra_worklist)
        maybe_erase_unused!(extra_worklist, compact, pop!(extra_worklist))
    end
    cfg = CFG(compact.result_bbs, Int[first(bb.stmts) for bb in compact.result_bbs[2:end]])
    return IRCode(compact.ir, compact.result, compact.result_types, compact.result_lines, cfg, NewNode[])
end

function compact!(code::IRCode)
    compact = IncrementalCompact(code)
    # Just run through the iterator without any processing
    state = start(compact)
    while !done(compact, state)
        _, state = next(compact, state)
    end
    return finish(compact)
end

struct BBIdxStmt
    ir::IRCode
end

bbidxstmt(ir) = BBIdxStmt(ir)

start(x::BBIdxStmt) = (1,1)
done(x::BBIdxStmt, (idx, bb)) = idx > length(x.ir.stmts)
function next(x::BBIdxStmt, (idx, bb))
    active_bb = x.ir.cfg.blocks[bb]
    next_bb = bb
    if idx == last(active_bb.stmts)
        next_bb += 1
    end
    return (bb, idx, x.ir.stmts[idx]), (idx + 1, next_bb)
end

max_stars_repo_name: NymanLauri/julia

max_issues_repo_name: NymanLauri/julia

max_forks_repo_name: NymanLauri/julia


text:
qq = `qq`
while true
    println("please input the passcode")
    passcode = readline()
    if passcode == "wozhenbang" || passcode == "1"
        break
    end
end
run(qq)

max_stars_repo_name: Basicconstruction/Easy-tools

max_issues_repo_name: Basicconstruction/Easy-tools

max_forks_repo_name: Basicconstruction/Easy-tools


text:
@testset "solve_mf.jl" begin

end


text:
#
# Newton and Newton Schulz Hybrid algoritm
#
# Reference:
# Nicholas J. Higham and Robert S. Schreiber, Fast Polar Decomposition
# of an arbitrary matrix, SIAM, J. Sci. Statist. Comput. Vol. 11, No. 4
# (1990) pp. 648-655.
#

type NewtonHybridAlg{T} <: PolarAlg
    maxiter::Int    # maximum number of iterations.
    verbose::Bool   # whether to show procedural information.
    tol::T    # convergence tolerance
    theta::T  # switch parameter

    function NewtonHybridAlg(;maxiter::Integer=100,
                             verbose::Bool=false,
                             tol::Real=cbrt(eps(T)),
                             theta::Real=convert(T, 0.6))
        maxiter > 1 || error("maxiter must  be greater than 1.")
        tol > 0 || error("tol must be positive.")
        theta > 0 || error("theta must be positive.")

        new(maxiter,
            verbose,
            tol,
            theta)
    end
end

function solve!{T}(alg::NewtonHybridAlg,
                X::Matrix{T}, U::Matrix{T}, H::Matrix{T})
    common_iter_hybr!(NewtonUpdater(true, eps(T)^(1/4)), NewtonSchulzUpdater(), X, U, H,
                      alg.maxiter, alg.verbose, alg.tol, alg.theta)
end


text:
using Multigraphs
import LightGraphs: nv, ne, outneighbors, inneighbors, neighbors, rem_edge!,
    add_edge!, has_edge

export AbstractZXDiagram
export spider_type, phase, spiders, rem_spider!, rem_spiders!

abstract type AbstractZXDiagram{T, P} end

nv(::AbstractZXDiagram) = 0
ne(::AbstractZXDiagram) = 0
outneighbors(::AbstractZXDiagram{T, P}, v) where {T, P} = T[]
inneighbors(::AbstractZXDiagram{T, P}, v) where {T, P} = T[]
neighbors(::AbstractZXDiagram{T, P}, v) where {T, P} = T[]
rem_edge!(::AbstractZXDiagram, x...) = false
rem_edge!(::AbstractZXDiagram, v1, v2) = false
add_edge!(::AbstractZXDiagram, x...) = false
add_edge!(::AbstractZXDiagram, v1, v2) = false
has_edge(::AbstractZXDiagram, vs...) = false

spider_type(::AbstractZXDiagram, v) = nothing
phase(::AbstractZXDiagram, v) = nothing
spiders(::AbstractZXDiagram{T, P}) where {T, P} = T[]
rem_spider!(::AbstractZXDiagram, v) = false
rem_spiders!(::AbstractZXDiagram, vs)= false


text:
# GCM-specific Sources
# This file contains helpers and lists currently available options

# Current options for GCM-specific sources:
"""
Defines a forcing that parametrises radiative and frictional effects using
Newtonian relaxation and Rayleigh friction, following Held and Suarez (1994)
"""
function held_suarez_forcing!(
    bl,
    source,
    state,
    diffusive,
    aux,
    t::Real,
    direction,
)
    FT = eltype(state)

    # Parameters
    T_ref = FT(255)

    # Extract the state
    ρ = state.ρ
    ρu = state.ρu
    ρe = state.ρe

    ts = recover_thermo_state(bl, state, aux)
    e_int = internal_energy(ts)
    T = air_temperature(ts)
    _R_d = FT(R_d(bl.param_set))
    _day = FT(day(bl.param_set))
    _grav = FT(grav(bl.param_set))
    _cp_d = FT(cp_d(bl.param_set))
    _cv_d = FT(cv_d(bl.param_set))
    _p0 = FT(MSLP(bl.param_set))

    # Held-Suarez parameters
    k_a = FT(1 / (40 * _day))
    k_f = FT(1 / _day)
    k_s = FT(1 / (4 * _day))
    ΔT_y = FT(60)
    Δθ_z = FT(10)
    T_equator = FT(315)
    T_min = FT(200)
    σ_b = FT(7 / 10)

    # Held-Suarez forcing
    φ = latitude(bl, aux)
    p = air_pressure(ts)

    #TODO: replace _p0 with dynamic surface pressure in Δσ calculations to account
    #for topography, but leave unchanged for calculations of σ involved in T_equil
    σ = p / _p0
    exner_p = σ^(_R_d / _cp_d)
    Δσ = (σ - σ_b) / (1 - σ_b)
    height_factor = max(0, Δσ)
    T_equil = (T_equator - ΔT_y * sin(φ)^2 - Δθ_z * log(σ) * cos(φ)^2) * exner_p
    T_equil = max(T_min, T_equil)
    k_T = k_a + (k_s - k_a) * height_factor * cos(φ)^4
    k_v = k_f * height_factor

    # Apply Held-Suarez forcing
    source.ρu -= k_v * projection_tangential(bl, aux, ρu)
    source.ρe -= k_T * ρ * _cv_d * (T - T_equil)

    return nothing
end

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA