text:
function checked_mod1(x::T, y::T) where T<:Base.BitInteger
    result = checked_mod(x, y)
    result = ifelse(result === zero(T), y, result)
    return result
end

function checked_fld1(x::T, y::T) where T<:Base.BitInteger
    d = checked_div(x, y)
    return d + (!signbit(x ⊻ y) & (d * y !== x))
end

for OP in (:div, :rem, :fld, :cld)
    for S in (:SafeSigned, :SafeUnsigned)
        @eval begin
            function $OP(x::T, y::T, ::typeof(RoundToZero)) where {T <: $S}
                xx = baseint(x)
                yy = baseint(y)
                result = $OP(xx, yy)
                return $S(result)
            end
        end
    end
end

for (OP, CHK) in ((:(+), :checked_add), (:(-), :checked_sub),
                  (:(*), :checked_mul), (:div, :checked_div),
                  (:fld, :checked_fld), (:cld, :checked_cld),
                  (:rem, :checked_rem), (:mod, :checked_mod),
                  (:mod1, :checked_mod1), (:fld1, :checked_fld1))
    @eval begin
        @inline function $OP(x::T, y::T) where T<:SafeSigned
            ix = baseint(x)
            iy = baseint(y)
            result = $CHK(ix, iy)
            return safeint(result)
        end

        @inline function $OP(x::T, y::T) where T<:SafeUnsigned
            ix = baseint(x)
            iy = baseint(y)
            result = $CHK(ix, iy)
            return safeint(result)
        end
    end
end

for (OP, CHK) in ((:(+), :checked_add), (:(-), :checked_sub),
    (:(*), :checked_mul), (:div, :checked_div),
    (:fld, :checked_fld), (:cld, :checked_cld),
    (:rem, :checked_rem), (:mod, :checked_mod),
    (:mod1, :checked_mod1), (:fld1, :checked_fld1))
    for (T1, T2) in PairedSafes
        @eval begin
            @inline function $OP(x::$T1, y::$T2)
                xx, yy = promote(x, y)
                ix = baseint(xx)
                iy = baseint(yy)
                result = $CHK(ix, iy)
                return safeint(result)
            end
        end
    end
end

for (OP, CHK) in ((:(+), :checked_add), (:(-), :checked_sub),
    (:(*), :checked_mul), (:div, :checked_div),
    (:fld, :checked_fld), (:cld, :checked_cld),
    (:rem, :checked_rem), (:mod, :checked_mod),
    (:mod1, :checked_mod1), (:fld1, :checked_fld1))
    for (T1, T2) in MixedInts
        @eval begin
            @inline function $OP(x::$T1, y::$T2)
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end
        end
    end
end

for (OP, CHK) in ((:(+), :checked_add), (:(-), :checked_sub),
    (:(*), :checked_mul), (:div, :checked_div),
    (:fld, :checked_fld), (:cld, :checked_cld),
    (:rem, :checked_rem), (:mod, :checked_mod),
    (:mod1, :checked_mod1), (:fld1, :checked_fld1))
    @eval begin
        @inline function $OP(x::T, y::Bool) where {T<:SafeInteger}
            xx, yy = promote(x, y)
            return $OP(xx, yy)
        end

        @inline function $OP(x::Bool, y::T) where {T<:SafeInteger}
            xx, yy = promote(x, y)
            return $OP(xx, yy)
        end
    end
end

