From fb03256db6258c1763b8c5da8c4503bab93d3ecf Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@juliahub.com>
Date: Wed, 14 Feb 2024 15:40:02 +0530
Subject: [PATCH] refactor: use `unknowns`, fix some tests

---
 src/Blocks/analysis_points.jl       |  2 +-
 test/Blocks/test_analysis_points.jl |  4 ++--
 test/Mechanical/multibody.jl        |  2 +-
 test/Mechanical/rotational.jl       | 10 +++++-----
 test/multi_domain.jl                |  4 ++--
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/Blocks/analysis_points.jl b/src/Blocks/analysis_points.jl
index b7690692d..903079919 100644
--- a/src/Blocks/analysis_points.jl
+++ b/src/Blocks/analysis_points.jl
@@ -13,7 +13,7 @@ function ap_var(sys)
         # collect to turn symbolic arrays into arrays of symbols
         return length(sys.u) == 1 ? sys.u : collect(sys.u)
     end
-    x = states(sys)
+    x = unknowns(sys)
     length(x) == 1 && return x[1]
     error("Could not determine the analysis-point variable in system $(nameof(sys)). To use an analysis point, apply it to a connection between two causal blocks containing connectors of type `RealInput/RealOutput` from ModelingToolkitStandardLibrary.Blocks.")
 end
diff --git a/test/Blocks/test_analysis_points.jl b/test/Blocks/test_analysis_points.jl
index ff940b720..436f69809 100644
--- a/test/Blocks/test_analysis_points.jl
+++ b/test/Blocks/test_analysis_points.jl
@@ -20,8 +20,8 @@ sys = ODESystem(eqs, t, systems = [P, C], name = :hej)
 ssys = structural_simplify(sys)
 prob = ODEProblem(ssys, [P.x => 1], (0, 10))
 sol = solve(prob, Rodas5())
-@test norm(sol[1]) >= 1
-@test norm(sol[end]) < 1e-6 # This fails without the feedback through C
+@test norm(sol.u[1]) >= 1
+@test norm(sol.u[end]) < 1e-6 # This fails without the feedback through C
 # plot(sol)
 
 matrices, _ = get_sensitivity(sys, ap)
diff --git a/test/Mechanical/multibody.jl b/test/Mechanical/multibody.jl
index 6b218f64e..e9076c5ac 100644
--- a/test/Mechanical/multibody.jl
+++ b/test/Mechanical/multibody.jl
@@ -22,7 +22,7 @@ eqs = [connect(link1.TX1, cart.flange) #, force.flange)
 @named model = ODESystem(eqs, t, [], []; systems = [link1, link2, cart, fixed])
 
 sys = structural_simplify(model)
-@test length(states(sys)) == 6
+@test length(unknowns(sys)) == 6
 
 # The below code does work...
 #=
diff --git a/test/Mechanical/rotational.jl b/test/Mechanical/rotational.jl
index 38cf3a129..4c2cd03ab 100644
--- a/test/Mechanical/rotational.jl
+++ b/test/Mechanical/rotational.jl
@@ -30,7 +30,7 @@ using OrdinaryDiffEq: ReturnCode.Success
     sol = solve(prob, Rodas4())
     @test SciMLBase.successful_retcode(sol)
 
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, Pair[], (0, 10.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0, Pair[], (0, 10.0))
     sol = solve(prob, DFBDF())
     @test SciMLBase.successful_retcode(sol)
     @test all(sol[inertia1.w] .== 0)
@@ -84,8 +84,8 @@ end
             sine,
         ])
     sys = structural_simplify(model)
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0,
-        [D(D(inertia2.phi)) => 1.0; D.(states(model)) .=> 0.0], (0, 10.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0,
+        [D(D(inertia2.phi)) => 1.0; D.(unknowns(model)) .=> 0.0], (0, 10.0))
     sol = solve(prob, DFBDF())
     @test SciMLBase.successful_retcode(sol)
 
@@ -212,7 +212,7 @@ end
             angle_sensor,
         ])
     sys = structural_simplify(model)
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, Pair[], (0, 10.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0, Pair[], (0, 10.0))
 
     sol = solve(prob, DFBDF())
     @test SciMLBase.successful_retcode(sol)
@@ -259,7 +259,7 @@ end
     @test all(sol[rel_speed_sensor.w_rel.u] .== sol[speed_sensor.w.u])
     @test all(sol[torque_sensor.tau.u] .== -sol[inertia1.flange_b.tau])
 
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, Pair[], (0, 10.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0, Pair[], (0, 10.0))
     sol = solve(prob, DFBDF())
     @test SciMLBase.successful_retcode(sol)
     @test all(sol[inertia1.w] .== 0)
diff --git a/test/multi_domain.jl b/test/multi_domain.jl
index cbb87ea84..80ae9840c 100644
--- a/test/multi_domain.jl
+++ b/test/multi_domain.jl
@@ -70,7 +70,7 @@ using OrdinaryDiffEq: ReturnCode.Success
     @test sol[inertia.w][idx_t]≈(dc_gain * [V_step; -tau_L_step])[2] rtol=1e-3
     @test sol[emf.i][idx_t]≈(dc_gain * [V_step; -tau_L_step])[1] rtol=1e-3
 
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, Pair[], (0, 6.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0, Pair[], (0, 6.0))
     sol = solve(prob, DFBDF())
     @test sol.retcode == Success
     # EMF equations
@@ -159,7 +159,7 @@ end
 
     @test all(sol[inertia.w] .== sol[speed_sensor.w.u])
 
-    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, Pair[], (0, 6.0))
+    prob = DAEProblem(sys, D.(unknowns(sys)) .=> 0.0, Pair[], (0, 6.0))
     sol = solve(prob, DFBDF())
 
     @test sol.retcode == Success