Add colorrange keyword with default to adjacency plot (#519)

* add init value to `t_stops` reduce * add `colorrange` kwarg with default to `adjacency`
Neuroblox · Jan 6, 2025 · 2cfd175 · 2cfd175
1 parent 664860c
commit 2cfd175
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/ext/MakieExtension.jl b/ext/MakieExtension.jl
@@ -20,7 +20,8 @@ import Neuroblox: powerspectrumplot, powerspectrumplot!
 
 @recipe(Adjacency, blox_or_graph) do scene
     Theme(
-        colormap = :grays
+        colormap = :grays,
+        colorrange = nothing
     )
 end
 
@@ -41,8 +42,14 @@ function Makie.plot!(p::Adjacency)
     hideydecorations!(ax, ticklabels = false, ticks = false)
 
     X, Y, D = findnz(adj.matrix)
+
+    colorrange = if isnothing(p.colorrange[])
+        (minimum(D), maximum(D))
+    else
+        p.colorrange[]
+    end
 
-    heatmap!(p, Y, X, D; colormap = p.colormap[])
+    heatmap!(p, Y, X, D; colormap = p.colormap[], colorrange)
 
     return p
 end

diff --git a/src/blox/reinforcement_learning.jl b/src/blox/reinforcement_learning.jl
@@ -242,7 +242,7 @@ function run_experiment!(agent::Agent, env::ClassificationEnvironment; verbose=f
         agent.problem = remake(agent.problem; tspan, p=init_params)
     end
 
-    t_stops = mapreduce(get_eval_times, union, values(learning_rules))
+    t_stops = mapreduce(get_eval_times, union, values(learning_rules); init=Float64[])
 
     action_selection = agent.action_selection 
     if !isnothing(action_selection)