Earth sci ml basev10

Project.toml

@@ -10,7 +10,6 @@ IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
-SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
@@ -23,8 +22,8 @@ GasChemExt = "GasChem"
 
 [compat]
 DifferentialEquations = "7"
-EarthSciMLBase = "0.6, 0.8"
-GasChem = "0.5"
+EarthSciMLBase = "0.10"
+GasChem = "0.6"
 IfElse = "0.1"
 ModelingToolkit = "8"
 SafeTestsets = "0.1"

docs/src/dry_deposition.md

@@ -18,7 +18,7 @@ model = DrydepositionG(t)
 ```
 Before running any simulations with the model we need to convert it into a system of differential equations.
 ```julia
-sys = structural_simplify(get_mtk(model))
+sys = structural_simplify(model)
 tspan = (0.0, 3600*24) 
 u0 = [2.0,10.0,5,5,2.34,0.15] # initial concentrations of SO₂, O₃, NO₂, NO, H₂O₂, CH₂O
 sol = solve(ODEProblem(sys, u0, tspan, []),AutoTsit5(Rosenbrock23()), saveat=10.0) # default parameters
@@ -34,7 +34,7 @@ using Unitful
 @parameters t [unit = u"s", description="Time"]
 model = DrydepositionG(t)
 
-sys = structural_simplify(get_mtk(model))
+sys = structural_simplify(model)
 tspan = (0.0, 3600*24) 
 u0 = [2.0,10.0,5,5,2.34,0.15] # initial concentrations of SO₂, O₃, NO₂, NO, H₂O₂, CH₂O
 sol = solve(ODEProblem(sys, u0, tspan, []),AutoTsit5(Rosenbrock23()), saveat=10.0) # default parameters

docs/src/emep.md

