%load_ext autoreload
%autoreload 2

import xarray as xr
import xgcm
import matplotlib.pyplot as plt

Tutorial: analyzing finite-volume ocean model budgets with xbudget

This tutorial shows how to use xbudget to close and manipulate mass and tracer budgets in finite-volume.

Load example MOM6 dataset from Zenodo

from load_example_model_grid import load_MOM6_coarsened_diagnostics
grid = load_MOM6_coarsened_diagnostics()

File 'MOM6_global_example_diagnostics_zlevels_v0_0_6.nc' already exists at ../data/MOM6_global_example_diagnostics_zlevels_v0_0_6.nc. Skipping download.

Load budget metadata from preset MOM6 dictionary

import xbudget
xbudget_dict = xbudget.load_preset_budget(model="MOM6").copy()

We load in the file xbudget/conventions/MOM6.yaml, which comprehensively details the mass, heat, and salt budgets in MOM6, and how they can be evaluated using various combinations of diagnostics. In MOM6, most mass and tracer diagnostics are output as thickness tendencies and thickness-weighted tracer tendencies, respectively, and must be converted into density-weighted integrals over a finite-volume cell by multiplying by a seawater density and grid cell area.

For example, consider the structure of the heat budget:

import json
print(json.dumps(xbudget_dict['heat'], sort_keys=True, indent=4))

{
    "lambda": "thetao",
    "lhs": {
        "sum": {
            "Eulerian_tendency": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "opottemptend",
                    "var": null
                },
                "var": null
            },
            "advection": {
                "sum": {
                    "interfacial": {
                        "product": {
                            "area": "areacello",
                            "sign": -1.0,
                            "tracer_content_tendency_per_unit_area": "Th_tendency_vert_remap",
                            "var": null
                        },
                        "var": null
                    },
                    "lateral": {
                        "product": {
                            "area": "areacello",
                            "sign": -1.0,
                            "tracer_content_tendency_per_unit_area": "T_advection_xy",
                            "var": null
                        },
                        "var": null
                    },
                    "var": null
                },
                "var": null
            },
            "surface_ocean_flux_advective_negative_lhs": {
                "product": {
                    "area": "areacello",
                    "density": 1035.0,
                    "lambda_mass": "tos",
                    "sign": -1.0,
                    "specific_heat_capacity": 3992.0,
                    "thickness_tendency": "boundary_forcing_h_tendency",
                    "var": null
                },
                "var": null
            },
            "var": null
        },
        "var": null
    },
    "rhs": {
        "sum": {
            "bottom_flux": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "internal_heat_heat_tendency",
                    "var": null
                },
                "var": null
            },
            "diffusion": {
                "sum": {
                    "interfacial": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "opottempdiff",
                            "var": null
                        },
                        "var": null
                    },
                    "lateral": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "opottemppmdiff",
                            "var": null
                        },
                        "var": null
                    },
                    "var": null
                },
                "var": null
            },
            "frazil_ice": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "frazil_heat_tendency",
                    "var": null
                },
                "var": null
            },
            "surface_exchange_flux": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "boundary_forcing_heat_tendency",
                    "var": null
                },
                "sum": {
                    "advective": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "heat_content_surfwater",
                            "var": null
                        },
                        "var": null
                    },
                    "nonadvective": {
                        "sum": {
                            "latent": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "hflso",
                                    "var": null
                                },
                                "var": null
                            },
                            "longwave": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "rlntds",
                                    "var": null
                                },
                                "var": null
                            },
                            "sensible": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "hfsso",
                                    "var": null
                                },
                                "var": null
                            },
                            "shortwave": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "rsdoabsorb",
                                    "var": null
                                },
                                "var": null
                            },
                            "var": null
                        },
                        "var": null
                    },
                    "var": null
                },
                "var": null
            },
            "surface_ocean_flux_advective_negative_rhs": {
                "product": {
                    "area": "areacello",
                    "density": 1035.0,
                    "lambda_mass": "tos",
                    "sign": -1.0,
                    "specific_heat_capacity": 3992.0,
                    "thickness_tendency": "boundary_forcing_h_tendency",
                    "var": null
                },
                "var": null
            },
            "var": null
        },
        "var": null
    },
    "surface_lambda": "tos"
}

The var: null entries in the dictionary denote where diagnostics are not directly available. In many cases, however, they can be reconstructed by taking the sum or product of available diagnostics.

We can use the xbudget.collect_budgets function to leverage the information in this metadata dictionary to fill in all of these gaps in the budget.

