AreaDefinition to_cf Prototype.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "279bcc41",
   "metadata": {},
   "source": [
    "## Prototype Pyresample AreaDefinition to_cf()\n",
    "\n",
    "This notebook is a prototype for a routine that is the counterpart of [load_cf_area()](https://pyresample.readthedocs.io/en/latest/howtos/geometry_utils.html#loading-from-netcdf-cf). It takes a pyresample AreaDefinition object as input and returns an xarray Dataset that holds the basic information needed to prepare netCDF/CF files for this particular grid. It returns a template that can be further worked upon (add relevant variables, more attributes, more dimensions, etc...).\n",
    "\n",
    "Status:\n",
    "* Sept 7th 2023 : Prototype ready for inspection. It does not support all possible CF projections. It also does not support user configuration (e.g. via a cf_info).\n",
    "* Sept 8th 2023 : Move encodings to the DataArrays (doesn't seem to work for coordinate variables lat/lon)\n",
    "* Sept 28th 2023 : Add a small chunk of code to decide the order of x and y dimensions from the shape of lat/lon\n",
    "\n",
    "TODO :\n",
    "* Investigate if the newer pyproj's [crs.cs_to_cf()](https://pyproj4.github.io/pyproj/latest/build_crs_cf.html) could be used to shorten this even more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b197e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pyresample as pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8de03e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/thomasl/mambaforge/envs/py310/lib/python3.10/site-packages/pyproj/crs/crs.py:1296: UserWarning: You will likely lose important projection information when converting to a PROJ string from another format. See: https://proj.org/faq.html#what-is-the-best-format-for-describing-coordinate-reference-systems\n",
      "  proj = self._crs.to_proj4(version=version)\n",
      "/home/thomasl/mambaforge/envs/py310/lib/python3.10/site-packages/pyproj/crs/crs.py:1296: UserWarning: You will likely lose important projection information when converting to a PROJ string from another format. See: https://proj.org/faq.html#what-is-the-best-format-for-describing-coordinate-reference-systems\n",
      "  proj = self._crs.to_proj4(version=version)\n"
     ]
    }
   ],
   "source": [
    "# load an area_def from a netCDF/CF file (could be any other solution to obtain an AreaDefinition)\n",
    "ds = xr.open_dataset('https://thredds.met.no/thredds/dodsC/osisaf/met.no/reprocessed/ice/conc_450a_files/monthly/2020/ice_conc_nh_ease2-250_cdr-v3p0_202011.nc')\n",
    "area_def, cf_info = pr.utils.load_cf_area(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ded997ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_cf_template(area_def, skip_lonlat=True, crs_encoding='pyproj_to_cf'):\n",
    "    \"\"\"Return a template xarray Dataset holding the structure of a netCDF/CF file for this grid.\"\"\"\n",
    "    \n",
    "    # prepare the crs object with pyproj.to_cf()\n",
    "    crs_cf = area_def.crs.to_cf()\n",
    "    type_of_grid_mapping = crs_cf['grid_mapping_name']\n",
    "    \n",
    "    # alter this set of attributes on user input\n",
    "    if crs_encoding == 'pyproj_to_cf':\n",
    "        pass\n",
    "    elif crs_encoding == 'only_cf_attrs':\n",
    "        crs_cf.pop('crs_wkt')\n",
    "    elif crs_encoding == 'only_wkt':\n",
    "        crs_cf = {'crs_wkt': crs_cf['crs_wkt']}\n",
    "    else:\n",
    "        raise ValueError(\"Unsupported value for crs_encoding=\")\n",
    "        \n",
    "    # prepare the x and y axis (1D)\n",
    "    xy = dict()\n",
    "    xy_dims = ('x', 'y')\n",
    "    for axis in xy_dims:\n",
    "        \n",
    "        # access the valid standard_names (also handle the 'default')\n",
    "        try:\n",
    "            valid_coord_standard_names = pr.utils.cf._valid_cf_coordinate_standardnames[type_of_grid_mapping][axis]\n",
    "        except KeyError:\n",
    "            valid_coord_standard_names = pr.utils.cf._valid_cf_coordinate_standardnames['default'][axis]\n",
    "        \n",
    "        xy[axis] = dict()\n",
    "        # CF wants the center of the grid cells\n",
    "        if axis == 'x':\n",
    "            xy[axis]['_coords'] = area_def.projection_x_coords\n",
    "        else:\n",
    "            xy[axis]['_coords'] = area_def.projection_y_coords     \n",
    "        # each axis requires a valid name, which depends on the type of projection\n",
    "        xy[axis]['standard_name'] = valid_coord_standard_names[0]\n",
    "        # CF recommendation to have axis= attribute\n",
    "        xy[axis]['axis'] = axis.upper()\n",
    "        # units\n",
    "        xy[axis]['units'] = 'm'\n",
    "    \n",
    "    # latitude and longitude (2D)\n",
    "    lons, lats = area_def.get_lonlats()\n",
    "    \n",
    "    # determine the order of the x, y dimensions that match the shape of lat/lon arrays\n",
    "    \n",
    "    if lons.shape == (len(xy['y']['_coords']),len(xy['x']['_coords'])):\n",
    "        xy_dims = ('y', 'x')\n",
    "    elif lons.shape == (len(xy['x']['_coords']),len(xy['y']['_coords'])):\n",
    "        xy_dims = ('x', 'y')\n",
    "    else:\n",
    "        raise ValueError(\"Incompatible shape for lon/lat {}, x {}, and y {}.\".format(lons.shape,\n",
    "                                                                                    len(xy['x']['_coords']),\n",
    "                                                                                    len(xy['y']['_coords'])))\n",
    "    \n",
    "    # define a Dataset as a template.\n",
    "    #   we cannot define a Dataset without a data variable. The strategy is to \n",
    "    #   create the Dataset with an (empty) variable, and the user can later \n",
    "    #   add his own variables, and finally drop our 'template' variable before\n",
    "    #   writing to file.\n",
    "    varn = 'template'\n",
    "    shape = lons.shape\n",
    "    da_empty_data = np.ones_like(lons) * np.nan\n",
    "    da_dims = list(xy_dims)\n",
    "    da_coords = {'x':('x',xy['x']['_coords']), 'y':('y',xy['y']['_coords']),}\n",
    "    if not skip_lonlat:\n",
    "        da_coords['lon']=(xy_dims, lons)\n",
    "        da_coords['lat']=(xy_dims, lats)\n",
    "    ds = xr.Dataset(data_vars={varn:(da_dims, da_empty_data),\n",
    "                               'crs':([], 0)},\n",
    "                    coords=da_coords,)\n",
    "    \n",
    "    # add CF attributes and encodings to the xarray template\n",
    "    \n",
    "    # x and y dims\n",
    "    for axis in xy_dims:\n",
    "        for attr in xy[axis].keys():\n",
    "            if attr.startswith('_'): continue\n",
    "            ds[axis].attrs[attr] = xy[axis][attr]\n",
    "        ds[axis].encoding = {'_FillValue':None}\n",
    "    \n",
    "    # crs object\n",
    "    ds['crs'].attrs = crs_cf\n",
    "    ds['crs'].encoding = {'dtype':'int32'}\n",
    "    \n",
    "    # latitude and longitude\n",
    "    if not skip_lonlat:\n",
    "        ds['lon'].attrs={'long_name': 'longitude coordinate', 'standard_name': 'longitude', 'units': 'degrees_east'}\n",
    "        ds['lat'].attrs={'long_name': 'latitude coordinate', 'standard_name': 'latitude', 'units': 'degrees_north'}\n",
    "        ds['lon'].encoding = {'_FillValue':None}\n",
    "        ds['lat'].encoding = {'_FillValue':None}\n",
    "        \n",
    "    # the empty variable itself\n",
    "    ds[varn].attrs['grid_mapping'] = 'crs'\n",
    "    \n",
    "    # global attributes\n",
    "    ds.attrs['Conventions'] = 'CF-1.8'\n",
    "    \n",
    "    return ds\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4526c183",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate the routine.\n",
    "\n",
    "# Use to_cf_template() to obtain a template for the netCDF/CF file\n",
    "ds = to_cf_template(area_def, skip_lonlat=False, )\n",
    "\n",
    "# Create a DataArray holding user data that we want to add to the file. \n",
    "myData = np.arange(0,ds['lon'].size).reshape(ds['lon'].shape).astype('float')\n",
    "da_var2 = xr.DataArray(myData, coords=ds['template'].coords, dims=ds['template'].dims,\n",
    "                       attrs=ds['template'].attrs, name='my_var')\n",
    "# Add some attributes\n",
    "da_var2.attrs['long_name'] = 'my variable defined on that particular grid'\n",
    "da_var2.attrs['units'] = 'K'\n",
    "\n",
    "# Add the variable to the template file\n",
    "ds = ds.merge(da_var2)\n",
    "\n",
    "# We don't need the template variable anymore, drop it.\n",
    "ds = ds.drop('template')\n",
    "\n",
    "# Write to file.\n",
    "o_nc = 'cf_template.nc'\n",
    "ds.to_netcdf('cf_template.nc', format='NETCDF4_CLASSIC')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c49d8ab9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "netcdf cf_template {\r\n",
      "dimensions:\r\n",
      "\tx = 432 ;\r\n",
      "\ty = 432 ;\r\n",
      "variables:\r\n",
      "\tint crs ;\r\n",
      "\t\tcrs:crs_wkt = \"PROJCRS[\\\"undefined\\\",BASEGEOGCRS[\\\"undefined\\\",DATUM[\\\"undefined\\\",ELLIPSOID[\\\"undefined\\\",6378137,298.257223563,LENGTHUNIT[\\\"metre\\\",1,ID[\\\"EPSG\\\",9001]]]],PRIMEM[\\\"Greenwich\\\",0,ANGLEUNIT[\\\"degree\\\",0.0174532925199433],ID[\\\"EPSG\\\",8901]]],CONVERSION[\\\"unknown\\\",METHOD[\\\"Lambert Azimuthal Equal Area\\\",ID[\\\"EPSG\\\",9820]],PARAMETER[\\\"Latitude of natural origin\\\",90,ANGLEUNIT[\\\"degree\\\",0.0174532925199433],ID[\\\"EPSG\\\",8801]],PARAMETER[\\\"Longitude of natural origin\\\",0,ANGLEUNIT[\\\"degree\\\",0.0174532925199433],ID[\\\"EPSG\\\",8802]],PARAMETER[\\\"False easting\\\",0,LENGTHUNIT[\\\"metre\\\",1],ID[\\\"EPSG\\\",8806]],PARAMETER[\\\"False northing\\\",0,LENGTHUNIT[\\\"metre\\\",1],ID[\\\"EPSG\\\",8807]]],CS[Cartesian,2],AXIS[\\\"(E)\\\",east,ORDER[1],LENGTHUNIT[\\\"metre\\\",1,ID[\\\"EPSG\\\",9001]]],AXIS[\\\"(N)\\\",north,ORDER[2],LENGTHUNIT[\\\"metre\\\",1,ID[\\\"EPSG\\\",9001]]]]\" ;\r\n",
      "\t\tcrs:semi_major_axis = 6378137. ;\r\n",
      "\t\tcrs:semi_minor_axis = 6356752.31424518 ;\r\n",
      "\t\tcrs:inverse_flattening = 298.257223563 ;\r\n",
      "\t\tcrs:reference_ellipsoid_name = \"undefined\" ;\r\n",
      "\t\tcrs:longitude_of_prime_meridian = 0. ;\r\n",
      "\t\tcrs:prime_meridian_name = \"Greenwich\" ;\r\n",
      "\t\tcrs:geographic_crs_name = \"undefined\" ;\r\n",
      "\t\tcrs:horizontal_datum_name = \"undefined\" ;\r\n",
      "\t\tcrs:projected_crs_name = \"undefined\" ;\r\n",
      "\t\tcrs:grid_mapping_name = \"lambert_azimuthal_equal_area\" ;\r\n",
      "\t\tcrs:latitude_of_projection_origin = 90. ;\r\n",
      "\t\tcrs:longitude_of_projection_origin = 0. ;\r\n",
      "\t\tcrs:false_easting = 0. ;\r\n",
      "\t\tcrs:false_northing = 0. ;\r\n",
      "\tdouble x(x) ;\r\n",
      "\t\tx:standard_name = \"projection_x_coordinate\" ;\r\n",
      "\t\tx:axis = \"X\" ;\r\n",
      "\t\tx:units = \"m\" ;\r\n",
      "\tdouble y(y) ;\r\n",
      "\t\ty:standard_name = \"projection_y_coordinate\" ;\r\n",
      "\t\ty:axis = \"Y\" ;\r\n",
      "\t\ty:units = \"m\" ;\r\n",
      "\tdouble lon(y, x) ;\r\n",
      "\t\tlon:_FillValue = NaN ;\r\n",
      "\t\tlon:long_name = \"longitude coordinate\" ;\r\n",
      "\t\tlon:standard_name = \"longitude\" ;\r\n",
      "\t\tlon:units = \"degrees_east\" ;\r\n",
      "\tdouble lat(y, x) ;\r\n",
      "\t\tlat:_FillValue = NaN ;\r\n",
      "\t\tlat:long_name = \"latitude coordinate\" ;\r\n",
      "\t\tlat:standard_name = \"latitude\" ;\r\n",
      "\t\tlat:units = \"degrees_north\" ;\r\n",
      "\tdouble my_var(y, x) ;\r\n",
      "\t\tmy_var:_FillValue = NaN ;\r\n",
      "\t\tmy_var:grid_mapping = \"crs\" ;\r\n",
      "\t\tmy_var:long_name = \"my variable defined on that particular grid\" ;\r\n",
      "\t\tmy_var:units = \"K\" ;\r\n",
      "\t\tmy_var:coordinates = \"lat lon\" ;\r\n",
      "\r\n",
      "// global attributes:\r\n",
      "\t\t:Conventions = \"CF-1.8\" ;\r\n",
      "}\r\n"
     ]
    }
   ],
   "source": [
    "!ncdump -ch cf_template.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "964d6690",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EQUALITY\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/thomasl/mambaforge/envs/py310/lib/python3.10/site-packages/pyproj/crs/crs.py:1296: UserWarning: You will likely lose important projection information when converting to a PROJ string from another format. See: https://proj.org/faq.html#what-is-the-best-format-for-describing-coordinate-reference-systems\n",
      "  proj = self._crs.to_proj4(version=version)\n"
     ]
    }
   ],
   "source": [
    "# Simple test: can we load the AreaDefinition from the file?\n",
    "new_area_def, new_cf_info = pr.utils.load_cf_area(o_nc)\n",
    "# If yes, is it the same AreaDefinition as we wrote in?\n",
    "if (area_def != new_area_def):\n",
    "    print(\"EQUALITY AREA_EXTENT: \", np.allclose(area_def.area_extent, new_area_def.area_extent))\n",
    "    print(\"EQUALITY SHAPE: \", area_def.shape == new_area_def.shape)\n",
    "    print(\"EQUALITY CRS: \", area_def.crs == new_area_def.crs)\n",
    "    lons, lats = area_def.get_lonlats()\n",
    "    nlons, nlats = new_area_def.get_lonlats()\n",
    "    print(\"EQUALITY LATs: \", abs(lats-nlats).max() < 1.e-4)\n",
    "    print(\"EQUALITY LONs: \", abs(lons-nlons).max() < 1.e-4)\n",
    "else:\n",
    "    print(\"EQUALITY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec3e7c3e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}