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

using Behavior
using Behavior:
    StepDefinitionMatcher,
    Executor, executescenario, @expect, SuccessfulStepExecution,
    issuccess, beforescenario, afterscenario, FromSourceExecutionEnvironment,
    FromMacroStepDefinitionMatcher, QuietRealTimePresenter
using Behavior.Gherkin.Experimental

mutable struct FakeExecutionEnvironment <: Behavior.ExecutionEnvironment
    afterscenariowasexecuted::Bool

    FakeExecutionEnvironment() = new(false)
end

function Behavior.beforescenario(::FakeExecutionEnvironment, context::StepDefinitionContext, scenario::Gherkin.Scenario)
    context[:beforescenariowasexecuted] = true
end

function Behavior.afterscenario(f::FakeExecutionEnvironment, context::StepDefinitionContext, scenario::Gherkin.Scenario)
    context[:afterscenariowasexecuted] = true
    f.afterscenariowasexecuted = true
end

struct SingleStepDefinitionMatcher <: StepDefinitionMatcher
    stepbody::Function
end

function Behavior.findstepdefinition(
        s::SingleStepDefinitionMatcher,
        step::Behavior.Gherkin.ScenarioStep)
    stepdefinition = (context, args) -> begin
        s.stepbody(context, args)
        SuccessfulStepExecution()
    end
    StepDefinitionMatch(StepDefinition("some text", stepdefinition, StepDefinitionLocation("", 0)))
end

@testset "Execution Environment" begin
    @testset "Scenario Execution Environment" begin
        @testset "beforescenario is defined; beforescenario is executed before the scenario" begin
            # Arrange
            env = FakeExecutionEnvironment()

            # This step definition tests that the symbol :beforescenariowasexecuted is present in
            # the context at execution time. This symbol should be set by the
            # FakeExecutionEnvironment.
            stepdefmatcher = SingleStepDefinitionMatcher((context, args) -> @assert context[:beforescenariowasexecuted])
            executor = Executor(stepdefmatcher; executionenv=env)
            scenario = Scenario("Description", String[], ScenarioStep[Given("")])

            # Act
            result = executescenario(executor, Background(), scenario)

            # Assert
            @test issuccess(result.steps[1])
        end

        @testset "beforescenario is the default noop; beforescenario is not executed before the scenario" begin
            # Arrange
            stepdefmatcher = SingleStepDefinitionMatcher((context, args) -> @assert !haskey(context, :beforescenariowasexecuted))
            executor = Executor(stepdefmatcher)
            scenario = Scenario("Description", String[], ScenarioStep[Given("")])

            # Act
            result = executescenario(executor, Background(), scenario)

            # Assert
            @test issuccess(result.steps[1])
        end

        @testset "afterscenario is defined; afterscenario has not been called at scenario execution" begin
            # Arrange
            env = FakeExecutionEnvironment()

            # This step definition tests that the symbol :beforescenariowasexecuted is present in
            # the context at execution time. This symbol should be set by the
            # FakeExecutionEnvironment.
            stepdefmatcher = SingleStepDefinitionMatcher((context, args) -> @assert !haskey(context, :afterscenariowasexecuted))
            executor = Executor(stepdefmatcher; executionenv=env)
            scenario = Scenario("Description", String[], ScenarioStep[Given("")])

            # Act
            result = executescenario(executor, Background(), scenario)

            # Assert
            @test issuccess(result.steps[1])
        end

        @testset "afterscenario is defined; afterscenario is called" begin
            # Arrange
            env = FakeExecutionEnvironment()

            # This step definition tests that the symbol :beforescenariowasexecuted is present in
            # the context at execution time. This symbol should be set by the
            # FakeExecutionEnvironment.
            stepdefmatcher = SingleStepDefinitionMatcher((context, args) -> nothing)
            executor = Executor(stepdefmatcher; executionenv=env)
            scenario = Scenario("Description", String[], ScenarioStep[Given("")])

            # Act
            result = executescenario(executor, Background(), scenario)

            # Assert
            @test env.afterscenariowasexecuted
        end
    end

    @testset "FromSourceExecutionEnvironment" begin
        @testset "beginscenario is defined in source; beginscenario is executed" begin
            # Arrange
            env = FromSourceExecutionEnvironment("""
                using Behavior

                @beforescenario() do context, scenario
                    context[:beforescenariowasexecuted] = true
                end
            """)

            context = StepDefinitionContext()
            scenario = Gherkin.Scenario("", String[], ScenarioStep[])

            # Act
            beforescenario(env, context, scenario)

            # Assert
            @test context[:beforescenariowasexecuted]
        end

        @testset "No beginscenario is defined in source; beginscenario is a noop" begin
            # Arrange
            env = FromSourceExecutionEnvironment("")

            context = StepDefinitionContext()
            scenario = Gherkin.Scenario("", String[], ScenarioStep[])

            # Act
            beforescenario(env, context, scenario)

            # Assert
            @test !haskey(context, :beforescenariowasexecuted)
        end

        @testset "afterscenario is defined in source; afterscenario is executed" begin
            # Arrange
            env = FromSourceExecutionEnvironment("""
                using Behavior

                @afterscenario() do context, scenario
                    context[:afterscenariowasexecuted] = true
                end
            """)

            context = StepDefinitionContext()
            scenario = Gherkin.Scenario("", String[], ScenarioStep[])

            # Act
            afterscenario(env, context, scenario)

            # Assert
            @test context[:afterscenariowasexecuted]
        end

        @testset "No afterscenario is defined in source; afterscenario is a noop" begin
            # Arrange
            env = FromSourceExecutionEnvironment("")

            context = StepDefinitionContext()
            scenario = Gherkin.Scenario("", String[], ScenarioStep[])

            # Act
            afterscenario(env, context, scenario)

            # Assert
            @test !haskey(context, :afterscenariowasexecuted)
        end
    end

    @testset "Before and After feature hooks" begin
        @testset "beforefeature saves the feature; Feature runs; Feature is found in list" begin
            # Arrange
            # Define @beforefeature
            env = FromSourceExecutionEnvironment("""
                using Behavior
                using Behavior.Gherkin

                features = Gherkin.Feature[]

                @beforefeature() do feature
                    push!(features, feature)
                end
            """)

            # Step definitions
            matcher = FromMacroStepDefinitionMatcher("""
                using Behavior

                @then("the current feature has description \\"{String}\\"") do context, description
                    currentfeature = Main.features[end]
                    @expect currentfeature.header.description == description
                end
            """)

            # Feature
            input = Experimental.ParserInput("""
                Feature: This is a feature description

                    Scenario: Check the feature description
                        Then the current feature has description "This is a feature description"
            """)
            parser = Experimental.FeatureFileParser()
            parserresult = parser(input)
            feature = parserresult.value

            executor = Executor(matcher, QuietRealTimePresenter(), executionenv=env)

            # Act
            featureresult = executefeature(executor, feature)

            # Assert
            steps = featureresult.scenarioresults[1].steps
            @test issuccess(steps[1])
        end

        @testset "afterfeature saves the feature; Feature runs; Feature is not found while running the steps" begin
            # Arrange
            # Define @afterfeature
            env = FromSourceExecutionEnvironment("""
                using Behavior
                using Behavior.Gherkin

                features = Gherkin.Feature[]

                @afterfeature() do feature
                    push!(features, feature)
                end
            """)

            # Step definitions
            matcher = FromMacroStepDefinitionMatcher("""
                using Behavior

                @then("the feature has not been added to the list") do context
                    @expect Main.features == []
                end
            """)

            # Feature
            input = Experimental.ParserInput("""
                Feature: This is a feature description

                    Scenario: Afterfeature adds the feature to a list
                        Then the feature has not been added to the list
            """)
            parser = Experimental.FeatureFileParser()
            parserresult = parser(input)
            feature = parserresult.value

            executor = Executor(matcher, QuietRealTimePresenter(), executionenv=env)

            # Act
            featureresult = executefeature(executor, feature)

            # Assert
            steps = featureresult.scenarioresults[1].steps
            @test issuccess(steps[1])
        end

        @testset "afterfeature saves the feature; Feature runs; Feature is found after execution" begin
            # Arrange
            # Define @afterfeature
            env = FromSourceExecutionEnvironment("""
                using Behavior
                using Behavior.Gherkin

                features = Gherkin.Feature[]

                @afterfeature() do feature
                    push!(features, feature)
                end
            """)

            # Step definitions
            matcher = FromMacroStepDefinitionMatcher("""
                using Behavior

                @then("the feature has not been added to the list") do context
                    @expect Main.features == []
                end
            """)

            # Feature
            input = Experimental.ParserInput("""
                Feature: This is a feature description

                    Scenario: Afterfeature adds the feature to a list
                        Then the feature has not been added to the list
            """)
            parser = Experimental.FeatureFileParser()
            parserresult = parser(input)
            feature = parserresult.value

            executor = Executor(matcher, QuietRealTimePresenter(), executionenv=env)

            # Act
            featureresult = executefeature(executor, feature)

            # Assert
            @test Main.features == [feature]
        end
    end
