Remove UnrolledUtilities

Too much compile time
CliMA · Apr 9, 2024 · 82e82c8 · 82e82c8 · Sbozzolo · Apr 9, 2024
1 parent 7f96d9b
commit 82e82c8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 42 deletions.
diff --git a/Project.toml b/Project.toml
@@ -10,7 +10,6 @@ ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
-UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [compat]
 Accessors = "0.1"
@@ -27,7 +26,6 @@ ProfileCanvas = "0.1"
 SafeTestsets = "0.1"
 SciMLBase = "1, 2"
 Test = "1"
-UnrolledUtilities = "0.1"
 julia = "1.9"
 
 [extras]

diff --git a/src/clima_diagnostics.jl b/src/clima_diagnostics.jl
@@ -1,8 +1,6 @@
 import Accessors
 import SciMLBase
 
-import UnrolledUtilities
-
 import .Callbacks:
     Callback, CallbackOrchestrator, DivisorSchedule, EveryDtSchedule
 import .Writers: write_field!, AbstractWriter
@@ -219,50 +217,47 @@ function DiagnosticsCallback(diagnostics_handler::DiagnosticsHandler)
     # dump them to disk. At the moment, they all end up in the same place, but we might want
     # to keep them separate
 
-    # UnrolledUtilities.unrolled_flatmap helps with type stability
-
-    callbacks = UnrolledUtilities.unrolled_flatmap(
-        diagnostics_handler.scheduled_diagnostics,
-    ) do diag
-        isa_time_reduction = !isnothing(diag.reduction_time_func)
-        if isa_time_reduction
-            compute_callback =
-                integrator -> begin
-                    compute_callback!(
-                        integrator,
-                        diagnostics_handler.accumulators[diag],
-                        diagnostics_handler.storage[diag],
-                        Ref(diagnostics_handler.counters[diag]),
-                        diag.variable.compute!,
-                        diag.reduction_time_func,
-                    )
-                end
-        else
-            compute_callback =
+    callbacks =
+        Iterators.flatmap(diagnostics_handler.scheduled_diagnostics) do diag
+            isa_time_reduction = !isnothing(diag.reduction_time_func)
+            if isa_time_reduction
+                compute_callback =
+                    integrator -> begin
+                        compute_callback!(
+                            integrator,
+                            diagnostics_handler.accumulators[diag],
+                            diagnostics_handler.storage[diag],
+                            Ref(diagnostics_handler.counters[diag]),
+                            diag.variable.compute!,
+                            diag.reduction_time_func,
+                        )
+                    end
+            else
+                compute_callback =
+                    integrator -> begin
+                        compute_callback!(
+                            integrator,
+                            diagnostics_handler.storage[diag],
+                            Ref(diagnostics_handler.counters[diag]),
+                            diag.variable.compute!,
+                        )
+                    end
+            end
+            output_callback =
                 integrator -> begin
-                    compute_callback!(
+                    output_callback!(
                         integrator,
+                        diagnostics_handler.accumulators,
                         diagnostics_handler.storage[diag],
-                        Ref(diagnostics_handler.counters[diag]),
-                        diag.variable.compute!,
+                        diag,
+                        diagnostics_handler.counters,
                     )
                 end
+            (
+                Callback(compute_callback, diag.compute_schedule_func),
+                Callback(output_callback, diag.output_schedule_func),
+            )
         end
-        output_callback =
-            integrator -> begin
-                output_callback!(
-                    integrator,
-                    diagnostics_handler.accumulators,
-                    diagnostics_handler.storage[diag],
-                    diag,
-                    diagnostics_handler.counters,
-                )
-            end
-        (
-            Callback(compute_callback, diag.compute_schedule_func),
-            Callback(output_callback, diag.output_schedule_func),
-        )
-    end
 
     return CallbackOrchestrator(callbacks)
 end