diff --git a/metcalcpy/contributed/spacetime/cross_spectra.py b/metcalcpy/contributed/spacetime/cross_spectra.py index a7adb497..689207e0 100644 --- a/metcalcpy/contributed/spacetime/cross_spectra.py +++ b/metcalcpy/contributed/spacetime/cross_spectra.py @@ -13,6 +13,7 @@ data sets. """ import argparse +import os import xarray as xr import numpy as np @@ -24,6 +25,28 @@ from spacetime import get_symmasymm from spacetime_utils import save_Spectra +import metcalcpy.util.read_env_vars_in_config as readconfig + + +# Read in the YAML config file +# user can use their own, if none specified at the command line, +# use the "default" example YAML config file, spectra_plot_coh2.py +# Using a custom YAML reader so we can use environment variables +cross_config_file = os.getenv("COMP_SPECTRA_YAML_CONFIG_NAME","spectra_comp.yaml") + +config_dict = readconfig.parse_config(cross_config_file) + +# Retrieve settings from config file +#pathdata is now set in the METplus conf file +pathout = config_dict['pathout'][0] +print("Output path ",pathout) +model = config_dict['model'] +print("Model ",model) + +# Make output directory if it does not exist +if not os.path.exists(pathout): + os.makedirs(pathout) + """ Set parameters for the spectra calculation. spd: Number of observations per day. @@ -35,62 +58,68 @@ latMin: Minimum latitude to include. latMax: Maximum latitude to include. """ -spd = 2 -nperseg = 46 * spd -segOverLap = -20 * spd -Symmetry = "symm" -latMin = -15. -latMax = 15. -datestrt = '2015-12-01' # plot start date, format: yyyy-mm-dd -datelast = '2016-03-31' # plot end date, format: yyyy-mm-dd - -parser = argparse.ArgumentParser() -parser.add_argument('--datapath', type=str, - help='data path') -parser.add_argument('--pathout', type=str, - help='output path') -args = parser.parse_args() - -# datapath = '../data/' -# path to the location where to save the output spectra -# pathout = '../data/' -datapath = args.datapath -pathout = args.pathout +#spd = 4 +#nperseg = 15 * spd +#segOverLap = -5 * spd +#Symmetry = "symm" +#latMin = -15. +#latMax = 15. +#datestrt = '2014-04-01T06:00:00' # plot start date, format: yyyy-mm-dd +#datelast = '2014-05-04T18:00:00' # plot end date, format: yyyy-mm-dd +spd = config_dict['spd'] +nperseg = config_dict['nperseg']*spd +segOverLap = config_dict['segOverLap']*spd +Symmetry = config_dict['Symmetry'] +latMin = config_dict['latMin'] +latMax = config_dict['latMax'] +datestrt = config_dict['datestrt'] +datelast = config_dict['datelast'] +print('spd: ', spd,', nperseg: ',nperseg,', segOverLap: ',segOverLap) +print('symmetry: ',Symmetry,', latMin: ',latMin,', latMax: ',latMax) +print('datestrt: ',datestrt,', datelast: ',datelast) print("reading data from file:") """ Read in data here. Example: """ -ds = xr.open_dataset(datapath+'precip.trmm.'+str(spd)+'x.1p0.v7a.fillmiss.comp.2014-2016.nc') +filenames = os.environ.get("COMP_SPECTRA_INPUT_FILE_NAMES","P_verif,P_model,Vlev1_model,Vlev2_model").split(",") +print("Filename ",filenames[1]) +ds = xr.open_dataset(filenames[0]) z = ds.precip z = z.sel(lat=slice(latMin, latMax)) z = z.sel(time=slice(datestrt, datelast)) z = z.squeeze() latC = ds.lat.sel(lat=slice(latMin, latMax)) -ds = xr.open_dataset(datapath+'precip.erai.sfc.1p0.'+str(spd)+'x.2014-2016.nc') -x = ds.precip +ds = xr.open_dataset(filenames[1]) +x = ds.prate x = x.sel(lat=slice(latMin, latMax)) x = x.sel(time=slice(datestrt, datelast)) x = x.squeeze() latA = ds.lat.sel(lat=slice(latMin, latMax)) ds.close() -ds = xr.open_dataset(datapath+'div.erai.850.1p0.'+str(spd)+'x.2014-2016.nc') -y = ds.div +ds = xr.open_dataset(filenames[2]) +ds = ds.rename({'latitude':'lat'}) +ds = ds.sortby('lat',ascending=True) +y = ds.u y = y.sel(lat=slice(latMin, latMax)) y = y.sel(time=slice(datestrt, datelast)) y = y.squeeze() latB = ds.lat.sel(lat=slice(latMin, latMax)) -ds = xr.open_dataset(datapath+'div.erai.200.1p0.'+str(spd)+'x.2014-2016.nc') -w = ds.div -w = y.sel(lat=slice(latMin, latMax)) -w = y.sel(time=slice(datestrt, datelast)) -w = y.squeeze() +ds = xr.open_dataset(filenames[3]) +ds = ds.rename({'latitude':'lat'}) +ds = ds.sortby('lat',ascending=True) +w = ds.u +w = w.sel(lat=slice(latMin, latMax)) +w = w.sel(time=slice(datestrt, datelast)) +w = w.squeeze() latB = ds.lat.sel(lat=slice(latMin, latMax)) - +print(x.shape) +print(y.shape) +print(z.shape) if any(latA - latB) != 0: print("Latitudes must be the same for both variables! Check latitude ordering.") @@ -132,7 +161,7 @@ freq = freq * spd wnum = result['wave'] # save spectra in netcdf file -fileout = 'SpaceTimeSpectra_ERAI_P_D850_' + Symmetry + '_' + str(spd) + 'spd' +fileout = 'SpaceTimeSpectra_'+model+'_P_D850_' + Symmetry + '_' + str(spd) + 'spd' print('saving spectra to file: ' + pathout + fileout + '.nc') save_Spectra(STC, freq, wnum, fileout, pathout) @@ -142,16 +171,17 @@ freq = freq * spd wnum = result['wave'] # save spectra in netcdf file -fileout = 'SpaceTimeSpectra_ERAI_P_D200_' + Symmetry + '_' + str(spd) + 'spd' +fileout = 'SpaceTimeSpectra_'+model+'_P_D200_' + Symmetry + '_' + str(spd) + 'spd' print('saving spectra to file: ' + pathout + fileout + '.nc') save_Spectra(STC, freq, wnum, fileout, pathout) -result = mjo_cross(X, Z, nperseg, segOverLap) +spd = 2 +result = mjo_cross(X.sel(time=Z.time), Z, nperseg, segOverLap) STC = result['STC'] # , freq, wave, number_of_segments, dof, prob, prob_coh2 freq = result['freq'] freq = freq * spd wnum = result['wave'] # save spectra in netcdf file -fileout = 'SpaceTimeSpectra_ERAI_TRMM_P_' + Symmetry + '_' + str(spd) + 'spd' +fileout = 'SpaceTimeSpectra_ERAI_'+model+'_P_' + Symmetry + '_' + str(spd) + 'spd' print('saving spectra to file: ' + pathout + fileout + '.nc') save_Spectra(STC, freq, wnum, fileout, pathout)