From 3ad6f7f60c73c76276b848107d486b9320aceab6 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Sun, 13 Oct 2024 21:38:18 -0500
Subject: [PATCH 1/5] WIP: Add option to load all data at once.

---
 Project.toml                 |  6 +--
 src/geosfp.jl                | 46 +++++++++------------
 src/load.jl                  | 25 +++++++-----
 src/nei2016monthly.jl        | 29 +++++++++-----
 src/netcdf_output.jl         |  4 +-
 test/geosfp_test.jl          | 78 +++++++++++++++++-------------------
 test/load_test.jl            | 38 +++++++++++++-----
 test/nei2016monthly_test.jl  | 37 ++++++++++-------
 test/netcdf_output_test.jl   |  7 ++--
 test/update_callback_test.jl | 25 ++++++------
 10 files changed, 165 insertions(+), 130 deletions(-)

diff --git a/Project.toml b/Project.toml
index cdbe8651..e3204dc6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "EarthSciData"
 uuid = "a293c155-435f-439d-9c11-a083b6b47337"
 authors = ["EarthSciML Authors"]
-version = "0.9.4"
+version = "0.10.0"
 
 [deps]
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
@@ -28,13 +28,13 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 [compat]
 AllocCheck = "0.1"
 DataInterpolations = "3, 4, 5, 6"
-DiffEqCallbacks = "2, 3"
+DiffEqCallbacks = "2, 3, 4"
 DifferentialEquations = "7"
 DocStringExtensions = "0.9"
 DomainSets = "0.7"
 Downloads = "1"
 DynamicQuantities = "0.13, 1"
-EarthSciMLBase = "0.16"
+EarthSciMLBase = "0.17"
 Interpolations = "0.15"
 LaTeXStrings = "1"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
diff --git a/src/geosfp.jl b/src/geosfp.jl
index 8b3dfade..c632140b 100644
--- a/src/geosfp.jl
+++ b/src/geosfp.jl
@@ -219,16 +219,13 @@ Domain options (as of 2022-01-30):
 - NATIVE
 - c720
 
-`coord_defaults` can be used to provide default values for the coordinates of the
-domain. For example if we want to perform a 2D simulation with a vertical dimension,
-we can set `coord_defaults = Dict(:lev => 1)`.
+The native data type for this dataset is Float32.
 
-`dtype` represents the desired data type of the interpolated values. The native data type
-for this dataset is Float32.
+`stream_data` specifies whether the data should be streamed in as needed or loaded all at once.
 
 See http://geoschemdata.wustl.edu/ExtData/ for current options.
 """
-function GEOSFP(domain::AbstractString; coord_defaults=Dict{Symbol,Number}(), spatial_ref="+proj=longlat +datum=WGS84 +no_defs", dtype=Float32, name=:GEOSFP, kwargs...)
+function GEOSFP(domain::AbstractString, domaininfo::DomainInfo; name=:GEOSFP, stream_data=true)
     filesets = Dict{String,GEOSFPFileSet}(
         "A1" => GEOSFPFileSet(domain, "A1"),
         "A3cld" => GEOSFPFileSet(domain, "A3cld"),
@@ -237,25 +234,24 @@ function GEOSFP(domain::AbstractString; coord_defaults=Dict{Symbol,Number}(), sp
         "A3mstE" => GEOSFPFileSet(domain, "A3mstE"),
         "I3" => GEOSFPFileSet(domain, "I3"))
 
-    sample_time = DateTime(2022, 5, 1) # Dummy time to get variable names and dimensions from data.
-    eqs = []
-    vars = []
+    starttime, endtime = EarthSciMLBase.tspan_datetime(domaininfo)
+    pvdict = Dict([Symbol(v) => v for v in EarthSciMLBase.pvars(domaininfo)]...)
+    eqs = Equation[]
+    vars = Num[]
     itps = []
     for (filename, fs) in filesets
-        for varname ∈ varnames(fs, sample_time)
-            itp = DataSetInterpolator{dtype}(fs, varname, sample_time; spatial_ref, kwargs...)
-            dims = dimnames(itp, sample_time)
+        for varname ∈ varnames(fs, starttime)
+            dt = EarthSciMLBase.dtype(domaininfo)
+            itp = DataSetInterpolator{dt}(fs, varname, starttime, endtime,
+                domaininfo.spatial_ref; stream_data=stream_data)
+            dims = dimnames(itp, starttime)
             coords = Num[]
-            for dim ∈ dims
+            for dim in dims
                 d = Symbol(dim)
-                if d ∈ keys(coord_defaults) # Set a default value for this coordinate.
-                    v = (@parameters $d = coord_defaults[d])[1]
-                else # No default value.
-                    v = (@parameters $d)[1]
-                end
-                push!(coords, v)
+                @assert d ∈ keys(pvdict) "GEOSFP coordinate $d not found in domaininfo coordinates ($(pvs))."
+                push!(coords, pvdict[d])
             end
-            eq = create_interp_equation(itp, filename, t, sample_time, coords)
+            eq = create_interp_equation(itp, filename, t, starttime, coords)
             push!(eqs, eq)
             push!(vars, eq.lhs)
             push!(itps, itp)
@@ -266,16 +262,12 @@ function GEOSFP(domain::AbstractString; coord_defaults=Dict{Symbol,Number}(), sp
     @constants P_unit = 1.0 [unit = u"Pa", description = "Unit pressure"]
     @variables P(t) [unit = u"Pa", description = "Pressure"]
     @variables I3₊PS(t) [unit = u"Pa", description = "Pressure at the surface"]
-    d = :lev
-    if d ∈ keys(coord_defaults) # Set a default value for this coordinate.
-        lev = (@parameters $d = coord_defaults[d])[1]
-    else # No default value.
-        lev = (@parameters $d)[1]
-    end
+    @assert :lev in keys(pvdict) "GEOSFP coordinate :lev not found in domaininfo coordinates ($(pvs))."
+    lev = pvdict[:lev]
     pressure_eq = P ~ P_unit * Ap(lev) + Bp(lev) * I3₊PS
     push!(eqs, pressure_eq)
 
-    sys = ODESystem(eqs, t, name=name,
+    sys = ODESystem(eqs, t, vars, [pvdict[:lon], pvdict[:lat], lev], name=name,
         metadata=Dict(:coupletype => GEOSFPCoupler))
     cb = UpdateCallbackCreator(sys, vars, itps)
     return sys, cb
diff --git a/src/load.jl b/src/load.jl
index 3d0ef489..5209bfc4 100644
--- a/src/load.jl
+++ b/src/load.jl
@@ -132,8 +132,7 @@ data records for the times immediately before and after the current time step.
 
 `varname` is the name of the variable to interpolate. `default_time` is the time to use when initializing
 the interpolator. `spatial_ref` is the spatial reference system that the simulation will be using.
-`cache_size` is the number of time steps that should be held in the cache at any given time
-(default=3; must be >=3).
+`stream_data` specifies whether the data should be streamed in as needed or loaded all at once.
 
 !!! warning
     WARNING: This interpolator includes a hack so that the function `deepcopy` does not copy the interpolator,
@@ -156,11 +155,19 @@ mutable struct DataSetInterpolator{To,N,N2,FT,ITPT}
     copyfinish::Channel{Int}
     lock::ReentrantLock
     initialized::Bool
-    kwargs
 
-    function DataSetInterpolator{To}(fs::FileSet, varname::AbstractString, default_time::DateTime; spatial_ref="+proj=longlat +datum=WGS84 +no_defs", cache_size=3, kwargs...) where {To<:Real}
-        cache_size >= 3 || throw(ArgumentError("cache_size must be at least 3"))
-        metadata = loadmetadata(fs, default_time, varname; kwargs...)
+    function DataSetInterpolator{To}(fs::FileSet, varname::AbstractString,
+        starttime::DateTime, enddtime::DateTime, spatial_ref; stream_data=true) where {To<:Real}
+        metadata = loadmetadata(fs, starttime, varname)
+
+        # Download all of the data we will need.
+        check_times = starttime:DataFrequencyInfo(fs, starttime).frequency:enddtime
+        cpts = unique(vcat([DataFrequencyInfo(fs, t).centerpoints for t in check_times]...))
+        cache_size = 3
+        if !stream_data
+            cache_size = max(sum([starttime <= t < enddtime for t in cpts]), 2)
+        end
+
         load_cache = zeros(To, repeat([1], length(metadata.varsize))...)
         data = zeros(To, repeat([2], length(metadata.varsize))..., cache_size) # Add a dimension for time.
         interp_cache = similar(data)
@@ -184,7 +191,7 @@ mutable struct DataSetInterpolator{To,N,N2,FT,ITPT}
         itp = new{To,N,N2,FT,ITPT}(fs, varname, data, interp_cache, itp2, load_cache,
             metadata, times, DateTime(1, 1, 1), coord_trans,
             Channel{DateTime}(0), Channel(1), Channel{Int}(0),
-            ReentrantLock(), false, kwargs)
+            ReentrantLock(), false)
         Threads.@spawn async_loader(itp)
         itp
     end
@@ -298,7 +305,7 @@ function async_loader(itp::DataSetInterpolator)
     for t ∈ itp.loadrequest
         if t != tt
             try
-                loadslice!(itp.load_cache, itp.fs, t, itp.varname; itp.kwargs...)
+                loadslice!(itp.load_cache, itp.fs, t, itp.varname)
                 tt = t
             catch err
                 @error err
@@ -311,7 +318,7 @@ function async_loader(itp::DataSetInterpolator)
         take!(itp.copyfinish)
         try
             tt = nexttimepoint(itp, tt)
-            loadslice!(itp.load_cache, itp.fs, tt, itp.varname; itp.kwargs...)
+            loadslice!(itp.load_cache, itp.fs, tt, itp.varname)
         catch err
             @error err
             rethrow(err)
diff --git a/src/nei2016monthly.jl b/src/nei2016monthly.jl
index 27e037a1..37cdda6a 100644
--- a/src/nei2016monthly.jl
+++ b/src/nei2016monthly.jl
@@ -149,26 +149,37 @@ NOTE: This is an interpolator that returns an emissions value by interpolating b
 centers of the nearest grid cells in the underlying emissions grid, so it may not exactly conserve the total
 emissions mass, especially if the simulation grid is coarser than the emissions grid.
 """
