Source code for esbmtk.post_processing

"""esbmtk: A general purpose Earth Science box model toolkit.

Copyright (C), 2020 Ulrich G.  Wortmann

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or (at
your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

from __future__ import annotations

import typing as tp

import numpy as np
import numpy.typing as npt

if tp.TYPE_CHECKING:
    from .esbmtk import GasReservoir, Reservoir, Species, SpeciesGroup

# declare numpy types
NDArrayFloat = npt.NDArray[np.float64]
NDArrayInt = npt.NDArray[np.int64]


[docs] def carbonate_system_1_pp(box_names: SpeciesGroup) -> None: """ Calculate carbonate system diagnostic variables from DIC and H⁺. Computes bicarbonate concentration (HCO3⁻), carbonate concentration (CO_3^2-), pH, and calcite saturation state (Ω) from previously calculated DIC and H⁺ concentrations. Parameters ---------- box_names : SpeciesGroup or list[SpeciesGroup] Reservoir group(s) containing initialized carbonate system variables. Each reservoir must contain at minimum: - ``DIC`` - ``Hplus`` - ``swc`` (seawater chemistry constants) Returns ------- None Results are stored as ``VectorData`` objects attached to each reservoir group. Notes ----- The following variables are created: - ``HCO3`` : bicarbonate concentration - ``CO3`` : carbonate concentration - ``pH`` : pH on the total scale - ``Omega`` : calcite saturation state - Concentrations must be expressed as mol kg⁻¹. - The model state variables must be stored as mol kg⁻¹. References ---------- Boudreau et al. (2010), https://doi.org/10.1029/2009GB003654 Follows et al. (2006), https://doi.org/10.1016/j.ocemod.2005.05.004 """ from esbmtk import VectorData if not isinstance(box_names, list): box_names = [box_names] for rg in box_names: k1 = rg.swc.K1 # K1 k2 = rg.swc.K2 # K2 k1k2 = rg.swc.K1K2 hplus = rg.Hplus.c dic = rg.DIC.c VectorData( name="HCO3", register=rg, species=rg.mo.HCO3, data=dic / (1 + hplus / k1 + k2 / hplus), label="HCO3-", plt_units=rg.mo.c_unit, ) VectorData( name="CO3", register=rg, species=rg.mo.CO3, data=dic / (1 + hplus / k2 + hplus**2 / k1k2), label="CO32-", plt_units=rg.mo.c_unit, ) VectorData( name="pH", register=rg, species=rg.mo.pH, data=-np.log10(hplus), label="pH", plt_units="total scale", ) VectorData( name="Omega", register=rg, species=rg.mo.pH, data=rg.swc.ca2 * rg.CO3.c / rg.swc.Ksp_ca, label=r"$\Omega$-Calcite", plt_units="", )
[docs] def carbonate_system_2_pp( bn: Reservoir | list, # 2 Reservoir handle export_fluxes: float | list, # 3 CaCO3 export flux as DIC zsat_min: float = 200, zmax: float = 10000, ) -> None: """ Carbonate chemistry post-processing diagnostics for CS2. Computes carbonate concentration, CO2aq concentration, pH, lysocline depth, carbonate compensation depth, dissolution flux, and burial flux from model DIC and H⁺ concentrations. Parameters ---------- bn : Reservoir or list[Reservoir] Reservoir(s) for post-processing. export_fluxes : float or array-like or list CaCO₃ export flux supplied to each reservoir. Scalar values are expanded to a full time series. zsat_min : float, default=200 Minimum saturation depth (m). zmax : float, default=10000 Maximum depth represented by the carbonate lookup tables (m). Returns ------- None Results are written as ``VectorData`` objects attached to each reservoir. Creates ------- CO3 Carbonate ion concentration. CO2aq Dissolved aqueous CO₂ concentration. pH pH on the total scale. zsat Saturation horizon depth (m). zcc Carbonate compensation depth (m). Fdiss Carbonate dissolution flux (mol yr⁻¹). Fburial Carbonate burial flux (mol yr⁻¹). CaCO3_export Export flux used in the calculation (mol yr⁻¹). Assumptions ----------- - Concentrations are expressed in mol kg⁻¹. - Model state variables are stored in mol kg⁻¹. - Lookup tables defining seafloor area and saturation carbonate concentrations have already been initialized. References ---------- Boudreau et al. (2010), https://doi.org/10.1029/2009GB003654 Follows et al. (2006), https://doi.org/10.1016/j.ocemod.2005.05.004 """ from math import log from esbmtk import VectorData, Q_ # ensure that all objects are lists if not isinstance(bn, list): bn = [bn] if not isinstance(export_fluxes, list): export_fluxes = [export_fluxes] # loop over boxes for i, rg in enumerate(bn): p = rg.cs2.function_params sp, cp, area_table, area_dz_table, Csat_table = p ksp0, kc, AD, zsat0, I_caco3, alpha, zsat_min, zmax, z0 = cp k1, k2, k1k2, KW, KB, ca2, boron, isotopes = sp hplus: NDArrayFloat = rg.Hplus.c dic: NDArrayFloat = rg.DIC.c zsnow: NDArrayInt = rg.zsnow.c.astype(int) export_data = export_fluxes[i] # make sure we have a vector if isinstance(export_data, float | int): export: NDArrayFloat = dic * 0 + export_data else: export = export_data # hco3 = dic / (1 + hplus / k1 + k2 / hplus) co3: NDArrayFloat = dic / (1 + hplus / k2 + hplus**2 / k1k2) co2aq: NDArrayFloat = dic / (1 + k1 / hplus + k1k2 / hplus**2) zsat: NDArrayInt = np.clip( zsat0 * np.log(ca2 * co3 / ksp0), zsat_min, zmax, ).astype(int) B_AD: NDArrayFloat = export / AD Fdiss: NDArrayFloat = co3 * 0 Fburial: NDArrayFloat = co3 * 0 zcc: NDArrayInt = co3.astype(int) * 0 for i, z in enumerate(zsat): zcc[i] = int( zsat0 * log(export[i] * ca2 / (ksp0 * AD * kc) + ca2 * co3[i] / ksp0) ) # eq3 if zcc[i] > zmax: zcc[i] = zmax print( f"Warning zcc > zmax, i = {i}, co3 = {co3[i] * 1e6} umol/kg, export = {export[i] / 1e12:.2f} Tmol/y" ) elif zcc[i] < z0: zcc[i] = z0 A_z0_zsat: float = area_table[z0] - area_table[z] A_zsat_zcc: float = area_table[z] - area_table[zcc[i]] A_zcc_zmax: float = area_table[zcc[i]] - area_table[zmax] BCC: float = A_zcc_zmax * B_AD[i] BNS: float = alpha * A_z0_zsat * B_AD[i] diff_co3: float = Csat_table[z : zcc[i]] - co3[i] area_p: NDArrayFloat = area_dz_table[z : zcc[i]] try: BDS_under: float = kc * area_p.dot(diff_co3) except: breakpoint() BDS_resp: float = alpha * (A_zsat_zcc * B_AD[i] - BDS_under) BDS: float = BDS_under + BDS_resp """ Note that we do not recalculate zsnow in post processing since it is already known from the original run """ if zsnow[i] <= zcc[i]: # reset zsnow # dzdt_zsnow: int = abs(zsnow - zcc) # zsnow[i] = zcc[i] BPDC: float = 0.0 else: # integrate saturation difference over area # if zsnow[i] > zmax: # zsnow[i] = zmax # integrate saturation difference over area diff: float = Csat_table[zcc[i] : zsnow[i]] - co3[i] area_p_snow: NDArrayFloat = area_dz_table[zcc[i] : zsnow[i]] BPDC = max(0, kc * area_p_snow.dot(diff)) # dzdt_zsnow = -BPDC / (area_dz_table[int(zsnow)] * I_caco3) Fdiss[i] = BDS + BCC + BNS + BPDC Fburial[i] = export[i] - Fdiss[i] VectorData( name="Fburial", register=rg, species=rg.mo.Fburial, data=Fburial, label="Fburial", plt_units=rg.mo.f_unit, ) VectorData( name="Fdiss", register=rg, species=rg.mo.Fdiss, data=Fdiss, label="Fdiss", plt_units=rg.mo.f_unit, ) VectorData( name="CO3", register=rg, species=rg.mo.CO3, data=co3, label="CO32-", plt_units=rg.mo.c_unit, ) VectorData( name="CO2aq", register=rg, species=rg.mo.CO2aq, data=co2aq, label="CO2aq", plt_units=rg.mo.c_unit, ) VectorData( name="pH", register=rg, species=rg.mo.pH, data=-np.log10(hplus), label="pH", plt_units="total scale", ) VectorData( name="zsat", register=rg, species=rg.mo.zsat, data=zsat, label="zsat", plt_units="m", ) VectorData( name="zcc", register=rg, species=rg.mo.zcc, data=zcc, label="zcc", plt_units="m", ) if isinstance(export, Q_): export = export.magnitude VectorData( name="CaCO3_export", register=rg, species=rg.mo.DIC, data=export, label="CaCO3_export", plt_units="mol/year", )
[docs] def carbonate_system_3_pp( bn: Reservoir | list, # 2 Reservoir handle export_fluxes: float | list, # 3 CaCO3 export flux as DIC zsat_min: float = 200, zmax: float = 10000, ) -> None: """ Carbonate chemistry post-processing diagnostics for CS3. Computes carbonate concentration, CO2aq concentration, pH, lysocline depth, carbonate compensation depth, dissolution flux, and burial flux from model DIC and H⁺ concentrations. Parameters ---------- bn : Reservoir or list[Reservoir] Reservoir(s) for post-processing. export_fluxes : float or array-like or list CaCO₃ export flux supplied to each reservoir. Scalar values are expanded to a full time series. zsat_min : float, default=200 Minimum saturation depth (m). zmax : float, default=10000 Maximum depth represented by the carbonate lookup tables (m). Returns ------- None Results are written as ``VectorData`` objects attached to each reservoir. Creates ------- CO3 Carbonate ion concentration. CO2aq Dissolved aqueous CO₂ concentration. pH pH on the total scale. zsat Saturation horizon depth (m). zcc Carbonate compensation depth (m). Fdiss Carbonate dissolution flux (mol yr⁻¹). Fburial Carbonate burial flux (mol yr⁻¹). CaCO3_export Export flux used in the calculation (mol yr⁻¹). Assumptions ----------- - Concentrations are expressed in mol kg⁻¹. - Model state variables are stored in mol kg⁻¹. - Lookup tables defining seafloor area and saturation carbonate concentrations have already been initialized. References ---------- Boudreau et al. (2010), https://doi.org/10.1029/2009GB003654 Follows et al. (2006), https://doi.org/10.1016/j.ocemod.2005.05.004 """ from math import log from esbmtk import VectorData # ensure that all objects are lists if not isinstance(bn, list): bn = [bn] if not isinstance(export_fluxes, list): export_fluxes = [export_fluxes] # loop over boxes for i, rg in enumerate(bn): p = rg.cs3.function_params sp, cp, area_table, area_dz_table, Csat_table = p ksp0, kc, AD, zsat0, I_caco3, alpha, zsat_min, zmax, z0 = cp k1, k2, k1k2, KW, KB, ca2, boron, isotopes = sp hplus: NDArrayFloat = rg.Hplus.c dic: NDArrayFloat = rg.DIC.c zsnow: NDArrayInt = rg.zsnow.c.astype(int) export_data = export_fluxes[i] # make sure we have a vector if isinstance(export_data, float | int): export: NDArrayFloat = dic * 0 + export_data else: export = export_data # hco3 = dic / (1 + hplus / k1 + k2 / hplus) co3: NDArrayFloat = dic / (1 + hplus / k2 + hplus**2 / k1k2) co2aq: NDArrayFloat = dic / (1 + k1 / hplus + k1k2 / hplus**2) zsat: NDArrayInt = np.clip( zsat0 * np.log(ca2 * co3 / ksp0), zsat_min, zmax, ).astype(int) B_AD: NDArrayFloat = export / AD Fdiss: NDArrayFloat = co3 * 0 Fburial: NDArrayFloat = co3 * 0 zcc: NDArrayInt = co3.astype(int) * 0 for i, z in enumerate(zsat): zcc[i] = int( zsat0 * log(export[i] * ca2 / (ksp0 * AD * kc) + ca2 * co3[i] / ksp0) ) # eq3 if zcc[i] > zmax: zcc[i] = zmax print( f"Warning zcc > zmax, i = {i}, co3 = {co3[i] * 1e6} umol/kg, export = {export[i] / 1e12:.2f} Tmol/y" ) elif zcc[i] < z0: zcc[i] = z0 A_z0_zsat: float = area_table[z0] - area_table[z] A_zsat_zcc: float = area_table[z] - area_table[zcc[i]] A_zcc_zmax: float = area_table[zcc[i]] - area_table[zmax] BCC: float = A_zcc_zmax * B_AD[i] BNS: float = alpha * A_z0_zsat * B_AD[i] diff_co3: float = Csat_table[z : zcc[i]] - co3[i] area_p: NDArrayFloat = area_dz_table[z : zcc[i]] try: BDS_under: float = kc * area_p.dot(diff_co3) except: breakpoint() BDS_resp: float = alpha * (A_zsat_zcc * B_AD[i] - BDS_under) BDS: float = BDS_under + BDS_resp """ Note that we do not recalculate zsnow in post processing since it is already known from the original run """ if zsnow[i] <= zcc[i]: # reset zsnow # dzdt_zsnow: int = abs(zsnow - zcc) # zsnow[i] = zcc[i] BPDC: float = 0.0 else: # integrate saturation difference over area # if zsnow[i] > zmax: # zsnow[i] = zmax # integrate saturation difference over area diff: float = Csat_table[zcc[i] : zsnow[i]] - co3[i] area_p_snow: NDArrayFloat = area_dz_table[zcc[i] : zsnow[i]] BPDC = max(0, kc * area_p_snow.dot(diff)) # dzdt_zsnow = -BPDC / (area_dz_table[int(zsnow)] * I_caco3) Fdiss[i] = BDS + BCC + BNS + BPDC Fburial[i] = export[i] - Fdiss[i] VectorData( name="Fburial", register=rg, species=rg.mo.Fburial, data=Fburial, label="Fburial", plt_units=rg.mo.f_unit, ) VectorData( name="Fdiss", register=rg, species=rg.mo.Fdiss, data=Fdiss, label="Fdiss", plt_units=rg.mo.f_unit, ) VectorData( name="CO3", register=rg, species=rg.mo.CO3, data=co3, label="CO32-", plt_units=rg.mo.c_unit, ) VectorData( name="CO2aq", register=rg, species=rg.mo.CO2aq, data=co2aq, label="CO2aq", plt_units=rg.mo.c_unit, ) VectorData( name="pH", register=rg, species=rg.mo.pH, data=-np.log10(hplus), label="pH", plt_units="total scale", ) VectorData( name="zsat", register=rg, species=rg.mo.zsat, data=zsat, label="zsat", plt_units="m", ) VectorData( name="zcc", register=rg, species=rg.mo.zcc, data=zcc, label="zcc", plt_units="m", ) VectorData( name="CaCO3_export", register=rg, species=rg.mo.DIC, data=export, label="CaCO3_export", plt_units="mol/year", )
[docs] def carbonate_system_4_pp( bn: Reservoir | list, export_fluxes: float | list, zsat_min: float = 200, zmax: float = 10000, ) -> None: """ Carbonate chemistry post-processing diagnostics for CS4. Computes carbonate concentration, CO2aq concentration, pH, lysocline depth, carbonate compensation depth, dissolution flux, and burial flux from model DIC and H⁺ concentrations. Parameters ---------- bn : Reservoir or list[Reservoir] Reservoir(s) for post-processing. export_fluxes : float or array-like or list CaCO₃ export flux supplied to each reservoir. Scalar values are expanded to a full time series. zsat_min : float, default=200 Minimum saturation depth (m). zmax : float, default=10000 Maximum depth represented by the carbonate lookup tables (m). Returns ------- None Results are written as ``VectorData`` objects attached to each reservoir. Creates ------- CO3 Carbonate ion concentration. CO2aq Dissolved aqueous CO₂ concentration. pH pH on the total scale. zsat Saturation horizon depth (m). zcc Carbonate compensation depth (m). Fdiss Carbonate dissolution flux (mol yr⁻¹). Fburial Carbonate burial flux (mol yr⁻¹). CaCO3_export Export flux used in the calculation (mol yr⁻¹). Notes ----- This is intended only for deep boxes. Surface and intermediate boxes should be post-processed with ``carbonate_system_1_pp``. Assumptions ----------- - Concentrations are expressed in mol kg⁻¹. - Model state variables are stored in mol kg⁻¹. - Lookup tables defining seafloor area and saturation carbonate concentrations have already been initialized. References ---------- Boudreau et al. (2010), https://doi.org/10.1029/2009GB003654 Follows et al. (2006), https://doi.org/10.1016/j.ocemod.2005.05.004 """ from math import log from esbmtk import VectorData # ensure list inputs if not isinstance(bn, list): bn = [bn] if not isinstance(export_fluxes, list): export_fluxes = [export_fluxes] for i, rg in enumerate(bn): # --- Pull CS4 parameters --- p = rg.cs4.function_params sp, cp, area_table, area_dz_table, Csat_table = p ksp0, kc, AD, zsat0, I_caco3, alpha, zsat_min, zmax, z0, zint = cp k1, k2, k1k2, KW, KB, ca2, boron, isotopes = sp # --- State variables --- hplus: NDArrayFloat = rg.Hplus.c dic: NDArrayFloat = rg.DIC.c zsnow: NDArrayInt = rg.zsnow.c.astype(int) export_data = export_fluxes[i] # ensure export is a vector if isinstance(export_data, float | int): export: NDArrayFloat = dic * 0 + export_data else: export = export_data # --- Carbonate chemistry --- co3: NDArrayFloat = dic / (1 + hplus / k2 + hplus**2 / k1k2) co2aq: NDArrayFloat = dic / (1 + k1 / hplus + k1k2 / hplus**2) zsat: NDArrayInt = np.clip( zsat0 * np.log(ca2 * co3 / ksp0), zsat_min, zmax, ).astype(int) B_AD: NDArrayFloat = export / AD Fdiss: NDArrayFloat = co3 * 0 Fburial: NDArrayFloat = co3 * 0 zcc: NDArrayInt = co3.astype(int) * 0 for i, z in enumerate(zsat): zcc[i] = int( zsat0 * log(export[i] * ca2 / (ksp0 * AD * kc) + ca2 * co3[i] / ksp0) ) if zcc[i] > zmax: zcc[i] = zmax print( f"Warning zcc > zmax, i = {i}, " f"co3 = {co3[i] * 1e6} umol/kg, " f"export = {export[i] / 1e12:.2f} Tmol/y" ) elif zcc[i] < z0: zcc[i] = z0 A_z0_zsat = area_table[z0] - area_table[z] A_zsat_zcc = area_table[z] - area_table[zcc[i]] A_zcc_zmax = area_table[zcc[i]] - area_table[zmax] BCC = A_zcc_zmax * B_AD[i] BNS = alpha * A_z0_zsat * B_AD[i] diff_co3 = Csat_table[z:zcc[i]] - co3[i] area_p = area_dz_table[z:zcc[i]] BDS_under = kc * area_p.dot(diff_co3) BDS_resp = alpha * (A_zsat_zcc * B_AD[i] - BDS_under) BDS = BDS_under + BDS_resp # snowline contribution if zsnow[i] <= zcc[i]: BPDC = 0.0 else: diff = Csat_table[zcc[i]:zsnow[i]] - co3[i] area_p_snow = area_dz_table[zcc[i]:zsnow[i]] BPDC = max(0, kc * area_p_snow.dot(diff)) Fdiss[i] = BDS + BCC + BNS + BPDC Fburial[i] = export[i] - Fdiss[i] # --- Save results --- VectorData( name="Fburial", register=rg, species=rg.mo.Fburial, data=Fburial, label="Fburial", plt_units=rg.mo.f_unit, ) VectorData( name="Fdiss", register=rg, species=rg.mo.Fdiss, data=Fdiss, label="Fdiss", plt_units=rg.mo.f_unit, ) VectorData( name="CO3", register=rg, species=rg.mo.CO3, data=co3, label="CO32-", plt_units=rg.mo.c_unit, ) VectorData( name="CO2aq", register=rg, species=rg.mo.CO2aq, data=co2aq, label="CO2aq", plt_units=rg.mo.c_unit, ) VectorData( name="pH", register=rg, species=rg.mo.pH, data=-np.log10(hplus), label="pH", plt_units="total scale", ) VectorData( name="zsat", register=rg, species=rg.mo.zsat, data=zsat, label="zsat", plt_units="m", ) VectorData( name="zcc", register=rg, species=rg.mo.zcc, data=zcc, label="zcc", plt_units="m", ) VectorData( name="CaCO3_export", register=rg, species=rg.mo.DIC, data=export, label="CaCO3_export", plt_units="mol/year", )
[docs] def gas_exchange_fluxes( liquid_reservoir: Species, gas_reservoir: GasReservoir, pv: str, ): """ Calculate air-sea gas exchange fluxes. Computes bidirectional gas exchange between a dissolved species in an ocean reservoir and its atmospheric reservoir using a piston velocity formula. Parameters ---------- liquid_reservoir : Species Reservoir species participating in gas exchange. Species that are currently supported are O2 and CO2. gas_reservoir : GasReservoir Atmospheric reservoir coupled to the liquid reservoir. pv : str or Quantity Piston velocity. E.g.: ``"4.8 m/d"`` Returns ------- ndarray Gas exchange flux time series returned by ``esbmtk.gas_exchange``. Raises ------ ValueError If ``pv`` cannot be interpreted as a quantity. ValueError If the liquid species is not ``DIC`` or ``O2``. """ from esbmtk import Q_, gas_exchange if isinstance(pv, str): pv = Q_(pv).to("meter/yr").magnitude elif isinstance(pv, Q_): pv = pv.to("meter/yr").magnitude else: raise ValueError("pv must be quantity or string") scale = liquid_reservoir.register.area * pv gas_c = gas_reservoir if liquid_reservoir.species.name == "DIC": solubility = liquid_reservoir.register.swc.SA_co2 g_c_aq = liquid_reservoir.register.CO2aq a_db = liquid_reservoir.register.swc.co2_a_db a_dg = liquid_reservoir.register.swc.co2_a_dg a_u = liquid_reservoir.register.swc.co2_a_u elif liquid_reservoir.species.name == "O2": solubility = liquid_reservoir.register.swc.SA_o2 g_c_aq = liquid_reservoir.register.O2 a_db = liquid_reservoir.register.swc.o2_a_db a_dg = liquid_reservoir.register.swc.o2_a_dg a_u = liquid_reservoir.register.swc.o2_a_u else: raise ValueError("flux calculation is only supported for DIC and O2") p = ( scale, solubility, a_db, a_dg, a_u, gas_reservoir.isotopes, liquid_reservoir.isotopes, ) return gas_exchange(gas_c.c, liquid_reservoir.c, g_c_aq.c, p)