xbudget 0.5.0__py3-none-any.whl

xbudget/__init__.py

@@ -0,0 +1,4 @@
+""" xbudget: xarray and xgcm-based functions for evaluating finite-volume budgets"""
+from .presets import *
+from .collect import *
+from .version import __version__

xbudget/collect.py

@@ -0,0 +1,381 @@
+from operator import mul
+from functools import reduce
+import copy
+import numpy as np
+import numbers
+import xarray as xr
+import xgcm
+
+import warnings
+
+def aggregate(xbudget_dict, decompose=[]):
+    """Aggregate xbudget dictionary into simpler root-level budgets.
+
+    Parameters
+    ----------
+    xbudget_dict : dictionary in xbudget-compatible format
+    decompose : str or list (default: [])
+        Name of variable type(s) to decompose into the summed parts
+
+    Examples
+    --------
+    >>> xbudget_dict = {
+        "heat": {
+            "rhs": {
+                "sum": {
+                    "advection": {
+                        "var":"advective_tendency"
+                    },
+                    "var": "heat_rhs_sum"
+                },
+                "var": "heat_rhs",
+            }
+        }
+    }
+    >>> xbudget.aggregate(xbudget_dict)
+    {'heat': {'rhs': {'advection': 'advective_tendency'}}}
+
+    >>>xbudget_dict = {
+        "heat": {
+            "rhs": {
+                "sum": {
+                    "advection": {
+                        "var":"advective_tendency",
+                        "sum": {
+                            "horizontal": {
+                                "var":"advective_tendency_h",
+                            },
+                            "vertical": {
+                                "var":"advective_tendency_v"
+                            },
+                            "var":"heat_rhs_sum_advection_sum"
+                        }
+                    },
+                    "var": "heat_rhs_sum"
+                },
+                "var": "heat_rhs",
+            }
+        }
+    }
+    >>> xbudget.aggregate(xbudget_dict)
+    {'heat': {'rhs': {'advection': 'advective_tendency'}}}
+
+    >>> xbudget.aggregate(xbudget_dict, decompose="advection")
+    {'heat': {'rhs': {'advection_horizontal': 'advective_tendency_h',
+    'advection_vertical': 'advective_tendency_v'}}}
+
+    See also
+    --------
+    disaggregate, deep_search, _deep_search
+    """
+    new_budgets = copy.deepcopy(xbudget_dict)
+    for tr, tr_xbudget_dict in xbudget_dict.items():
+        for side,terms in tr_xbudget_dict.items():
+            if side in ["lhs", "rhs"]:
+                new_budgets[tr][side] = deep_search(
+                    disaggregate(tr_xbudget_dict[side], decompose=decompose)
+                )
+    return new_budgets
+
+def disaggregate(b, decompose=[]):
+    """Disaggregate variable's provenance dictionary into summed parts
+
+    Parameters
+    ----------
+    b : xbudget sub-dictionary for a variable
+    decompose : str or list (default: [])
+        Name of variable type(s) to decompose into the summed parts
+
+    Examples
+    --------
+    >>> b = {
+        "sum": {
+            "advection": {
+                "var":"advective_tendency",
+                "sum": {
+                    "horizontal": {
+                        "var":"advective_tendency_h",
+                    },
+                    "vertical": {
+                        "var":"advective_tendency_v"
+                    },
+                    "var":"heat_rhs_sum_advection_sum"
+                }
+            },
+            "var": "heat_rhs_sum"
+        },
+        "var": "heat_rhs",
+    }
+    >>> {'advection': 'advective_tendency'}
+    {'advection': 'advective_tendency'}
+
+    >>> xbudget.disaggregate(b, decompose="advection")
+    {'advection': {'horizontal': 'advective_tendency_h',
+    'vertical': 'advective_tendency_v'}}
+    
+    See also
+    --------
+    aggregate
+    """
+    if "sum" in b:
+        bsum_novar = {k:v for (k,v) in b["sum"].items() if (k!="var") and (v is not None)}
+        sum_dict = dict((k,v["var"]) if ("var" in v) else (k,v) for k,v in bsum_novar.items())
+        b_recurse = {}
+        for (k,v) in sum_dict.items():
+            if k not in decompose:
+                b_recurse[k] = v
+            else:
+                v_dict = disaggregate(b["sum"][k], decompose=decompose)
+                if "product" in v_dict.keys():
+                    b_recurse[k] = v_dict["var"]
+                else:
+                    b_recurse[k] = v_dict
+        return b_recurse
+    return b
+
+def deep_search(b):
+    """Utility function for searching for variables in xbudget dictionary.
+    
+    See also
+    --------
+    aggregate, _deep_search
+    """
+    return _deep_search(b, new_b={}, k_last=None)
+
+def _deep_search(b, new_b={}, k_last=None):
+    """Recursive function for searching for variables in xbudget dictionary.
+    
+    See also
+    --------
+    aggregate, deep_search
+    """
+    if type(b) is str:
+        new_b[k_last] = b
+    elif type(b) is dict:
+        for (k, v) in b.items():
+            if k_last is not None:
+                k = f"{k_last}_{k}"
+            _deep_search(v, new_b=new_b, k_last=k)
+        return new_b
+
+def collect_budgets(ds, xbudget_dict):
+    """Fills xbudget dictionary with all tracer content tendencies
+
+    Parameters
+    ----------
+    ds : xr.Dataset containing budget diagnostics
+    xbudget_dict : dictionary in xbudget-compatible format
+        Example format:
+        >>> xbudget_dict = {
+            "heat": {
+                "rhs": {
+                    "sum": {
+                        "advection": {
+                            "var":"advective_tendency"
+                        },
+                        "var": "heat_rhs_sum"
+                    },
+                    "var": "heat_rhs",
+                }
+            }
+        }
+    """
+    for eq, v in xbudget_dict.items():
+        for side in ["lhs", "rhs"]:
+            if side in v:
+                budget_fill_dict(ds, v[side], f"{eq}_{side}")
+
+def budget_fill_dict(data, xbudget_dict, namepath):
+    """Recursively fill xbudget dictionary
+
+    Parameters
+    ----------
+    data : xgcm.grid or xr.Dataset
+    xbudget_dict : dictionary in xbudget-compatible format containing variable in namepath
+    namepath : name of variable in dataset (data._ds or data)
+    """
+    if type(data)==xgcm.grid.Grid:
+        grid = data
+        ds = grid._ds
+    else:
+        ds = data
+        grid = None
+    
+    var_pref = None
+
+    if ((xbudget_dict["var"] is not None) and
+        (xbudget_dict["var"] in ds)       and
+        (namepath not in ds)):
+        var_rename = ds[xbudget_dict["var"]].rename(namepath)
+        var_rename.attrs['provenance'] = xbudget_dict["var"]
+        ds[namepath] = ds[xbudget_dict["var"]]
+        var_pref = ds[namepath]
+
+    for k,v in xbudget_dict.items():
+        if k in ['sum', 'product']:
+            op_list = []
+            for k_term, v_term in v.items():
+                if isinstance(v_term, dict): # recursive call to get this variable
+                    v_term_recursive = budget_fill_dict(data, v_term, f"{namepath}_{k}_{k_term}")
+                    if v_term_recursive is not None:
+                        op_list.append(v_term_recursive)
+                elif isinstance(v_term, numbers.Number):
+                    op_list.append(v_term)
+                elif isinstance(v_term, str):
+                    if v_term in ds:
+                        op_list.append(ds[v_term])
+                    else:
+                        warnings.warn(f"Variable {v_term} is missing from the dataset `ds`, so it is being skipped. To suppress this warning, remove {v_term} from the `xbudget_dict`.")
+                        if k=="product":
+                            op_list.append(0.)
+
+            # Compute variable from sum or product operation
+            if (
+                (len(op_list) == 0) |
+                all([e is None for e in op_list]) |
+                any([e is None for e in op_list])
+            ):
+                return None
+            else:
+                var = sum(op_list) if k=="sum" else reduce(mul, op_list, 1)
+                if not isinstance(var, xr.DataArray):
+                    continue
+
+            # Variable metadata
+            var_name = f"{namepath}_{k}"
+            var = var.rename(var_name)
+            var_provenance = [o.name if isinstance(o, xr.DataArray) else o for o in op_list]
+            var.attrs["provenance"] = var_provenance
+            ds[var_name] = var
+            if (xbudget_dict[k]["var"] is None):
+                xbudget_dict[k]["var"] = var_name
+
+            if (xbudget_dict["var"] is None):
+                var_copy = var.copy()
+                var_copy.attrs["provenance"] = var_name
+                xbudget_dict["var"] = namepath
+                if namepath not in ds:
+                    ds[namepath] = var_copy
+
+            # keep record of the first-listed variable
+            if var_pref is None:
+                var_pref = var.copy()
+                
+        if k == "difference":
+            if grid is not None:
+                staggered_axes = {
+                    axn:c for axn,ax in grid.axes.items()
+                    for pos,c in ax.coords.items()
+                    if pos!="center"
+                }
+                v_term = [v_term for k_term,v_term in v.items() if k_term!="var"][0]
+                if v_term not in ds:
+                    warnings.warn(f"Variable {v_term} is missing from the dataset `ds`, so it is being skipped. To suppress this warning, remove {v_term} from the `xbudget_dict`.")
+                    continue
+                candidate_axes = [axn for (axn,c) in staggered_axes.items() if c in ds[v_term].dims]
+                if len(candidate_axes) == 1:
+                    axis = candidate_axes[0]
+                else:
+                    raise ValueError("Flux difference inconsistent with finite volume discretization.")
+                var = grid.diff(ds[v_term].fillna(0.), axis)
+                var_name = f"{namepath}_difference"
+                var = var.rename(var_name)
+                var_provenance = v_term
+                var.attrs["provenance"] = var_provenance
+                ds[var_name] = var
+                if var_pref is None:
+                    var_pref = var.copy()
+            else:
+                raise ValueError("Input `ds` must be `xgcm.Grid` instance if using `difference` operations.")
+
+    return var_pref
+
+def get_vars(xbudget_dict, terms):
+    """Get xbudget sub-dictionaries for specified terms.
+
+    Parameters
+    ----------
+    xbudget_dict : dictionary in xbudget-compatible format
+    terms : str or list of str
+
+    Examples
+    -------
+    >>> xbudget_dict = {
+        "heat": {
+            "rhs": {
+                "sum": {
+                    "advection": {
+                        "var":"advective_tendency"
+                    },
+                    "var": "heat_rhs_sum"
+                },
+                "var": "heat_rhs",
+            }
+        }
+    }
+    >>> xbudget.get_vars(xbudget_dict, "heat_rhs_sum")
+    {'var': 'heat_rhs_sum', 'sum': ['advective_tendency']}
+    """
+    return _get_vars(xbudget_dict, terms)
+
+def _get_vars(b, terms, k_long=""):
+    """Recursive version of _get_vars for determining variable provenance tree.
+    
+    Parameters
+    ----------
+    b : dictionary
+    terms : str or list of str
+    k_long : variable name suffix
+    
+    See also
+    --------
+    get_vars
+    """
+    if isinstance(terms, (list, np.ndarray)):
+        return [_get_vars(b, term) for term in terms]
+    elif type(terms) is str:
+        for k,v in b.items():
+            if type(v) is str:
+                k_short = k_long.replace("_sum", "").replace("_product", "")
+                if v==terms:
+                    decomps = {"var": v}
+                    if len(terms) > len("_sum"):
+                        if (terms[-len("_sum"):] == "_sum") and ("sum" in b):
+                            ts = {kk:vv for (kk,vv) in b["sum"].items() if kk!="var"}
+                            decomps["sum"] = [vv["var"] if type(vv) is dict else vv for (kk,vv) in ts.items()]
+                        elif (terms[-len("_sum"):] == "_sum"):
+                            ts = {kk:vv for (kk,vv) in b.items() if kk!="var"}
+                            decomps["sum"] = [vv["var"] if type(vv) is dict else vv for (kk,vv) in ts.items()]
+                    if len(terms) > len("_product"):
+                        if (terms[-len("_product"):] == "_product") and ("product" in b):
+                            ts = {kk:vv for (kk,vv) in b["product"].items() if kk!="var"}
+                            decomps["product"] = [vv["var"] if type(vv) is dict else vv for (kk,vv) in ts.items()]
+                        elif (terms[-len("_product"):] == "_product"):
+                            ts = {kk:vv for (kk,vv) in b.items() if kk!="var"}
+                            decomps["product"] = [vv["var"] if type(vv) is dict else vv for (kk,vv) in ts.items()]
+                    return decomps
+
+                if k!="var":
+                    k_short+="_"+k
+                if k_short==terms:
+                    return v
+            elif type(v) is dict:
+                if k_long=="":
+                    new_k = k
+                elif len(k_long)>0:
+                    new_k = f"{k_long}_{k}"
+                var = _get_vars(v, terms, k_long=new_k)
+                if var is not None:
+                    return var
+
+def flatten(container):
+    for i in container:
+        if isinstance(i, (list,tuple)):
+            for j in flatten(i):
+                yield j
+        else:
+            yield i
+            
+def flatten_lol(lol):
+    """Flatten a list of lists into a single list."""
+    return list(flatten(lol))

