Vertical Coordinate Transformation

Transforming Vertical Coordinates¶

Description¶

Recipe showing how to conservatively transform the vertical coordinate system on which a variable is stored using annual-mean outputs from the National Oceanography Centre Near-Present-Day global eORCA1 configuration of NEMO forced using ERA-5 climatologically adjusted atmospheric forcing from 1976-2024.

For more details on this model configuration and the available outputs, users can explore the Near-Present-Day documentation here.

In [1]:

Copied!





# -- Import required packages -- #
import numpy as np
import xarray as xr
from nemo_cookbook import NEMODataTree
# -- Import required packages -- #
import numpy as np
import xarray as xr
from nemo_cookbook import NEMODataTree

Optional: Connect Client to Dask Local Cluster to run analysis in parallel.¶

Note: Although using Dask is not strictly necessary for this simple example using eORCA1, if we wanted to generalise this recipe to eORCA025 or eORCA12 outputs, using Dask would be essential to avoid unnecessary slow calculations using only a single process.

In [ ]:

Copied!





# -- Initialise Dask Local Cluster -- #
import os
import dask
from dask.distributed import Client, LocalCluster

# Update temporary directory for Dask workers:
dask.config.set({'temporary_directory': f"{os.getcwd()}/dask_tmp",
                 'local_directory': f"{os.getcwd()}/dask_tmp"
                 })

# Create Local Cluster:
cluster = LocalCluster(n_workers=4, threads_per_worker=3, memory_limit='5GB')
client = Client(cluster)
client
# -- Initialise Dask Local Cluster -- #
import os
import dask
from dask.distributed import Client, LocalCluster

# Update temporary directory for Dask workers:
dask.config.set({'temporary_directory': f"{os.getcwd()}/dask_tmp",
                 'local_directory': f"{os.getcwd()}/dask_tmp"
                 })

# Create Local Cluster:
cluster = LocalCluster(n_workers=4, threads_per_worker=3, memory_limit='5GB')
client = Client(cluster)
client

Accessing NEMO Model Data¶

Let's begin by loading the grid variables for our eORCA1 NEMO model from the JASMIN Object Store.

Alternatively, you can replace the domain_filepath below with a file path to your domain_cfg.nc file and read this with xarray's open_dataset() function.

In [2]:

Copied!

# Define directory path to ancillary files:
domain_filepath = "https://noc-msm-o.s3-ext.jc.rl.ac.uk/npd-eorca1-jra55v1/domain"

# Open eORCA1 NEMO model domain_cfg:
ds_domain = xr.open_zarr(domain_filepath, consolidated=True, chunks={})

ds_domain
# Define directory path to ancillary files:
domain_filepath = "https://noc-msm-o.s3-ext.jc.rl.ac.uk/npd-eorca1-jra55v1/domain"

# Open eORCA1 NEMO model domain_cfg:
ds_domain = xr.open_zarr(domain_filepath, consolidated=True, chunks={})

ds_domain

Out[2]:

<xarray.Dataset> Size: 669MB
Dimensions:        (y: 331, x: 360, t: 1, z: 75)
Dimensions without coordinates: y, x, t, z
Data variables: (12/58)
    atlmsk         (y, x) float32 477kB dask.array<chunksize=(331, 360), meta=np.ndarray>
    bathy_metry    (t, y, x) float32 477kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
    bottom_level   (t, y, x) int32 477kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
    e1f            (t, y, x) float64 953kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
    e1t            (t, y, x) float64 953kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
    e1u            (t, y, x) float64 953kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
    ...             ...
    nav_lev        (z) float32 300B dask.array<chunksize=(75,), meta=np.ndarray>
    nav_lon        (y, x) float32 477kB dask.array<chunksize=(331, 360), meta=np.ndarray>
    pacmsk         (y, x) float32 477kB dask.array<chunksize=(331, 360), meta=np.ndarray>
    socmsk         (y, x) float32 477kB dask.array<chunksize=(331, 360), meta=np.ndarray>
    time_counter   (t) float64 8B dask.array<chunksize=(1,), meta=np.ndarray>
    top_level      (t, y, x) int32 477kB dask.array<chunksize=(1, 331, 360), meta=np.ndarray>