xbudget.collect_budgets(grid, xbudget_dict)

All of these reconstructed budget terms are also automatically (and lazily) added to the original dataset with a straight-forward naming convention, so they can be used directly to evaluate the budget.

import json
print(json.dumps(xbudget_dict['heat'], sort_keys=True, indent=4))

{
    "lambda": "thetao",
    "lhs": {
        "sum": {
            "Eulerian_tendency": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "opottemptend",
                    "var": "heat_lhs_sum_Eulerian_tendency_product"
                },
                "var": "heat_lhs_sum_Eulerian_tendency"
            },
            "advection": {
                "sum": {
                    "interfacial": {
                        "product": {
                            "area": "areacello",
                            "sign": -1.0,
                            "tracer_content_tendency_per_unit_area": "Th_tendency_vert_remap",
                            "var": "heat_lhs_sum_advection_sum_interfacial_product"
                        },
                        "var": "heat_lhs_sum_advection_sum_interfacial"
                    },
                    "lateral": {
                        "product": {
                            "area": "areacello",
                            "sign": -1.0,
                            "tracer_content_tendency_per_unit_area": "T_advection_xy",
                            "var": "heat_lhs_sum_advection_sum_lateral_product"
                        },
                        "var": "heat_lhs_sum_advection_sum_lateral"
                    },
                    "var": "heat_lhs_sum_advection_sum"
                },
                "var": "heat_lhs_sum_advection"
            },
            "surface_ocean_flux_advective_negative_lhs": {
                "product": {
                    "area": "areacello",
                    "density": 1035.0,
                    "lambda_mass": "tos",
                    "sign": -1.0,
                    "specific_heat_capacity": 3992.0,
                    "thickness_tendency": "boundary_forcing_h_tendency",
                    "var": "heat_lhs_sum_surface_ocean_flux_advective_negative_lhs_product"
                },
                "var": "heat_lhs_sum_surface_ocean_flux_advective_negative_lhs"
            },
            "var": "heat_lhs_sum"
        },
        "var": "heat_lhs"
    },
    "rhs": {
        "sum": {
            "bottom_flux": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "internal_heat_heat_tendency",
                    "var": "heat_rhs_sum_bottom_flux_product"
                },
                "var": "heat_rhs_sum_bottom_flux"
            },
            "diffusion": {
                "sum": {
                    "interfacial": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "opottempdiff",
                            "var": "heat_rhs_sum_diffusion_sum_interfacial_product"
                        },
                        "var": "heat_rhs_sum_diffusion_sum_interfacial"
                    },
                    "lateral": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "opottemppmdiff",
                            "var": "heat_rhs_sum_diffusion_sum_lateral_product"
                        },
                        "var": "heat_rhs_sum_diffusion_sum_lateral"
                    },
                    "var": "heat_rhs_sum_diffusion_sum"
                },
                "var": "heat_rhs_sum_diffusion"
            },
            "frazil_ice": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "frazil_heat_tendency",
                    "var": "heat_rhs_sum_frazil_ice_product"
                },
                "var": "heat_rhs_sum_frazil_ice"
            },
            "surface_exchange_flux": {
                "product": {
                    "area": "areacello",
                    "tracer_content_tendency_per_unit_area": "boundary_forcing_heat_tendency",
                    "var": "heat_rhs_sum_surface_exchange_flux_product"
                },
                "sum": {
                    "advective": {
                        "product": {
                            "area": "areacello",
                            "tracer_content_tendency_per_unit_area": "heat_content_surfwater",
                            "var": "heat_rhs_sum_surface_exchange_flux_sum_advective_product"
                        },
                        "var": "heat_rhs_sum_surface_exchange_flux_sum_advective"
                    },
                    "nonadvective": {
                        "sum": {
                            "latent": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "hflso",
                                    "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_latent_product"
                                },
                                "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_latent"
                            },
                            "longwave": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "rlntds",
                                    "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_longwave_product"
                                },
                                "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_longwave"
                            },
                            "sensible": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "hfsso",
                                    "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_sensible_product"
                                },
                                "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_sensible"
                            },
                            "shortwave": {
                                "product": {
                                    "area": "areacello",
                                    "tracer_content_tendency_per_unit_area": "rsdoabsorb",
                                    "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_shortwave_product"
                                },
                                "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_shortwave"
                            },
                            "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum"
                        },
                        "var": "heat_rhs_sum_surface_exchange_flux_sum_nonadvective"
                    },
                    "var": "heat_rhs_sum_surface_exchange_flux_sum"
                },
                "var": "heat_rhs_sum_surface_exchange_flux"
            },
            "surface_ocean_flux_advective_negative_rhs": {
                "product": {
                    "area": "areacello",
                    "density": 1035.0,
                    "lambda_mass": "tos",
                    "sign": -1.0,
                    "specific_heat_capacity": 3992.0,
                    "thickness_tendency": "boundary_forcing_h_tendency",
                    "var": "heat_rhs_sum_surface_ocean_flux_advective_negative_rhs_product"
                },
                "var": "heat_rhs_sum_surface_ocean_flux_advective_negative_rhs"
            },
            "var": "heat_rhs_sum"
        },
        "var": "heat_rhs"
    },
    "surface_lambda": "tos"
}

