separate out ODESystem construction

Author: wsphillips

demo/STGneuron.jl

@@ -26,8 +26,8 @@ dynamics = HodgkinHuxley(channels, gradients);
 @named neuron = CompartmentSystem(Vₘ, dynamics;
                                   geometry = geo,
                                   extensions = [calcium_conversion]);
-t_total = 5000.0
-sim = Simulation(neuron, t_total * ms)
+
+sim = Simulation(neuron, (0ms, 4000ms))
 solution = solve(sim, Rosenbrock23());
 
 # Plot at 5kHz sampling

demo/hodgkinhuxley.jl

@@ -1,7 +1,6 @@
 # Classic Hodgkin Huxley neuron with a "current pulse" stimulus
 using Conductor, Unitful, ModelingToolkit, OrdinaryDiffEq, Plots
 import Unitful: mV, mS, cm, µm, pA, nA, mA, µA, ms
-import Conductor: Na, K # shorter aliases for Sodium/Potassium
 
 Vₘ = ParentScope(MembranePotential(-65mV))
 
@@ -24,15 +23,13 @@ kdr_kinetics = [
 @named leak = IonChannel(Leak, max_g = 0.3mS / cm^2)
 
 channels = [NaV, Kdr, leak];
-reversals = Equilibria([Na => 50.0mV, K => -77.0mV, Leak => -54.4mV])
-
+reversals = Equilibria([Sodium => 50.0mV, Potassium => -77.0mV, Leak => -54.4mV])
 @named pulse_stim = PulseTrain(amplitude = 400.0pA, duration = 100ms, delay = 100ms)
-
 dynamics = HodgkinHuxley(channels, reversals);
 
 @named neuron = Compartment(Vₘ, dynamics; geometry = Sphere(radius = 20µm),
-                         stimuli = [pulse_stim])
+                            stimuli = [pulse_stim])
 
-sim = Simulation(neuron, 300ms)
+sim = Simulation(neuron, (0ms,300ms))
 solution = solve(sim, Rosenbrock23(), abstol = 0.01, reltol = 0.01, dtmax = 100.0);
 plot(solution; size = (1200, 800))

src/Conductor.jl

@@ -117,6 +117,11 @@ abstract type AbstractConductanceSystem <: AbstractTimeDependentSystem end
 abstract type AbstractCompartmentSystem <: AbstractTimeDependentSystem end
 abstract type AbstractNeuronalNetworkSystem <: AbstractTimeDependentSystem end
 
+const AbstractConductorSystem = Union{AbstractCurrentSystem,
+                                      AbstractConductanceSystem,
+                                      AbstractCompartmentSystem,
+                                      AbstractNeuronalNetworkSystem}
+
 # Model properties
 abstract type ConductanceModel end
 

src/simulation.jl

@@ -1,9 +1,12 @@
-function simplify_simulation(sys, time)
+
+# for now, fallback to ODESystem conversion
+function ModelingToolkit.ODESystem(sys::AbstractConductorSystem; simplify = true)
     odesys = convert(ODESystem, sys)
-    t_val = ustrip(Float64, ms, time)
-    return t_val, structural_simplify(odesys)
+    return simplify ? structural_simplify(odesys) : odesys
 end
 
+# simplified = settunable(simplified, tunable)
+
 """
 $(TYPEDSIGNATURES)
 
@@ -12,16 +15,11 @@ duration, `time`.
 
 If `return_system == true`, returns a simplified `ODESystem` instead.
 """
-function Simulation(neuron::AbstractCompartmentSystem, time::Time; return_system = false,
-                    jac = false, sparse = false,
-                    parallel = Symbolics.SerialForm())
-    t_val, simplified = simplify_simulation(neuron, time)
-    if return_system
-        return simplified
-    else
-        @info repr("text/plain", simplified)
-        return ODEProblem(simplified, [], (0.0, t_val), []; jac, sparse, parallel)
-    end
+function Simulation(neuron::AbstractCompartmentSystem, tspan;
+                    simplify = true, parallel = Symbolics.SerialForm(), kwargs...)
+    simplified = ODESystem(neuron; simplify)
+    tstart, tstop = time_span(tspan)
+    return ODEProblem(simplified, [], (tstart, tstop), []; parallel, kwargs...)
 end
 
 struct NetworkParameters{T}
