From d5782894dde8bb7bcbd17226e351b13fbd251f5d Mon Sep 17 00:00:00 2001
From: Alexis Renchon <a.renchon@gmail.com>
Date: Thu, 5 Sep 2024 07:40:36 -0700
Subject: [PATCH] wip

---
 experiments/integrated/fluxnet/run_fluxnet.jl | 14 ++++-----
 src/diagnostics/diagnostic.jl                 | 29 ++++++++++---------
 2 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/experiments/integrated/fluxnet/run_fluxnet.jl b/experiments/integrated/fluxnet/run_fluxnet.jl
index 5b03a8a226..9765b141a7 100644
--- a/experiments/integrated/fluxnet/run_fluxnet.jl
+++ b/experiments/integrated/fluxnet/run_fluxnet.jl
@@ -278,7 +278,7 @@ outdir = joinpath(pkgdir(ClimaLand),
                  )
 output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path(outdir)
 
-d_writer = ClimaDiagnostics.Writers.DictWriter(output_dir)
+d_writer = ClimaDiagnostics.Writers.DictWriter()
 
 ref_time = DateTime(2005) # random. not sure what it should be
 
@@ -287,26 +287,26 @@ diags = ClimaLand.default_diagnostics(
     t0,
     ref_time;
     output_writer = d_writer,
-    output_vars = :long,
+    output_vars = :short,
     average_period = :hourly,
    )
 
-diagnostic_handler = ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0, dt = dt)
+diagnostic_handler = ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0, dt = dt);
 
-diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
+diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler);
 
 drivers = ClimaLand.get_drivers(land)
 updatefunc = ClimaLand.make_update_drivers(drivers)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 
-sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+sol = SciMLBase.solve(
     prob,
     ode_algo;
     dt = dt,
     callback = SciMLBase.CallbackSet(driver_cb, diag_cb),
-)
-
+);
 
+ClimaLand.Diagnostics.diagnostic_as_vectors(d_writer, "gpp_1h_average")
 
 
 
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index b9e8a8af62..22fb7d3746 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -1,9 +1,9 @@
-# ClimaLand diagnostics contains # a dictionary `ALL_DIAGNOSTICS` with all the
-# diagnostics we know how to compute, keyed over their short name. If you want
-# to add more diagnostics, look at the included files. You can add your own file
-# if you want to define several new diagnostics that are conceptually related.
-# The dictionary `ALL_DIAGNOSTICS` should be considered an implementation
-# detail, use the getters/setters.
+ ClimaLand diagnostics contains # a dictionary `ALL_DIAGNOSTICS` with all the
+ diagnostics we know how to compute, keyed over their short name. If you want
+ to add more diagnostics, look at the included files. You can add your own file
+ if you want to define several new diagnostics that are conceptually related.
+ The dictionary `ALL_DIAGNOSTICS` should be considered an implementation
+ detail, use the getters/setters.
 
 const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
 
@@ -24,7 +24,7 @@ If possible, please follow the naming scheme outline in
 https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb
 
 Keyword arguments
-=================
+ =================
 
 - `short_name`: Name used to identify the variable in the output files and in the file
                 names. Short but descriptive. `ClimaLand` diagnostics are identified by the
@@ -88,13 +88,13 @@ function get_diagnostic_variable(short_name)
     return ALL_DIAGNOSTICS[short_name]
 end
 
-# General helper functions for undefined diagnostics for a particular model
+ General helper functions for undefined diagnostics for a particular model
 error_diagnostic_variable(variable, land_model::T) where {T} =
     error("Cannot compute $variable with model = $(nameof(T))")
 
-# with_error is a helper macro that generates the error message
-# when the user tries calling something that is incompatible with the model.
-# It should be called when defining compute functions
+ with_error is a helper macro that generates the error message
+ when the user tries calling something that is incompatible with the model.
+ It should be called when defining compute functions
 macro with_error(compute_function_expr)
     compute_function_expr.head == :function ||
         error("Cannot parse this function, head is not a :function")
@@ -132,17 +132,18 @@ Extract `diagnostic` from given `writer` as tuple of vectors (time and value).
 
 `diagnostic` is typically a string with the short name of the diagnostic.
 """
+import ClimaCore
 function diagnostic_as_vectors(writer::DictWriter, diagnostic)
-    axes(writer[diagnostic]) isa ClimaCore.Spaces.PointSpace || error("diagnostic_as_vectors works only on PointSpaces")
+    axes(first(values(writer[diagnostic]))) isa ClimaCore.Spaces.PointSpace || error("diagnostic_as_vectors works only on PointSpaces")
 
     # writer[diagnostic] is a dictionary with keys the times and with values Fields. We need
     # to be a little careful because dictionaries are not ordered, so we have to sort them
     # by time.
     times = collect(keys(writer[diagnostic]))
     sort_indices = sortperm(times)
-    values = first.(parent.(values(writer[diagnostic]))[sort_indices])
+    values_1 = first.(parent.(values(writer[diagnostic]))[sort_indices])
 
-    return times, values
+    return times, values_1
 end
 
 # Do you want to define more diagnostics? Add them here