tweak lif neuron demo

demo/lifneuron.jl

@@ -3,47 +3,33 @@ using Conductor, ModelingToolkit, OrdinaryDiffEq,Graphs, SparseArrays, Plots
 import Unitful.ms
 
 stim_dynamics = Conductor.LIF(-75.0, tau_membrane = 10.0, tau_synaptic = 10.0,
-                              threshold = -55.0, resistance = 0.1, stimulus = 500)
+                              threshold = -55.0, resistance = 0.1, stimulus = 220.01)
 @named stim_neuron = CompartmentSystem(stim_dynamics)
-sim_neuron = Simulation(stim_neuron, time = 300ms)
+
+sim_neuron = Simulation(stim_neuron, time = 1000ms)
 sol_neuron = solve(sim_neuron, Euler(); dt = 0.25);
-plot(sol_neuron[stim_neuron.V] + sol_neuron[stim_neuron.S]*(70)) # single neuron ok
 
+S = observed(stim_neuron)[1].lhs # hack until bug-fix in MTK
+Plots.plot(sol_neuron[stim_neuron.V] + sol_neuron[S]*(70))
 
+# Network simulation, constant spiking
+n = 100 # number of neurons
+w = 400.0 # synaptic weight (homogenous)
+t_total = 500.0 # total time of simulation
 dynamics = Conductor.LIF(-75.0, tau_membrane = 10.0, tau_synaptic = 10.0, threshold = -55.0,
-                         resistance = 0.1, stimulus = 0)
-neuron_pop = [CompartmentSystem(dynamics) for _ in 1:100]
-neuron_pop[5] = stim_neuron
-topology = NetworkTopology(neuron_pop, random_regular_digraph(100, 20, dir=:in);
-                           default_weight=100.0);
-
-# NOTE: This is configured to causes network-wide spiking but it will crash even with
-# completely silent network (set `default_weight = 80.0` for subthreshold connectivity)
-@named network = NeuronalNetworkSystem(topology)
-sim_network = Simulation(network, time = 300ms);
-
-# I have tried several solvers and callback triggering mechanisms. Euler is currently least
-# likely to crash out.
-@timed sol_network = solve(sim_network, Euler(), dt=1); # careful--this might end your day
-
-deduplicated_grid = sol_network(0.0:1.0:300.0)
-
-# Symbolic solution indexing with observables will also cause issues.
-# It hangs or crashes for large neuron populations, so handle data manually for now
-# Example, with a passthrough state, no problem
-#series = [deduplicated_grid[network[x].V]' for x in 1:50]
-# ...but with an observable (RGF eval for each subsystem) can cause lockup or total crash
-#deduplicated_grid[network[1].S] # just one usually works, but be careful
-#series = [deduplicated_grid[network[x].S] for x in 1:100] # will 100% hang or terminate process
-#spike_map = sparse(hcat((sol_network[network[x].S] for x in 1:50)...))
-
-spike_map = sparse(map(sol_network[1:2:200,:]) do x
-             x >= -55.0
-         end)
-
+                         resistance = 0.1, stimulus = 0);
+neuron_pop = [CompartmentSystem(dynamics) for _ in 1:n]; # should use Population instead
+neuron_pop[5] = stim_neuron;
+topology = NetworkTopology(neuron_pop, random_regular_digraph(n, 20, dir=:in);
+                           default_weight=w);
+@named network = NeuronalNetworkSystem(topology);
+sim_network = Simulation(network, time = t_total*ms);
+@time sol_network = solve(sim_network, Rosenbrock23(), tstops=0.0:1.0:t_total);
+dedup_grid = sol_network(0.0:1.0:300.0) # deduplicate timestep saving from callbacks
+# NOTE: Symbolic solution indexing with can cause issues, so handle data manually for now
+spike_map = sparse(map(x -> x >= -55.0, dedup_grid[1:2:2*n,:]))
+
+# To see a raster plot:
 using GLMakie, Makie
-spy(spike_map, markersize = 4, marker = :circle)
-# plot(sol_network, vars = [network[x].S for x in 1:100])
-# plot(sol_network[network[5].V])
-# 
+Makie.spy(spike_map, markersize = 4, marker = :circle)
 

src/compartment.jl