-function NEI2016MonthlyEmis(sector, x, y, lev; spatial_ref="+proj=longlat +datum=WGS84 +no_defs", dtype=Float32, scale=1.0, name=:NEI2016MonthlyEmis, kwargs...)
+function NEI2016MonthlyEmis(sector::AbstractString, domaininfo::DomainInfo; scale=1.0,
+    name=:NEI2016MonthlyEmis, stream_data=true)
     fs = NEI2016MonthlyEmisFileSet(sector)
-    sample_time = DateTime(2016, 5, 1) # Dummy time to get variable names and dimensions from data.
-    eqs = []
+    starttime, endtime = EarthSciMLBase.tspan_datetime(domaininfo)
+    pvdict = Dict([Symbol(v) => v for v in EarthSciMLBase.pvars(domaininfo)]...)
+    @assert :x in keys(pvdict) || :lon in keys(pvdict) "x or lon must be specified in the domaininfo"
+    @assert :y in keys(pvdict) || :lat in keys(pvdict) "y or lat must be specified in the domaininfo"
+    @assert :lev in keys(pvdict) "lev must be specified in the domaininfo"
+    x = :x in keys(pvdict) ? pvdict[:x] : pvdict[:lon]
+    y = :y in keys(pvdict) ? pvdict[:y] : pvdict[:lat]
+    lev = pvdict[:lev]
     @parameters(
         Δz = 60.0, [unit = u"m", description = "Height of the first vertical grid layer"],
     )
-    vars = []
+    eqs = Equation[]
+    vars = Num[]
     itps = []
