From 222188b523115b68763d8ffabc159f38a35e64cb Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Thu, 12 Sep 2024 16:38:04 -0400
Subject: [PATCH] Fix n_calls_per_step

---
 Project.toml                              |  2 +-
 ext/ClimaTimeSteppersBenchmarkToolsExt.jl | 12 ++++++------
 src/solvers/rosenbrock.jl                 |  1 +
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/Project.toml b/Project.toml
index d055127b..04444666 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.35"
+version = "0.7.36"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
diff --git a/ext/ClimaTimeSteppersBenchmarkToolsExt.jl b/ext/ClimaTimeSteppersBenchmarkToolsExt.jl
index 0cb28551..9a08994d 100644
--- a/ext/ClimaTimeSteppersBenchmarkToolsExt.jl
+++ b/ext/ClimaTimeSteppersBenchmarkToolsExt.jl
@@ -39,16 +39,16 @@ n_calls_per_step(::CTS.ARS343, max_newton_iters) = Dict(
     "post_implicit!" => 4,
     "step!" => 1,
 )
-function n_calls_per_step(alg::CTS.RosenbrockAlgorithm{CTS.RosenbrockTableau{N}}) where {N}
+function n_calls_per_step(alg::CTS.RosenbrockAlgorithm)
     return Dict(
         "Wfact" => 1,
-        "ldiv!" => N,
-        "T_imp!" => N,
-        "T_exp_T_lim!" => N,
+        "ldiv!" => CTS.n_stages(alg.tableau),
+        "T_imp!" => CTS.n_stages(alg.tableau),
+        "T_exp_T_lim!" => CTS.n_stages(alg.tableau),
         "lim!" => 0,
-        "dss!" => N,
+        "dss!" => CTS.n_stages(alg.tableau),
         "post_explicit!" => 0,
-        "post_implicit!" => N,
+        "post_implicit!" => CTS.n_stages(alg.tableau),
         "step!" => 1,
     )
 end
diff --git a/src/solvers/rosenbrock.jl b/src/solvers/rosenbrock.jl
index 1a2a24d4..78f8d9e9 100644
--- a/src/solvers/rosenbrock.jl
+++ b/src/solvers/rosenbrock.jl
@@ -27,6 +27,7 @@ struct RosenbrockTableau{N, SM <: SMatrix{N, N}, SM1 <: SMatrix{N, 1}}
     """m = b Γ⁻¹, used to compute the increments k"""
     m::SM1
 end
+n_stages(::RosenbrockTableau{N}) where {N} = N
 
 function RosenbrockTableau(α::SMatrix{N, N}, Γ::SMatrix{N, N}, b::SMatrix{1, N}) where {N}
     A = α / Γ