xbudget/conventions/MOM6.yaml

@@ -0,0 +1,349 @@
+---
+mass: # finite-volume mass budget in units of kg/s
+  lambda: "density"
+  thickness: "thkcello"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          thickness_tendency: "dhdt"
+          density: 1035.
+          area: "areacello"
+  rhs:
+    var: null
+    sum:
+      var: null
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              thickness_tendency: "dynamics_h_tendency"
+              density: 1035.
+              area: "areacello"
+            sum:
+              var: null
+              zonal_convergence:
+                var: null
+                product:
+                  var: null
+                  zonal_divergence:
+                    var: null
+                    difference:
+                      var: null
+                      zonal_mass_transport: "umo"
+                  sign: -1.
+              meridional_convergence:
+                var: null
+                product:
+                  var: null
+                  meridional_divergence:
+                    var: null
+                    difference:
+                      var: null
+                      meridional_mass_transport: "vmo"
+                  sign: -1.
+          interfacial:
+            var: null
+            product:
+              var: null
+              thickness_tendency: "vert_remap_h_tendency"
+              density: 1035.
+              area: "areacello"
+      surface_exchange_flux:
+        var: null
+        product:
+          var: null
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+        sum:
+          var: null
+          rain_and_ice:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "prlq"
+              area: "areacello"
+          snow:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "prsn"
+              area: "areacello"
+          evaporation:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "evs"
+              area: "areacello"
+          rivers:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "friver"
+              area: "areacello"
+          icebergs:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "ficeberg"
+              area: "areacello"
+          sea_ice_melt:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "fsitherm"
+              area: "areacello"
+          virtual_precip_restoring:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "vprec"
+              area: "areacello"
+
+heat: # finite-volume heat budget in units of J/s
+  lambda: "thetao"
+  surface_lambda: "tos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "opottemptend"
+          area: "areacello"
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "T_advection_xy"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "Th_tendency_vert_remap"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          specific_heat_capacity: 3992.
+          lambda_mass: "tos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "opottemppmdiff"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "opottempdiff"
+              area: "areacello"
+            
+      surface_exchange_flux:
+        var:
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "boundary_forcing_heat_tendency"
+          area: "areacello"
+        sum:
+          var: null
+          nonadvective:
+            var: null
+            sum:
+              var: null # This sum is currently broken because the shortwave terms are 3D while the other terms are 2D!
+              latent:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "hflso"
+                  area: "areacello"
+              sensible:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "hfsso"
+                  area: "areacello"
+              longwave:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "rlntds"
+                  area: "areacello"
+              shortwave:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "rsdoabsorb"
+                  area: "areacello"
+                # Not yet supported!
+                #var: "rsdoabsorb"
+                #finite_difference:
+                #  Z: "rsdo"
+          advective:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "heat_content_surfwater"
+              area: "areacello"
+            #Alternative method:
+            #product:
+            #  var: null
+            #  specific_heat_capacity: 3992.
+            #  lambda_mass: "tos"
+            #  thickness_tendency: "boundary_forcing_h_tendency"
+            #  density: 1035.
+            #  area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          specific_heat_capacity: 3992.
+          lambda_mass: "tos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+      bottom_flux:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "internal_heat_heat_tendency"
+          area: "areacello"
+      frazil_ice:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "frazil_heat_tendency"
+          area: "areacello"
+
+salt: # finite-volume salt budget in units of kg/s
+  lambda: "so"
+  surface_lambda: "sos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "osalttend"
+          area: "areacello"
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "S_advection_xy"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "Sh_tendency_vert_remap"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          unit_conversion: 0.001
+          lambda_mass: "sos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+      
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "osaltpmdiff"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "osaltdiff"
+              area: "areacello"
+      surface_exchange_flux:
+        var:
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "boundary_forcing_salt_tendency"
+          area: "areacello"
+        sum:
+          var: null
+          nonadvective:
+            var: null
+            sum:
+              var: null
+              basal:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "sfdsi"
+                  area: "areacello"
+          advective:
+            var: null
+            product:
+              var: null
+              unit_conversion: 0.001
+              lambda_mass: 0.
+              thickness_tendency: "boundary_forcing_h_tendency"
+              density: 1035.
+              area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          unit_conversion: 0.001
+          lambda_mass: "sos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"