-    for varname ∈ varnames(fs, sample_time)
-        itp = DataSetInterpolator{dtype}(fs, varname, sample_time; spatial_ref, kwargs...)
-        @constants zero_emis = 0 [unit = units(itp, sample_time) / u"m"]
-        eq = create_interp_equation(itp, "", t, sample_time, [x, y],
+    for varname ∈ varnames(fs, starttime)
+        dt = EarthSciMLBase.dtype(domaininfo)
+        itp = DataSetInterpolator{dt}(fs, varname, starttime, endtime,
+            domaininfo.spatial_ref; stream_data=stream_data)
+        @constants zero_emis = 0 [unit = units(itp, starttime) / u"m"]
+        zero_emis = ModelingToolkit.unwrap(zero_emis) # Unsure why this is necessary.
+        eq = create_interp_equation(itp, "", t, starttime, [x, y],
             wrapper_f=(eq) -> ifelse(lev < 2, eq / Δz * scale, zero_emis),
         )
         push!(eqs, eq)
         push!(vars, eq.lhs)
         push!(itps, itp)
     end
-    sys = ODESystem(eqs, t; name=name,
+    sys = ODESystem(eqs, t, vars, [x, y, lev, Δz]; name=name,
         metadata=Dict(:coupletype => NEI2016MonthlyEmisCoupler))
     cb = UpdateCallbackCreator(sys, vars, itps)
     return sys, cb
diff --git a/src/netcdf_output.jl b/src/netcdf_output.jl
index 8bf99179..1eadd86c 100644
--- a/src/netcdf_output.jl
+++ b/src/netcdf_output.jl
@@ -71,7 +71,7 @@ function EarthSciMLBase.init_callback(nc::NetCDFOutputter, s::Simulator)
     nc.tvar = nctvar
     nc.grid = s.grid
     nc.h = 1
-    start, finish = EarthSciMLBase.time_range(s.domaininfo)
+    start, finish = EarthSciMLBase.tspan(s.domaininfo)
     return PresetTimeCallback(start:nc.time_interval:finish,
         (integrator) -> affect!(nc, integrator),
         finalize=(c, u, t, integrator) -> close(nc.file),
@@ -107,4 +107,4 @@ function affect!(nc::NetCDFOutputter, integrator)
     nc.tvar[nc.h] = integrator.t
     nc.h += 1
     return false
-end
\ No newline at end of file
+end
diff --git a/test/geosfp_test.jl b/test/geosfp_test.jl
index ce75d250..8d072545 100644
--- a/test/geosfp_test.jl
+++ b/test/geosfp_test.jl
@@ -7,12 +7,20 @@ using ModelingToolkit: t, D
 using DynamicQuantities
 import NCDatasets
 
+domain = DomainInfo(DateTime(2022, 1, 1), DateTime(2022, 1, 3);
+    latrange=deg2rad(-85.0f0):deg2rad(2):deg2rad(85.0f0),
+    lonrange=deg2rad(-180.0f0):deg2rad(2.5):deg2rad(175.0f0),
+    levelrange=1:10, dtype=Float64)
+
 @testset "GEOS-FP" begin
-    @parameters lev
-    @parameters lon [unit = u"rad"]
-    @parameters lat [unit = u"rad"]
+    lon, lat, lev = EarthSciMLBase.pvars(domain)
     @constants c_unit = 6.0 [unit = u"rad" description = "constant to make units cancel out"]
-    geosfp, _ = GEOSFP("4x5"; dtype=Float64)
+    geosfp, _ = GEOSFP("4x5", domain)
+
+    @test Symbol.(parameters(geosfp)) == [:lon, :lat, :lev]
+
+    domain2 = EarthSciMLBase.add_partial_derivative_func(domain,
+        partialderivatives_δPδlev_geosfp(geosfp))
 
     function Example()
         @variables c(t) = 5.0 [unit = u"mol/m^3"]
@@ -20,49 +28,37 @@ import NCDatasets
     end
     examplesys = Example()
 
-    domain = DomainInfo(
-        [
-            partialderivatives_δxyδlonlat,
-            partialderivatives_δPδlev_geosfp(geosfp),
-        ],
-        constIC(0.0, t ∈ Interval(Dates.datetime2unix(DateTime(2022, 1, 1)), Dates.datetime2unix(DateTime(2022, 1, 3)))),
-        zerogradBC(lat ∈ Interval(deg2rad(-85.0f0), deg2rad(85.0f0))),
-        periodicBC(lon ∈ Interval(deg2rad(-180.0f0), deg2rad(175.0f0))),
-        zerogradBC(lev ∈ Interval(1.0f0, 10.0f0)),
-    )
-
-    composed_sys = couple(examplesys, domain, Advection(), geosfp)
+    composed_sys = couple(examplesys, domain2, Advection(), geosfp)
     pde_sys = convert(PDESystem, composed_sys)
 
     eqs = equations(pde_sys)
 
     want_terms = [
-        "MeanWind₊v_lon(t, lat, lon, lev)", "GEOSFP₊A3dyn₊U(t, lat, lon, lev)",
-        "MeanWind₊v_lat(t, lat, lon, lev)", "GEOSFP₊A3dyn₊V(t, lat, lon, lev)",
-        "MeanWind₊v_lev(t, lat, lon, lev)", "GEOSFP₊A3dyn₊OMEGA(t, lat, lon, lev)",
-        "GEOSFP₊A3dyn₊U(t, lat, lon, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, U}, t, lon, lat, lev)",
-        "GEOSFP₊A3dyn₊OMEGA(t, lat, lon, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, OMEGA}, t, lon, lat, lev)",
-        "GEOSFP₊A3dyn₊V(t, lat, lon, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, V}, t, lon, lat, lev)",
-        "Differential(t)(ExampleSys₊c(t, lat, lon, lev))", "Differential(lon)(ExampleSys₊c(t, lat, lon, lev)",
-        "MeanWind₊v_lon(t, lat, lon, lev)", "lon2m",
-        "Differential(lat)(ExampleSys₊c(t, lat, lon, lev)",
-        "MeanWind₊v_lat(t, lat, lon, lev)", "lat2meters",
+        "MeanWind₊v_lon(t, lon, lat, lev)", "GEOSFP₊A3dyn₊U(t, lon, lat, lev)",
+        "MeanWind₊v_lat(t, lon, lat, lev)", "GEOSFP₊A3dyn₊V(t, lon, lat, lev)",
+        "MeanWind₊v_lev(t, lon, lat, lev)", "GEOSFP₊A3dyn₊OMEGA(t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊U(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, U}, t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊OMEGA(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, OMEGA}, t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊V(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, V}, t, lon, lat, lev)",
+        "Differential(t)(ExampleSys₊c(t, lon, lat, lev))", "Differential(lon)(ExampleSys₊c(t, lon, lat, lev)",
+        "MeanWind₊v_lon(t, lon, lat, lev)", "lon2m",
+        "Differential(lat)(ExampleSys₊c(t, lon, lat, lev)",
+        "MeanWind₊v_lat(t, lon, lat, lev)", "lat2meters",
         "sin(ExampleSys₊c_unit*lat)", "sin(ExampleSys₊c_unit*lon)",
-        "ExampleSys₊c(t, lat, lon, lev)", "t",
-        "Differential(lev)(ExampleSys₊c(t, lat, lon, lev))",
-        "MeanWind₊v_lev(t, lat, lon, lev)", "P_unit",
+        "ExampleSys₊c(t, lon, lat, lev)", "t",
+        "Differential(lev)(ExampleSys₊c(t, lon, lat, lev))",
+        "MeanWind₊v_lev(t, lon, lat, lev)", "P_unit",
     ]
     have_eqs = string.(eqs)
-    have_eqs = replace.(have_eqs, ("Main."=>"",))
+    have_eqs = replace.(have_eqs, ("Main." => "",))
     for term ∈ want_terms
         @test any(occursin.((term,), have_eqs))
     end
 end
 
 @testset "GEOS-FP pressure levels" begin
-    @parameters lat, [unit=u"rad"], lon, [unit=u"rad"], lev
-    geosfp, updater = GEOSFP("4x5"; dtype=Float64,
-        coord_defaults=Dict(:lev => 1.0, :lat => deg2rad(39.1), :lon => deg2rad(-155.7)))
+    @parameters lat, [unit = u"rad"], lon, [unit = u"rad"], lev
+    geosfp, updater = GEOSFP("4x5", domain)
 
     # Rearrange pressure equation so it can be evaluated for P.
     iips = findfirst((x) -> x == :I3₊PS, [Symbolics.tosymbol(eq.lhs, escape=false) for eq in equations(geosfp)])
@@ -82,7 +78,7 @@ end
     dp = partialderivatives_δPδlev_geosfp(geosfp)
 
     # Check level coordinate index
-    ff = dp([lat, lon, lev])
+    ff = dp([lon, lat, lev])
     @test all(keys(ff) .=== [3])
 
     fff = ModelingToolkit.subs_constants(ff[3])
@@ -97,15 +93,14 @@ end
 
 @testset "GEOS-FP new day" begin
     @parameters(
-        lon = 0.0, [unit=u"rad"],
-        lat = 0.0, [unit=u"rad"],
+        lon = 0.0, [unit = u"rad"],
+        lat = 0.0, [unit = u"rad"],
         lev = 1.0,
     )
     starttime = datetime2unix(DateTime(2022, 5, 1, 23, 58))
     endtime = datetime2unix(DateTime(2022, 5, 2, 0, 3))
 
-    geosfp, updater = GEOSFP("4x5"; dtype=Float64,
-        coord_defaults=Dict(:lon => 0.0, :lat => 0.0, :lev => 1.0))
+    geosfp, updater = GEOSFP("4x5", domain)
 
     iips = findfirst((x) -> x == :I3₊PS, [Symbolics.tosymbol(eq.lhs, escape=false) for eq in equations(geosfp)])
     pseq = equations(geosfp)[iips]
@@ -117,13 +112,12 @@ end
 
 @testset "GEOS-FP wrong year" begin
     @parameters(
-        lon = 0.0, [unit=u"rad"],
-        lat = 0.0, [unit=u"rad"],
+        lon = 0.0, [unit = u"rad"],
+        lat = 0.0, [unit = u"rad"],
         lev = 1.0,
     )
     starttime = datetime2unix(DateTime(5000, 1, 1))
 
-    geosfp, updater = GEOSFP("4x5"; dtype=Float64,
-        coord_defaults=Dict(:lon => 0.0, :lat => 0.0, :lev => 1.0))
+    geosfp, updater = GEOSFP("4x5", domain)
     @test_throws Base.Exception EarthSciData.lazyload!(updater, starttime)
 end
diff --git a/test/load_test.jl b/test/load_test.jl
index eb480b0b..781d5b4b 100644
--- a/test/load_test.jl
+++ b/test/load_test.jl
@@ -10,6 +10,8 @@ using Test
 
 fs = EarthSciData.GEOSFPFileSet("4x5", "A3dyn")
 t = DateTime(2022, 5, 1)
+te = DateTime(2022, 5, 2)
+spatial_ref = "+proj=longlat +datum=WGS84 +no_defs"
 @test EarthSciData.url(fs, t) == "http://geoschemdata.wustl.edu/ExtData/GEOS_4x5/GEOS_FP/2022/05/GEOSFP.20220501.A3dyn.4x5.nc"
 
 @test endswith(EarthSciData.localpath(fs, t), joinpath("GEOS_4x5", "GEOS_FP", "2022", "05", "GEOSFP.20220501.A3dyn.4x5.nc"))
@@ -24,7 +26,7 @@ metadata = EarthSciData.loadmetadata(fs, t, "U")
 @test metadata.varsize == [72, 46, 72]
 @test metadata.dimnames == ["lon", "lat", "lev"]
 
-itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", t)
+itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", t, te, spatial_ref)
 
 @test String(latexify(itp)) == "\$\\mathrm{GEOSFPFileSet}\\left( U \\right)\$"
 
@@ -80,10 +82,12 @@ end
     fs = DummyFileSet(DateTime(2022, 4, 30), DateTime(2022, 5, 4))
 
     @testset "big cache" begin
-        @test_nowarn EarthSciData.DataSetInterpolator{Float32}(fs, "U", fs.start; cache_size=100)
+        @test_nowarn EarthSciData.DataSetInterpolator{Float32}(fs, "U", DateTime(2022, 5, 1), DateTime(2022, 5, 3),
+            spatial_ref; stream_data=false)
     end
 
-    itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", fs.start; cache_size=5)
+    itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", DateTime(2022, 5, 1), DateTime(2022, 5, 3),
+        spatial_ref; stream_data=true)
     dfi = EarthSciData.DataFrequencyInfo(fs, fs.start)
 
     answerdata = [tv(fs, t) * v for t ∈ dfi.centerpoints, v ∈ [1.0, 0.5, 2.0]]
@@ -105,10 +109,10 @@ end
 
     @test uvals ≈ answers
 
-    @test length(itp.times) == 5
-    @test itp.times == [DateTime("2022-05-02T19:30:00"), DateTime("2022-05-02T22:30:00"),
-        DateTime("2022-05-03T01:30:00"), DateTime("2022-05-03T04:30:00"),
-        DateTime("2022-05-03T07:30:00")]
+    interp!(itp, times[end], xs[end])
+    @test length(itp.times) == 3
+    @test itp.times == [DateTime("2022-05-02T22:30:00"), DateTime("2022-05-03T01:30:00"),
+        DateTime("2022-05-03T04:30:00")]
 
     uvals = zeros(Float32, length(times), length(xs))
     answers = zeros(Float32, length(times), length(xs))
@@ -121,10 +125,26 @@ end
         end
     end
     @test uvals ≈ answers
+
+    @testset "no stream" begin
+        itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", DateTime(2022, 5, 1),
+            DateTime(2022, 5, 2), spatial_ref; stream_data=false)
+
+        uvals = zeros(Float32, length(times), length(xs))
+        answers = zeros(Float32, length(times), length(xs))
+        for (i, tt) ∈ enumerate(times)
+            for (j, x) ∈ enumerate(xs)
+                uvals[i, j] = interp!(itp, tt, x)
+                answers[i, j] = answer_itp(datetime2unix(tt), x)
+            end
+        end
+
+        @test uvals ≈ answers
+    end
 end
 
 @testset "allocations" begin
-    itp = EarthSciData.DataSetInterpolator{Float64}(fs, "U", t)
+    itp = EarthSciData.DataSetInterpolator{Float64}(fs, "U", t, te, spatial_ref)
     tt = DateTime(2022, 5, 1)
     interp!(itp, tt, 1.0, 0.0, 1.0)
 
@@ -138,7 +158,7 @@ end
             rethrow(err)
         end
 
-        itp2 = EarthSciData.DataSetInterpolator{Float32}(fs, "U", t)
+        itp2 = EarthSciData.DataSetInterpolator{Float32}(fs, "U", t, te, spatial_ref)
         interp!(itp2, tt, 1.0f0, 0.0f0, 1.0f0)
         checkf(itp2, tt, 1.0f0, 0.0f0, 1.0f0)
         true
diff --git a/test/nei2016monthly_test.jl b/test/nei2016monthly_test.jl
index 95737b06..73d2c24d 100644
--- a/test/nei2016monthly_test.jl
+++ b/test/nei2016monthly_test.jl
@@ -7,8 +7,18 @@ using DifferentialEquations
 import Proj
 using AllocCheck
 
-@parameters lat, [unit = u"rad"], lon, [unit = u"rad"], lev
-emis, updater = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", lon, lat, lev; dtype=Float64)
+domain = DomainInfo(DateTime(2016, 5, 1), DateTime(2016, 5, 2);
+    latrange=deg2rad(-85.0f0):deg2rad(2):deg2rad(85.0f0),
+    lonrange=deg2rad(-180.0f0):deg2rad(2.5):deg2rad(175.0f0),
+    levelrange=1:10, dtype=Float64)
+lon, lat, lev = EarthSciMLBase.pvars(domain)
+
+ts, te = EarthSciMLBase.tspan_datetime(domain)
+sample_time = ts
+
+spatial_ref="+proj=longlat +datum=WGS84 +no_defs"
+
+emis, updater = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
 fileset = EarthSciData.NEI2016MonthlyEmisFileSet("mrggrid_withbeis_withrwc")
 
 eqs = equations(emis)
@@ -18,30 +28,30 @@ eqs = equations(emis)
 sample_time = DateTime(2016, 5, 1)
 
 @testset "can't deepcopy" begin
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time; spatial_ref="+proj=longlat +datum=WGS84 +no_defs")
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     deepcopy(itp) === itp
 end
 
 @testset "correct projection" begin
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time; spatial_ref="+proj=longlat +datum=WGS84 +no_defs")
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     @test interp!(itp, sample_time, deg2rad(-97.0f0), deg2rad(40.0f0)) ≈ 9.211331f-10
 end
 
 @testset "incorrect projection" begin
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time;
-        spatial_ref="+proj=axisswap +order=2,1 +step +proj=longlat +datum=WGS84 +no_defs")
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te,
+        "+proj=axisswap +order=2,1 +step +proj=longlat +datum=WGS84 +no_defs")
     @test_throws Proj.PROJError interp!(itp, sample_time, deg2rad(-97.0f0), deg2rad(40.0f0))
 end
 
 @testset "Out of domain" begin
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time)
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     @test_throws BoundsError interp!(itp, sample_time, deg2rad(0.0f0), deg2rad(40.0f0))
 end
 
 
 @testset "monthly frequency" begin
     sample_time = DateTime(2016, 5, 1)
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time; spatial_ref="+proj=longlat +datum=WGS84 +no_defs")
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     EarthSciData.initialize!(itp, sample_time)
     ti = EarthSciData.DataFrequencyInfo(itp.fs, sample_time)
     @test month(itp.times[1]) == 4
