Merge pull request #38 from PALEOtoolkit/simplify_carbchem

Simplify carbonate chemistry
PALEOtoolkit · Dec 17, 2024 · 2bab0b2 · 2bab0b2
2 parents 23d5043 + 14208df
commit 2bab0b2
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 28 deletions.
diff --git a/src/CarbChem.jl b/src/CarbChem.jl
@@ -264,7 +264,7 @@ function prepare_do_carbchem(m::PB.ReactionMethod, (vars, input_concs, outputs))
     # create per-thread buffers for temporary results
     Constants_tbuf  = [Vector{BufType}(undef, length(PALEOcarbchem.CNameIdx)) for t in 1:Threads.nthreads()]    
     res_tbuf    = [Vector{BufType}(undef, length(PALEOcarbchem.RNameIdx)) for t in 1:Threads.nthreads()]
-    input_concs_cell_tbuf = [PB.IteratorUtils.named_tuple_ref(keys(input_concs), BufType) for t in 1:Threads.nthreads()]
+    input_concs_cell_tbuf = [Vector{BufType}(undef, length(input_concs)) for t in 1:Threads.nthreads()]
 
     buffers = (Constants_tbuf, res_tbuf, input_concs_cell_tbuf)
 
@@ -290,7 +290,7 @@ function do_carbchem(
 
     Constants           = Constants_tbuf[Threads.threadid()] # get per thread buffer. threadid()=1 for a single-threaded app
     res                 = res_tbuf[Threads.threadid()]
-    input_concs_cell    = input_concs_cell_tbuf[Threads.threadid()]
+    input_concs_cell_v  = input_concs_cell_tbuf[Threads.threadid()]
 
     BufType = eltype(Constants)
 
@@ -314,15 +314,12 @@ function do_carbchem(
         )
 
         # Set 'input_concs_cell' with 'input_concs' for this cell
-        # NB: a Julia bug (v1.7), map allocates so can't just create a NamedTuple
-        # workaround using a 'mutable NamedTuple' buffer 'input_concs_cell' which has same fields as 'input_concs'
-        for (ic_cell, ic) in PB.IteratorUtils.zipstrict(values(input_concs_cell), values(input_concs))
-           ic_cell[] = r_rhofac*PB.SIMDutils.vgatherind(BufType, ic, i)
+        # To avoid allocations, first write into a Vector, then create a NamedTuple
+        for (j, ic) in enumerate(values(input_concs))
+           input_concs_cell_v[j] = r_rhofac*PB.SIMDutils.vgatherind(BufType, ic, i)
         end
-        # TODO allocates ? (Julia 1.9.0-rc2)
-        # PB.IteratorUtils.foreach_tuple(values(input_concs_cell), values(input_concs)) do ic_cell, ic
-        #    ic_cell[] = r_rhofac*PB.SIMDutils.vgatherind(BufType, ic, i)
-        # end
+        input_concs_cell = NamedTuple{keys(input_concs), NTuple{length(input_concs), BufType}}(input_concs_cell_v)
+
 
         if solve_pH == Val(true)  # this will be a compile-time switch as solve_pH Type is part of method signature
                     #   mol       / m^3         / kg m-3    
@@ -366,8 +363,8 @@ function do_carbchem(
                     pHfree, 
                     Val(3),  # scalein - Free
                     Val(1),  # scaleout - Total
-                    input_concs_cell.TS[], 
-                    input_concs_cell.TF[],
+                    input_concs_cell.TS, 
+                    input_concs_cell.TF,
                 ),
                 vars.pHtot, 
                 i,
@@ -389,4 +386,5 @@ function do_carbchem(
     return nothing
 end
 
+
 end
diff --git a/src/PALEOcarbchem/PALEOcarbchem.jl b/src/PALEOcarbchem/PALEOcarbchem.jl
@@ -38,7 +38,7 @@ julia> modern_concs = PALEOcarbchem.calc_modern_default_concs(35.0, Options=opti
 julia> map(x -> @sprintf("%.14e", x), modern_concs)
 (TF = "6.83258396883673e-05", TS = "2.82354341328601e-02", TB = "4.15700000000000e-04", Ca = "1.02845697008497e-02")
 
-julia> input_concs = (TCi=[2000e-6], TS=[modern_concs.TS], TF=[modern_concs.TF], TSi=[1e-3], TP=[1e-6],  TB=[modern_concs.TB], TH2S=[1e-6], TNH3=[1e-6], Ca=[modern_concs.Ca]);  # all in mol kg-1
+julia> input_concs = (TCi=2000e-6, TS=modern_concs.TS, TF=modern_concs.TF, TSi=1e-3, TP=1e-6,  TB=modern_concs.TB, TH2S=1e-6, TNH3=1e-6, Ca=modern_concs.Ca);  # all in mol kg-1
 
 julia> res = zeros(length(PALEOcarbchem.ResultNames));
 

diff --git a/src/PALEOcarbchem/Solvers.jl b/src/PALEOcarbchem/Solvers.jl
@@ -98,7 +98,7 @@ Calculate TAlk, and speciation, given pH and conserved concentrations (total DIC
 # Arguments:
 - `res`:  (output) Vector `res` of length length([`ResultNames`](@ref)) with details of TA contributions etc
 - `C`: constants from [`calc_constants!`](@ref). NB: must be on Free pH scale.
-- `concs::NamedTuple`: (mol kg-sw, NB: each element should be a Ref or length 1 Vector)  total concentrations for sulphate, fluoride, and each optional component of alkalinity enabled in `C`
+- `concs::NamedTuple`: (mol kg-sw)  total concentrations for sulphate, fluoride, and each optional component of alkalinity enabled in `C`
 - `pHfree`: pH on free scale
 
 # Implementation
@@ -131,7 +131,7 @@ function calculateTAfromTCpHfree!(
 
     # sulphate 
     if Options.Components.S == Val(true)
-        TS      = concs.TS[]
+        TS      = concs.TS
         HSO4    = TS/(1 + C[CI.KS]/H);# ' since KS is on the free scale
         TA      += -HSO4
         r[RI.TS]= TS; r[RI.HSO4]=HSO4       
@@ -143,7 +143,7 @@ function calculateTAfromTCpHfree!(
 
     # fluoride
     if Options.Components.F == Val(true)
-        TF      = concs.TF[]
+        TF      = concs.TF
         HF      = TF/(1 + C[CI.KF]/H);# ' since KF is on the free scale
         TA      += -HF
         r[RI.TF] = TF; r[RI.HF] = HF
@@ -155,7 +155,7 @@ function calculateTAfromTCpHfree!(
 
     # Inorganic carbon
     if Options.Components.Ci == Val(true)
-        TCi = concs.TCi[]
+        TCi = concs.TCi
         # Inorganic carbon
         CDenom      = (C[CI.K1]*H + H*H + C[CI.K1]*C[CI.K2]);
         HCO3      = TCi*C[CI.K1]*H  / CDenom;
@@ -184,7 +184,7 @@ function calculateTAfromTCpHfree!(
 
     # Boron
     if Options.Components.B == Val(true)
-        TB        = concs.TB[]
+        TB        = concs.TB
         BAlk      = TB*C[CI.KB]/(C[CI.KB] + H)
         TA        += BAlk
         r[RI.TB]=TB; r[RI.BAlk]=BAlk
@@ -196,7 +196,7 @@ function calculateTAfromTCpHfree!(
 
     # Phosphate
     if Options.Components.P == Val(true)
-        TP = concs.TP[]
+        TP = concs.TP
         PhosBot     = H*H*H + C[CI.KP1]*H*H + C[CI.KP1]*C[CI.KP2]*H + C[CI.KP1]*C[CI.KP2]*C[CI.KP3]
         H3PO4     = TP*H*H*H/PhosBot;        # neutral    Alk -1
         H2PO4     = TP*C[CI.KP1]*H*H/PhosBot;       # 1-         zero level of protons
@@ -227,7 +227,7 @@ function calculateTAfromTCpHfree!(
 
     # Silicate
     if Options.Components.Si == Val(true)
-        TSi = concs.TSi[]
+        TSi = concs.TSi
         SiAlk     = TSi*C[CI.KSi]/(C[CI.KSi] + H)
         TA      += SiAlk
         if do_dTAdpH == Val(true)
@@ -239,7 +239,7 @@ function calculateTAfromTCpHfree!(
 
     # Sulphide
     if Options.Components.H2S == Val(true)
-        TH2S = concs.TH2S[]
+        TH2S = concs.TH2S
         HSAlk     = TH2S*C[CI.KH2S]/(C[CI.KH2S] + H)
         H2S       = TH2S - HSAlk
         TA       += HSAlk
@@ -252,7 +252,7 @@ function calculateTAfromTCpHfree!(
 
     # Ammonia
     if Options.Components.NH3 == Val(true)
-        TNH3      = concs.TNH3[]
+        TNH3      = concs.TNH3
         NH3Alk    = TNH3*C[CI.KNH3]/(C[CI.KNH3] + H)
         NH4       = TNH3 - NH3Alk
         TA        += NH3Alk
@@ -265,7 +265,7 @@ function calculateTAfromTCpHfree!(
 
     # Carbonate saturation (not a contribution to alkalinity !)
     if Options.Components.Omega == Val(true)
-        Ca      = concs.Ca[]
+        Ca      = concs.Ca
         (OmegaCA, OmegaAR) = calculateOmega(C, CO3, Ca)
         r[RI.Ca]=Ca; r[RI.OmegaCA]=OmegaCA; r[RI.OmegaAR]=OmegaAR
     end

diff --git a/test/PALEOcarbchem/runCarbChembench.jl b/test/PALEOcarbchem/runCarbChembench.jl
@@ -51,7 +51,7 @@ function bench_TAfromTCpHfree()
     PALEOcarbchem.calc_constants!(C, TempC, Pdbar, Sal, Options=Options)
     modernC = PALEOcarbchem.calc_modern_default_concs(Sal, Options=Options)
 
-    Concs = (TCi=[2000e-6], TS=[modernC.TS], TF=[modernC.TF], TSi=[1e-3], TP=[1e-6],  TB=[modernC.TB], TH2S=[1e-6], TNH3=[1e-6], Ca=[modernC.Ca])
+    Concs = (TCi=2000e-6, TS=modernC.TS, TF=modernC.TF, TSi=1e-3, TP=1e-6,  TB=modernC.TB, TH2S=1e-6, TNH3=1e-6, Ca=modernC.Ca)
 
     resarray = Vector{Float64}(undef, length(PALEOcarbchem.RNameIdx))
     println("do_dTAdpH=false")
@@ -68,7 +68,7 @@ function bench_pHfromTATC()
     PALEOcarbchem.calc_constants!(C, TempC, Pdbar, Sal, Options=Options)
     modernC = PALEOcarbchem.calc_modern_default_concs(Sal, Options=Options)
 
-    Concs = (TCi=[2000e-6], TS=[modernC.TS], TF=[modernC.TF], TSi=[1e-3], TP=[1e-6],  TB=[modernC.TB], TH2S=[1e-6], TNH3=[1e-6], Ca=[modernC.Ca])
+    Concs = (TCi=2000e-6, TS=modernC.TS, TF=modernC.TF, TSi=1e-3, TP=1e-6,  TB=modernC.TB, TH2S=1e-6, TNH3=1e-6, Ca=modernC.Ca)
 
     resarray = Vector{Float64}(undef, length(PALEOcarbchem.RNameIdx))
 
@@ -116,7 +116,7 @@ function bench_TAfromTCpHfree_simd()
     PALEOcarbchem.calc_constants!(C, TempC, Pdbar, Sal, Options=Options)
     modernC = PALEOcarbchem.calc_modern_default_concs(Sal, Options=Options)
 
-    Concs = (TCi=[ES(2000e-6)], TS=[modernC.TS], TF=[modernC.TF], TSi=[ES(1e-3)], TP=[ES(1e-6)],  TB=[modernC.TB], TH2S=[ES(1e-6)], TNH3=[ES(1e-6)], Ca=[modernC.Ca])
+    Concs = (TCi=ES(2000e-6), TS=modernC.TS, TF=modernC.TF, TSi=ES(1e-3), TP=ES(1e-6),  TB=modernC.TB, TH2S=ES(1e-6), TNH3=ES(1e-6), Ca=modernC.Ca)
 
     pHstart = ES(8.0)
     resarray = Vector{ES}(undef, length(PALEOcarbchem.RNameIdx))
@@ -135,7 +135,7 @@ function bench_pHfromTATC_simd()
     PALEOcarbchem.calc_constants!(C, TempC, Pdbar, Sal, Options=Options)
     modernC = PALEOcarbchem.calc_modern_default_concs(Sal, Options=Options)
 
-    Concs = (TCi=[ES(2000e-6)], TS=[modernC.TS], TF=[modernC.TF], TSi=[ES(1e-3)], TP=[ES(1e-6)],  TB=[modernC.TB], TH2S=[ES(1e-6)], TNH3=[ES(1e-6)], Ca=[modernC.Ca])
+    Concs = (TCi=ES(2000e-6), TS=modernC.TS, TF=modernC.TF, TSi=ES(1e-3), TP=ES(1e-6),  TB=modernC.TB, TH2S=ES(1e-6), TNH3=ES(1e-6), Ca=modernC.Ca)
 
     resarray = Vector{ES}(undef, length(PALEOcarbchem.RNameIdx))
 

diff --git a/test/PALEOcarbchem/runCarbChemtests.jl b/test/PALEOcarbchem/runCarbChemtests.jl
@@ -133,7 +133,7 @@ end
                 xCO2dryinp= 0.000600944914072491)
 
     CI = PALEOcarbchem.CNameIdx
-    Concs = (TCi=[2000e-6], TS=modernC.TS, TF=modernC.TF, TSi=[1e-3], TP=[1e-6],  TB=modernC.TB, TH2S=[1e-6], TNH3=[1e-6], Ca=modernC.Ca)
+    Concs = (TCi=2000e-6, TS=modernC.TS, TF=modernC.TF, TSi=1e-3, TP=1e-6,  TB=modernC.TB, TH2S=1e-6, TNH3=1e-6, Ca=modernC.Ca)
 
     resarray = Vector{Float64}(undef, length(PALEOcarbchem.RNameIdx))
     (TA, dTAdpH) = PALEOcarbchem.calculateTAfromTCpHfree!(resarray, C, Options, Concs, 8.0, do_dTAdpH=Val(true))