for T in UnpairedSafes
  @eval begin

    function (/)(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        checked_div(ix, iy)
        result = ix / iy
        return result
    end

    function (\)(x::S, y::S) where S<:$T
        ix = baseint(y)
        iy = baseint(x)
        checked_div(iy, ix)
        result = ix / iy
        return result
    end

    function divrem(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        return safeint(div(ix, iy)), safeint(rem(ix, iy)) # div, rem already are checked
    end

    function fldmod(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        return safeint(fld(ix, iy)), safeint(mod(ix, iy)) # fld, mod already are checked
    end

    function fldmod1(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        return safeint(fld1(ix, iy)), safeint(mod1(ix, iy)) # fld1, mod1 already are checked
    end

    function gcd(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        return safeint(gcd(ix, iy))
    end

    function lcm(x::S, y::S) where S<:$T
        ix = baseint(x)
        iy = baseint(y)
        return safeint(lcm(ix, iy))
    end

    function divgcd(x::S, y::S) where S<:$T
        g = gcd(x,y)
        return div(x,g), div(y,g)
    end

  end
end

for OP in (:divrem, :fldmod, :fldmod1, :divgcd, :gcd, :lcm)
    for T in UnpairedSafes
        @eval begin
            @inline function $OP(x::S, y::Bool) where {S<:$T}
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end

            @inline function $OP(x::Bool, y::S) where {S<:$T}
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end
        end
    end
end

for OP in (:(/), :(\))
    for T in UnpairedSafes
        @eval begin
            @inline function $OP(x::S, y::Bool) where {S<:$T}
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end

            @inline function $OP(x::Bool, y::S) where {S<:$T}
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end
        end
    end
end

for OP in (:divrem, :fldmod, :fldmod1, :divgcd, :gcd, :lcm)
    for (T1, T2) in MixedPairs
        @eval begin
            @inline function $OP(x::TX, y::TY) where {TX<:$T1, TY<:$T2}
                xx, yy = promote(x, y)
                return $OP(xx, yy)
            end
        end
    end
end

for (T1, T2) in MixedPairs
    @eval begin
        @inline function /(x::$T1, y::$T2)
            xx, yy = promote(x, y)
            return /(xx, yy)
        end
        @inline function \(x::$T1, y::$T2)
            xx, yy = promote(x, y)
            return \(xx, yy)
        end
    end
end


text:
# avoid creating a progress bar with @withprogress if progress logging is disabled
# and add a custom progress logger if the current logger does not seem to be able to handle
# progress logs
macro ifwithprogresslogger(progress, exprs...)
    return quote
        if $progress
            if $hasprogresslevel($Logging.current_logger())
                $ProgressLogging.@withprogress $(exprs...)
            else
                $with_progresslogger($Base.@__MODULE__, $Logging.current_logger()) do
                    $ProgressLogging.@withprogress $(exprs...)
                end
            end
        else
            $(exprs[end])
        end
    end |> esc
end

# improved checks?
function hasprogresslevel(logger)
    return Logging.min_enabled_level(logger) ≤ ProgressLogging.ProgressLevel
end

# filter better, e.g., according to group?
function with_progresslogger(f, _module, logger)
    logger1 = LoggingExtras.EarlyFilteredLogger(progresslogger()) do log
        log._module === _module && log.level == ProgressLogging.ProgressLevel
    end
    logger2 = LoggingExtras.EarlyFilteredLogger(logger) do log
        log._module !== _module || log.level != ProgressLogging.ProgressLevel
    end

    Logging.with_logger(f, LoggingExtras.TeeLogger(logger1, logger2))
end

function progresslogger()
    # detect if code is running under IJulia since TerminalLogger does not work with IJulia
    # https://github.com/JuliaLang/IJulia.jl#detecting-that-code-is-running-under-ijulia
    if (Sys.iswindows() && VERSION < v"1.5.3") || (isdefined(Main, :IJulia) && Main.IJulia.inited)
        return ConsoleProgressMonitor.ProgressLogger()
    else
        return TerminalLoggers.TerminalLogger()
    end
end

##

using ProgressLogging: _asprogress, ProgressString, ProgressLevel, uuid4
using Base.Meta: isexpr

const _id_var_children = gensym(:progress_id_children)
const _name_var_children = gensym(:progress_name_children)

macro progressid_child(N)
    id_err = "`@progressid_child` must be used inside `@withprogress_children`"

    quote
        $Base.@isdefined($_id_var_children) ? $_id_var_children[$N] : $error($id_err)
    end |> esc
end

macro logprogress_child(args...)
    _logprogress_child(args...)
end

function _logprogress_child(N, name, progress = nothing, args...)
    name_expr = :($Base.@isdefined($_name_var_children) ? $_name_var_children[$N] : "")
    if progress == nothing
        # Handle: @logprogress progress
        kwargs = (:(progress = $name), args...)
        progress = name
        name = name_expr
    elseif isexpr(progress, :(=)) && progress.args[1] isa Symbol
        # Handle: @logprogress progress key1=val1 ...
        kwargs = (:(progress = $name), progress, args...)
        progress = name
        name = name_expr
    else
        # Otherwise, it's: @logprogress name progress key1=val1 ...
        kwargs = (:(progress = $progress), args...)
    end

    id_err = "`@logprogress_child` must be used inside `@withprogress_children`"
    id_expr = :($Base.@isdefined($_id_var_children) ? $_id_var_children[$N] : $error($id_err))

    @gensym id_tmp
    # Emitting progress log record as old/open API (i.e., using
    # `progress` key) and _also_ as new API based on `Progress` type.
    msgexpr = :($ProgressString($_asprogress(
        $name,
        $id_tmp,
        $(ProgressLogging._id_var);
        progress = $progress,
    )))
    quote
        $id_tmp = $id_expr
        $Logging.@logmsg($ProgressLevel, $msgexpr, $(kwargs...), _id = $id_tmp)
    end |> esc
end

macro withprogress_children(N, exprs...)
    _withprogress_children(N, exprs...)
end

function _withprogress_children(N, exprs...)
    length(exprs) == 0 &&
        throw(ArgumentError("`@withprogress_children` requires at least one number and one expression."))

    m = ProgressLogging.@__MODULE__

    kwargs = Dict{Symbol,Any}(:names => :(["" for i in 1:$N]))
    unsupported = []
    for kw in exprs[1:end-1]
        if isexpr(kw, :(=)) && length(kw.args) == 2 && haskey(kwargs, kw.args[1])
            kwargs[kw.args[1]] = kw.args[2]
        else
            push!(unsupported, kw)
        end
    end

    # Error on invalid input expressions:
    if !isempty(unsupported)
        msg = sprint() do io
            println(io, "Unsupported optional arguments:")
            for kw in unsupported
                println(io, kw)
            end
            print(io, "`@withprogress_children` supports only following keyword arguments: ")
            join(io, keys(kwargs), ", ")
        end
        throw(ArgumentError(msg))
    end

    ex = exprs[end]
    id_err = "`@withprogress_children` must be used inside `@withprogress`"

    i_var = gensym()
    quote
        let
            $Base.@isdefined($(ProgressLogging._id_var)) || $error($id_err)
            $_id_var_children = [ $uuid4() for i in 1:$N ]
            $_name_var_children = $(kwargs[:names])

            for $i_var in 1:$N
                @logprogress_child $i_var nothing
            end

            try
                $ex
            finally
                for $i_var in 1:$N
                    @logprogress_child $i_var nothing
                end
            end
        end
    end |> esc
end

macro ifwithprogresslogger_children(progress, exprs...)
    return quote
        if $progress
            if $hasprogresslevel($Logging.current_logger())
                @withprogress_children $(exprs...)
            else
                $with_progresslogger($Base.@__MODULE__, $Logging.current_logger()) do
                    @withprogress_children $(exprs...)
                end
            end
        else
            $(exprs[end])
        end
    end |> esc
end

ext: jl


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 430d5134-92af-11eb-1fcc-d71f5216d62a
#Set up the working environment
begin
	using Pkg
	Pkg.activate("..")
	using Plots, PlutoUI,LaTeXStrings, DifferentialEquations
end

# ╔═╡ 5adc7e34-92af-11eb-2cfe-67899a4a0291
md"""
## Attempting a more realistic model

Next, we'll try to make the model more realistic by including **reinfection**, **vaccination** and **mortalitiy**.

##### Mortality
We'll only consider deaths of infected people at a mortality rate $m$: $$I \xrightarrow{m} D$$

##### Vaccination & Reinfection
To account for vaccination, we'll assume that
1. The rate of infection is different for vaccinated people $V$ and susceptible people $S$, represented by a *vaccination immunity factor*, $i_v < 1$

$$S \xrightarrow{r_s} I \xleftarrow{r_s \cdot i_v} V$$

2. Likewise, susceptible people $S$ and recovered people $R$ become infected at different rates, represented by an *recovered immunity factor*, $i_r < 1$ that lowers the infection risk

$$S \xrightarrow{r_s} I \xleftarrow{r_s \cdot i_r} R$$

3. We assume that recovered people $R$ and susceptible people $S$ are vaccinated at the same rate $v$

$$S \xrightarrow{v} V \xleftarrow{v} R$$

Taken together, this yields the following set of differential equations:

$$\begin{align}
	\frac{dS}{dt} = -rSI -vS\\
	\frac{dV}{dt} = vS + vR - i_v r VI\\
	\frac{dD}{dt} = mI\\
	\frac{dR}{dt} = aI - i_r rRI - vR\\
	\frac{dI}{dt} = rSI + i_r rRI + i_v r VI - mI - aI
\end{align}$$

"""

# ╔═╡ 7a19b238-92af-11eb-30fd-b189d06a0882
md""" ## Simulation
*Change the sliders to adjust the corresponding values*:

| parameter                   | value                     |
|-----------------------------|---------------------------|
| transmission rate $r$       | $(@bind r2 Slider(0:0.005:1, default=0.125, show_value=true)) |
| recovery rate $a$           | $(@bind a2 Slider(0:0.001:0.25, default=0.1, show_value=true))|
| rate of vaccination $v$     | $(@bind v2 Slider(0:0.0001:0.01, default=0.0008, show_value=true)) |
| mortality $m$               | $(@bind m2 Slider(0:0.0001:0.05, default=0.0006, show_value=true)) |
| % immunity recovered        | $(@bind ir Slider(0:1.0:100.0, default=60, show_value=true)) |
| % immunity vaccinated       | $(@bind iv Slider(0:1.0:100.0, default=90, show_value=true)) |
|  $I(0)$                     | $(@bind I02 Slider(0.0001:0.001:0.1, default=0.004, show_value=true)) |
|  $R(0)$                     | $(@bind R02 Slider(0:0.001:0.1, default=0.055,show_value=true)) |
| $V(0)$                      | $(@bind V02 Slider(0:0.001:0.1, default=0.048,show_value=true)) |
| $D(0)$                      | $(@bind D02 Slider(0:0.001:0.1, default=0.001,show_value=true)) |
| timespan $t_{max}$          | $(@bind tmax2 Slider(10.0:10.0:2000.0, default=100.0,show_value=true)) days |
| zoom y-axis                 | $(@bind ymax2 Slider(0.001:0.001:1.0, default=1.0,show_value=true)) |
| Population size             | $(@bind popsize NumberField(500000:10000000000, default=9044650)) |
"""

# ╔═╡ 8f867f9a-92af-11eb-3671-43904aa74d8b
function getres(solution, i, t, p)
	v = solution(t)[i]*p
	return round(maximum([0.0 v]))
end

# ╔═╡ 96d8a5c2-92af-11eb-1d3f-9bdd035eafec
#S, I, R, V, D
function sirvd!(du,u,p,t)
	du[1] = -u[1] * ( p[1] * u[2] + p[3])
	du[2] = u[2] * (p[1] * (u[1] + p[5] * u[4] + p[4] * u[3]) - p[6] - p[2])
	du[3] = p[2] * u[2] - u[3] * ( p[1] * p[4] *  u[2] + p[3])
	du[4] = p[3] * (u[1] + u[3]) - u[4] * u[2] * p[1] * p[5]
	du[5] = p[6] * u[2]
end

# ╔═╡ 9c0dfe82-92af-11eb-3192-49c2deac6e9b
p2 = [r2, a2, v2, (100 - ir)/100, (100 - iv)/100, m2]

# ╔═╡ a1e4ca5e-92af-11eb-1c40-410a8d3abaca
u02 = [1-I02-R02-V02-D02, I02, R02, V02, D02]

# ╔═╡ ada6253e-92af-11eb-0e32-d9760ab88271
tspan2 = (0.0, 2000)

# ╔═╡ b08f47b2-92af-11eb-367b-979567aebb4d
prob2 = ODEProblem(sirvd!,u02,tspan2,p2)

# ╔═╡ 88b411f0-92af-11eb-2abf-151b2bfb313b
begin
	sol2 = solve(prob2, dt=1.0, dense = true)
	infected = sol2[2,:]*popsize
	dead = infected + sol2[5,:]*popsize
	recovered = dead + sol2[3,:]*popsize
	vaccinated = recovered + sol2[4,:]*popsize
	susceptible = vaccinated + sol2[1,:]*popsize
	icu = infected*0.015
	icu_cap = [1000 for t in sol2.t]

	sim_stackarea = plot(sol2.t, [susceptible, vaccinated, recovered, dead, infected, icu],
				fill = true,
				color = ["gray89" "skyblue3" "paleturquoise3" "gray39" "lightcoral" "maroon"],
				fillalpha = 1,
				label = ["S(t)" "V(t)" "R(t)" "D(t)" "I(t)" L"$I(t)_{icu}$"],
				xlabel = "t / days",
				title = "Timecourse (stacked area)",
				xlims = (0, tmax2),
				ylims = (0, ymax2*popsize))

	sim_timecourse = plot(sol2,
				color = ["gray89" "lightcoral" "paleturquoise3" "skyblue3" "gray39"],
				label = ["S(t)" "I(t)" "R(t)" "V(t)" "D(t)"],
				xlabel = "t / days",
				title = "Timecourse",
				xlims = (0, tmax2),
				ylims = (0, ymax2))

	sim_rate_tc = plot(0:tmax2, sol2(0:tmax2, Val{1}, idxs=2),
				color = ["lightcoral"],
				label = [L"\frac{dI}{dt}"],
				xlabel = "t / days",
				title = "New Infections",
				xlims = (0, tmax2))

	timecourses = plot(sim_timecourse, sim_rate_tc, layout = (2,1))

	plot(sim_stackarea, timecourses, layout = (1, 2),fontfamily="Computer Modern")
end

# ╔═╡ 7ec0712a-92af-11eb-2144-15d33af27c3d
md"""
| t (days) | I(t)                               | R(t)                               | V(t)                               | D(t)                               |
|----------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|
| 0        | $(getres(sol2, 2, 0, popsize))     | $(getres(sol2, 3, 0, popsize))     | $(getres(sol2, 4, 0, popsize))     | $(getres(sol2, 5, 0, popsize))     |
| $tmax2   | $(getres(sol2, 2, tmax2, popsize)) | $(getres(sol2, 3, tmax2, popsize)) | $(getres(sol2, 4, tmax2, popsize)) | $(getres(sol2, 5, tmax2, popsize)) |

Base reprodution rate = $(round((r2*(1-I02-R02-V02-D02)/a2), sigdigits = 3)),
Infectious period = $(round(1/a2, sigdigits = 3))

"""

# ╔═╡ 230527fa-92b3-11eb-308f-ffab3777029c
sol2(0:tmax2, Val{1}, idxs=2)

# ╔═╡ Cell order:
# ╟─430d5134-92af-11eb-1fcc-d71f5216d62a
# ╟─5adc7e34-92af-11eb-2cfe-67899a4a0291
# ╟─7a19b238-92af-11eb-30fd-b189d06a0882
# ╟─7ec0712a-92af-11eb-2144-15d33af27c3d
# ╟─88b411f0-92af-11eb-2abf-151b2bfb313b
# ╟─230527fa-92b3-11eb-308f-ffab3777029c
# ╟─8f867f9a-92af-11eb-3671-43904aa74d8b
# ╟─96d8a5c2-92af-11eb-1d3f-9bdd035eafec
# ╟─9c0dfe82-92af-11eb-3192-49c2deac6e9b
# ╟─a1e4ca5e-92af-11eb-1c40-410a8d3abaca
# ╟─ada6253e-92af-11eb-0e32-d9760ab88271
# ╟─b08f47b2-92af-11eb-367b-979567aebb4d

max_stars_repo_name: el-uhu/intro-to-modelling

max_issues_repo_name: el-uhu/intro-to-modelling

max_forks_repo_name: el-uhu/intro-to-modelling


text:
"""
`ϕs = tinterpolate(ϕsindex, tindex, nstack)` uses linear
interpolation/extrapolation in time to "fill out" to times `1:nstack`
a deformation defined intermediate times `tindex` . Note that
`ϕs[tindex] == ϕsindex`.
"""
function tinterpolate(ϕsindex, tindex, nstack)
    ϕs = Vector{eltype(ϕsindex)}(undef, nstack)
    # Before the first tindex
    k = 0
    for i in 1:tindex[1]-1
        ϕs[k+=1] = ϕsindex[1]
    end
    # Within tindex
    for i in 2:length(tindex)
        ϕ1 = ϕsindex[i-1]
        ϕ2 = ϕsindex[i]
        Δt = tindex[i] - tindex[i-1]
        for j = 0:Δt-1
            α = convert(eltype(eltype(ϕ1.u)), j/Δt)
            ϕs[k+=1] = GridDeformation((1-α)*ϕ1.u + α*ϕ2.u, ϕ2.nodes)
        end
    end
    # After the last tindex
    for i in tindex[end]:nstack
        ϕs[k+=1] = ϕsindex[end]
    end
    return ϕs
end

"""
    ϕs′ = medfilt(ϕs)

Perform temporal median-filtering on a sequence of deformations. This is a form of smoothing
that does not "round the corners" on sudden (but persistent) shifts.
"""
function medfilt(ϕs::AbstractVector{D}, n) where D<:AbstractDeformation
    nhalf = n>>1
    2nhalf+1 == n || error("filter size must be odd")
    T = eltype(eltype(D))
    v = Array{T}(undef, ndims(D), n)
    vs = ntuple(d->view(v, d, :), ndims(D))
    ϕ1 = copy(ϕs[1])
    ϕout = Vector{typeof(ϕ1)}(undef, length(ϕs))
    ϕout[1] = ϕ1
    _medfilt!(ϕout, ϕs, v, vs)  # function barrier due to instability of vs
end

@noinline function _medfilt!(ϕout, ϕs, v, vs::NTuple{N,T}) where {N,T}
    n = size(v,2)
    nhalf = n>>1
    tmp = Vector{eltype(T)}(undef, N)
    u1 = ϕout[1].u
    for i = 1+nhalf:length(ϕs)-nhalf
        u = similar(u1)
        for I in eachindex(ϕs[i].u)
            for j = -nhalf:nhalf
                utmp = ϕs[i+j].u[I]
                for d = 1:N
                    v[d, j+nhalf+1] = utmp[d]
                end
            end
            for d = 1:N
                tmp[d] = median!(vs[d])
            end
            u[I] = tmp
        end
        ϕout[i] = GridDeformation(u, ϕs[i].nodes)
    end
    # Copy the beginning and end
    for i = 2:nhalf  # we did [1] in medfilt
        ϕout[i] = copy(ϕs[i])
    end
    for i = length(ϕs)-nhalf+1:length(ϕs)
        ϕout[i] = copy(ϕs[i])
    end
    ϕout
end


text:
function TDMFile()
    return TDMFile(())
end

function add_observations_block(obj::TDMFile)
    return jcall(obj, "addObservationsBlock", void, ())
end

function check_time_systems(obj::TDMFile)
    return jcall(obj, "checkTimeSystems", void, ())
end

function get_observations_blocks(obj::TDMFile)
    return jcall(obj, "getObservationsBlocks", List, ())
end

function set_observations_blocks(obj::TDMFile, arg0::List)
    return jcall(obj, "setObservationsBlocks", void, (List,), arg0)
end


text:
@testset "GridTools" begin
    nt = length(T_all)
    n = nt * nb
    x₀ = p₀.xgeo * ones(n)
    @testset "convert to 3D array" for xᵢ in state_to_tracers(x₀, nb, nt)
        xᵢ3D = rearrange_into_3Darray(xᵢ, wet3D)
        @test size(xᵢ3D) == size(wet3D)
        @test size(rearrange_into_1Dvector(xᵢ3D, wet3D)) == size(xᵢ)
    end
end


text:
# Fourier spectral discretisation with periodic boundary conditions

struct FourierDiscretisation{N} <: Sampling{N}
    n::Int
    h::Float64
end

function D!(y, d::FourierDiscretisation{2}, a, u, n, axis)
    iks = 2π * complex.(0, fftfreq(d.n, 1 / d.h))
    if axis == 2
        iks = transpose(iks)
    end
    buf = copy(u)
    fft!(buf, axis)
    @. buf *= iks^(n)
    ifft!(buf, axis)
    @. y += a * buf
end

# TODO test this with sin or something that's smooth at the boundary
interpolant(d::FourierDiscretisation{2}, r) = finterp(d, r, 0.0)

# Convolution is easy in a Fourier basis
finterp(d::FourierDiscretisation{2}, r, a) =
    finterp1(d, real(r), a) .* transpose(finterp1(d, imag(r), a))

function finterp1(d, r::Float64, a = 0.0)
    a < eps() && return finterp0(d, r)
    x = first(daxes(d))[:]
    # Centre on a grid point in case a « h
    j = argmin(@. abs(x - r))
    x0 = x[j]
    u = similar(x, Complex{Float64})
    @. u = exp(-(x - x0)^2 / 2a^2)
    fft!(u)
    ks = 2π * fftfreq(d.n, 1 / d.h)
    @. u *= exp(-1im * ks * (r - x0))
    ifft!(u)
    real(u)
end

# sinc shortcut for a = 0
function finterp0(d, r::Float64)
    x = first(daxes(d))[:]
    y = @. sin(π*(x-r)/d.h)/d.n/tan(π*(x-r)/d.n/d.h)
    @. y[abs(x-r) < eps()] = 1
    y
end

function dmats(d::FourierDiscretisation{2})
    # differentiate delta functions

    u = zeros(Complex{Float64}, d.n, 1)
    u[1] = 1
    D1 = Array{Complex{Float64}}(undef, d.n, d.n)
    y = zeros(Complex{Float64}, d.n, 1)
    dif!(y, d, 1, u; axis=1)
    for j = 1:d.n
        D1[:, j] = circshift(y, j-1)
    end
    D2 = Array{Complex{Float64}}(undef, d.n, d.n)
    y = zeros(Complex{Float64}, d.n, 1)
    dif2!(y, d, 1, u; axis=1)
    for j = 1:d.n
        D2[:, j] = circshift(y, j-1)
    end
    D1, D2
end

"""
    matrices(s::Superfluid{2}, d::FourierDiscretisation{2}, Ω, ψ, syms::Vararg{Symbol})

Like operators, but as matrices instead of functions.

The argument ψ is an order parameter to linearise U about.
This is the kind of place that TensArrays will be useful.

TODO Combine the FDDiscretisation and FourierDiscretisation methods
into a Sampling method (rename Sampling to Collocation?)
"""
function matrices(s::Superfluid{2}, d::FourierDiscretisation{2}, Ω, ψ, syms::Vararg{Symbol})
    x, y = daxes(d)
    ∂, ∂² = dmats(d)
    V = diagm(0 => sample(s.V, d)[:])
    eye = Matrix(I, d.n, d.n)
    ρ = diagm(0 => abs2.(ψ[:]))

    T = -s.hbm * kron(eye, ∂²) / 2 - s.hbm * kron(∂², eye) / 2
    U = s.C / d.h^2 * ρ
    J = -1im * (repeat(x, 1, d.n)[:] .* kron(∂, eye) - repeat(y, d.n, 1)[:] .* kron(eye, ∂))
    L = T + V + U - Ω * J
    H = T + V + U / 2 - Ω * J

    ops = Dict(:V => V, :T => T, :U => U, :J => J, :L => L, :H => H)
    isempty(syms) ? ops : [ops[j] for j in syms]
end


text:
module ModelExploration

export Product, Literal, ModelSpace, PullbackSpace, ModelHom, SliceDiagHom, select,
       unfold, to_model_hom, SliceCat, InduceHomIn, InduceHomOut, LitFG, LitNG,
       PushoutSpace, LitModelHom

using Catlab.CategoricalAlgebra
using Catlab.CategoricalAlgebra: pullback
using Catlab.Graphs, Catlab.Present
using Catlab.Theories: attrtype

#############
const SliceCat{Hom} = TypeCat{Slice{Hom}, SliceHom{Hom}}

abstract type ModelSpace{T,Base} end
abstract type FunctorGen end
abstract type NatGen end

"""
Base case: an explicit diagram considered as a model space

E.g.

1. Shape: •, and diagram maps this object to the petri net 'infectious_type'
2. Shape: • •, and diagram maps these to SIR and SIS.
3. Shape: • •, and diagram maps these to SIR->I.T. and SIS->I.T morphisms
"""
struct Literal{T,Base} <: ModelSpace{T,Base}
  lit::Diagram{T,Base}
end

"""
A recipe for constructing a modelspace (diagram) homomorphism.

E.g.
dom = Literal example #2
codom = Literal example #1
functor F = uniquely determined
nat= homomorphisms from SIR and SIS to Infectious_type (explained below)

Dom => F;Codom :: Functors from shape(dom) -> Petri
data: map from ob(shape(dom)) -> Hom(Petri). αᵢ:: F(i) -> G(i)
only morphisms are identities, so naturality is automatically satisfied.
Therefore, αᵢ merely requires a homomorphism from SIR and SIS to Infectious_type

Therefore, this is equivalent to Literal example #3.
"""
struct ModelHom{T,Base} <: ModelSpace{T,Base}
  dom::ModelSpace{T,Base}
  codom::ModelSpace{T,Base}
  fun::FunctorGen
  nat::NatGen
end

"""
Model space generated by pullback of two diagrams
(sliced over the same base diagram)
"""
struct PullbackSpace{T,Base} <: ModelSpace{T,Base}
  g1::ModelSpace{T,Base}
  g2::ModelSpace{T,Base}
end

struct PushoutSpace{T,Base} <: ModelSpace{T,Base}
  g1::ModelSpace{T,Base}
  g2::ModelSpace{T,Base}
end
function LitModelHom(d::DiagramHom)
  ModelHom(Literal(dom(d)), Literal(codom(d)),
           LitFG(shape_map(d)),
           LitNG(Dict(collect(components(diagram_map(d))))))
end
struct LitFG <: FunctorGen
  lit::Functor
end
struct LitNG <: NatGen
  lit::AbstractDict
end
struct UniqueFG <: FunctorGen end
struct UniqueNG <: NatGen end
struct AllFG <: FunctorGen end
struct AllNG <: NatGen end
# idea: constrain the search for functors/nat transformations declarative
# with positive/negative application conditions.

"""
Assume that the model space given as an argument is diagram in a slice category.

Generates a diagram hom between diagrams in the underlying category of the slice
category.

For example, Diagram{T,Petri/X} is converted into DiagramHom{T,Petri}, where
the codomain has shape • and sends that to the Petri net X.
"""
struct SliceDiagHom{T,BaseHom} <: ModelSpace{T,SliceCat{BaseHom}}
  slice_diagram::ModelSpace{T,SliceCat{BaseHom}}
end

"""
One may have a diagram homomorphism into A, but want to apply it to some
composite of limits/colimits involving A.

For example, a map X->A would induce a map into X->A+B
And a map X->A would induce a map XxB -> AxB
"""
struct InduceHomIn
  F::ModelHom
  P::ModelSpace
end

"""
One may have a diagram homomorphism out of A, but want to apply it to some
composite of limits/colimits involving A.

For example, a map A->X would induce a map into A+B->X+B
And a map A->X would induce a map AxB -> X
"""
struct InduceHomOut
  F::ModelHom
  P::ModelSpace
end

############################################################################

@present ThOne(FreeSchema) begin
  X::Ob
end
One = FinCat(ThOne)
##########################################################################
unfold(g::Literal)::Diagram = g.lit

function unfold(g::SliceDiagHom)::Diagram
  res = unfold(g.slice_diagram)
  for ob in values(ob_map(diagram(res)))
    is_natural(ob.slice) || error("Not natural")
  end
  return res
end

"""
Generate the data of a diagram homomorphism from a slice hom
"""
function to_model_hom(g::SliceDiagHom{T,BaseHom}) where {T,BaseHom}
  Base = TypeCat{T, BaseHom}
  slicehom = diagram(unfold(g)) # underlying functor of diagram
  d_ob = Dict(map(collect(ob_map(slicehom))) do (k, v)
    k => dom(v.slice)
  end)
  d_hom = Dict(map(collect(hom_map(slicehom))) do (k, v)
    k => v.f
  end)
  X = codom(last(first(collect(ob_map(slicehom)))).slice)
  d = Diagram{T}(FinDomFunctor(d_ob,d_hom,dom(slicehom), Base()))
  cod =Diagram{T}(FinDomFunctor(Dict(:X=>X),nothing,One, Base()))
  shape_m_ob = Dict([k=>:X for k in keys(d_ob)])
  shape_m_hom = Dict([h=>id(One[:X]) for h in hom_generators(dom(slicehom))])
  shape_m = FinDomFunctor(shape_m_ob, shape_m_hom, dom(slicehom), One)
  diagram_m = Dict(map(collect(ob_map(slicehom))) do (k, v)
    k => v.slice
  end)
  ModelHom{T,Base}(Literal{T,Base}(d), Literal{T,Base}(cod),
                   LitFG(shape_m), LitNG(diagram_m),)
end

unfold(l::LitFG) = l.lit
unfold(l::LitNG) = l.lit

unfold(g::ModelHom{T,Base}) where {T,Base} = DiagramHom{T}(
    unfold(g.fun), unfold(g.nat), unfold(g.dom), unfold(g.codom))

#=function unfold(g::SliceLit; loose::Bool=true)::Vector{ACSetTransformation}
  transformations = ACSetTransformation[]
  S = getS(g.base)
  b = map(g.base; Dict([at => x->nothing for at in attrtype(S)])...)
  for (model, constraint) in g.models
    append!(transformations,
            homomorphisms(model, b; loose=loose, initial=constraint))
  end
  return transformations
end=#

"""Assumes slice Generators. TODO: search homomorphisms otherwise"""
#=function unfold(g::Pullback)::Vector{StructACSet}
  models = StructACSet[]
  for x in unfold(g.g1)
    for y in unfold(g.g2)
      push!(models, ob(pullback(x, y)))
    end
  end
  return models
end=#

# returns a diag limit
function unfold(g::PullbackSpace{T,Base}) where {T,Base}
  pb = pullback(Multicospan([unfold(g.g1), unfold(g.g2)]))
  return pb
end

function unfold(g::PushoutSpace{T,Base}) where {T,Base}
  pb = pushout(Multispan([unfold(g.g1), unfold(g.g2)]))
  return pb
end

function select(g::ModelSpace, lossFn::Function)::StructACSet

  best = typemax(Float64)
  best_m = nothing
  for (i,m) in enumerate(unfold(g)[1:9])
      println("i $i")
      lss = lossFn(m)
      println("\tloss $lss")
      if lss < best
        best_m = m
      end
  end
  return best_m
end

end # module


text:
function initialize_inner!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::PSY.DynamicInverter{C, O, PSY.VoltageModeControl, DC, P, F},
) where {
    C <: PSY.Converter,
    O <: PSY.OuterControl,
    DC <: PSY.DCSource,
    P <: PSY.FrequencyEstimator,
    F <: PSY.Filter,
}

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.VoltageModeControl)
    Ir_filter = device_states[external_ix[1]]
    Ii_filter = device_states[external_ix[2]]
    Ir_cnv = device_states[external_ix[3]]
    Ii_cnv = device_states[external_ix[4]]
    Vr_filter = device_states[external_ix[5]] #TODO: Should be inner reference after initialization
    Vi_filter = device_states[external_ix[6]] #TODO: Should be inner reference after initialization

    #Obtain inner variables for component
    ω_oc = PSY.get_ω_ref(dynamic_device)
    θ0_oc = get_inner_vars(dynamic_device)[θ_oc_var]
    Vdc = get_inner_vars(dynamic_device)[Vdc_var]

    #Obtain output of converter
    Vr_cnv0 = get_inner_vars(dynamic_device)[Vr_cnv_var]
    Vi_cnv0 = get_inner_vars(dynamic_device)[Vi_cnv_var]

    #Get Voltage Controller parameters
    inner_control = PSY.get_inner_control(dynamic_device)
    filter = PSY.get_filter(dynamic_device)
    kpv = PSY.get_kpv(inner_control)
    kiv = PSY.get_kiv(inner_control)
    kffi = PSY.get_kffi(inner_control)
    cf = PSY.get_cf(filter)
    rv = PSY.get_rv(inner_control)
    lv = PSY.get_lv(inner_control)

    #Get Current Controller parameters
    kpc = PSY.get_kpc(inner_control)
    kic = PSY.get_kic(inner_control)
    kffv = PSY.get_kffv(inner_control)
    lf = PSY.get_lf(filter)
    ωad = PSY.get_ωad(inner_control)
    kad = PSY.get_kad(inner_control)

    function f!(out, x)
        θ_oc = x[1]
        v_refr = x[2]
        ξ_d = x[3]
        ξ_q = x[4]
        γ_d = x[5]
        γ_q = x[6]
        ϕ_d = x[7]
        ϕ_q = x[8]

        #Reference Frame Transformations
        I_dq_filter = ri_dq(θ_oc + pi / 2) * [Ir_filter; Ii_filter]
        I_dq_cnv = ri_dq(θ_oc + pi / 2) * [Ir_cnv; Ii_cnv]
        V_dq_filter = ri_dq(θ_oc + pi / 2) * [Vr_filter; Vi_filter]
        V_dq_cnv0 = ri_dq(θ_oc + pi / 2) * [Vr_cnv0; Vi_cnv0]

        #Voltage controller references
        Vd_filter_ref = (v_refr - rv * I_dq_filter[d] + ω_oc * lv * I_dq_filter[q])
        Vq_filter_ref = (-rv * I_dq_filter[q] - ω_oc * lv * I_dq_filter[d])

        #Current controller references
        Id_cnv_ref = (
            kpv * (Vd_filter_ref - V_dq_filter[d]) + kiv * ξ_d -
            cf * ω_oc * V_dq_filter[q] + kffi * I_dq_filter[d]
        )
        Iq_cnv_ref = (
            kpv * (Vq_filter_ref - V_dq_filter[q]) +
            kiv * ξ_q +
            cf * ω_oc * V_dq_filter[d] +
            kffi * I_dq_filter[q]
        )

        #References for Converter Output Voltage
        Vd_cnv_ref = (
            kpc * (Id_cnv_ref - I_dq_cnv[d]) + kic * γ_d - ω_oc * lf * I_dq_cnv[q] +
            kffv * V_dq_filter[d] - kad * (V_dq_filter[d] - ϕ_d)
        )
        Vq_cnv_ref = (
            kpc * (Iq_cnv_ref - I_dq_cnv[q]) +
            kic * γ_q +
            ω_oc * lf * I_dq_cnv[d] +
            kffv * V_dq_filter[q] - kad * (V_dq_filter[q] - ϕ_q)
        )

        out[1] = Vd_filter_ref - V_dq_filter[d]
        out[2] = Vq_filter_ref - V_dq_filter[q]
        out[3] = Id_cnv_ref - I_dq_cnv[d]
        out[4] = Iq_cnv_ref - I_dq_cnv[q]
        out[5] = V_dq_filter[d] - ϕ_d
        out[6] = V_dq_filter[q] - ϕ_q
        out[7] = Vd_cnv_ref - V_dq_cnv0[d]
        out[8] = Vq_cnv_ref - V_dq_cnv0[q]
    end
    x0 = [θ0_oc, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    sol = NLsolve.nlsolve(f!, x0)
    if !NLsolve.converged(sol)
        @warn("Initialization in Inner Control failed")
    else
        sol_x0 = sol.zero
        #Update angle:
        get_inner_vars(dynamic_device)[θ_oc_var] = sol_x0[1]
        outer_ix = get_local_state_ix(dynamic_device, O)
        outer_states = @view device_states[outer_ix]
        #Assumes that angle is in second position
        outer_states[2] = sol_x0[1]
        #Update V_ref (#TODO)
        PSY.get_ext(dynamic_device)[CONTROL_REFS][V_ref_index] = sol_x0[2]
        PSY.set_V_ref!(
            PSY.get_reactive_power(PSY.get_outer_control(dynamic_device)),
            sol_x0[2],
        )
        get_inner_vars(dynamic_device)[V_oc_var] = sol_x0[2]
        #Update Converter modulation
        m0_dq = (ri_dq(sol_x0[1] + pi / 2) * [Vr_cnv0; Vi_cnv0]) ./ Vdc
        get_inner_vars(dynamic_device)[md_var] = m0_dq[d]
        get_inner_vars(dynamic_device)[mq_var] = m0_dq[q]
        #Update states
        inner_ix = get_local_state_ix(dynamic_device, PSY.VoltageModeControl)
        inner_states = @view device_states[inner_ix]
        inner_states[1] = sol_x0[3] #ξ_d
        inner_states[2] = sol_x0[4] #ξ_q
        inner_states[3] = sol_x0[5] #γ_d
        inner_states[4] = sol_x0[6] #γ_q
        inner_states[5] = sol_x0[7] #ϕ_d
        inner_states[6] = sol_x0[8] #ϕ_q
    end
end

function initialize_inner!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::PSY.DynamicInverter{C, O, PSY.CurrentModeControl, DC, P, F},
) where {
    C <: PSY.Converter,
    O <: PSY.OuterControl,
    DC <: PSY.DCSource,
    P <: PSY.FrequencyEstimator,
    F <: PSY.Filter,
}

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.CurrentModeControl)
    # Ir_filter = device_states[external_ix[1]]
    # Ii_filter = device_states[external_ix[2]]
    Ir_cnv = device_states[external_ix[3]]
    Ii_cnv = device_states[external_ix[4]]
    Vr_filter = device_states[external_ix[5]] #TODO: Should be inner reference after initialization
    Vi_filter = device_states[external_ix[6]] #TODO: Should be inner reference after initialization

    #Obtain inner variables for component
    ω_oc = PSY.get_ω_ref(dynamic_device)
    θ0_oc = get_inner_vars(dynamic_device)[θ_freq_estimator_var]
    Vdc = get_inner_vars(dynamic_device)[Vdc_var]
    Id_cnv_ref = get_inner_vars(dynamic_device)[Id_oc_var]
    Iq_cnv_ref = get_inner_vars(dynamic_device)[Iq_oc_var]

    #Obtain output of converter
    Vr_cnv0 = get_inner_vars(dynamic_device)[Vr_cnv_var]
    Vi_cnv0 = get_inner_vars(dynamic_device)[Vi_cnv_var]

    #Get Current Controller parameters
    inner_control = PSY.get_inner_control(dynamic_device)
    filter = PSY.get_filter(dynamic_device)
    kpc = PSY.get_kpc(inner_control)
    kic = PSY.get_kic(inner_control)
    kffv = PSY.get_kffv(inner_control)
    lf = PSY.get_lf(filter)

    function f!(out, x)
        γ_d = x[1]
        γ_q = x[2]

        #Reference Frame Transformations
        I_dq_cnv = ri_dq(θ0_oc + pi / 2) * [Ir_cnv; Ii_cnv]
        V_dq_filter = ri_dq(θ0_oc + pi / 2) * [Vr_filter; Vi_filter]
        V_dq_cnv0 = ri_dq(θ0_oc + pi / 2) * [Vr_cnv0; Vi_cnv0]

        #References for Converter Output Voltage
        Vd_cnv_ref = (
            kpc * (Id_cnv_ref - I_dq_cnv[d]) + kic * γ_d - ω_oc * lf * I_dq_cnv[q] +
            kffv * V_dq_filter[d]
        )
        Vq_cnv_ref = (
            kpc * (Iq_cnv_ref - I_dq_cnv[q]) +
            kic * γ_q +
            ω_oc * lf * I_dq_cnv[d] +
            kffv * V_dq_filter[q]
        )

        out[1] = Vd_cnv_ref - V_dq_cnv0[d]
        out[2] = Vq_cnv_ref - V_dq_cnv0[q]
    end
    x0 = [0.0, 0.0]
    sol = NLsolve.nlsolve(f!, x0)
    if !NLsolve.converged(sol)
        @warn("Initialization in Inner Control failed")
    else
        sol_x0 = sol.zero
        #Update Converter modulation
        m0_dq = (ri_dq(θ0_oc + pi / 2) * [Vr_cnv0; Vi_cnv0]) ./ Vdc
        get_inner_vars(dynamic_device)[md_var] = m0_dq[d]
        get_inner_vars(dynamic_device)[mq_var] = m0_dq[q]
        #Update states
        inner_ix = get_local_state_ix(dynamic_device, PSY.CurrentModeControl)
        inner_states = @view device_states[inner_ix]
        inner_states[1] = sol_x0[1] #γ_d
        inner_states[2] = sol_x0[2] #γ_q
    end
end


text:
@testset "Expectation" begin
    Mod.@polyvar x[1:3]
    p = x[3] - 2x[1]*x[2]^2 + 3x[3]*x[1] - 5x[1]^3
    v = (1,2,3)
    m = dirac(monomials(p), x=>v)
    @test MultivariateMoments.expectation(m, p) == p(x=>v) == MultivariateMoments.expectation(p, m)
    @test_throws ErrorException dot(x[1] * x[2] * x[3], m)
    @test dot(0.5 * x[1] * x[2]^2, m) == 2.0
    @test dot(m, x[1] * x[3]) == 3

    @testset "Expectation - MB" begin
        Mod.@polyvar x[1:2]
        p = x[2] - 2x[1]*x[2]^2 + 3x[2]*x[1] - 5x[1]^3
        m = measure([1, 0, 2, 3, 0, 4, 0, 2, 5, 0], maxdegree_basis(ChebyshevBasis, x, 3))
        @test MultivariateMoments.expectation(m, p) == MultivariateMoments.expectation(p, m) == 4.25
        @test_throws ErrorException dot((x[1] * x[2])^2, m)
        @test dot(0.5 * x[1] * x[2]^2, m) == 1.0
        @test dot(m, x[1] * x[2]) == 4
    end

end


text:
"""
A transient version of the `Gridap` `FEOperator` that depends on time
"""
abstract type TransientFEOperator{C<:OperatorType} <: GridapType end

"""
Returns the test space
"""
function get_test(op::TransientFEOperator)
  @abstractmethod
end

"""
Returns the (possibly) time-dependent trial space
"""
function get_trial(op::TransientFEOperator)
  @abstractmethod # time dependent
end

function allocate_residual(op::TransientFEOperator,uh,cache)
  @abstractmethod
end

function allocate_jacobian(op::TransientFEOperator,uh,cache)
  @notimplemented
end

"""
Idem as `residual!` of `ODEOperator`
"""
function residual!(b::AbstractVector,op::TransientFEOperator,t,uh,uht,cache)
  @abstractmethod
end

"""
Idem as `residual!` of `ODEOperator`
"""
function jacobian!(A::AbstractMatrix,op::TransientFEOperator,t,uh,uht,cache)
  @abstractmethod
end

"""
Idem as `jacobian_t!` of `ODEOperator`
"""
function jacobian_t!(A::AbstractMatrix,op::TransientFEOperator,t,uh,uht,duht_du,cache)
  @abstractmethod
end

"""
Idem as `jacobian_and_jacobian_t!` of `ODEOperator`
"""
function jacobian_and_jacobian_t!(A::AbstractMatrix,op::TransientFEOperator,t,uh,uht,duht_du,cache)
  @abstractmethod
end

"""
Returns the assembler, which is constant for all time steps for a given FE
operator.

Note: adaptive FE spaces involve to generate new FE spaces and
corresponding operators, due to the ummutable approach in `Gridap`
"""
get_assembler(feop::TransientFEOperator) = @abstractmethod

# Default API

"""
Returns a `ODEOperator` wrapper of the `TransientFEOperator` that can be
straightforwardly used with the `ODETools` module.
"""
function get_algebraic_operator(feop::TransientFEOperator{C}) where C
  ODEOpFromFEOp{C}(feop)
end

OperatorType(::Type{<:TransientFEOperator{C}}) where C = C

# @fverdugo This function is just in case we need to override it in the future for some specialization.
# This default implementation is enough for the moment.
function allocate_cache(op::TransientFEOperator)
  nothing
end

function update_cache!(cache::Nothing,op::TransientFEOperator,t::Real)
  nothing
end

# Specializations

"""
Transient FE operator that is defined by a transient Weak form
"""
struct TransientFEOperatorFromWeakForm{C} <: TransientFEOperator{C}
  res::Function
  jac::Function
  jac_t::Function
  assem_t::Assembler
  trial
  trial_t
  test::FESpace
end

function TransientConstantFEOperator(m::Function,a::Function,b::Function,
  trial,test)
  res(t,u,ut,v) = m(ut,v) + a(u,v) - b(v)
  jac(t,u,ut,du,v) = a(du,v)
  jac_t(t,u,ut,dut,v) = m(dut,v)
  assem_t = SparseMatrixAssembler(trial,test)
  TransientFEOperatorFromWeakForm{Constant}(res,jac,jac_t,assem_t,trial,∂t(trial),test)
end
function TransientConstantFEOperator(res::Function,trial,test)
  # function with autodiff not implemented
  @notimplemented()
end

function TransientAffineFEOperator(m::Function,a::Function,b::Function,
  trial,test)
  res(t,u,ut,v) = m(t,ut,v) + a(t,u,v) - b(t,v)
  jac(t,u,ut,du,v) = a(t,du,v)
  jac_t(t,u,ut,dut,v) = m(t,dut,v)
  assem_t = SparseMatrixAssembler(trial,test)
  TransientFEOperatorFromWeakForm{Affine}(res,jac,jac_t,assem_t,trial,∂t(trial),test)
end
function TransientAffineFEOperator(res::Function,trial,test)
  # function with autodiff not implemented
  @notimplemented()
end

function TransientFEOperator(res::Function,jac::Function,jac_t::Function,
  trial,test)
  assem_t = SparseMatrixAssembler(trial,test)
  TransientFEOperatorFromWeakForm{Nonlinear}(res,jac,jac_t,assem_t,trial,∂t(trial),test)
end
function TransientFEOperator(res::Function,trial,test)
  # function with autodiff not implemented
  @notimplemented()
end

function SparseMatrixAssembler(
  trial::Union{TransientTrialFESpace,TransientMultiFieldTrialFESpace},
  test::FESpace)
  SparseMatrixAssembler(evaluate(trial,nothing),test)
end

get_assembler(op::TransientFEOperatorFromWeakForm) = op.assem_t

get_test(op::TransientFEOperatorFromWeakForm) = op.test

get_trial(op::TransientFEOperatorFromWeakForm) = op.trial

function allocate_residual(op::TransientFEOperatorFromWeakForm,uh::FEFunction,cache)
  v = get_cell_shapefuns(get_test(op))
  vecdata = collect_cell_vector(op.res(0.0,uh,uh,v))
  allocate_vector(op.assem_t,vecdata)
end

function residual!(b::AbstractVector,op::TransientFEOperatorFromWeakForm,
  t::Real,uh::FEFunction,uh_t::FEFunction,cache)
  v = get_cell_shapefuns(get_test(op))
  vecdata = collect_cell_vector(op.res(t,uh,uh_t,v))
  assemble_vector!(b,op.assem_t,vecdata)
  b
end

function allocate_jacobian(op::TransientFEOperatorFromWeakForm,uh::FEFunction,cache)
  matdata = matdata_jacobian(op,0.0,uh,uh)
  allocate_matrix(op.assem_t,matdata)
end

function jacobian!(A::AbstractMatrix,op::TransientFEOperatorFromWeakForm,
  t::Real,uh::FEFunction,uh_t::FEFunction,cache)
  matdata = matdata_jacobian(op,t,uh,uh_t)
  assemble_matrix_add!(A,op.assem_t, matdata)
  A
end

function jacobian_t!(A::AbstractMatrix,op::TransientFEOperatorFromWeakForm,
  t::Real,uh,uh_t,duht_du::Real,cache)
  matdata = matdata_jacobian_t(op,t,uh,uh_t,duht_du)
  assemble_matrix_add!(A,op.assem_t, matdata)
  A
end

function jacobian_and_jacobian_t!(A::AbstractMatrix,op::TransientFEOperatorFromWeakForm,
  t::Real,uh,uh_t,duht_du::Real,cache)
  matdata_j = matdata_jacobian(op,t,uh,uh_t)
  matdata_jt = matdata_jacobian_t(op,t,uh,uh_t,duht_du)
  matdata = vcat_matdata(matdata_j,matdata_jt)
  assemble_matrix_add!(A,op.assem_t, matdata)
  A
end

function vcat_matdata(matdata_j,matdata_jt)
  term_to_cellmat_j, term_to_cellidsrows_j, term_to_cellidscols_j = matdata_j
  term_to_cellmat_jt, term_to_cellidsrows_jt, term_to_cellidscols_jt = matdata_jt

  term_to_cellmat = vcat(term_to_cellmat_j,term_to_cellmat_jt)
  term_to_cellidsrows = vcat(term_to_cellidsrows_j,term_to_cellidsrows_jt)
  term_to_cellidscols = vcat(term_to_cellidscols_j,term_to_cellidscols_jt)

  matdata = (term_to_cellmat,term_to_cellidsrows, term_to_cellidscols)
end

function matdata_jacobian(op::TransientFEOperatorFromWeakForm,t::Real,uh::FEFunction,uh_t::FEFunction)
  Uh = evaluate(get_trial(op),nothing)
  du = get_cell_shapefuns_trial(Uh)
  v = get_cell_shapefuns(get_test(op))
  matdata = collect_cell_matrix(op.jac(t,uh,uh_t,du,v))
end

function matdata_jacobian_t(op::TransientFEOperatorFromWeakForm,t::Real,uh::FEFunction,uh_t::FEFunction,duht_du::Real)
  Uh = evaluate(get_trial(op),nothing)
  du_t = get_cell_shapefuns_trial(Uh)
  v = get_cell_shapefuns(get_test(op))
  matdata = collect_cell_matrix(duht_du*op.jac_t(t,uh,uh_t,du_t,v))
end

# Tester

function test_transient_fe_operator(op::TransientFEOperator,uh)
  odeop = get_algebraic_operator(op)
  @test isa(odeop,ODEOperator)
  cache = allocate_cache(op)
  V = get_test(op)
  @test isa(V,FESpace)
  U = get_trial(op)
  U0 = U(0.0)
  @test isa(U0,FESpace)
  r = allocate_residual(op,uh,cache)
  @test isa(r,AbstractVector)
  residual!(r,op,0.0,uh,uh,cache)
  @test isa(r,AbstractVector)
  J = allocate_jacobian(op,uh,cache)
  @test isa(J,AbstractMatrix)
  jacobian!(J,op,0.0,uh,uh,cache)
  @test isa(J,AbstractMatrix)
  jacobian_t!(J,op,0.0,uh,uh,1.0,cache)
  @test isa(J,AbstractMatrix)
  jacobian_and_jacobian_t!(J,op,0.0,uh,uh,1.0,cache)
  @test isa(J,AbstractMatrix)
  cache = update_cache!(cache,op,0.0)
  true
end


text:
# Utility functions

"""
    seconds_milliseconds(x) -> ::Dates.CompoundPeriod

Convert a real value `x` into a `Dates.CompoundPeriod` of seconds and
milliseconds, truncating to the nearest millisecond.

# Example
julia> using SeisRequests

julia> SeisRequests.seconds_milliseconds(5.1019)
5 seconds, 101 milliseconds
"""
function seconds_milliseconds(x)
    s, ms = divrem(x, 1)
    ms *= 1000
    Second(floor(Int, s)) + Millisecond(floor(Int, ms))
end

"""
    split_channel_code(code) -> (network, station, location, channel)

Split up a single string `code` giving a channel code in the form
`"⟨network⟩.⟨station⟩.⟨channel⟩.⟨location⟩"` into its component parts,
and return a named tuple.
"""
function split_channel_code(code)
    tokens = split(code, '.')
    length(tokens) == 4 ||
        throw(ArgumentError("incorrect number of fields in code \"$code\""))
    (network=tokens[1], station=tokens[2], location=tokens[3], channel=tokens[4])
end


text:
## Market risk calculation (not product-specific)

function calculate_scenario_pls(pvs:: AbstractVector{<: CurrencyAmount})
    base_pv = last(pvs)
    scenario_pvs = (@view pvs[1:end-1]) .- base_pv
    return scenario_pvs
end

function calculate_scenario_pls(product:: AbstractProduct, scenarios:: AbstractMarketDataContainer)
    scenario_pvs = price_time_series(product, scenarios)
    scenario_pls = calculate_scenario_pls(scenario_pvs)
    return scenario_pls
end

function value_at_risk(scenario_pls:: AbstractVector, quantile_value:: Real; alpha=1.0, beta=alpha)
    0.0 < quantile_value < 1.0 || error("quantile must be between 0 and 1")
    var = -quantile(scenario_pls, one(quantile_value)-quantile_value; alpha, beta)
    return var
end

function value_at_risk(product:: AbstractProduct, container:: FXForwardMarketDataContainer, valuation_date:: Date;
    quantile_value:: Number=0.99, time_horizon:: Integer = 1, lookback:: Integer = 250)

    scenarios = create_historical_scenarios(container; valuation_date, time_horizon, lookback)
    scenario_pls = calculate_scenario_pls(product, scenarios)
    var = value_at_risk(scenario_pls, quantile_value)
    return var
end

function value_at_risk(products:: AbstractVector{<: AbstractProduct}, container:: FXForwardMarketDataContainer, spot_rates:: Dict, valuation_date:: Date;
    quantile_value:: Number=0.99, time_horizon:: Integer = 1, lookback:: Integer = 250, target_currency:: Currency = Currency{:EUR}())
    scenarios = create_historical_scenarios(container; valuation_date, time_horizon, lookback)

    scenario_pls = calculate_scenario_pls.(products, scenarios)
    scenario_pls_all = sum(convert.(target_currency, x, Ref(values(spot_rates))) for x in scenario_pls)

    var = value_at_risk(scenario_pls_all, quantile_value)
    return var
end

max_issues_repo_name: lungben/MarketRiskPrototype

max_forks_repo_name: lungben/MarketRiskPrototype


text:
module MovcolN

using Polynomials
using Devectorize
using DASSL
using ArrayViews
using Iterators

include("polynomials.jl")       # polynomial related functions

export Equation, movcol_solve

# Equation should consist of functions F,G,Bl,Br,M,g,tau,Bxl,Bxr,gamma
abstract Equation

immutable TemporaryArrays
    Hul :: Vector{Float64}
    Hur :: Vector{Float64}
    res :: Vector{Float64}
    ux  :: Matrix{Float64}
    utx :: Matrix{Float64}
end

function TemporaryArrays(nd,nu)
    TemporaryArrays(Array(Float64,nd),
                    Array(Float64,nd),
                    Array(Float64,nd+1),
                    Array(Float64,nd+1,nu),
                    Array(Float64,nd+1,nu))
end

immutable CollocationPoint
    left  :: Matrix{Float64}
    right :: Matrix{Float64}
    s     :: Float64
end

immutable MeshPoint
    x  :: Float64
    xt :: Float64
    u  :: AbstractMatrix{Float64}
    ut :: AbstractMatrix{Float64}
    nd :: Int
    nu :: Int
end

# CollocationData is the same for every mesh point and it depends on
# the number of collocation points (ns) only.
immutable CollocationData # only works with Float64
    ns :: Int        # number of derivatives provided at mesh points

    # vectors of tuples (Qleft, Qright, s)
    gauss   :: Vector{CollocationPoint}
    lobatto :: Vector{CollocationPoint}

    AB :: Matrix{Float64}
end

# @todo tests...
function CollocationData(ns)
    QsGauss   = map(s->generateQats(ns,s),  gauss(ns)[1])
    QsLobatto = map(s->generateQats(ns,s),lobatto(ns+1)[1])
    AB = generateAB(ns)
    return CollocationData(ns,QsGauss,QsLobatto,AB)
end

# create matrix AB=inv(A)*B with A and B such that A*Fi = B*Gi/Hi
# (i.e. approximating F=∂ₓG), assuming that Fi is evaluated at Gauss
# points and Gi is evaluated at Lobatto points.  The dimensions of A
# and B are n×n and n×n+1 respectively.
function generateAB(n; T=Float64)
    sg, _ = gauss(n)
    sl, _ = lobatto(n+1)
    x    = Poly(T[0,1])
    Lint = polyint(Lagrange(sg,x))
    # A = ∫Li dx over the interval [sl[j],sl[j+1]]
    A  = -T[polyval(Lint[k],sl[j+1])-polyval(Lint[k],sl[j])
           for j=1:n, k=1:n]
    B  = spdiagm((fill(1,n),fill(-1,n)),(0,1))
    AB = inv(A)*B
    return AB
end

# generates an array of matrices Q such that
# HiᵏUₓ⁽ᵏ⁾(xi)=Q[k+1]*(H.*U12), where U12 is a vector
# [u1,u1x,u1xx,...,u2,u2x,u2xx,...] of length 2(n+1),
# H=[1,Hi,Hi^2,...,Hi^n,1,Hi,Hi^2,...,Hi^n] and xi are fixed points
# (in practice, these are either Gauss or Lobatto points)
#
# @todo this wont work for integer xi, because it will try to convert
# Rational into integer
function generateQats{T}(n,s::T;kmax=n+1)
    L = hermite01(n)
    Qleft  = T[polyval(polyder(L[1,j],i),s) for i=0:kmax-1, j=1:n]
    Qright = T[polyval(polyder(L[2,j],i),s) for i=0:kmax-1, j=1:n]
    return CollocationPoint(Qleft,Qright,s)
end

# given the vectors ul=[u,ux,uxx,...]_(x=x[N]) and
# ur=[u,ux,uxx,...]_(x=x[N+1]) generate a vector
# ux=[u,ux,uxx,...]_(x=x[N]+s*h), where s is prescribed in Qs (this
# descriptions of ul, ur and ux are oversimplification as each of them
# holds the derivatives of more than just one function,
# e.g. ul[j,i]=∂ⱼuᵢ(x[N]) ).  Generation of ux is realized by
# performing linear operations on ul and ur via the matrices Qleft and
# Qright.  The vector H is equal to [h^j for j=0:nd+1].
function computeux!{T}(ux::Array{T,2},Q::CollocationPoint,H::Vector{T},ul::AbstractMatrix{T},ur::AbstractMatrix{T},tmp)
    nd = size(ul,1)
    nu = size(ul,2)

    for i=1:nu
        @devec tmp.Hul[:] = H[1:nd].*ul[:,i]
        @devec tmp.Hur[:] = H[1:nd].*ur[:,i]
        mult_lr!(tmp.res,Q.left,Q.right,tmp.Hul,tmp.Hur)
        @devec ux[:,i] = tmp.res./H
    end
end

# Function similar to computeux! but computing
# utx=[ut,utx,utxx,...]. It requires more data (Ht,xt,utl,utr,ux)
# because it has to account for the mesh translation and rescaling in
# time.
function computeutx!{T}(utx::Array{T,2},Q::CollocationPoint,H::Vector{T},Ht::Vector{T},xt::T,utl,utr,ux,ul,ur,tmp)
    nd = size(utl,1)
    nu = size(utl,2)
    h  = H[2]
    ht = Ht[2]
    xt = (xt+Q.s*ht)

    for i = 1:nu
        @devec tmp.Hul[:] = H[1:nd].*utl[:,i]+Ht[1:nd].*ul[:,i]
        @devec tmp.Hur[:] = H[1:nd].*utr[:,i]+Ht[1:nd].*ur[:,i]
        mult_lr!(tmp.res,Q.left,Q.right,tmp.Hul,tmp.Hur)
        # upscale spatial derivatives
        scaling_utx!(utx,tmp.res,xt,ux,H,Ht,i)
    end
end

function mult_lr!(res,Qleft,Qright,Hul,Hur)
    nd = length(Hul)
    for i = 1:nd+1
        res[i] = 0
        for j = 1:nd
            @inbounds res[i] += Hul[j]*Qleft[i,j]+Hur[j]*Qright[i,j]
        end
    end
end

# correct for the mesh rescaling, dilation and shift of the mesh
function scaling_utx!(utx,res,xt,ux,H,Ht,i)
    for j = 1:length(res)-1
        @inbounds utx[j,i] = (res[j]-Ht[j]*ux[j,i])/H[j]-xt*ux[j+1,i]
    end
end

# For a given function F(t,u,ux,uxx,...,ut,utx,utxx,...) return the
# value F_(x=x[N]+s*h).  The spatial and temporal derivatives at
# x=x[N]+s*h are computed via computeux! and computeutx! functions.
function getcollocationvalues!{T}(FVal,Qs::Vector{CollocationPoint},F,t::T,left,right,H,Ht,tmp)
    h = H[2]

    for j = 1:size(FVal,1)
        s = Qs[j].s

        if s == zero(T)
            F(view(FVal,j,:),t,left.x,left.u,left.ut)
        elseif s == zero(T)
            F(view(FVal,j,:),t,right.x,right.u,right.ut)
        else
            x = left.x; xt = left.xt
            xs  = x+s*h
             computeux!(tmp.ux, Qs[j],H,left.u,right.u,tmp)
            computeutx!(tmp.utx,Qs[j],H,Ht,xt,left.ut,right.ut,tmp.ux,left.u,right.u,tmp)
            F(view(FVal,j,:),t,xs,tmp.ux,tmp.utx)

        end

    end
end

# memory alignment
#
#      derivative (nd)
#     /   point (nx)
#    /   /
# u[i,j,k]
#      \
#       function (nu)
#

function computeresu{T}(pde :: Equation,
                        coldata :: CollocationData,
                        t  :: T,
                        x  :: Vector{T}, xt :: Vector{T},
                        u  :: Array{T,3},ut :: Array{T,3};
                        args...)
    nd = size(u,1)
    nu = size(u,2)
    nx = size(u,3)
    ns = nd
    ns2 = div(ns,2)

    FGauss   = Array(T,ns  ,nu)
    GLobatto = Array(T,ns+1,nu)
    resu     = Array(T,ns  ,nu,nx)
    tmp      = TemporaryArrays(nd,nu)
    H        = Array(T,ns+1)
    Ht       = Array(T,ns+1)

    for i = 1:nx-1

        left  = MeshPoint(x[i  ],xt[i  ],view(u,:,:,i  ),view(ut,:,:,i  ), nd, nu)
        right = MeshPoint(x[i+1],xt[i+1],view(u,:,:,i+1),view(ut,:,:,i+1), nd, nu)
        h = x[i+1]-x[i];      ht = xt[i+1]-xt[i]

        for j = 0:nd
            H[j+1]  = h^j
            Ht[j+1] = j*ht*h^(j-1)
        end

        getcollocationvalues!(FGauss,  coldata.gauss,  pde.F,t,left,right,H,Ht,tmp)
        getcollocationvalues!(GLobatto,coldata.lobatto,pde.G,t,left,right,H,Ht,tmp)

        # equivalent to
        #
        # resu[:,:,i] = FGauss-coldata.AB*GLobatto/h
        #
        # for small AB the fillresu! is significantly faster than the
        # build-in matrix multiplication
        fillresu!(resu,FGauss,GLobatto,coldata.AB,h,i)

    end

    # residua for the boundary conditions, Bl and Br should each
    # return the vector of size at least nu*ns/2, the remaining
    # boundary conditions are discarded
    resBl = pde.Bl(t,x[1  ],xt[1  ],view(u,:,:,1 ),view(ut,:,:,1 ))::Array{Float64,1}
    resBr = pde.Br(t,x[end],xt[end],view(u,:,:,nx),view(ut,:,:,nx))::Array{Float64,1}
    resu[    1:ns2,:,nx] = reshape(resBl,nu,ns2)
    resu[ns2+1:ns, :,nx] = reshape(resBr,nu,ns2)

    return resu

end

function fillresu!(resu,F,G,AB,h,i)
    nG = size(G,1)
    nd = size(resu,1)
    nu = size(resu,2)
    for l = 1:nu
        for k = 1:nd
            resu[k,l,i] = F[k,l]
        end
    end
    for j = 1:nG
        for l = 1:nu
            for k = 1:nd
                resu[k,l,i] -= AB[k,j]*G[j,l]/h
            end
        end
    end
end

# compute residua for the mesh
function computeresx(pde :: Equation,
                     coldata :: CollocationData,
                     t :: Float64,
                     x :: Vector{Float64}, xt :: Vector{Float64},
                     u :: Array{Float64,3},ut :: Array{Float64,3};
                     args...)

    nd,nu,nx = size(u)
    ns = nd                     # number of collocation points
    ns2 = div(ns,2)

    uhalf    = Array(Float64,ns+1,nu)
    uthalf   = Array(Float64,ns+1,nu)
    resx     = Array(Float64,nx)
    Mhalf    = Array(Float64,nx-1)
    Yhalf    = Array(Float64,nx+1)
    tmp      = TemporaryArrays(nd,nu)

    Qshalf = coldata.lobatto[ns2+1]

    for i = 1:nx-1
        xl = x[i  ]; xtl = xt[i  ]; ul = view(u,:,:,i  ); utl = view(ut,:,:,i  )
        xr = x[i+1]; xtr = xt[i+1]; ur = view(u,:,:,i+1); utr = view(ut,:,:,i+1)
        h  = x[i+1]-x[i];      ht = xt[i+1]-xt[i]
        H  = [h^j for j=0:nd]
        Ht = [j*ht*h^(j-1) for j=0:nd]

        if i==1 || i==nx-1
            Mhalf[i] = ((pde.M(t,xl,ul)::Float64+pde.M(t,xr,ur)::Float64)/2)
        else
            xhalf  = (xl+xr)/2
            xthalf = (xtl+xtr)/2
             computeux!(uhalf, Qshalf,H,ul,ur,tmp)
            computeutx!(uthalf,Qshalf,H,Ht,xthalf,utl,utr,uhalf,ul,ur,tmp)
            Mhalf[i] = pde.M(t,xhalf,uhalf)::Float64
        end

        Yhalf[i+1] = -pde.tau*(xtr-xtl)/(xr-xl)^2+1/(xr-xl)

        if i == 1
            resx[ 1] = pde.Bxl(t,xl,xtl,ul,utl)::Float64
        end
        if i == nx-1
            resx[nx] = pde.Bxr(t,xr,xtr,ur,utr)::Float64
        end
    end

    # boundary conditions for Yhalf
    Yhalf[ 1 ] = Yhalf[  2  ]
    Yhalf[end] = Yhalf[end-1]

    Mhalf = smoothen(Mhalf;args...)

    # compute the residua for mesh movement
    for i = 2:nx-1
        g=pde.gamma*(pde.gamma+1)
        # mmpde6?
        # resx[i] = (Yhalf[i+1]-g*(Yhalf[i+2]-2*Yhalf[i+1]+Yhalf[i]))/Mhalf[i]-(Yhalf[i]-g*(Yhalf[i+1]-2*Yhalf[i]+Yhalf[i-1]))/Mhalf[i-1]
        # mmpde4?
        resx[i] = pde.tau*(xt[i+1]-2*xt[i]+xt[i-1])+Mhalf[i]*(x[i+1]-x[i])-Mhalf[i-1]*(x[i]-x[i-1])
    end

    return resx

end

# simple gaussian filter to smoothen the monitor function
function smoothen(M; passes = 4, args...)
    if passes == 0
        return M
    else
        Msmooth = zero(M)
        @simd for i = 2:length(M)-1
            @inbounds Msmooth[i]=(M[i-1]+2M[i]+M[i+1])/4
        end
        Msmooth[1]=(M[1]+M[2])/2
        Msmooth[end]=(M[end-1]+M[end])/2
        return smoothen(Msmooth, passes=passes-1)
    end
end

# creates a relaxed mesh
function meshinit(pde :: Equation,
                  coldata :: CollocationData,
                  nx :: Int;
                  mesherr = 1e-5,
                  mesh_maxsteps = Inf,
                  args...)

    ns = coldata.ns
    x0 = linspace(pde.xspan...,nx)
    x0t = zero(x0)

    # Create a new pde with fixed mesh endpoints
    xpde = deepcopy(pde)
    xpde.Bxl=(t,x,xt,u,ut)->xt
    xpde.Bxr=(t,x,xt,u,ut)->xt

    # wrapper function for DASSL
    function dae(t,x,xt)
        u    = reshape(vcat(map(xpde.u0, x)...),ns*xpde.nu,1,nx)
        resx = computeresx(xpde,coldata,xpde.t0,x,xt,u,u;args...)
        return resx
    end

    nstp = 0
    for (t,x,xt) in dasslIterator(dae,x0,0.0;reltol=mesherr*1e-2,abstol=mesherr*1e-2)
        nstp += 1

        if norm(xt,Inf) < mesherr
            info("Mesh converged after $nstp steps")
            x0  = x
            break
        elseif nstp > mesh_maxsteps
            error("Unable to converge in $maxsteps steps")
            break
        end
    end

    return x0

end

function movcol_solve(pde :: Equation,
                      nx :: Int,
                      nd :: Int;
                      args...)

    coldata = CollocationData(nd)
    if isodd(nd); error("nd should be even"); end
    if nx < 2; error("nx should be at least 2"); end
    nu = pde.nu
    x0 = meshinit(pde,coldata,nx;args...)
    u0 = reshape(vcat(map(pde.u0, x0)...),nd,nu,nx)

    # wrapper function for DASSL
    function dae(t,y,yt)
        res = Array(Float64,nd*nu*nx+nx)
        # @todo xt = g(t)*xt, ut = g(t)*ut and make tau=tau(t)
         x =  y[1:nx];  u = reshape( y[nx+1:end],nd,nu,nx)
        xt = yt[1:nx]; ut = reshape(yt[nx+1:end],nd,nu,nx)
        resu = computeresu(pde,coldata,t,x,xt,u,ut;args...)
        resx = computeresx(pde,coldata,t,x,xt,u,ut;args...)
        res[1:nx] = resx
        res[nx+1:end] = resu
        return res
    end

    function get_results(t,y,yt)
         x =  y[1:nx];  u = reshape( y[nx+1:end],nd,nu,nx)
        xt = yt[1:nx]; ut = reshape(yt[nx+1:end],nd,nu,nx)
        return t, x, u, xt, ut
    end

    y0  = vcat(x0,vec(u0))
    sol = dasslIterator(dae,y0,pde.t0;args...)

    return imap(z->get_results(z...),sol)

end

end # module


text:
##############################################################################################

#Setting up helper functions
function countmemb(itr)
    d = Dict{String, Int}()
    for val in itr
        if isa(val, Number) && isnan(val)
            continue
        end
        d[string(val)] = get!(d, string(val), 0) + 1
    end
    return d
end

##############################################################################################
#does weighted average of V, but highest weighted particle for world state. also prints world state and weight for each particle
function print_Vs_and_Rs_to_file(V_file, ws_file, traces, num_particles::Int64, params::Video_Params, v::Int64, order::Vector{Int64}, last_time::Bool=false)
    print(V_file, v, " & ")
    print(ws_file, v, " & ")

    print(V_file, order[v], " & ")
    print(ws_file, order[v], " & ")

    avg_v, world_states, best_world_state = process(traces, num_particles, params, v)

    for i = 1:params.n_possible_objects
        print(V_file, avg_v[i,1], "&")
        print(V_file, avg_v[i,2], "&")
    end

    v_matrix = zeros(params.n_possible_objects, 2)
    for j = 1:num_particles
        choices = get_choices(traces[j])
        weight = get_score(traces[j])
        for i = 1:params.n_possible_objects
            v_matrix[i,1] = choices[:videos => v => :v_matrix => :lambda_fa => i => :fa]
            v_matrix[i,2] = choices[:videos => v => :v_matrix => :miss_rate => i => :miss]
            print(V_file, v_matrix[i,1], "&")
            print(V_file, v_matrix[i,2], "&")
        end
        print(V_file, weight, "&")
    end

    print(V_file, best_world_state)
    ###########################################
    #switch to ws file
    print(ws_file, best_world_state, "&")

    for j = 1:num_particles
        choices = get_choices(traces[j])
        world_state = choices[:videos => v => :init_scene]
        weight = get_score(traces[j])
        if j==num_particles && last_time #if this is the last thing printing, don't put &
            print(ws_file, world_state, "&")
            print(ws_file, weight)
        else
            print(ws_file, world_state, "&")
            print(ws_file, weight, "&")
        end
    end
    print(V_file, "\n")
    print(ws_file, "\n")
    return (best_world_state, avg_v) #return the mode reality and average v
end

##############################################################################################
#returns the weighted average v_matrix and the highest-weighted world state for a set particles
function process(traces, num_particles::Int64, params::Video_Params, v::Int64)

    weights = Array{Float64}(undef, num_particles)
    v_matrixes = Array{Matrix{Float64}}(undef, num_particles)
    world_states = Array{Array{Any,1}}(undef, num_particles)

    for i = 1:num_particles
        weights[i] = get_score(traces[i])
        choices = get_choices(traces[i])
        v_matrix = zeros(params.n_possible_objects, 2)
        for j = 1:params.n_possible_objects
            v_matrix[j,1] = choices[:videos => v => :v_matrix => :lambda_fa => j => :fa]
            v_matrix[j,2] = choices[:videos => v => :v_matrix => :miss_rate => j => :miss]
        end
        v_matrixes[i] = v_matrix
        world_states[i] = choices[:videos => v => :init_scene]
    end
    normalized_log_weights = weights ./ (sum(weights))
    avg_v = sum(exp.(normalized_log_weights)./sum(exp.(normalized_log_weights)) .* v_matrixes)#change out of log
    best_world_state = world_states[findmax(weights)[2]]

    return avg_v, world_states, best_world_state
end

export print_Vs_and_Rs_to_file_new
export countmemb

max_stars_repo_name: zhangir-azerbayev/MetaGen

max_issues_repo_name: zhangir-azerbayev/MetaGen

max_forks_repo_name: zhangir-azerbayev/MetaGen


text:
# Swing Option pricing for natural gas resources using LSM.

struct PriceProcess{T<:AbstractFloat}
    initialprice::T
    mean_reversion_rate::T
    volatility::T
    mean_reversion_level::T
    riskfree::T
end

struct SwingContract{I<:Integer,R<:AbstractFloat}
    settlement::Date
    maturity::Date
    strike::R
    max_rights::I
    daily_contract::I
    min_contract::I
    max_contract::I
end

struct Valuation
    order::Int
    runs::Int
end
struct SwingOption
    process::PriceProcess
    contract::SwingContract
    val::Valuation
end

function price_option(opt::SwingOption)
    # Discount Factor
    days = Dates.value(opt.contract.maturity - opt.contract.settlement)
    stepsize = 1.0 / days
    discount_rate = exp(-opt.process.riskfree * stepsize)

    # Swing inventory values in MMBtu
    swingup = opt.contract.max_contract - opt.contract.daily_contract
    swingdown = opt.contract.min_contract - opt.contract.daily_contract

    # Initialize retval with the simulated prices. Note that retval is 3D.
    retval = zeros(runs,days,opt.contract.max_rights)
    prices = simulate_price(opt)

    # Obtain price spread between strike and simulated prices
    retval .=  prices .- opt.contract.strike

    # Obtain exercise values for each run and day
    retval .= (retval .> 0.0) .* swingup .* retval .+ (retval .<= 0.0) .* swingdown .* retval

    # value iteration. Termainal value set by the exercise values above.
    @inbounds for epoch in (days-1):-1:1

        # Initialize contuation values that carry over between swing rights
        next_continue_vals = zeros(length(runs))

        for right in 1:opt.contract.max_rights
            exercise_vals = @view retval[:,epoch,right]
            discounted_vals = discount_rate * retval[:,epoch+1,right]

            # power series matrix for regression coefficients
            power_matrix = power_series(prices,opt.val.order)
            reg_coeff = power_matrix \ discounted_vals

            # Expected continuation values
            continue_vals = power_matrix * reg_coeff

            # next_continue_vals are the continuation values after exercise
            next_continue_vals = right == 1 ? zeros(length(continue_vals)) : retval[:,epoch,right-1]

            # Decision
            retval[:, epoch, right] .=

        end

    end

end

function simulate_price(opt::SwingOption)
    # Rename with locals option assumptions
    mean_rev_speed = opt.process.mean_reversion_rate
    mean_rev_level = opt.process.mean_reversion_level
    vol = opt.process.volatility
    runs = opt.val.runs

    # Initialize stepsize
    days = Dates.value(opt.contract.maturity - opt.contract.settlement)
    stepsize = 1.0 / days

    # Initialize retval
    retval = zeros(runs,days)
    initial_logprice = log(opt.process.initialprice)

    # Simulate paths along rows of log-prices
    for run in 1:runs, epoch in 1:days
        last_price = epoch == 1 ? initial_logprice : retval[run,epoch-1]
        logprice = last_price + mean_rev_speed*(mean_rev_level-last_price)*stepsize +
                    vol * sqrt(stepsize) * randn()
        retval[run,epoch] = logprice
    end
    return exp.(retval)
end