@@ -69,7 +79,7 @@ end
     @check_allocs checkf(itp, t, loc1, loc2) = EarthSciData.interp_unsafe(itp, t, loc1, loc2)
 
     sample_time = DateTime(2016, 5, 1)
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time)
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     interp!(itp, sample_time, deg2rad(-97.0f0), deg2rad(40.0f0))
     # If there is an error, it should occur in the proj library.
     # https://github.com/JuliaGeo/Proj.jl/issues/104
@@ -81,7 +91,7 @@ end
         contains(string(s), "libproj.proj_trans")
     end
 
-    itp2 = EarthSciData.DataSetInterpolator{Float64}(fileset, "NOX", sample_time)
+    itp2 = EarthSciData.DataSetInterpolator{Float64}(fileset, "NOX", ts, te, spatial_ref)
     interp!(itp2, sample_time, deg2rad(-97.0), deg2rad(40.0))
     try # If there is an error, it should occur in the proj library.
         checkf(itp2, sample_time, deg2rad(-97.0), deg2rad(40.0))
@@ -93,17 +103,16 @@ end
 end
 
 @testset "Coupling with GEOS-FP" begin
-    gfp = GEOSFP("4x5"; dtype=Float64,
-        coord_defaults=Dict(:lon => 0.0, :lat => 0.0, :lev => 1.0))
+    gfp = GEOSFP("4x5", domain)
 
     eqs = equations(convert(ODESystem, couple(emis, gfp)))
 
-    @test occursin("NEI2016MonthlyEmis₊lat(t) ~ GEOSFP₊lat", string(eqs))
+    @test occursin("NEI2016MonthlyEmis₊lat(t) ~ lat", string(eqs))
 end
 
 @testset "wrong year" begin
     sample_time = DateTime(2016, 5, 1)
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", sample_time)
+    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     sample_time = DateTime(2017, 5, 1)
     @test_throws AssertionError EarthSciData.initialize!(itp, sample_time)
 end
diff --git a/test/netcdf_output_test.jl b/test/netcdf_output_test.jl
index 5f5aaf2f..8b8a7772 100644
--- a/test/netcdf_output_test.jl
+++ b/test/netcdf_output_test.jl
@@ -24,19 +24,20 @@ domain = DomainInfo(
     constIC(0.0, t ∈ Interval(0.0, 2.0)),
     constBC(16.0, x ∈ Interval(-1.0, 1.0),
         y ∈ Interval(-2.0, 2.0),
-        lev ∈ Interval(1.0, 3.0)))
+        lev ∈ Interval(1.0, 3.0)),
+        grid_spacing = [0.1, 0.1, 1])
 
 file = tempname() * ".nc"
 
 csys = couple(sys, domain)
 
 o = NetCDFOutputter(file, 1.0; extra_vars=[
-    structural_simplify(convert(ODESystem, csys)).Test₊sys.v
+    structural_simplify(convert(ODESystem, csys)).Test₊sys₊v
 ])
 
 csys = couple(csys, o)
 
-sim = Simulator(csys, [0.1, 0.1, 1])
+sim = Simulator(csys)
 st = SimulatorStrangThreads(Tsit5(), Euler(), 0.01)
 
 run!(sim, st)
diff --git a/test/update_callback_test.jl b/test/update_callback_test.jl
index 9af36a94..f5732c68 100644
--- a/test/update_callback_test.jl
+++ b/test/update_callback_test.jl
@@ -8,6 +8,17 @@ using DifferentialEquations
 using DynamicQuantities
 
 @parameters lat = deg2rad(40.0) lon = deg2rad(-97.0) lev = 0.0
+starttime = datetime2unix(DateTime(2016, 3, 1))
+endtime = datetime2unix(DateTime(2016, 5, 2))
+domain = DomainInfo(
+    constIC(16.0, t ∈ Interval(starttime, endtime)),
+    constBC(16.0,
+        lon ∈ Interval(deg2rad(-115), deg2rad(-68.75)),
+        lat ∈ Interval(deg2rad(25), deg2rad(53.71875)),
+        lev ∈ Interval(1, 2)
+    ),
+    grid_spacing = [deg2rad(15.0), deg2rad(15.0), 1])
+
 @variables ACET(t) = 0.0 [unit = u"kg*m^-3"]
 @constants c = 1000 [unit = u"s"]
 
@@ -23,21 +34,11 @@ function EarthSciMLBase.couple2(sys::SysCoupler, emis::EarthSciData.NEI2016Month
     operator_compose(sys, emis)
 end
 
-emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", lon, lat, lev; dtype=Float64)
-
-starttime = datetime2unix(DateTime(2016, 3, 1))
-endtime = datetime2unix(DateTime(2016, 5, 2))
-domain = DomainInfo(
-    constIC(16.0, t ∈ Interval(starttime, endtime)),
-    constBC(16.0,
-        lon ∈ Interval(deg2rad(-115), deg2rad(-68.75)),
-        lat ∈ Interval(deg2rad(25), deg2rad(53.71875)),
-        lev ∈ Interval(1, 2)
-    ))
+emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
 
 csys = couple(sys, emis, domain)
 
-sim = Simulator(csys, [deg2rad(15.0), deg2rad(15.0), 1])
+sim = Simulator(csys)
 
 st = SimulatorStrangSerial(Tsit5(), Euler(), 100.0)
 

From e9338eb5b2530a710712f116363046492a6348f8 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Sun, 27 Oct 2024 16:27:55 -0500
Subject: [PATCH 2/5] Update to EarthSciMLBase v0.18

---
 Project.toml                 |  4 ++--
 bench/runbenchmarks.jl       | 20 ++++++++++++--------
 src/netcdf_output.jl         | 28 +++++++++++++++++-----------
 src/update_callback.jl       | 12 +++++++-----
 test/geosfp_test.jl          |  4 ++--
 test/nei2016monthly_test.jl  |  8 +++++---
 test/netcdf_output_test.jl   | 31 ++++++++++++++++---------------
 test/update_callback_test.jl | 14 ++++++++------
 8 files changed, 69 insertions(+), 52 deletions(-)

diff --git a/Project.toml b/Project.toml
index e3204dc6..ed4592c5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -26,7 +26,7 @@ Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
-AllocCheck = "0.1"
+AllocCheck = "0.2"
 DataInterpolations = "3, 4, 5, 6"
 DiffEqCallbacks = "2, 3, 4"
 DifferentialEquations = "7"
@@ -34,7 +34,7 @@ DocStringExtensions = "0.9"
 DomainSets = "0.7"
 Downloads = "1"
 DynamicQuantities = "0.13, 1"
-EarthSciMLBase = "0.17"
+EarthSciMLBase = "0.18"
 Interpolations = "0.15"
 LaTeXStrings = "1"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
diff --git a/bench/runbenchmarks.jl b/bench/runbenchmarks.jl
index c3130105..0d102c9f 100644
--- a/bench/runbenchmarks.jl
+++ b/bench/runbenchmarks.jl
@@ -44,7 +44,7 @@ function EarthSciMLBase.couple2(sys::SysCoupler, emis::EarthSciData.NEI2016Month
     sys, emis = sys.sys, emis.sys
     operator_compose(sys, emis)
 end
-function nei_simulator()
+function nei_simulator(st)
     @parameters lat = deg2rad(40.0) lon = deg2rad(-97.0) lev = 0.0
     @variables ACET(t) = 0.0 [unit = u"kg*m^-3"]
     @constants c = 1000 [unit = u"s"]
@@ -64,22 +64,26 @@ function nei_simulator()
             lon ∈ Interval(deg2rad(-115), deg2rad(-68.75)),
             lat ∈ Interval(deg2rad(25), deg2rad(53.7)),
             lev ∈ Interval(1, 2)
-        ))
+        );
+        grid_spacing=[deg2rad(1.25), deg2rad(1.2), 1])
 
     csys = couple(sys, emis, domain)
 
-    Simulator(csys, [deg2rad(1.25), deg2rad(1.2), 1])
+    ODEProblem(csys, st)
 end
 
 suite["NEI Simulator"] = BenchmarkGroup()
-sim = nei_simulator()
 st = SimulatorStrangSerial(Tsit5(), Euler(), 100.0)
-suite["NEI Simulator"]["Serial"] = @benchmarkable run!($sim, $st)
+sim = nei_simulator(st)
+suite["NEI Simulator"]["Serial"] = @benchmarkable solve($sim, dt=100.0,
+    save_on=false, save_start=false, save_end=false, initialize_save=false)
 
 st = SimulatorStrangThreads(Tsit5(), Euler(), 100.0)
-suite["NEI Simulator"]["Threads"] = @benchmarkable run!($sim, $st)
+sim = nei_simulator(st)
+suite["NEI Simulator"]["Threads"] = @benchmarkable solve($sim, dt=100.0,
+    save_on=false, save_start=false, save_end=false, initialize_save=false)
 
-tune!(suite, verbose = true)
-results = run(suite, verbose = true)
+tune!(suite, verbose=true)
+results = run(suite, verbose=true)
 
 BenchmarkTools.save("output.json", median(results))
diff --git a/src/netcdf_output.jl b/src/netcdf_output.jl
index 1eadd86c..aad125fe 100644
--- a/src/netcdf_output.jl
+++ b/src/netcdf_output.jl
@@ -35,15 +35,18 @@ mutable struct NetCDFOutputter
 end
 
 "Set up the output file and the callback function."
-function EarthSciMLBase.init_callback(nc::NetCDFOutputter, s::Simulator)
+function EarthSciMLBase.init_callback(nc::NetCDFOutputter, sys::EarthSciMLBase.CoupledSystem,
+    sys_mtk, obs_eqs, dom::EarthSciMLBase.DomainInfo)
+
     rm(nc.filepath, force=true)
     ds = NCDataset(nc.filepath, "c")