@@ -31,46 +29,49 @@ end
 
 Base.getindex(x::NetworkParameters, i) = x.ps[i]
 topology(x::NetworkParameters) = getfield(x, :topology)
+
 function get_weights(integrator, model)
     topo = topology(integrator.p)
     return graph(topo)[model]
 end
-function Simulation(network::NeuronalNetworkSystem, time::Time; return_system = false,
-        jac = false, sparse = false, parallel = Symbolics.SerialForm(), continuous_events = false,
-        refractory = true)
-    t_val, simplified = simplify_simulation(network, time)
-    return_system && return simplified
+
+function Simulation(network::NeuronalNetworkSystem, tspan; simplify = true,
+                    parallel = Symbolics.SerialForm(), continuous_events = false,
+                    refractory = true, kwargs...)
+    odesys = ODESystem(network; simplify)
+    tstart, tstop = time_span(tspan)
     if !any(iseventbased.(synaptic_systems(network)))
-        return ODEProblem(simplified, [], (0.0, t_val), []; jac, sparse, parallel)
+        return ODEProblem(odesys, [], (tstart, tstop), []; parallel, kwargs...)
     else
-        cb = generate_callback(network, simplified; continuous_events, refractory)
-        prob = ODEProblem(simplified, [], (0.0, t_val), []; callback = cb, jac, sparse, parallel)
-        remake(prob; p = NetworkParameters(prob.p, get_topology(network) ))
+        cb = generate_callback(network, odesys; continuous_events, refractory)
+        prob = ODEProblem(odesys, [], (tstart, tstop), [];
+                          callback = cb, parallel, kwargs...)
+        remake(prob; p = NetworkParameters(prob.p, get_topology(network)))
     end
 end
 
 # if continuous, condition has vector cb signature: cond(out, u, t, integrator)
-function generate_callback_condition(network, simplified; continuous_events, refractory)
-    voltage_indices = map_voltage_indices(network, simplified; roots_only = true)
+function generate_callback_condition(network, odesys; continuous_events, refractory)
+    voltage_indices = map_voltage_indices(network, odesys; roots_only = true)
     if continuous_events
         return ContinuousSpikeDetection(voltage_indices)
     else # discrete condition for each compartment
         return [DiscreteSpikeDetection(voltage_index, refractory) for voltage_index in voltage_indices]
     end
 end
 
-function generate_callback_affects(network, simplified)
+function generate_callback_affects(network, odesys)
     spike_affects = []
     for sys in synaptic_systems(network)
-        push!(spike_affects, SpikeAffect(sys, network, simplified))
+        push!(spike_affects, SpikeAffect(sys, network, odesys))
     end
     tailcall = nothing # placeholder for voltage reset
     return NetworkAffects(spike_affects, tailcall)
 end
 
-function generate_callback(network, simplified; continuous_events, refractory)
-    cb_condition = generate_callback_condition(network, simplified; continuous_events, refractory)
-    cb_affect = generate_callback_affects(network, simplified)
+function generate_callback(network, odesys; continuous_events, refractory)
+    cb_condition = generate_callback_condition(network, odesys; continuous_events, refractory)
+    cb_affect = generate_callback_affects(network, odesys)
     if continuous_events
         return VectorContinuousCallback(cb_condition, cb_affect,
                                         length(cb_condition.voltage_indices)) 

src/util.jl

@@ -53,6 +53,14 @@ function build_toplevel(system)
     build_toplevel!(dvs, ps, eqs, defs, system)
 end
 
+function time_span(tspan::Tuple{Time,Time})
+    ustrip(Float64, ms, tspan[1]), ustrip(Float64, ms, tspan[2])
+end
+
+time_span(tspan::Tuple{Float64,Float64}) = tspan
+time_span(tspan::Time) = zero(Float64), ustrip(Float64, ms, tspan)
+time_span(tspan::Real) = zero(Float64), Float64(tspan)
+
 heaviside(x) = ifelse(x > zero(x), one(x), zero(x))
 @register_symbolic heaviside(x)
 ModelingToolkit.get_unit(op::typeof(heaviside), args) = ms^-1