xbudget/conventions/MOM6_3Donly.yaml

@@ -0,0 +1,230 @@
+---
+mass: # finite-volume mass budget in units of kg/s
+  lambda: "density"
+  thickness: "thkcello"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          thickness_tendency: null
+          density: 1035.
+          area: "areacello"
+  rhs:
+    var: null
+    sum:
+      var: null
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            sum:
+              var: null
+              zonal_convergence:
+                var: null
+                product:
+                  var: null
+                  zonal_divergence:
+                    var: null
+                    difference:
+                      var: null
+                      zonal_mass_transport: "umo"
+                  sign: -1.
+              meridional_convergence:
+                var: null
+                product:
+                  var: null
+                  meridional_divergence:
+                    var: null
+                    difference:
+                      var: null
+                      meridional_mass_transport: "vmo"
+          interfacial:
+            var: null
+            product:
+              var: null
+              thickness_tendency: "vert_remap_h_tendency"
+              density: 1035.
+              area: "areacello"
+      surface_exchange_flux:
+        var: null
+        product:
+          var: null
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+
+heat: # finite-volume heat budget in units of J/s
+  lambda: "thetao"
+  surface_lambda: "tos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "opottemptend"
+          area: "areacello"
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "T_advection_xy"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "Th_tendency_vert_remap"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          specific_heat_capacity: 3992.
+          lambda_mass: "tos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "opottemppmdiff"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "opottempdiff"
+              area: "areacello"
+            
+      surface_exchange_flux:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "boundary_forcing_heat_tendency"
+          area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          specific_heat_capacity: 3992.
+          lambda_mass: "tos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+      bottom_flux:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "internal_heat_heat_tendency"
+          area: "areacello"
+      frazil_ice:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "frazil_heat_tendency"
+          area: "areacello"
+
+salt: # finite-volume salt budget in units of kg/s
+  lambda: "so"
+  surface_lambda: "sos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "osalttend"
+          area: "areacello"
+      advection:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "S_advection_xy"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              tracer_content_tendency_per_unit_area: "Sh_tendency_vert_remap"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          unit_conversion: 0.001
+          lambda_mass: "sos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"
+      
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null
+        sum:
+          var: null
+          lateral:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "osaltpmdiff"
+              area: "areacello"
+          interfacial:
+            var: null
+            product:
+              var: null
+              tracer_content_tendency_per_unit_area: "osaltdiff"
+              area: "areacello"
+      surface_exchange_flux:
+        var:
+        product:
+          var: null
+          tracer_content_tendency_per_unit_area: "boundary_forcing_salt_tendency"
+          area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          unit_conversion: 0.001
+          lambda_mass: "sos"
+          thickness_tendency: "boundary_forcing_h_tendency"
+          density: 1035.
+          area: "areacello"