end


text:
module FieldDefaults

using FieldMetadata, Setfield
using FieldMetadata: @default, default, units
using Base: tail

export @default_kw, @udefault_kw, default_kw, udefault_kw

macro default_kw(ex)
    default_kw_unknown(ex, :default_kw, false)
end

macro default_kw(typ, ex)
    default_kw_block(typ, ex, :default_kw)
end

macro udefault_kw(ex)
    default_kw_unknown(ex, :udefault_kw, false)
end

macro udefault_kw(typ, ex)
    default_kw_block(typ, ex, :udefault_kw)
end

default_kw_block(typ, ex, func) = begin
    quote
        import FieldDefaults.default
        $(FieldMetadata.funcs_from_block(typ, ex, :default, Any))
        $(esc(typ))(;kwargs...) = $func($(esc(typ)); kwargs...)
    end
end

default_kw_unknown(ex, func, update) = begin
    typ = FieldMetadata.firsthead(ex, :struct) do typ_ex
        FieldMetadata.namify(typ_ex.args[2])
    end
    quote
        import FieldDefaults.default
        $(FieldMetadata.funcs_from_unknown(ex, :default, Any; update=update))
        $(esc(typ))(;kwargs...) = $func($(esc(typ)); kwargs...)
    end
end

insert_kwargs(kwargs, T::Type) = insert_kwargs(fn -> get_default(T, fn), kwargs, T)
insert_kwargs(defaultfunc, kwargs, T::Type) =
    insert_kwargs(defaultfunc, Tuple(fieldnames(T)), keys(kwargs.data), Tuple(kwargs.data))
insert_kwargs(defaultfunc, fieldnames::Tuple, keys, vals) = begin
    fn = fieldnames[1]
    ind = findfirst(n -> n == fn, keys)
    def = ind == nothing ? defaultfunc(fn) : vals[ind]
    (def, insert_kwargs(defaultfunc, tail(fieldnames), keys, vals)...)
end
insert_kwargs(defaultfunc, fieldnames::Tuple{}, keys, vals) = ()

default_kw(::Type{T}; kwargs...) where T = T(insert_kwargs(kwargs, T)...)

# Combined default() and units()
udefault_kw(::Type{T}; kwargs...) where T =
    T(insert_kwargs(fn -> maybeaddunits(get_default(T, fn), units(T, fn)), kwargs, T)...)

# Add units to Real, otherwise return as-is
maybeaddunits(def::Union{Real,AbstractArray}, units) = def * units
maybeaddunits(def, units) = def

# Somewhere to add custom methods to modify `default`
get_default(args...) = default(args...)

end # module

ext: jl


text:
module IES

import Base: iterate, size
using LinearAlgebra
using Distances: pairwise, Euclidean
using Combinatorics: powerset
using Arpack: eigs
using SparseArrays

include("diffmaps.jl")
include("helper.jl")
include("manifold.jl")
include("ies.jl")

export rlap, DiffMap, gaussian_similarity, rmetric, ies

end


text:
using DiffEqProblemLibrary, Sundials, Test
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems;
importodeproblems();
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear

prob = prob_ode_2Dlinear
integrator = init(prob, CVODE_BDF())
step!(integrator)
integrator(integrator.t, Val{1})

for i in integrator
    @info i.t
end

integrator.t == 1.0
step!(integrator)
integrator.t > 1.0

reinit!(integrator)
solve!(integrator)
u1 = copy(integrator.u)
reinit!(integrator)
solve!(integrator)
u2 = copy(integrator.u)
@test u1 != vec(prob.u0)
@test u1 == u2


text:
# StepSimulator
# maintained by @zsunberg

mutable struct StepSimulator <: Simulator
    rng::AbstractRNG
    max_steps::Union{Nothing,Any}
    spec
end
function StepSimulator(spec; rng=Random.GLOBAL_RNG, max_steps=nothing)
    return StepSimulator(rng, max_steps, spec)
end

function simulate(sim::StepSimulator, mdp::MDP{S}, policy::Policy, init_state::S=initialstate(mdp, sim.rng)) where {S}
    symtuple = convert_spec(sim.spec, MDP)
    if sim.max_steps == nothing
        max_steps = typemax(Int64)
    else
        max_steps = sim.max_steps
    end
    return MDPSimIterator(symtuple, mdp, policy, sim.rng, init_state, max_steps)
end

function simulate(sim::StepSimulator, pomdp::POMDP, policy::Policy, bu::Updater=updater(policy))
    dist = initialstate_distribution(pomdp)
    return simulate(sim, pomdp, policy, bu, dist)
end

function simulate(sim::StepSimulator, pomdp::POMDP, policy::Policy, bu::Updater, dist::Any, is=initialstate(pomdp, sim.rng))
    initial_belief = initialize_belief(bu, dist)
    symtuple = convert_spec(sim.spec, POMDP)
    if sim.max_steps == nothing
        max_steps = typemax(Int64)
    else
        max_steps = sim.max_steps
    end
    return POMDPSimIterator(symtuple, pomdp, policy, bu, sim.rng, initial_belief, is, max_steps)
end

struct MDPSimIterator{SPEC, M<:MDP, P<:Policy, RNG<:AbstractRNG, S}
    mdp::M
    policy::P
    rng::RNG
    init_state::S
    max_steps::Int
end

function MDPSimIterator(spec::Union{Tuple, Symbol}, mdp::MDP, policy::Policy, rng::AbstractRNG, init_state, max_steps::Int)
    return MDPSimIterator{spec, typeof(mdp), typeof(policy), typeof(rng), typeof(init_state)}(mdp, policy, rng, init_state, max_steps)
end

Base.IteratorSize(::Type{<:MDPSimIterator}) = Base.SizeUnknown()

function Base.iterate(it::MDPSimIterator, is::Tuple{Int, S}=(1, it.init_state)) where S
    if isterminal(it.mdp, is[2]) || is[1] > it.max_steps
        return nothing
    end
    t = is[1]
    s = is[2]
    a, ai = action_info(it.policy, s)
    sp, r, i = generate_sri(it.mdp, s, a, it.rng)
    return (out_tuple(it, (s, a, r, sp, t, i, ai)), (t+1, sp))
end

struct POMDPSimIterator{SPEC, M<:POMDP, P<:Policy, U<:Updater, RNG<:AbstractRNG, B, S}
    pomdp::M
    policy::P
    updater::U
    rng::RNG
    init_belief::B
    init_state::S
    max_steps::Int
end
function POMDPSimIterator(spec::Union{Tuple,Symbol}, pomdp::POMDP, policy::Policy, up::Updater, rng::AbstractRNG, init_belief, init_state, max_steps::Int)
    return POMDPSimIterator{spec,
                            typeof(pomdp),
                            typeof(policy),
                            typeof(up),
                            typeof(rng),
                            typeof(init_belief),
                            typeof(init_state)}(pomdp,
                                                policy,
                                                up,
                                                rng,
                                                init_belief,
                                                init_state,
                                                max_steps)
end

Base.IteratorSize(::Type{<:POMDPSimIterator}) = Base.SizeUnknown()

function Base.iterate(it::POMDPSimIterator, is::Tuple{Int,S,B} = (1, it.init_state, it.init_belief)) where {S,B}
    if isterminal(it.pomdp, is[2]) || is[1] > it.max_steps
        return nothing
    end
    t = is[1]
    s = is[2]
    b = is[3]
    a, ai = action_info(it.policy, b)
    sp, o, r, i = generate_sori(it.pomdp, s, a, it.rng)
    bp, ui = update_info(it.updater, b, a, o)
    return (out_tuple(it, (s, a, r, sp, t, i, ai, b, o, bp, ui)), (t+1, sp, bp))
end

const sym_to_ind = Dict(sym=>i for (i, sym) in enumerate(COMPLETE_POMDP_STEP))

@generated function out_tuple(it::Union{MDPSimIterator, POMDPSimIterator}, all::Tuple)
    spec = it.parameters[1]
    if isa(spec, Tuple)
        calls = []
        for sym in spec
            push!(calls, :($sym = all[$(sym_to_ind[sym])]))
        end

        return quote
            return ($(calls...),)
        end
    else
        @assert isa(spec, Symbol) "Invalid specification: $spec is not a Symbol or Tuple."
        return quote
            return all[$(sym_to_ind[spec])]
        end
    end
end