@@ -16,7 +16,7 @@ model = Wetdeposition(t)
 
 Before running any simulations with the model we need to convert it into a system of differential equations.
 ```julia 
-sys = structural_simplify(get_mtk(model))
+sys = structural_simplify(model)
 tspan = (0.0, 3600*24)
 u0 = [2.0,10.0,5,1400,275,50,0.15]  # initial concentration of SO₂, O₃, NO₂, CH₄, CO, DMS, ISOP
 prob = ODEProblem(sys, u0, tspan, [])
@@ -33,7 +33,7 @@ using Unitful
 @parameters t [unit = u"s", description="Time"]
 model = Wetdeposition(t)
 
-sys = structural_simplify(get_mtk(model))
+sys = structural_simplify(model)
 tspan = (0.0, 3600*24)
 u0 = [2.0,10.0,5,1400,275,50,0.15]  # initial concentration of SO₂, O₃, NO₂, CH₄, CO, DMS, ISOP
 prob = ODEProblem(sys, u0, tspan, [])

ext/GasChemExt.jl

@@ -1,10 +1,49 @@
 module GasChemExt
 
-using AtmosphericDeposition, GasChem, EarthSciMLBase
-Base.:(+)(w::Wetdeposition, b::SuperFast) = operator_compose(b, w)
-Base.:(+)(b::SuperFast, w::Wetdeposition) = w + b
+using AtmosphericDeposition, GasChem, EarthSciMLBase, Unitful, ModelingToolkit
 
-Base.:(+)(d::DrydepositionG, b::SuperFast) = operator_compose(b, d)
-Base.:(+)(b::SuperFast, d::DrydepositionG) = d + b
+export register_couplings_ext
+
+# Use a global flag to track initialization and ensure that the coupling between SuperFast and NEI Emission is only initialized once
+const couplings_registered_ext = Ref(false)
+
+function register_couplings_ext()
+    if couplings_registered_ext[]
+        println("Couplings have already been registered.")
+        return
+    end
+
+    println("Registering couplings in ext")
+    @parameters t [unit = u"s"]
+
+    register_coupling(SuperFast(t), Wetdeposition(t)) do c, e
+        operator_compose(convert(ODESystem, c), e, Dict(
+            c.SO2 => e.SO2,
+            c.NO2 => e.NO2,
+            c.O3 => e.O3,
+            c.CH4 => e.CH4,
+            c.CO => e.CO,
+            c.DMS => e.DMS,
+            c.ISOP => e.ISOP))
+    end
+
+    register_coupling(SuperFast(t), DrydepositionG(t)) do c, e
+        operator_compose(convert(ODESystem, c), e, Dict(
+            c.SO2 => e.SO2,
+            c.NO2 => e.NO2,
+            c.O3 => e.O3,
+            c.NO => e.NO,
+            c.H2O2 => e.H2O2,
+            c.CH2O => e.CH2O))
+    end
+
+    couplings_registered_ext[] = true
+    println("Coupling registry after registration: ", EarthSciMLBase.coupling_registry)
+end
+
+function __init__()
+    println("Initializing EmissionExt module")
+    register_couplings_ext()
+end
 
 end

src/dry_deposition.jl

@@ -236,41 +236,38 @@ Build Drydeposition model (gas)
 	d = DrydepositionG(t)
 ```
 """
-struct DrydepositionG <: EarthSciMLODESystem
-    sys::ODESystem
-    function DrydepositionG(t)
-        iSeason = 1
-        iLandUse = 10
-        rain = false
-        dew = false
-        @parameters z = 50 [unit = u"m", description = "top of the surface layer"]
-        @parameters z₀ = 0.04 [unit = u"m", description = "roughness lenght"]
-        @parameters u_star = 0.44 [unit = u"m/s", description = "friction velocity"]
-        @parameters L = 0 [unit = u"m", description = "Monin-Obukhov length"]
-        @parameters ρA = 1.2 [unit = u"kg*m^-3", description = "air density"]
-        @parameters G = 300 [unit = u"W*m^-2", description = "solar irradiation"]
-        @parameters T = 298 [unit = u"K", description = "temperature"]
-        @parameters θ = 0 [description = "slope of the local terrain, in unit radians"]
+function DrydepositionG(t)
+    iSeason = 1
+    iLandUse = 10
+    rain = false
+    dew = false
+    @parameters z = 50 [unit = u"m", description = "top of the surface layer"]
+    @parameters z₀ = 0.04 [unit = u"m", description = "roughness lenght"]
+    @parameters u_star = 0.44 [unit = u"m/s", description = "friction velocity"]
+    @parameters L = 0 [unit = u"m", description = "Monin-Obukhov length"]
+    @parameters ρA = 1.2 [unit = u"kg*m^-3", description = "air density"]
+    @parameters G = 300 [unit = u"W*m^-2", description = "solar irradiation"]
+    @parameters T = 298 [unit = u"K", description = "temperature"]
+    @parameters θ = 0 [description = "slope of the local terrain, in unit radians"]
 
-        D = Differential(t)
+    D = Differential(t)
 
-        @variables SO2(t) [unit = u"ppb"]
-        @variables O3(t) [unit = u"ppb"]
-        @variables NO2(t) [unit = u"ppb"]
-        @variables NO(t) [unit = u"ppb"]
-        @variables H2O2(t) [unit = u"ppb"]
-        @variables CH2O(t) [unit = u"ppb"]
+    @variables SO2(t) [unit = u"ppb"]
+    @variables O3(t) [unit = u"ppb"]
+    @variables NO2(t) [unit = u"ppb"]
+    @variables NO(t) [unit = u"ppb"]
+    @variables H2O2(t) [unit = u"ppb"]
+    @variables CH2O(t) [unit = u"ppb"]
 
-        eqs = [
-            D(SO2) ~ -DryDepGas(z, z₀, u_star, L, ρA, So2Data, G, T, θ, iSeason, iLandUse, rain, dew, true, false) / z * SO2
-            D(O3) ~ -DryDepGas(z, z₀, u_star, L, ρA, O3Data, G, T, θ, iSeason, iLandUse, rain, dew, false, true) / z * O3
-            D(NO2) ~ -DryDepGas(z, z₀, u_star, L, ρA, No2Data, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * NO2
-            D(NO) ~ -DryDepGas(z, z₀, u_star, L, ρA, NoData, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * NO
-            D(H2O2) ~ -DryDepGas(z, z₀, u_star, L, ρA, H2o2Data, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * H2O2
-            D(CH2O) ~ -DryDepGas(z, z₀, u_star, L, ρA, HchoData, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * CH2O
-        ]
+    eqs = [
+        D(SO2) ~ -DryDepGas(z, z₀, u_star, L, ρA, So2Data, G, T, θ, iSeason, iLandUse, rain, dew, true, false) / z * SO2
+        D(O3) ~ -DryDepGas(z, z₀, u_star, L, ρA, O3Data, G, T, θ, iSeason, iLandUse, rain, dew, false, true) / z * O3
+        D(NO2) ~ -DryDepGas(z, z₀, u_star, L, ρA, No2Data, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * NO2
+        D(NO) ~ -DryDepGas(z, z₀, u_star, L, ρA, NoData, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * NO
+        D(H2O2) ~ -DryDepGas(z, z₀, u_star, L, ρA, H2o2Data, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * H2O2
+        D(CH2O) ~ -DryDepGas(z, z₀, u_star, L, ρA, HchoData, G, T, θ, iSeason, iLandUse, rain, dew, false, false) / z * CH2O
+    ]
 
-        new(ODESystem(eqs, t, [SO2, O3, NO2, NO, H2O2, CH2O], [z, z₀, u_star, L, ρA, G, T, θ]; name=:DrydepositionG))
-    end
+    ODESystem(eqs, t, [SO2, O3, NO2, NO, H2O2, CH2O], [z, z₀, u_star, L, ρA, G, T, θ]; name=:DrydepositionG)
 end
 

src/wet_deposition.jl

@@ -57,34 +57,31 @@ Build Wetdeposition model
 	wd = Wetdeposition(t)
 ```
 """