xbudget/conventions/MOM6_drift.yaml

@@ -0,0 +1,70 @@
+---
+mass: # finite-volume mass budget in units of kg/s
+  lambda: "density"
+  thickness: "thkcello"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      advection:
+        var: null
+      surface_exchange_flux:
+        var: null
+
+heat: # finite-volume heat budget in units of J/s
+  lambda: "thetao"
+  surface_lambda: "tos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null 
+      surface_exchange_flux:
+        var: null
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+      bottom_flux:
+        var: null
+      frazil_ice:
+        var: null
+
+salt: # finite-volume salt budget in units of kg/s
+  lambda: "so"
+  surface_lambda: "sos"
+  lhs:
+    var: null
+    sum:
+      var: null
+      Eulerian_tendency:
+        var: null
+      advection:
+        var: null
+        sum:
+          var: null
+      surface_ocean_flux_advective_negative_lhs:
+        var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      diffusion:
+        var: null
+      surface_exchange_flux:
+        var: null
+      surface_ocean_flux_advective_negative_rhs:
+        var: null

xbudget/conventions/MOM6_surface.yaml

@@ -0,0 +1,214 @@
+---
+mass: # finite-volume mass budget in units of kg/s
+  lambda: "density"
+  thickness: "thkcello"
+  lhs:
+    var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      surface_exchange_flux:
+        var: null
+        sum:
+          var: null
+          rain_and_ice:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "prlq"
+              area: "areacello"
+          snow:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "prsn"
+              area: "areacello"
+          evaporation:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "evs"
+              area: "areacello"
+          rivers:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "friver"
+              area: "areacello"
+          icebergs:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "ficeberg"
+              area: "areacello"
+          sea_ice_melt:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "fsitherm"
+              area: "areacello"
+          virtual_precip_restoring:
+            var: null
+            product:
+              var: null
+              mass_tendency_per_unit_area: "vprec"
+              area: "areacello"
+
+heat: # finite-volume heat budget in units of J/s
+  lambda: "thetao"
+  surface_lambda: "tos"
+  lhs:
+    var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      surface_exchange_flux:
+        var:
+        sum:
+          var: null
+          nonadvective:
+            var: null
+            sum:
+              var: null
+              latent:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "hflso"
+                  area: "areacello"
+              sensible:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "hfsso"
+                  area: "areacello"
+              longwave:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "rlntds"
+                  area: "areacello"
+              shortwave:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "rsntds"
+                  area: "areacello"
+          advective:
+            var: null
+            product:
+              var: null
+              specific_heat_capacity: 3992.
+              lambda_mass: "tos"
+              mass_tendency_per_unit_area: "wfo"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        product:
+          var: null
+          sign: -1.
+          specific_heat_capacity: 3992.
+          lambda_mass: "tos"
+          mass_tendency_per_unit_area: "wfo"
+          area: "areacello"
+
+salt: # finite-volume salt budget in units of kg/s
+  lambda: "so"
+  surface_lambda: "sos"
+  lhs:
+    var: null
+  rhs:
+    var: null
+    sum:
+      var: null
+      surface_exchange_flux:
+        var:
+        sum:
+          var: null
+          nonadvective:
+            var: null
+            sum:
+              var: null
+              basal:
+                var: null
+                product:
+                  var: null
+                  tracer_content_tendency_per_unit_area: "sfdsi"
+                  area: "areacello"
+          advective:
+            var: null
+            product:
+              var: null
+              unit_conversion: 0.001
+              lambda_mass: 0.
+              mass_tendency_per_unit_area: "wfo"
+              area: "areacello"
+      surface_ocean_flux_advective_negative_rhs:
+        var: null
+        sum:
+          var: null
+          rain_and_ice:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "prlq"
+              area: "areacello"
+          snow:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "prsn"
+              area: "areacello"
+          evaporation:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "evs"
+              area: "areacello"
+          rivers:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "friver"
+              area: "areacello"
+          icebergs:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "ficeberg"
+              area: "areacello"
+          sea_ice_melt:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "fsitherm"
+              area: "areacello"
+          virtual_precip_restoring:
+            var: null
+            product:
+              var: null
+              sign: -1.
+              unit_conversion: 0.001
+              lambda_mass: "sos"
+              mass_tendency_per_unit_area: "vprec"
+              area: "areacello"