@@ -221,17 +221,17 @@ function LIF(voltage = 0.0; tau_membrane = 10.0, tau_synaptic = 10.0, threshold
 end
 
 function CompartmentSystem(dynamics::LIF, defaults, extensions, name, parent)
-    (; V, τ_mem, τ_syn, V_th, R, inputs, stimuli) = dynamics
+    (; V, τ_mem, τ_syn, V_th, R, stimuli) = dynamics
     Iₑ = stimuli[1]
     @variables I(t) = 0.0 S(t) = false
     @parameters V_rest = MTK.getdefault(V)
-    gen = GeneratedCollections(dvs = Set((V, I, S)),
+    gen = GeneratedCollections(dvs = Set((V, I)),
                                ps = Set((τ_mem, τ_syn, V_th, R, V_rest, Iₑ)),
-                               eqs = [D(V) ~ (-(V-V_rest)/τ_mem) + (R*(I + Iₑ))/τ_mem,
-                                      S ~ V >= V_th,
+                               eqs = [D(V) ~ (-(V-V_rest)/τ_mem) + (R*I + R*Iₑ)/τ_mem,
                                       D(I) ~ -I/τ_syn])
 
     (; eqs, dvs, ps, observed, systems, defs) = gen
+    push!(observed, S ~ V >= V_th)
     merge!(defs, defaults)
     return CompartmentSystem(dynamics, eqs, t, collect(dvs), collect(ps), observed, name,
                              systems, defs, extensions, parent)
@@ -243,15 +243,38 @@ function Base.convert(::Type{ODESystem}, compartment::CompartmentSystem{LIF}; wi
     ps  = get_ps(compartment)
     eqs = get_eqs(compartment)
     defs = get_defaults(compartment)
+    obs = get_observed(compartment)
     syss = convert.(ODESystem, get_systems(compartment))
     V = @nonamespace compartment.V
     S = @nonamespace compartment.S
     V_th = @nonamespace compartment.V_th
     V_rest = @nonamespace compartment.V_rest
 
     cb = with_cb ? MTK.SymbolicDiscreteCallback(V >= V_th, [V~V_rest]) : MTK.SymbolicDiscreteCallback[]
+
     return ODESystem(eqs, t, dvs, ps; systems = syss, defaults = defs,
-                     name = nameof(compartment), discrete_events = cb)
+                     name = nameof(compartment), discrete_events = cb,
+                    observed = obs)
+end
+
+function Base.convert(::Type{SDESystem}, compartment::CompartmentSystem{LIF}; with_cb = false)
+
+    dvs = get_states(compartment)
+    ps  = get_ps(compartment)
+    eqs = get_eqs(compartment)
+    defs = get_defaults(compartment)
+    obs = get_observed(compartment)
+    syss = convert.(ODESystem, get_systems(compartment))
+    V = @nonamespace compartment.V
+    S = @nonamespace compartment.S
+    V_th = @nonamespace compartment.V_th
+    V_rest = @nonamespace compartment.V_rest
+
+    cb = with_cb ? MTK.SymbolicDiscreteCallback(V >= V_th, [V~V_rest]) : MTK.SymbolicDiscreteCallback[]
+
+    noise_eqs = [0.0, 1.0]
+    return SDESystem(eqs, noise_eqs, t, dvs, ps; systems = syss, defaults = defs,
+                     name = nameof(compartment), discrete_events = cb, observed = obs)
 end
 
 function CompartmentSystem(dynamics::HodgkinHuxley, defaults, extensions, name, parent)

src/simulation.jl

@@ -19,7 +19,7 @@ function Simulation(neuron::AbstractCompartmentSystem; time::Time, return_system
         return simplified
     else
         @info repr("text/plain", simplified)
-        return ODEProblem(simplified, [], (0., t_val), []; jac, sparse)
+        return ODEProblem(simplified, [], (0., t_val), []) # ignore kwargs until Symbolics bugfix
     end
 end
 
@@ -31,7 +31,7 @@ function Simulation(neuron::CompartmentSystem{LIF}; time::Time, return_system =
         return simplified
     else
         @info repr("text/plain", simplified)
-        return ODEProblem(simplified, [], (0., t_val), []; jac, sparse)
+        return ODEProblem(simplified, [], (0., t_val), []) # ignore kwargs until Symbolics bugfix
     end
 end
 
@@ -42,7 +42,7 @@ function Simulation(network::NeuronalNetworkSystem; time::Time, return_system =
         return simplified
     else
         @info repr("text/plain", simplified)
-        return ODEProblem(simplified, [], (0., t_val), []; jac, sparse)
+        return ODEProblem(simplified, [], (0., t_val), []) # ignore kwargs until Symbolics bugfix
     end
 end