convert_spec(spec, T::Type{POMDP}) = convert_spec(spec, Set(tuple(:sp, :bp, :s, :a, :r, :b, :o, :i, :ai, :ui, :t)))
convert_spec(spec, T::Type{MDP}) = convert_spec(spec, Set(tuple(:sp, :s, :a, :r, :i, :ai, :t)))

function convert_spec(spec, recognized::Set{Symbol})
    conv = convert_spec(spec)
    for s in (isa(conv, Tuple) ? conv : tuple(conv))
        if !(s in recognized)
            @warn("uncrecognized symbol $s in step iteration specification $spec.")
        end
    end
    return conv
end

function convert_spec(spec::String)
    syms = [Symbol(m.match) for m in eachmatch(r"(sp|bp|ai|ui|s|a|r|b|o|i|t)", spec)]
    if length(syms) == 0
        error("$spec does not contain any valid symbols for step iterator output. Valid symbols are sp, bp, ai, ui, s, a, r, b, o, i, t")
    end
    if length(syms) == 1
        return Symbol(first(syms))
    else
        return tuple(syms...)
    end
end

function convert_spec(spec::Tuple)
    for s in spec
        @assert isa(s, Symbol)
    end
    return spec
end

convert_spec(spec::Symbol) = spec

"""
    stepthrough(problem, policy, [spec])
    stepthrough(problem, policy, [spec], [rng=rng], [max_steps=max_steps])
    stepthrough(mdp::MDP, policy::Policy, [init_state], [spec]; [kwargs...])
    stepthrough(pomdp::POMDP, policy::Policy, [up::Updater, [initial_belief, [initial_state]]], [spec]; [kwargs...])

Create a simulation iterator. This is intended to be used with for loop syntax to output the results of each step *as the simulation is being run*.

Example:

    pomdp = BabyPOMDP()
    policy = RandomPolicy(pomdp)

    for (s, a, o, r) in stepthrough(pomdp, policy, "s,a,o,r", max_steps=10)
        println("in state \$s")
        println("took action \$o")
        println("received observation \$o and reward \$r")
    end

The optional `spec` argument can be a string, tuple of symbols, or single symbol and follows the same pattern as [`eachstep`](@ref) called on a `SimHistory` object.

Under the hood, this function creates a `StepSimulator` with `spec` and returns a `[PO]MDPSimIterator` by calling simulate with all of the arguments except `spec`. All keyword arguments are passed to the `StepSimulator` constructor.
"""
function stepthrough end # for documentation

function stepthrough(mdp::MDP, policy::Policy, spec::Union{String, Tuple, Symbol}=COMPLETE_MDP_STEP; kwargs...)
    sim = StepSimulator(spec; kwargs...)
    return simulate(sim, mdp, policy)
end

function stepthrough(mdp::MDP{S},
                     policy::Policy,
                     init_state::S,
                     spec::Union{String, Tuple, Symbol}=COMPLETE_MDP_STEP;
                     kwargs...) where {S}
    sim = StepSimulator(spec; kwargs...)
    return simulate(sim, mdp, policy, init_state)
end

function stepthrough(pomdp::POMDP, policy::Policy, args...; kwargs...)
    spec_included=false
    if !isempty(args) && isa(last(args), Union{String, Tuple, Symbol})
        spec = last(args)
        spec_included = true
        if spec isa statetype(pomdp) && length(args) == 3
            error("Ambiguity between `initial_state` and `spec` arguments in stepthrough. Please explicitly specify the initial state and spec.")
        end
    else
        spec = COMPLETE_POMDP_STEP
    end
    sim = StepSimulator(spec; kwargs...)
    return simulate(sim, pomdp, policy, args[1:end-spec_included]...)
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

"""
Examples
--------

julia> problems = get_problems()
julia> solver = Solver(Modal)
julia> push!(solver, problems...)
julia> solver()

"""
type Modal <: AbstractSolver
    geometric_stiffness :: Bool
    eigvals :: Vector
    eigvecs :: Matrix
    nev :: Int
    which :: Symbol
end

function Modal(nev=10, which=:SM)
    solver = Modal(false, Vector(), Matrix(), nev, which)
end