xbudget/presets.py

@@ -0,0 +1,42 @@
+import yaml
+from os import path
+
+def load_preset_budget(model="MOM6"):
+    """Loads preset xbudget dictionary from yaml file for supported models.
+
+    Note: Does not yet check whether or not the format of the yaml file (or
+    the corresponding Python dictionary) is in an xbudget-compatible format!
+
+    Parameters
+    ----------
+    model : str (default "MOM6")
+      Currently supported models: ["MOM6"]
+      Please open an Issue if you would like to contribute an xbudget yaml
+      file for a new model–see /conventions/ for examples.
+
+    Returns
+    -------
+    Python dictionary
+    """
+    return load_yaml(f"{path.dirname(__file__)}/conventions/{model}.yaml")
+    
+def load_yaml(filepath):
+    """Loads a yaml file as a Python dictionary.
+
+    Note: Does not yet check whether or not the format of the yaml file (or
+    the corresponding Python dictionary) is in an xbudget-compatible format!
+
+    Parameters
+    ----------
+    filepath : path to yaml file, as str
+
+    Returns
+    -------
+    Python dictionary
+    """
+    with open(filepath, "r") as stream:
+        try:
+            xbudget_dict = yaml.safe_load(stream)
+        except yaml.YAMLError as exc:
+            print(exc)
+    return xbudget_dict