-struct Wetdeposition <: EarthSciMLODESystem
-    sys::ODESystem
-    function Wetdeposition(t)
-        @parameters cloudFrac = 0.5 [description = "fraction of grid cell covered by clouds"]
-        @parameters qrain = 0.5 [description = "rain mixing ratio"]
-        @parameters ρ_air = 1.204 [unit = u"kg*m^-3", description = "air density"]
-        @parameters Δz = 200 [unit = u"m", description = "fall distance"]
+function Wetdeposition(t)
+    @parameters cloudFrac = 0.5 [description = "fraction of grid cell covered by clouds"]
+    @parameters qrain = 0.5 [description = "rain mixing ratio"]
+    @parameters ρ_air = 1.204 [unit = u"kg*m^-3", description = "air density"]
+    @parameters Δz = 200 [unit = u"m", description = "fall distance"]
 
-        D = Differential(t)
+    D = Differential(t)
 
-        @variables SO2(t) [unit = u"ppb"]
-        @variables O3(t) [unit = u"ppb"]
-        @variables NO2(t) [unit = u"ppb"]
-        @variables CH4(t) [unit = u"ppb"]
-        @variables CO(t) [unit = u"ppb"]
-        @variables DMS(t) [unit = u"ppb"]
-        @variables ISOP(t) [unit = u"ppb"]
+    @variables SO2(t) [unit = u"ppb"]
+    @variables O3(t) [unit = u"ppb"]
+    @variables NO2(t) [unit = u"ppb"]
+    @variables CH4(t) [unit = u"ppb"]
+    @variables CO(t) [unit = u"ppb"]
+    @variables DMS(t) [unit = u"ppb"]
+    @variables ISOP(t) [unit = u"ppb"]
 
-        eqs = [
-            D(SO2) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[2] * SO2
-            D(O3) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * O3
-            D(NO2) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * NO2
-            D(CH4) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * CH4
-            D(CO) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * CO
-            D(DMS) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * DMS
-            D(ISOP) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * ISOP
-        ]
+    eqs = [
+        D(SO2) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[2] * SO2
+        D(O3) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * O3
+        D(NO2) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * NO2
+        D(CH4) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * CH4
+        D(CO) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * CO
+        D(DMS) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * DMS
+        D(ISOP) ~ -WetDeposition(cloudFrac, qrain, ρ_air, Δz)[3] * ISOP
+    ]
 
-        new(ODESystem(eqs, t, [SO2, O3, NO2, CH4, CO, DMS, ISOP], [cloudFrac, qrain, ρ_air, Δz]; name=:Wetdeposition))
-    end
-end
+    ODESystem(eqs, t, [SO2, O3, NO2, CH4, CO, DMS, ISOP], [cloudFrac, qrain, ρ_air, Δz]; name=:WetDeposition)
+end

test/connector_test.jl

@@ -5,8 +5,12 @@ using Test, Unitful, ModelingToolkit, GasChem, Dates, EarthSciMLBase
     ModelingToolkit.check_units(eqs...) = nothing
     start = Dates.datetime2unix(Dates.DateTime(2016, 5, 1))
     @parameters t
-    composed_ode = SuperFast(t) + FastJX(t) + DrydepositionG(t) + Wetdeposition(t)
+    composed_ode = couple(SuperFast(t), FastJX(t), DrydepositionG(t), Wetdeposition(t))
     tspan = (start, start+3600*24*3)
     sys = structural_simplify(get_mtk(composed_ode))
     @test length(states(sys)) ≈ 18
+
+    eqs = string(equations(sys))
+    wanteqs = ["Differential(t)(superfast₊O3(t)) ~ superfast₊DrydepositionG_ddt_O3ˍt(t) + superfast₊WetDeposition_ddt_O3ˍt(t)"]
+    @test contains(string(eqs), wanteqs[1])
 end