From 680dd7ca419b31fc259521f5f8ed7b4e854fc352 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Mon, 29 Apr 2024 18:00:38 -0700 Subject: [PATCH 01/15] in progress ... --- pcmdi_metrics/mjo/lib/__init__.py | 6 +- pcmdi_metrics/mjo/lib/lib_mjo.py | 39 ++++++- pcmdi_metrics/mjo/lib/mjo_metric_calc.py | 126 +++++++++++++++-------- pcmdi_metrics/utils/__init__.py | 2 + 4 files changed, 126 insertions(+), 47 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/__init__.py b/pcmdi_metrics/mjo/lib/__init__.py index 5f4352ff6..8af9fa945 100644 --- a/pcmdi_metrics/mjo/lib/__init__.py +++ b/pcmdi_metrics/mjo/lib/__init__.py @@ -7,11 +7,13 @@ decorate_2d_array_axes, generate_axes_and_decorate, get_daily_ano_segment, - interp2commonGrid, + #interp2commonGrid, + interp2commonGrid_xcdat, mjo_metrics_to_json, output_power_spectra, space_time_spectrum, - subSliceSegment, + # subSliceSegment, + subSliceSegment_xcdat, taper, unit_conversion, write_netcdf_output, diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 7e3cb58dd..e2e5d676f 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -15,6 +15,12 @@ from scipy import signal import pcmdi_metrics +import xarray as xr + +from typing import Union +from pcmdi_metrics.io import get_time_key +from pcmdi_metrics.utils import create_target_grid, regrid +from pcmdi_metrics.io import select_subset def interp2commonGrid(d, dlat, debug=False): @@ -34,7 +40,20 @@ def interp2commonGrid(d, dlat, debug=False): return d2 -def subSliceSegment(d, year, mon, day, length): +def interp2commonGrid_xcdat(ds, data_var, dlat, dlon=None, debug=False): + if dlon is None: + dlon = dlat + nlat = int(180/dlat) + nlon = int(360/dlon) + grid = create_target_grid(target_grid_resolution=f"{dlat}x{dlon}") + ds_regrid = regrid(ds, data_var, grid) + ds_regrid_subset = select_subset(ds, lat=(-10, 10)) + if debug: + print("debug: ds_regrid_subset[data_var] shape:", ds_regrid_subset[data_var].shape) + return ds_regrid_subset + + +def subSliceSegment(ds, data_var, year, mon, day, length): """ Note: From given cdms array (3D: time and spatial 2D) Subslice to get segment with given length starting from given time. @@ -56,6 +75,24 @@ def subSliceSegment(d, year, mon, day, length): return d2 +def subSliceSegment_xcdat(ds: Union[xr.Dataset, xr.DataArray], year: int, mon: int, day:int, length: int) -> Union[xr.Dataset, xr.DataArray]: + """ + Note: From given cdms array (3D: time and spatial 2D) + Subslice to get segment with given length starting from given time. + input + - ds: xarray dataset or dataArray + - year: segment starting year (integer) + - mon: segment starting month (integer) + - day: segement starting day (integer) + - length: segment length (integer) + """ + + time_key = get_time_key(ds) + n = list(ds[time_key].values).index(ds.sel(time=f'{year:04}-{mon:02}-{day:02}')[time_key]) + + return ds.isel(time=slice(n, n + length)) # slie 180 time steps starting from above index + + def Remove_dailySeasonalCycle(d, d_cyc): """ Note: Remove daily seasonal cycle diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py index 88b1a16e2..73cf3fed2 100644 --- a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py +++ b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py @@ -1,26 +1,33 @@ import os -import cdms2 -import cdtime +#import cdms2 +#import cdtime import MV2 import numpy as np +from datetime import datetime + from pcmdi_metrics.mjo.lib import ( Remove_dailySeasonalCycle, calculate_ewr, generate_axes_and_decorate, get_daily_ano_segment, - interp2commonGrid, + #interp2commonGrid, + interp2commonGrid_xcdat, output_power_spectra, space_time_spectrum, - subSliceSegment, - unit_conversion, + #subSliceSegment, + subSliceSegment_xcdat, + #unit_conversion, write_netcdf_output, ) from .debug_chk_plot import debug_chk_plot from .plot_wavenumber_frequency_power import plot_power +from pcmdi_metrics.io import xcdat_open, get_time, get_latitude, get_longitude, get_time_key +from pcmdi_metrics.utils import adjust_units + def mjo_metric_ewr_calculation( mip, @@ -34,38 +41,56 @@ def mjo_metric_ewr_calculation( degX, UnitsAdjust, inputfile, - var, - startYear, - endYear, - segmentLength, - dir_paths, - season="NDJFMA", + data_var: str, + startYear: int, + endYear: int, + segmentLength: int, + dir_paths: str, + season: str="NDJFMA", ): # Open file to read daily dataset if debug: print("debug: open file") - f = cdms2.open(inputfile) - d = f[var] - tim = d.getTime() - comTim = tim.asComponentTime() - lat = d.getLatitude() - lon = d.getLongitude() + #f = cdms2.open(inputfile) + #d = f[data_var] + ds = xcdat_open(inputfile) + + #tim = d.getTime() + #comTim = tim.asComponentTime() + #lat = d.getLatitude() + #lon = d.getLongitude() + + #tim = get_time(ds) + lat = get_latitude(ds) + lon = get_longitude(ds) # Get starting and ending year and month if debug: print("debug: check time") - first_time = comTim[0] - last_time = comTim[-1] + + #first_time = comTim[0] + #last_time = comTim[-1] + + time_key = get_time_key(ds) + first_time = ds.indexes[time_key].to_datetimeindex()[0].to_pydatetime() + last_time = ds.indexes[time_key].to_datetimeindex()[-1].to_pydatetime() if season == "NDJFMA": # Adjust years to consider only when continuous NDJFMA is available + """ if first_time > cdtime.comptime(startYear, 11, 1): startYear += 1 if last_time < cdtime.comptime(endYear, 4, 30): endYear -= 1 - + """ + if first_time > datetime(startYear, 11, 1): + startYear += 1 + if last_time < datetime(endYear, 4, 30): + endYear -= 1 + # Number of grids for 2d fft input - NL = len(d.getLongitude()) # number of grid in x-axis (longitude) + #NL = len(d.getLongitude()) # number of grid in x-axis (longitude) + NL = len(lon.values) # number of grid in x-axis (longitude) if cmmGrid: NL = int(360 / degX) NT = segmentLength # number of time step for each segment (need to be an even number) @@ -88,35 +113,46 @@ def mjo_metric_ewr_calculation( # Store each year's segment in a dictionary: segment[year] segment = {} segment_ano = {} - daSeaCyc = MV2.zeros((NT, d.shape[1], d.shape[2])) + #daSeaCyc = MV2.zeros((NT, d.shape[1], d.shape[2])) + daSeaCyc = np.zeros((NT, ds[data_var].shape[1], ds[data_var].shape[2])) + + # Loop over years for year in range(startYear, endYear): print(year) - segment[year] = subSliceSegment(d, year, mon, day, NT) + #segment[year] = subSliceSegment(d, year, mon, day, NT) + segment[year] = subSliceSegment_xcdat(ds, year, mon, day, NT) # units conversion - segment[year] = unit_conversion(segment[year], UnitsAdjust) + #segment[year] = unit_conversion(segment[year], UnitsAdjust) + segment[year][data_var] = adjust_units(segment[year][data_var], UnitsAdjust) # Get climatology of daily seasonal cycle - daSeaCyc = MV2.add(MV2.divide(segment[year], float(numYear)), daSeaCyc) + #daSeaCyc = MV2.add(MV2.divide(segment[year], float(numYear)), daSeaCyc) + daSeaCyc = np.add(np.divide(segment[year][data_var].values, float(numYear)), daSeaCyc) + # Remove daily seasonal cycle from each segment if numYear > 1: + # Loop over years for year in range(startYear, endYear): - segment_ano[year] = Remove_dailySeasonalCycle(segment[year], daSeaCyc) + #segment_ano[year] = Remove_dailySeasonalCycle(segment[year], daSeaCyc) + segment_ano[year] = segment[year] - daSeaCyc else: segment_ano[year] = segment[year] + # Assign lat/lon to arrays - daSeaCyc.setAxis(1, lat) - daSeaCyc.setAxis(2, lon) - segment_ano[year].setAxis(1, lat) - segment_ano[year].setAxis(2, lon) - - """ Space-time power spectra - - Handle each segment (i.e. each year) separately. - 1. Get daily time series (3D: time and spatial 2D) - 2. Meridionally average (2D: time and spatial, i.e., longitude) - 3. Get anomaly by removing time mean of the segment - 4. Proceed 2-D FFT to get power. - Then get multi-year averaged power after the year loop. - """ + # daSeaCyc.setAxis(1, lat) + # daSeaCyc.setAxis(2, lon) + # segment_ano[year].setAxis(1, lat) + # segment_ano[year].setAxis(2, lon) + + # ----------------------------------------------------------------- + # Space-time power spectra + # ----------------------------------------------------------------- + # Handle each segment (i.e. each year) separately. + # 1. Get daily time series (3D: time and spatial 2D) + # 2. Meridionally average (2D: time and spatial, i.e., longitude) + # 3. Get anomaly by removing time mean of the segment + # 4. Proceed 2-D FFT to get power. + # Then get multi-year averaged power after the year loop. + # ----------------------------------------------------------------- # Define array for archiving power from each year segment Power = np.zeros((numYear, NT + 1, NL + 1), np.float) @@ -129,7 +165,8 @@ def mjo_metric_ewr_calculation( d_seg = segment_ano[year] # Regrid: interpolation to common grid if cmmGrid: - d_seg = interp2commonGrid(d_seg, degX, debug=debug) + #d_seg = interp2commonGrid(d_seg, degX, debug=debug) + d_seg = interp2commonGrid_xcdat(d_seg, degX, debug=debug) # Subregion, meridional average, and remove segment time mean d_seg_x_ano = get_daily_ano_segment(d_seg) # Compute space-time spectrum @@ -166,9 +203,9 @@ def mjo_metric_ewr_calculation( os.makedirs(dir_paths["graphics"], exist_ok=True) fout = os.path.join(dir_paths["graphics"], output_filename) if model == "obs": - title = f"OBS ({run})\n{var.capitalize()}, {season} {startYear}-{endYear}" + title = f"OBS ({run})\n{data_var.capitalize()}, {season} {startYear}-{endYear}" else: - title = f"{mip.upper()}: {model} ({run})\n{var.capitalize()}, {season} {startYear}-{endYear}" + title = f"{mip.upper()}: {model} ({run})\n{data_var.capitalize()}, {season} {startYear}-{endYear}" if cmmGrid: title += ", common grid (2.5x2.5deg)" @@ -189,5 +226,6 @@ def mjo_metric_ewr_calculation( d_seg_x_ano, Power, OEE, segment[year], daSeaCyc, segment_ano[year] ) - f.close() + #f.close() + ds.close() return metrics_result diff --git a/pcmdi_metrics/utils/__init__.py b/pcmdi_metrics/utils/__init__.py index d870b03c1..10c2d6e31 100644 --- a/pcmdi_metrics/utils/__init__.py +++ b/pcmdi_metrics/utils/__init__.py @@ -1,3 +1,5 @@ +from .adjust_units import adjust_units + from .custom_season import ( custom_season_average, custom_season_departure, From 8236449fda11572db575b8fb3d6bad0040aa1a58 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 1 May 2024 14:55:20 -0700 Subject: [PATCH 02/15] update --- pcmdi_metrics/mjo/lib/__init__.py | 20 +- pcmdi_metrics/mjo/lib/debug_chk_plot.py | 28 +-- pcmdi_metrics/mjo/lib/lib_mjo.py | 226 ++++++++++++++++-- pcmdi_metrics/mjo/lib/mjo_metric_calc.py | 151 ++++++------ .../lib/plot_wavenumber_frequency_power.py | 9 +- 5 files changed, 315 insertions(+), 119 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/__init__.py b/pcmdi_metrics/mjo/lib/__init__.py index 8af9fa945..2ab62ad3c 100644 --- a/pcmdi_metrics/mjo/lib/__init__.py +++ b/pcmdi_metrics/mjo/lib/__init__.py @@ -3,20 +3,26 @@ from .dict_merge import dict_merge # noqa from .lib_mjo import ( # noqa Remove_dailySeasonalCycle, - calculate_ewr, + # calculate_ewr, + calculate_ewr_xcdat, decorate_2d_array_axes, - generate_axes_and_decorate, - get_daily_ano_segment, - #interp2commonGrid, + # generate_axes_and_decorate, + generate_axes_and_decorate_xcdat, + # get_daily_ano_segment, + get_daily_ano_segment_xcdat, + # interp2commonGrid, interp2commonGrid_xcdat, mjo_metrics_to_json, - output_power_spectra, - space_time_spectrum, + # output_power_spectra, + output_power_spectra_xcdat, + # space_time_spectrum, + space_time_spectrum_xcdat, # subSliceSegment, subSliceSegment_xcdat, taper, unit_conversion, - write_netcdf_output, + # write_netcdf_output, + write_netcdf_output_xcdat, ) from .mjo_metric_calc import mjo_metric_ewr_calculation # noqa from .plot_wavenumber_frequency_power import plot_power # noqa diff --git a/pcmdi_metrics/mjo/lib/debug_chk_plot.py b/pcmdi_metrics/mjo/lib/debug_chk_plot.py index 656f545e9..75db53cd1 100644 --- a/pcmdi_metrics/mjo/lib/debug_chk_plot.py +++ b/pcmdi_metrics/mjo/lib/debug_chk_plot.py @@ -4,29 +4,15 @@ import matplotlib.pyplot as plt from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter +from pcmdi_metrics.io import get_latitude, get_longitude + def debug_chk_plot(d_seg_x_ano, Power, OEE, segment_year, daSeaCyc, segment_ano_year): os.makedirs("debug", exist_ok=True) - """ FIX ME --- - x = vcs.init() - x.plot(d_seg_x_ano) - x.png('debug/d_seg_x_ano.png') - - x.clear() - x.plot(Power) - x.png('debug/power.png') - - x.clear() - x.plot(OEE) - x.png('debug/OEE.png') - """ - print("type(segment_year)", type(segment_year)) print("segment_year.shape:", segment_year.shape) - print(segment_year.getAxis(0)) - print(segment_year.getAxis(1)) - print(segment_year.getAxis(2)) + plot_map(segment_year[0], "debug/segment.png") print("type(daSeaCyc)", type(daSeaCyc)) @@ -35,16 +21,14 @@ def debug_chk_plot(d_seg_x_ano, Power, OEE, segment_year, daSeaCyc, segment_ano_ print("type(segment_ano_year)", type(segment_ano_year)) print("segment_ano_year.shape:", segment_ano_year.shape) - print(segment_ano_year.getAxis(0)) - print(segment_ano_year.getAxis(1)) - print(segment_ano_year.getAxis(2)) + plot_map(segment_ano_year[0], "debug/segment_ano.png") def plot_map(data, filename): fig = plt.figure(figsize=(10, 6)) - lons = data.getLongitude() - lats = data.getLatitude() + lons = get_longitude(data) + lats = get_latitude(data) ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=180)) im = ax.contourf(lons, lats, data, transform=ccrs.PlateCarree(), cmap="viridis") ax.coastlines() diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index e2e5d676f..de4efff3b 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -7,20 +7,19 @@ https://doi.org/10.1007/s00382-017-3558-4 """ +from typing import Union + import cdms2 import cdtime import cdutil import MV2 import numpy as np +import xarray as xr from scipy import signal import pcmdi_metrics -import xarray as xr - -from typing import Union -from pcmdi_metrics.io import get_time_key +from pcmdi_metrics.io import get_time_key, select_subset from pcmdi_metrics.utils import create_target_grid, regrid -from pcmdi_metrics.io import select_subset def interp2commonGrid(d, dlat, debug=False): @@ -41,19 +40,21 @@ def interp2commonGrid(d, dlat, debug=False): def interp2commonGrid_xcdat(ds, data_var, dlat, dlon=None, debug=False): - if dlon is None: + if dlon is None: dlon = dlat - nlat = int(180/dlat) - nlon = int(360/dlon) + # nlat = int(180 / dlat) + # nlon = int(360 / dlon) grid = create_target_grid(target_grid_resolution=f"{dlat}x{dlon}") ds_regrid = regrid(ds, data_var, grid) - ds_regrid_subset = select_subset(ds, lat=(-10, 10)) + ds_regrid_subset = select_subset(ds_regrid, lat=(-10, 10)) if debug: - print("debug: ds_regrid_subset[data_var] shape:", ds_regrid_subset[data_var].shape) + print( + "debug: ds_regrid_subset[data_var] shape:", ds_regrid_subset[data_var].shape + ) return ds_regrid_subset -def subSliceSegment(ds, data_var, year, mon, day, length): +def subSliceSegment(d, year, mon, day, length): """ Note: From given cdms array (3D: time and spatial 2D) Subslice to get segment with given length starting from given time. @@ -75,7 +76,9 @@ def subSliceSegment(ds, data_var, year, mon, day, length): return d2 -def subSliceSegment_xcdat(ds: Union[xr.Dataset, xr.DataArray], year: int, mon: int, day:int, length: int) -> Union[xr.Dataset, xr.DataArray]: +def subSliceSegment_xcdat( + ds: Union[xr.Dataset, xr.DataArray], year: int, mon: int, day: int, length: int +) -> Union[xr.Dataset, xr.DataArray]: """ Note: From given cdms array (3D: time and spatial 2D) Subslice to get segment with given length starting from given time. @@ -86,11 +89,15 @@ def subSliceSegment_xcdat(ds: Union[xr.Dataset, xr.DataArray], year: int, mon: i - day: segement starting day (integer) - length: segment length (integer) """ - + time_key = get_time_key(ds) - n = list(ds[time_key].values).index(ds.sel(time=f'{year:04}-{mon:02}-{day:02}')[time_key]) - - return ds.isel(time=slice(n, n + length)) # slie 180 time steps starting from above index + n = list(ds[time_key].values).index( + ds.sel(time=f"{year:04}-{mon:02}-{day:02}")[time_key] + ) + + return ds.isel( + time=slice(n, n + length) + ) # slie 180 time steps starting from above index def Remove_dailySeasonalCycle(d, d_cyc): @@ -133,6 +140,33 @@ def get_daily_ano_segment(d_seg): return d_seg_x_ano +def get_daily_ano_segment_xcdat(d_seg, data_var): + """ + Note: 1. Get daily time series (3D: time and spatial 2D) + 2. Meridionally average (2D: time and spatial, i.e., longitude) + 3. Get anomaly by removing time mean of the segment + input + - d_seg: xarray dataset + - data_var: name of variable + output + - d_seg_x_ano: xarray dataset that contains 2d output array + """ + # sub region + d_seg = select_subset(d_seg, lat=(-10, 10)) + + # Get meridional average (3d (t, y, x) to 2d (t, y)) + d_seg_x = d_seg.spatial.average(data_var, axis=["Y"]) + + # Get time-average in the segment on each longitude grid + d_seg_x_ave = d_seg_x.temporal.average(data_var) + + # Remove time mean for each segment + d_seg_x_ano = d_seg.copy() + d_seg_x_ano[data_var] = d_seg_x[data_var] - d_seg_x_ave[data_var] + + return d_seg_x_ano + + def space_time_spectrum(d_seg_x_ano): """ input @@ -169,6 +203,42 @@ def space_time_spectrum(d_seg_x_ano): return p +def space_time_spectrum_xcdat(d_seg_x_ano, data_var): + """ + input + - d: 2d cdms MV2 array (t (time), n (space)) + output + - p: 2d array for power + NOTE: Below code taken from + https://github.com/CDAT/wk/blob/2b953281c7a4c5d0ac2d79fcc3523113e31613d5/WK/process.py#L188 + """ + # Number of grid in longitude axis, and timestep for each segment + NTSub = d_seg_x_ano[data_var].shape[0] # NTSub + NL = d_seg_x_ano[data_var].shape[1] # NL + # Tapering + d_seg_x_ano[data_var] = taper(d_seg_x_ano[data_var]) + # Power sepctrum analysis + EE = np.fft.fft2(d_seg_x_ano[data_var], axes=(1, 0)) / float(NL) / float(NTSub) + # Now the array EE(n,t) contains the (complex) space-time spectrum. + """ + Create array PEE(NL+1,NT/2+1) which contains the (real) power spectrum. + Note how the PEE array is arranged into a different order to EE. + In this code, PEE is "Power", and its multiyear average will be "power" + """ + # OK NOW THE LITTLE MAGIC WITH REORDERING ! + A = np.absolute(EE[0 : NTSub // 2 + 1, 1 : NL // 2 + 1]) ** 2 + B = np.absolute(EE[NTSub // 2 : NTSub, 1 : NL // 2 + 1]) ** 2 + C = np.absolute(EE[NTSub // 2 : NTSub, 0 : NL // 2 + 1]) ** 2 + D = np.absolute(EE[0 : NTSub // 2 + 1, 0 : NL // 2 + 1]) ** 2 + # Define returning array + p = np.zeros((NTSub + 1, NL + 1), np.float) + p[NTSub // 2 :, : NL // 2] = A[:, ::-1] + p[: NTSub // 2, : NL // 2] = B[:, ::-1] + p[NTSub // 2 + 1 :, NL // 2 :] = C[::-1, :] + p[: NTSub // 2 + 1, NL // 2 :] = D[::-1, :] + return p + + def taper(data): """ Note: taper first and last 45 days with cosine window, using scipy.signal function @@ -177,7 +247,7 @@ def taper(data): output: - data: tapered data """ - window = signal.tukey(len(data)) + window = signal.windows.tukey(len(data)) data2 = data.copy() for i in range(0, len(data)): data2[i] = MV2.multiply(data[i][:], window[i]) @@ -254,6 +324,49 @@ def generate_axes_and_decorate(Power, NT, NL): return Power, ff, ss +def generate_axes_and_decorate_xcdat(Power, NT, NL): + """ + Note: Generates axes for the decoration + input + - Power: 2d numpy array + - NT: integer, number of time step + - NL: integer, number of spatial grid + output + - Power: decorated 2d cdms array + - ff: frequency axis + - ss: wavenumber axis + """ + # frequency + ff = [] + for t in range(0, NT + 1): + ff.append(float(t - NT / 2) / float(NT)) + ff = np.array(ff) + + # wave number + ss = [] + for n in range(0, NL + 1): + ss.append(float(n) - float(NL / 2)) + ss = np.array(ss) + + # Add name attributes to x and y coordinates + x_coords = xr.IndexVariable( + "zonalwavenumber", ss, attrs={"name": "zonalwavenumber", "units": "-"} + ) + y_coords = xr.IndexVariable( + "frequency", ff, attrs={"name": "frequency", "units": "cycles per day"} + ) + + # Create an xarray DataArray + da = xr.DataArray( + Power, + coords={"frequency": y_coords, "zonalwavenumber": x_coords}, + dims=["frequency", "zonalwavenumber"], + name="power", + ) + + return da + + def output_power_spectra(NL, NT, Power, ff, ss): """ Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) @@ -287,6 +400,54 @@ def output_power_spectra(NL, NT, Power, ff, ss): return OEE +def output_power_spectra_xcdat(NL, NT, Power): + """ + Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) + """ + # The corresponding frequencies, ff, and wavenumbers, ss, are:- + PEE = Power + + ff = Power.frequency + ss = Power.zonalwavenumber + + OEE = np.zeros((21, 11)) + for n in range(int(NL / 2), int(NL / 2) + 1 + 10): + nn = n - int(NL / 2) + for t in range(int(NT / 2) - 10, int(NT / 2 + 1 + 10)): + tt = -(int(NT / 2) + 1) + 11 + t + OEE[tt, nn] = PEE[t, n] + a = list((ff[i] for i in range(int(NT / 2) - 10, int(NT / 2) + 1 + 10))) + b = list((ss[i] for i in range(int(NL / 2), int(NL / 2) + 1 + 10))) + a = np.array(a) + b = np.array(b) + # Decoration + + # Add name attributes to x and y coordinates + x_coords = xr.IndexVariable( + "zonalwavenumber", b, attrs={"name": "zonalwavenumber", "units": "-"} + ) + y_coords = xr.IndexVariable( + "frequency", a, attrs={"name": "frequency", "units": "cycles per day"} + ) + + # Create an xarray DataArray + OEE = xr.DataArray( + OEE, + coords={"frequency": y_coords, "zonalwavenumber": x_coords}, + dims=["frequency", "zonalwavenumber"], + name="power", + ) + + # Transpose for visualization + # OEE = np.transpose(OEE, (1, 0)) + print("before transpose, OEE.shape:", OEE.shape) + transposed_OEE = OEE.transpose() + print("after transpose, transposed_OEE.shape:", transposed_OEE.shape) + return transposed_OEE + + # return OEE + + def write_netcdf_output(d, fname): """ Note: write array in a netcdf file @@ -299,6 +460,17 @@ def write_netcdf_output(d, fname): fo.close() +def write_netcdf_output_xcdat(da, fname): + """ + Note: write array in a netcdf file + input + - d: array + - fname: string. directory path and name of the netcd file, without .nc + """ + ds = xr.Dataset({da.name: da}) + ds.to_netcdf(fname + ".nc") + + def calculate_ewr(OEE): """ According to DK's gs script (MSD/level_2/sample/stps/e.w.ratio.gs.sample), @@ -315,6 +487,26 @@ def calculate_ewr(OEE): return ewr, eastPower, westPower +def calculate_ewr_xcdat(OEE): + """ + According to DK's gs script (MSD/level_2/sample/stps/e.w.ratio.gs.sample), + E/W ratio is calculated as below: + 'd amean(power,x=14,x=17,y=2,y=4)/aave(power,x=5,x=8,y=2,y=4)' + where x for frequency and y for wavenumber. + Actual ranges of frequency and wavenumber have been checked and applied. + """ + east_power_domain = OEE.sel( + zonalwavenumber=slice(1, 3), frequency=slice(0.0166667, 0.0333333) + ) + west_power_domain = OEE.sel( + zonalwavenumber=slice(1, 3), frequency=slice(-0.0333333, -0.0166667) + ) + eastPower = np.average(east_power_domain) + westPower = np.average(west_power_domain) + ewr = eastPower / westPower + return ewr, eastPower, westPower + + def unit_conversion(data, UnitsAdjust): """ Convert unit following given tuple using MV2 diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py index 73cf3fed2..c062f9dd3 100644 --- a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py +++ b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py @@ -1,33 +1,28 @@ import os +from datetime import datetime -#import cdms2 -#import cdtime -import MV2 +# import cdms2 +# import cdtime +# import MV2 import numpy as np +import xarray as xr -from datetime import datetime - -from pcmdi_metrics.mjo.lib import ( - Remove_dailySeasonalCycle, - calculate_ewr, - generate_axes_and_decorate, - get_daily_ano_segment, - #interp2commonGrid, +from pcmdi_metrics.io import get_latitude, get_longitude, get_time_key, xcdat_open +from pcmdi_metrics.mjo.lib import ( # calculate_ewr,; generate_axes_and_decorate,; get_daily_ano_segment,; interp2commonGrid,; output_power_spectra,; space_time_spectrum,; subSliceSegment,; unit_conversion,; Remove_dailySeasonalCycle,; write_netcdf_output, + calculate_ewr_xcdat, + generate_axes_and_decorate_xcdat, + get_daily_ano_segment_xcdat, interp2commonGrid_xcdat, - output_power_spectra, - space_time_spectrum, - #subSliceSegment, + output_power_spectra_xcdat, + space_time_spectrum_xcdat, subSliceSegment_xcdat, - #unit_conversion, - write_netcdf_output, + write_netcdf_output_xcdat, ) +from pcmdi_metrics.utils import adjust_units from .debug_chk_plot import debug_chk_plot from .plot_wavenumber_frequency_power import plot_power -from pcmdi_metrics.io import xcdat_open, get_time, get_latitude, get_longitude, get_time_key -from pcmdi_metrics.utils import adjust_units - def mjo_metric_ewr_calculation( mip, @@ -46,50 +41,33 @@ def mjo_metric_ewr_calculation( endYear: int, segmentLength: int, dir_paths: str, - season: str="NDJFMA", + season: str = "NDJFMA", ): # Open file to read daily dataset if debug: - print("debug: open file") - #f = cdms2.open(inputfile) - #d = f[data_var] + print(f"debug: open file: {inputfile}") + ds = xcdat_open(inputfile) - - #tim = d.getTime() - #comTim = tim.asComponentTime() - #lat = d.getLatitude() - #lon = d.getLongitude() - - #tim = get_time(ds) + lat = get_latitude(ds) lon = get_longitude(ds) # Get starting and ending year and month if debug: print("debug: check time") - - #first_time = comTim[0] - #last_time = comTim[-1] - + time_key = get_time_key(ds) first_time = ds.indexes[time_key].to_datetimeindex()[0].to_pydatetime() - last_time = ds.indexes[time_key].to_datetimeindex()[-1].to_pydatetime() + last_time = ds.indexes[time_key].to_datetimeindex()[-1].to_pydatetime() if season == "NDJFMA": # Adjust years to consider only when continuous NDJFMA is available - """ - if first_time > cdtime.comptime(startYear, 11, 1): - startYear += 1 - if last_time < cdtime.comptime(endYear, 4, 30): - endYear -= 1 - """ if first_time > datetime(startYear, 11, 1): startYear += 1 if last_time < datetime(endYear, 4, 30): endYear -= 1 - + # Number of grids for 2d fft input - #NL = len(d.getLongitude()) # number of grid in x-axis (longitude) NL = len(lon.values) # number of grid in x-axis (longitude) if cmmGrid: NL = int(360 / degX) @@ -110,39 +88,58 @@ def mjo_metric_ewr_calculation( mon = 5 numYear = endYear - startYear + 1 day = 1 + # Store each year's segment in a dictionary: segment[year] segment = {} segment_ano = {} - #daSeaCyc = MV2.zeros((NT, d.shape[1], d.shape[2])) - daSeaCyc = np.zeros((NT, ds[data_var].shape[1], ds[data_var].shape[2])) - + + daSeaCyc = xr.DataArray( + np.zeros((NT, ds[data_var].shape[1], ds[data_var].shape[2])), + dims=["day", "lat", "lon"], + coords={"day": np.arange(180), "lat": lat, "lon": lon}, + ) + daSeaCyc_values = daSeaCyc.values.copy() + + if debug: + print("debug: before year loop: daSeaCyc.shape:", daSeaCyc.shape) + # Loop over years for year in range(startYear, endYear): print(year) - #segment[year] = subSliceSegment(d, year, mon, day, NT) segment[year] = subSliceSegment_xcdat(ds, year, mon, day, NT) # units conversion - #segment[year] = unit_conversion(segment[year], UnitsAdjust) segment[year][data_var] = adjust_units(segment[year][data_var], UnitsAdjust) + if debug: + print( + "debug: year, segment[year][data_var].shape:", + year, + segment[year][data_var].shape, + ) # Get climatology of daily seasonal cycle - #daSeaCyc = MV2.add(MV2.divide(segment[year], float(numYear)), daSeaCyc) - daSeaCyc = np.add(np.divide(segment[year][data_var].values, float(numYear)), daSeaCyc) - + # daSeaCyc_values = np.add( + # np.divide(segment[year][data_var].values, float(numYear)), daSeaCyc_values + # ) + daSeaCyc_values = ( + segment[year][data_var].values / float(numYear) + ) + daSeaCyc_values + + daSeaCyc.values = daSeaCyc_values + + if debug: + print("debug: after year loop: daSeaCyc.shape:", daSeaCyc.shape) + # Remove daily seasonal cycle from each segment if numYear > 1: # Loop over years for year in range(startYear, endYear): - #segment_ano[year] = Remove_dailySeasonalCycle(segment[year], daSeaCyc) - segment_ano[year] = segment[year] - daSeaCyc + # segment_ano[year] = Remove_dailySeasonalCycle(segment[year], daSeaCyc) + segment_ano[year] = segment[year].copy() + segment_ano[year][data_var].values = ( + segment[year][data_var].values - daSeaCyc.values + ) else: segment_ano[year] = segment[year] - - # Assign lat/lon to arrays - # daSeaCyc.setAxis(1, lat) - # daSeaCyc.setAxis(2, lon) - # segment_ano[year].setAxis(1, lat) - # segment_ano[year].setAxis(2, lon) - + # ----------------------------------------------------------------- # Space-time power spectra # ----------------------------------------------------------------- @@ -165,24 +162,30 @@ def mjo_metric_ewr_calculation( d_seg = segment_ano[year] # Regrid: interpolation to common grid if cmmGrid: - #d_seg = interp2commonGrid(d_seg, degX, debug=debug) - d_seg = interp2commonGrid_xcdat(d_seg, degX, debug=debug) + d_seg = interp2commonGrid_xcdat(d_seg, data_var, degX, debug=debug) # Subregion, meridional average, and remove segment time mean - d_seg_x_ano = get_daily_ano_segment(d_seg) + d_seg_x_ano = get_daily_ano_segment_xcdat(d_seg, data_var) # Compute space-time spectrum if debug: print("debug: compute space-time spectrum") - Power[n, :, :] = space_time_spectrum(d_seg_x_ano) + Power[n, :, :] = space_time_spectrum_xcdat(d_seg_x_ano, data_var) # Multi-year averaged power Power = np.average(Power, axis=0) + # Generates axes for the decoration - Power, ff, ss = generate_axes_and_decorate(Power, NT, NL) + Power = generate_axes_and_decorate_xcdat(Power, NT, NL) + # Output for wavenumber-frequency power spectra - OEE = output_power_spectra(NL, NT, Power, ff, ss) + OEE = output_power_spectra_xcdat(NL, NT, Power) + + if debug: + print("OEE:", OEE) + print("OEE.shape:", OEE.shape) # E/W ratio - ewr, eastPower, westPower = calculate_ewr(OEE) + ewr, eastPower, westPower = calculate_ewr_xcdat(OEE) + print("ewr: ", ewr) print("east power: ", eastPower) print("west power: ", westPower) @@ -196,14 +199,16 @@ def mjo_metric_ewr_calculation( if nc_out: os.makedirs(dir_paths["diagnostic_results"], exist_ok=True) fout = os.path.join(dir_paths["diagnostic_results"], output_filename) - write_netcdf_output(OEE, fout) + write_netcdf_output_xcdat(OEE, fout) # Plot if plot: os.makedirs(dir_paths["graphics"], exist_ok=True) fout = os.path.join(dir_paths["graphics"], output_filename) if model == "obs": - title = f"OBS ({run})\n{data_var.capitalize()}, {season} {startYear}-{endYear}" + title = ( + f"OBS ({run})\n{data_var.capitalize()}, {season} {startYear}-{endYear}" + ) else: title = f"{mip.upper()}: {model} ({run})\n{data_var.capitalize()}, {season} {startYear}-{endYear}" @@ -223,9 +228,13 @@ def mjo_metric_ewr_calculation( # Debug checking plot if debug and plot: debug_chk_plot( - d_seg_x_ano, Power, OEE, segment[year], daSeaCyc, segment_ano[year] + d_seg_x_ano, + Power, + OEE, + segment[year][data_var], + daSeaCyc, + segment_ano[year][data_var], ) - #f.close() ds.close() return metrics_result diff --git a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py index d60683fd3..8184fb237 100644 --- a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py +++ b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py @@ -8,8 +8,13 @@ def plot_power(d, title, fout, ewr=None): - y = d.getAxis(0)[:] - x = d.getAxis(1)[:] + # y = d.getAxis(0)[:] + # x = d.getAxis(1)[:] + + # y, x = d.indexes.values() + # x, y = d.indexes.values() + x = d["frequency"] + y = d["zonalwavenumber"] # adjust font size SMALL_SIZE = 8 From aad15b84a1bc651ed6585bee4051213a9534ad3c Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 1 May 2024 16:26:37 -0700 Subject: [PATCH 03/15] update and clean up --- pcmdi_metrics/mjo/lib/__init__.py | 27 +- pcmdi_metrics/mjo/lib/lib_mjo.py | 284 +----------------- pcmdi_metrics/mjo/lib/mjo_metric_calc.py | 42 ++- .../lib/plot_wavenumber_frequency_power.py | 5 - 4 files changed, 35 insertions(+), 323 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/__init__.py b/pcmdi_metrics/mjo/lib/__init__.py index 2ab62ad3c..c1587b174 100644 --- a/pcmdi_metrics/mjo/lib/__init__.py +++ b/pcmdi_metrics/mjo/lib/__init__.py @@ -2,27 +2,16 @@ from .debug_chk_plot import debug_chk_plot # noqa from .dict_merge import dict_merge # noqa from .lib_mjo import ( # noqa - Remove_dailySeasonalCycle, - # calculate_ewr, - calculate_ewr_xcdat, - decorate_2d_array_axes, - # generate_axes_and_decorate, - generate_axes_and_decorate_xcdat, - # get_daily_ano_segment, - get_daily_ano_segment_xcdat, - # interp2commonGrid, - interp2commonGrid_xcdat, + calculate_ewr, + generate_axes_and_decorate, + get_daily_ano_segment, + interp2commonGrid, mjo_metrics_to_json, - # output_power_spectra, - output_power_spectra_xcdat, - # space_time_spectrum, - space_time_spectrum_xcdat, - # subSliceSegment, - subSliceSegment_xcdat, + output_power_spectra, + space_time_spectrum, + subSliceSegment, taper, - unit_conversion, - # write_netcdf_output, - write_netcdf_output_xcdat, + write_netcdf_output, ) from .mjo_metric_calc import mjo_metric_ewr_calculation # noqa from .plot_wavenumber_frequency_power import plot_power # noqa diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index de4efff3b..6ebefdcac 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -9,37 +9,15 @@ from typing import Union -import cdms2 -import cdtime -import cdutil -import MV2 import numpy as np import xarray as xr from scipy import signal -import pcmdi_metrics -from pcmdi_metrics.io import get_time_key, select_subset +from pcmdi_metrics.io import base, get_time_key, select_subset from pcmdi_metrics.utils import create_target_grid, regrid -def interp2commonGrid(d, dlat, debug=False): - """ - input - - d: cdms array - - dlat: resolution (i.e. grid distance) in degree - output - - d2: d interpolated to dlat resolution grid - """ - nlat = int(180 / dlat) - grid = cdms2.createGaussianGrid(nlat, xorigin=0.0, order="yx") - d2 = d.regrid(grid, regridTool="regrid2", mkCyclic=True) - d2 = d2(latitude=(-10, 10)) - if debug: - print("debug: d2.shape:", d2.shape) - return d2 - - -def interp2commonGrid_xcdat(ds, data_var, dlat, dlon=None, debug=False): +def interp2commonGrid(ds, data_var, dlat, dlon=None, debug=False): if dlon is None: dlon = dlat # nlat = int(180 / dlat) @@ -54,29 +32,7 @@ def interp2commonGrid_xcdat(ds, data_var, dlat, dlon=None, debug=False): return ds_regrid_subset -def subSliceSegment(d, year, mon, day, length): - """ - Note: From given cdms array (3D: time and spatial 2D) - Subslice to get segment with given length starting from given time. - input - - d: cdms array - - year: segment starting year (integer) - - mon: segment starting month (integer) - - day: segement starting day (integer) - - length: segment length (integer) - """ - tim = d.getTime() - comTim = tim.asComponentTime() - h = comTim[0].hour - m = comTim[0].minute - s = comTim[0].second - cptime = cdtime.comptime(year, mon, day, h, m, s) # start date of segment - n = comTim.index(cptime) # time dimension index of above start date - d2 = d.subSlice((n, n + length)) # slie 180 time steps starting from above index - return d2 - - -def subSliceSegment_xcdat( +def subSliceSegment( ds: Union[xr.Dataset, xr.DataArray], year: int, mon: int, day: int, length: int ) -> Union[xr.Dataset, xr.DataArray]: """ @@ -100,47 +56,7 @@ def subSliceSegment_xcdat( ) # slie 180 time steps starting from above index -def Remove_dailySeasonalCycle(d, d_cyc): - """ - Note: Remove daily seasonal cycle - input - - d: cdms array - - d_cyc: numpy array - output - - d2: cdms array - """ - d2 = MV2.subtract(d, d_cyc) - # Preserve Axes - for i in range(len(d.shape)): - d2.setAxis(i, d.getAxis(i)) - # Preserve variable id (How to preserve others?) - d2.id = d.id - return d2 - - -def get_daily_ano_segment(d_seg): - """ - Note: 1. Get daily time series (3D: time and spatial 2D) - 2. Meridionally average (2D: time and spatial, i.e., longitude) - 3. Get anomaly by removing time mean of the segment - input - - d_seg: cdms2 data - output - - d_seg_x_ano: 2d array - """ - cdms2.setAutoBounds("on") - # sub region - d_seg = d_seg(latitude=(-10, 10)) - # Get meridional average (3d (t, y, x) to 2d (t, y)) - d_seg_x = cdutil.averager(d_seg, axis="y", weights="weighted") - # Get time-average in the segment on each longitude grid - d_seg_x_ave = cdutil.averager(d_seg_x, axis="t") - # Remove time mean for each segment - d_seg_x_ano = MV2.subtract(d_seg_x, d_seg_x_ave) - return d_seg_x_ano - - -def get_daily_ano_segment_xcdat(d_seg, data_var): +def get_daily_ano_segment(d_seg, data_var): """ Note: 1. Get daily time series (3D: time and spatial 2D) 2. Meridionally average (2D: time and spatial, i.e., longitude) @@ -167,43 +83,7 @@ def get_daily_ano_segment_xcdat(d_seg, data_var): return d_seg_x_ano -def space_time_spectrum(d_seg_x_ano): - """ - input - - d: 2d cdms MV2 array (t (time), n (space)) - output - - p: 2d array for power - NOTE: Below code taken from - https://github.com/CDAT/wk/blob/2b953281c7a4c5d0ac2d79fcc3523113e31613d5/WK/process.py#L188 - """ - # Number of grid in longitude axis, and timestep for each segment - NTSub = d_seg_x_ano.shape[0] # NTSub - NL = d_seg_x_ano.shape[1] # NL - # Tapering - d_seg_x_ano = taper(d_seg_x_ano) - # Power sepctrum analysis - EE = np.fft.fft2(d_seg_x_ano, axes=(1, 0)) / float(NL) / float(NTSub) - # Now the array EE(n,t) contains the (complex) space-time spectrum. - """ - Create array PEE(NL+1,NT/2+1) which contains the (real) power spectrum. - Note how the PEE array is arranged into a different order to EE. - In this code, PEE is "Power", and its multiyear average will be "power" - """ - # OK NOW THE LITTLE MAGIC WITH REORDERING ! - A = np.absolute(EE[0 : NTSub // 2 + 1, 1 : NL // 2 + 1]) ** 2 - B = np.absolute(EE[NTSub // 2 : NTSub, 1 : NL // 2 + 1]) ** 2 - C = np.absolute(EE[NTSub // 2 : NTSub, 0 : NL // 2 + 1]) ** 2 - D = np.absolute(EE[0 : NTSub // 2 + 1, 0 : NL // 2 + 1]) ** 2 - # Define returning array - p = np.zeros((NTSub + 1, NL + 1), np.float) - p[NTSub // 2 :, : NL // 2] = A[:, ::-1] - p[: NTSub // 2, : NL // 2] = B[:, ::-1] - p[NTSub // 2 + 1 :, NL // 2 :] = C[::-1, :] - p[: NTSub // 2 + 1, NL // 2 :] = D[::-1, :] - return p - - -def space_time_spectrum_xcdat(d_seg_x_ano, data_var): +def space_time_spectrum(d_seg_x_ano, data_var): """ input - d: 2d cdms MV2 array (t (time), n (space)) @@ -250,81 +130,11 @@ def taper(data): window = signal.windows.tukey(len(data)) data2 = data.copy() for i in range(0, len(data)): - data2[i] = MV2.multiply(data[i][:], window[i]) + data2[i] = np.multiply(data[i][:], window[i]) return data2 -def decorate_2d_array_axes( - a, y, x, a_id=None, y_id=None, x_id=None, y_units=None, x_units=None -): - """ - Note: Decorate array with given axes - input - - a: 2d cdms MV2 or numpy array to decorate axes - - y: list of numbers to be axis 0 - - x: list of numbers to be axis 1 - - a_id: id of variable a - - y_id, x_id: id of axes, string - - y_units, x_units: units of axes - output - - return the given array, a, with axes attached - """ - y = MV2.array(y) - x = MV2.array(x) - # Create the frequencies axis - Y = cdms2.createAxis(y) - Y.id = y_id - Y.units = y_units - # Create the wave numbers axis - X = cdms2.createAxis(x) - X.id = x_id - X.units = x_units - # Makes it an MV2 with axis and id (id come sfrom orignal data id) - a = MV2.array(a, axes=(Y, X), id=a_id) - return a - - def generate_axes_and_decorate(Power, NT, NL): - """ - Note: Generates axes for the decoration - input - - Power: 2d numpy array - - NT: integer, number of time step - - NL: integer, number of spatial grid - output - - Power: decorated 2d cdms array - - ff: frequency axis - - ss: wavenumber axis - """ - # frequency - ff = [] - for t in range(0, NT + 1): - ff.append(float(t - NT / 2) / float(NT)) - ff = MV2.array(ff) - ff.id = "frequency" - ff.units = "cycles per day" - # wave number - ss = [] - for n in range(0, NL + 1): - ss.append(float(n) - float(NL / 2)) - ss = MV2.array(ss) - ss.id = "zonalwavenumber" - ss.units = "-" - # Decoration - Power = decorate_2d_array_axes( - Power, - ff, - ss, - a_id="power", - y_id=ff.id, - x_id=ss.id, - y_units=ff.units, - x_units=ss.units, - ) - return Power, ff, ss - - -def generate_axes_and_decorate_xcdat(Power, NT, NL): """ Note: Generates axes for the decoration input @@ -367,40 +177,7 @@ def generate_axes_and_decorate_xcdat(Power, NT, NL): return da -def output_power_spectra(NL, NT, Power, ff, ss): - """ - Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) - """ - # The corresponding frequencies, ff, and wavenumbers, ss, are:- - PEE = Power - OEE = np.zeros((21, 11)) - for n in range(int(NL / 2), int(NL / 2) + 1 + 10): - nn = n - int(NL / 2) - for t in range(int(NT / 2) - 10, int(NT / 2 + 1 + 10)): - tt = -(int(NT / 2) + 1) + 11 + t - OEE[tt, nn] = PEE[t, n] - a = list((ff[i] for i in range(int(NT / 2) - 10, int(NT / 2) + 1 + 10))) - b = list((ss[i] for i in range(int(NL / 2), int(NL / 2) + 1 + 10))) - a = MV2.array(a) - b = MV2.array(b) - # Decoration - OEE = decorate_2d_array_axes( - OEE, - a, - b, - a_id="power", - y_id=ff.id, - x_id=ss.id, - y_units=ff.units, - x_units=ss.units, - ) - # Transpose for visualization - OEE = MV2.transpose(OEE, (1, 0)) - OEE.id = "power" - return OEE - - -def output_power_spectra_xcdat(NL, NT, Power): +def output_power_spectra(NL, NT, Power): """ Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) """ @@ -448,19 +225,7 @@ def output_power_spectra_xcdat(NL, NT, Power): # return OEE -def write_netcdf_output(d, fname): - """ - Note: write array in a netcdf file - input - - d: array - - fname: string. directory path and name of the netcd file, without .nc - """ - fo = cdms2.open(fname + ".nc", "w") - fo.write(d) - fo.close() - - -def write_netcdf_output_xcdat(da, fname): +def write_netcdf_output(da, fname): """ Note: write array in a netcdf file input @@ -472,22 +237,6 @@ def write_netcdf_output_xcdat(da, fname): def calculate_ewr(OEE): - """ - According to DK's gs script (MSD/level_2/sample/stps/e.w.ratio.gs.sample), - E/W ratio is calculated as below: - 'd amean(power,x=14,x=17,y=2,y=4)/aave(power,x=5,x=8,y=2,y=4)' - where x for frequency and y for wavenumber. - Actual ranges of frequency and wavenumber have been checked and applied. - """ - east_power_domain = OEE(zonalwavenumber=(1, 3), frequency=(0.0166667, 0.0333333)) - west_power_domain = OEE(zonalwavenumber=(1, 3), frequency=(-0.0333333, -0.0166667)) - eastPower = np.average(np.array(east_power_domain)) - westPower = np.average(np.array(west_power_domain)) - ewr = eastPower / westPower - return ewr, eastPower, westPower - - -def calculate_ewr_xcdat(OEE): """ According to DK's gs script (MSD/level_2/sample/stps/e.w.ratio.gs.sample), E/W ratio is calculated as below: @@ -507,26 +256,11 @@ def calculate_ewr_xcdat(OEE): return ewr, eastPower, westPower -def unit_conversion(data, UnitsAdjust): - """ - Convert unit following given tuple using MV2 - input: - - data: cdms array - - UnitsAdjust: tuple with 4 elements - e.g.: (True, 'multiply', 86400., 'mm d-1'): e.g., kg m-2 s-1 to mm d-1 - (False, 0, 0, 0): no unit conversion - """ - if UnitsAdjust[0]: - data = getattr(MV2, UnitsAdjust[1])(data, UnitsAdjust[2]) - data.units = UnitsAdjust[3] - return data - - def mjo_metrics_to_json( outdir, json_filename, result_dict, model=None, run=None, cmec_flag=False ): # Open JSON - JSON = pcmdi_metrics.io.base.Base(outdir, json_filename) + JSON = base.Base(outdir, json_filename) # Dict for JSON if model is None and run is None: result_dict_to_json = result_dict diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py index c062f9dd3..01aa91c29 100644 --- a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py +++ b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py @@ -1,22 +1,19 @@ import os from datetime import datetime -# import cdms2 -# import cdtime -# import MV2 import numpy as np import xarray as xr from pcmdi_metrics.io import get_latitude, get_longitude, get_time_key, xcdat_open -from pcmdi_metrics.mjo.lib import ( # calculate_ewr,; generate_axes_and_decorate,; get_daily_ano_segment,; interp2commonGrid,; output_power_spectra,; space_time_spectrum,; subSliceSegment,; unit_conversion,; Remove_dailySeasonalCycle,; write_netcdf_output, - calculate_ewr_xcdat, - generate_axes_and_decorate_xcdat, - get_daily_ano_segment_xcdat, - interp2commonGrid_xcdat, - output_power_spectra_xcdat, - space_time_spectrum_xcdat, - subSliceSegment_xcdat, - write_netcdf_output_xcdat, +from pcmdi_metrics.mjo.lib import ( + calculate_ewr, + generate_axes_and_decorate, + get_daily_ano_segment, + interp2commonGrid, + output_power_spectra, + space_time_spectrum, + subSliceSegment, + write_netcdf_output, ) from pcmdi_metrics.utils import adjust_units @@ -106,7 +103,7 @@ def mjo_metric_ewr_calculation( # Loop over years for year in range(startYear, endYear): print(year) - segment[year] = subSliceSegment_xcdat(ds, year, mon, day, NT) + segment[year] = subSliceSegment(ds, year, mon, day, NT) # units conversion segment[year][data_var] = adjust_units(segment[year][data_var], UnitsAdjust) if debug: @@ -116,9 +113,6 @@ def mjo_metric_ewr_calculation( segment[year][data_var].shape, ) # Get climatology of daily seasonal cycle - # daSeaCyc_values = np.add( - # np.divide(segment[year][data_var].values, float(numYear)), daSeaCyc_values - # ) daSeaCyc_values = ( segment[year][data_var].values / float(numYear) ) + daSeaCyc_values @@ -132,7 +126,7 @@ def mjo_metric_ewr_calculation( if numYear > 1: # Loop over years for year in range(startYear, endYear): - # segment_ano[year] = Remove_dailySeasonalCycle(segment[year], daSeaCyc) + # Remove daily Seasonal Cycle segment_ano[year] = segment[year].copy() segment_ano[year][data_var].values = ( segment[year][data_var].values - daSeaCyc.values @@ -162,29 +156,29 @@ def mjo_metric_ewr_calculation( d_seg = segment_ano[year] # Regrid: interpolation to common grid if cmmGrid: - d_seg = interp2commonGrid_xcdat(d_seg, data_var, degX, debug=debug) + d_seg = interp2commonGrid(d_seg, data_var, degX, debug=debug) # Subregion, meridional average, and remove segment time mean - d_seg_x_ano = get_daily_ano_segment_xcdat(d_seg, data_var) + d_seg_x_ano = get_daily_ano_segment(d_seg, data_var) # Compute space-time spectrum if debug: print("debug: compute space-time spectrum") - Power[n, :, :] = space_time_spectrum_xcdat(d_seg_x_ano, data_var) + Power[n, :, :] = space_time_spectrum(d_seg_x_ano, data_var) # Multi-year averaged power Power = np.average(Power, axis=0) # Generates axes for the decoration - Power = generate_axes_and_decorate_xcdat(Power, NT, NL) + Power = generate_axes_and_decorate(Power, NT, NL) # Output for wavenumber-frequency power spectra - OEE = output_power_spectra_xcdat(NL, NT, Power) + OEE = output_power_spectra(NL, NT, Power) if debug: print("OEE:", OEE) print("OEE.shape:", OEE.shape) # E/W ratio - ewr, eastPower, westPower = calculate_ewr_xcdat(OEE) + ewr, eastPower, westPower = calculate_ewr(OEE) print("ewr: ", ewr) print("east power: ", eastPower) @@ -199,7 +193,7 @@ def mjo_metric_ewr_calculation( if nc_out: os.makedirs(dir_paths["diagnostic_results"], exist_ok=True) fout = os.path.join(dir_paths["diagnostic_results"], output_filename) - write_netcdf_output_xcdat(OEE, fout) + write_netcdf_output(OEE, fout) # Plot if plot: diff --git a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py index 8184fb237..b71353c4a 100644 --- a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py +++ b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py @@ -8,11 +8,6 @@ def plot_power(d, title, fout, ewr=None): - # y = d.getAxis(0)[:] - # x = d.getAxis(1)[:] - - # y, x = d.indexes.values() - # x, y = d.indexes.values() x = d["frequency"] y = d["zonalwavenumber"] From ebc9b506a033c178c9cea9146b7aa217f26d0ca4 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 1 May 2024 16:27:42 -0700 Subject: [PATCH 04/15] update and clean up --- pcmdi_metrics/mjo/{lib => scripts}/dict_merge.py | 0 pcmdi_metrics/mjo/{lib => scripts}/post_process_plot.py | 0 .../mjo/{lib => scripts}/post_process_plot_ensemble_mean.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename pcmdi_metrics/mjo/{lib => scripts}/dict_merge.py (100%) rename pcmdi_metrics/mjo/{lib => scripts}/post_process_plot.py (100%) rename pcmdi_metrics/mjo/{lib => scripts}/post_process_plot_ensemble_mean.py (100%) diff --git a/pcmdi_metrics/mjo/lib/dict_merge.py b/pcmdi_metrics/mjo/scripts/dict_merge.py similarity index 100% rename from pcmdi_metrics/mjo/lib/dict_merge.py rename to pcmdi_metrics/mjo/scripts/dict_merge.py diff --git a/pcmdi_metrics/mjo/lib/post_process_plot.py b/pcmdi_metrics/mjo/scripts/post_process_plot.py similarity index 100% rename from pcmdi_metrics/mjo/lib/post_process_plot.py rename to pcmdi_metrics/mjo/scripts/post_process_plot.py diff --git a/pcmdi_metrics/mjo/lib/post_process_plot_ensemble_mean.py b/pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py similarity index 100% rename from pcmdi_metrics/mjo/lib/post_process_plot_ensemble_mean.py rename to pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py From 953f1a5b20b2b5d5e264df0f11a8db8264564aaf Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Thu, 2 May 2024 09:40:40 -0700 Subject: [PATCH 05/15] clean up --- pcmdi_metrics/mjo/lib/lib_mjo.py | 32 +++++++++++-------- .../lib/plot_wavenumber_frequency_power.py | 9 +++--- .../mjo/scripts/post_process_plot.py | 8 ++--- .../post_process_plot_ensemble_mean.py | 17 ++++++---- 4 files changed, 37 insertions(+), 29 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 6ebefdcac..3420a36f0 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -2,6 +2,9 @@ Code written by Jiwoo Lee, LLNL. Feb. 2019 Inspired by Daehyun Kim and Min-Seop Ahn's MJO metrics. +Code update history +2024-05 converted to use xcdat as base building block (Jiwoo Lee) + Reference: Ahn, MS., Kim, D., Sperber, K.R. et al. Clim Dyn (2017) 49: 4023. https://doi.org/10.1007/s00382-017-3558-4 @@ -36,7 +39,7 @@ def subSliceSegment( ds: Union[xr.Dataset, xr.DataArray], year: int, mon: int, day: int, length: int ) -> Union[xr.Dataset, xr.DataArray]: """ - Note: From given cdms array (3D: time and spatial 2D) + Note: From given array (3D: time and spatial 2D) Subslice to get segment with given length starting from given time. input - ds: xarray dataset or dataArray @@ -56,7 +59,7 @@ def subSliceSegment( ) # slie 180 time steps starting from above index -def get_daily_ano_segment(d_seg, data_var): +def get_daily_ano_segment(d_seg: xr.Dataset, data_var: str) -> xr.Dataset: """ Note: 1. Get daily time series (3D: time and spatial 2D) 2. Meridionally average (2D: time and spatial, i.e., longitude) @@ -83,12 +86,13 @@ def get_daily_ano_segment(d_seg, data_var): return d_seg_x_ano -def space_time_spectrum(d_seg_x_ano, data_var): +def space_time_spectrum(d_seg_x_ano: xr.Dataset, data_var: str) -> np.ndarray: """ input - - d: 2d cdms MV2 array (t (time), n (space)) + - d: xarray dataset that contains 2d DataArray (t (time), n (space)) named as `data_var` + - data_var: name of the 2d DataArray output - - p: 2d array for power + - p: 2d numpy array for power NOTE: Below code taken from https://github.com/CDAT/wk/blob/2b953281c7a4c5d0ac2d79fcc3523113e31613d5/WK/process.py#L188 """ @@ -123,7 +127,7 @@ def taper(data): """ Note: taper first and last 45 days with cosine window, using scipy.signal function input - - data: cdms 2d array (t, n) t: time, n: space (meridionally averaged) + - data: 2d array (t, n) t: time, n: space (meridionally averaged) output: - data: tapered data """ @@ -134,7 +138,7 @@ def taper(data): return data2 -def generate_axes_and_decorate(Power, NT, NL): +def generate_axes_and_decorate(Power, NT: int, NL: int) -> xr.DataArray: """ Note: Generates axes for the decoration input @@ -142,9 +146,7 @@ def generate_axes_and_decorate(Power, NT, NL): - NT: integer, number of time step - NL: integer, number of spatial grid output - - Power: decorated 2d cdms array - - ff: frequency axis - - ss: wavenumber axis + - xr.DataArray that contains Power 2d DataArray that has frequency and zonalwavenumber axes """ # frequency ff = [] @@ -177,7 +179,7 @@ def generate_axes_and_decorate(Power, NT, NL): return da -def output_power_spectra(NL, NT, Power): +def output_power_spectra(NL: int, NT: int, Power): """ Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) """ @@ -225,12 +227,14 @@ def output_power_spectra(NL, NT, Power): # return OEE -def write_netcdf_output(da, fname): +def write_netcdf_output(da: xr.DataArray, fname): """ Note: write array in a netcdf file input - - d: array - - fname: string. directory path and name of the netcd file, without .nc + - d: xr.DataArray object + - fname: string of filename. Directory path that includes file name without .nc + output + - None """ ds = xr.Dataset({da.name: da}) ds.to_netcdf(fname + ".nc") diff --git a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py index b71353c4a..e1e11397a 100644 --- a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py +++ b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py @@ -1,13 +1,13 @@ import copy import os -import cdms2 import matplotlib.cm import matplotlib.pyplot as plt +import xarray as xr from matplotlib.patches import Rectangle -def plot_power(d, title, fout, ewr=None): +def plot_power(d: xr.DataArray, title: str, fout: str, ewr=None): x = d["frequency"] y = d["zonalwavenumber"] @@ -132,8 +132,9 @@ def plot_power(d, title, fout, ewr=None): imgdir = "." - f = cdms2.open(os.path.join(datadir, ncfile)) - d = f("power") + ds = xr.open_dataset(os.path.join(datadir, ncfile)) + d = ds["power"] + fout = os.path.join(imgdir, pngfilename) plot_power(d, title, fout, ewr=ewr) diff --git a/pcmdi_metrics/mjo/scripts/post_process_plot.py b/pcmdi_metrics/mjo/scripts/post_process_plot.py index bea7ca873..44aced650 100644 --- a/pcmdi_metrics/mjo/scripts/post_process_plot.py +++ b/pcmdi_metrics/mjo/scripts/post_process_plot.py @@ -1,7 +1,7 @@ import glob import os -import cdms2 +import xarray as xr from lib_mjo import calculate_ewr from plot_wavenumber_frequency_power import plot_power @@ -48,10 +48,9 @@ def main(): ncfile = ( "_".join([mip, model, exp, run, "mjo", period, "cmmGrid"]) + ".nc" ) - f = cdms2.open(os.path.join(datadir, ncfile)) - d = f("power") + ds = xr.open_dataset(os.path.join(datadir, ncfile)) + d = ds["power"] d_runs.append(d) - f.close() title = ( mip.upper() + ": " @@ -69,6 +68,7 @@ def main(): fout = os.path.join(imgdir, pngfilename) # plot plot_power(d, title, fout, ewr) + ds.close() except Exception: print(model, run, "cannnot load") pass diff --git a/pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py b/pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py index 83f9012c0..df9b432bd 100644 --- a/pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py +++ b/pcmdi_metrics/mjo/scripts/post_process_plot_ensemble_mean.py @@ -1,8 +1,8 @@ import glob import os -import cdms2 -import MV2 +import numpy as np +import xarray as xr from lib_mjo import calculate_ewr from plot_wavenumber_frequency_power import plot_power @@ -62,18 +62,21 @@ def main(): ) + ".nc" ) - f = cdms2.open(os.path.join(datadir, ncfile)) - d = f("power") + + ds = xr.open_dataset(os.path.join(datadir, ncfile)) + d = ds["power"] + d_runs.append(d) - f.close() + except Exception as err: print(model, run, "cannnot load:", err) pass + if run == runs_list[-1]: num_runs = len(d_runs) # ensemble mean - d_avg = MV2.average(d_runs, axis=0) - d_avg.setAxisList(d.getAxisList()) + d_avg = np.average(d_runs, axis=0) + # d_avg.setAxisList(d.getAxisList()) title = ( mip.upper() + ": " From a56b94ce09d1ca4c0fa2a6847b5ceaaf55771898 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Thu, 2 May 2024 09:46:39 -0700 Subject: [PATCH 06/15] pre-commit fix --- pcmdi_metrics/utils/__init__.py | 1 - pcmdi_metrics/utils/adjust_units.py | 27 +++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 pcmdi_metrics/utils/adjust_units.py diff --git a/pcmdi_metrics/utils/__init__.py b/pcmdi_metrics/utils/__init__.py index 10c2d6e31..dc4a935f0 100644 --- a/pcmdi_metrics/utils/__init__.py +++ b/pcmdi_metrics/utils/__init__.py @@ -1,5 +1,4 @@ from .adjust_units import adjust_units - from .custom_season import ( custom_season_average, custom_season_departure, diff --git a/pcmdi_metrics/utils/adjust_units.py b/pcmdi_metrics/utils/adjust_units.py new file mode 100644 index 000000000..ee88d3d4a --- /dev/null +++ b/pcmdi_metrics/utils/adjust_units.py @@ -0,0 +1,27 @@ +import xarray as xr + + +def adjust_units(da: xr.DataArray, adjust_tuple: tuple) -> xr.DataArray: + """Convert unit following information in the given tuple + + Parameters + ---------- + da : xr.DataArray + input data array + adjust_tuple : tuple with at least 3 elements (4th element is optional for units) + e.g.: (True, 'multiply', 86400., 'mm d-1'): e.g., kg m-2 s-1 to mm d-1 + (False, 0, 0, 0): no unit conversion + + Returns + ------- + xr.DataArray + data array that contains converted values and attributes + """ + action_dict = {"multiply": "*", "divide": "/", "add": "+", "subtract": "-"} + if adjust_tuple[0]: + print("Converting units by ", adjust_tuple[1], adjust_tuple[2]) + cmd = " ".join(["da", str(action_dict[adjust_tuple[1]]), str(adjust_tuple[2])]) + da = eval(cmd) + if len(adjust_tuple) > 3: + da.assign_attrs(units=adjust_tuple[3]) + return da From 087e30e85284a973faa98adeb39628f4fe055dda Mon Sep 17 00:00:00 2001 From: lee1043 Date: Thu, 2 May 2024 14:27:33 -0700 Subject: [PATCH 07/15] move dict_merge script to lib --- pcmdi_metrics/mjo/{scripts => lib}/dict_merge.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename pcmdi_metrics/mjo/{scripts => lib}/dict_merge.py (100%) diff --git a/pcmdi_metrics/mjo/scripts/dict_merge.py b/pcmdi_metrics/mjo/lib/dict_merge.py similarity index 100% rename from pcmdi_metrics/mjo/scripts/dict_merge.py rename to pcmdi_metrics/mjo/lib/dict_merge.py From 493076af8f0cd5a7bd3376b51c8ed791eee1e1c0 Mon Sep 17 00:00:00 2001 From: lee1043 Date: Wed, 8 May 2024 15:53:00 -0700 Subject: [PATCH 08/15] use more straitforward code because otherwise error occurs with non-standard calendar --- pcmdi_metrics/mjo/lib/mjo_metric_calc.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py index 01aa91c29..8ad7a458e 100644 --- a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py +++ b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py @@ -54,8 +54,18 @@ def mjo_metric_ewr_calculation( print("debug: check time") time_key = get_time_key(ds) - first_time = ds.indexes[time_key].to_datetimeindex()[0].to_pydatetime() - last_time = ds.indexes[time_key].to_datetimeindex()[-1].to_pydatetime() + + # Get first time step date + first_time_year = ds[time_key][0].item().year + first_time_month = ds[time_key][0].item().month + first_time_day = ds[time_key][0].item().day + first_time = datetime(first_time_year, first_time_month, first_time_day) + + # Get last time step date + last_time_year = ds[time_key][-1].item().year + last_time_month = ds[time_key][-1].item().month + last_time_day = ds[time_key][-1].item().day + last_time = datetime(last_time_year, last_time_month, last_time_day) if season == "NDJFMA": # Adjust years to consider only when continuous NDJFMA is available From 6a88bd3cb6cb3167128c8441d87c78b6c7cfcfb7 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 8 May 2024 17:55:12 -0700 Subject: [PATCH 09/15] use gaussian grid as common grid to be more consistent with the CDMS/CDAT version of the code --- pcmdi_metrics/mjo/lib/lib_mjo.py | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 3420a36f0..53786e60e 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -14,24 +14,43 @@ import numpy as np import xarray as xr +import xcdat as xc from scipy import signal from pcmdi_metrics.io import base, get_time_key, select_subset -from pcmdi_metrics.utils import create_target_grid, regrid +from pcmdi_metrics.utils import regrid + +# from pcmdi_metrics.utils import create_target_grid def interp2commonGrid(ds, data_var, dlat, dlon=None, debug=False): if dlon is None: dlon = dlat - # nlat = int(180 / dlat) - # nlon = int(360 / dlon) - grid = create_target_grid(target_grid_resolution=f"{dlat}x{dlon}") + + # grid = create_target_grid(target_grid_resolution=f"{dlat}x{dlon}") + nlat = int(180 / dlat) + grid = xc.create_gaussian_grid(nlat) + + # If the longitude values include 0 and 360, then remove 360 to avoid having repeating grid + if 0 in grid.lon.values and 360 in grid.lon.values: + min_lon = grid.lon.values[0] # 0 + # max_lon = grid.lon.values[-1] # 360 + second_max_lon = grid.lon.values[-2] # 360-dlat + grid = grid.sel(lon=slice(min_lon, second_max_lon)) + + # Reverse latitude if needed + if grid.lat.values[0] > grid.lat.values[-1]: + grid = grid.isel(lat=slice(None, None, -1)) + + # Regrid ds_regrid = regrid(ds, data_var, grid) ds_regrid_subset = select_subset(ds_regrid, lat=(-10, 10)) + if debug: print( "debug: ds_regrid_subset[data_var] shape:", ds_regrid_subset[data_var].shape ) + return ds_regrid_subset From 550f4d837f59c5857283aab09e1320f0ad78b439 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 8 May 2024 18:26:46 -0700 Subject: [PATCH 10/15] clean up; potential bug fix that could happen when segment length is other than 180 --- pcmdi_metrics/mjo/lib/mjo_metric_calc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py index 8ad7a458e..0b38f28b6 100644 --- a/pcmdi_metrics/mjo/lib/mjo_metric_calc.py +++ b/pcmdi_metrics/mjo/lib/mjo_metric_calc.py @@ -94,6 +94,7 @@ def mjo_metric_ewr_calculation( elif season == "MJJASO": mon = 5 numYear = endYear - startYear + 1 + day = 1 # Store each year's segment in a dictionary: segment[year] @@ -103,7 +104,7 @@ def mjo_metric_ewr_calculation( daSeaCyc = xr.DataArray( np.zeros((NT, ds[data_var].shape[1], ds[data_var].shape[2])), dims=["day", "lat", "lon"], - coords={"day": np.arange(180), "lat": lat, "lon": lon}, + coords={"day": np.arange(NT), "lat": lat, "lon": lon}, ) daSeaCyc_values = daSeaCyc.values.copy() From 84f5730b0a0980941c438254b210fd3a7f966a9c Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 8 May 2024 18:46:39 -0700 Subject: [PATCH 11/15] add gaussian grid option to utils function and refer to it --- pcmdi_metrics/mjo/lib/lib_mjo.py | 23 +++++-------------- pcmdi_metrics/utils/grid.py | 38 +++++++++++++++++++++++++++----- 2 files changed, 37 insertions(+), 24 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 53786e60e..1196f237f 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -14,33 +14,20 @@ import numpy as np import xarray as xr -import xcdat as xc from scipy import signal from pcmdi_metrics.io import base, get_time_key, select_subset -from pcmdi_metrics.utils import regrid - -# from pcmdi_metrics.utils import create_target_grid +from pcmdi_metrics.utils import create_target_grid, regrid def interp2commonGrid(ds, data_var, dlat, dlon=None, debug=False): if dlon is None: dlon = dlat - # grid = create_target_grid(target_grid_resolution=f"{dlat}x{dlon}") - nlat = int(180 / dlat) - grid = xc.create_gaussian_grid(nlat) - - # If the longitude values include 0 and 360, then remove 360 to avoid having repeating grid - if 0 in grid.lon.values and 360 in grid.lon.values: - min_lon = grid.lon.values[0] # 0 - # max_lon = grid.lon.values[-1] # 360 - second_max_lon = grid.lon.values[-2] # 360-dlat - grid = grid.sel(lon=slice(min_lon, second_max_lon)) - - # Reverse latitude if needed - if grid.lat.values[0] > grid.lat.values[-1]: - grid = grid.isel(lat=slice(None, None, -1)) + # Generate grid + grid = create_target_grid( + target_grid_resolution=f"{dlat}x{dlon}", grid_type="uniform" + ) # Regrid ds_regrid = regrid(ds, data_var, grid) diff --git a/pcmdi_metrics/utils/grid.py b/pcmdi_metrics/utils/grid.py index 4de4d677a..2184f33e8 100644 --- a/pcmdi_metrics/utils/grid.py +++ b/pcmdi_metrics/utils/grid.py @@ -17,6 +17,7 @@ def create_target_grid( lon1: float = 0.0, lon2: float = 360.0, target_grid_resolution: str = "2.5x2.5", + grid_type: str = "uniform", ) -> xr.Dataset: """Generate a uniform grid for given latitude/longitude ranges and resolution @@ -32,6 +33,8 @@ def create_target_grid( Starting latitude, by default 360. target_grid_resolution : str, optional grid resolution in degree for lat and lon, by default "2.5x2.5" + grid_type : str, optional + type of the grid ('uniform' or 'gaussian'), by default "uniform" Returns ------- @@ -46,11 +49,11 @@ def create_target_grid( Global uniform grid: - >>> t_grid = create_target_grid(-90, 90, 0, 360, target_grid="5x5") + >>> grid = create_target_grid(-90, 90, 0, 360, target_grid="5x5") Regional uniform grid: - >>> t_grid = create_target_grid(30, 50, 100, 150, target_grid="0.5x0.5") + >>> grid = create_target_grid(30, 50, 100, 150, target_grid="0.5x0.5") """ # generate target grid res = target_grid_resolution.split("x") @@ -60,10 +63,33 @@ def create_target_grid( start_lon = lon1 + lon_res / 2.0 end_lat = lat2 - lat_res / 2 end_lon = lon2 - lon_res / 2 - t_grid = xc.create_uniform_grid( - start_lat, end_lat, lat_res, start_lon, end_lon, lon_res - ) - return t_grid + + if grid_type == "uniform": + grid = xc.create_uniform_grid( + start_lat, end_lat, lat_res, start_lon, end_lon, lon_res + ) + elif grid_type == "gaussian": + nlat = int(180 / lat_res) + grid = xc.create_gaussian_grid(nlat) + + # If the longitude values include 0 and 360, then remove 360 to avoid having repeating grid + if 0 in grid.lon.values and 360 in grid.lon.values: + min_lon = grid.lon.values[0] # 0 + # max_lon = grid.lon.values[-1] # 360 + second_max_lon = grid.lon.values[-2] # 360-dlat + grid = grid.sel(lon=slice(min_lon, second_max_lon)) + + # Reverse latitude if needed + if grid.lat.values[0] > grid.lat.values[-1]: + grid = grid.isel(lat=slice(None, None, -1)) + + grid = grid.sel(lat=slice(start_lat, end_lat), lon=slice(start_lon, end_lon)) + else: + raise ValueError( + f"grid_type {grid_type} is undefined. Please use either `uniform` or `gaussian" + ) + + return grid def __haversine(lat1, lon1, lat2, lon2): From 5b61d85540f4404ce7ca3b23efcba4e8bd037335 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 8 May 2024 18:51:08 -0700 Subject: [PATCH 12/15] clean up --- pcmdi_metrics/utils/grid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pcmdi_metrics/utils/grid.py b/pcmdi_metrics/utils/grid.py index 2184f33e8..968cbd055 100644 --- a/pcmdi_metrics/utils/grid.py +++ b/pcmdi_metrics/utils/grid.py @@ -86,7 +86,7 @@ def create_target_grid( grid = grid.sel(lat=slice(start_lat, end_lat), lon=slice(start_lon, end_lon)) else: raise ValueError( - f"grid_type {grid_type} is undefined. Please use either `uniform` or `gaussian" + f"grid_type {grid_type} is undefined. Please use either 'uniform' or 'gaussian'" ) return grid From f45530781b579664c3e98c3df9171213b71e9233 Mon Sep 17 00:00:00 2001 From: lee1043 Date: Tue, 14 May 2024 14:04:35 -0700 Subject: [PATCH 13/15] adjust criteria for frequency domain selection to make the selected frequency domain to be consistent with that of cdms version --- pcmdi_metrics/mjo/lib/lib_mjo.py | 4 ++-- pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 1196f237f..07a389b0b 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -255,10 +255,10 @@ def calculate_ewr(OEE): Actual ranges of frequency and wavenumber have been checked and applied. """ east_power_domain = OEE.sel( - zonalwavenumber=slice(1, 3), frequency=slice(0.0166667, 0.0333333) + zonalwavenumber=slice(1, 3), frequency=slice(0.016, 0.034) ) west_power_domain = OEE.sel( - zonalwavenumber=slice(1, 3), frequency=slice(-0.0333333, -0.0166667) + zonalwavenumber=slice(1, 3), frequency=slice(-0.034, -0.016) ) eastPower = np.average(east_power_domain) westPower = np.average(west_power_domain) diff --git a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py index e1e11397a..564021507 100644 --- a/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py +++ b/pcmdi_metrics/mjo/lib/plot_wavenumber_frequency_power.py @@ -87,8 +87,8 @@ def plot_power(d: xr.DataArray, title: str, fout: str, ewr=None): currentAxis = plt.gca() currentAxis.add_patch( Rectangle( - (0.0166667, 1), - 0.0333333 - 0.0166667, + (0.016, 1), + 0.034 - 0.016, 2, edgecolor="black", ls="--", @@ -97,8 +97,8 @@ def plot_power(d: xr.DataArray, title: str, fout: str, ewr=None): ) currentAxis.add_patch( Rectangle( - (-0.0333333, 1), - 0.0333333 - 0.0166667, + (-0.034, 1), + 0.034 - 0.016, 2, edgecolor="black", ls="--", From 2e95b742180f2709011e658291a7bd8d1670a49e Mon Sep 17 00:00:00 2001 From: lee1043 Date: Tue, 14 May 2024 22:28:25 -0700 Subject: [PATCH 14/15] cleaned up --- pcmdi_metrics/mjo/lib/lib_mjo.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 07a389b0b..721c8c661 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -185,7 +185,7 @@ def generate_axes_and_decorate(Power, NT: int, NL: int) -> xr.DataArray: return da -def output_power_spectra(NL: int, NT: int, Power): +def output_power_spectra(NL: int, NT: int, Power, debug: bool = False): """ Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) """ @@ -224,10 +224,13 @@ def output_power_spectra(NL: int, NT: int, Power): ) # Transpose for visualization - # OEE = np.transpose(OEE, (1, 0)) - print("before transpose, OEE.shape:", OEE.shape) + if debug: + print("before transpose, OEE.shape:", OEE.shape) + transposed_OEE = OEE.transpose() - print("after transpose, transposed_OEE.shape:", transposed_OEE.shape) + + if debug: + print("after transpose, transposed_OEE.shape:", transposed_OEE.shape) return transposed_OEE # return OEE From 177f616f09e084f68f4cd9a401d86f1f3a5bda5e Mon Sep 17 00:00:00 2001 From: lee1043 Date: Tue, 14 May 2024 22:29:40 -0700 Subject: [PATCH 15/15] clean up --- pcmdi_metrics/mjo/lib/lib_mjo.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pcmdi_metrics/mjo/lib/lib_mjo.py b/pcmdi_metrics/mjo/lib/lib_mjo.py index 721c8c661..53a2678b4 100644 --- a/pcmdi_metrics/mjo/lib/lib_mjo.py +++ b/pcmdi_metrics/mjo/lib/lib_mjo.py @@ -185,7 +185,7 @@ def generate_axes_and_decorate(Power, NT: int, NL: int) -> xr.DataArray: return da -def output_power_spectra(NL: int, NT: int, Power, debug: bool = False): +def output_power_spectra(NL: int, NT: int, Power, debug: bool = False) -> xr.DataArray: """ Below code taken and modified from Daehyun Kim's Fortran code (MSD/level_2/sample/stps/stps.sea.f.sample) """ @@ -205,7 +205,6 @@ def output_power_spectra(NL: int, NT: int, Power, debug: bool = False): b = list((ss[i] for i in range(int(NL / 2), int(NL / 2) + 1 + 10))) a = np.array(a) b = np.array(b) - # Decoration # Add name attributes to x and y coordinates x_coords = xr.IndexVariable(