@@ -74,3 +82,22 @@ function set_symarray_metadata(x, ctx, val)
     end
 end
 
+function settunable(sys, ps)
+    new_ps = setmetadata.(ps, ModelingToolkit.VariableTunable, true) 
+    @set! sys.ps = union(new_ps, parameters(sys))
+    return sys
+end
+
+function setbounds(sys, p_dists)
+    new_ps = parameters(sys)
+    for pair in p_dists
+        i = findfirst(isequal(first(pair)), new_ps)
+        i == nothing && throw("Parameter $(first(pair)) not found.")
+        new_ps[i] = setmetadata.(new_ps, ModelingToolkit.VariableBounds, second(pair))
+    end
+    @set! sys.ps = union(new_ps, parameters(sys))
+    return sys
+end
+
+function setdistributions(sys, ps) end
+

test/hodgkinhuxley.jl

@@ -54,8 +54,7 @@ using Unitful: mV, mS, cm, µm, µA, ms, pA
                       states(neuron),
                       parameters(neuron)]) == [11, 11, 8]
 
-    time = 300.0
-    sim_sys = Simulation(neuron, time * ms; return_system = true)
+    sim_sys = ODESystem(neuron)
 
     @test length.([equations(sim_sys),
                       states(sim_sys),
@@ -136,6 +135,7 @@ using Unitful: mV, mS, cm, µm, µA, ms, pA
         return nothing
     end
 
+    time = 300.0
     byhand_prob = ODEProblem{true}(hodgkin_huxley!, u0, (0.0, 300.0), p)
     mtk_prob = ODEProblem(sim_sys, [], (0.0, time), [])
 

test/prinzneuron.jl

@@ -24,8 +24,7 @@ include(normpath(@__DIR__, "..", "demo", "prinz_kinetics.jl"))
                       states(neuron),
                       parameters(neuron)]) == [31, 31, 17]
 
-    time = 2000
-    simul_sys = Simulation(neuron, time * ms; return_system = true)
+    simul_sys = ODESystem(neuron)
 
     @test length.([equations(simul_sys),
                       states(simul_sys),
@@ -184,12 +183,13 @@ include(normpath(@__DIR__, "..", "demo", "prinz_kinetics.jl"))
         end
     end
 
+    time = 2000.0
     byhand_prob = ODEProblem{true}(prinz_neuron!, u0, (0.0, time), p)
     mtk_prob = ODEProblem(simul_sys, [], (0.0, time), [])
     byhand_sol = solve(byhand_prob, Rosenbrock23(), reltol = 1e-8, abstol = 1e-8)
     current_mtk_sol = solve(mtk_prob, Rosenbrock23(), reltol = 1e-8, abstol = 1e-8)
 
-    tsteps = 0.0:0.025:2000.0
+    tsteps = 0.0:0.025:time
     byhand_out = Array(byhand_sol(tsteps, idxs = 2))
     current_mtk_out = current_mtk_sol(tsteps)[Vₘ]
 

test/simplesynapse.jl

@@ -64,8 +64,7 @@ import Conductor: Na, K
                       states(network),
                       parameters(network)]) == [24, 24, 20]
 
-    ttot = 250.0
-    simul_sys = Simulation(network, ttot * ms; return_system = true)
+    simul_sys = ODESystem(network)
 
     @test length.([equations(simul_sys),
                       states(simul_sys),
@@ -235,8 +234,9 @@ import Conductor: Na, K
     end
 
     # Solve and check for invariance
+    ttot = 250.0
     byhand_prob = ODEProblem{true}(simple_synapse!, u0, (0.0, ttot), p)
-    sim = Simulation(network, ttot * ms)
+    sim = Simulation(network, (0.0ms, ttot * ms))
     byhand_sol = solve(byhand_prob, Rosenbrock23(), reltol = 1e-9, abstol = 1e-9,
                        saveat = 0.025)
     current_mtk_sol = solve(sim, Rosenbrock23(), reltol = 1e-9, abstol = 1e-9,