Because searching through this data structure by eye can be tedious, we include helpers functions to help retrieve the variable name corresponding to each term in the budget.

from xbudget import get_vars
grid._ds[get_vars(xbudget_dict, "heat_lhs_sum_advection")['var']]

<xarray.DataArray 'heat_lhs_sum_advection' (time: 1, z_l: 35, yh: 180, xh: 240)> Size: 12MB
array([[[[0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         ...,
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.]],

        [[0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         ...,
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.]],

        [[0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         ...,
...
         ...,
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.]],

        [[0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         ...,
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.]],

        [[0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         ...,
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.],
         [0., 0., 0., ..., 0., 0., 0.]]]], shape=(1, 35, 180, 240))
Coordinates:
  * time       (time) object 8B 2000-07-01 00:00:00
  * z_l        (z_l) float64 280B 2.5 10.0 20.0 32.5 ... 5.5e+03 6e+03 6.5e+03
  * yh         (yh) int64 1kB 0 1 2 3 4 5 6 7 ... 173 174 175 176 177 178 179
  * xh         (xh) int64 2kB 0 1 2 3 4 5 6 7 ... 233 234 235 236 237 238 239
    geolon     (yh, xh) float64 346kB ...
    lon        (yh, xh) float64 346kB ...
    geolat     (yh, xh) float64 346kB ...
    lat        (yh, xh) float64 346kB ...
    deptho     (yh, xh) float32 173kB ...
    wet        (yh, xh) float32 173kB ...
    areacello  (yh, xh) float64 346kB ...
Attributes:
    cell_measures:  volume: volcello area: areacello
    time_avg_info:  average_T1,average_T2,average_DT
    standard_name:  cell_area
    note:           We ignore land cells in partially wet cells when coarseni...
    provenance:     heat_lhs_sum_advection_sum

xarray.DataArray

'heat_lhs_sum_advection'

• time: 1
• z_l: 35
• yh: 180
• xh: 240

• 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0

  array([[[[0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           ...,
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.]],
  
          [[0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           ...,
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.]],
  
          [[0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           ...,
  ...
           ...,
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.]],
  
          [[0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           ...,
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.]],
  
          [[0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           ...,
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.],
           [0., 0., 0., ..., 0., 0., 0.]]]], shape=(1, 35, 180, 240))