xbudget/version.py

@@ -0,0 +1,3 @@
+"""xbudget: version information"""
+
+__version__ = "0.5.0"

xbudget-0.5.0.dist-info/METADATA

@@ -0,0 +1,38 @@
+Metadata-Version: 2.4
+Name: xbudget
+Version: 0.5.0
+Summary: Helper functions and meta-data conventions for wrangling finite-volume ocean model budgets
+Project-URL: Homepage, https://github.com/hdrake/xbudget
+Project-URL: Bugs/Issues/Features, https://github.com/hdrake/xbudget/issues
+Author-email: "Henri F. Drake" <hfdrake@uci.edu>
+License-File: LICENSE
+Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Requires-Dist: numpy
+Requires-Dist: xarray
+Requires-Dist: xgcm>=0.9.0
+Description-Content-Type: text/markdown
+
+# xbudget
+Helper functions and meta-data conventions for wrangling finite-volume ocean model budgets.
+
+Quick Start Guide
+-----------------
+
+**For users: minimal installation within an existing environment**
+```bash
+pip install git+https://github.com/hdrake/xbudget.git@main
+```
+
+**For developers: installing from scratch using `conda`**
+```bash
+git clone git@github.com:hdrake/xbudget.git
+cd xbudget
+conda env create -f docs/environment.yml
+conda activate docs_env_xbudget
+pip install -e .
+python -m ipykernel install --user --name docs_env_xbudget --display-name "docs_env_xbudget"
+jupyter-lab
+```

