#!/usr/bin/env python import sys, os.path, os, glob import xarray as xr, numpy as np, pandas as pd #import matplotlib.pyplot as plt #more imports from numpy import pi as Pi, cos, sin import cftime # # #start from here #year_new = 2000# use new time axis in year 2000 negative_natural = True # rotate local (x,y) coordinate to negative natural coordinate (TC moves to the x-negative direction) """ data_name = 'precip' scale_factor = 24*3600 units = 'mm/day' data_name = 'slp' scale_factor = 1 units = 'mb' """ data_name1 = 'v_ref' scale_factor = 1 units = 'm/s' data_name2 = 'u_ref' scale_factor = 1 units = 'm/s' data_name = 'windspeed' odata_name = f'{data_name}_xy' R = 6370 #km, earth radius L1 = 2*Pi*R/360 #km, distance of one degree longitude at the equator (or latitude at any location) x = xr.DataArray(np.arange(-400, 401, 20), dims='x', attrs=dict(units='km')) y = xr.DataArray(np.arange(-400, 401, 20), dims='y', attrs=dict(units='km')) if negative_natural: x = x.assign_attrs(long_name='negative natural s') y = y.assign_attrs(long_name='negative natural n') # tc tracks ds = xr.open_dataset("/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/netcdf/HIRAM.CTL1990s_v201910_tigercpu_intelmpi_18_540PE.tc_tracks.TS.0101-0200.nc") # interpolation loop zz = np.zeros((ds.year.size, ds.storm.size, ds.stage.size, y.size, x.size)) + np.nan angle_tcmotion = np.zeros((ds.year.size, ds.storm.size, ds.stage.size)) + np.nan for year in range(200,201): iyear = year-101 # composite variable ifile = f'/tigress/wenchang/MODEL_OUT/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/POSTP/{year:04d}0101.atmos_4xdaily.nc' da_u = xr.open_dataset(ifile)['u_ref'] \ .pipe(lambda x: x*scale_factor) \ .assign_attrs(units=units) \ .rename(grid_xt='lon', grid_yt='lat') da_v = xr.open_dataset(ifile)['v_ref'] \ .pipe(lambda x: x*scale_factor) \ .assign_attrs(units=units) \ .rename(grid_xt='lon', grid_yt='lat') da = (da_u**2 + da_v**2)**0.5 # time axis adjustment # change year so that time range is 2000-01-01 00:00:00 to 2001-01-01 00:00:00 time = da.time time_new = [cftime.DatetimeGregorian(yr-year+2000, month, day, hour) for yr,month,day,hour in zip(time.dt.year.values, time.dt.month.values, time.dt.day.values, time.dt.hour.values)] da = da.assign_coords(time=time_new) # no values on Dec 31: shift time index one day backward from Mar 1st to the end of year if not ds.hour.where((ds.month==12)&(ds.day==31)).max()>0: time = da.time times = np.array([f'{yr:04d}-{month:02d}-{day:02d} {hour:02d}:00:00' for yr,month,day,hour in zip(time.dt.year.values, time.dt.month.values, time.dt.day.values, time.dt.hour.values)]) time_new = da.indexes['time'].where( times<'2000-03-01 00:00:00', other=da.indexes['time'].shift(-1, 'D') ) da = da.assign_coords(time=time_new) da = da.assign_coords(time=da.indexes['time'].shift(-1, 'H')) # so that no 24:00:00 (not recognized by python) for istorm in range(13, ds.storm.size): #print('istorm:', istorm) if np.isnan(ds.isel(year=iyear, storm=istorm, stage=0).lon.item()): break for istage in range(44, ds.stage.size): print('istorm:', year, istorm, istage) ds0 = ds.isel(storm=istorm, stage=istage, year=iyear) lon0 = ds0.lon.item() if np.isnan(lon0): break lat0 = ds0.lat.item() month0 = ds0.month.item() day0 = ds0.day.item() hour0 = ds0.hour.item() - 1 # minus 1 to be consistent with da.time # track point time stamp t0 = f'2000-{month0:02d}-{day0:02d}T{hour0:02d}' # local coordinates # rotate local coordinate so that TC moves to the left (direction of negative x axis): negative natural coordinate if negative_natural: # get dlat and dlon if istage == 0: # first track point: forward difference ds_next = ds.isel(year=iyear, storm=istorm, stage=1) dlat = ds_next.lat.item() - ds0.lat.item() dlon = ds_next.lon.item() - ds0.lon.item() elif istage == 119 or np.isnan(ds.isel(year=iyear, storm=istorm, stage=istage+1).lon.item()): # last track point: backward difference ds_prev = ds.isel(year=iyear, storm=istorm, stage=istage-1) dlat = ds0.lat.item() - ds_prev.lat.item() dlon = ds0.lon.item() - ds_prev.lon.item() else: #central difference ds_next = ds.isel(year=iyear, storm=istorm, stage=istage+1) ds_prev = ds.isel(year=iyear, storm=istorm, stage=istage-1) dlat = ( ds_next.lat.item() - ds_prev.lat.item() )/2 dlon = ( ds_next.lon.item() - ds_prev.lon.item() )/2 angle_tcmotion[iyear, istorm, istage] = np.angle(dlon*cos(lat0*Pi/180) + dlat*1j, deg=True)# angle (in degree) between TC motion and zonal direction theta = np.angle(-dlon*cos(lat0*Pi/180) - dlat*1j) # angle between x-negative in rotating axis and zonal direction # x-negative natural coordinate (x,y) projected onto local (x_, y_) coordinate x_ = x*cos(theta) - y*sin(theta) y_ = x*sin(theta) + y*cos(theta) lat_xy = (y_+x_)*0 + lat0 + y_/L1 lon_xy = (y_+x_)*0 + lon0 + x_/(L1*cos(lat_xy*Pi/180)) else: lat_xy = (y+x)*0 + lat0 + y/L1 lon_xy = (y+x)*0 + lon0 + x/(L1*cos(lat_xy*Pi/180)) print(da.sel(time=t0)) print(t0) da0 = da.sel(time=t0).squeeze() \ .interp(lon=lon_xy, lat=lat_xy) \ .drop(['lon', 'lat', 'time']) da.assign_coords(time=da.indexes['time'].shift(1, 'H')).sel(time='2000-03-01T00').plot() zz[iyear, istorm, istage, :, :] = da0.values da_ = xr.DataArray(zz, dims=['year','storm', 'stage', 'y', 'x'], coords=[ds.year, ds.storm, ds.stage, y, x], attrs=dict(units=units)) ds[odata_name] = da_ if negative_natural: angle_tcmotion = xr.DataArray(angle_tcmotion, dims=['year','storm', 'stage'], coords=[ds.storm, ds.stage], attrs=dict(units='deg', long_name='angle between TC motion and zonal east')) ds['angle_tcmotion'] = angle_tcmotion ds.to_netcdf(__file__.replace('.py', '.nc'), encoding={odata_name: {'zlib': True, 'complevel': 1, 'dtype': 'float32'}})