• Coordinates: (11)
  • time
    (time)
    object
    2000-07-01 00:00:00

    long_name :

    historical time

    axis :

    T

    calendar_type :

    NOLEAP

    bounds :

    time_bnds

    cell_methods :

    time:mean

    array([cftime.DatetimeNoLeap(2000, 7, 1, 0, 0, 0, 0, has_year_zero=True)],
          dtype=object)

  • z_l
    (z_l)
    float64
    2.5 10.0 20.0 ... 6e+03 6.5e+03

    array([2.5000e+00, 1.0000e+01, 2.0000e+01, 3.2500e+01, 5.1250e+01, 7.5000e+01,
           1.0000e+02, 1.2500e+02, 1.5625e+02, 2.0000e+02, 2.5000e+02, 3.1250e+02,
           4.0000e+02, 5.0000e+02, 6.0000e+02, 7.0000e+02, 8.0000e+02, 9.0000e+02,
           1.0000e+03, 1.1000e+03, 1.2000e+03, 1.3000e+03, 1.4000e+03, 1.5375e+03,
           1.7500e+03, 2.0625e+03, 2.5000e+03, 3.0000e+03, 3.5000e+03, 4.0000e+03,
           4.5000e+03, 5.0000e+03, 5.5000e+03, 6.0000e+03, 6.5000e+03])

  • yh
    (yh)
    int64
    0 1 2 3 4 5 ... 175 176 177 178 179

    long_name :

    cell center y-index (nominally latitude)

    cell_methods :

    yh:point

    array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
            14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
            28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
            42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
            56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
            70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
            84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
            98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
           112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
           126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
           140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
           154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
           168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179])

  • xh
    (xh)
    int64
    0 1 2 3 4 5 ... 235 236 237 238 239

    long_name :

    cell center x-index (nominally longitude)

    cell_methods :

    xh:point

    array([  0,   1,   2, ..., 237, 238, 239], shape=(240,))

  • geolon
    (yh, xh)
    float64
    ...

    long_name :

    Longitude of tracer (T) points

    units :

    degrees_east

    cell_methods :

    time: point

    [43200 values with dtype=float64]

  • lon
    (yh, xh)
    float64
    ...

    long_name :

    Longitude of tracer (T) points

    units :

    degrees_east

    cell_methods :

    time: point

    [43200 values with dtype=float64]

  • geolat
    (yh, xh)
    float64
    ...

    long_name :

    Latitude of tracer (T) points

    units :

    degrees_north

    cell_methods :

    time: point

    [43200 values with dtype=float64]

  • lat
    (yh, xh)
    float64
    ...

    long_name :

    Latitude of tracer (T) points

    units :

    degrees_north

    cell_methods :

    time: point

    [43200 values with dtype=float64]

  • deptho
    (yh, xh)
    float32
    ...

    long_name :

    Sea Floor Depth

    units :

    m

    cell_methods :

    xh:mean yh:mean time:point

    cell_measures :

    area: areacello

    standard_name :

    sea_floor_depth_below_geoid

    [43200 values with dtype=float32]

  • wet
    (yh, xh)
    float32
    ...

    long_name :

    0 if land, 1 if ocean at tracer points

    cell_methods :

    xh:mean yh:mean time:point

    cell_measures :

    area: areacello

    note :

    Can be between 0 and 1 if coarse cell includes both wet and dry sub-cells.

    [43200 values with dtype=float32]

  • areacello
    (yh, xh)
    float64
    ...

    long_name :

    Ocean Grid-Cell Area

    units :

    m2

    cell_methods :

    area:sum yh:sum xh:sum time: point

    standard_name :

    cell_area

    note :

    We ignore land cells in partially wet cells when coarsening, so that tracer content can be accurately reconstructed by multiplying coarsened area-averaged tendencies by it. Fully wet (`wet==1.0`) and fully dry (`wet==0.0`) cells should be unaffected, and will just represent the total cell area. For the partially wet cells, total cell area can be derived from the ocean area by divding `areacello` by `wet`.

    [43200 values with dtype=float64]

• Attributes: (5)

  cell_measures :

  volume: volcello area: areacello

  time_avg_info :

  average_T1,average_T2,average_DT

  standard_name :

  cell_area

  note :

  We ignore land cells in partially wet cells when coarsening, so that tracer content can be accurately reconstructed by multiplying coarsened area-averaged tendencies by it. Fully wet (`wet==1.0`) and fully dry (`wet==0.0`) cells should be unaffected, and will just represent the total cell area. For the partially wet cells, total cell area can be derived from the ocean area by divding `areacello` by `wet`.

  provenance :

  heat_lhs_sum_advection_sum

If the variable is not immediately available as a diagnostic–but instead derived from other diagnostics–the provenance attribute points to the origin of the variable.

For variables that are the sum or product of other terms, get_vars also returns a list of the variables that feed into the operation:

get_vars(xbudget_dict, "heat_lhs_sum_advection_sum")

{'var': 'heat_lhs_sum_advection_sum',
 'sum': ['heat_lhs_sum_advection_sum_lateral',
  'heat_lhs_sum_advection_sum_interfacial']}

The terms from which the variable is derived are not necessarily other xr.DataArray objects–they can also be constant numbers.

get_vars(xbudget_dict, "heat_lhs_sum_surface_ocean_flux_advective_negative_lhs_product")

{'var': 'heat_lhs_sum_surface_ocean_flux_advective_negative_lhs_product',
 'product': [-1.0,
  3992.0,
  'tos',
  'boundary_forcing_h_tendency',
  1035.0,
  'areacello']}

The get_vars functions takes either a string or a list of strings:

get_vars(xbudget_dict, ["heat_lhs", "heat_rhs_sum_diffusion_sum", "boundary_forcing_heat_tendency"])