xbudget-0.5.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+xbudget/__init__.py,sha256=JyiQGuRTHCkpfa70e7xum_Z5VSTr1airEqR-fZJroj4,164
+xbudget/collect.py,sha256=3ihE7S87OfE_k3ylpx8HzPiGSkkaIsu_Pl3-vp7DcJw,13316
+xbudget/presets.py,sha256=VQKWF1CDm-ei3AwEIgulwo6UBLFaag0HFW-eYVRCuhk,1221
+xbudget/version.py,sha256=XsM9KNFRxEVBi_aTXVM-KXGt0OhCm_HdvZq9o0W0DGI,58
+xbudget/conventions/MOM6.yaml,sha256=mV5rRyv9wYRrEm13E7G-F1ADXjlm2Gmp4CTP9smkAQI,9807
+xbudget/conventions/MOM6_3Donly.yaml,sha256=AllZ2KahJ3mXtIoOYyRvFA-wYOcfiOVs3Phw2-3CZxE,6113
+xbudget/conventions/MOM6_drift.yaml,sha256=JJGh2anx0Ok54yNqFUz4ODCiWSJKqbVZL7QEEA7rZnI,1345
+xbudget/conventions/MOM6_surface.yaml,sha256=qyVw9vKX2wvWck8S1VjzfKl4VhmJ1L7fsFVtNURI3cM,5840
+xbudget-0.5.0.dist-info/METADATA,sha256=7Pav0qRzbTxf_oEWyX0s71eDZoYUSbBm5p9T5__1e00,1260
+xbudget-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xbudget-0.5.0.dist-info/licenses/LICENSE,sha256=e8O6mejrFQyNqfJmLJNKyjBZrH-dprkeLdqqKxnY7dE,1071
+xbudget-0.5.0.dist-info/RECORD,,

xbudget-0.5.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

xbudget-0.5.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Henri F. Drake
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.