function call(solver::Solver{Modal}; show_info=true, debug=false, bc_invertible=false)
    show_info && info(repeat("-", 80))
    show_info && info("Starting natural frequency solver")
    show_info && info("Increment time t=$(round(solver.time, 3))")
    show_info && info(repeat("-", 80))
    initialize!(solver)
    # assemble all field problems
    info("Assembling problems ...")
    tic()
    for problem in get_field_problems(solver)
        assemble!(problem, solver.time)
        assemble!(problem, solver.time, Val{:mass_matrix})
    end
    for problem in get_boundary_problems(solver)
        assemble!(problem, solver.time)
    end
    t1 = round(toq(), 2)
    info("Assembled in $t1 seconds.")
    M, K, Kg, f = get_field_assembly(solver)
    Kb, C1, C2, D, fb, g = get_boundary_assembly(solver)
    K = K + Kb
    f = f + fb
    if solver.properties.geometric_stiffness
        K += Kg
    end

    @assert nnz(D) == 0
    @assert C1 == C2

    tic()

    if bc_invertible
        P, h = create_projection(C1, g, Val{:invertible})
    else
        P, h = create_projection(C1, g)
    end
    K_red = P'*K*P
    M_red = P'*M*P
    # make sure matrices are symmetric
    K_red = 1/2*(K_red + K_red')
    M_red = 1/2*(M_red + M_red')

#=
    ndim = size(C1,1)
    nz = get_nonzero_rows(C1)
    nz = setdiff(collect(1:ndim), nz)
    g = zeros(ndim)
    P = spzeros(ndim, ndim)
    for j in nz
        P[j,j] = 1.0
    end
    K_red = P'*K*P
    M_red = P'*M*P
    # make sure matrices are symmetric
    K_red = 1/2*(K_red + K_red')
    M_red = 1/2*(M_red + M_red')

    #=
    K_red = K[nz,nz]
    M_red = M[nz,nz]
    # make sure matrices are symmetric
    K_red = 1/2*(K_red + K_red')
    M_red = 1/2*(M_red + M_red')
    =#

    #=
    K_red = copy(K)
    M_red = copy(M)
    for j=1:size(K_red)
        j in nz && continue
        K_red[j,:] = 0.0
        K_red[:,j] = 0.0
        M_red[j,:] = 0.0
        M_red[:,j] = 0.0
    end
    =#
=#

    t1 = round(toq(), 2)
    info("Eliminated dirichlet boundaries in $t1 seconds.")

    nz = get_nonzero_rows(K_red)
    ndofs = solver.ndofs
    props = solver.properties
    info("Calculate $(props.nev) eigenvalues...")
    if debug && length(nz) < 100
        info("Stiffness matrix:")
        dump(round(full(K[nz, nz])))
        info("Mass matrix:")
        dump(round(full(M[nz, nz])))
    end

    tic()
    om2 = nothing
    X = nothing
    try
        om2, X = eigs(K_red[nz,nz], M_red[nz,nz]; nev=props.nev, which=props.which)
    catch
        info("failed to calculate eigenvalues")
        info("reduced system")
        info("is K symmetric? ", issym(K_red[nz,nz]))
        info("is M symmetric? ", issym(M_red[nz,nz]))
        info("is K positive definite? ", isposdef(K_red[nz,nz]))
        info("is M positive definite? ", isposdef(M_red[nz,nz]))
        k1 = maximum(abs(K_red[nz,nz] - K_red[nz,nz]'))
        m1 = maximum(abs(M_red[nz,nz] - M_red[nz,nz]'))
        info("K 'skewness' (max(abs(K - K'))) = ", k1)
        info("M 'skewness' (max(abs(M - M'))) = ", m1)

        info("original matrix")
        info("is K symmetric? ", issym(K[nz,nz]))
        info("is M symmetric? ", issym(M[nz,nz]))
        info("is K positive definite? ", isposdef(K[nz,nz]))
        info("is M positive definite? ", isposdef(M[nz,nz]))
        k1 = maximum(abs(K[nz,nz] - K[nz,nz]'))
        m1 = maximum(abs(M[nz,nz] - M[nz,nz]'))
        info("K 'skewness' (max(abs(K - K'))) = ", k1)
        info("M 'skewness' (max(abs(M - M'))) = ", m1)

        rethrow()
    end
    info("Eigenvalues computed in $t1 seconds. Eigenvalues: $om2")

    props.eigvals = om2
    props.eigvecs = zeros(ndofs, length(om2))
    v = zeros(ndofs)
    for i=1:length(om2)
        fill!(v, 0.0)
        v[nz] = X[:,i]
        props.eigvecs[:,i] = P*v + g
    end
    t1 = round(toq(), 2)

    for i=1:length(om2)
        freq = real(sqrt(om2[i])/(2.0*pi))
        u = props.eigvecs[:,i]
        field_dim = get_unknown_field_dimension(solver)
        field_name = get_unknown_field_name(solver)
        if field_dim != 1
            nnodes = round(Int, length(u)/field_dim)
            u = reshape(u, field_dim, nnodes)
            u = Vector{Float64}[u[:,i] for i in 1:nnodes]
        end
        for problem in get_problems(solver)
            for element in get_elements(problem)
                connectivity = get_connectivity(element)
                update!(element, field_name, freq => u[connectivity])
            end
        end
    end

    return true
end


text:
"""
    removePointFromIntervalTree!(rootNode::IntervalTreeNode, value::T) where T<:Real

Remove a value from the tree by splitting any interval containing that point
Return the new root of the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root of the tree
- `value::T`: the value to be removed
"""
function removePointFromIntervalTree!(rootNode::IntervalTreeNode, value::T) where T<:Real
        oldIntervals= IntervalTreeNode[]
        findIntersectingIntervals!(Interval(value,value,1),rootNode,oldIntervals)
        newIntervals = IntervalTreeNode[]
        for interval in oldIntervals
                splitIntervalOnPoint!(value, newIntervals, interval)
                rootNode = deleteNodeFromIntervalTree!(interval)
        end
        for interval in newIntervals
                if isnothing(rootNode)
                        rootNode = interval
                else
                        addNodeToIntervalTree!(rootNode, interval)
                end
        end
        return rootNode
end

"""
    removeIntervalFromIntervalTree!(rootNode::IntervalTreeNode, interval::Interval)

Remove a value from the tree by splitting any interval instersecting that interval
Return the new root of the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root of the tree
- `interval::Interval`: the interval to be removed
"""
function removeIntervalFromIntervalTree!(rootNode::IntervalTreeNode, interval::Interval)
        #don't do anything on a bad interval
        if interval.min > interval.max
                return
        end

        oldIntervals= IntervalTreeNode[]
        secondPass = IntervalTreeNode[]
        findIntersectingIntervals!(interval,rootNode,oldIntervals)
        newIntervals = IntervalTreeNode[]
        for oldInterval in oldIntervals
                if interval.stepSize == oldInterval.stepSize
                        splitIntervalOnInterval!(interval, newIntervals, oldInterval)
                        rootNode = deleteNodeFromIntervalTree!(oldInterval)
                else
                        push!(secondPass, oldInterval)
                end
        end

        for oldInterval in secondPass
                i = interval.min
                while i<=interval.max
                        rootNode = removePointFromIntervalTree!(rootNode, i)
                        i += interval.stepSize
                end
        end

        for newInterval in newIntervals
                addNodeToIntervalTree!(rootNode, newInterval)
        end

        return rootNode
end

"""
    splitIntervalOnPoint!(value::T, newIntervals::Array{IntervalTreeNode}, oldInterval::IntervalTreeNode) where T<:Real

Split an interval on a point, excludes that point from the new intervals, does not remove the old interval

# Arguments
- `value::T`: the value to split on
- `newIntervals::Array{IntervalTreeNode}`: the list of intervals to add the new intervals to
- `oldInterval::IntervalTreeNode`: the interval to split
"""
function splitIntervalOnPoint!(value::T, newIntervals::Array{IntervalTreeNode}, oldInterval::IntervalTreeNode) where T<:Real
        if !isADivisibleByB(value-oldInterval.min, oldInterval.stepSize)
                #add the left side interval
                if value>oldInterval.min
                        push!(newIntervals, IntervalTreeNode(oldInterval.min, value - roundedMod(value - oldInterval.min, oldInterval.stepSize), oldInterval.stepSize))
                end
                #add the right side of the interval
                if value<oldInterval.max
                        push!(newIntervals, IntervalTreeNode(value + roundedMod(oldInterval.max-value, oldInterval.stepSize),oldInterval.max, oldInterval.stepSize))
                end
        else
                #add the left side interval
                if value>oldInterval.min
                        push!(newIntervals, IntervalTreeNode(oldInterval.min, value-oldInterval.stepSize, oldInterval.stepSize))
                end
                #add the right side of the interval
                if value<oldInterval.max
                        push!(newIntervals, IntervalTreeNode(value+oldInterval.stepSize, oldInterval.max, oldInterval.stepSize))
                end
        end
end

"""
    splitIntervalOnInterval!(interval::Interval, newIntervals::Array{IntervalTreeNode}, oldInterval::IntervalTreeNode)

Split an interval on an interval, excludes that interval from the new intervals, does not remove the old interval

# Arguments
- `interval::Interval`: the interval to split on
- `newIntervals::Array{IntervalTreeNode}`: the list of intervals to add the new intervals to
- `oldInterval::IntervalTreeNode`: the interval to split
"""
function splitIntervalOnInterval!(interval::Interval, newIntervals::Array{IntervalTreeNode}, oldInterval::IntervalTreeNode)
        if !isADivisibleByB(abs(interval.min-oldInterval.min),oldInterval.stepSize)
                #add the left side interval
                if interval.min>oldInterval.min
                        push!(newIntervals, IntervalTreeNode(oldInterval.min, interval.min - roundedMod(interval.min - oldInterval.min,oldInterval.stepSize), oldInterval.stepSize))
                end
                #add the right side of the interval
                if interval.max<oldInterval.max
                        push!(newIntervals, IntervalTreeNode(interval.max + roundedMod(oldInterval.max-interval.max,oldInterval.stepSize),oldInterval.max, oldInterval.stepSize))
                end
        else
                #add the left side interval
                if interval.min>oldInterval.min
                        push!(newIntervals, IntervalTreeNode(oldInterval.min, interval.min-oldInterval.stepSize, oldInterval.stepSize))
                end
                #add the right side of the interval
                if interval.max<oldInterval.max
                        push!(newIntervals, IntervalTreeNode(interval.max+oldInterval.stepSize, oldInterval.max, oldInterval.stepSize))
                end
        end
end

"""
    imposeMaximumOnIntervalTree!(rootNode::IntervalTreeNode, max::T) where T<:Real

Remove all values above the input max value from the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root node
- `max::T`: the maximum value to impose
"""
function imposeMaximumOnIntervalTree!(rootNode::IntervalTreeNode, max::T) where T<:Real
        if findIntervalTreeMax(rootNode)>max
                intersections = IntervalTreeNode[]
                findIntersectingIntervals!(Interval(max, findIntervalTreeMax(rootNode),1), rootNode, intersections)
                for intersection in intersections
                        rootNode = deleteNodeFromIntervalTree!(intersection)
                        if intersection.min <= max
                                addNodeToIntervalTree!(rootNode, IntervalTreeNode(intersection.min,max - roundedMod(max-intersection.min,intersection.stepSize),intersection.stepSize))
                        end
                end
        end
        return rootNode
end

"""
    imposeMinimumOnIntervalTree!(rootNode::IntervalTreeNode, min::T) where T<:Real

Remove all values below the input min value from the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root node
- `min::T`: the minimum value to impose
"""
function imposeMinimumOnIntervalTree!(rootNode::IntervalTreeNode, min::T) where T<:Real
        if findIntervalTreeMin(rootNode)<min
                intersections = IntervalTreeNode[]
                findIntersectingIntervals!(Interval(findIntervalTreeMin(rootNode), min,1), rootNode, intersections)
                for intersection in intersections
                        rootNode = deleteNodeFromIntervalTree!(intersection)
                        if intersection.max >= min
                                addNodeToIntervalTree!(rootNode, IntervalTreeNode(min + roundedMod(intersection.max-min,intersection.stepSize),intersection.max ,intersection.stepSize))
                        end
                end
        end
        return rootNode
end

"""
    imposeOpenMaximumOnIntervalTree!(rootNode::IntervalTreeNode, max::T) where T<:Real

Remove all values >= the input max value from the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root node
- `max::T`: the maximum value to impose
"""
function imposeOpenMaximumOnIntervalTree!(rootNode::IntervalTreeNode, max::T) where T<:Real
        if findIntervalTreeMax(rootNode)>=max
                intersections = IntervalTreeNode[]
                findIntersectingIntervals!(Interval(max, findIntervalTreeMax(rootNode),1), rootNode, intersections)
                for intersection in intersections
                        rootNode = deleteNodeFromIntervalTree!(intersection)
                        if intersection.min < max
                                if !isADivisibleByB(max-intersection.min, intersection.stepSize)
                                        addNodeToIntervalTree!(rootNode, IntervalTreeNode(intersection.min,max - roundedMod(max-intersection.min,intersection.stepSize),intersection.stepSize))
                                else
                                        addNodeToIntervalTree!(rootNode, IntervalTreeNode(intersection.min,max - intersection.stepSize,intersection.stepSize))
                                end
                        end
                end
        end
        return rootNode
end

"""
    imposeOpenMinimumOnIntervalTree!(rootNode::IntervalTreeNode, min::T) where T<:Real

Remove all values <= the input min value from the tree

# Arguments
- `rootNode::IntervalTreeNode`: the root node
- `min::T`: the minimum value to impose
"""
function imposeOpenMinimumOnIntervalTree!(rootNode::IntervalTreeNode, min::T) where T<:Real
        if findIntervalTreeMin(rootNode)<=min
                intersections = IntervalTreeNode[]
                findIntersectingIntervals!(Interval(findIntervalTreeMin(rootNode), min,1), rootNode, intersections)
                for intersection in intersections
                        rootNode = deleteNodeFromIntervalTree!(intersection)
                        if intersection.max > min
                                if !isADivisibleByB(intersection.max-min, intersection.stepSize)
                                        addNodeToIntervalTree!(rootNode, IntervalTreeNode(min + roundedMod(intersection.max-min,intersection.stepSize),intersection.max ,intersection.stepSize))
                                else
                                        addNodeToIntervalTree!(rootNode, IntervalTreeNode(min + intersection.stepSize,intersection.max ,intersection.stepSize))
                                end
                        end
                end
        end
        return rootNode
end

"""
    addIntervalToIntervalTreeWithMerging!(rootNode::IntervalTreeNode,newNode::Interval)

Add an interval to the tree, but first remove all nodes it overlaps with and merge them into one node
Returns the root node (which may have changed)

# Arguments
- `rootNode::IntervalTreeNode`: the root node
- `newNode::Interval`: the Interval to add to the tree
"""
function addIntervalToIntervalTreeWithMerging!(rootNode::IntervalTreeNode,newNode::Interval)
        intersections = IntervalTreeNode[]
        findIntersectingIntervals!(Interval(newNode.min,newNode.max,newNode.stepSize),rootNode,intersections)

        for intersection in intersections
                if ((intersection.stepSize == newNode.stepSize) && isADivisibleByB(abs(intersection.min-newNode.min),newNode.stepSize))
                        rootNode = deleteNodeFromIntervalTree!(intersection)
                        newNode.min = min(newNode.min,intersection.min)
                        newNode.max = max(newNode.max, intersection.max)
                end
        end
        if rootNode !== nothing
                addNodeToIntervalTree!(rootNode, IntervalTreeNode(newNode.min, newNode.max, newNode.stepSize))
        else
                rootNode = IntervalTreeNode(newNode.min, newNode.max, newNode.stepSize)
        end
        return rootNode
end

"""
    deepCopyIntervalTreeNodes(rootNode::Union{IntervalTreeNode,Nothing}, isMerging::Bool)

Deep copy an interval tree, but the order of the nodes is not maintained (the new tree will still be a proper ordering though)

# Arguments
- `rootNode::Union{IntervalTreeNode,Nothing}`: the root node
- `isMerging::Bool`: if true, the copied tree will eliminate repeat information as it copies
"""
function deepCopyIntervalTreeNodes(rootNode::Union{IntervalTreeNode,Nothing}, isMerging::Bool)
        if rootNode === nothing
                return nothing
        end

        return recursiveIntervalTreeCopy!(rootNode,IntervalTreeNode(rootNode.min, rootNode.max, rootNode.stepSize),isMerging,true)
end

"""
    recursiveIntervalTreeCopy!(currentNode::Union{Nothing, IntervalTreeNode}, newRoot::IntervalTreeNode,isMerging::Bool,isRoot::Bool)

An internal method for recursively copying an interval tree onto a specified root node

# Arguments
- `currentNode::Union{Nothing, IntervalTreeNode}`: the current node to copy
- `newRoot::IntervalTreeNode`: the root of the copy
- `isMerging::Bool`: if true, the copy will merge overlapping nodes
- `isRoot::Bool`: External calls should always be true, inside the method it is always false
"""
function recursiveIntervalTreeCopy!(currentNode::Union{Nothing, IntervalTreeNode}, newRoot::IntervalTreeNode,isMerging::Bool,isRoot::Bool)
        if isRoot
                if isnothing(currentNode.leftNode) && isnothing(currentNode.rightNode)
                        if isMerging
                                newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(currentNode.min, currentNode.max,currentNode.stepSize))
                        else
                                addNodeToIntervalTree!(newRoot, IntervalTreeNode(currentNode.min, currentNode.max,currentNode.stepSize))
                        end
                        return newRoot
                else
                        if isMerging
                                newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(currentNode.min, currentNode.max,currentNode.stepSize))
                        else
                                addNodeToIntervalTree!(newRoot, IntervalTreeNode(currentNode.min, currentNode.max,currentNode.stepSize))
                        end
                        newRoot = recursiveIntervalTreeCopy!(currentNode.leftNode,newRoot, isMerging,false)
                        newRoot = recursiveIntervalTreeCopy!(currentNode.rightNode,newRoot, isMerging,false)
                        return newRoot
                end
        elseif currentNode !== nothing
                if !isMerging
                        addNodeToIntervalTree!(newRoot, IntervalTreeNode(currentNode.min, currentNode.max,currentNode.stepSize))
                else
                        newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(currentNode.min, currentNode.max,currentNode.stepSize))
                end
                newRoot = recursiveIntervalTreeCopy!(currentNode.leftNode,newRoot, isMerging,false)
                newRoot = recursiveIntervalTreeCopy!(currentNode.rightNode,newRoot, isMerging,false)
        else
                return newRoot
        end
end

"""
    unionIntervalTrees(rootA::IntervalTreeNode, rootB::IntervalTreeNode, isMerging::Bool)

Return the union of two interval trees

# Arguments
- `rootA::IntervalTreeNode`: the root of one of the trees
- `rootB::IntervalTreeNode`: the root of a different tree
- `isMerging::Bool`: if true, the union will merge overlapping nodes
"""
function unionIntervalTrees(rootA::IntervalTreeNode, rootB::IntervalTreeNode, isMerging::Bool)
        return recursiveIntervalTreeCopy!(rootA, deepCopyIntervalTreeNodes(rootB, isMerging), isMerging, true)
end

"""
    intersectIntervalTrees(rootA::Union{Nothing, IntervalTreeNode}, rootB::Union{IntervalTreeNode, Nothing})

Return the intersection of two interval trees
O(m log|n|)  where m is the size of tree A and n is the size of tree B
Significantly faster when the step sizes of overlapping intervals are the same

# Arguments
- `rootA::IntervalTreeNode`: the root of one of the trees
- `rootB::IntervalTreeNode`: the root of a different tree
"""
function intersectIntervalTrees(rootA::Union{Nothing, IntervalTreeNode}, rootB::Union{IntervalTreeNode, Nothing})
        return recursiveIntervalTreeIntersect!(rootA, nothing,rootB)
end

"""
    recursiveIntervalTreeIntersect!(currentNode::Union{Nothing, IntervalTreeNode}, newRoot::Union{Nothing, IntervalTreeNode},rootB::Union{Nothing, IntervalTreeNode})

Internal method that recursively constructs the intersection of two interval trees
O(m log|n|)  where m is the size of tree A and n is the size of tree B
Significantly faster when the step sizes of overlapping intervals are the same

# Arguments
- `currentNode::Union{Nothing, IntervalTreeNode}`: the current node from treeA to test
- `newRoot::Union{Nothing, IntervalTreeNode}`: the resulting interval tree, if calling externally make the argument: nothing
- `rootB::IntervalTreeNode`: the root of a different tree
"""
function recursiveIntervalTreeIntersect!(currentNode::Union{Nothing, IntervalTreeNode}, newRoot::Union{Nothing, IntervalTreeNode},rootB::Union{Nothing, IntervalTreeNode})
        if currentNode === nothing
                return newRoot
        else
                intersections = IntervalTreeNode[]
                findIntersectingIntervals!(Interval(currentNode.min,currentNode.max,currentNode.stepSize),rootB,intersections)
                for intersection in intersections
                        #if they overlap with the same step size, and the steps line up
                        if ((currentNode.stepSize == intersection.stepSize) && isADivisibleByB(abs(intersection.min-currentNode.min),currentNode.stepSize))
                                if newRoot !== nothing
                                        newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(max(currentNode.min, intersection.min),min(currentNode.max, intersection.max),currentNode.stepSize))
                                else
                                        newRoot = IntervalTreeNode(max(currentNode.min, intersection.min),min(currentNode.max, intersection.max),currentNode.stepSize)
                                end
                        #else if they have different step sizes
                        elseif (currentNode.stepSize != intersection.stepSize)
                                #increment both intervals until you find the lowest intersection point
                                min1 = currentNode.min
                                min2 = intersection.min
                                while (min1 != min2) && (min1 <= min(currentNode.max, intersection.max)) && (min2 <= min(currentNode.max, intersection.max))
                                        if min1 < min2
                                                min1+=currentNode.stepSize
                                                min1 = round(min1, digits = 8)

                                        else
                                                min2+= intersection.stepSize
                                                min2 = round(min2, digits = 8)
                                        end
                                end
                                #if there is a valid starting point
                                if (min1 == min2) && (min1<=min(currentNode.max, intersection.max)) && (min2 <= min(currentNode.max, intersection.max))
                                        #find the gap that allows for overlap (like lCM but with fractions allowed)
                                        gap = max(currentNode.stepSize, intersection.stepSize)
                                        while (!isADivisibleByB(gap, min(currentNode.stepSize, intersection.stepSize)) && (min1+gap <= min(currentNode.max, intersection.max)))
                                                gap += max(currentNode.stepSize, intersection.stepSize)
                                        end
                                        #if there is a valid gap
                                        if min1+gap <= min(currentNode.max, intersection.max)
                                                if newRoot !== nothing
                                                        newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(min1,min1+(trunc((min(currentNode.max, intersection.max)-min1)/gap)*gap),gap))
                                                else
                                                        newRoot = IntervalTreeNode(min1,min1+(trunc((min(currentNode.max, intersection.max)-min1)/gap)*gap),gap)
                                                end
                                        else # enter in a single point
                                                if newRoot !== nothing
                                                        newRoot = addIntervalToIntervalTreeWithMerging!(newRoot, Interval(min1,min1,min(currentNode.stepSize,intersection.stepSize)))
                                                else
                                                        newRoot = IntervalTreeNode(min1,min1,min(currentNode.stepSize,intersection.stepSize))
                                                end
                                        end
                                end
                        end
                end
                newRoot = recursiveIntervalTreeIntersect!(currentNode.leftNode,newRoot, rootB)
                return recursiveIntervalTreeIntersect!(currentNode.rightNode,newRoot, rootB)
        end
end

"""
    disjunctiveUnion(rootA::IntervalTreeNode, rootB::IntervalTreeNode)

Return the disjunctive union of two interval trees as an ordered touple.
The first elements of the touple contains the elements not in rootA, the second contains the elements not in rootB

# Arguments
- `rootA::IntervalTreeNode`: the root of one of the trees
- `rootB::IntervalTreeNode`: the root of a different tree
"""
function disjunctiveUnion(rootA::IntervalTreeNode, rootB::IntervalTreeNode)
        intersections = intervalTreeToSet(intersectIntervalTrees(rootA,rootB))
        #create disjoint set of rootA
        disjointA = Set{Real}()
        elements = intervalTreeToSet(rootB)
        for e in elements
                if !(e in intersections)
                        push!(disjointA, e)
                end
        end
        #create disjoint set of rootB
        disjointB = Set{Real}()
        elements = intervalTreeToSet(rootA)
        for e in elements
                if !(e in intersections)
                        push!(disjointB, e)
                end
        end

        return (disjointA, disjointB)
end


text:
using .PyCall
using .Distributions
using .Tullio
using .FFTW
using .StatsBase

"""
Wrap scipy's inverse short time fourier transform.

    scipy.signal.istft(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None, nfft=None, input_onesided=True, boundary=True, time_axis=- 1, freq_axis=- 2)
"""
function istft(Zxx; kwargs...)
    signal.istft(Zxx; kwargs...)
end
export istft

function stft(x; kwargs...)
    signal.stft(x; kwargs...)
end
export stft

function linearnonstationaryspectrum(g::Function, Δt, T, windows::Int=20, normalise=true; doplot=false)
    #g(x, t) is a power function
    # Run a quick stft to get the right samplings and arrays
    t⃗ = 0.0:Δt:T
    x = randn(length(t⃗))
    f, t, Z = stft(x; fs=1.0/Δt, nperseg=floor(length(t⃗)/windows)-1)

    R = zeros(size(Z))
    Φ = rand(Uniform(0, 2π), size(Z)) # Uniform random phase

    for ti ∈ 1:size(Z, 2)
        ff = g.(f, (t[ti],))
        if isinf(ff[1])
            ff[1] = 0.0 # So we don't diverge at 0 frequency for power laws
        end
        if normalise
            E = sum(ff.*(f[2]-f[1])) # Approximate energy
            ff = ff./E # So the power spectrum has unit energy
        end
        R[:, ti] = sqrt.(ff) # Fourier amplitudes
    end
    if doplot
        p = heatmap(t[1:end], log10.(f[2:end]), log10.(R[2:end, :].^2), xguide="t", yguide="log₁₀(f)", colorbar_title="log₁₀(S)", yticks=-3:1, framestyle=:box, clims=(-1.5, 1.5))
        display(p)
    end
    ZZ = R.*exp.(Φ.*im)
    return istft(ZZ; fs=1.0/Δt, nperseg=floor(length(t⃗)/windows)-1)[end]
end
export linearnonstationaryspectrum

function selfAffine(β, Δt, T, args...)
    g(x, t) = 1/(x^(β(t)))
    linearnonstationaryspectrum(g, Δt, T, args...)
end
export selfAffine

"""
    Generate a locally self-affine time series, which has a spectrum that varies over time according to a power-law parameterised by β(t)

Similar to Theiler1992, who uses an inverse Fourier transform for random-phase (linear) surrogates from a given power spectrum. However, this process uses the inverse short time fourier transform to construct a non-stationary signal with a time-varying power spectrum.
"""
function selfAffine(P::Process)
    # Just one parameter, β
    seed(P.solver_rng) # Maybe this doesn't propagate to python...
    β = parameter_function(P)
    Δt = getΔt(P)
    T = gettmax(P) - gettransient_t0(P)
    return selfAffine(β, Δt, T)
end

selfAffineSim = Process(
    process = selfAffine,
    parameter_profile = ramp,
    parameter_profile_parameters = (0.0, 1.0, 0.0, 2000.0),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.01,
    savedt = 0.1,
    tmax = 2000.0)
export selfAffineSim

"""
    fouriersurrogate(x::AbstractVector, t::AbstractVector; f::Function=x->x, g::Function=x->x; kwargs...)

Generate a non-stationary surrogate time series from an input `x` after performing manipulations in (short-time) Fourier space. The function `g(A, t)` is applied to the amplitudes of the Fourier coefficients whereas `h(𝜑, t)` is applied to the phases. After modifying the short-time fourier representation of the time series, the inverse short-time Fourier transform is used to construct the surrogate time series. `x` should be regularly sampled, and the time coordinates in `t` are crucial for correctly evaluating `g` and `h`. Optional `kwargs` are passed to python's `stft` and `istft` functions.
"""
function fouriersurrogate(x::AbstractVector, t::AbstractVector; g::Function=(𝑓, 𝑡, A)->A, h::Function=(𝑓, 𝑡, 𝜑)->𝜑, nperseg=1000, kwargs...)::AbstractVector
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    𝑓, 𝑡, 𝑍 = stft(x; fs, nperseg, kwargs...)
    𝑡 .+= t[1] # Offset the start time to match input signal
    gg .= g((𝑓,), 𝑡, (abs.(𝑍[:, j]),), (angle.(𝑍[:, j]),))
    hh .= h((𝑓,), 𝑡, (abs.(𝑍[:, j]),), (angle.(𝑍[:, j]),))
    𝑍 = gg.*cis.(hh)
    𝑡, x̂ = istft(𝑍; fs, nperseg, kwargs...)
    @assert (𝑡[2] - 𝑡[1]) == (t[2] - t[1])
    x̂ = x̂[1:length(x)] # The stft is zero-padded by python to fit evenly into the windows, so the istft is longer than the input x even though it has the same sampling frequency.
end

function fouriersurrogate(x::DimArray; kwargs...)
    t = timeDims(x)
    x̂ = fouriersurrogate(Array(x), t; kwargs...)
    DimArray(x̂, (Ti(t),))
end

export fouriersurrogate

"""
Function for randomising (from a uniform distribution between 0 and 2π) an angle with a probability of 𝜂.
"""
function corruptangle(𝜑::Number, 𝜂)
    𝜑 %= 2π
    rand() < 𝜂 ? rand()*2π : 𝜑
end
corruptangle(𝜑::Vector, 𝜂) = corruptangle.(𝜑, (𝜂,))

"""
Set phases below a threshold to 0 (synchronised)
Should prevent any issues with the desnity of randomisation interfering with power spectrum
"""
function thresholdsynchronise(𝑓::Vector, 𝜑::Number, ν)
    𝑓 < ν ? 0.0 : 𝜑
end
thresholdsynchronise(𝑓::Vector, 𝜑::Vector, ν) = thresholdsynchronise(𝑓, 𝜑, (ν,))

# """
# Take a Process and use the dark magic to produce a corrupted version, which has an extra parameter controlling the probability of the phase of each fourier coefficient being randomised. If planning to save and load this process, an instance of it must first be loaded so that the simulating function is exported
# """
# function corruptphase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0)
#     # In this case, savedt should really be a multiple of dt
#     S = P()
#     ps = getparameter_profile_parameters(P)
#     if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
#         S.parameter_profile = (getparameter_profile(P), parameter_profile)
#         ps = [ps]
#     else
#         S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
#     end
#     S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
#     pr = string(getprocess(P))
#     fname = Symbol("phaseCorrupted"*titlecase(pr))
#     pr = Symbol(pr)
#     @eval begin # Cry me a river
#         function ($fname)(S::Process)
#             D = S()
#             D.parameter_profile = getparameter_profile(S)[1:end-1]
#             if length(D.parameter_profile) == 1
#                 D.parameter_profile = D.parameter_profile[1]
#             end
#             D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
#             if length(D.parameter_profile_parameters) == 1
#                 D.parameter_profile_parameters = D.parameter_profile_parameters[1]
#             end
#             downsample = floor(D.savedt/D.dt) |> Int
#             D.savedt = D.dt # To keep accuracy for the transforms
#             D.process = $pr
#             x = timeseries(D, transient=true)
#             𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
#             ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
#             x = hcat([fouriersurrogate(y, times(D, transient=true); h=(𝑓, 𝑡, 𝜑)->corruptangle(𝜑, 𝜂(𝑡)), nperseg=1000*downsample) for y ∈ ys]...)
#             x = x[1:downsample:end, :]
#         end
#         export $fname
#     end
#     S.process = @eval $fname
#     return S
# end
# export corruptphase

"""
Take a Process and synchronise phases according to a frequency cutoff ν
"""
function synchronisephase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("phaseSynchronised"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            downsample = floor(D.savedt/D.dt) |> Int
            D.savedt = D.dt # To keep accuracy for the transforms
            D.process = $pr
            x = timeseries(D, transient=true)
            ν = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([fouriersurrogate(y, times(D, transient=true); h=(𝑓, 𝑡, A, 𝜑)->thresholdsynchronise(𝑓, 𝜑, ν(𝑡)), nperseg=1000*downsample) for y ∈ ys]...)
            x = x[1:downsample:end, :]
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export synchronisephase

"""
Now we switch gears and go for a simpler 'windowed fourier transform'; it's just the stft, but with non-overlapping, very wide windows that will match exactly the estimation windows we use. This is basically just Carl, since we won't worry about any nonsense like discontinuities at window edges.
"""
function wft(x, t; nwindows=20)
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    remainder = length(x)%nwindows
    remainder > 1 ? (@warn "The time series does not divide well into the number of windows supplied. The remainder is $remainder") : nothing
    𝑓 = rfftfreq((length(x) - remainder)÷nwindows, fs)
    𝑥, 𝑡 = [reshape(i[1:end-remainder], ((length(i)-remainder)÷nwindows, nwindows)) for i ∈ (x, t)]
    # Have to be careful with floating points here. Start by converting all times to integers, since we know they are equally spaced
    𝑡ᵢ = Int.(round.(𝑡./Δt))
    𝑡 = mean(𝑡ᵢ, dims=1)[:].*Δt # No (fewer?) floating point errors
    𝑍 = rfft(𝑥, 1)
    return (𝑓, 𝑡, 𝑍)
end
export wft

"""
Inverse of wft transformation. Just simple windows. If the original time series did not fit well into the number of windows, give a remainder here and it will append that many samples onto the end of the time series, which all have the same value as the last point in the reconstructed time series (shouldn't be a massive deal if the remainder is only 1 or 2 and the sampling period is small).
"""
function iwft(𝑡, 𝑍; remainder=0)
    if length(𝑡) == 1
        x̂ = irfft(𝑍, 2*size(𝑍, 1)-1, 1) # Not usually even in this case
        x̂ = reshape(x̂, length(x̂))
        t̂ = 1:length(x̂) # Can't do a whole lot better with one window time
    else
        nwindows = size(𝑍, 2)
        @assert nwindows == length(𝑡) # 𝑡 has window centres
        Δ𝑡 = (𝑡[2] - 𝑡[1])
        @assert all(diff(𝑡, dims=1) .≈ Δ𝑡) # Equally spaced window centres
        x̂ = irfft(𝑍, 2*size(𝑍, 1)-2, 1)
        x̂ = reshape(x̂, length(x̂))
        Δt̂ = Δ𝑡/((length(x̂))÷nwindows)
        N = Δ𝑡/Δt̂
        t̂₀ = 𝑡[1]-Δt̂*(N-1)/2
        append!(x̂, fill(x̂[end], remainder))
        t̂ = t̂₀:Δt̂:Δt̂*(N*nwindows-1 + remainder)
    end
    return (x̂, t̂)
end
iwft(𝑓, 𝑡, 𝑍; kwargs...) = iwft(𝑡, 𝑍; kwargs...) # 𝑓 not needed, but in case you want to pass wft result directly to iwft
iwft(𝑍::AbstractDimArray; kwargs...) = iwft(timeDims(𝑍), Array(𝑍); kwargs...)
export iwft

"""
Randomise phases of frequencies that constitute a propotion `p` of power (from the high to low frequencies)
"""
function thresholdcorrupt(𝜑, A, 𝑝) # ! SO P IS THE SQRT ROOT OF PROPORTIONAL POWER!!!
    𝜑 .= (𝜑 .+ 2π).%2π
    A[1] = 0 # The first amplitude is just the offset, which we don't care about (i.e. we want power of a mean-centred signal)
    psd = (A.^2)./sum(A.^2)
    cpsd = cumsum(psd)
    # display(plot(cpsd))
    # display(plot(A))
    idx = findfirst(cpsd .> 1-𝑝^2)
    isnothing(idx) ? (return 𝜑) : nothing
    phi = deepcopy(𝜑)
    phi[idx:end] .= rand(length(phi[idx:end])).*2π
    #display(plot([𝜑, phi]))
    return phi
end
thresholdcorrupt(𝑝) = (𝑓, 𝜑, 𝑍) -> thresholdcorrupt(𝜑, 𝑍, 𝑝)

function windowedfouriersurrogate(x::AbstractVector, t::AbstractVector; g::Function=(𝑓, 𝑡, A, 𝜑)->A, h::Function=(𝑓, 𝑡, A, 𝜑)->𝜑, nwindows=20)::AbstractVector
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    𝑓, 𝑡, 𝑍 = wft(x, t; nwindows)
    #display(heatmap(log10.(abs.(𝑍[2:end-1, :])), scale=:log))
    @tullio gg[j] := g(𝑓, 𝑡[j], abs.(𝑍[:, j]), angle.(𝑍[:, j]))
    @tullio hh[j] := h(𝑓, 𝑡[j], abs.(𝑍[:, j]), angle.(𝑍[:, j]))
    gg = hcat(gg...)
    hh = hcat(hh...)
    𝑍 = gg.*cis.(hh)
    #display(plot(angle.(𝑍)))
    #@tullio 𝑍[i, j] = g(𝑓[i], 𝑡[j], abs(𝑍[i, j]), angle(𝑍[i, j]))*cis(h(𝑓[i], 𝑡[j], abs(𝑍[i, j]), angle(𝑍[i, j])))
    #display(heatmap(log10.(abs.(𝑍[2:end-1, :])), scale=:log))
    remainder = Int(length(x)%nwindows)
    x̂, 𝑡 = iwft(𝑡, 𝑍, remainder=remainder) # For most situations in which this function is a good idea, this remainder will be 1
    @assert nwindows == 1 || (𝑡[2] - 𝑡[1]) ≈ (t[2] - t[1])
    return x̂
end

function windowedfouriersurrogate(x::DimArray; kwargs...)
    t = timeDims(x)
    x̂ = windowedfouriersurrogate(Array(x), t; kwargs...)
    DimArray(x̂, (Ti(t),))
end

export windowedfouriersurrogate

"""
Take a Process and produce a corrupted version, which has an extra parameter controlling the probability of the phase of each fourier coefficient being randomised. If planning to save and load this process, an instance of it must first be loaded so that the simulating function is exported
"""
function corruptphase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0; originalres=false, nwindows=20)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("phaseCorrupted"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin # Cry me a river
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            if $originalres
                downsample = floor(D.savedt/D.dt) |> Int
                D.savedt = D.dt # To keep accuracy for the transforms
            else
                downsample = 1
            end
            D.process = $pr
            x = timeseries(D, transient=false)
            Y = timeseries(D, transient=true)
            𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([windowedfouriersurrogate(y, times(D, transient=false); h=(𝑓, 𝑡, A, 𝜑)->thresholdcorrupt(𝜑, A, 𝜂(𝑡)), nwindows=$nwindows) for y ∈ ys]...)
            idx = findfirst(D.t0 .== times(D, transient=true))
            Y[idx:end, :] = x[1:downsample:end, :]
            return Y
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export corruptphase

"""
E.g. Lorenz attractor with corrupted phases
"""
phaseCorruptedLorenzSim = corruptphase(lorenzSim(
    X0 = [0.0, -0.01, 9.0],
    parameter_profile = (constant, constant, constant),
    parameter_profile_parameters = (10.0, 28.0, 8/3),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.05,
    tmax = 1000.0,
    alg = AutoVern7(Rodas5()),
    solver_opts = Dict(:adaptive => true, :reltol => 1e-15)), rampInterval, (0.0, 0.25, 0.0, 1000.0))

"""
E.g. Lorenz attractor with synchronised phases
"""
phaseSynchronisedLorenzSim = synchronisephase(lorenzSim(
    X0 = [0.0, -0.01, 9.0],
    parameter_profile = (constant, constant, constant),
    parameter_profile_parameters = (10.0, 28.0, 8/3),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.05,
    tmax = 1000.0,
    alg = AutoVern7(Rodas5()),
    solver_opts = Dict(:adaptive => true, :reltol => 1e-15)), rampInterval, (0.0, 1.0, 0.0, 1000.0))

function lowpass(𝑓, x::Vector, p)
    idxs = 𝑓 .> p
    xx = deepcopy(x)
    xx[idxs] .= 0.0
    return xx
end

"""
Pass a process through a low pass fourier filter
"""
function lowpass(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0; nwindows=20)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("lowpass"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            D.process = $pr
            x = timeseries(D, transient=false)
            Y = timeseries(D, transient=true)
            𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([windowedfouriersurrogate(y, times(D, transient=false); g=(𝑓, 𝑡, A, 𝜑)->lowpass(𝑓, A, 𝜂(𝑡)), h=(𝑓, 𝑡, A, 𝜑)->lowpass(𝑓, 𝜑, 𝜂(𝑡)), nwindows=$nwindows) for y ∈ ys]...)
            idx = findfirst(D.t0 .== times(D, transient=true))
            Y[idx:end, :] = x
            return Y
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export lowpass


text:
# function Mpx(px_pml::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
#     lambda = modExpand(lambda, ext, iflag)
#     (m, n) = size(lambda)
#     a1 = 9/8;   a2 = -1/24;
#     c1 = a1/dx; c2 = a2/dx;
#     C3 = zeros(m*n)
#     denum = zeros(m*n)
#     for ix = 1 : n
#         for iz = 1 : m
#             a = 1/dt + px_pml[iz,ix]/2
#             b = 1/dt - px_pml[iz,ix]/2
#             denum[(ix-1)*m+iz] = lambda[iz,ix] / a
#             C3[(ix-1)*m+iz]= b / a
#         end
#     end
#     tmp = spzeros(n, n)
#     tmp[1,1] =  c1; tmp[1,2] = c2;
#     tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
#     for ix = 3: n-1
#         tmp[ix,ix-1] = -c1; tmp[ix, ix  ] = c1;
#         tmp[ix,ix-2] = -c2; tmp[ix, ix+1] = c2;
#     end
#     tmp[n,n-1] = -c1; tmp[n,n] = c1; tmp[n,n-2] = -c2;
#     MpxBpx = spdiagm(C3)
#     MpxBvx = spdiagm(denum) * kron(tmp, speye(m))
#     return  MpxBpx, MpxBvx
# end

function Mpx(px_pml::SparseMatrixCSC{Float64,Int64}, v::Array{Float64,2}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
    v = modExpand(v, ext, iflag)
    (m, n) = size(px_pml)
    a1 = 9/8  ; a2 = -1/24;
    c1 = a1/dx; c2 = a2/dx;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt + px_pml[iz,ix]/2
            b = 1/dt - px_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = (v[iz,ix])^2 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(n, n)
    tmp[1,1] =  c1; tmp[1,2] = c2;
    tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
    for ix = 3: n-1
        tmp[ix,ix-1] = -c1; tmp[ix, ix  ] = c1;
        tmp[ix,ix-2] = -c2; tmp[ix, ix+1] = c2;
    end
    tmp[n,n-1] = -c1; tmp[n,n] = c1; tmp[n,n-2] = -c2;
    MpxBpx = spdiagm(C3)
    MpxBvx = spdiagm(denum) * kron(tmp, speye(m))
    return MpxBpx, MpxBvx
end

function Mpx_back(px_pml::SparseMatrixCSC{Float64,Int64}, v::Array{Float64,2}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
    v = modExpand(v, ext, iflag)
    (m, n) = size(px_pml)
    a1 = 9/8  ; a2 = -1/24;
    c1 = a1/dx; c2 = a2/dx;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt - px_pml[iz,ix]/2
            b = 1/dt + px_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = (v[iz,ix])^2 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(n, n)
    tmp[1,1] =  c1; tmp[1,2] = c2;
    tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
    for ix = 3: n-1
        tmp[ix,ix-1] = -c1; tmp[ix, ix  ] = c1;
        tmp[ix,ix-2] = -c2; tmp[ix, ix+1] = c2;
    end
    tmp[n,n-1] = -c1; tmp[n,n] = c1; tmp[n,n-2] = -c2;
    MpxBpx = spdiagm(C3)
    MpxBvx = -spdiagm(denum) * kron(tmp, speye(m))
    return MpxBpx, MpxBvx
end

max_stars_repo_name: WneleiGao/AcousticWaveLab

max_issues_repo_name: WneleiGao/AcousticWaveLab

max_forks_repo_name: WneleiGao/AcousticWaveLab


text:
using TargetedLearning.Common

facts("Testing automatic influence curves") do
    #simulate scalar estimates for EY1, EY0
    srand(1234)
    n = 50
    a = rand(n)
    b = rand(n)
    ey1 = ScalarEstimate(mean(a), a .- mean(a), "EY1")
    ey0 = ScalarEstimate(mean(b), b .- mean(b), "EY0")
    @fact (ey1 - ey0).ic --> roughly(ey1.ic .- ey0.ic)

    #see http://onlinelibrary.wiley.com/doi/10.1002/sim.3445/abstract appendix
    # for influence curves for logRR and logOR
    @fact log(ey1/ey0).ic --> roughly((1/ey1.psi) .* ey1.ic .- (1/ey0.psi) .* ey0.ic)

    logor = log((ey1/(1-ey1))/(ey0/(1-ey0)))
    @fact logor.ic --> roughly((1/ey1.psi + 1/(1-ey1.psi)) .* ey1.ic - (1/ey0.psi + 1/(1-ey0.psi)) .* ey0.ic)

    logor2 = log(ey1) - log(1-ey1) - log(ey0) + log(1-ey0)
    @fact logor.ic --> roughly(logor2.ic)

    ey1tx1 = ey1 * ey1 * ey1 / sqrt(ey1)
    ey1tx2 = ey1 ^ 2.5
    @fact ey1tx1.psi --> roughly(ey1tx2.psi)
    @fact ey1tx1.ic --> roughly(ey1tx2.ic)
end


text:
using StructJuMP

m = StructuredModel(num_scenarios=p)

n = 10
@variable(m, x[1:n] >= 0)
@variable(m, y[1:n] >= 0)
for i in 1:n
    @constraint(m, x[i] - y[i] <= i)
end

p = 3
q = 2
for s = 1:p
    bl = StructuredModel(parent=m, num_scenarios=2, id=s)
    @variable(bl, z[1:n] >= 0)
    @constraint(bl, x + sum(z[i] for i=1:n) .== 1)
    for t = 1:q
        bll = StructuredModel(parent=bl, id=t)
        @variable(bll, w[1:n] >= 0)
        @constraint(bll, sum(w[i] for i=1:n if iseven(i)) == 1)
        par = parent(bll)
        @constraint(bll, sum(z[i] for i=1:n) - sum(w[i] for i=1:n) >= 0)
        @constraint(bll, sum(z[i] for i=1:n) >= 0)
        @constraint(bll, 0 >= sum(w[i] for i=1:n))
    end
end

@constraint(m, sum(z for bl in children(m), z=getVar(bl,:z)) == 1)

#sum(indx for cond(indx)) (models[indx].y[1] + x[1]) = 0