[{'var': 'heat_lhs'},
 {'var': 'heat_rhs_sum_diffusion_sum',
  'sum': ['heat_rhs_sum_diffusion_sum_lateral',
   'heat_rhs_sum_diffusion_sum_interfacial']},
 {'var': 'boundary_forcing_heat_tendency'}]

Verification that the budgets close

We compare the left-hand-side (LHS) tendency terms against the sum of all of the right-hand-side (RHS) terms.

for eq, vmax in zip(["mass", "heat", "salt"], [1.e4, 1e12, 1.e5]):
    plt.figure(figsize=(10, 4))

    plt.subplot(1,2,1)
    pc = grid._ds[get_vars(xbudget_dict, f"{eq}_lhs")['var']].isel(z_l = 0).isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
    plt.title(f"LHS {eq} tendency")

    plt.subplot(1,2,2)
    pc = grid._ds[get_vars(xbudget_dict, f"{eq}_rhs")['var']].isel(z_l = 0).isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
    plt.title(f"RHS {eq} tendency")

    plt.tight_layout()

Decompose RHS budget terms with xbudget

Since the full budget is relatively clunky and can be decomposed in a number of different ways, we include helper functions for collecting all of the high-level terms and optionally dissaggreating terms of interest:

# Default high-level aggregate budget
simple_budgets = xbudget.aggregate(xbudget_dict)
simple_budgets

{'mass': {'lambda': 'density',
  'thickness': 'thkcello',
  'lhs': {'Eulerian_tendency': 'mass_lhs_sum_Eulerian_tendency'},
  'rhs': {'advection': 'mass_rhs_sum_advection',
   'surface_exchange_flux': 'mass_rhs_sum_surface_exchange_flux'}},
 'heat': {'lambda': 'thetao',
  'surface_lambda': 'tos',
  'lhs': {'Eulerian_tendency': 'heat_lhs_sum_Eulerian_tendency',
   'advection': 'heat_lhs_sum_advection',
   'surface_ocean_flux_advective_negative_lhs': 'heat_lhs_sum_surface_ocean_flux_advective_negative_lhs'},
  'rhs': {'diffusion': 'heat_rhs_sum_diffusion',
   'surface_exchange_flux': 'heat_rhs_sum_surface_exchange_flux',
   'surface_ocean_flux_advective_negative_rhs': 'heat_rhs_sum_surface_ocean_flux_advective_negative_rhs',
   'bottom_flux': 'heat_rhs_sum_bottom_flux',
   'frazil_ice': 'heat_rhs_sum_frazil_ice'}},
 'salt': {'lambda': 'so',
  'surface_lambda': 'sos',
  'lhs': {'Eulerian_tendency': 'salt_lhs_sum_Eulerian_tendency',
   'advection': 'salt_lhs_sum_advection',
   'surface_ocean_flux_advective_negative_lhs': 'salt_lhs_sum_surface_ocean_flux_advective_negative_lhs'},
  'rhs': {'diffusion': 'salt_rhs_sum_diffusion',
   'surface_exchange_flux': 'salt_rhs_sum_surface_exchange_flux',
   'surface_ocean_flux_advective_negative_rhs': 'salt_rhs_sum_surface_ocean_flux_advective_negative_rhs'}}}

vmax = 1.e-4

for eq, vmax in zip(["mass", "heat", "salt"], [1.e6, 1e12, 1e4]):
    N = len(simple_budgets[eq]['rhs'])
    plt.figure(figsize=(5*N, 4))
    for i, (k,v) in enumerate(simple_budgets[eq]['rhs'].items(), start=1):
        plt.subplot(1,N, i)
        if "z_l" in grid._ds[v].dims:
            grid._ds[v].isel(z_l = 0).isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
        else:
            grid._ds[v].isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
        plt.title(f"{eq} {k}")

    plt.tight_layout()

Custom budget decompositions

Suppose we want to have a better understanding of the processes driving both the surface_flux terms and interior diffusion terms. We can simply pass the list of these process names to the decompose keyword argument. We can also further breakdown the nonadvective part of the surface_flux term.

decomposed_budgets = xbudget.aggregate(xbudget_dict, decompose=["surface_exchange_flux", "nonadvective", "diffusion"])
decomposed_budgets

