From 11a35ee0480b8462d50e9ca7a8b687e29fb33915 Mon Sep 17 00:00:00 2001 From: Raphael Dussin Date: Wed, 11 May 2022 17:26:03 -0400 Subject: [PATCH] create obs data and use supergrid lon/lat --- data/windstress_OM4.py | 117 +++++++++++++-- om4labs/catalogs/OMIP2_catalog_gfdl.yml | 31 +++- om4labs/diags/stress_curl/stress_curl.py | 173 +++++++++++++++++------ 3 files changed, 267 insertions(+), 54 deletions(-) diff --git a/data/windstress_OM4.py b/data/windstress_OM4.py index 4b31e2a..40f1b6f 100644 --- a/data/windstress_OM4.py +++ b/data/windstress_OM4.py @@ -1,19 +1,118 @@ import xarray as xr from glob import glob +from xgcm import Grid +import xesmf +import subprocess as sp -ppdir = "/archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ts/annual" -files = glob(f"{ppdir}/5yr/*tauuo.nc") -files += glob(f"{ppdir}/5yr/*tauvo.nc") -files += glob(f"{ppdir}/1yr/*tauuo.nc") -files += glob(f"{ppdir}/1yr/*tauvo.nc") +# outputdir = "/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4" +outputdir = "." + + +ppdir_OM4p25 = "/archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/" + +ppdir_OM4p125 = "/archive/Raphael.Dussin/FMS2019.01.03_devgfdl_20210706/CM4_piControl_c192_OM4p125_v6_alt1/gfdl.ncrc4-intel18-prod-openmp/pp/ocean_annual" + +# Read the files from the 1/4 degree +files = glob(f"{ppdir_OM4p25}/ts/annual/5yr/*tauuo.nc") +files += glob(f"{ppdir_OM4p25}/ts/annual/5yr/*tauvo.nc") +files += glob(f"{ppdir_OM4p25}/ts/annual/1yr/*tauuo.nc") +files += glob(f"{ppdir_OM4p25}/ts/annual/1yr/*tauvo.nc") +files += glob(f"{ppdir_OM4p25}/ocean_annual.static.nc") + +sp.check_call( + "dmget /archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ts/annual/5yr/*tauuo.nc /archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ts/annual/5yr/*tauvo.nc /archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ts/annual/1yr/*tauuo.nc /archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ts/annual/1yr/*tauvo.nc /archive/Raphael.Dussin/xanadu_esm4_20190304_mom6_2019.08.08/OM4p25_JRA55do1.4_0netfw_cycle6/gfdl.ncrc4-intel16-prod/pp/ocean_annual/ocean_annual.static.nc", + shell=True, +) ds = xr.open_mfdataset(files, chunks={"time": 5}) -ds.sel(time=slice("1999", "2018")).mean(dim="time").to_netcdf( - "/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstress_OM4p25_JRA1.4_cycle6_1999-2018.nc" +# build xgcm grid +grid = Grid( + ds, + coords={ + "X": {"center": "xh", "right": "xq"}, + "Y": {"center": "yh", "right": "yq"}, + }, + periodic=["X"], +) + +# build wind stress curl +stress_curl = -grid.diff( + ds["tauuo"].fillna(0.0) * ds["dxCu"], "Y", boundary="fill" +) + grid.diff(ds["tauvo"].fillna(0.0) * ds["dyCv"], "X", boundary="fill") + +stress_curl = stress_curl / (ds["areacello_bu"] * 1035.0) +stress_curl = stress_curl.where(ds["wet_c"] == 1) +stress_curl = stress_curl.assign_coords({"lon": ds["geolon_c"], "lat": ds["geolon_c"]}) + +# write means on the 1/4 degree grid +stress_curl_1999_2018 = ( + stress_curl.sel(time=slice("1999", "2018")).mean(dim="time").to_dataset(name="curl") +) +stress_curl_1959_1978 = ( + stress_curl.sel(time=slice("1959", "1978")).mean(dim="time").to_dataset(name="curl") +) + + +stress_curl_1999_2018.rename({"lon": "geolon_c", "lat": "geolat_c"}).to_netcdf( + f"{outputdir}/windstresscurl_OM4p25_JRA1.4_cycle6_1999-2018.nc", + format="NETCDF3_64BIT", ) -ds.sel(time=slice("1959", "1978")).mean(dim="time").to_netcdf( - "/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstress_OM4p25_JRA1.4_cycle6_1959-1978.nc" +stress_curl_1959_1978.rename({"lon": "geolon_c", "lat": "geolat_c"}).to_netcdf( + f"{outputdir}/windstresscurl_OM4p25_JRA1.4_cycle6_1959_1978.nc", + format="NETCDF3_64BIT", ) + +# prepare grids for remapping +prepare_grids = False + +if prepare_grids: + # OM4p25 + grid_OM4p25 = xr.open_dataset(f"{ppdir_OM4p25}/ocean_annual.static.nc") + grid_src = xr.Dataset() + grid_src["lon"] = grid_OM4p25["geolon_c"] + grid_src["lat"] = grid_OM4p25["geolat_c"] + + # OM4p125 + grid_OM4p125 = xr.open_dataset(f"{ppdir_OM4p125}/ocean_annual.static.nc") + grid_dst = xr.Dataset() + grid_dst["lon"] = grid_OM4p125["geolon_c"] + grid_dst["lat"] = grid_OM4p125["geolat_c"] + + grid_src.drop_vars(["xq", "yq"]).to_netcdf(f"{outputdir}/grid_cm4p25_corners.nc") + grid_dst.drop_vars(["xq", "yq"]).to_netcdf(f"{outputdir}/grid_cm4p125_corners.nc") +else: + grid_src = xr.open_dataset(f"{outputdir}/grid_cm4p25_corners.nc") + grid_dst = xr.open_dataset(f"{outputdir}/grid_cm4p125_corners.nc") + + +# Create weights with command line tool... sigh... +sp.check_call( + f"ESMF_RegridWeightGen -s {outputdir}/grid_cm4p25_corners.nc -d {outputdir}/grid_cm4p125_corners.nc -m neareststod -w {outputdir}/remap_wgts_nn_p25_to_p125.nc", + shell=True, +) + + +# Remap using the pre-computed weights +remap = xesmf.Regridder( + grid_src, + grid_dst, + "nearest_s2d", + periodic=True, + reuse_weights=True, + weights=f"{outputdir}/remap_wgts_nn_p25_to_p125.nc", +) + +stress_curl_1999_2018_remapped = remap(stress_curl_1999_2018.chunk(dict(xq=-1, yq=-1))) +stress_curl_1959_1978_remapped = remap(stress_curl_1959_1978.chunk(dict(xq=-1, yq=-1))) + +stress_curl_1999_2018_remapped.rename({"lon": "geolon_c", "lat": "geolat_c"}).to_netcdf( + f"{outputdir}/windstresscurl_OM4p25_JRA1.4_cycle6_1999-2018_remapped_p125.nc", + format="NETCDF3_64BIT", +) +stress_curl_1959_1978_remapped.rename({"lon": "geolon_c", "lat": "geolat_c"}).to_netcdf( + f"{outputdir}/windstresscurl_OM4p25_JRA1.4_cycle6_1959_1978_remapped_p125.nc", + format="NETCDF3_64BIT", +) diff --git a/om4labs/catalogs/OMIP2_catalog_gfdl.yml b/om4labs/catalogs/OMIP2_catalog_gfdl.yml index c9b8e4a..6bfed2b 100644 --- a/om4labs/catalogs/OMIP2_catalog_gfdl.yml +++ b/om4labs/catalogs/OMIP2_catalog_gfdl.yml @@ -4,24 +4,47 @@ plugins: - module: intake_xarray sources: - OM4_windstress_1959-1978: + OM4_windstresscurl_1959-1978_gridOM4: description: "Wind stress from OM4p25 cycle6 forced by JRA1.4 from OMIP2" driver: netcdf args: - urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstress_OM4p25_JRA1.4_cycle6_1959-1978.nc' + urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstresscurl_OM4p25_JRA1.4_cycle6_1959_1978.nc' chunks: {'time': 1} xarray_kwargs: decode_times: True metadata: origin_url: '' - OM4_windstress_1999-2018: + OM4_windstresscurl_1999-2018_gridOM4: description: "Wind stress from OM4p25 cycle6 forced by JRA1.4 from OMIP2" driver: netcdf args: - urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstress_OM4p25_JRA1.4_cycle6_1999-2018.nc' + urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstresscurl_OM4p25_JRA1.4_cycle6_1999-2018.nc' chunks: {'time': 1} xarray_kwargs: decode_times: True metadata: origin_url: '' + + OM4_windstresscurl_1959-1978_gridOM4p125: + description: "Wind stress from OM4p25 cycle6 forced by JRA1.4 from OMIP2" + driver: netcdf + args: + urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstresscurl_OM4p25_JRA1.4_cycle6_1959_1978_remapped_p125.nc' + chunks: {'time': 1} + xarray_kwargs: + decode_times: True + metadata: + origin_url: '' + + OM4_windstresscurl_1999-2018_gridOM4p125: + description: "Wind stress from OM4p25 cycle6 forced by JRA1.4 from OMIP2" + driver: netcdf + args: + urlpath: '/archive/Raphael.Dussin/datasets/wind_stress_omip2/om4/windstresscurl_OM4p25_JRA1.4_cycle6_1999-2018_remapped_p125.nc' + chunks: {'time': 1} + xarray_kwargs: + decode_times: True + metadata: + origin_url: '' + diff --git a/om4labs/diags/stress_curl/stress_curl.py b/om4labs/diags/stress_curl/stress_curl.py index 33c59a4..88a5a95 100644 --- a/om4labs/diags/stress_curl/stress_curl.py +++ b/om4labs/diags/stress_curl/stress_curl.py @@ -4,6 +4,7 @@ import cmocean import xcompare import xarray as xr +import xesmf from xgcm import Grid from xcompare import plot_three_panel @@ -33,17 +34,33 @@ def parse(cliargs=None, template=False): description = " " - exclude = ["basin", "hgrid", "topog", "gridspec", "config"] + exclude = ["basin", "hgrid", "topog", "gridspec", "static"] parser = default_diag_parser( description=description, template=template, exclude=exclude ) + parser.add_argument( + "-s", + "--static", + type=str, + required=True, + help="Path to ocean.static.nc file", + ) + + parser.add_argument( + "-g", + "--hgrid", + type=str, + required=True, + help="Path to ocean_hgrid.nc file", + ) + parser.add_argument( "--dataset", type=str, required=False, - default="OM4_windstress", + default="OM4_windstresscurl", help="Name of the observational dataset, \ as provided in intake catalog", ) @@ -56,7 +73,6 @@ def parse(cliargs=None, template=False): help="Time period for OMIP2 dataset, available are 1959-1978 and 1999-2018", ) - if template is True: return parser.parse_args(None).__dict__ else: @@ -65,20 +81,26 @@ def parse(cliargs=None, template=False): def read(dictArgs): ds_model = xr.open_mfdataset(dictArgs["infile"], use_cftime=True) - dates = date_range(ds_model) + dates = date_range(ds_model) if dictArgs["obsfile"] is not None: # priority to user-provided obs file - dsobs = xr.open_mfdataset( + ds_ref_curl = xr.open_mfdataset( dictArgs["obsfile"], combine="by_coords", decode_times=False ) else: # use dataset from catalog, either from command line or default cat = open_intake_catalog(dictArgs["platform"], "OMIP2") - obsdataset = f"{dictArgs['dataset']}_{dictArgs['period']}" - ds_ref_tau = cat[obsdataset].to_dask() + obsdataset = ( + f"{dictArgs['dataset']}_{dictArgs['period']}_grid{dictArgs['config']}" + ) + ds_ref_curl = cat[obsdataset].to_dask() + + if "time" not in ds_ref_curl.dims: + ds_ref_curl = ds_ref_curl.expand_dims(dim="time") - ds_static = xr.open_mfdataset(dictArgs["static"]) + ds_static = xr.open_dataset(dictArgs["static"]) + ds_hgrid = xr.open_dataset(dictArgs["hgrid"]) # replace the nominal xq and yq by indices so that Xarray does not get confused. # Confusion arises since there are inconsistencies between static file grid and @@ -88,14 +110,36 @@ def read(dictArgs): # doing the curl operation on the stress. ds_model["xq"] = xr.DataArray(np.arange(len(ds_model["xq"])), dims=["xq"]) ds_model["yq"] = xr.DataArray(np.arange(len(ds_model["yq"])), dims=["yq"]) - ds_ref_tau["xq"] = xr.DataArray(np.arange(len(ds_ref_tau["xq"])), dims=["xq"]) - ds_ref_tau["yq"] = xr.DataArray(np.arange(len(ds_ref_tau["yq"])), dims=["yq"]) + ds_ref_curl["xq"] = xr.DataArray(np.arange(len(ds_ref_curl["xq"])), dims=["xq"]) + ds_ref_curl["yq"] = xr.DataArray(np.arange(len(ds_ref_curl["yq"])), dims=["yq"]) ds_static["xq"] = xr.DataArray(np.arange(len(ds_static["xq"])), dims=["xq"]) ds_static["yq"] = xr.DataArray(np.arange(len(ds_static["yq"])), dims=["yq"]) - ds_model.attrs = {"date_range":dates} + # override lon/lat at q-points + ds_static["geolon_c"] = xr.DataArray( + ds_hgrid["x"][::2, ::2].values, dims=("yq", "xq") + ) + ds_static["geolat_c"] = xr.DataArray( + ds_hgrid["y"][::2, ::2].values, dims=("yq", "xq") + ) + + assert (np.equal(ds_model["xq"], ds_static["xq"])).all() + assert (np.equal(ds_model["yq"], ds_static["yq"])).all() + assert (np.equal(ds_ref_curl["xq"], ds_static["xq"])).all() + assert (np.equal(ds_ref_curl["yq"], ds_static["yq"])).all() + + # override lon/lat at q-points + ds_ref_curl["geolon_c"] = xr.DataArray( + ds_hgrid["x"][::2, ::2].values, dims=("yq", "xq") + ) + ds_ref_curl["geolat_c"] = xr.DataArray( + ds_hgrid["y"][::2, ::2].values, dims=("yq", "xq") + ) + + ds_model.attrs = {"date_range": dates} + + return ds_model, ds_ref_curl, ds_static - return ds_model, ds_ref_tau, ds_static def calc_curl_stress( ds, @@ -168,60 +212,106 @@ def calc_curl_stress( return stress_curl -def calculate(ds_model, ds_ref_tau, ds_static, dictArgs): +def calculate(ds_model, ds_ref_curl, ds_static, dictArgs): curl_model = xr.Dataset( { "stress_curl": calc_curl_stress(ds_model, ds_static), "areacello_bu": ds_static["areacello_bu"], + "geolon_c": ds_static["geolon_c"], + "geolat_c": ds_static["geolat_c"], } ) curl_ref = xr.Dataset( { - "stress_curl": calc_curl_stress(ds_ref_tau, ds_static), + "stress_curl": ds_ref_curl["curl"], "areacello_bu": ds_static["areacello_bu"], + "geolon_c": ds_ref_curl["geolon_c"], + "geolat_c": ds_ref_curl["geolat_c"], } ) - results = xcompare.compare_datasets(curl_model, curl_ref) + return curl_model, curl_ref - return results +def plot(dictArgs, curl_model, curl_ref): -def plot(dictArgs, results): + results = xcompare.compare_datasets(curl_model, curl_ref) - fig = xcompare.plot_three_panel(results, "stress_curl", cmap=cmocean.cm.delta, - projection=ccrs.Robinson(), coastlines = False, - vmin=-3e-10, vmax=3e-10, diffvmin=-3e-10, diffvmax=3e-10, - labels=[dictArgs["label"],"OMIP reference"], - lon_range=(-180,180),lat_range=(-90,90)) + fig = xcompare.plot_three_panel( + results, + "stress_curl", + cmap=cmocean.cm.delta, + projection=ccrs.Robinson(), + coastlines=False, + vmin=-3e-10, + vmax=3e-10, + diffvmin=-3e-10, + diffvmax=3e-10, + labels=[dictArgs["label"], "OMIP reference"], + lon_range=(-180, 180), + lat_range=(-90, 90), + ) return fig def plot_old( - field, + curl_model, + curl_ref, vmin=-3e-10, vmax=3e-10, lat_lon_ext=[-180, 180, -85.0, 90.0], lon="geolon_c", lat="geolat_c", + fieldname="stress_curl", cmap=cmocean.cm.delta, title="Curl of stress (N/m**2) acting on ocean surface", ): # convert xarray to ordinary numpy arrays - if isinstance(field, xr.DataArray): - geolon = field[lon].values - geolat = field[lat].values - field = field.values - - fig = plt.figure(figsize=[22, 8]) - ax = fig.add_subplot(projection=ccrs.Robinson(), facecolor="grey") - cb = ax.pcolormesh( - geolon, - geolat, - field, + if isinstance(curl_model, xr.Dataset): + geolon_model = curl_model[lon].values + geolat_model = curl_model[lat].values + data_model = curl_model[fieldname].values + + if isinstance(curl_ref, xr.Dataset): + geolon_ref = curl_ref[lon].values + geolat_ref = curl_ref[lat].values + data_ref = curl_ref[fieldname].values + + # fig = plt.figure(figsize=[22, 8]) + fig, axs = plt.subplots( + ncols=1, + nrows=3, + subplot_kw={"projection": ccrs.Robinson(), "facecolor": "grey"}, + ) + # ax = fig.add_subplot(projection=ccrs.Robinson(), facecolor="grey") + # ax = fig.add_subplot(projection=ccrs.Robinson(), facecolor="grey") + cb = axs[0].pcolormesh( + geolon_model, + geolat_model, + data_model.squeeze(), + transform=ccrs.PlateCarree(), + vmin=vmin, + vmax=vmax, + cmap=cmap, + ) + + cb = axs[1].pcolormesh( + geolon_ref, + geolat_ref, + data_ref.squeeze(), + transform=ccrs.PlateCarree(), + vmin=vmin, + vmax=vmax, + cmap=cmap, + ) + + cb = axs[2].pcolormesh( + geolon_ref, + geolat_ref, + data_model.squeeze() - data_ref.squeeze(), transform=ccrs.PlateCarree(), vmin=vmin, vmax=vmax, @@ -229,13 +319,13 @@ def plot_old( ) # add separate colorbar - cb = plt.colorbar(cb, ax=ax, format="%.1e", extend="both", shrink=0.6) - cb.ax.tick_params(labelsize=12) + # cb = plt.colorbar(cb, ax=ax, format="%.1e", extend="both", shrink=0.6) + # cb.ax.tick_params(labelsize=12) # add gridlines and extent of lat/lon - ax.gridlines(color="black", alpha=0.5, linestyle="--") - ax.set_extent(lat_lon_ext, crs=ccrs.PlateCarree()) - _ = plt.title(title, fontsize=14) + # ax.gridlines(color="black", alpha=0.5, linestyle="--") + # ax.set_extent(lat_lon_ext, crs=ccrs.PlateCarree()) + # _ = plt.title(title, fontsize=14) return fig @@ -247,8 +337,9 @@ def run(dictArgs): ds_model, ds_ref_tau, ds_static = read(dictArgs) - results = calculate(ds_model, ds_ref_tau, ds_static, dictArgs) - figs = plot(dictArgs, results) + curl_model, curl_ref = calculate(ds_model, ds_ref_tau, ds_static, dictArgs) + # figs = plot(dictArgs, results) + figs = plot(dictArgs, curl_model, curl_ref) figs = [figs] if not isinstance(figs, list) else figs assert isinstance(figs, list), "Figures must be inside a list object"