{ "cells": [ { "cell_type": "markdown", "id": "1728a25c-6bba-46c2-a9d4-c2f1e7696914", "metadata": {}, "source": [ "# Analyzing and visualizing CAM-SE output in Python\n", "\n", "## Overview\n", "\n", "We demonstrate a variety of options for analyzing and visualizing output from the Community Atmosphere Model (CAM) with the spectral element (SE) grid in Python. This notebook was developed for the ESDS [Collaborative Work Time on Unstructured Grids](https://ncar.github.io/esds/posts/2023/unstructured-grid-collab-1/), which took place on April 17, 2023. A recap of the related CAM-SE discussion can be found [here](https://ncar.github.io/esds/posts/2023/unstructured-grid-collab-1/#katie-dagon-cgd).\n", "\n", "## Contents\n", "1. [Regrid CAM-SE output using map file](#section1)\n", "2. [Use a CAM-SE remap function to scale up](#section2)\n", "3. [Regrid CAM-SE output using xESMF](#section3)\n", "4. [Direct sparse matrix multiply-add](#section4)" ] }, { "cell_type": "markdown", "id": "d2e28c5e-becf-4c5d-a7f7-e578a6212b96", "metadata": {}, "source": [ "## Imports" ] }, { "cell_type": "code", "execution_count": 1, "id": "5449ce5a-74ec-4312-92e7-3ec4653a0cbf", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "sys : 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:08:06) [GCC 11.3.0]\n", "xesmf : 0.7.1\n", "distributed : 2023.4.1\n", "json : 2.0.9\n", "ncar_jobqueue: 2021.4.14\n", "xarray : 2023.3.0\n", "numpy : 1.23.5\n", "scipy : 1.10.1\n", "opt_einsum : 3.3.0\n", "\n" ] } ], "source": [ "%load_ext watermark\n", "\n", "import distributed\n", "import ncar_jobqueue\n", "import numpy as np\n", "import opt_einsum\n", "import scipy as sp\n", "import xarray as xr\n", "import xesmf\n", "\n", "%watermark -iv" ] }, { "cell_type": "markdown", "id": "724159c5-1ef1-4828-b820-f71dd1a626ac", "metadata": {}, "source": [ "## Spin up a Cluster" ] }, { "cell_type": "code", "execution_count": 2, "id": "63bdae88-aaf2-46e9-ba1d-b7c753e08c26", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
3d08bca5
\n", "\n", " Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/kdagon/proxy/8787/status\n", " | \n", "\n", " Workers: 0\n", " | \n", "
\n", " Total threads: 0\n", " | \n", "\n", " Total memory: 0 B\n", " | \n", "
Scheduler-9a2e31b9-bbf4-4978-8e18-8c1af8bb535a
\n", "\n", " Comm: tcp://10.12.206.45:39748\n", " | \n", "\n", " Workers: 0\n", " | \n", "
\n", " Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/kdagon/proxy/8787/status\n", " | \n", "\n", " Total threads: 0\n", " | \n", "
\n", " Started: Just now\n", " | \n", "\n", " Total memory: 0 B\n", " | \n", "
Client-b37f404b-50c0-11ee-8170-3cecef19f78e
\n", "Connection method: Cluster object | \n", "Cluster type: dask_jobqueue.PBSCluster | \n", " \n", "
\n", " Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/kdagon/proxy/8787/status\n", " | \n", "\n", " |
3d08bca5
\n", "\n", " Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/kdagon/proxy/8787/status\n", " | \n", "\n", " Workers: 0\n", " | \n", "
\n", " Total threads: 0\n", " | \n", "\n", " Total memory: 0 B\n", " | \n", "
Scheduler-9a2e31b9-bbf4-4978-8e18-8c1af8bb535a
\n", "\n", " Comm: tcp://10.12.206.45:39748\n", " | \n", "\n", " Workers: 0\n", " | \n", "
\n", " Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/kdagon/proxy/8787/status\n", " | \n", "\n", " Total threads: 0\n", " | \n", "
\n", " Started: Just now\n", " | \n", "\n", " Total memory: 0 B\n", " | \n", "
<xarray.Dataset>\n", "Dimensions: (src_grid_rank: 1, dst_grid_rank: 2, n_a: 777602,\n", " n_b: 884736, nv_a: 3, nv_b: 4, n_s: 2654208)\n", "Dimensions without coordinates: src_grid_rank, dst_grid_rank, n_a, n_b, nv_a,\n", " nv_b, n_s\n", "Data variables: (12/19)\n", " src_grid_dims (src_grid_rank) int32 ...\n", " dst_grid_dims (dst_grid_rank) int32 ...\n", " yc_a (n_a) float64 ...\n", " yc_b (n_b) float64 ...\n", " xc_a (n_a) float64 ...\n", " xc_b (n_b) float64 ...\n", " ... ...\n", " area_b (n_b) float64 ...\n", " frac_a (n_a) float64 ...\n", " frac_b (n_b) float64 ...\n", " col (n_s) int32 ...\n", " row (n_s) int32 ...\n", " S (n_s) float64 ...\n", "Attributes:\n", " title: ESMF Offline Regridding Weight Generator\n", " normalization: destarea\n", " map_method: Bilinear remapping\n", " ESMF_regrid_method: Bilinear\n", " conventions: NCAR-CSM\n", " domain_a: /glade/p/cgd/ccr/people/nanr/mapfiles/ne120.nc\n", " domain_b: /glade/p/cgd/ccr/people/nanr/mapfiles/0.23x0.31.nc\n", " grid_file_src: /glade/p/cgd/ccr/people/nanr/mapfiles/ne120.nc\n", " grid_file_dst: /glade/p/cgd/ccr/people/nanr/mapfiles/0.23x0.31.nc\n", " CVS_revision: 6.3.0r
<xarray.Dataset>\n", "Dimensions: (lev: 30, ilev: 31, ncol: 777602, time: 2920, nbnd: 2)\n", "Coordinates:\n", " * lev (lev) float64 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6\n", " * ilev (ilev) float64 2.255 5.032 10.16 18.56 ... 967.5 985.1 1e+03\n", " * time (time) object 2000-01-01 03:00:00 ... 2001-01-01 00:00:00\n", "Dimensions without coordinates: ncol, nbnd\n", "Data variables: (12/31)\n", " hyam (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " hybm (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " P0 float64 ...\n", " hyai (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " hybi (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " lat (ncol) float64 dask.array<chunksize=(50000,), meta=np.ndarray>\n", " ... ...\n", " n2ovmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " f11vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " f12vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " sol_tsi (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " nsteph (time) int32 dask.array<chunksize=(50,), meta=np.ndarray>\n", " TS (time, ncol) float32 dask.array<chunksize=(50, 50000), meta=np.ndarray>\n", "Attributes:\n", " np: 4\n", " ne: 120\n", " Conventions: CF-1.0\n", " source: CAM\n", " case: b.e13.B20TRC5CN.ne120_g16.003\n", " title: UNSET\n", " logname: \n", " host: \n", " Version: $Name$\n", " revision_Id: $Id$\n", " initial_file: /projects/ccsm/inputdata/atm/cam/inic/homme/cami-mam3_0...\n", " topography_file: /projects/ccsm/inputdata/atm/cam/topo/USGS-gtopo30_ne12...
<xarray.DataArray 'TS' (time: 2920, ncol: 777602)>\n", "dask.array<open_dataset-bb5b27d0821608d3fc5037525bd2f985TS, shape=(2920, 777602), dtype=float32, chunksize=(50, 50000), chunktype=numpy.ndarray>\n", "Coordinates:\n", " * time (time) object 2000-01-01 03:00:00 ... 2001-01-01 00:00:00\n", "Dimensions without coordinates: ncol\n", "Attributes:\n", " units: K\n", " long_name: Surface temperature (radiative)
<xarray.DataArray (lon: 1152, lat: 768)>\n", "array([[244.04267883, 243.94797035, 243.71354076, ..., 231.61040262,\n", " 231.07736837, 231.86198425],\n", " [244.04267883, 243.94817018, 243.71394044, ..., 231.61271724,\n", " 231.07852571, 231.86198425],\n", " [244.04267883, 243.94837698, 243.71435405, ..., 231.61500008,\n", " 231.07966716, 231.86198425],\n", " ...,\n", " [244.04267883, 243.95208424, 243.72176851, ..., 231.57748313,\n", " 231.06090852, 231.86198425],\n", " [244.04267883, 243.95070604, 243.71901211, ..., 231.58848746,\n", " 231.06641072, 231.86198425],\n", " [244.04267883, 243.94933473, 243.7162695 , ..., 231.59946073,\n", " 231.07189739, 231.86198425]])\n", "Coordinates:\n", " * lon (lon) float64 0.0 0.3125 0.625 0.9375 ... 358.8 359.1 359.4 359.7\n", " * lat (lat) float64 -90.0 -89.77 -89.53 -89.3 ... 89.3 89.53 89.77 90.0
<xarray.Dataset>\n", "Dimensions: (lev: 30, ilev: 31, ncol: 777602, time: 2920, nbnd: 2)\n", "Coordinates:\n", " * lev (lev) float64 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6\n", " * ilev (ilev) float64 2.255 5.032 10.16 18.56 ... 967.5 985.1 1e+03\n", " * time (time) object 2000-01-01 03:00:00 ... 2001-01-01 00:00:00\n", "Dimensions without coordinates: ncol, nbnd\n", "Data variables: (12/31)\n", " hyam (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " hybm (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " P0 float64 ...\n", " hyai (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " hybi (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " lat (ncol) float64 dask.array<chunksize=(777602,), meta=np.ndarray>\n", " ... ...\n", " n2ovmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " f11vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " f12vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " sol_tsi (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " nsteph (time) int32 dask.array<chunksize=(50,), meta=np.ndarray>\n", " TS (time, ncol) float32 dask.array<chunksize=(50, 777602), meta=np.ndarray>\n", "Attributes:\n", " np: 4\n", " ne: 120\n", " Conventions: CF-1.0\n", " source: CAM\n", " case: b.e13.B20TRC5CN.ne120_g16.003\n", " title: UNSET\n", " logname: \n", " host: \n", " Version: $Name$\n", " revision_Id: $Id$\n", " initial_file: /projects/ccsm/inputdata/atm/cam/inic/homme/cami-mam3_0...\n", " topography_file: /projects/ccsm/inputdata/atm/cam/topo/USGS-gtopo30_ne12...
<xarray.Dataset>\n", "Dimensions: (src_grid_rank: 1, dst_grid_rank: 2, n_a: 777602,\n", " n_b: 884736, nv_a: 3, nv_b: 4, n_s: 2654208)\n", "Dimensions without coordinates: src_grid_rank, dst_grid_rank, n_a, n_b, nv_a,\n", " nv_b, n_s\n", "Data variables: (12/19)\n", " src_grid_dims (src_grid_rank) int32 ...\n", " dst_grid_dims (dst_grid_rank) int32 ...\n", " yc_a (n_a) float64 ...\n", " yc_b (n_b) float64 ...\n", " xc_a (n_a) float64 ...\n", " xc_b (n_b) float64 ...\n", " ... ...\n", " area_b (n_b) float64 ...\n", " frac_a (n_a) float64 ...\n", " frac_b (n_b) float64 ...\n", " col (n_s) int32 ...\n", " row (n_s) int32 ...\n", " S (n_s) float64 ...\n", "Attributes:\n", " title: ESMF Offline Regridding Weight Generator\n", " normalization: destarea\n", " map_method: Bilinear remapping\n", " ESMF_regrid_method: Bilinear\n", " conventions: NCAR-CSM\n", " domain_a: /glade/p/cgd/ccr/people/nanr/mapfiles/ne120.nc\n", " domain_b: /glade/p/cgd/ccr/people/nanr/mapfiles/0.23x0.31.nc\n", " grid_file_src: /glade/p/cgd/ccr/people/nanr/mapfiles/ne120.nc\n", " grid_file_dst: /glade/p/cgd/ccr/people/nanr/mapfiles/0.23x0.31.nc\n", " CVS_revision: 6.3.0r
<xarray.Dataset>\n", "Dimensions: (lev: 30, ilev: 31, nbnd: 2, lat: 768, lon: 1152)\n", "Coordinates:\n", " * lev (lev) float64 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6\n", " * ilev (ilev) float64 2.255 5.032 10.16 18.56 ... 967.5 985.1 1e+03\n", " time object 2000-01-01 03:00:00\n", " * lat (lat) float64 -90.0 -89.77 -89.53 -89.3 ... 89.53 89.77 90.0\n", " * lon (lon) float64 0.0 0.3125 0.625 0.9375 ... 359.1 359.4 359.7\n", "Dimensions without coordinates: nbnd\n", "Data variables: (12/28)\n", " hyam (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " hybm (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " P0 float64 ...\n", " hyai (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " hybi (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " ntrm int32 ...\n", " ... ...\n", " n2ovmr float64 dask.array<chunksize=(), meta=np.ndarray>\n", " f11vmr float64 dask.array<chunksize=(), meta=np.ndarray>\n", " f12vmr float64 dask.array<chunksize=(), meta=np.ndarray>\n", " sol_tsi float64 dask.array<chunksize=(), meta=np.ndarray>\n", " nsteph int32 dask.array<chunksize=(), meta=np.ndarray>\n", " TS (lat, lon) float64 244.0 244.0 244.0 ... 231.9 231.9 231.9\n", "Attributes:\n", " np: 4\n", " ne: 120\n", " Conventions: CF-1.0\n", " source: CAM\n", " case: b.e13.B20TRC5CN.ne120_g16.003\n", " title: UNSET\n", " logname: \n", " host: \n", " Version: $Name$\n", " revision_Id: $Id$\n", " initial_file: /projects/ccsm/inputdata/atm/cam/inic/homme/cami-mam3_0...\n", " topography_file: /projects/ccsm/inputdata/atm/cam/topo/USGS-gtopo30_ne12...
<xarray.Dataset>\n", "Dimensions: (lev: 30, ilev: 31, time: 2920, nbnd: 2, lat: 768, lon: 1152)\n", "Coordinates:\n", " * lev (lev) float64 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6\n", " * ilev (ilev) float64 2.255 5.032 10.16 18.56 ... 967.5 985.1 1e+03\n", " * time (time) object 2000-01-01 03:00:00 ... 2001-01-01 00:00:00\n", " * lat (lat) float64 -90.0 -89.77 -89.53 -89.3 ... 89.53 89.77 90.0\n", " * lon (lon) float64 0.0 0.3125 0.625 0.9375 ... 359.1 359.4 359.7\n", "Dimensions without coordinates: nbnd\n", "Data variables: (12/29)\n", " hyam (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " hybm (lev) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n", " P0 float64 ...\n", " hyai (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " hybi (ilev) float64 dask.array<chunksize=(31,), meta=np.ndarray>\n", " ntrm int32 ...\n", " ... ...\n", " f11vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " f12vmr (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " sol_tsi (time) float64 dask.array<chunksize=(50,), meta=np.ndarray>\n", " nsteph (time) int32 dask.array<chunksize=(50,), meta=np.ndarray>\n", " area (lat, lon) float64 dask.array<chunksize=(768, 1152), meta=np.ndarray>\n", " TS (time, lat, lon) float32 dask.array<chunksize=(50, 768, 1152), meta=np.ndarray>\n", "Attributes:\n", " np: 4\n", " ne: 120\n", " Conventions: CF-1.0\n", " source: CAM\n", " case: b.e13.B20TRC5CN.ne120_g16.003\n", " title: UNSET\n", " logname: \n", " host: \n", " Version: $Name$\n", " revision_Id: $Id$\n", " initial_file: /projects/ccsm/inputdata/atm/cam/inic/homme/cami-mam3_0...\n", " topography_file: /projects/ccsm/inputdata/atm/cam/topo/USGS-gtopo30_ne12...