{'mass': {'lambda': 'density',
  'thickness': 'thkcello',
  'lhs': {'Eulerian_tendency': 'mass_lhs_sum_Eulerian_tendency'},
  'rhs': {'advection': 'mass_rhs_sum_advection',
   'surface_exchange_flux_rain_and_ice': 'mass_rhs_sum_surface_exchange_flux_sum_rain_and_ice',
   'surface_exchange_flux_snow': 'mass_rhs_sum_surface_exchange_flux_sum_snow',
   'surface_exchange_flux_evaporation': 'mass_rhs_sum_surface_exchange_flux_sum_evaporation',
   'surface_exchange_flux_rivers': 'mass_rhs_sum_surface_exchange_flux_sum_rivers',
   'surface_exchange_flux_icebergs': 'mass_rhs_sum_surface_exchange_flux_sum_icebergs',
   'surface_exchange_flux_sea_ice_melt': 'mass_rhs_sum_surface_exchange_flux_sum_sea_ice_melt',
   'surface_exchange_flux_virtual_precip_restoring': 'mass_rhs_sum_surface_exchange_flux_sum_virtual_precip_restoring'}},
 'heat': {'lambda': 'thetao',
  'surface_lambda': 'tos',
  'lhs': {'Eulerian_tendency': 'heat_lhs_sum_Eulerian_tendency',
   'advection': 'heat_lhs_sum_advection',
   'surface_ocean_flux_advective_negative_lhs': 'heat_lhs_sum_surface_ocean_flux_advective_negative_lhs'},
  'rhs': {'diffusion_lateral': 'heat_rhs_sum_diffusion_sum_lateral',
   'diffusion_interfacial': 'heat_rhs_sum_diffusion_sum_interfacial',
   'surface_exchange_flux_nonadvective_latent': 'heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_latent',
   'surface_exchange_flux_nonadvective_sensible': 'heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_sensible',
   'surface_exchange_flux_nonadvective_longwave': 'heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_longwave',
   'surface_exchange_flux_nonadvective_shortwave': 'heat_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_shortwave',
   'surface_exchange_flux_advective': 'heat_rhs_sum_surface_exchange_flux_sum_advective',
   'surface_ocean_flux_advective_negative_rhs': 'heat_rhs_sum_surface_ocean_flux_advective_negative_rhs',
   'bottom_flux': 'heat_rhs_sum_bottom_flux',
   'frazil_ice': 'heat_rhs_sum_frazil_ice'}},
 'salt': {'lambda': 'so',
  'surface_lambda': 'sos',
  'lhs': {'Eulerian_tendency': 'salt_lhs_sum_Eulerian_tendency',
   'advection': 'salt_lhs_sum_advection',
   'surface_ocean_flux_advective_negative_lhs': 'salt_lhs_sum_surface_ocean_flux_advective_negative_lhs'},
  'rhs': {'diffusion_lateral': 'salt_rhs_sum_diffusion_sum_lateral',
   'diffusion_interfacial': 'salt_rhs_sum_diffusion_sum_interfacial',
   'surface_exchange_flux_nonadvective_basal': 'salt_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_basal',
   'surface_exchange_flux_advective': 'salt_rhs_sum_surface_exchange_flux_sum_advective',
   'surface_ocean_flux_advective_negative_rhs': 'salt_rhs_sum_surface_ocean_flux_advective_negative_rhs'}}}

decomposed_budgets[eq]['rhs']

{'diffusion_lateral': 'salt_rhs_sum_diffusion_sum_lateral',
 'diffusion_interfacial': 'salt_rhs_sum_diffusion_sum_interfacial',
 'surface_exchange_flux_nonadvective_basal': 'salt_rhs_sum_surface_exchange_flux_sum_nonadvective_sum_basal',
 'surface_exchange_flux_advective': 'salt_rhs_sum_surface_exchange_flux_sum_advective',
 'surface_ocean_flux_advective_negative_rhs': 'salt_rhs_sum_surface_ocean_flux_advective_negative_rhs'}

vmax = 1.e-4
import numpy as np

for eq, vmax in zip(["mass", "heat", "salt"], [1e6, 1e12, 1e4]):
    nrow = np.int64(np.ceil(len(decomposed_budgets[eq]['rhs'])/2))
    plt.figure(figsize=(12, nrow*4))
    for i, (k,v) in enumerate(decomposed_budgets[eq]['rhs'].items(), start=1):
        plt.subplot(nrow, 2, i)
        if "z_l" in grid._ds[v].dims:
            grid._ds[v].isel(z_l = 0).isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
        else:
            grid._ds[v].isel(time=0).plot(vmin=-vmax, vmax=vmax, cmap="RdBu_r")
        plt.title(f"{eq} {k}")

    plt.tight_layout()
    plt.show()