-    pv = EarthSciMLBase.pvars(s.sys.domaininfo)
+    grid = EarthSciMLBase.grid(dom)
+    pv = EarthSciMLBase.pvars(dom)
     @assert length(pv) == 3 "Currently only 3D simulations are supported."
-    @assert length(s.grid) == 3 "Currently only 3D simulations are supported."
+    @assert length(grid) == 3 "Currently only 3D simulations are supported."
     pvstr = [String(Symbol(p)) for p in pv]
     for (i, p) in enumerate(pvstr)
-        ds.dim[p] = length(s.grid[i])
+        ds.dim[p] = length(grid[i])
     end
     ds.dim["time"] = Inf
     function makencvar(v, dims)
@@ -53,7 +56,7 @@ function EarthSciMLBase.init_callback(nc::NetCDFOutputter, s::Simulator)
         ncvar.attrib["units"] = string(DynamicQuantities.dimension(ModelingToolkit.get_unit(v)))
         ncvar
     end
-    ncvars = [makencvar(v, [pvstr..., "time"]) for v in vcat(unknowns(s.sys_mtk), nc.extra_vars)]
+    ncvars = [makencvar(v, [pvstr..., "time"]) for v in vcat(unknowns(sys_mtk), nc.extra_vars)]
     nctvar = defVar(ds, "time", Float64, ("time",))
     nctvar.attrib["description"] = "Time"
     nctvar.attrib["units"] = "seconds since 1970-1-1"
@@ -61,17 +64,20 @@ function EarthSciMLBase.init_callback(nc::NetCDFOutputter, s::Simulator)
         d = defVar(ds, p, nc.dtype, (p,))
         d.attrib["description"] = ModelingToolkit.getdescription(pv[i])
         d.attrib["units"] = string(DynamicQuantities.dimension(ModelingToolkit.get_unit(pv[i])))
-        d[:] = s.grid[i]
+        d[:] = grid[i]
     end
-    for j in eachindex(nc.extra_vars)
-        push!(nc.extra_var_fs, s.obs_fs[s.obs_fs_idx[nc.extra_vars[j]]])
+    if length(nc.extra_vars) > 0
+        obs_fs = EarthSciMLBase.obs_functions(obs_eqs, dom)
+        for j in eachindex(nc.extra_vars)
+            push!(nc.extra_var_fs, obs_fs(nc.extra_vars[j]))
+        end
     end
     nc.file = ds
     nc.vars = ncvars
     nc.tvar = nctvar
-    nc.grid = s.grid
+    nc.grid = grid
     nc.h = 1