Attributes:
    NCO:      netCDF Operators version 5.2.2 (Homepage = http://nco.sf.net, C...
    history:  Fri Oct  4 10:58:13 2024: ncks -O --4 --no_abc --cnk_plc=xpl --...

Next, we need to import the sea water conservative temperature and absolute salinity stored at T-points in a single dataset.

In [3]:

Copied!





# Define directory path to model output files:
output_dir = "https://noc-msm-o.s3-ext.jc.rl.ac.uk/npd-eorca1-jra55v1/T1y"

# Construct NEMO model grid dataset, including seawater conservative temperature (C) and absolute salinity (g/kg):
ds_gridT = xr.merge([xr.open_zarr(f"{output_dir}/{var}", consolidated=True, chunks={'deptht': 75})[var] for var in ['thetao_con', 'e3t']], compat="override")

# Subsetting the time_counter dimension: 
ds_gridT = ds_gridT.sel(time_counter=slice('2000-01', '2010-12'))
# Define directory path to model output files:
output_dir = "https://noc-msm-o.s3-ext.jc.rl.ac.uk/npd-eorca1-jra55v1/T1y"

# Construct NEMO model grid dataset, including seawater conservative temperature (C) and absolute salinity (g/kg):
ds_gridT = xr.merge([xr.open_zarr(f"{output_dir}/{var}", consolidated=True, chunks={'deptht': 75})[var] for var in ['thetao_con', 'e3t']], compat="override")

# Subsetting the time_counter dimension: 
ds_gridT = ds_gridT.sel(time_counter=slice('2000-01', '2010-12'))

Creating a NEMODataTree¶

Next, let's create a NEMODataTree to store our domain and T- & V-grid variables for the eORCA1 model.

In [4]:

Copied!





# Define dictionary of grid datasets defining eORCA1 parent model domain with no child/grand-child nests:
# Note: domain_cfg z-dimension is expected to be named 'nav_lev'.
datasets = {"parent": {"domain": ds_domain.rename({"z": "nav_lev"}), "gridT": ds_gridT}}

# Initialise a new NEMODataTree whose parent domain is zonally periodic & north-folding on F-points:
nemo = NEMODataTree.from_datasets(datasets=datasets, iperio=True, nftype="F")

nemo
# Define dictionary of grid datasets defining eORCA1 parent model domain with no child/grand-child nests:
# Note: domain_cfg z-dimension is expected to be named 'nav_lev'.
datasets = {"parent": {"domain": ds_domain.rename({"z": "nav_lev"}), "gridT": ds_gridT}}

# Initialise a new NEMODataTree whose parent domain is zonally periodic & north-folding on F-points:
nemo = NEMODataTree.from_datasets(datasets=datasets, iperio=True, nftype="F")

nemo

Out[4]:

<xarray.DatasetView> Size: 176B
Dimensions:        (time_counter: 11)
Coordinates:
    time_centered  (time_counter) datetime64[ns] 88B dask.array<chunksize=(1,), meta=np.ndarray>
  * time_counter   (time_counter) datetime64[ns] 88B 2000-07-02 ... 2010-07-0...
Data variables:
    *empty*
Attributes:
    nftype:   F
    iperio:   True

Exploring our eORCA1 NEMO model vertical coordinate system¶

Let's begin by visualising the conservative temperature in the original vertical coordinates of our eORCA1 NPD simulation by plotting a vertical profile at a location in the North Atlantic Ocean:

By examining the vertical profile & the size of the k grid coordinate variable, we can see that our eORCA1 configuration has 75 unevenly-spaced z*-coordinate levels.

In [5]:

Copied!

# Plot an example time-mean vertical conservative temperature profile in the subpolar North Atlantic:
nemo['/gridT']['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht', ylim=(3600, 0), marker='o')

# Size of the vertical coordinate dimension:
print(nemo['/gridT']['k'].size)
# Plot an example time-mean vertical conservative temperature profile in the subpolar North Atlantic:
nemo['/gridT']['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht', ylim=(3600, 0), marker='o')

# Size of the vertical coordinate dimension:
print(nemo['/gridT']['k'].size)

No description has been provided for this image

Transforming the eORCA1 vertical coordinate system¶

Next, let's apply a conservative vertical coordinate transformation to remap our conservative temperature output to a new regularly-spaced (200 m) vertical grid:

In [6]:

Copied!

# Define our target vertical grid coordinate:
e3t_target = xr.DataArray(np.repeat(200.0, 30), dims=['k_new'])

e3t_target
# Define our target vertical grid coordinate:
e3t_target = xr.DataArray(np.repeat(200.0, 30), dims=['k_new'])

e3t_target

Out[6]:

<xarray.DataArray (k_new: 30)> Size: 240B
array([200., 200., 200., 200., 200., 200., 200., 200., 200., 200., 200.,
       200., 200., 200., 200., 200., 200., 200., 200., 200., 200., 200.,
       200., 200., 200., 200., 200., 200., 200., 200.])
Dimensions without coordinates: k_new

In [7]:

Copied!





# Transform eORCA1 3-dimensional conservative temperature field to new vertical coordinate system:
ds_k_transform = nemo.transform_vertical_grid(grid = '/gridT',
                                              var = 'thetao_con',
                                              e3_new = e3t_target
                                              )

ds_k_transform
# Transform eORCA1 3-dimensional conservative temperature field to new vertical coordinate system:
ds_k_transform = nemo.transform_vertical_grid(grid = '/gridT',
                                              var = 'thetao_con',
                                              e3_new = e3t_target
                                              )

ds_k_transform

Out[7]:

<xarray.Dataset> Size: 631MB
Dimensions:        (time_counter: 11, j: 331, i: 360, k_new: 30)
Coordinates:
  * time_counter   (time_counter) datetime64[ns] 88B 2000-07-02 ... 2010-07-0...
    time_centered  (time_counter) datetime64[ns] 88B dask.array<chunksize=(1,), meta=np.ndarray>
    gphit          (j, i) float64 953kB dask.array<chunksize=(331, 360), meta=np.ndarray>
    glamt          (j, i) float64 953kB dask.array<chunksize=(331, 360), meta=np.ndarray>
  * j              (j) int64 3kB 1 2 3 4 5 6 7 8 ... 325 326 327 328 329 330 331
  * i              (i) int64 3kB 1 2 3 4 5 6 7 8 ... 354 355 356 357 358 359 360
    deptht_new     (k_new) float64 240B 200.0 400.0 600.0 ... 5.8e+03 6e+03
Dimensions without coordinates: k_new
Data variables:
    thetao_con     (time_counter, k_new, j, i) float64 315MB dask.array<chunksize=(1, 30, 331, 360), meta=np.ndarray>
    e3t_new        (time_counter, j, i, k_new) float64 315MB dask.array<chunksize=(1, 331, 360, 30), meta=np.ndarray>

Notice that the output above returns a Dataset containing two DataArrays: the vertically remapped conservative temperature thetao_con and the vertical grid cell thicknesses e3t_out (accounting for partial grid cell at the sea floor).

Since both of these DataArrays contain dask arrays, we haven't actually computed anything yet. To do this, we need to call the .compute() method:

In [8]:

Copied!

ds_k_transform = ds_k_transform.compute()
ds_k_transform = ds_k_transform.compute()

Visualising the vertically transformed conservative temperature field¶

Let's plot the original and vertically transformed conservative temperature profiles together.

In [9]:

Copied!

# Plot our original time-mean vertical conservative temperature profile in the subpolar North Atlantic:
nemo['/gridT']['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht', ylim=(3600, 0), marker='o', color='dodgerblue')

# Plot our vertically transformed time-mean vertical conservative temperature profile in the subpolar North Atlantic:
ds_k_transform['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht_new', ylim=(3600, 0), marker='o', color='coral')
# Plot our original time-mean vertical conservative temperature profile in the subpolar North Atlantic:
nemo['/gridT']['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht', ylim=(3600, 0), marker='o', color='dodgerblue')

# Plot our vertically transformed time-mean vertical conservative temperature profile in the subpolar North Atlantic:
ds_k_transform['thetao_con'].isel(i=260, j=260, time_counter=0).plot(y='deptht_new', ylim=(3600, 0), marker='o', color='coral')

Out[9]:

[<matplotlib.lines.Line2D at 0x14ffd4f50>]

We can also verify that the product of the vertical grid cell thickness (m) and conservative temperature (C) is conserved following the transformation.

Note that we need to transform all variables to a consistent dtype (in this case float64) before making the comparison.

In [10]:

Copied!

# Calculate the sum of the product of the transformed conservative temperature and cell thickness at our example location:
print((ds_k_transform['thetao_con'].isel(time_counter=0, i=260, j=260).astype(np.float64) * ds_k_transform['e3t_new'].isel(time_counter=0, i=260, j=260).astype(np.float64)).sum().compute())

# Calculate the sum of the product of the original conservative temperature and cell thickness at our example location:
print((nemo['/gridT']['thetao_con'].isel(time_counter=0, i=260, j=260).astype(np.float64) * nemo['/gridT']['e3t'].isel(time_counter=0, i=260, j=260).astype(np.float64)).sum().compute())
# Calculate the sum of the product of the transformed conservative temperature and cell thickness at our example location:
print((ds_k_transform['thetao_con'].isel(time_counter=0, i=260, j=260).astype(np.float64) * ds_k_transform['e3t_new'].isel(time_counter=0, i=260, j=260).astype(np.float64)).sum().compute())

# Calculate the sum of the product of the original conservative temperature and cell thickness at our example location:
print((nemo['/gridT']['thetao_con'].isel(time_counter=0, i=260, j=260).astype(np.float64) * nemo['/gridT']['e3t'].isel(time_counter=0, i=260, j=260).astype(np.float64)).sum().compute())

<xarray.DataArray ()> Size: 8B
array(14272.51116587)
Coordinates:
    time_counter   datetime64[ns] 8B 2000-07-02
    time_centered  datetime64[ns] 8B 2000-07-02
    gphit          float64 8B 51.34
    glamt          float64 8B -26.44
    j              int64 8B 261
    i              int64 8B 261
<xarray.DataArray ()> Size: 8B
array(14272.51116587)
Coordinates:
    time_counter   datetime64[ns] 8B 2000-07-02
    time_centered  datetime64[ns] 8B 2000-07-02
    gphit          float64 8B 51.34
    glamt          float64 8B -26.44
    j              int64 8B 261
    i              int64 8B 261