Merge branch 'master' of gitlab.dkrz.de:m300602/pyicon

pyicon/__init__.py

@@ -22,6 +22,7 @@ from .pyicon_plotting import *
 from .pyicon_accessor import *
 #print('-----simulation')
 from .pyicon_simulation import *
+from .pyicon_thermo import *
 
 # --- import pyicon.view
 #print('-----view')

pyicon/pyicon_calc_xr.py

 import warnings
-#print('numpy')
 import numpy as np
-#print('xarray')
 import xarray as xr
 from itertools import product
-#print('Done modules calc.')
+
 
 def convert_tgrid_data(ds_tg, check_previous_conversion=True,
                        set_dim_order=None, old_dim_behaviour=None):
@@ -36,6 +34,7 @@ def convert_tgrid_data(ds_tg, check_previous_conversion=True,
     ds_IcD : xr.Dataset
         A tgrid dataset compatible with pyicon functions
 
+
     Notes
     -----
     Open classical ICON grid file by:
@@ -181,24 +180,33 @@ def convert_tgrid_data(ds_tg, check_previous_conversion=True,
 
     return ds_IcD
 
-def print_verbose(verbose=1, message="", verbose_stage=1):
+
+def _print_verbose(verbose=1, message="", verbose_stage=1):
+    """ Prints message depending on verbosity
+    """
     if verbose >= verbose_stage:
         print(message)
     return
 
+
 def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
     """ Crop a grid file.
 
-    Input:
-    ------
-    ds_tg: xarray Dataset, which contains the grid file
-    ireg_c: numpy index list of cell-points which should by in cropped domain
+    Parameters
+    ----------
+    ds_tg : xr.Dataset
+        dataset containing the grid file
+    
+    ireg_c : np.array
+        list of cell indices which should by in cropped domain
+
 
-    Output:
-    ------
+    Returns
+    -------
     ds_tg_cut: xarray Dataset, which contains (most of) the cropped grid variables.
 
-    Example usage:
+
+    Example usage
     --------------
     ds_tg = xr.open_mfdataset(fpath_tgrid)
     clon = ds_tg.clon.compute().data * 180./np.pi
@@ -218,14 +226,19 @@ def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
     else: offset = 1
 
     # --- find edges and vertices belonging to cells of cutted domain
-    print_verbose(verbose, "find edges")
+    _print_verbose(verbose, "find edges")
     vertex_of_cell = ds_tg.vertex_of_cell.isel(cell=ireg_c).compute().data - offset
     edge_of_cell = ds_tg.edge_of_cell.isel(cell=ireg_c).compute().data - offset
     ireg_e, inde = np.unique(edge_of_cell, return_index=True)
     ireg_v, indv = np.unique(vertex_of_cell, return_index=True)
 
+    ireg_c = ireg_c.astype("int32")
+    ireg_e = ireg_e.astype("int32")
+    ireg_v = ireg_v.astype("int32")
+    
+    
     # --- new dataset with cutted coordinates
-    print_verbose(verbose, "cut coordinates")
+    _print_verbose(verbose, "cut coordinates")
     ds_tg_cut = xr.Dataset(coords=dict(
         clon=ds_tg['clon'][ireg_c],
         clat=ds_tg['clat'][ireg_c],
@@ -239,7 +252,7 @@ def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
     ds_tg_cut['ireg_v'] = xr.DataArray(ireg_v, dims=['vertex'])
 
     # --- re-index
-    print_verbose(verbose, "reindex")
+    _print_verbose(verbose, "reindex")
     reindex_c = np.zeros_like(ds_tg.clon, dtype='int32') - offset
     reindex_c[ireg_c] = np.arange(ireg_c.size, dtype='int32')
     reindex_e = np.zeros_like(ds_tg.elon, dtype='int32') - offset
@@ -249,37 +262,37 @@ def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
 
     var = 'vertex_of_cell'
     da = ds_tg[var].isel(cell=ireg_c) - offset
-    data = reindex_v[da.data.flatten()].reshape(da.shape)
+    data = reindex_v[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = "edge_vertices"
     da = ds_tg[var].isel(edge=ireg_e) - offset
-    data = reindex_v[da.data.flatten()].reshape(da.shape)
+    data = reindex_v[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = 'vertices_of_vertex'
     da = ds_tg[var].isel(vertex=ireg_v) - offset
-    data = reindex_v[da.data.flatten()].reshape(da.shape)
+    data = reindex_v[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = 'edge_of_cell'
     da = ds_tg[var].isel(cell=ireg_c) - offset
-    data = reindex_e[da.data.flatten()].reshape(da.shape)
+    data = reindex_e[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = 'edges_of_vertex'
     da = ds_tg[var].isel(vertex=ireg_v) - offset
-    data = reindex_e[da.data.flatten()].reshape(da.shape)
+    data = reindex_e[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = 'adjacent_cell_of_edge'
     da = ds_tg[var].isel(edge=ireg_e) - offset
-    data = reindex_c[da.data.flatten()].reshape(da.shape)
+    data = reindex_c[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     var = 'cells_of_vertex'
     da = ds_tg[var].isel(vertex=ireg_v) - offset
-    data = reindex_c[da.data.flatten()].reshape(da.shape)
+    data = reindex_c[da.data.flatten().compute().astype("int32")].reshape(da.shape)
     ds_tg_cut[var] = xr.DataArray(data + offset, dims=da.dims)
 
     reindex_vars = [
@@ -326,8 +339,36 @@ def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
                 ds_tg_cut[var] = ds_tg[var]
     return ds_tg_cut
 
+
 ## Functions to map between 3D Cartesian and 2D local vectors
 def xr_calc_2dlocal_from_3d(ds_IcD, p_vn_c):
+    """ Transform vector from cartesian to spherical basis at a cell center
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    p_vn_c : xr.Dataset
+        dataset containing cartesian representation of
+        vector. Should have the dimension 'cart'
+
+
+    Returns
+    -------
+    uo : xr.DataArray
+        zonal component of vector
+
+    vo : xr.DataArray
+        meridional component of vector
+
+
+    Notes
+    -----
+    The 3D vector passed is not (u, v, w) where w is the local
+    vertical.
+    
+    """
     sinLon = np.sin(ds_IcD.clon*np.pi/180.)
     cosLon = np.cos(ds_IcD.clon*np.pi/180.)
     sinLat = np.sin(ds_IcD.clat*np.pi/180.)
@@ -344,7 +385,34 @@ def xr_calc_2dlocal_from_3d(ds_IcD, p_vn_c):
     vo = -(u1*cosLon + u2*sinLon)*sinLat + u3*cosLat
     return uo, vo
 
+
 def xr_calc_3d_from_2dlocal(ds_IcD, uo, vo):
+    """ Transform vector from spherical to cartesian basis at a cell center
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+    
+    uo : xr.DataArray
+        zonal component
+        
+    vo : xr.DataArray
+        meridional component
+
+    Returns
+    -------
+    p_vn_c : xr.DataArray
+        representation of the horizontal vector in a cartesian basis.
+        
+
+    Notes
+    -----
+    The components of the returned vector *are not* zonal, meridional and vertical.
+    This function transforms a locally horizontal vector from a spherical representation
+    to a cartesian representation.
+    
+    """
     sinLon = np.sin(ds_IcD.clon*np.pi/180.)
     cosLon = np.cos(ds_IcD.clon*np.pi/180.)
     sinLat = np.sin(ds_IcD.clat*np.pi/180.)
@@ -358,10 +426,23 @@ def xr_calc_3d_from_2dlocal(ds_IcD, uo, vo):
     p_vn_c = xr.concat([u1,u2,u3], dim='cart', coords='minimal').transpose(*new_dims)
     return p_vn_c
 
-## Mapping between cells and edges
 
+## Mapping between cells and edges
 def xr_calc_edge2cell_coeff_cc_t(ds_IcD):
-    dist_vector = ds_IcD.edge_cart_vec - ds_IcD.cell_cart_vec.isel(cell=ds_IcD.adjacent_cell_of_edge)#).transpose('edge', 'nc', 'cart')
+    """ Calculates the cell to edge coefficients
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    Returns
+    -------
+    edge2cell_coeff_cc_t : xr.DataArray
+        coefficients used in mapping cells to edges
+    """
+    dist_vector = ds_IcD.edge_cart_vec - ds_IcD.cell_cart_vec.isel(cell=ds_IcD.adjacent_cell_of_edge)
     orientation = (dist_vector*ds_IcD.edge_prim_norm).sum(dim='cart')
     dist_vector *= np.sign(orientation)
     edge2cell_coeff_cc_t = (  ds_IcD.edge_prim_norm*ds_IcD.grid_sphere_radius
@@ -370,9 +451,30 @@ def xr_calc_edge2cell_coeff_cc_t(ds_IcD):
     edge2cell_coeff_cc_t = edge2cell_coeff_cc_t.transpose('edge', 'nc_e', 'cart')
     return edge2cell_coeff_cc_t
 
+
 def xr_cell2edges(ds_IcD, p_vn_c, edge2cell_coeff_cc_t=None):
+    """ Remaps vector from cell center to edges
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    p_vn_c : xr.DataArray
+        locally horizontal vector on cell center in cartesian
+        representation.
+
+    edge2cell_coeff_cc_t : xr.DataArray or None
+        coefficients used in mapping cells to edges
+
+
+    Returns
+    -------
+    ptp_vn : xr.DataArray
+        vector p_vn_c remapped to edges
+    """
     if edge2cell_coeff_cc_t is None:
-      edge2cell_coeff_cc_t = xr_calc_edge2cell_coeff_cc_t(ds_IcD)
+        edge2cell_coeff_cc_t = xr_calc_edge2cell_coeff_cc_t(ds_IcD)
     ic0 = ds_IcD.adjacent_cell_of_edge.isel(nc_e=0).data
     ic1 = ds_IcD.adjacent_cell_of_edge.isel(nc_e=1).data
     ptp_vn = (
@@ -388,9 +490,22 @@ def xr_cell2edges(ds_IcD, p_vn_c, edge2cell_coeff_cc_t=None):
     )
     return ptp_vn
 
-## Mapping between edges and cells
 
+## Mapping between edges and cells
 def xr_calc_fixed_volume_norm(ds_IcD):
+    """ Calculates the fixed volume
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    Returns
+    -------
+    fixed_vol_norm : xr.DataArray
+        volume of grid cell
+    """
     dist_vector = (
         ds_IcD.edge_cart_vec.isel(edge=ds_IcD.edge_of_cell) 
         - ds_IcD.cell_cart_vec
@@ -404,7 +519,21 @@ def xr_calc_fixed_volume_norm(ds_IcD):
     fixed_vol_norm = fixed_vol_norm.sum(dim='ne_c')
     return fixed_vol_norm
 
+
 def xr_calc_edge2cell_coeff_cc(ds_IcD):
+    """ Calculates the edge to cell coefficients
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    Returns
+    -------
+    edge2cell_coeff_cc : xr.DataArray
+        coefficients used in mapping edges to cells
+    """
     dist_vector = (
         ds_IcD.edge_cart_vec.isel(edge=ds_IcD.edge_of_cell) - ds_IcD.cell_cart_vec
     )
@@ -414,7 +543,37 @@ def xr_calc_edge2cell_coeff_cc(ds_IcD):
     edge2cell_coeff_cc = edge2cell_coeff_cc.compute()
     return edge2cell_coeff_cc
 
+
 def xr_edges2cell(ds_IcD, ve, dze, dzc, edge2cell_coeff_cc=None, fixed_vol_norm=None):
+    """ Remaps vector from edges to cell
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    ve : xr.DataArray
+        vector on edges
+
+    dxe : xr.DataArray
+        edge length
+
+    dzc : xr.DataArray
+        vertical grid spacing at cell centre
+
+    edge2cell_coeff_cc : xr.DataArray or None
+        coefficients used in mapping edges to cells
+
+    fixed_vol_norm : xr.DataArray or None
+        volume of grid cell
+    
+
+    Returns
+    -------
+    p_vn_c : xr.DataArray
+        cartesian representation of vector ve on cell centres
+    """
     if fixed_vol_norm is None:
         fixed_vol_norm = xr_calc_fixed_volume_norm(ds_IcD)
     if edge2cell_coeff_cc is None:
@@ -438,34 +597,102 @@ def xr_edges2cell(ds_IcD, ve, dze, dzc, edge2cell_coeff_cc=None, fixed_vol_norm=
     return p_vn_c
 
 ## Mapping between edges and edges
-
 def xr_calc_edge2edge_viacell_coeff(ds_IcD):
+    """ Calculate coefficients for mapping edge vectors to edge vectors
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+    """
     raise NotImplementedError("")
     # FIXME: Continue here
     edge2edge_viacell_coeff = ()
     return edge2edge_viacell_coeff
 
+
 def xr_edges2edges_via_cell(ds_IcD, vn_e, dze='const'):
+    """ Maps edges to edges via a cell
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    vn_e : xr.DataArray
+        vector at edge
+
+    dze : xr.DataArray or 'const'
+    """
     raise NotImplementedError("")
     # FIXME: Continue here
     out_vn_e = ()
     return out_vn_e
 
-def xr_edges2edges_via_cell(ds_IcD, vn_e, scalar, dze='const'):
+
+def xr_edges2edges_via_cell_scalar(ds_IcD, vn_e, scalar, dze='const'):
+    """ Calculates flux of scalar at edges
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    vn_e : xr.DataArray
+        vector at edge
+
+    scalar : xr.DataArray
+        scalar at cell centre
+
+    dze : xr.DataArray or 'const'    
+    """
     raise NotImplementedError("")
     # FIXME: Continue here
     out_vn_e = ()
     return out_vn_e
 
-## Divergence
 
+## Divergence
 def xr_calc_div_coeff(ds_IcD):
+    """ Calculates coefficients for calculating divergence
+    
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    Returns
+    -------
+    div_coeff : xr.DataArray
+        coefficients for calculating divergence
+    """
     div_coeff = (
         ds_IcD.edge_length.isel(edge=ds_IcD.edge_of_cell) * ds_IcD.orientation_of_normal / ds_IcD.cell_area
     )
     return div_coeff
 
+
 def xr_calc_div(ds_IcD, vector, div_coeff=None):
+    """ Calculates coefficients for calculating divergence
+    
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    vector : xr.DataArray
+        vector at cell edges
+
+    div_coeff : xr.DataArray
+        coefficients for calculating divergence
+    
+    Returns
+    -------
+    div_of_vector : xr.DataArray
+        divergence of vector at cell centers
+    
+    """
     if div_coeff is None:
         div_coeff = xr_calc_div_coeff(ds_IcD)
     div_of_vector = (
@@ -475,14 +702,46 @@ def xr_calc_div(ds_IcD, vector, div_coeff=None):
     return div_of_vector
 
 ## Gradient
-
 def xr_calc_grad_coeff(ds_IcD):
+    """ Calculates coefficients for calculating gradient
+    
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+
+    Returns
+    -------
+    grad_coeff : xr.DataArray
+        coefficients for calculating gradient
+    """
     grad_coeff = (
         1./ds_IcD.dual_edge_length
     )
     return grad_coeff
 
+
 def xr_calc_grad(ds_IcD, scalar, grad_coeff=None):
+    """ Calculates coefficients for calculating gradient
+    
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
+
+    scalar : xr.DataArray
+        scalar at cell center
+
+    grad_coeff : xr.DataArray
+        coefficients for calculating gradient
+
+    
+    Returns
+    -------
+    grad_of_scalar : xr.DataArray
+        horizontal gradient of scalar at edges
+    """
     if grad_coeff is None:
         grad_coeff = xr_calc_grad_coeff(ds_IcD)
     grad_of_scalar = (
@@ -493,15 +752,19 @@ def xr_calc_grad(ds_IcD, scalar, grad_coeff=None):
 
 
 # Curl
+def xr_calc_rot_coeff(ds_IcD):
+    """ Calculates coefficients used in calculating the curl
+    
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        pyicon dataset containing coordinate info
 
 
-def xr_calc_rot_coeff(ds_IcD):
-    """
-    Input:
-    ------
-    ds_IcD: xarray Dataset, which contains the grid file and was modified by
-            convert_tgrid_data.
-            
+    Returns
+    -------
+    curl_coeffs : xr.DataArray
+        coefficients for calculating curls    
     """
     curl_coeffs = ds_IcD["edge_orientation"] * ds_IcD["dual_edge_length"].isel(edge=ds_IcD["edges_of_vertex"]) / ds_IcD["dual_area"]
     return curl_coeffs
@@ -521,11 +784,13 @@ def xr_calc_curl(ds_IcD, vector, rot_coeff=None):
     rot_coeff : xr.DataArray or None
         Array containing dims ("vertex", "ne_v")
 
+
     Returns
     -------
     curl_vec : xr.DataArray
         vertical component of the curl of the vector defined on vertex points
 
+
     Notes
     -----
     We calculate the curl through the use of Stokes'/Green's theorem. A similar

pyicon/pyicon_thermo.py

+import numpy as np
+import xarray as xr
+
+
+def _calculate_eos_pressure(da_depth, da_zos=0, da_stretch_c=1):
+    """calculates hydrostatic pressure for the equation of state
+
+    Parameters
+    ----------
+    da_depth : xr.DataArray
+        depth in the water column in metres
+        
+    da_zos : xr.DataArray or 0, optional
+        sea surface height in metres
+        
+    da_stretch_c : xr.DataArray or 1, optional
+        stretch factor from zstar
+
+        
+    Returns
+    -------
+    p : xr.DataArray
+        hydrostatic pressure as required by the equation of state in dBar
+
+        
+    Notes
+    -----
+    The pressure does not include atmospheric pressure. This is neglected by
+    the equation of state used in ICON.
+    """
+    SItodBar = 1e-4
+    rho_ref = 1025.022
+
+    p = (da_stretch_c * da_depth + da_zos) * rho_ref * SItodBar
+    return p
+
+
+def _calculate_insitu_temperature(sal, t_pot, p):
+    """calculates in-situ temperature from model variables and pressure.
+
+    Parameters
+    ----------
+    sal : xr.DataArray
+        salinity in PSU
+    
+    t_pot : xr.DataArray
+        potential temperature (what ICON outputs) in deg C
+    
+    p : xr.DataArray
+        pressure in dBar
+
+
+    Returns
+    -------
+    t_insitu : xr.DataArray
+        in-situ temperature in deg C
+    """
+    a_a1 = 3.6504e-4
+    a_a2 = 8.3198e-5
+    a_a3 = 5.4065e-7
+    a_a4 = 4.0274e-9
+    
+    a_b1 = 1.7439e-5
+    a_b2 = 2.9778e-7
+    
+    a_c1 = 8.9309e-7
+    a_c2 = 3.1628e-8
+    a_c3 = 2.1987e-10
+    
+    a_d = 4.1057e-9
+    
+    a_e1 = 1.6056e-10
+    a_e2 = 5.0484e-12
+    
+    z_sref = 35e0
+    
+    
+    qnq = -p * (-a_a3 + p * a_c3)
+    qn3 = -p * a_a4
+    qvs = (p * (a_b1 - a_d * p)) * (sal - z_sref) + p * (a_a1 + p * (a_c1 - a_e1 * p))
+    dvs = (a_b2 * p) * (sal - z_sref) + 1 + p * (-a_a2 + p * (a_c2 - a_e2 * p))
+
+    t   = (t_pot + qvs) / dvs
+    fne = - qvs + t * (dvs + t * (qnq + t * qn3)) - t_pot
+    fst = dvs + t * (2 * qnq + 3 * qn3 * t)
+
+    t_insitu   = t - fne / fst
+    return t_insitu
+
+
+def _calculate_mpiom_density(sal, t_insitu, p):
+    """calculates potential density
+
+    Parameters
+    ----------
+    sal : xr.DataArray
+        salinity in PSU
+    
+    t_insitu : xr.DataArray
+        in-situ temperature in deg C
+    
+    p : xr.DataArray
+        pressure in dBar
+
+
+    Returns
+    -------
+    rho : xr.DataArray
+        density in kg / m^3
+    """
+    r_a0=999.842594
+    r_a1=6.793952e-2
+    r_a2=-9.095290e-3
+    r_a3=1.001685e-4
+    r_a4=-1.120083e-6
+    r_a5=6.536332e-9
+    
+    r_b0 = 8.24493e-1
+    r_b1 = -4.0899e-3
+    r_b2 = 7.6438e-5
+    r_b3 = -8.2467e-7
+    r_b4 = 5.3875e-9
+    
+    r_c0 = -5.72466e-3
+    r_c1 = 1.0227e-4
+    r_c2 = -1.6546e-6
+    
+    r_d0=4.8314e-4
+    
+    r_e1 = 148.4206
+    r_e0 = 19652.21
+    r_e2 = -2.327105
+    r_e3 = 1.360477e-2
+    r_e4 = -5.155288e-5
+    
+    r_f0 = 54.6746
+    r_f1 = -0.603459
+    r_f2 = 1.09987e-2
+    r_f3 = -6.1670e-5
+    
+    r_g0 = 7.944e-2
+    r_g1 = 1.6483e-2
+    r_g2 = -5.3009e-4
+      
+    r_h0 = 3.239908
+    r_h1 = 1.43713e-3
+    r_h2 = 1.16092e-4
+    r_h3 = -5.77905e-7
+    
+    r_ai0 = 2.2838e-3
+    r_ai1 = -1.0981e-5
+    r_ai2 = -1.6078e-6
+    
+    r_aj0 = 1.91075e-4
+    
+    r_ak0 = 8.50935e-5
+    r_ak1 = -6.12293e-6
+    r_ak2 = 5.2787e-8
+    
+    r_am0 = -9.9348e-7
+    r_am1 = 2.0816e-8
+    r_am2 = 9.1697e-10
+
+    t = t_insitu
+    s = xr.where(sal > 0.0, sal, 0.0)
+    s_2 = np.square(s)
+    s3h  = s * np.sqrt(s)
+    
+    rho = r_a0 + t * (r_a1 + t * (r_a2 + t * (r_a3 + t * (r_a4 + t * r_a5)))) \
+        + s * (r_b0 + t * (r_b1 + t * (r_b2 + t * (r_b3 + t * r_b4)))) \
+        + r_d0 * s_2 + s3h * (r_c0 + t * (r_c1 + r_c2 * t))
+
+    denom = 1.0 - p / (p * (r_h0 + t * (r_h1 + t * (r_h2 + t * r_h3)) \
+        + s * (r_ai0 + t * (r_ai1 + r_ai2 * t)) + r_aj0 * s3h + (r_ak0 \
+        + t * (r_ak1 + t * r_ak2) \
+        + s * (r_am0 + t * (r_am1 + t * r_am2))) * p) + r_e0 \
+        + t * (r_e1 + t * (r_e2 + t * (r_e3 + t * r_e4))) \
+        + s * (r_f0 + t * (r_f1 + t * (r_f2 + t * r_f3))) \
+        + s3h * (r_g0 + t * (r_g1 + r_g2 * t)))
+    rho = rho/denom
+
+    return rho
+
+
+def calculate_density(so, to, zos=0, stretch_c=1, depth=None, eos_type="mpiom"):
+    """calculates density from model variables
+
+    Parameters
+    ----------
+    to : xr.DataArray
+        potential temperature in deg C
+    
+    so : xr.DataArray
+        salinity in PSU
+    
+    zos : xr.DataArray or 0, optional
+        sea surface height in metres, by default 0
+    
+    stretch_c : xr.DataArray or 1, optional
+        zos stretch factor, by default 1
+    
+    depth : xr.DataArray, optional
+        Array containing depth in metres at prism centres, by default taken
+        from to. Depth should be increasingly positive with distance below the
+        surface.
+    
+    eos_type : str, optional
+        which equation of state to use, by default "mpiom"
+
+
+    Returns
+    -------
+    rho : xr.DataArray
+        potential density
+
+
+    Raises
+    ------
+    TypeError
+        raised when depth not provided as a kwarg and also not found as a 
+        coordinate on to
+    
+    NotImplementedError
+        raised when an unrecognised equation of state is requested
+
+        
+    Notes
+    -----
+    Equations of state are not necessarily linear. When this is the case, the
+    time mean density can not be obtained from the the time mean temperature
+    and salinity. To minimise error you should either calculate the density
+    online at model run time, or use the highest temporal frequency of output
+    available in your calculations.
+    """    
+    if depth is None:
+        try:
+            depth = to["depth"]
+        except KeyError:
+            raise TypeError(
+                "depth not found in to. Please provide depth as a kwarg."
+                )
+    
+    if eos_type == "mpiom" or eos_type == "gill":
+        # Then this is mpiom equation of state
+        p = _calculate_eos_pressure(depth, da_zos=zos, da_stretch_c=stretch_c)
+        t_insitu = _calculate_insitu_temperature(so, to, p)
+        rho = _calculate_mpiom_density(so, t_insitu, p)
+        
+    else:
+        raise NotImplementedError("The requested eos is not implemented")
+    
+    return rho
+        
\ No newline at end of file

pyicon/tests/test_pyicon_thermo.py

+import numpy as np
+from pyicon.pyicon_thermo import (
+    _calculate_mpiom_density,
+    _calculate_insitu_temperature
+)
+
+def test_mpiom_density():
+    """the test values come from A3.2 of Gill (1982), Atmospheric & Ocean
+    Dynamics
+    """
+    assert np.allclose(_calculate_mpiom_density(0, 5, 0), 999.96675, atol=1e-5)
+    assert np.allclose(_calculate_mpiom_density(35, 5, 0), 1027.67547, atol=1e-5)
+    assert np.allclose(_calculate_mpiom_density(35, 25, 1000), 1062.53817, atol=1e-5)
+    
+def test_insitu_temperature():
+    """the test values come from A3.5 of Gill (1982), Atmospheric & Ocean
+    Dynamics
+    """
+    assert np.allclose(_calculate_insitu_temperature(25, 8.4678516, 1000), 10, atol=1e-7)
\ No newline at end of file