-    start, finish = EarthSciMLBase.tspan(s.domaininfo)
+    start, finish = EarthSciMLBase.tspan(dom)
     return PresetTimeCallback(start:nc.time_interval:finish,
         (integrator) -> affect!(nc, integrator),
         finalize=(c, u, t, integrator) -> close(nc.file),
@@ -81,7 +87,7 @@ function EarthSciMLBase.init_callback(nc::NetCDFOutputter, s::Simulator)
 end
 
 """
-Write the current state of the `Simulator` to the NetCDF file.
+Write the current state of the system to the NetCDF file.
 """
 function affect!(nc::NetCDFOutputter, integrator)
     u = reshape(integrator.u, length(nc.vars) - length(nc.extra_vars), [length(g) for g in nc.grid]...)
diff --git a/src/update_callback.jl b/src/update_callback.jl
index 3eb8595b..7537e275 100644
--- a/src/update_callback.jl
+++ b/src/update_callback.jl
@@ -17,12 +17,14 @@ function lazyload!(uc::UpdateCallbackCreator, t::AbstractFloat)
 end
 
 """
-Create a callback for this simulator. We only want to update the interpolators
+Create a callback for this system. We only want to update the interpolators
 that are actually used in the system.
 """
-function EarthSciMLBase.init_callback(uc::UpdateCallbackCreator, s::Simulator)
+function EarthSciMLBase.init_callback(uc::UpdateCallbackCreator, sys::EarthSciMLBase.CoupledSystem,
+    sys_mtk, obs_eqs, dom::EarthSciMLBase.DomainInfo)
+
     sysname = nameof(uc.sys)
-    needvars = Symbol.([unknowns(s.sys_mtk); [eq.lhs for eq in observed(s.sys_mtk)]])
+    needvars = Symbol.([unknowns(sys_mtk); [eq.lhs for eq in observed(sys_mtk)]])
     itps = []
     for (v, itp) in zip(uc.variables, uc.interpolators)
         nv = Symbol(sysname, "₊", v)
@@ -46,7 +48,7 @@ function EarthSciMLBase.init_callback(uc::UpdateCallbackCreator, s::Simulator)
     DiscreteCallback(
         (u, t, integrator) -> true, # TODO(CT): Could change to only run when we know interpolators need to be updated.
         update_callback,
-        initialize = initialize,
-        save_positions = (false, false),
+        initialize=initialize,
+        save_positions=(false, false),
     )
 end
diff --git a/test/geosfp_test.jl b/test/geosfp_test.jl
index 8d072545..92ac8600 100644
--- a/test/geosfp_test.jl
+++ b/test/geosfp_test.jl
@@ -1,8 +1,8 @@
-using Main.EarthSciData
+using EarthSciData
 using EarthSciMLBase
 using Test
 using Dates
-using ModelingToolkit, DomainSets
+using ModelingToolkit
 using ModelingToolkit: t, D
 using DynamicQuantities
 import NCDatasets
diff --git a/test/nei2016monthly_test.jl b/test/nei2016monthly_test.jl
index 73d2c24d..0df049c3 100644
--- a/test/nei2016monthly_test.jl
+++ b/test/nei2016monthly_test.jl
@@ -1,4 +1,4 @@
-using Main.EarthSciData
+using EarthSciData
 using Test
 using DynamicQuantities, EarthSciMLBase, ModelingToolkit
 using ModelingToolkit: t
@@ -105,9 +105,11 @@ end
 @testset "Coupling with GEOS-FP" begin
     gfp = GEOSFP("4x5", domain)
 
-    eqs = equations(convert(ODESystem, couple(emis, gfp)))
+    sys = convert(ODESystem, couple(emis, gfp))
+    structural_simplify(sys)
+    eqs = equations(sys)
 
-    @test occursin("NEI2016MonthlyEmis₊lat(t) ~ lat", string(eqs))
+    @test occursin("NEI2016MonthlyEmis₊lat(t) ~ GEOSFP₊lat", string(eqs))
 end
 
 @testset "wrong year" begin
diff --git a/test/netcdf_output_test.jl b/test/netcdf_output_test.jl
index 8b8a7772..995e2dad 100644
--- a/test/netcdf_output_test.jl
+++ b/test/netcdf_output_test.jl
@@ -18,7 +18,7 @@ eqs = [
     v ~ (x + y) * lev
 ]
 
-sys = ODESystem(eqs, t; name=:Test₊sys)
+sys = ODESystem(eqs, t; name=:sys)
 
 domain = DomainInfo(
     constIC(0.0, t ∈ Interval(0.0, 2.0)),
@@ -32,26 +32,27 @@ file = tempname() * ".nc"
 csys = couple(sys, domain)
 
 o = NetCDFOutputter(file, 1.0; extra_vars=[
-    structural_simplify(convert(ODESystem, csys)).Test₊sys₊v
+    structural_simplify(convert(ODESystem, csys)).sys₊v
 ])
 
 csys = couple(csys, o)
 
-sim = Simulator(csys)
-st = SimulatorStrangThreads(Tsit5(), Euler(), 0.01)
+dt = 0.01
+st = SolverStrangThreads(Tsit5(), dt)
+prob = ODEProblem(csys, st)
 
-run!(sim, st)
+solve(prob, Euler(), dt=dt)
 
 ds = NCDataset(file, "r")
 
-@test size(ds["Test₊sys₊u"], 4) == 3
-@test all(isapprox.(ds["Test₊sys₊u"][:, :, :, 1], 1.0, atol=0.011))
-@test sum(abs.(ds["Test₊sys₊v"][:, :, :, 1])) ≈ 5754.0f0
-@test all(isapprox.(ds["Test₊sys₊u"][:, :, :, 2], 2.0, atol=0.011))
-@test sum(abs.(ds["Test₊sys₊v"][:, :, :, 2])) ≈ 5754.0f0
-@test all(isapprox.(ds["Test₊sys₊u"][:, :, :, 3], 3.0, atol=0.011))
-@test sum(abs.(ds["Test₊sys₊v"][:, :, :, 3])) ≈ 5754.0f0
-@test size(ds["Test₊sys₊u"]) == (21, 41, 3, 3)
+@test size(ds["sys₊u"], 4) == 3
+@test all(isapprox.(ds["sys₊u"][:, :, :, 1], 1.0, atol=0.011))
+@test sum(abs.(ds["sys₊v"][:, :, :, 1])) ≈ 5754.0f0
+@test all(isapprox.(ds["sys₊u"][:, :, :, 2], 2.0, atol=0.011))
+@test sum(abs.(ds["sys₊v"][:, :, :, 2])) ≈ 5754.0f0
+@test all(isapprox.(ds["sys₊u"][:, :, :, 3], 3.0, atol=0.011))
+@test sum(abs.(ds["sys₊v"][:, :, :, 3])) ≈ 5754.0f0
+@test size(ds["sys₊u"]) == (21, 41, 3, 3)
 
 @test ds["time"][:] == [DateTime("1970-01-01T00:00:00"), DateTime("1970-01-01T00:00:01"), DateTime("1970-01-01T00:00:02")]
 @test ds["x"][:] ≈ -1.0:0.1:1.0
@@ -61,7 +62,7 @@ ds = NCDataset(file, "r")
 @test ds["x"].attrib["description"] == "x coordinate"
 @test ds["x"].attrib["units"] == "kg"
 
-@test ds["Test₊sys₊u"].attrib["description"] == "u value"
-@test ds["Test₊sys₊u"].attrib["units"] == "kg"
+@test ds["sys₊u"].attrib["description"] == "u value"
+@test ds["sys₊u"].attrib["units"] == "kg"
 
 rm(file, force=true)
diff --git a/test/update_callback_test.jl b/test/update_callback_test.jl
index f5732c68..7448652c 100644
--- a/test/update_callback_test.jl
+++ b/test/update_callback_test.jl
@@ -38,16 +38,18 @@ emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
 
 csys = couple(sys, emis, domain)
 
-sim = Simulator(csys)
+dt = 100.0
+st = SolverStrangSerial(Tsit5(), dt)
+prob = ODEProblem(csys, st)
 
-st = SimulatorStrangSerial(Tsit5(), Euler(), 100.0)
-
-sol = run!(sim, st)
+sol = solve(prob, Euler(), dt=dt)
 
 @test sum(sol.u[end]) ≈ 3.3756746955152187e-6
 
-st = SimulatorStrangThreads(Tsit5(), Euler(), 100.0)
+st = SolverStrangThreads(Tsit5(), dt)
+
+prob = ODEProblem(csys, st)
 
-sol = run!(sim, st)
+sol = solve(prob, Euler(), dt=dt)
 
 @test sum(sol.u[end]) ≈ 3.3756746955152187e-6

From c2e536413c5a90af263deade1b0d6eaac809b256 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Sun, 27 Oct 2024 22:11:50 -0500
Subject: [PATCH 3/5] Get rid of update callback

---
 docs/src/geosfp.md                            |  9 +---
 docs/src/nei2016.md                           |  2 +-
 src/EarthSciData.jl                           |  1 -
 src/geosfp.jl                                 | 11 ++--
 src/load.jl                                   | 36 ++++---------
 src/nei2016monthly.jl                         |  9 ++--
 src/update_callback.jl                        | 54 -------------------
 test/geosfp_test.jl                           | 23 ++++----
 test/load_test.jl                             |  4 +-
 test/nei2016monthly_test.jl                   |  8 +--
 test/runtests.jl                              |  3 +-
 ...{update_callback_test.jl => solve_test.jl} |  6 +--
 12 files changed, 39 insertions(+), 127 deletions(-)
 delete mode 100644 src/update_callback.jl
 rename test/{update_callback_test.jl => solve_test.jl} (91%)

diff --git a/docs/src/geosfp.md b/docs/src/geosfp.md
index 33455fb8..81855197 100644
--- a/docs/src/geosfp.md
+++ b/docs/src/geosfp.md
@@ -18,9 +18,7 @@ using DynamicQuantities: dimension
     lon, [unit=u"rad"], 
     lat, [unit=u"rad"],
 )
-geosfp, geosfp_updater = GEOSFP("4x5")
-
-geosfp
+geosfp = GEOSFP("4x5")
 ```
 
 Note that the [`GEOSFP`](@ref) function returns to things, an equation system and an object that can used to update the time in the underlying data loaders. 
@@ -59,13 +57,10 @@ domain = DomainInfo(
     zerogradBC(lev ∈ Interval(1.0f0, 11.0f0)),
 )
 
-composed_sys = couple(examplesys, domain, geosfp, geosfp_updater)
+composed_sys = couple(examplesys, domain, geosfp)
 pde_sys = convert(PDESystem, composed_sys)
 ```
 
-You can see above that we add both `geosfp` and `geosfp_updater` to our coupled system. 
-If we didn't include the updater, the resulting model would not give the correct results.
-
 Now, finally, we can run the simulation and plot the GEOS-FP wind fields in the result:
 
 (The code below is commented out because it is very slow right now. A faster solution is coming soon!)
diff --git a/docs/src/nei2016.md b/docs/src/nei2016.md
index 9f6705f2..2c23cf83 100644
--- a/docs/src/nei2016.md
+++ b/docs/src/nei2016.md
@@ -24,7 +24,7 @@ using EarthSciData, ModelingToolkit, DynamicQuantities, DataFrames
 using ModelingToolkit: t
 using DynamicQuantities: dimension
 @parameters lat, [unit=u"rad"], lon, [unit=u"rad"], lev [unit=u"rad"]
-emis, emis_updater = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", lon, lat, lev)
+emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", lon, lat, lev)
 ```
 
 ## Variables
diff --git a/src/EarthSciData.jl b/src/EarthSciData.jl
index ec64f92d..44047461 100644
--- a/src/EarthSciData.jl
+++ b/src/EarthSciData.jl
@@ -12,7 +12,6 @@ using Scratch
 # General utilities
 include("load.jl")
 include("utils.jl")
-include("update_callback.jl")
 
 # Specific data sets
 include("netcdf.jl")
diff --git a/src/geosfp.jl b/src/geosfp.jl
index c632140b..66c498c4 100644
--- a/src/geosfp.jl
+++ b/src/geosfp.jl
@@ -223,7 +223,7 @@ The native data type for this dataset is Float32.
 
 `stream_data` specifies whether the data should be streamed in as needed or loaded all at once.
 
-See http://geoschemdata.wustl.edu/ExtData/ for current options.
+See http://geoschemdata.wustl.edu/ExtData/ for current data domain options.
 """
 function GEOSFP(domain::AbstractString, domaininfo::DomainInfo; name=:GEOSFP, stream_data=true)
     filesets = Dict{String,GEOSFPFileSet}(
@@ -238,23 +238,21 @@ function GEOSFP(domain::AbstractString, domaininfo::DomainInfo; name=:GEOSFP, st
     pvdict = Dict([Symbol(v) => v for v in EarthSciMLBase.pvars(domaininfo)]...)
     eqs = Equation[]
     vars = Num[]
-    itps = []
     for (filename, fs) in filesets
         for varname ∈ varnames(fs, starttime)
             dt = EarthSciMLBase.dtype(domaininfo)
             itp = DataSetInterpolator{dt}(fs, varname, starttime, endtime,
                 domaininfo.spatial_ref; stream_data=stream_data)
-            dims = dimnames(itp, starttime)
+            dims = dimnames(itp)
             coords = Num[]
             for dim in dims
                 d = Symbol(dim)
                 @assert d ∈ keys(pvdict) "GEOSFP coordinate $d not found in domaininfo coordinates ($(pvs))."
                 push!(coords, pvdict[d])
             end
-            eq = create_interp_equation(itp, filename, t, starttime, coords)
+            eq = create_interp_equation(itp, filename, t, coords)
             push!(eqs, eq)
             push!(vars, eq.lhs)
-            push!(itps, itp)
         end
     end
 
@@ -269,8 +267,7 @@ function GEOSFP(domain::AbstractString, domaininfo::DomainInfo; name=:GEOSFP, st
 
     sys = ODESystem(eqs, t, vars, [pvdict[:lon], pvdict[:lat], lev], name=name,
         metadata=Dict(:coupletype => GEOSFPCoupler))
-    cb = UpdateCallbackCreator(sys, vars, itps)
-    return sys, cb
+    return sys
 end
 
 function EarthSciMLBase.couple2(mw::EarthSciMLBase.MeanWindCoupler, g::GEOSFPCoupler)
diff --git a/src/load.jl b/src/load.jl
index 5209bfc4..f870595e 100644
--- a/src/load.jl
+++ b/src/load.jl
@@ -133,11 +133,6 @@ data records for the times immediately before and after the current time step.
 `varname` is the name of the variable to interpolate. `default_time` is the time to use when initializing
 the interpolator. `spatial_ref` is the spatial reference system that the simulation will be using.
 `stream_data` specifies whether the data should be streamed in as needed or loaded all at once.
-
-!!! warning
-    WARNING: This interpolator includes a hack so that the function `deepcopy` does not copy the interpolator,
-    it just returns the interpolator.
-    This is necessary to allow us to update the interpolator from a callback, but it may cause unintended side effects.
 """
 mutable struct DataSetInterpolator{To,N,N2,FT,ITPT}
     fs::FileSet
@@ -192,7 +187,6 @@ mutable struct DataSetInterpolator{To,N,N2,FT,ITPT}
             metadata, times, DateTime(1, 1, 1), coord_trans,
             Channel{DateTime}(0), Channel(1), Channel{Int}(0),
             ReentrantLock(), false)
-        Threads.@spawn async_loader(itp)
         itp
     end
 end
@@ -205,10 +199,6 @@ function Base.show(io::IO, itp::DataSetInterpolator)
     print(io, "DataSetInterpolator{$(typeof(itp.fs)), $(itp.varname)}")
 end
 
-# FIXME: WARNING: This is a hack to make sure that the interpolator is not copied
-# so that we can update it from a callback. It may cause unintended side effects.
-Base.deepcopy_internal(itp::DataSetInterpolator, dict::IdDict) = itp
-
 """ Return the units of the data. """
 ModelingToolkit.get_unit(itp::DataSetInterpolator) = itp.metadata.units
 
@@ -330,6 +320,7 @@ function initialize!(itp::DataSetInterpolator, t::DateTime)
     if itp.initialized == false
         itp.load_cache = zeros(eltype(itp.load_cache), itp.metadata.varsize...)
         itp.data = zeros(eltype(itp.data), itp.metadata.varsize..., size(itp.data, length(size(itp.data)))) # Add a dimension for time.
+        Threads.@spawn async_loader(itp)
         itp.initialized = true
     end
     times = interp_cache_times!(itp, t) # Figure out which times we need.
@@ -382,30 +373,21 @@ $(SIGNATURES)
 
 Return the dimension names associated with this interpolator.
 """
-function dimnames(itp::DataSetInterpolator, t::DateTime)
-    lazyload!(itp, t)
-    itp.metadata.dimnames
-end
+dimnames(itp::DataSetInterpolator) = itp.metadata.dimnames
 
 """
 $(SIGNATURES)
 
 Return the units of the data associated with this interpolator.
 """
-function units(itp::DataSetInterpolator, t::DateTime)
-    lazyload!(itp, t)
-    itp.metadata.units
-end
+units(itp::DataSetInterpolator) = itp.metadata.units
 
 """
 $(SIGNATURES)
 
 Return the description of the data associated with this interpolator.
 """
-function description(itp::DataSetInterpolator, t::DateTime)
-    lazyload!(itp, t)
-    itp.metadata.description
-end
+description(itp::DataSetInterpolator) = itp.metadata.description
 
 """
 $(SIGNATURES)
@@ -472,20 +454,20 @@ Create an equation that interpolates the given dataset at the given time and loc
 `wrapper_f` can specify a function to wrap the interpolated value, for example `eq -> eq / 2`
 to divide the interpolated value by 2.
 """
-function create_interp_equation(itp::DataSetInterpolator, filename, t, sample_time, coords; wrapper_f=v -> v)
+function create_interp_equation(itp::DataSetInterpolator, filename, t, coords; wrapper_f=v -> v)
     # Create right hand side of equation.
     if length(coords) == 3
-        rhs = wrapper_f(interp_unsafe(itp, t, coords[1], coords[2], coords[3]))
+        rhs = wrapper_f(interp!(itp, t, coords[1], coords[2], coords[3]))
     elseif length(coords) == 2
-        rhs = wrapper_f(interp_unsafe(itp, t, coords[1], coords[2]))
+        rhs = wrapper_f(interp!(itp, t, coords[1], coords[2]))
     elseif length(coords) == 1
-        rhs = wrapper_f(interp_unsafe(itp, t, coords[1]))
+        rhs = wrapper_f(interp!(itp, t, coords[1]))
     else
         error("Unexpected number of coordinates: $(length(coords))")
     end
 
     # Create left hand side of equation.
-    desc = description(itp, sample_time)
+    desc = description(itp)
     uu = ModelingToolkit.get_unit(rhs)
     n = length(filename) > 0 ? Symbol("$(filename)₊$(itp.varname)") : Symbol("$(itp.varname)")
     lhs = only(@variables $n(t) [unit = uu, description = desc])
diff --git a/src/nei2016monthly.jl b/src/nei2016monthly.jl
index 37cdda6a..c5bee188 100644
--- a/src/nei2016monthly.jl
+++ b/src/nei2016monthly.jl
@@ -165,22 +165,19 @@ function NEI2016MonthlyEmis(sector::AbstractString, domaininfo::DomainInfo; scal
     )
     eqs = Equation[]
     vars = Num[]
-    itps = []
     for varname ∈ varnames(fs, starttime)
         dt = EarthSciMLBase.dtype(domaininfo)
         itp = DataSetInterpolator{dt}(fs, varname, starttime, endtime,
             domaininfo.spatial_ref; stream_data=stream_data)
-        @constants zero_emis = 0 [unit = units(itp, starttime) / u"m"]
+        @constants zero_emis = 0 [unit = units(itp) / u"m"]
         zero_emis = ModelingToolkit.unwrap(zero_emis) # Unsure why this is necessary.
-        eq = create_interp_equation(itp, "", t, starttime, [x, y],
+        eq = create_interp_equation(itp, "", t, [x, y],
             wrapper_f=(eq) -> ifelse(lev < 2, eq / Δz * scale, zero_emis),
         )
         push!(eqs, eq)
         push!(vars, eq.lhs)
-        push!(itps, itp)
     end
     sys = ODESystem(eqs, t, vars, [x, y, lev, Δz]; name=name,
         metadata=Dict(:coupletype => NEI2016MonthlyEmisCoupler))
-    cb = UpdateCallbackCreator(sys, vars, itps)
-    return sys, cb
+    return sys
 end
diff --git a/src/update_callback.jl b/src/update_callback.jl
deleted file mode 100644
index 7537e275..00000000
--- a/src/update_callback.jl
+++ /dev/null
@@ -1,54 +0,0 @@
-
-"""
-This struct holds information that can be used to create a callback function that updates the
-states of the interpolators.
-"""
-struct UpdateCallbackCreator
-    sys::ODESystem
-    variables
-    interpolators
-end
-
-function lazyload!(uc::UpdateCallbackCreator, t::AbstractFloat)
-    dt = unix2datetime(t)
-    for itp in uc.interpolators
-        lazyload!(itp, dt)
-    end
-end
-
-"""
-Create a callback for this system. We only want to update the interpolators
-that are actually used in the system.
-"""
-function EarthSciMLBase.init_callback(uc::UpdateCallbackCreator, sys::EarthSciMLBase.CoupledSystem,
-    sys_mtk, obs_eqs, dom::EarthSciMLBase.DomainInfo)
-
-    sysname = nameof(uc.sys)
-    needvars = Symbol.([unknowns(sys_mtk); [eq.lhs for eq in observed(sys_mtk)]])
-    itps = []
-    for (v, itp) in zip(uc.variables, uc.interpolators)
-        nv = Symbol(sysname, "₊", v)
-        if nv ∈ needvars
-            push!(itps, itp)
-        end
-    end
-    function initialize(c, u, t, integrator)
-        dt = unix2datetime(integrator.t)
-        for itp in itps
-            itp.initialized = false
-            lazyload!(itp, dt)
-        end
-    end
-    function update_callback(integrator)
-        dt = unix2datetime(integrator.t)
-        for itp in itps
-            lazyload!(itp, dt)
-        end
-    end
-    DiscreteCallback(
-        (u, t, integrator) -> true, # TODO(CT): Could change to only run when we know interpolators need to be updated.
-        update_callback,
-        initialize=initialize,
-        save_positions=(false, false),
-    )
-end
diff --git a/test/geosfp_test.jl b/test/geosfp_test.jl
index 92ac8600..8b2a3a49 100644
--- a/test/geosfp_test.jl
+++ b/test/geosfp_test.jl
@@ -15,7 +15,7 @@ domain = DomainInfo(DateTime(2022, 1, 1), DateTime(2022, 1, 3);
 @testset "GEOS-FP" begin
     lon, lat, lev = EarthSciMLBase.pvars(domain)
     @constants c_unit = 6.0 [unit = u"rad" description = "constant to make units cancel out"]
-    geosfp, _ = GEOSFP("4x5", domain)
+    geosfp = GEOSFP("4x5", domain)
 
     @test Symbol.(parameters(geosfp)) == [:lon, :lat, :lev]
 
@@ -37,9 +37,9 @@ domain = DomainInfo(DateTime(2022, 1, 1), DateTime(2022, 1, 3);
         "MeanWind₊v_lon(t, lon, lat, lev)", "GEOSFP₊A3dyn₊U(t, lon, lat, lev)",
         "MeanWind₊v_lat(t, lon, lat, lev)", "GEOSFP₊A3dyn₊V(t, lon, lat, lev)",
         "MeanWind₊v_lev(t, lon, lat, lev)", "GEOSFP₊A3dyn₊OMEGA(t, lon, lat, lev)",
-        "GEOSFP₊A3dyn₊U(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, U}, t, lon, lat, lev)",
-        "GEOSFP₊A3dyn₊OMEGA(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, OMEGA}, t, lon, lat, lev)",
-        "GEOSFP₊A3dyn₊V(t, lon, lat, lev)", "EarthSciData.interp_unsafe(DataSetInterpolator{EarthSciData.GEOSFPFileSet, V}, t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊U(t, lon, lat, lev)", "EarthSciData.interp!(DataSetInterpolator{EarthSciData.GEOSFPFileSet, U}, t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊OMEGA(t, lon, lat, lev)", "EarthSciData.interp!(DataSetInterpolator{EarthSciData.GEOSFPFileSet, OMEGA}, t, lon, lat, lev)",
+        "GEOSFP₊A3dyn₊V(t, lon, lat, lev)", "EarthSciData.interp!(DataSetInterpolator{EarthSciData.GEOSFPFileSet, V}, t, lon, lat, lev)",
         "Differential(t)(ExampleSys₊c(t, lon, lat, lev))", "Differential(lon)(ExampleSys₊c(t, lon, lat, lev)",
         "MeanWind₊v_lon(t, lon, lat, lev)", "lon2m",
         "Differential(lat)(ExampleSys₊c(t, lon, lat, lev)",
@@ -58,7 +58,7 @@ end
 
 @testset "GEOS-FP pressure levels" begin
     @parameters lat, [unit = u"rad"], lon, [unit = u"rad"], lev
-    geosfp, updater = GEOSFP("4x5", domain)
+    geosfp = GEOSFP("4x5", domain)
 
     # Rearrange pressure equation so it can be evaluated for P.
     iips = findfirst((x) -> x == :I3₊PS, [Symbolics.tosymbol(eq.lhs, escape=false) for eq in equations(geosfp)])
@@ -67,7 +67,6 @@ end
     peq = substitute(equations(geosfp)[iip], pseq.lhs => pseq.rhs)
 
     # Check Pressure levels
-    EarthSciData.lazyload!(updater, datetime2unix(DateTime(2022, 5, 1)))
     P = ModelingToolkit.subs_constants(peq.rhs)
     P_expr = build_function(P, [t, lon, lat, lev])
     mypf = eval(P_expr)
@@ -100,13 +99,12 @@ end
     starttime = datetime2unix(DateTime(2022, 5, 1, 23, 58))
     endtime = datetime2unix(DateTime(2022, 5, 2, 0, 3))
 
-    geosfp, updater = GEOSFP("4x5", domain)
+    geosfp = GEOSFP("4x5", domain)
 
     iips = findfirst((x) -> x == :I3₊PS, [Symbolics.tosymbol(eq.lhs, escape=false) for eq in equations(geosfp)])
     pseq = equations(geosfp)[iips]
     PS_expr = build_function(pseq.rhs, t, lon, lat, lev)
     psf = eval(PS_expr)
-    EarthSciData.lazyload!(updater, starttime)
     psf(starttime, 0.0, 0.0, 1.0)
 end
 
@@ -118,6 +116,11 @@ end
     )
     starttime = datetime2unix(DateTime(5000, 1, 1))
 
-    geosfp, updater = GEOSFP("4x5", domain)
-    @test_throws Base.Exception EarthSciData.lazyload!(updater, starttime)
+    geosfp = GEOSFP("4x5", domain)
+
+    iips = findfirst((x) -> x == :I3₊PS, [Symbolics.tosymbol(eq.lhs, escape=false) for eq in equations(geosfp)])
+    pseq = equations(geosfp)[iips]
+    PS_expr = build_function(pseq.rhs, t, lon, lat, lev)
+    psf = eval(PS_expr)
+    @test_throws Base.Exception psf(starttime, 0.0, 0.0, 1.0)
 end
diff --git a/test/load_test.jl b/test/load_test.jl
index 781d5b4b..bcae6508 100644
--- a/test/load_test.jl
+++ b/test/load_test.jl
@@ -1,4 +1,4 @@
-using Main.EarthSciData
+using EarthSciData
 using Dates
 using ModelingToolkit
 using Random
@@ -30,7 +30,7 @@ itp = EarthSciData.DataSetInterpolator{Float32}(fs, "U", t, te, spatial_ref)
 
 @test String(latexify(itp)) == "\$\\mathrm{GEOSFPFileSet}\\left( U \\right)\$"
 
-@test EarthSciData.dimnames(itp, t) == ["lon", "lat", "lev"]
+@test EarthSciData.dimnames(itp) == ["lon", "lat", "lev"]
 @test issetequal(EarthSciData.varnames(fs, t), ["U", "OMEGA", "RH", "DTRAIN", "V"])
 
 @testset "interpolation" begin
diff --git a/test/nei2016monthly_test.jl b/test/nei2016monthly_test.jl
index 0df049c3..dd73d200 100644
--- a/test/nei2016monthly_test.jl
+++ b/test/nei2016monthly_test.jl
@@ -18,7 +18,7 @@ sample_time = ts
 
 spatial_ref="+proj=longlat +datum=WGS84 +no_defs"
 
-emis, updater = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
+emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
 fileset = EarthSciData.NEI2016MonthlyEmisFileSet("mrggrid_withbeis_withrwc")
 
 eqs = equations(emis)
@@ -27,11 +27,6 @@ eqs = equations(emis)
 
 sample_time = DateTime(2016, 5, 1)
 
-@testset "can't deepcopy" begin
-    itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
-    deepcopy(itp) === itp
-end
-
 @testset "correct projection" begin
     itp = EarthSciData.DataSetInterpolator{Float32}(fileset, "NOX", ts, te, spatial_ref)
     @test interp!(itp, sample_time, deg2rad(-97.0f0), deg2rad(40.0f0)) ≈ 9.211331f-10
@@ -69,7 +64,6 @@ end
     sys = extend(ODESystem([eq], t, [], []; name=:test_sys), emis)
     sys = structural_simplify(sys)
     tt = Dates.datetime2unix(sample_time)
-    EarthSciData.lazyload!(updater, tt)
     prob = ODEProblem(sys, zeros(1), (tt, tt + 60.0), [lat => deg2rad(40.0), lon => deg2rad(-97.0), lev => 1.0])
     sol = solve(prob)
     @test 2 > sol.u[end][end] > 1
diff --git a/test/runtests.jl b/test/runtests.jl
index 847564e8..1adbf404 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,11 +1,10 @@
 using EarthSciData
 using Test, SafeTestsets
 
-
 @testset "EarthSciData.jl" begin
     @safetestset "load" begin include("load_test.jl") end
     @safetestset "geosfp" begin include("geosfp_test.jl") end
     @safetestset "nei2016monthly" begin include("nei2016monthly_test.jl") end
     @safetestset "NetCDFOutputter" begin include("netcdf_output_test.jl") end
-    @safetestset "Update Callback" begin include("update_callback_test.jl") end
+    @safetestset "Update Callback" begin include("solve_test.jl") end
 end
diff --git a/test/update_callback_test.jl b/test/solve_test.jl
similarity index 91%
rename from test/update_callback_test.jl
rename to test/solve_test.jl
index 7448652c..07762e62 100644
--- a/test/update_callback_test.jl
+++ b/test/solve_test.jl
@@ -1,4 +1,4 @@
-using Main.EarthSciData
+using EarthSciData
 
 using Test
 using EarthSciMLBase, ModelingToolkit
@@ -44,7 +44,7 @@ prob = ODEProblem(csys, st)
 
 sol = solve(prob, Euler(), dt=dt)
 
-@test sum(sol.u[end]) ≈ 3.3756746955152187e-6
+@test sum(sol.u[end]) ≈ 2.7782704702240893e-5
 
 st = SolverStrangThreads(Tsit5(), dt)
 
@@ -52,4 +52,4 @@ prob = ODEProblem(csys, st)
 
 sol = solve(prob, Euler(), dt=dt)
 
-@test sum(sol.u[end]) ≈ 3.3756746955152187e-6
+@test sum(sol.u[end]) ≈ 2.7782704702240893e-5

From d17dfa045d6b3706bffb71fc81cf9ff2f72b46f2 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Sun, 27 Oct 2024 22:21:55 -0500
Subject: [PATCH 4/5] Update docs compat

---
 docs/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index b4e86ff9..4fc14803 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -16,7 +16,7 @@ DifferentialEquations = "7"
 Documenter = "1"
 DomainSets = "0.7"
 DynamicQuantities = "0.13, 1"
-EarthSciMLBase = "0.16"
+EarthSciMLBase = "0.18"
 MethodOfLines = "0.11"
 ModelingToolkit = "9"
 Plots = "1"

From d58bae0110eea5dc8afcd26ce89138fb66489d28 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Mon, 28 Oct 2024 08:34:44 -0500
Subject: [PATCH 5/5] Add domain to docs

---
 docs/src/geosfp.md  | 24 +++++++++---------------
 docs/src/nei2016.md | 16 +++++++++++++---
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/docs/src/geosfp.md b/docs/src/geosfp.md
index 81855197..901f4e50 100644
--- a/docs/src/geosfp.md
+++ b/docs/src/geosfp.md
@@ -6,19 +6,19 @@ First, let's initialize some packages and set up the [GEOS-FP](@ref GEOSFP) equa
 
 ```@example geosfp
 using EarthSciData, EarthSciMLBase
-using DomainSets, ModelingToolkit, MethodOfLines, DifferentialEquations
+using ModelingToolkit, MethodOfLines, DifferentialEquations
 using ModelingToolkit: t, D
 using Dates, Plots, DataFrames
 using DynamicQuantities
 using DynamicQuantities: dimension
 
 # Set up system
-@parameters(
-    lev, 
-    lon, [unit=u"rad"], 
-    lat, [unit=u"rad"],
-)
-geosfp = GEOSFP("4x5")
+domain = DomainInfo(DateTime(2022, 1, 1), DateTime(2022, 1, 3);
+    latrange=deg2rad(-85.0f0):deg2rad(2):deg2rad(85.0f0),
+    lonrange=deg2rad(-180.0f0):deg2rad(2.5):deg2rad(175.0f0),
+    levelrange=1:11, dtype=Float32)
+
+geosfp = GEOSFP("4x5", domain)
 ```
 
 Note that the [`GEOSFP`](@ref) function returns to things, an equation system and an object that can used to update the time in the underlying data loaders. 
@@ -40,6 +40,8 @@ To fix this, we create another equation system that is an ODE.
 
 ```@example geosfp
 function Example()
+    @parameters lat=0.0 [unit=u"rad"]
+    @parameters lon=0.0 [unit=u"rad"]
     @variables c(t) = 5.0 [unit=u"s"]
     ODESystem([D(c) ~ sin(lat * 6) + sin(lon * 6)], t, name=:Docs₊Example)
 end
@@ -49,14 +51,6 @@ examplesys = Example()
 Now, let's couple these two systems together, and also add in advection and some information about the domain:
 
 ```@example geosfp
-domain = DomainInfo(
-    partialderivatives_δxyδlonlat,
-    constIC(0.0, t ∈ Interval(Dates.datetime2unix(DateTime(2022, 1, 1)), Dates.datetime2unix(DateTime(2022, 1, 3)))),
-    zerogradBC(lat ∈ Interval(deg2rad(-80.0f0), deg2rad(80.0f0))),
-    periodicBC(lon ∈ Interval(deg2rad(-180.0f0), deg2rad(180.0f0))),
-    zerogradBC(lev ∈ Interval(1.0f0, 11.0f0)),
-)
-
 composed_sys = couple(examplesys, domain, geosfp)
 pde_sys = convert(PDESystem, composed_sys)
 ```
diff --git a/docs/src/nei2016.md b/docs/src/nei2016.md
index 2c23cf83..9a96c654 100644
--- a/docs/src/nei2016.md
+++ b/docs/src/nei2016.md
@@ -20,11 +20,21 @@ export JULIA_NO_VERIFY_HOSTS=gaftp.epa.gov
 This is what its equation system looks like:
 
 ```@example nei2016
-using EarthSciData, ModelingToolkit, DynamicQuantities, DataFrames
+using EarthSciData, EarthSciMLBase
+using ModelingToolkit, DynamicQuantities, DataFrames
 using ModelingToolkit: t
 using DynamicQuantities: dimension
-@parameters lat, [unit=u"rad"], lon, [unit=u"rad"], lev [unit=u"rad"]
-emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", lon, lat, lev)
+using Dates
+
+domain = DomainInfo(
+        DateTime(2016, 5, 1), DateTime(2016, 5, 2);
+        lonrange=deg2rad(-115):deg2rad(2.5):deg2rad(-68.75),
+        latrange=deg2rad(25):deg2rad(2):deg2rad(53.7),
+        levelrange=1:10, 
+        dtype=Float32,
+)
+
+emis = NEI2016MonthlyEmis("mrggrid_withbeis_withrwc", domain)
 ```
 
 ## Variables