From 2778c8ce632231ddd0e4946b498a20aa34514f05 Mon Sep 17 00:00:00 2001
From: Reuben Gardos Reid <5456207+ReubenJ@users.noreply.github.com>
Date: Tue, 18 Feb 2025 13:42:31 +0100
Subject: [PATCH] Add QN-specific `reinit!` function

This fix allows the use of `Attractors.jl` to find steady states in QNs.
---
 Project.toml                |  4 ++++
 src/qualitative_networks.jl | 25 +++++++++++++++++++------
 test/Project.toml           |  1 +
 test/test-qn.jl             | 18 ++++++++++++++++--
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/Project.toml b/Project.toml
index 70bd502..92e872b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -16,7 +16,9 @@ HerbSearch = "3008d8e8-f9aa-438a-92ed-26e9c7b4829f"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SoleLogics = "b002da8f-3cb3-4d91-bbe3-2953433912b5"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
@@ -32,6 +34,8 @@ HerbSearch = "0.4.1"
 MLStyle = "0.4.17"
 MetaGraphsNext = "0.7"
 Random = "1.10"
+SciMLBase = "2.74.1"
 SoleLogics = "0.12.0"
+StaticArrays = "1.9.12"
 Statistics = "1.10"
 julia = "1.6"
diff --git a/src/qualitative_networks.jl b/src/qualitative_networks.jl
index f296c69..fa8c18e 100644
--- a/src/qualitative_networks.jl
+++ b/src/qualitative_networks.jl
@@ -1,13 +1,15 @@
 import DynamicalSystemsBase: get_state, set_state!
 
 using AbstractTrees: Leaves
-using DynamicalSystemsBase: ArbitrarySteppable
+using DynamicalSystemsBase: ArbitrarySteppable, current_parameters, initial_state
 using HerbConstraints: addconstraint!, DomainRuleNode, VarNode, Ordered, Forbidden
 using HerbCore: AbstractGrammar, RuleNode, get_rule
 using HerbGrammar: add_rule!, rulenode2expr, @csgrammar
 using HerbSearch: rand
 using MLStyle: @match
 using MetaGraphsNext: MetaGraph, SimpleDiGraph, add_edge!, nv, labels
+import SciMLBase
+using StaticArrays: MVector
 
 base_qn_grammar = @csgrammar begin
     Val = Val + Val
@@ -202,11 +204,11 @@ with an [`ArbitrarySteppable`](https://juliadynamics.github.io/DynamicalSystems.
 from [`DynamicalSystems`](https://juliadynamics.github.io/DynamicalSystems.jl/stable/).
 See [`aqn`](@ref) for an example.
 """
-struct QualitativeNetwork
+struct QualitativeNetwork{N,C}
     "Graph containing the topology and target functions of the network"
     graph::MetaGraph
     "State of the network"
-    state::AbstractVector{Int}
+    state::MVector{C,Int}
     "The maximum activation level/state value of any component"
     N::Int
 
@@ -215,7 +217,7 @@ struct QualitativeNetwork
             error("All values in state must be <= N (N=$N)")
         end
 
-        return new(g, s, N)
+        return new{N,length(s)}(g, s, N)
     end
 end
 
@@ -306,7 +308,7 @@ end
 """
     $(TYPEDSIGNATURES)
 """
-function limit_change(prev_value, next_value, N::Int)
+function limit_change(prev_value, next_value, N::Integer)
     if next_value > prev_value
         limited_value = min(prev_value + 1, N)
     elseif next_value < prev_value
@@ -315,7 +317,7 @@ function limit_change(prev_value, next_value, N::Int)
         limited_value = next_value
     end
 
-    return round(Int, limited_value)
+    return limited_value
 end
 
 """
@@ -335,6 +337,17 @@ extract_state(model::QN) = model.state
 extract_parameters(model::QN) = model.graph
 reset_model!(model::QN, u, _) = model.state .= u
 
+function SciMLBase.reinit!(
+    ds::ArbitrarySteppable{<:AbstractVector{<:Real},<:QualitativeNetwork},
+    u::AbstractVector{<:Real} = initial_state(ds);
+    p = current_parameters(ds),
+    t0 = 0, # t0 is not used but required for downstream.
+)
+    ds.reinit(ds.model, u, p)
+    ds.t[] = 0
+    return ds
+end
+
 """
     $(TYPEDSIGNATURES)
 
diff --git a/test/Project.toml b/test/Project.toml
index 5f8ccc2..609f792 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Attractors = "f3fd9213-ca85-4dba-9dfd-7fc91308fec7"
 DynamicalSystemsBase = "6e36e845-645a-534a-86f2-f5d4aa5a06b4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 HerbCore = "2b23ba43-8213-43cb-b5ea-38c12b45bd45"
diff --git a/test/test-qn.jl b/test/test-qn.jl
index 5ff0d24..b223cd9 100644
--- a/test/test-qn.jl
+++ b/test/test-qn.jl
@@ -1,7 +1,10 @@
+using Attractors: AttractorsViaRecurrences, basins_of_attraction
 using DynamicalSystemsBase: step!, get_state, set_state!
 using Graphs: ne, nv
 using Random: seed!
 
+seed!(42)
+
 @testset "QN Grammar Creation" begin
     entities = [:a, :b, :c]
     constants = [i for i = 1:10]
@@ -65,11 +68,22 @@ end
 end
 
 @testset "Async QN" begin
-    seed!(42)
-
     for i = 1:100
         async_qn = aqn(network, N + i)
         step!(async_qn, 10)
         @test all(get_state(async_qn.model) .<= N + i)
     end
 end
+
+@testset "Get attractors" begin
+    n_entities = 3
+    qn = sample_qualitative_network(n_entities, 2)
+
+    async_qn = aqn(qn, 1)
+
+    grid = Tuple(range(0, 1) for _ = 1:n_entities)
+
+    mapper = AttractorsViaRecurrences(async_qn, grid)
+
+    basins = basins_of_attraction(mapper, grid)
+end