diff --git a/.github/workflows/test-B01_SL_load_single_file.yml b/.github/workflows/test-B01_SL_load_single_file.yml index 28004887..cfbf9d7a 100644 --- a/.github/workflows/test-B01_SL_load_single_file.yml +++ b/.github/workflows/test-B01_SL_load_single_file.yml @@ -8,8 +8,8 @@ jobs: strategy: fail-fast: false matrix: - version: ['3.9'] - runs-on: ubuntu-latest + version: ['3.11'] + runs-on: ubuntu-22.04 steps: - name: install mpi run: sudo apt update && sudo apt-get install openmpi-bin openmpi-doc libopenmpi-dev @@ -21,11 +21,15 @@ jobs: cache: "pip" - name: install icesat2-tracks using pip run: pip install . + - name: List dependencies + run: pip list - name: first step B01_SL_load_single_file run: python src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py 20190502052058_05180312_005_01 SH_testSLsinglefile2 True - name: second step make_spectra run: python src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py SH_20190502_05180312 SH_testSLsinglefile2 True - name: third step plot_spectra run: python src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py SH_20190502_05180312 SH_testSLsinglefile2 True - - name: fourth step plot_spectra + - name: fourth step IOWAGA thredds run: python src/icesat2_tracks/analysis_db/A02c_IOWAGA_thredds_prior.py SH_20190502_05180312 SH_testSLsinglefile2 True + - name: Fifth step B04_angle + run: python src/icesat2_tracks/analysis_db/B04_angle.py SH_20190502_05180312 SH_testSLsinglefile2 True diff --git a/analysis_db/B04_angle.py b/analysis_db/B04_angle.py deleted file mode 100644 index 365239fb..00000000 --- a/analysis_db/B04_angle.py +++ /dev/null @@ -1,818 +0,0 @@ -# %% -import os, sys -#execfile(os.environ['PYTHONSTARTUP']) - -""" -This file open a ICEsat2 track applied filters and corections and returns smoothed photon heights on a regular grid in an .nc file. -This is python 3 -""" - -exec(open(os.environ['PYTHONSTARTUP']).read()) -exec(open(STARTUP_2021_IceSAT2).read()) - -#%matplotlib inline - -import ICEsat2_SI_tools.convert_GPS_time as cGPS -import h5py -import ICEsat2_SI_tools.io as io -import ICEsat2_SI_tools.spectral_estimates as spec - -import imp -import copy -import spicke_remover -import datetime -import concurrent.futures as futures - -from numba import jit - -from ICEsat2_SI_tools import angle_optimizer -import ICEsat2_SI_tools.wave_tools as waves -import concurrent.futures as futures - -import time - -from contextlib import contextmanager -col.colormaps2(21) - -@contextmanager -def suppress_stdout(): - with open(os.devnull, "w") as devnull: - old_stdout = sys.stdout - sys.stdout = devnull - try: - yield - finally: - sys.stdout = old_stdout - -col_dict = col.rels -#import s3fs -# %% -track_name, batch_key, test_flag = io.init_from_input(sys.argv) # loads standard experiment -#track_name, batch_key, test_flag = '20190605061807_10380310_004_01', 'SH_batch01', False -#track_name, batch_key, test_flag = '20190601094826_09790312_004_01', 'SH_batch01', False -#track_name, batch_key, test_flag = '20190207111114_06260210_004_01', 'SH_batch02', False -#track_name, batch_key, test_flag = '20190219073735_08070210_004_01', 'SH_batch02', False -#track_name, batch_key, test_flag = '20190215184558_07530210_004_01', 'SH_batch02', False - -# good track -#track_name, batch_key, test_flag = '20190502021224_05160312_004_01', 'SH_batch02', False -#track_name, batch_key, test_flag = '20190502050734_05180310_004_01', 'SH_batch02', False -#track_name, batch_key, test_flag = '20190210143705_06740210_004_01', 'SH_batch02', False -#track_name, batch_key, test_flag = 'NH_20190301_09580203', 'NH_batch05', True -#track_name, batch_key, test_flag = 'SH_20190210_06740210', 'SH_publish', True -#track_name, batch_key, test_flag = 'SH_20190219_08070212', 'SH_batchminimal', True - - - -#print(track_name, batch_key, test_flag) -hemis, batch = batch_key.split('_') -#track_name= '20190605061807_10380310_004_01' -ATlevel= 'ATL03' - -save_path = mconfig['paths']['work'] +batch_key+'/B04_angle/' -save_name = 'B04_'+track_name - -plot_path = mconfig['paths']['plot'] + '/'+hemis+'/'+batch_key+'/' + track_name + '/' -MT.mkdirs_r(plot_path) -MT.mkdirs_r(save_path) -bad_track_path =mconfig['paths']['work'] +'bad_tracks/'+ batch_key+'/' -# %% - -all_beams = mconfig['beams']['all_beams'] -high_beams = mconfig['beams']['high_beams'] -low_beams = mconfig['beams']['low_beams'] -beam_groups = mconfig['beams']['groups'] - -#Gfilt = io.load_pandas_table_dict(track_name + '_B01_regridded', load_path) # rhis is the rar photon data - -load_path = mconfig['paths']['work'] +batch_key+'/B01_regrid/' -#G_binned = io.load_pandas_table_dict(track_name + '_B01_binned' , load_path) # -G_binned_store = h5py.File(load_path +'/'+track_name + '_B01_binned.h5', 'r') -G_binned = dict() -for b in all_beams: - - G_binned[b] = io.get_beam_hdf_store(G_binned_store[b]) -G_binned_store.close() - -load_path = mconfig['paths']['work'] +batch_key+'/B02_spectra/' -Gx = xr.load_dataset(load_path+ '/B02_'+track_name + '_gFT_x.nc' ) # -Gk = xr.load_dataset(load_path+ '/B02_'+track_name + '_gFT_k.nc' ) # - - -# %% load prior information -load_path = mconfig['paths']['work'] +batch_key+'/A02_prior/' -#track_name = '20190208104534_06410210_004_01' -try: - Prior = MT.load_pandas_table_dict('/A02_'+track_name, load_path)['priors_hindcast'] -except: - print('Prior not founds exit') - MT.json_save('B04_fail', plot_path, {'time':time.asctime( time.localtime(time.time()) ) , 'reason': 'Prior not found'}) - print('exit()') - exit() - - -#### Define Prior -# Use partitions -# Prior2 = Prior.loc[['ptp0','ptp1','ptp2','ptp3','ptp4','ptp5']]['mean'] -# dominat_period = Prior2[Prior2.max() ==Prior2] -# aa = Prior.loc[['pdp0','pdp1','pdp2','pdp3','pdp4','pdp5']]['mean'].astype('float') -# dominant_dir = waves.get_ave_amp_angle(aa *0+1,aa )[1] -# dominant_dir_spread = Prior.loc[['pspr0','pspr1','pspr2','pspr3','pspr4','pspr5']]['mean'].median() -# -# prior_sel= {'alpha': ( dominant_dir *np.pi/180 , dominant_dir_spread *np.pi/180) } # to radiens -#prior_sel= {'alpha': ( -60 *np.pi/180 , dominant_dir_spread *np.pi/180) } # to radiens - -Prior - -Pperiod = Prior.loc[['ptp0','ptp1','ptp2','ptp3','ptp4','ptp5']]['mean'] -Pdir = Prior.loc[['pdp0','pdp1','pdp2','pdp3','pdp4','pdp5']]['mean'].astype('float') -Pspread = Prior.loc[['pspr0','pspr1','pspr2','pspr3','pspr4','pspr5']]['mean'] - -Pperiod = Pperiod[~np.isnan(list(Pspread))] -Pdir = Pdir[~np.isnan(list(Pspread))] -Pspread = Pspread[~np.isnan(list(Pspread))] - - -# reset dirs: -Pdir[Pdir > 180] = Pdir[Pdir > 180] - 360 -Pdir[Pdir < -180] = Pdir[Pdir < -180] + 360 - -# reorder dirs -dir_best = [0] -for dir in Pdir: - ip = np.argmin([ abs(dir_best[-1] - dir), abs(dir_best[-1] - (dir - 360 )), abs(dir_best[-1] - (dir + 360 )) ] ) - new_dir = np.array([ dir, (dir - 360 ) , (dir + 360 ) ])[ip] - dir_best.append(new_dir) -dir_best = np.array(dir_best[1:]) - -# %% - -if len(Pperiod) == 0: - print('constant peak wave number') - kk = Gk.k - Pwavenumber = kk*0 + (2 * np.pi / (1/ Prior.loc['fp']['mean']) )**2 / 9.81 - dir_best = kk*0 + Prior.loc['dp']['mean'] - #dir_interp = np.interp(kk, Pwavenumber[Pwavenumber.argsort()] , dir_best[Pwavenumber.argsort()] ) - dir_interp_smth = dir_interp = kk*0 + Prior.loc['dp']['mean'] - spread_smth = spread_interp = kk*0 + Prior.loc['spr']['mean'] - #spread_interp = np.interp(kk, Pwavenumber[Pwavenumber.argsort()] , Pspread[Pwavenumber.argsort()].astype('float') ) - #spread_smth = M.runningmean(spread_interp, 30, tailcopy= True) - #spread_smth[-1] = spread_smth[-2] - -else: - Pwavenumber = (2 * np.pi / Pperiod )**2 / 9.81 - kk = Gk.k - dir_interp = np.interp(kk, Pwavenumber[Pwavenumber.argsort()] , dir_best[Pwavenumber.argsort()] ) - dir_interp_smth = M.runningmean(dir_interp, 30, tailcopy= True) - dir_interp_smth[-1] = dir_interp_smth[-2] - - spread_interp = np.interp(kk, Pwavenumber[Pwavenumber.argsort()] , Pspread[Pwavenumber.argsort()].astype('float') ) - spread_smth = M.runningmean(spread_interp, 30, tailcopy= True) - spread_smth[-1] = spread_smth[-2] - - -font_for_pres() - -F = M.figure_axis_xy(5, 4.5, view_scale= 0.5) -plt.subplot(2, 1, 1) -plt.title('Prior angle smoothed\n'+ track_name, loc ='left') - - -plt.plot( Pwavenumber , dir_best, '.r', markersize = 8) -plt.plot( kk , dir_interp, '-', color= 'red', linewidth = 0.8, zorder=11) -plt.plot( kk , dir_interp_smth , color=col.green1) - -plt.fill_between(kk, dir_interp_smth -spread_smth, dir_interp_smth +spread_smth, zorder= 1, color=col.green1, alpha = 0.2 ) -plt.ylabel('Angle (deg)') -#plt.xlabel('wavenumber ($2 \pi/\lambda$)') - -ax2 = plt.subplot(2, 1, 2) -plt.title('Prior angle adjusted ', loc ='left') - -# adjust angle def: -dir_interp_smth[dir_interp_smth> 180] = dir_interp_smth[dir_interp_smth> 180]- 360 -dir_interp_smth[dir_interp_smth< -180] = dir_interp_smth[dir_interp_smth< -180]+ 360 - -plt.fill_between(kk, dir_interp_smth -spread_smth, dir_interp_smth +spread_smth, zorder= 1, color=col.green1, alpha = 0.2 ) -plt.plot( kk , dir_interp_smth , '.', markersize = 1 , color=col.green1) - -ax2.axhline(85, color='gray', linewidth= 2) -ax2.axhline(-85, color='gray', linewidth= 2) - -plt.ylabel('Angle (deg)') -plt.xlabel('wavenumber ($2 \pi/\lambda$)') - -F.save_light(path= plot_path, name = 'B04_prior_angle') - - -# save -dir_interp_smth = xr.DataArray(data=dir_interp_smth * np.pi/180 , dims='k', coords ={'k':kk}, name='Prior_direction') -spread_smth = xr.DataArray(data=spread_smth* np.pi/180 , dims='k', coords ={'k':kk}, name='Prior_spread') -Prior_smth = xr.merge([dir_interp_smth, spread_smth]) - -# %% -prior_angle =Prior_smth.Prior_direction * 180/np.pi -if (abs(prior_angle) > 80).all(): - print('Prior angle is ', prior_angle.min().data, prior_angle.max().data, '. quit.') - dd_save = {'time' : time.asctime( time.localtime(time.time()) ), - 'angle': list([ float(prior_angle.min().data), float(prior_angle.max().data), float(prior_angle.median()) ]) } - MT.json_save('B04_fail', plot_path, dd_save) - print('exit()') - exit() - -# Use fake -#prior_sel= {'alpha': ( 0.6 , dominant_dir_spread *np.pi/180) } # to radiens - -# Use mean direction -#prior_sel= {'alpha': ( Prior.loc['dp']['mean'] *np.pi/180 , Prior.loc['spr']['mean'] *np.pi/180) } - - -# define paramater range -params_dict = {'alpha': [ -0.85 * np.pi /2, 0.85 * np.pi /2, 5], - 'phase':[ 0 , 2*np.pi , 10]} - -alpha_dx = 0.02 -max_wavenumbers = 25 - -sample_flag = True -optimize_flag = False -brute_flag = False - -plot_flag = False - -Nworkers = 6 -N_sample_chain = 300 -N_sample_chain_burn = 30 - -max_x_pos = 8 -x_pos_jump = 2 - -def make_fake_data(xi,group ): - ki= Gk.k[0:2] - - bins = np.arange(params_dict['alpha'][0], params_dict['alpha'][1]+alpha_dx, alpha_dx) - bins_pos = (bins[0:-1] + np.diff(bins)/2) - marginal_stack = xr.DataArray( np.nan* np.vstack([bins_pos, bins_pos]).T, dims= ('angle', 'k'), coords = {'angle':bins_pos, 'k':ki.data } ) - - group_name = str('group' + group[0].split('gt')[1].split('l')[0]) - marginal_stack.coords['beam_group'] = group_name - marginal_stack.coords['x'] = xi - marginal_stack.name = 'marginals' - marginal_stack.expand_dims(dim = 'x', axis = 2).expand_dims(dim = 'beam_group', axis = 3) - return marginal_stack - -def define_wavenumber_weights_tot_var(dd, m = 3, variance_frac = 0.33, k_upper_lim= None, verbose=False): - - """ - return peaks of a power spectrum dd that in the format such that they can be used as weights for the frequencies based fitting - - inputs: - dd xarray with PSD as data amd coordindate wavenumber k - m running mean half-width in gridpoints - variance_frac (0 to 1) How much variance should be explained by the returned peaks - verbose if true it plots some stuff - - - return: - mask size of dd. where True the data is identified as having significant amplitude - k wanumbers where mask is true - dd_rm smoothed version of dd - positions postions where of significant data in array - """ - - if len(dd.shape) == 2: - dd_use = dd.mean('beam') - - if m is None: - dd_rm = dd_use.data#M.runningmean(dd, m, tailcopy=True) - else: - dd_rm = M.runningmean(dd_use, m, tailcopy=True) - - k = dd_use.k[~np.isnan(dd_rm)].data - dd_rm = dd_rm[~np.isnan(dd_rm)] - - orders = dd_rm.argsort()[::-1] - var_mask = dd_rm[orders].cumsum()/dd_rm.sum() < variance_frac - pos_cumsum = orders[var_mask] - #k_list = k[pos_cumsum] - #dd_list = dd_rm[pos_cumsum] - mask = var_mask[orders.argsort()] - if k_upper_lim is not None: - mask = (k < k_upper_lim) & mask - - if verbose: - - plt.plot(dd.k, dd, '-', color = col_dict[str(amp_data.beam[0].data)], markersize= 20, alpha = 0.6) - plt.plot(k, dd_rm, '-k', markersize= 20) - #print(k_list, dd_list) - plt.plot(k[mask], dd_rm[mask], '.r', markersize= 10, zorder=12) - if k_upper_lim is not None: - plt.gca().axvline(k_upper_lim, color= 'black') - - return mask, k, dd_rm, pos_cumsum - -def define_wavenumber_weights_threshold(dd, m = 3, Nstd= 2, verbose=False): - - if m is None: - dd_rm = dd#M.runningmean(dd, m, tailcopy=True) - else: - dd_rm = M.runningmean(dd, m, tailcopy=True) - - k = dd.k[~np.isnan(dd_rm)] - dd_rm = dd_rm[~np.isnan(dd_rm)] - - treshold = np.nanmean(dd_rm) + np.nanstd(dd_rm) *Nstd - mask = dd_rm > treshold - - - if verbose: - plt.plot(dd.k, dd, '-k', markersize= 20) - plt.plot(k, dd_rm, '-b', markersize= 20) - - k_list = k[mask] - dd_list = dd_rm[mask] - #print(k_list, dd_list) - plt.plot(k_list, dd_list, '.r', markersize= 10, zorder=12) - - return mask, k, dd_rm, np.arange(0, mask.size)[mask] - -def plot_instance(z_model, fargs , key, SM, non_dim=False, title_str= None , brute=False, optimze= False, sample= False, view_scale = 0.3): - - x_concat, y_concat, z_concat = fargs - - F = M.figure_axis_xy(5,6, view_scale = view_scale, container = True) - plt.suptitle(title_str) - gs = GridSpec(4, 3, wspace=0.4, hspace=1.2)#figure=fig, - F.gs = gs - # y_offset= 0.5 - # plt.plot(Gm.eta, Gm.y_model_normed+y_offset * Gm.dist_y/np.diff(Gm.dist_y), **model_plot_karg) - # plt.plot(Gm.eta, Gm.y_model_normed+y_offset * Gm.dist_y/np.diff(Gm.dist_y), **data_plot_karg) - # plt.xlim(-1000, 1000) - - - import itertools - col_list = itertools.cycle([col.cascade2, col.rascade2, col.cascade1, col.rascade1, col.cascade3, col.rascade3]) - - beam_list = list(set(y_concat)) - for y_pos, pos in zip(beam_list, [ gs[0, :] , gs[1, :] ] ): - - F.ax2 = F.fig.add_subplot(pos) - - plt.title( str(y_pos) ) - #plt.plot(GG['x_prime'], GG['z_model'], '.' ) - - - plt.plot(x_concat[y_concat == y_pos], z_concat[y_concat == y_pos] , c=col.gray, linewidth = 1) - plt.plot(x_concat[y_concat == y_pos], z_model[y_concat == y_pos] , '-', c=next(col_list)) - plt.xlim(x_concat[y_concat == y_pos][0], x_concat[y_concat == y_pos][-1]) - - - plt.xlabel('meter') - F.ax3 = F.fig.add_subplot(gs[2:, 0:-1]) - if brute is True: - plt.title('Brute-force costs', loc='left') - SM.plot_brute(marker= '.', color ='blue', markersize=15, label= 'Brute', zorder=10) - # plt.colorbar(orientation='horizontal') - if optimze is True: - SM.plot_optimze(color= 'r', markersize=10, zorder=12, label= 'Dual Annealing') - - if sample is True: - SM.plot_sample(markersize= 2, linewidth= 0.8, alpha= 0.2, color= 'black', zorder=8) - - F.ax4 = F.fig.add_subplot(gs[2:, -1]) - return F - - -# isolate x positions with data -data_mask = Gk.gFT_PSD_data.mean('k') -data_mask.coords['beam_group'] = ('beam', ['beam_group'+g[2] for g in data_mask.beam.data]) -data_mask_group = data_mask.groupby('beam_group').mean(skipna=False) -# these stancils are actually used -data_sel_mask = data_mask_group.sum('beam_group') !=0 - -x_list = data_sel_mask.x[data_sel_mask] # iterate over these x posistions -x_list_flag = ~np.isnan(data_mask_group.sel(x = x_list) )# flag that is False if there is no data - -#### limit number of x coordinates - -x_list = x_list[::x_pos_jump] -if len(x_list) > max_x_pos: - x_list = x_list[0:max_x_pos] -x_list_flag= x_list_flag.sel(x =x_list) - -# plot -font_for_print() -F = M.figure_axis_xy(5.5, 3, view_scale= 0.8) -plt.suptitle(track_name) -ax1 = plt.subplot(2, 1, 1) -plt.title('Data in Beam', loc= 'left') -plt.pcolormesh(data_mask.x/1e3, data_mask.beam, data_mask, cmap= plt.cm.OrRd) -for i in np.arange(1.5, 6, 2): - ax1.axhline(i, color= 'black', linewidth =0.5) -plt.xlabel('Distance from Ice Edge') - -ax2 = plt.subplot(2, 1, 2) -plt.title('Data in Group', loc= 'left') -plt.pcolormesh(data_mask.x/1e3, data_mask_group.beam_group, data_mask_group, cmap= plt.cm.OrRd) - -for i in np.arange(0.5, 3, 1): - ax2.axhline(i, color= 'black', linewidth =0.5) - -plt.plot( x_list/1e3, x_list*0 +0, '.', markersize= 2, color= col.cascade1 ) -plt.plot( x_list/1e3, x_list*0 +1, '.', markersize= 2, color= col.cascade1 ) -plt.plot( x_list/1e3, x_list*0 +2, '.', markersize= 2, color= col.cascade1 ) - -plt.xlabel('Distance from Ice Edge') - -F.save_pup(path= plot_path, name = 'B04_data_avail') - - - -# %% -Marginals = dict() -L_collect = dict() - -group_number = np.arange(len(beam_groups)) -ggg, xxx = np.meshgrid(group_number , x_list.data ) - -for gi in zip(ggg.flatten(), xxx.flatten()): - print(gi) - - group, xi = beam_groups[gi[0]], gi[1] - - if bool(x_list_flag.sel(x= xi).isel(beam_group= gi[0]).data) is False: - print('no data, fill with dummy') - ikey = str(xi) +'_' + '_'.join(group) - Marginals[ikey] = make_fake_data(xi, group) - #print(Marginals[ikey].angle.data[::20]) - continue - - GGx = Gx.sel(beam= group).sel(x = xi) - GGk = Gk.sel(beam= group).sel(x = xi) - - ### define data - # normalize data - key = 'y_data' - amp_Z = (GGx[key] - GGx[key].mean(['eta'])) /GGx[key].std(['eta']) - N = amp_Z.shape[0] - - # define x,y positions - eta_2d = GGx.eta + GGx.x_coord - GGx.x_coord.mean() - nu_2d = GGx.eta * 0 + GGx.y_coord - GGx.y_coord.mean() - - # repack as np arrays - x_concat = eta_2d.data.T.flatten() - y_concat = nu_2d.data.T.flatten() - z_concat = amp_Z.data.flatten() - - x_concat= x_concat[~np.isnan(z_concat)] - y_concat= y_concat[~np.isnan(z_concat)] - z_concat= z_concat[~np.isnan(z_concat)] - N_data = x_concat.size - - if np.isnan(z_concat).sum() != 0: - raise ValueError('There are still nans') - - mean_dist = (nu_2d.isel(beam= 0) - nu_2d.isel(beam= 1)).mean().data - k_upper_lim = 2 *np.pi / ( mean_dist *1 ) - - print('k_upper_lim ', k_upper_lim) - # threshold method - #mask, k, weights, positions = define_wavenumber_weights_threshold( Gi.mean('dist_y')['gFT_PSD_data'], 3 , verbose= True) - #plt.plot(k[mask], weights[mask], 'g*', markersize=20) - # plt.show() - - #variance method - amp_data = np.sqrt(GGk.gFT_cos_coeff**2 + GGk.gFT_sin_coeff**2) - mask, k, weights, positions = define_wavenumber_weights_tot_var(amp_data, m= 1, k_upper_lim= k_upper_lim, variance_frac = 0.20 , verbose= False) - #plt.xlim( k[mask].min()*0.8 ,max(k_upper_lim, k[mask].max()*1.2) ) - #plt.xlim( k[mask].min()*0.8 ,k[mask].max()*1.4 ) - #plt.show() - - if (len(k[mask]) ==0): - print('no good k found, fill with dummy') - ikey = str(xi) +'_' + '_'.join(group) - Marginals[ikey] = make_fake_data(xi, group) - continue - - - # # testing objective functions - # imp.reload(angle_optimizer) - # pars = SM.params - # pars['alpha'].value = 0.3 - # K_abs = pars['K_prime'] / np.cos(pars['alpha']) - # - # k = K_abs * np.cos(pars['alpha']) - # l = K_abs * np.sin(pars['alpha']) - # - # z1 = angle_optimizer.get_z_model(x_concat, y_concat, pars['K_prime'], pars['K_amp'], pars['alpha'],pars['phase']) - # %timeit z1 = angle_optimizer.get_z_model(x_concat, y_concat, pars['K_prime'], pars['K_amp'], pars['alpha'],pars['phase']) - # - # z2 = angle_optimizer.wavemodel_single_wave(x_concat, y_concat, k,l, pars['K_amp'].value, group_phase= pars['phase'].value ) - # %timeit z2 = angle_optimizer.get_z_model_single_wave(x_concat, y_concat, k,l, pars['K_amp'].value, group_phase= pars['phase'].value ) - - #### prepare loop - # init object and test - #imp.reload(angle_optimizer) - - SM = angle_optimizer.sample_with_mcmc(params_dict) - SM.set_objective_func(angle_optimizer.objective_func) - - SM.fitting_args = fitting_args = (x_concat, y_concat, z_concat) - - # test: - k_prime_max= 0.02 #[mask][0] # chose a test wavenumber - amp_Z= 1 - prior_sel= {'alpha': ( Prior_smth.sel(k =k_prime_max, method='nearest').Prior_direction.data, - Prior_smth.sel(k =k_prime_max, method='nearest').Prior_spread.data) } - SM.fitting_kargs = fitting_kargs = {'prior': prior_sel , 'prior_weight' : 3 } - # test if it works - SM.params.add('K_prime', k_prime_max , vary=False , min=k_prime_max*0.5, max=k_prime_max*1.5) - SM.params.add('K_amp', amp_Z , vary=False , min=amp_Z*.0 , max=amp_Z*5) - try: - SM.test_objective_func() - except: - raise ValueError('Objective function test fails') - - #%timeit SM.test_objective_func() - - #k_prime_max, Z_max = k_list[0], weight_list[0] - def get_instance(k_pair): - - #print('##') - #print(k_pair) - k_prime_max,Z_max = k_pair - #print(np.round(k_prime_max, 4)) - - prior_sel= {'alpha': ( Prior_smth.sel(k =k_prime_max, method='nearest').Prior_direction.data, - Prior_smth.sel(k =k_prime_max, method='nearest').Prior_spread.data) } - - #print(prior_sel) - SM.fitting_kargs = fitting_kargs = {'prior': prior_sel , 'prior_weight' : 2 } - #SM.fitting_kargs = fitting_kargs = {'prior': None , 'prior_weight' : 1 } - - amp_Z = 1##z_concat.var()#Z_max#0.5 #amp_enhancement * abs(Z_max)**2 /N - - SM.params.add('K_prime', k_prime_max , vary=False , min=k_prime_max*0.5, max=k_prime_max*1.5) - SM.params.add('K_amp' , amp_Z , vary=False , min=amp_Z*.0 , max=amp_Z*5) - #print(SM.params.pretty_print()) - L_sample_i = None - L_optimize_i = None - L_brute_i = None - if sample_flag: - #with suppress_stdout(): - SM.sample(verbose= False, steps=N_sample_chain,progress= False, workers= None) - L_sample_i = list(SM.fitter.params.valuesdict().values()) # mcmc results - - elif optimize_flag: - SM.optimize(verbose= False) - L_optimize_i = list(SM.fitter_optimize.params.valuesdict().values()) # mcmc results - - elif brute_flag: - SM.brute(verbose= False) - L_brute_i = list(SM.fitter_brute.params.valuesdict().values()) # mcmc results - else: - raise ValueError('non of sample_flag,optimize_flag, or brute_flag are True') - - y_hist, bins, bins_pos = SM.get_marginal_dist('alpha', alpha_dx, burn = N_sample_chain_burn, plot_flag= False) - fitter = SM.fitter # MCMC results - z_model = SM.objective_func(fitter.params, *fitting_args , test_flag= True) - cost = (fitter.residual**2).sum()/(z_concat**2).sum() - #cost_list.append( (fitter.residual**2).sum()/(z_concat**2).sum() ) - - if plot_flag: - - F = plot_instance(z_model, fitting_args, 'y_data_normed' , SM, brute = brute_flag , optimze= optimize_flag, sample= sample_flag, title_str = 'k='+ str(np.round(k_prime_max, 4)), view_scale = 0.6) - - if (fitting_kargs['prior'] is not None): - F.ax3.axhline(prior_sel['alpha'][0], color='green', linewidth = 2, label ='Prior') - #F.ax3.axhspan(prior_sel['alpha'][0]- prior_sel['alpha'][1], prior_sel['alpha'][0]+ prior_sel['alpha'][1], color='gray', alpha=0.3) - F.ax3.axhline(prior_sel['alpha'][0]- prior_sel['alpha'][1], color='green', linewidth = 0.7) - F.ax3.axhline(prior_sel['alpha'][0]+ prior_sel['alpha'][1], color='green', linewidth = 0.7) - - F.ax3.axhline(fitter.params['alpha'].min, color='gray', linewidth = 2, alpha = 0.6) - F.ax3.axhline(fitter.params['alpha'].max, color='gray', linewidth = 2, alpha = 0.6) - - plt.sca(F.ax3) - plt.legend() - plt.xlabel('Phase') - plt.ylabel('Angle') - plt.xlim(0, np.pi*2) - - plt.sca(F.ax4) - plt.xlabel('Density') - plt.stairs(y_hist, bins, orientation='horizontal', color='k') - - F.ax4.axhline(fitter.params['alpha'].min, color='gray', linewidth = 2, alpha = 0.6) - F.ax4.axhline(fitter.params['alpha'].max, color='gray', linewidth = 2, alpha = 0.6) - - F.ax3.set_ylim(min( -np.pi /2, prior_sel['alpha'][0]- 0.2 ) , max(np.pi /2, prior_sel['alpha'][0] + 0.2 ) ) - F.ax4.set_ylim(min( -np.pi /2, prior_sel['alpha'][0]- 0.2 ) , max(np.pi /2, prior_sel['alpha'][0]+ 0.2 ) ) - - - plt.show() - F.save_light(path= plot_path, name = track_name + '_fit_k' + str(k_prime_max)) - - marginal_stack_i = xr.DataArray( y_hist, dims= ('angle'), coords = {'angle':bins_pos } ) - marginal_stack_i.coords['k'] = np.array(k_prime_max) #( ('k'), np.array(k_prime_max) ) - - rdict = { - 'marginal_stack_i': marginal_stack_i, - 'L_sample_i': L_sample_i, - 'L_optimize_i': L_optimize_i, - 'L_brute_i': L_brute_i, - 'cost':cost - } - return k_prime_max, rdict - - k_list, weight_list = k[mask], weights[mask] - print('# of wavenumber: ' , len(k_list)) - if len(k_list) > max_wavenumbers: - print('cut wavenumber list to 20') - k_list = k_list[0:max_wavenumbers] - weight_list = weight_list[0:max_wavenumbers] - - # A= dict() - # for k_pair in zip(k_list, weight_list): - # kk, I = get_instance(k_pair) - # A[kk] = I - - with futures.ProcessPoolExecutor(max_workers=Nworkers) as executor: - A = dict( executor.map(get_instance, zip(k_list, weight_list) )) - - cost_stack = dict() - marginal_stack =dict() - #fitting_kargs = {'size' :1} - L_sample = pd.DataFrame(index=['alpha', 'group_phase', 'K_prime', 'K_amp'] ) - L_optimize = pd.DataFrame(index=['alpha', 'group_phase', 'K_prime', 'K_amp'] ) - L_brute = pd.DataFrame(index=['alpha', 'group_phase', 'K_prime', 'K_amp'] ) - - for kk,I in A.items(): - L_sample[kk] = I['L_sample_i'] - L_optimize[kk] = I['L_optimize_i'] - L_brute[kk] = I['L_brute_i'] - - marginal_stack[kk] = I['marginal_stack_i'] - cost_stack[kk] = I['cost'] - - # ## add beam_group dimension - marginal_stack = xr.concat(marginal_stack.values(), dim='k' ).sortby('k') - L_sample = L_sample.T.sort_values('K_prime') - L_optimize = L_optimize.T.sort_values('K_prime') - L_brute = L_brute.T.sort_values('K_prime') - - #print(marginal_stack.angle.data[::20]) - - print('done with ', group, xi/1e3) - - # collect - ikey = str(xi) +'_' + '_'.join(group) - - #marginal_stack.coords['cost'] = (('k'), np.expand_dims(np.expand_dims(list(cost_stack.values()), 1), 2) ) - marginal_stack.name = 'marginals' - marginal_stack = marginal_stack.to_dataset() - marginal_stack['cost'] = (('k'), list(cost_stack.values()) ) - marginal_stack['weight'] = (('k'), weight_list ) - - group_name = str('group' + group[0].split('gt')[1].split('l')[0]) - marginal_stack.coords['beam_group'] = group_name - marginal_stack.coords['x'] = xi - - Marginals[ikey] = marginal_stack.expand_dims(dim = 'x', axis = 0).expand_dims(dim = 'beam_group', axis = 1) - Marginals[ikey].coords['N_data'] = ( ('x', 'beam_group'), np.expand_dims(np.expand_dims(N_data, 0), 1) ) - # L_brute - # L_optimize - - L_sample['cost'] = cost_stack - L_sample['weight'] = weight_list - L_collect[group_name, str(int(xi))] = L_sample#pd.concat(L_collect_per_xi) - - - -# %% -#list(Marginals.values())[0] -MM = xr.merge( Marginals.values()) -MM =xr.merge([ MM, Prior_smth]) -MM.to_netcdf(save_path + save_name + '_marginals.nc') - -try: - LL = pd.concat(L_collect) - MT.save_pandas_table({'L_sample':LL} ,save_name+ '_res_table', save_path) -except: - pass - -# %% plot -font_for_print() -F = M.figure_axis_xy(6, 5.5, view_scale= 0.7, container = True) - -gs = GridSpec(4,6, wspace=0.2, hspace=.8)#figure=fig, - -ax0 = F.fig.add_subplot(gs[0:2, -1]) -ax0.tick_params(labelleft=False) - -#klims = k_list.min()*0.2 , k_list.max()*1.2 - -klims = 0, LL['K_prime'].max()*1.2 - - -for g in MM.beam_group: - MMi = MM.sel(beam_group= g) - plt.plot( MMi.weight.T,MMi.k, '.', color= col_dict[str(g.data)], markersize= 3, linewidth = 0.8) - -plt.xlabel('Power') -plt.ylim(klims) - -ax1 = F.fig.add_subplot(gs[0:2 , 0:-1]) - -for g in MM.beam_group: - Li = LL.loc[str(g.data)] - - angle_list = np.array(Li['alpha']) * 180 /np.pi - kk_list = np.array(Li['K_prime']) - weight_list_i = np.array(Li['weight']) - - plt.scatter( angle_list, kk_list, s= (weight_list_i*8e1)**2 , c=col_dict[str(g.data)], label ='mode ' + str(g.data) ) - - -# lflag= 'paritions ww3' -# for i in np.arange(6): -# i_dir, i_period = Prior.loc['pdp'+ str(i)]['mean'], Prior.loc['ptp'+ str(i)]['mean'] -# i_k = (2 * np.pi/ i_period)**2 / 9.81 -# i_dir = [i_dir -360 if i_dir > 180 else i_dir][0] -# i_dir = [i_dir +360 if i_dir < -180 else i_dir][0] -# -# plt.plot(i_dir, i_k, '.', markersize = 6, color= col.red, label= lflag) -# plt.plot(i_dir, i_k, '-', linewidth = 0.8, color= col.red) -# -# lflag = None - -dir_best[dir_best> 180] = dir_best[dir_best> 180] -360 -plt.plot(dir_best, Pwavenumber , '.r', markersize = 6) - -dir_interp[dir_interp> 180] = dir_interp[dir_interp> 180] -360 -plt.plot(dir_interp, Gk.k, '-', color= 'red', linewidth = 0.3, zorder=11) - - -#ax1.axvline( best_guess * 180/ np.pi , color=col.blue, linewidth = 1.5, label ='best guess fitting') - -# ax1.axvline( (prior_sel['alpha'][0]) * 180 /np.pi, color='k', linewidth = 1.5, label ='prior') -# ax1.axvline( (prior_sel['alpha'][0]- prior_sel['alpha'][1]) * 180 /np.pi, color='k', linewidth = 0.7, label ='prior uncertrainty') -# ax1.axvline( (prior_sel['alpha'][0]+ prior_sel['alpha'][1]) * 180 /np.pi , color='k', linewidth = 0.7) - -plt.fill_betweenx(Gk.k, (dir_interp_smth -spread_smth)* 180 /np.pi, (dir_interp_smth +spread_smth)* 180 /np.pi, zorder= 1, color=col.green1, alpha = 0.2 ) -plt.plot(dir_interp_smth * 180 /np.pi, Gk.k , '.', markersize = 1 , color=col.green1) - -ax1.axvline(85, color='gray', linewidth= 2) -ax1.axvline(-85, color='gray', linewidth= 2) - - -plt.legend() -plt.ylabel('wavenumber (deg)') -plt.xlabel('Angle (deg)') - -#plt.xlim(- 170, 170) -#plt.xlim(- 90, 90) -plt.ylim(klims) - -prior_angle_str =str(np.round( (prior_sel['alpha'][0]) * 180 /np.pi)) -plt.title(track_name + '\nprior=' + prior_angle_str + 'deg', loc= 'left' ) - -plt.xlim( min( [ -90, np.nanmin(dir_best)] ), max( [np.nanmax(dir_best), 90]) ) - - -ax3 = F.fig.add_subplot(gs[2 , 0:-1]) - -for g in MM.beam_group: - MMi = MM.sel(beam_group= g) - wegihted_margins = ( (MMi.marginals * MMi.weight).sum(['x','k'] )/MMi.weight.sum(['x', 'k']) ) - plt.plot( MMi.angle * 180/ np.pi, wegihted_margins , '.', color= col_dict[str(g.data)], markersize= 2, linewidth = 0.8) - -plt.ylabel('Density') -plt.title('weight margins', loc='left') - -#plt.plot(marginal_stack.angle * 180 /np.pi, marginal_stack.T , c=col.gray, label ='weighted mean BF') - -#plt.plot(cost_wmean.angle * 180 /np.pi, cost_wmean , c=col.rascade3, label ='weighted mean BF') -plt.xlim(- 90, 90) -#plt.xlim(- 125, 125) - -ax3 = F.fig.add_subplot(gs[-1 , 0:-1]) - -for g in MM.beam_group: - MMi = MM.sel(beam_group= g) - wegihted_margins = MMi.marginals.mean(['x','k'] )# ( (MMi.marginals * MMi.weight).sum(['x','k'] )/MMi.weight.sum(['x', 'k']) ) - plt.plot( MMi.angle * 180/ np.pi, wegihted_margins , '.', color= col_dict[str(g.data)], markersize= 2, linewidth = 0.8) - -plt.ylabel('Density') -plt.xlabel('Angle (deg)') -plt.title('unweighted margins', loc='left') - -#plt.plot(marginal_stack.angle * 180 /np.pi, marginal_stack.T , c=col.gray, label ='weighted mean BF') - -#plt.plot(cost_wmean.angle * 180 /np.pi, cost_wmean , c=col.rascade3, label ='weighted mean BF') -plt.xlim(- 90, 90) -#plt.xlim(- 125, 125) - -F.save_pup(path= plot_path, name = 'B04_marginal_distributions') - -MT.json_save('B04_success', plot_path, {'time':'time.asctime( time.localtime(time.time()) )'}) diff --git a/pyproject.toml b/pyproject.toml index 8cd27848..bc61136a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,17 +105,22 @@ classifiers = [ # Optional # https://packaging.python.org/discussions/install-requires-vs-requirements/ # noqa # THIS PROJECT NEEDS MPI PRE-installed IN OUR LOCAL SYSTEM dependencies = [ # Optional - "piecewise_regression", - "numba", - "xarray", - "mpi4py", - "icesat2-toolkit==1.0.0.22", - "geopandas", - "sliderule==4.1.0", - "ipyleaflet", - "lmfit", - "tables", - "siphon" + "piecewise_regression>=1.3.0, <2.0.0", + "numba >=0.58.1, <1.0.0", + "xarray >=2024.1.0", # this doesn't look like semmatic versioning + "mpi4py >=3.1.4, <4.0.0", + "numpy >=1.26.3, <2.0.0", + "icesat2-toolkit >=1.0.0.22, <2.0.0", + "pandas >=2.2.0, <3.0.0", + "geopandas >=0.14.2, <1.0.0", + "sliderule >=4.2.3, <5.0.0", + "ipyleaflet >=0.18.1, <0.19.0", + "lmfit >=1.2.2, <2.0.0", + "tables >=3.9.2, <4.0.0", + "emcee >=3.1.4, <4.0.0", + "netCDF4 >=1.6.5, <2.0.0", + "siphon >=0.9, <1.0.0", + "h5py >=3.5.0, < 4.0.0" ] # List additional groups of dependencies here (e.g. development diff --git a/src/icesat2_tracks/ICEsat2_SI_tools/angle_optimizer.py b/src/icesat2_tracks/ICEsat2_SI_tools/angle_optimizer.py index fa167fcc..7ec54aa2 100644 --- a/src/icesat2_tracks/ICEsat2_SI_tools/angle_optimizer.py +++ b/src/icesat2_tracks/ICEsat2_SI_tools/angle_optimizer.py @@ -206,9 +206,11 @@ def test_objective_func(self): def sample(self, fitting_args= None , method='emcee', steps=100, verbose= True, **kargs): fitting_args, fitting_kargs = self.fitting_args, self.fitting_kargs + # TODO: this funciton throws an error in CI. The nan_policy='omit' policiy was added to avoid this issue + # according to the guidelines in https://lmfit.github.io/lmfit-py/faq.html#i-get-errors-from-nan-in-my-fit-what-can-i-do self.fitter = self.LM.minimize(self.objective_func, self.params, method=method, args=fitting_args, kws=fitting_kargs , - nwalkers=self.nwalkers, steps=steps, pos= self.seeds, **kargs) + nwalkers=self.nwalkers, steps=steps, pos= self.seeds,nan_policy='omit' , **kargs) if verbose: print(self.LM.report_fit(self.fitter)) print('results at self.fitter') diff --git a/src/icesat2_tracks/ICEsat2_SI_tools/generalized_FT.py b/src/icesat2_tracks/ICEsat2_SI_tools/generalized_FT.py index 32280284..3e88f22a 100644 --- a/src/icesat2_tracks/ICEsat2_SI_tools/generalized_FT.py +++ b/src/icesat2_tracks/ICEsat2_SI_tools/generalized_FT.py @@ -2,6 +2,7 @@ import icesat2_tracks.ICEsat2_SI_tools.spectral_estimates as spec import icesat2_tracks.ICEsat2_SI_tools.lanczos as lanczos +import matplotlib.pyplot as plt def rebin(data, dk): @@ -14,7 +15,6 @@ def rebin(data, dk): Gmean["k_bins"] = k_low[0:-1] Gmean = Gmean.rename({"k_bins": "k"}) return Gmean, k_low_limits - # define weight function @@ -29,9 +29,13 @@ def smooth_data_to_weight(dd, m=150): dd_fake[2 * m : -2 * m] = dd weight = lanczos.lanczos_filter_1d_wrapping(np.arange(dd_fake.size), dd_fake, m) + weight = weight[2 * m : -2 * m] weight = weight / weight.max() + # ensure non-negative weights + weight[weight < 0.02] = 0.02 + return weight @@ -65,7 +69,7 @@ def get_weights_from_data( pars = Spec_fft.set_parameters(flim=np.sqrt(9.81 * k[-1]) / 2 / np.pi) k_max = (pars["f_max"].value * 2 * np.pi) ** 2 / 9.81 - if method == "gaussian": + if method is "gaussian": # simple gaussian weight def gaus(x, x_0, amp, sigma_g): return amp * np.exp(-0.5 * ((x - x_0) / sigma_g) ** 2) @@ -73,7 +77,7 @@ def gaus(x, x_0, amp, sigma_g): weight = gaus(k, k_max, 1, 0.02) ** (1 / 2) params = None - elif method == "parametric": + elif method is "parametric": # JONSWAP weight f = np.sqrt(9.81 * k) / (2 * np.pi) weight = Spec_fft.create_weight(freq=f, plot_flag=False, max_nfev=max_nfev) @@ -82,14 +86,11 @@ def gaus(x, x_0, amp, sigma_g): Spec_fft.fitter.params.pretty_print() params = Spec_fft.fitter.params - + else: raise ValueError(" 'method' must be either 'gaussian' or 'parametric' ") if plot_flag: - import matplotlib.pyplot as plt - - plt.plot( k_fft[1:], Spec_fft.data, c="gray", label="FFT for Prior", linewidth=0.5 ) @@ -97,7 +98,6 @@ def gaus(x, x_0, amp, sigma_g): k, weight, zorder=12, c="black", label="Fitted model to FFT", linewidth=0.5 ) plt.xlim(k[0], k[-1]) - # add pemnalty floor weight = weight + weight.max() * 0.1 @@ -108,6 +108,18 @@ def gaus(x, x_0, amp, sigma_g): return weight, params +def define_weight_shutter(weight, k, Ncut=3): + "creates masking function to lower high wavenumber weights, Ncut is the numnber by which the spectral peak is multiplied" + # Limit high wavenumbers weights + weight_shutter = weight * 0 + 1 + ki_cut = weight.argmax() * Ncut # k of peak + N_res = weight.size - ki_cut + if N_res < 1: + return weight_shutter + weight_shutter[ki_cut:] = np.linspace(1, 0, N_res) + return weight_shutter + + def make_xarray_from_dict(D, name, dims, coords): import xarray as xr @@ -124,26 +136,26 @@ def define_weights(stancil, prior, x, y, dx, k, max_nfev, plot_flag=False): return weights normalized to 1, prior_pars used for the next iteration """ - if (type(prior[0]) is bool) and not prior[ - 0 - ]: # prior = (False, None), this is the first iteration + if (type(prior[0]) is bool) and not prior[0]: # fit function to data weight, prior_pars = get_weights_from_data( x, y, dx, stancil, k, max_nfev, plot_flag=plot_flag, method="parametric" ) + weight_name = "$P_{init}$ from FFT" - elif ( - type(prior) is tuple - ): + elif type(prior) is tuple: + # combine old and new weights weight = 0.2 * smooth_data_to_weight(prior[0]) + 0.8 * prior[1] - weight_name = "smth. $P_{i-1}$" prior_pars = {"alpha": None, "amp": None, "f_max": None, "gamma": None} - else: + else: # prior = weight, this is all other iterations weight = smooth_data_to_weight(prior) weight_name = "smth. from data" prior_pars = {"alpha": None, "amp": None, "f_max": None, "gamma": None} + # Limit high wavenumbers weights + weight = weight * define_weight_shutter(weight, k, Ncut=3) + if plot_flag: import matplotlib.pyplot as plt @@ -176,9 +188,7 @@ def __init__(self, x, data, L, dx, wavenumber, data_error=None, ov=None): """ self.Lmeters = L - self.ov = ( - int(L / 2) if ov is None else ov - ) # when not defined in create_chunk_boundaries then L/2 + self.ov = int(L / 2) if ov is None else ov self.x = x self.dx = dx @@ -228,9 +238,9 @@ def cal_spectrogram( Lmeters, dk = self.Lmeters, self.dk Lpoints = self.Lpoints Lpoints_full = int(Lmeters / self.dx) + self.xlims = (np.round(X.min()), X.max()) if xlims is None else xlims - # init Lomb scargle object with noise as nummy data () def calc_gFT_apply(stancil, prior): """ windows the data accoding to stencil and applies LS spectrogram @@ -243,11 +253,11 @@ def calc_gFT_apply(stancil, prior): ta = time.perf_counter() x_mask = (stancil[0] <= X) & (X <= stancil[-1]) - print(stancil[1]) x = X[x_mask] if ( - x.size / Lpoints < 0.1 + x.size / Lpoints < 0.40 ): # if there are not enough photos set results to nan + print(" -- data density to low, skip stancil") return { "stancil_center": stancil[1], "p_hat": np.concatenate([self.k * np.nan, self.k * np.nan]), @@ -265,7 +275,6 @@ def calc_gFT_apply(stancil, prior): y_var = y.var() FT = generalized_Fourier(x, y, self.k) - if plot_flag: import matplotlib.pyplot as plt @@ -282,14 +291,39 @@ def calc_gFT_apply(stancil, prior): # define error err = ERR[x_mask] if ERR is not None else 1 - print("weights : ", time.perf_counter() - ta) + print("compute time weights : ", time.perf_counter() - ta) + ta = time.perf_counter() - FT.define_problem(weight, err) # 1st arg is Penalty, 2nd is error + FT.define_problem(weight, err) # solve problem: p_hat = FT.solve() - print("solve : ", time.perf_counter() - ta) + if np.isnan(np.mean(p_hat)): + print(" -- inversion nan!") + print(" -- data fraction", x.size / Lpoints) + print( + " -- weights:", + np.mean(weight), + "err:", + np.mean(err), + "y:", + np.mean(y), + ) + print(" -- skip stancil") + return { + "stancil_center": stancil[1], + "p_hat": np.concatenate([self.k * np.nan, self.k * np.nan]), + "inverse_stats": np.nan, + "y_model_grid": np.nan, + "y_data_grid": np.nan, + "x_size": x.size, + "PSD": False, + "weight": False, + "spec_adjust": np.nan, + } + + print("compute time solve : ", time.perf_counter() - ta) ta = time.perf_counter() x_pos = (np.round((x - stancil[0]) / self.dx, 0)).astype("int") @@ -301,7 +335,7 @@ def calc_gFT_apply(stancil, prior): y_data_grid = np.copy(eta) * np.nan y_data_grid[x_pos] = y - inverse_stats = FT.get_stats(self.dk, Lpoints_full, print_flag=True) + inverse_stats = FT.get_stats(self.dk, Lpoints_full, print_flag=plot_flag) # add fitting parameters of Prior to stats dict for k, I in prior_pars.items(): try: @@ -309,19 +343,17 @@ def calc_gFT_apply(stancil, prior): except: inverse_stats[k] = np.nan - print("stats : ", time.perf_counter() - ta) - - # multiply with the standard deviation of the data to get dimensions right - PSD = power_from_model( - p_hat, dk, self.k.size, x.size, Lpoints - ) + print("compute time stats : ", time.perf_counter() - ta) - if self.k.size * 2 > x.size: - col = "red" - else: - col = "blue" + # multiply with the standard deviation of the data to get dimensions right + PSD = power_from_model(p_hat, dk, self.k.size, x.size, Lpoints) if plot_flag: + if self.k.size * 2 > x.size: + col = "red" + else: + col = "blue" + plt.plot(self.k, PSD, color=col, label="GFT fit", linewidth=0.5) plt.title( "non-dim Spectral Segment Models, 2M=" @@ -357,28 +389,34 @@ def calc_gFT_apply(stancil, prior): return return_dict - # % derive L2 stancil - self.stancil_iter = spec.create_chunk_boundaries_unit_lengths( - Lmeters, self.xlims, ov=self.ov, iter_flag=True + # derive L2 stancil + self.stancil_iter_list = spec.create_chunk_boundaries_unit_lengths( + Lmeters, self.xlims, ov=self.ov, iter_flag=False ) - + self.stancil_iter = iter(self.stancil_iter_list.T.tolist()) + # apply func to all stancils Spec_returns = list() + # form: PSD_from_GFT, weight_used in inversion prior = False, False + N_stencil = len(self.stancil_iter_list.T) + Ni = 1 for ss in copy.copy(self.stancil_iter): + print(Ni, "/", N_stencil, "Stancils") + # prior step if prior[0] is False: # make NL fit of piors do not exist - print("1st step with NL-fit") + print("1st step: with NL-fit") I_return = calc_gFT_apply(ss, prior=prior) - prior = I_return["PSD"], I_return["weight"] - + prior = I_return["PSD"], I_return["weight"] # 2nd step if prior[0] is False: - print("priors still false skip 2nd step") + print("1st GFD failed (priors[0]=false), skip 2nd step") else: - print("2nd step use set priors:", type(prior[0]), type(prior[0])) + print("2nd step: use set priors:", type(prior[0]), type(prior[1])) + print(prior[0][0:3], prior[1][0:3]) I_return = calc_gFT_apply(ss, prior=prior) - prior = I_return["PSD"], I_return["weight"] + prior = I_return["PSD"], I_return["weight"] Spec_returns.append( dict( @@ -395,7 +433,9 @@ def calc_gFT_apply(stancil, prior): ) ) ) - + + Ni += 1 + # unpack resutls of the mapping: GFT_model = dict() Z_model = dict() @@ -441,28 +481,26 @@ def calc_gFT_apply(stancil, prior): self.N_per_stancil = N_per_stancil chunk_positions = np.array(list(D_specs.keys())) - self.N_stancils = len(chunk_positions) # number of spectral realizations + self.N_stancils = len(chunk_positions) # number of spectral realizatiobs # repack data, create xarray # 1st LS spectal estimates - G_power_data = make_xarray_from_dict( D_specs, "gFT_PSD_data", ["k"], {"k": self.k} ) - G_power_data = xr.concat(G_power_data.values(), dim="x").T # .to_dataset() + G_power_data = xr.concat(G_power_data.values(), dim="x").T G_power_model = make_xarray_from_dict( D_specs_model, "gFT_PSD_model", ["k"], {"k": self.k} ) - G_power_model = xr.concat(G_power_model.values(), dim="x").T # .to_dataset() + G_power_model = xr.concat(G_power_model.values(), dim="x").T self.G = G_power_model self.G.name = "gFT_PSD_model" - # 2nd FFT(Y_model) G_model_Z = make_xarray_from_dict(Z_model, "Z_hat", ["k"], {"k": self.k}) - G_model_Z = xr.concat(G_model_Z.values(), dim="x").T + G_model_Z = xr.concat(G_model_Z.values(), dim="x").T # 3rd GFT_model_coeff_A = dict() @@ -475,24 +513,19 @@ def calc_gFT_apply(stancil, prior): I[1], dims=["k"], coords={"k": self.k, "x": xi}, name="gFT_sin_coeff" ) - GFT_model_coeff_A = xr.concat( - GFT_model_coeff_A.values(), dim="x" - ).T - GFT_model_coeff_B = xr.concat( - GFT_model_coeff_B.values(), dim="x" - ).T + GFT_model_coeff_A = xr.concat(GFT_model_coeff_A.values(), dim="x").T + GFT_model_coeff_B = xr.concat(GFT_model_coeff_B.values(), dim="x").T # add weights to the data weights_k = make_xarray_from_dict(weights, "weight", ["k"], {"k": self.k}) - weights_k = xr.concat(weights_k.values(), dim="x").T + weights_k = xr.concat(weights_k.values(), dim="x").T # .to_dataset() - # 4th: model in real space eta = np.arange(0, self.Lmeters + self.dx, self.dx) - self.Lmeters / 2 y_model_eta = make_xarray_from_dict(y_model, "y_model", ["eta"], {"eta": eta}) y_data_eta = make_xarray_from_dict(y_data, "y_data", ["eta"], {"eta": eta}) - y_model_eta = xr.concat(y_model_eta.values(), dim="x").T - y_data_eta = xr.concat(y_data_eta.values(), dim="x").T + y_model_eta = xr.concat(y_model_eta.values(), dim="x").T + y_data_eta = xr.concat(y_data_eta.values(), dim="x").T # merge wavenumber datasets self.GG = xr.merge( @@ -517,7 +550,7 @@ def calc_gFT_apply(stancil, prior): for k, I in Pars.items(): if I is not np.nan: PP2[k] = I - + keys = Pars[next(iter(PP2))].keys() keys_DF = list(set(keys) - set(["model_error_k", "model_error_x"])) params_dataframe = pd.DataFrame(index=keys_DF) @@ -542,19 +575,18 @@ def calc_gFT_apply(stancil, prior): ) sta, ste = xi - self.Lmeters / 2, xi + self.Lmeters / 2 + x_pos = (np.round((X[(sta <= X) & (X <= ste)] - sta) / self.dx)).astype( "int" ) x_err = np.copy(eta) * np.nan - # check sizes and adjust if necessary. + # check sizes and adjust if necessary. if x_pos.size > I["model_error_x"].size: x_pos = x_pos[0 : I["model_error_x"].size] print("adjust x") elif x_pos.size < I["model_error_x"].size: - I["model_error_x"] = I["model_error_x"][ - 0:-1 - ] + I["model_error_x"] = I["model_error_x"][0:-1] print("adjust y") x_err[x_pos] = I["model_error_x"] @@ -605,6 +637,7 @@ def parceval(self, add_attrs=True, weight_data=False): import copy DATA = self.data + L = self.Lmeters X = self.x def get_stancil_var_apply(stancil): @@ -712,7 +745,7 @@ def power_from_model(p_hat, dk, M, N_x, N_x_full): """ Z = complex_represenation(p_hat, M, N_x_full) - spec, _ = Z_to_power_gFT(Z, dk, N_x, N_x_full) + spec, _ = Z_to_power_gFT(Z, dk, N_x, N_x_full) # use spec_incomplete # spectral density respesenting the incomplete data return spec @@ -783,7 +816,6 @@ def solve(self): return p_hat - def model(self): "returns the model dimensional units" if "p_hat" not in self.__dict__: @@ -831,8 +863,9 @@ def parceval(self, dk, Nx_full): return pars def get_stats(self, dk, Nx_full, print_flag=False): - residual = self.ydata - self.model() + + Lmeters = self.x[-1] - self.x[0] pars = { "data_var": self.ydata.var(), "model_var": self.model().var(), @@ -863,8 +896,6 @@ def get_stats(self, dk, Nx_full, print_flag=False): class get_prior_spec: def __init__(self, freq, data): - """ """ - import numpy as np import lmfit as LM self.LM = LM @@ -969,7 +1000,7 @@ def runningmean(self, var, m, tailcopy=False): print("0 Dimension is smaller then averaging length") return rr = np.asarray(var) * np.nan - + var_range = np.arange(m, int(s[0]) - m - 1, 1) for i in var_range[np.isfinite(var[m : int(s[0]) - m - 1])]: rr[int(i)] = np.nanmean(var[i - m : i + m]) diff --git a/src/icesat2_tracks/ICEsat2_SI_tools/sliderule_converter_tools.py b/src/icesat2_tracks/ICEsat2_SI_tools/sliderule_converter_tools.py index 31b274bf..a370c3cc 100644 --- a/src/icesat2_tracks/ICEsat2_SI_tools/sliderule_converter_tools.py +++ b/src/icesat2_tracks/ICEsat2_SI_tools/sliderule_converter_tools.py @@ -1,4 +1,3 @@ - from ipyleaflet import basemaps @@ -7,46 +6,62 @@ def correct_and_remove_height(Gi, height_limit): """ corrects the height and removes the points with height +- height_limit """ - h_mean_corrected = (Gi['h_mean'] - Gi['dem_h']) - false_height_mask = abs(h_mean_corrected) > height_limit - Gi['h_mean']=h_mean_corrected - Gi.drop(columns=['dem_h'], inplace=True) + h_mean_corrected = Gi["h_mean"] - Gi["dem_h"] + false_height_mask = abs(h_mean_corrected) > height_limit + Gi["h_mean"] = h_mean_corrected + Gi.drop(columns=["dem_h"], inplace=True) return Gi[~false_height_mask] - # polygon tools def make_plot_polygon(poly_test, color="green"): - """ create a plot polygon from the given coordinates""" + """create a plot polygon from the given coordinates""" - from ipyleaflet import Polygon - bb = [poly_test[0]['lon'], poly_test[0]['lat'], - poly_test[2]['lon'], poly_test[2]['lat']] - return Polygon( - locations=poly_test, - color=color) + from ipyleaflet import Polygon + + bb = [ + poly_test[0]["lon"], + poly_test[0]["lat"], + poly_test[2]["lon"], + poly_test[2]["lat"], + ] + return Polygon(locations=poly_test, color=color) # # make function to plot with basemaps def plot_polygon(poly_test, basemap=basemaps.Esri.WorldImagery, zoom=3): - """ plots polygon in the map""" - - from ipyleaflet import Map, GeoData, LayersControl,Rectangle, basemaps, basemap_to_tiles, TileLayer, SplitMapControl, Polygon + """plots polygon in the map""" + + from ipyleaflet import ( + Map, + GeoData, + LayersControl, + Rectangle, + basemaps, + basemap_to_tiles, + TileLayer, + SplitMapControl, + Polygon, + ) # icepx will want a bounding box with LL lon/lat, UR lon/lat import numpy as np polygon_plot = make_plot_polygon(poly_test, color="green") - center = [np.mean(list(set([p['lat'] for p in poly_test]))) - , np.mean(list(set([p['lon'] for p in poly_test])))] + center = [ + np.mean(list(set([p["lat"] for p in poly_test]))), + np.mean(list(set([p["lon"] for p in poly_test]))), + ] m = Map(basemap=basemap, center=center, zoom=zoom) - m.add_layer(polygon_plot); + m.add_layer(polygon_plot) return m + # gemetric tools def haversine(lon1, lat1, lon2, lat2, arc=False): from math import radians, cos, sin, asin, sqrt + """ Calculate the great circle distance between two points on the earth (specified in decimal degrees) @@ -59,61 +74,73 @@ def haversine(lon1, lat1, lon2, lat2, arc=False): # haversine formula dlon = lon2 - lon1 dlat = lat2 - lat1 - a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2 + a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2 c = 2 * asin(sqrt(a)) - r = 6371 # Radius of earth in kilometers. Use 3956 for miles + r = 6371 # Radius of earth in kilometers. Use 3956 for miles if arc: return c else: return c * r + # point closest to the equator in polygon def get_min_eq_dist(ppoly): """ returns the minimum distance of the polygon to the equator in meters """ - min_eq_dist= list() + min_eq_dist = list() for point in ppoly: - point['x_atc'] = haversine(point['lon'], 0, point['lon'], point['lat'], arc=False) - min_eq_dist.append(point['x_atc']) - return min(min_eq_dist)*1e3 # in meters. needed for defining the x-axis - - + point["x_atc"] = haversine( + point["lon"], 0, point["lon"], point["lat"], arc=False + ) + min_eq_dist.append(point["x_atc"]) + return min(min_eq_dist) * 1e3 # in meters. needed for defining the x-axis def create_polygons(latR, lonR): from shapely.geometry import Polygon, Point + latR.sort() lonR.sort() - poly_list=[{'lat':latR[ii], 'lon':lonR[jj]} for ii, jj in zip([1, 1, 0, 0, 1], [1, 0, 0, 1, 1])] - polygon_shapely= Polygon([(item['lon'], item['lat']) for item in poly_list]) - return { 'list':poly_list, - 'shapely':polygon_shapely, - 'lons':lonR, - 'lats':latR} + poly_list = [ + {"lat": latR[ii], "lon": lonR[jj]} + for ii, jj in zip([1, 1, 0, 0, 1], [1, 0, 0, 1, 1]) + ] + polygon_shapely = Polygon([(item["lon"], item["lat"]) for item in poly_list]) + return {"list": poly_list, "shapely": polygon_shapely, "lons": lonR, "lats": latR} + # RGT table manipulations: def find_lowest_point_on_RGT(Gs, RGT): - ipos = abs(Gs[Gs['RGT'] == RGT].geometry.y).argmin() - return Gs[Gs['RGT'] == RGT].iloc[ipos] + ipos = abs(Gs[Gs["RGT"] == RGT].geometry.y).argmin() + return Gs[Gs["RGT"] == RGT].iloc[ipos] + def find_highest_point_on_RGT(Gs, RGT): - ipos = abs(Gs[Gs['RGT'] == RGT].geometry.y).argmax() - return Gs[Gs['RGT'] == RGT].iloc[ipos] + ipos = abs(Gs[Gs["RGT"] == RGT].geometry.y).argmax() + return Gs[Gs["RGT"] == RGT].iloc[ipos] -#find_lowest_point_on_RGT(Gs, 2).geometry -def get_RGT_start_points(Gs, RGT ='RGT'): + +# find_lowest_point_on_RGT(Gs, 2).geometry +def get_RGT_start_points(Gs, RGT="RGT"): import pandas as pd import geopandas as gpd - G_lowest = pd.concat([find_lowest_point_on_RGT(Gs, rgt).T for rgt in Gs[RGT].unique()], axis=1).T.reset_index() - G_lowest = G_lowest.T.drop('index').T + + G_lowest = pd.concat( + [find_lowest_point_on_RGT(Gs, rgt).T for rgt in Gs[RGT].unique()], axis=1 + ).T.reset_index() + G_lowest = G_lowest.T.drop("index").T return gpd.GeoDataFrame(G_lowest, crs=Gs.crs) -def get_RGT_end_points(Gs, RGT= 'RGT'): + +def get_RGT_end_points(Gs, RGT="RGT"): import pandas as pd import geopandas as gpd - G_lowest = pd.concat([find_highest_point_on_RGT(Gs, rgt).T for rgt in Gs[RGT].unique()], axis=1).T.reset_index() - G_lowest = G_lowest.T.drop('index').T + + G_lowest = pd.concat( + [find_highest_point_on_RGT(Gs, rgt).T for rgt in Gs[RGT].unique()], axis=1 + ).T.reset_index() + G_lowest = G_lowest.T.drop("index").T return gpd.GeoDataFrame(G_lowest, crs=Gs.crs) @@ -124,24 +151,21 @@ def ascending_test(track_data): # get the first and last point first_point = abs(track_data.iloc[0].geometry.y) last_point = abs(track_data.iloc[-1].geometry.y) - # get the time - # first_time = cGPS.convert_GPS_time(first_point.index, first_point['rgt'], first_point['orbit_number']) - # last_time = cGPS.convert_GPS_time(last_point['delta_time'], last_point['rgt'], last_point['orbit_number']) - # test if ascending or descending if first_point < last_point: return True else: return False + def ascending_test_distance(track_data): """ test if the track is ascending or descending based on 'x_atc' column """ # get the first and last point - first_point = abs(track_data.iloc[track_data['x_atc'].argmin()].geometry.y) - last_point = abs(track_data.iloc[track_data['x_atc'].argmax()].geometry.y) - + first_point = abs(track_data.iloc[track_data["x_atc"].argmin()].geometry.y) + last_point = abs(track_data.iloc[track_data["x_atc"].argmax()].geometry.y) + # test if ascending or descending if first_point < last_point: return True else: @@ -149,7 +173,9 @@ def ascending_test_distance(track_data): ## tools for calculating minimal distance to the equator -def define_reference_distance_with_RGT(Gtrack_lowest, rgt, acending, ground_track_length = 20160e3): +def define_reference_distance_with_RGT( + Gtrack_lowest, rgt, acending, ground_track_length=20160e3 +): """ returns reference distance for the given 'rgt' groundtrack based on table of lowest points (Gtrack_lowest). This calculated the distance from the equator to the lowest point of the ground track. @@ -162,30 +188,35 @@ def define_reference_distance_with_RGT(Gtrack_lowest, rgt, acending, ground_trac ground_track_length: the length of the ground track in meters. Default is 20160e3 (the length of a half orbit) this is not very exact """ - start_point = Gtrack_lowest[Gtrack_lowest['RGT'] == rgt] + start_point = Gtrack_lowest[Gtrack_lowest["RGT"] == rgt] if len(start_point) == 0: - raise ValueError('RGT '+ str(rgt) +', no data in ground track') - - start_point_single =start_point.iloc[0] - start_point_dist = 1e3 * haversine(start_point_single.geometry.x, 0,start_point_single.geometry.x, start_point_single.geometry.y, arc=False) + raise ValueError("RGT " + str(rgt) + ", no data in ground track") + + start_point_single = start_point.iloc[0] + start_point_dist = 1e3 * haversine( + start_point_single.geometry.x, + 0, + start_point_single.geometry.x, + start_point_single.geometry.y, + arc=False, + ) if start_point.geometry.y.iloc[0] < 0: # SH if acending: - start_point_dist= start_point_dist + ground_track_length + start_point_dist = start_point_dist + ground_track_length else: - start_point_dist = 2 * ground_track_length - start_point_dist + start_point_dist = 2 * ground_track_length - start_point_dist else: # NH if acending: - start_point_dist= start_point_dist + start_point_dist = start_point_dist else: start_point_dist = ground_track_length - start_point_dist return start_point_dist, start_point - def plot_reference_point_coordinates(tmp, start_point_dist, start_point): """ plots the reference point and the data @@ -196,29 +227,39 @@ def plot_reference_point_coordinates(tmp, start_point_dist, start_point): """ import matplotlib.pyplot as plt - rgt= tmp['rgt'].unique()[0] - spoint_color = 'black' - data_col= 'blue' - accending_color= 'black' if ascending_test_distance(tmp) else 'red' + rgt = tmp["rgt"].unique()[0] + spoint_color = "black" + data_col = "blue" + accending_color = "black" if ascending_test_distance(tmp) else "red" - fig, axx = plt.subplots(1,2, figsize=(8,4)) + fig, axx = plt.subplots(1, 2, figsize=(8, 4)) ax = axx[0] - ax.plot(tmp.geometry.x, tmp.geometry.y, '.', markersize=0.5, c=data_col) + ax.plot(tmp.geometry.x, tmp.geometry.y, ".", markersize=0.5, c=data_col) - ax.plot(start_point.geometry.x, start_point.geometry.y, '.', c = spoint_color) + ax.plot(start_point.geometry.x, start_point.geometry.y, ".", c=spoint_color) ax.grid() - ax.set_xlabel('Longitude') - ax.set_ylabel('Latitude') - ax.set_title('RGT: ' + str(rgt), loc='left') + ax.set_xlabel("Longitude") + ax.set_ylabel("Latitude") + ax.set_title("RGT: " + str(rgt), loc="left") ax = axx[1] - ax.plot(tmp['x_atc'], tmp.geometry.y ,'.', markersize=0.5, c=data_col, label ='data') - ax.plot(start_point_dist, start_point.geometry.y ,'.', c =spoint_color, label='reference point') - - ax.set_xlabel('Distance from Equator') - ax.set_ylabel('Latitude') - ax.set_title('acending: ' + str(ascending_test_distance(tmp)), loc='right', color=accending_color) + ax.plot(tmp["x_atc"], tmp.geometry.y, ".", markersize=0.5, c=data_col, label="data") + ax.plot( + start_point_dist, + start_point.geometry.y, + ".", + c=spoint_color, + label="reference point", + ) + + ax.set_xlabel("Distance from Equator") + ax.set_ylabel("Latitude") + ax.set_title( + "acending: " + str(ascending_test_distance(tmp)), + loc="right", + color=accending_color, + ) plt.legend() ax.grid() return fig, axx @@ -226,58 +267,98 @@ def plot_reference_point_coordinates(tmp, start_point_dist, start_point): def plot_data_in_domain(gdf2, polygon_list): """ - makes two panel figure: + makes two panel figure: - on the left: plot the lon-lat photon postions for all track & the box - on the right: the hoffmoeller diagram with lon-time positions at the most southern points. Colors show ascending or descending tracks. inputs: gdf: GeoDataFrame with the down """ import matplotlib.pyplot as plt + # make two panel figure, on the left plot the photon postions, on the the hoffmoeller diagram - fig, axx = plt.subplots(1,2, figsize=(8,4)) + fig, axx = plt.subplots(1, 2, figsize=(8, 4)) - plt.suptitle('Data overview and distribution over domain') + plt.suptitle("Data overview and distribution over domain") ax = axx[0] ax.set_title("ATL Points") - ax.set_aspect('equal') + ax.set_aspect("equal") # FIX THIS TO SHOW EACH RGT value by color - gdf2.plot(ax=ax, column='rgt', label='RGT', c=gdf2['rgt'], markersize= 0.5, cmap='winter') - #ax.legend(loc="upper left") + gdf2.plot( + ax=ax, column="rgt", label="RGT", c=gdf2["rgt"], markersize=0.5, cmap="winter" + ) # Prepare coordinate lists for plotting the region of interest polygon region_lon = [e["lon"] for e in polygon_list] region_lat = [e["lat"] for e in polygon_list] - ax.plot(region_lon, region_lat, linewidth=1, color='green'); + ax.plot(region_lon, region_lat, linewidth=1, color="green") # labels - ax.set_xlabel('Longitude') - ax.set_ylabel('Latitude') + ax.set_xlabel("Longitude") + ax.set_ylabel("Latitude") # add a red dot for each 1st point in the RGT - for rgt in gdf2['rgt'].unique()[0:20]: - tmp = gdf2[gdf2['rgt']==rgt] + for rgt in gdf2["rgt"].unique()[0:20]: + tmp = gdf2[gdf2["rgt"] == rgt] # 1st point - ax.plot(tmp.iloc[0].geometry.x, tmp.iloc[0].geometry.y, 'rv', markersize= 5, label='RGT') + ax.plot( + tmp.iloc[0].geometry.x, + tmp.iloc[0].geometry.y, + "rv", + markersize=5, + label="RGT", + ) # last point - ax.plot(tmp.iloc[-1].geometry.x, tmp.iloc[-1].geometry.y, 'ks', markersize= 5, label='RGT') + ax.plot( + tmp.iloc[-1].geometry.x, + tmp.iloc[-1].geometry.y, + "ks", + markersize=5, + label="RGT", + ) # line between first and last if ascending_test_distance(tmp): - axx[1].plot(tmp.geometry.x.mean(), tmp.index.mean(), 'ko', markersize= 5, label='ascending') - ax.plot([tmp.iloc[0].geometry.x, tmp.iloc[-1].geometry.x], [tmp.iloc[0].geometry.y, tmp.iloc[-1].geometry.y], '-', color='black', linewidth=1) + axx[1].plot( + tmp.geometry.x.mean(), + tmp.index.mean(), + "ko", + markersize=5, + label="ascending", + ) + ax.plot( + [tmp.iloc[0].geometry.x, tmp.iloc[-1].geometry.x], + [tmp.iloc[0].geometry.y, tmp.iloc[-1].geometry.y], + "-", + color="black", + linewidth=1, + ) else: - axx[1].plot( tmp.geometry.x.mean(), tmp.index.mean(), 'o', color ='orange', markersize= 5, zorder=10, label='decending') - ax.plot([tmp.iloc[0].geometry.x, tmp.iloc[-1].geometry.x], [tmp.iloc[0].geometry.y, tmp.iloc[-1].geometry.y], '-', color='orange', linewidth=2.5) - + axx[1].plot( + tmp.geometry.x.mean(), + tmp.index.mean(), + "o", + color="orange", + markersize=5, + zorder=10, + label="decending", + ) + ax.plot( + [tmp.iloc[0].geometry.x, tmp.iloc[-1].geometry.x], + [tmp.iloc[0].geometry.y, tmp.iloc[-1].geometry.y], + "-", + color="orange", + linewidth=2.5, + ) ax = axx[1] - ax.plot(gdf2.geometry.x , gdf2.index, '.', markersize=0.5) + ax.plot(gdf2.geometry.x, gdf2.index, ".", markersize=0.5) # add labels - ax.set_xlabel('Longitude') - ax.set_ylabel('time') + ax.set_xlabel("Longitude") + ax.set_ylabel("time") return fig, axx + def check_RGT_in_domain(Gtrack_lowest, gdf): """ checks if the RGT in the gdf are found in Gtrack list @@ -288,19 +369,18 @@ def check_RGT_in_domain(Gtrack_lowest, gdf): result: set of RGT that are in both Gtrack_lowest and gdf """ import collections - #result = collections.Counter(gdf2['rgt']) & collections.Counter(Gtrack_lowest['RGT']) - gdf_list = list(gdf['rgt'].unique()) - result = set(gdf_list).intersection(set(Gtrack_lowest['RGT'])) - #result = set(gdf2['rgt']).intersection(set([0])) + + gdf_list = list(gdf["rgt"].unique()) + result = set(gdf_list).intersection(set(Gtrack_lowest["RGT"])) interect_list = list(result) interect_list.sort() - print('RGT in domain: ', len(Gtrack_lowest['RGT'].unique())) - print('RGT with data found: ', len(gdf_list)) - print('RGT in both: ', len(interect_list) ) + print("RGT in domain: ", len(Gtrack_lowest["RGT"].unique())) + print("RGT with data found: ", len(gdf_list)) + print("RGT in both: ", len(interect_list)) if len(interect_list) != len(gdf_list): - print('RGT not in both: ', list(set(gdf_list) - result) ) + print("RGT not in both: ", list(set(gdf_list) - result)) return interect_list @@ -319,19 +399,20 @@ def define_x_coordinate_in_polygon(table_data, polygon, round=True): table_data: GeoDataFrame with the data and the 'x' coordinate """ if round: - min_eq_dist = np.round(get_min_eq_dist(polygon)/1e2)*1e2 + min_eq_dist = np.round(get_min_eq_dist(polygon) / 1e2) * 1e2 else: min_eq_dist = np.round(get_min_eq_dist(polygon)) - + if ascending_test(table_data): - table_data['x'] = table_data['x_atc'] - min_eq_dist + table_data["x"] = table_data["x_atc"] - min_eq_dist else: - table_data['x'] = ((np.pi * 6371*1e3) - min_eq_dist) - table_data['x_atc'] + # print('descending') + table_data["x"] = ((np.pi * 6371 * 1e3) - min_eq_dist) - table_data["x_atc"] return table_data + def define_x_coordinate_with_RGT(table_data, Gtrack_lowest): - """ returns table with 'x' coordinate for a reference point based on the RGT intersection with the polygon taken from 'Gtrack_lowest' the reference point is defined as the most equatorward point of the polygon. @@ -341,21 +422,24 @@ def define_x_coordinate_with_RGT(table_data, Gtrack_lowest): returns: table_data: GeoDataFrame with the data and the 'x' coordinate """ - if len(table_data['rgt'].unique()) != 1: - raise ValueError('table_data must contain only one RGT') + if len(table_data["rgt"].unique()) != 1: + raise ValueError("table_data must contain only one RGT") - rgt = table_data['rgt'].unique()[0] + rgt = table_data["rgt"].unique()[0] acending = ascending_test_distance(table_data) - - start_point_dist, start_point = define_reference_distance_with_RGT(Gtrack_lowest, rgt, acending) + + start_point_dist, start_point = define_reference_distance_with_RGT( + Gtrack_lowest, rgt, acending + ) if acending: - table_data['x'] = table_data['x_atc'] - start_point_dist + table_data["x"] = table_data["x_atc"] - start_point_dist else: - table_data['x'] = start_point_dist - table_data['x_atc'] + table_data["x"] = start_point_dist - table_data["x_atc"] return table_data + def define_x_coordinate_from_data(table_data): """ define x coordinate from data @@ -365,14 +449,14 @@ def define_x_coordinate_from_data(table_data): table_data: GeoDataFrame with the data and x coordinate """ acending = ascending_test_distance(table_data) - + if acending: - start_point_dist = table_data['x_atc'].min() - table_data['x'] = table_data['x_atc'] - start_point_dist + start_point_dist = table_data["x_atc"].min() + table_data["x"] = table_data["x_atc"] - start_point_dist else: - start_point_dist = table_data['x_atc'].max() - table_data['x'] = start_point_dist - table_data['x_atc'] - table_data.sort_values(by='x', inplace=True) + start_point_dist = table_data["x_atc"].max() + table_data["x"] = start_point_dist - table_data["x_atc"] + table_data.sort_values(by="x", inplace=True) table_data.reset_index(inplace=True) return table_data @@ -385,6 +469,16 @@ def create_ID_name(tt, segment=None): returns: ID_name: string with the ID name """ - hemis= 'NH' if tt.geometry.y > 0 else 'SH' - segment = '00' if segment is None else str(segment) - return hemis + '_' + str(tt.name.year) + str(tt.name.month).zfill(2) + str(tt.name.day).zfill(2) + '_' + str(tt.rgt).zfill(4) + str(tt.cycle).zfill(2) + segment \ No newline at end of file + hemis = "NH" if tt.geometry.y > 0 else "SH" + segment = "00" if segment is None else str(segment) + return ( + hemis + + "_" + + str(tt.name.year) + + str(tt.name.month).zfill(2) + + str(tt.name.day).zfill(2) + + "_" + + str(tt.rgt).zfill(4) + + str(tt.cycle).zfill(2) + + segment + ) diff --git a/src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py b/src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py index 0f04f155..b558fd3f 100644 --- a/src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py +++ b/src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py @@ -44,10 +44,7 @@ ATL03_track_name = "ATL03_" + track_name + ".h5" # Configure SL Session -sliderule.authenticate("brown", ps_username="mhell", ps_password="Oijaeth9quuh") -icesat2.init( - "slideruleearth.io", organization="brown", desired_nodes=1, time_to_live=90 -) # minutes +icesat2.init("slideruleearth.io") # plot the ground tracks in geographic location diff --git a/src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py b/src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py index b23dcfea..00da0ff5 100644 --- a/src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py +++ b/src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py @@ -3,30 +3,35 @@ This is python 3 """ -import copy -import datetime -import h5py -import time import sys +import matplotlib.pyplot as plt +from icesat2_tracks.config.IceSAT2_startup import mconfig + +from threadpoolctl import threadpool_info, threadpool_limits +from pprint import pprint import numpy as np import xarray as xr -from pprint import pprint -from scipy.ndimage.measurements import label -from threadpoolctl import threadpool_info, threadpool_limits -import icesat2_tracks.ICEsat2_SI_tools.generalized_FT as gFT +import h5py import icesat2_tracks.ICEsat2_SI_tools.io as io import icesat2_tracks.ICEsat2_SI_tools.spectral_estimates as spec -import icesat2_tracks.local_modules.m_general_ph3 as M -import icesat2_tracks.local_modules.m_spectrum_ph3 as spicke_remover + +import time +import imp +import copy +import icesat2_tracks.ICEsat2_SI_tools.spicke_remover as spicke_remover +import datetime +import icesat2_tracks.ICEsat2_SI_tools.generalized_FT as gFT +from scipy.ndimage import label import icesat2_tracks.local_modules.m_tools_ph3 as MT -from icesat2_tracks.config.IceSAT2_startup import mconfig, plt +from icesat2_tracks.local_modules import m_general_ph3 as M + import tracemalloc -def linear_gap_fill(F, key_lead, key_int): +def linear_gap_fill(F,key_lead, key_int): """ F pd.DataFrame key_lead key in F that determined the independent coordindate @@ -43,9 +48,7 @@ def linear_gap_fill(F, key_lead, key_int): track_name, batch_key, test_flag = io.init_from_input( sys.argv ) # loads standard experiment - hemis, batch = batch_key.split("_") - ATlevel = "ATL03" load_path = mconfig["paths"]["work"] + "/" + batch_key + "/B01_regrid/" @@ -62,28 +65,26 @@ def linear_gap_fill(F, key_lead, key_int): + batch_key + "/" + track_name - + "/B_spectra/" + + "/B03_spectra/" ) MT.mkdirs_r(plot_path) MT.mkdirs_r(save_path) bad_track_path = mconfig["paths"]["work"] + "bad_tracks/" + batch_key + "/" - all_beams = mconfig["beams"]["all_beams"] high_beams = mconfig["beams"]["high_beams"] low_beams = mconfig["beams"]["low_beams"] N_process = 4 +print("N_process=", N_process) - -# open with hdf5 Gd = h5py.File(load_path + "/" + track_name + "_B01_binned.h5", "r") # test amount of nans in the data nan_fraction = list() for k in all_beams: - heights_c_std = io.get_beam_var_hdf_store(Gd[k], "dist") + heights_c_std = io.get_beam_var_hdf_store(Gd[k], "x") nan_fraction.append(np.sum(np.isnan(heights_c_std)) / heights_c_std.shape[0]) del heights_c_std @@ -95,7 +96,7 @@ def linear_gap_fill(F, key_lead, key_int): Ib = Gd[group[1]] ratio = Ia["x"][:].size / Ib["x"][:].size if (ratio > 10) | (ratio < 0.1): - # print("bad data ratio ", ratio, 1 / ratio) # TODO: add logger + print("bad data ratio ", ratio, 1 / ratio) bad_ratio_flag = True if (np.array(nan_fraction).mean() > 0.95) | bad_ratio_flag: @@ -114,10 +115,12 @@ def linear_gap_fill(F, key_lead, key_int): exit() # test LS with an even grid where missing values are set to 0 +imp.reload(spec) print(Gd.keys()) Gi = Gd[list(Gd.keys())[0]] # to select a test beam -dist = io.get_beam_var_hdf_store(Gd[list(Gd.keys())[0]], "dist") +dist = io.get_beam_var_hdf_store(Gd[list(Gd.keys())[0]], "x") +# make dataframe form hdf5 # derive spectal limits # Longest deserved period: T_max = 40 # sec @@ -129,6 +132,7 @@ def linear_gap_fill(F, key_lead, key_int): Lpoints = int(np.round(min_datapoint) * 10) Lmeters = Lpoints * dx + print("L number of gridpoint:", Lpoints) print("L length in km:", Lmeters / 1e3) print("approx number windows", 2 * dist.iloc[-1] / Lmeters - 1) @@ -137,6 +141,7 @@ def linear_gap_fill(F, key_lead, key_int): lambda_min = 9.81 * T_min**2 / (2 * np.pi) flim = 1 / T_min + oversample = 2 dlambda = Lmeters * oversample dk = 2 * np.pi / dlambda @@ -148,31 +153,22 @@ def linear_gap_fill(F, key_lead, key_int): print("define global xlims") dist_list = np.array([np.nan, np.nan]) for k in all_beams: - # print(k) # TODO: add logger - hkey = "heights_c_weighted_mean" - x = Gd[k + "/dist"][:] - # print(x[0], x[-1]) # TODO: add logger + print(k) + x = Gd[k + "/x"][:] + print(x[0], x[-1]) dist_list = np.vstack([dist_list, [x[0], x[-1]]]) xlims = np.nanmin(dist_list[:, 0]) - dx, np.nanmin(dist_list[:, 1]) -dist_lim = 2000e3 # maximum distanc in the sea ice tha tis analysed: - - -if (xlims[1] - xlims[0]) > dist_lim: - xlims = xlims[0], xlims[0] + dist_lim - print("-reduced xlims length to ", xlims[0] + dist_lim, "m") - for k in all_beams: - dist_i = io.get_beam_var_hdf_store(Gd[k], "dist") + dist_i = io.get_beam_var_hdf_store(Gd[k], "x") x_mask = (dist_i > xlims[0]) & (dist_i < xlims[1]) - # print(k, sum(x_mask["dist"]) / (xlims[1] - xlims[0])) # TODO: add logger + print(k, sum(x_mask["x"]) / (xlims[1] - xlims[0])) print("-reduced frequency resolution") kk = kk[::2] - print("set xlims: ", xlims) print( "Loop start: ", @@ -180,37 +176,35 @@ def linear_gap_fill(F, key_lead, key_int): tracemalloc.get_traced_memory()[1] / 1e6, ) - G_gFT = dict() G_gFT_x = dict() G_rar_fft = dict() Pars_optm = dict() -k = all_beams[0] + +k = all_beams[1] +# sliderule version +hkey = "h_mean" +hkey_sigma = "h_sigma" for k in all_beams: tracemalloc.start() # ------------------------------- use gridded data - hkey = "heights_c_weighted_mean" Gi = io.get_beam_hdf_store(Gd[k]) - x_mask = (Gi["dist"] > xlims[0]) & (Gi["dist"] < xlims[1]) + x_mask = (Gi["x"] > xlims[0]) & (Gi["x"] < xlims[1]) if sum(x_mask) / (xlims[1] - xlims[0]) < 0.005: - # print("------------------- not data in beam found; skip") # TODO: add logger - continue + print("------------------- not data in beam found; skip") Gd_cut = Gi[x_mask] - x = Gd_cut["dist"] + x = Gd_cut["x"] del Gi # cut data: x_mask = (x >= xlims[0]) & (x <= xlims[1]) x = x[x_mask] - dd = np.copy(Gd_cut[hkey]) - dd_error = np.copy(Gd_cut["heights_c_std"]) + dd_error = np.copy(Gd_cut[hkey_sigma]) + dd_error[np.isnan(dd_error)] = 100 - # plt.hist(1/dd_weight, bins=40) - F = M.figure_axis_xy(6, 3) - plt.subplot(2, 1, 1) # compute slope spectra !! dd = np.gradient(dd) @@ -222,12 +216,7 @@ def linear_gap_fill(F, key_lead, key_int): x_no_nans = x[~dd_nans] dd_error_no_nans = dd_error[~dd_nans] - plt.plot( - x_no_nans, dd_no_nans, ".", color="black", markersize=1, label="slope (m/m)" - ) - plt.legend() - plt.show() - + print("gFT") with threadpool_limits(limits=N_process, user_api="blas"): pprint(threadpool_info()) @@ -330,6 +319,7 @@ def linear_gap_fill(F, key_lead, key_int): plt.ylim(ylims[0], ylims[-1]) plt.show() + # add x-mean spectal error estimate to xarray S.parceval(add_attrs=True, weight_data=False) # assign beam coordinate @@ -340,10 +330,10 @@ def linear_gap_fill(F, key_lead, key_int): GG.coords["spec_adjust"] = (("x", "beam"), np.expand_dims(GG["spec_adjust"], 1)) # add more coodindates to the Dataset - x_coord_no_gaps = linear_gap_fill(Gd_cut, "dist", "x") - y_coord_no_gaps = linear_gap_fill(Gd_cut, "dist", "y") + x_coord_no_gaps = linear_gap_fill(Gd_cut, "x", "x") + y_coord_no_gaps = linear_gap_fill(Gd_cut, "x", "y") mapped_coords = spec.sub_sample_coords( - Gd_cut["dist"], x_coord_no_gaps, y_coord_no_gaps, S.stancil_iter, map_func=None + Gd_cut["x"], x_coord_no_gaps, y_coord_no_gaps, S.stancil_iter, map_func=None ) GG.coords["x_coord"] = GG_x.coords["x_coord"] = ( @@ -362,10 +352,10 @@ def linear_gap_fill(F, key_lead, key_int): GG.coords["x_coord"][nan_mask] = np.nan GG.coords["y_coord"][nan_mask] = np.nan - lons_no_gaps = linear_gap_fill(Gd_cut, "dist", "lons") - lats_no_gaps = linear_gap_fill(Gd_cut, "dist", "lats") + lons_no_gaps = linear_gap_fill(Gd_cut, "x", "lons") + lats_no_gaps = linear_gap_fill(Gd_cut, "x", "lats") mapped_coords = spec.sub_sample_coords( - Gd_cut["dist"], lons_no_gaps, lats_no_gaps, S.stancil_iter, map_func=None + Gd_cut["x"], lons_no_gaps, lats_no_gaps, S.stancil_iter, map_func=None ) GG.coords["lon"] = GG_x.coords["lon"] = ( @@ -445,10 +435,11 @@ def get_stancil_nans(stancil): label="mean FFT", ) plt.legend() - plt.show() + except: pass time.sleep(3) + plt.close("all") del Gd_cut @@ -458,7 +449,7 @@ def get_stancil_nans(stancil): MT.save_pandas_table(Pars_optm, save_name + "_params", save_path) -# repack data +# repack data def repack_attributes(DD): attr_dim_list = list(DD.keys()) for k in attr_dim_list: diff --git a/src/icesat2_tracks/analysis_db/B04_angle.py b/src/icesat2_tracks/analysis_db/B04_angle.py new file mode 100644 index 00000000..d4067845 --- /dev/null +++ b/src/icesat2_tracks/analysis_db/B04_angle.py @@ -0,0 +1,904 @@ +import os, sys + +""" +This file open a ICEsat2 track applied filters and corections and returns smoothed photon heights on a regular grid in an .nc file. +This is python 3 +""" +from icesat2_tracks.config.IceSAT2_startup import ( + mconfig, + color_schemes, + plt, + font_for_print, + font_for_pres, +) + + +import h5py +import icesat2_tracks.ICEsat2_SI_tools.io as io +import xarray as xr +import numpy as np + + +from matplotlib.gridspec import GridSpec + +from numba import jit + +from icesat2_tracks.ICEsat2_SI_tools import angle_optimizer + +import icesat2_tracks.local_modules.m_tools_ph3 as MT +import icesat2_tracks.local_modules.m_general_ph3 as M + +import pandas as pd + + +import time + +from contextlib import contextmanager + +color_schemes.colormaps2(21) + + +col_dict = color_schemes.rels + +track_name, batch_key, test_flag = io.init_from_input( + sys.argv +) # loads standard experiment + +hemis, batch = batch_key.split("_") + +ATlevel = "ATL03" + +save_path = mconfig["paths"]["work"] + batch_key + "/B04_angle/" +save_name = "B04_" + track_name + +plot_path = ( + mconfig["paths"]["plot"] + "/" + hemis + "/" + batch_key + "/" + track_name + "/" +) +MT.mkdirs_r(plot_path) +MT.mkdirs_r(save_path) +bad_track_path = mconfig["paths"]["work"] + "bad_tracks/" + batch_key + "/" + +all_beams = mconfig["beams"]["all_beams"] +high_beams = mconfig["beams"]["high_beams"] +low_beams = mconfig["beams"]["low_beams"] +beam_groups = mconfig["beams"]["groups"] + +load_path = mconfig["paths"]["work"] + batch_key + "/B01_regrid/" +G_binned_store = h5py.File(load_path + "/" + track_name + "_B01_binned.h5", "r") +G_binned = dict() +for b in all_beams: + G_binned[b] = io.get_beam_hdf_store(G_binned_store[b]) +G_binned_store.close() + +load_path = mconfig["paths"]["work"] + batch_key + "/B02_spectra/" +Gx = xr.load_dataset(load_path + "/B02_" + track_name + "_gFT_x.nc") # +Gk = xr.load_dataset(load_path + "/B02_" + track_name + "_gFT_k.nc") # + + +# load prior information +load_path = mconfig["paths"]["work"] + batch_key + "/A02_prior/" + +try: + Prior = MT.load_pandas_table_dict("/A02_" + track_name, load_path)[ + "priors_hindcast" + ] +except: + print("Prior not found. exit") + MT.json_save( + "B04_fail", + plot_path, + { + "time": time.asctime(time.localtime(time.time())), + "reason": "Prior not found", + }, + ) + exit() + +if np.isnan(Prior["mean"]["dir"]): + print("Prior failed, entries are nan. exit.") + MT.json_save( + "B04_fail", + plot_path, + { + "time": time.asctime(time.localtime(time.time())), + "reason": "Prior not found", + }, + ) + exit() + +#### Define Prior +Pperiod = Prior.loc[["ptp0", "ptp1", "ptp2", "ptp3", "ptp4", "ptp5"]]["mean"] +Pdir = Prior.loc[["pdp0", "pdp1", "pdp2", "pdp3", "pdp4", "pdp5"]]["mean"].astype( + "float" +) +Pspread = Prior.loc[["pspr0", "pspr1", "pspr2", "pspr3", "pspr4", "pspr5"]]["mean"] + +Pperiod = Pperiod[~np.isnan(list(Pspread))] +Pdir = Pdir[~np.isnan(list(Pspread))] +Pspread = Pspread[~np.isnan(list(Pspread))] + + +# this is a hack since the current data does not have a spread +Pspread[Pspread == 0] = 20 + +# reset dirs: +Pdir[Pdir > 180] = Pdir[Pdir > 180] - 360 +Pdir[Pdir < -180] = Pdir[Pdir < -180] + 360 + +# reorder dirs +dir_best = [0] +for dir in Pdir: + ip = np.argmin( + [ + abs(dir_best[-1] - dir), + abs(dir_best[-1] - (dir - 360)), + abs(dir_best[-1] - (dir + 360)), + ] + ) + new_dir = np.array([dir, (dir - 360), (dir + 360)])[ip] + dir_best.append(new_dir) +dir_best = np.array(dir_best[1:]) + +if len(Pperiod) == 0: + print("constant peak wave number") + kk = Gk.k + Pwavenumber = kk * 0 + (2 * np.pi / (1 / Prior.loc["fp"]["mean"])) ** 2 / 9.81 + dir_best = kk * 0 + Prior.loc["dp"]["mean"] + dir_interp_smth = dir_interp = kk * 0 + Prior.loc["dp"]["mean"] + spread_smth = spread_interp = kk * 0 + Prior.loc["spr"]["mean"] + + +else: + Pwavenumber = (2 * np.pi / Pperiod) ** 2 / 9.81 + kk = Gk.k + dir_interp = np.interp( + kk, Pwavenumber[Pwavenumber.argsort()], dir_best[Pwavenumber.argsort()] + ) + dir_interp_smth = M.runningmean(dir_interp, 30, tailcopy=True) + dir_interp_smth[-1] = dir_interp_smth[-2] + + spread_interp = np.interp( + kk, + Pwavenumber[Pwavenumber.argsort()], + Pspread[Pwavenumber.argsort()].astype("float"), + ) + spread_smth = M.runningmean(spread_interp, 30, tailcopy=True) + spread_smth[-1] = spread_smth[-2] + + +font_for_pres() + +F = M.figure_axis_xy(5, 4.5, view_scale=0.5) +plt.subplot(2, 1, 1) +plt.title("Prior angle smoothed\n" + track_name, loc="left") + + +plt.plot(Pwavenumber, dir_best, ".r", markersize=8) +plt.plot(kk, dir_interp, "-", color="red", linewidth=0.8, zorder=11) +plt.plot(kk, dir_interp_smth, color=color_schemes.green1) + +plt.fill_between( + kk, + dir_interp_smth - spread_smth, + dir_interp_smth + spread_smth, + zorder=1, + color=color_schemes.green1, + alpha=0.2, +) +plt.ylabel("Angle (deg)") + +ax2 = plt.subplot(2, 1, 2) +plt.title("Prior angle adjusted ", loc="left") + +# adjust angle def: +dir_interp_smth[dir_interp_smth > 180] = dir_interp_smth[dir_interp_smth > 180] - 360 +dir_interp_smth[dir_interp_smth < -180] = dir_interp_smth[dir_interp_smth < -180] + 360 + +plt.fill_between( + kk, + dir_interp_smth - spread_smth, + dir_interp_smth + spread_smth, + zorder=1, + color=color_schemes.green1, + alpha=0.2, +) +plt.plot(kk, dir_interp_smth, ".", markersize=1, color=color_schemes.green1) + +ax2.axhline(85, color="gray", linewidth=2) +ax2.axhline(-85, color="gray", linewidth=2) + +plt.ylabel("Angle (deg)") +plt.xlabel("wavenumber ($2 \pi/\lambda$)") + +F.save_light(path=plot_path, name="B04_prior_angle") + +# save +dir_interp_smth = xr.DataArray( + data=dir_interp_smth * np.pi / 180, + dims="k", + coords={"k": kk}, + name="Prior_direction", +) +spread_smth = xr.DataArray( + data=spread_smth * np.pi / 180, dims="k", coords={"k": kk}, name="Prior_spread" +) +Prior_smth = xr.merge([dir_interp_smth, spread_smth]) + +prior_angle = Prior_smth.Prior_direction * 180 / np.pi +if (abs(prior_angle) > 80).all(): + print("Prior angle is ", prior_angle.min().data, prior_angle.max().data, ". quit.") + dd_save = { + "time": time.asctime(time.localtime(time.time())), + "angle": list( + [ + float(prior_angle.min().data), + float(prior_angle.max().data), + float(prior_angle.median()), + ] + ), + } + MT.json_save("B04_fail", plot_path, dd_save) + print("exit()") + exit() + +# define paramater range +params_dict = { + "alpha": [-0.85 * np.pi / 2, 0.85 * np.pi / 2, 5], + "phase": [0, 2 * np.pi, 10], +} + +alpha_dx = 0.02 +max_wavenumbers = 25 + +sample_flag = True +optimize_flag = False +brute_flag = False + +plot_flag = False + +Nworkers = 6 +N_sample_chain = 300 +N_sample_chain_burn = 30 + +max_x_pos = 8 +x_pos_jump = 2 + + +def make_fake_data(xi, group): + ki = Gk.k[0:2] + + bins = np.arange( + params_dict["alpha"][0], params_dict["alpha"][1] + alpha_dx, alpha_dx + ) + bins_pos = bins[0:-1] + np.diff(bins) / 2 + marginal_stack = xr.DataArray( + np.nan * np.vstack([bins_pos, bins_pos]).T, + dims=("angle", "k"), + coords={"angle": bins_pos, "k": ki.data}, + ) + + group_name = str("group" + group[0].split("gt")[1].split("l")[0]) + marginal_stack.coords["beam_group"] = group_name + marginal_stack.coords["x"] = xi + marginal_stack.name = "marginals" + marginal_stack.expand_dims(dim="x", axis=2).expand_dims(dim="beam_group", axis=3) + return marginal_stack + + +def define_wavenumber_weights_tot_var( + dd, m=3, variance_frac=0.33, k_upper_lim=None, verbose=False +): + """ + return peaks of a power spectrum dd that in the format such that they can be used as weights for the frequencies based fitting + + inputs: + dd xarray with PSD as data amd coordindate wavenumber k + m running mean half-width in gridpoints + variance_frac (0 to 1) How much variance should be explained by the returned peaks + verbose if true it plots some stuff + + + return: + mask size of dd. where True the data is identified as having significant amplitude + k wanumbers where mask is true + dd_rm smoothed version of dd + positions postions where of significant data in array + """ + + if len(dd.shape) == 2: + dd_use = dd.mean("beam") + + if m is None: + dd_rm = dd_use.data + else: + dd_rm = M.runningmean(dd_use, m, tailcopy=True) + + k = dd_use.k[~np.isnan(dd_rm)].data + dd_rm = dd_rm[~np.isnan(dd_rm)] + + orders = dd_rm.argsort()[::-1] + var_mask = dd_rm[orders].cumsum() / dd_rm.sum() < variance_frac + pos_cumsum = orders[var_mask] + mask = var_mask[orders.argsort()] + if k_upper_lim is not None: + mask = (k < k_upper_lim) & mask + + if verbose: + plt.plot( + dd.k, + dd, + "-", + color=col_dict[str(amp_data.beam[0].data)], + markersize=20, + alpha=0.6, + ) + plt.plot(k, dd_rm, "-k", markersize=20) + + plt.plot(k[mask], dd_rm[mask], ".r", markersize=10, zorder=12) + if k_upper_lim is not None: + plt.gca().axvline(k_upper_lim, color="black") + + return mask, k, dd_rm, pos_cumsum + + +def define_wavenumber_weights_threshold(dd, m=3, Nstd=2, verbose=False): + if m is None: + dd_rm = dd + else: + dd_rm = M.runningmean(dd, m, tailcopy=True) + + k = dd.k[~np.isnan(dd_rm)] + dd_rm = dd_rm[~np.isnan(dd_rm)] + + treshold = np.nanmean(dd_rm) + np.nanstd(dd_rm) * Nstd + mask = dd_rm > treshold + + if verbose: + plt.plot(dd.k, dd, "-k", markersize=20) + plt.plot(k, dd_rm, "-b", markersize=20) + + k_list = k[mask] + dd_list = dd_rm[mask] + + plt.plot(k_list, dd_list, ".r", markersize=10, zorder=12) + + return mask, k, dd_rm, np.arange(0, mask.size)[mask] + + +def plot_instance( + z_model, + fargs, + key, + SM, + non_dim=False, + title_str=None, + brute=False, + optimze=False, + sample=False, + view_scale=0.3, +): + x_concat, y_concat, z_concat = fargs + + F = M.figure_axis_xy(5, 6, view_scale=view_scale, container=True) + plt.suptitle(title_str) + gs = GridSpec(4, 3, wspace=0.4, hspace=1.2) + F.gs = gs + + import itertools + + col_list = itertools.cycle( + [ + color_schemes.cascade2, + color_schemes.rascade2, + color_schemes.cascade1, + color_schemes.rascade1, + color_schemes.cascade3, + color_schemes.rascade3, + ] + ) + + beam_list = list(set(y_concat)) + for y_pos, pos in zip(beam_list, [gs[0, :], gs[1, :]]): + F.ax2 = F.fig.add_subplot(pos) + + plt.title(str(y_pos)) + + plt.plot( + x_concat[y_concat == y_pos], + z_concat[y_concat == y_pos], + c=color_schemes.gray, + linewidth=1, + ) + plt.plot( + x_concat[y_concat == y_pos], + z_model[y_concat == y_pos], + "-", + c=next(col_list), + ) + plt.xlim(x_concat[y_concat == y_pos][0], x_concat[y_concat == y_pos][-1]) + + plt.xlabel("meter") + F.ax3 = F.fig.add_subplot(gs[2:, 0:-1]) + if brute is True: + plt.title("Brute-force costs", loc="left") + SM.plot_brute(marker=".", color="blue", markersize=15, label="Brute", zorder=10) + + if optimze is True: + SM.plot_optimze(color="r", markersize=10, zorder=12, label="Dual Annealing") + + if sample is True: + SM.plot_sample(markersize=2, linewidth=0.8, alpha=0.2, color="black", zorder=8) + + F.ax4 = F.fig.add_subplot(gs[2:, -1]) + return F + + +# isolate x positions with data +data_mask = Gk.gFT_PSD_data.mean("k") +data_mask.coords["beam_group"] = ( + "beam", + ["beam_group" + g[2] for g in data_mask.beam.data], +) +data_mask_group = data_mask.groupby("beam_group").mean(skipna=False) +# these stancils are actually used +data_sel_mask = data_mask_group.sum("beam_group") != 0 + +x_list = data_sel_mask.x[data_sel_mask] # iterate over these x posistions +x_list_flag = ~np.isnan( + data_mask_group.sel(x=x_list) +) # flag that is False if there is no data + +#### limit number of x coordinates + +x_list = x_list[::x_pos_jump] +if len(x_list) > max_x_pos: + x_list = x_list[0:max_x_pos] +x_list_flag = x_list_flag.sel(x=x_list) + +# plot +font_for_print() +F = M.figure_axis_xy(5.5, 3, view_scale=0.8) +plt.suptitle(track_name) +ax1 = plt.subplot(2, 1, 1) +plt.title("Data in Beam", loc="left") +plt.pcolormesh(data_mask.x / 1e3, data_mask.beam, data_mask, cmap=plt.cm.OrRd) +for i in np.arange(1.5, 6, 2): + ax1.axhline(i, color="black", linewidth=0.5) +plt.xlabel("Distance from Ice Edge") + +ax2 = plt.subplot(2, 1, 2) +plt.title("Data in Group", loc="left") +plt.pcolormesh( + data_mask.x / 1e3, data_mask_group.beam_group, data_mask_group, cmap=plt.cm.OrRd +) + +for i in np.arange(0.5, 3, 1): + ax2.axhline(i, color="black", linewidth=0.5) + +plt.plot(x_list / 1e3, x_list * 0 + 0, ".", markersize=2, color=color_schemes.cascade1) +plt.plot(x_list / 1e3, x_list * 0 + 1, ".", markersize=2, color=color_schemes.cascade1) +plt.plot(x_list / 1e3, x_list * 0 + 2, ".", markersize=2, color=color_schemes.cascade1) + +plt.xlabel("Distance from Ice Edge") + +F.save_pup(path=plot_path, name="B04_data_avail") + +Marginals = dict() +L_collect = dict() + +group_number = np.arange(len(beam_groups)) +ggg, xxx = np.meshgrid(group_number, x_list.data) + +for gi in zip(ggg.flatten(), xxx.flatten()): + print(gi) + + group, xi = beam_groups[gi[0]], gi[1] + + if bool(x_list_flag.sel(x=xi).isel(beam_group=gi[0]).data) is False: + print("no data, fill with dummy") + ikey = str(xi) + "_" + "_".join(group) + Marginals[ikey] = make_fake_data(xi, group) + continue + + GGx = Gx.sel(beam=group).sel(x=xi) + GGk = Gk.sel(beam=group).sel(x=xi) + + ### define data + # normalize data + key = "y_data" + amp_Z = (GGx[key] - GGx[key].mean(["eta"])) / GGx[key].std(["eta"]) + N = amp_Z.shape[0] + + # define x,y positions + eta_2d = GGx.eta + GGx.x_coord - GGx.x_coord.mean() + nu_2d = GGx.eta * 0 + GGx.y_coord - GGx.y_coord.mean() + + # repack as np arrays + x_concat = eta_2d.data.T.flatten() + y_concat = nu_2d.data.T.flatten() + z_concat = amp_Z.data.flatten() + + x_concat = x_concat[~np.isnan(z_concat)] + y_concat = y_concat[~np.isnan(z_concat)] + z_concat = z_concat[~np.isnan(z_concat)] + N_data = x_concat.size + + if np.isnan(z_concat).sum() != 0: + raise ValueError("There are still nans") + + mean_dist = (nu_2d.isel(beam=0) - nu_2d.isel(beam=1)).mean().data + k_upper_lim = 2 * np.pi / (mean_dist * 1) + + print("k_upper_lim ", k_upper_lim) + + # variance method + amp_data = np.sqrt(GGk.gFT_cos_coeff**2 + GGk.gFT_sin_coeff**2) + mask, k, weights, positions = define_wavenumber_weights_tot_var( + amp_data, m=1, k_upper_lim=k_upper_lim, variance_frac=0.20, verbose=False + ) + + if len(k[mask]) == 0: + print("no good k found, fill with dummy") + ikey = str(xi) + "_" + "_".join(group) + Marginals[ikey] = make_fake_data(xi, group) + continue + + SM = angle_optimizer.sample_with_mcmc(params_dict) + SM.set_objective_func(angle_optimizer.objective_func) + nan_list = np.isnan(x_concat) | np.isnan(y_concat) | np.isnan(y_concat) + x_concat[nan_list] = [] + y_concat[nan_list] = [] + z_concat[nan_list] = [] + SM.fitting_args = fitting_args = (x_concat, y_concat, z_concat) + + # test: + k_prime_max = 0.02 + amp_Z = 1 + prior_sel = { + "alpha": ( + Prior_smth.sel(k=k_prime_max, method="nearest").Prior_direction.data, + Prior_smth.sel(k=k_prime_max, method="nearest").Prior_spread.data, + ) + } + SM.fitting_kargs = fitting_kargs = {"prior": prior_sel, "prior_weight": 3} + # test if it works + SM.params.add( + "K_prime", k_prime_max, vary=False, min=k_prime_max * 0.5, max=k_prime_max * 1.5 + ) + SM.params.add("K_amp", amp_Z, vary=False, min=amp_Z * 0.0, max=amp_Z * 5) + try: + SM.test_objective_func() + except: + raise ValueError("Objective function test fails") + + def get_instance(k_pair): + k_prime_max, Z_max = k_pair + + prior_sel = { + "alpha": ( + Prior_smth.sel(k=k_prime_max, method="nearest").Prior_direction.data, + Prior_smth.sel(k=k_prime_max, method="nearest").Prior_spread.data, + ) + } + + SM.fitting_kargs = fitting_kargs = {"prior": prior_sel, "prior_weight": 2} + + amp_Z = 1 + SM.params.add( + "K_prime", + k_prime_max, + vary=False, + min=k_prime_max * 0.5, + max=k_prime_max * 1.5, + ) + SM.params.add("K_amp", amp_Z, vary=False, min=amp_Z * 0.0, max=amp_Z * 5) + + L_sample_i = None + L_optimize_i = None + L_brute_i = None + if sample_flag: + SM.sample(verbose=False, steps=N_sample_chain, progress=False, workers=None) + L_sample_i = list(SM.fitter.params.valuesdict().values()) # mcmc results + + elif optimize_flag: + SM.optimize(verbose=False) + L_optimize_i = list( + SM.fitter_optimize.params.valuesdict().values() + ) # mcmc results + + elif brute_flag: + SM.brute(verbose=False) + L_brute_i = list( + SM.fitter_brute.params.valuesdict().values() + ) # mcmc results + else: + raise ValueError( + "non of sample_flag,optimize_flag, or brute_flag are True" + ) + + y_hist, bins, bins_pos = SM.get_marginal_dist( + "alpha", alpha_dx, burn=N_sample_chain_burn, plot_flag=False + ) + fitter = SM.fitter # MCMC results + z_model = SM.objective_func(fitter.params, *fitting_args, test_flag=True) + cost = (fitter.residual**2).sum() / (z_concat**2).sum() + + if plot_flag: + F = plot_instance( + z_model, + fitting_args, + "y_data_normed", + SM, + brute=brute_flag, + optimze=optimize_flag, + sample=sample_flag, + title_str="k=" + str(np.round(k_prime_max, 4)), + view_scale=0.6, + ) + + if fitting_kargs["prior"] is not None: + F.ax3.axhline( + prior_sel["alpha"][0], color="green", linewidth=2, label="Prior" + ) + F.ax3.axhline( + prior_sel["alpha"][0] - prior_sel["alpha"][1], + color="green", + linewidth=0.7, + ) + F.ax3.axhline( + prior_sel["alpha"][0] + prior_sel["alpha"][1], + color="green", + linewidth=0.7, + ) + + F.ax3.axhline( + fitter.params["alpha"].min, color="gray", linewidth=2, alpha=0.6 + ) + F.ax3.axhline( + fitter.params["alpha"].max, color="gray", linewidth=2, alpha=0.6 + ) + + plt.sca(F.ax3) + plt.legend() + plt.xlabel("Phase") + plt.ylabel("Angle") + plt.xlim(0, np.pi * 2) + + plt.sca(F.ax4) + plt.xlabel("Density") + plt.stairs(y_hist, bins, orientation="horizontal", color="k") + + F.ax4.axhline( + fitter.params["alpha"].min, color="gray", linewidth=2, alpha=0.6 + ) + F.ax4.axhline( + fitter.params["alpha"].max, color="gray", linewidth=2, alpha=0.6 + ) + + F.ax3.set_ylim( + min(-np.pi / 2, prior_sel["alpha"][0] - 0.2), + max(np.pi / 2, prior_sel["alpha"][0] + 0.2), + ) + F.ax4.set_ylim( + min(-np.pi / 2, prior_sel["alpha"][0] - 0.2), + max(np.pi / 2, prior_sel["alpha"][0] + 0.2), + ) + + plt.show() + F.save_light(path=plot_path, name=track_name + "_fit_k" + str(k_prime_max)) + + marginal_stack_i = xr.DataArray( + y_hist, dims=("angle"), coords={"angle": bins_pos} + ) + marginal_stack_i.coords["k"] = np.array(k_prime_max) + + rdict = { + "marginal_stack_i": marginal_stack_i, + "L_sample_i": L_sample_i, + "L_optimize_i": L_optimize_i, + "L_brute_i": L_brute_i, + "cost": cost, + } + return k_prime_max, rdict + + k_list, weight_list = k[mask], weights[mask] + print("# of wavenumber: ", len(k_list)) + if len(k_list) > max_wavenumbers: + print("cut wavenumber list to 20") + k_list = k_list[0:max_wavenumbers] + weight_list = weight_list[0:max_wavenumbers] + + # # parallel version tends to fail... + # Let's keep this in case we decided to work on parallellize it + # with futures.ProcessPoolExecutor(max_workers=Nworkers) as executor: + # A = dict( executor.map(get_instance, zip(k_list, weight_list) )) + + A = dict() + for k_pair in zip(k_list, weight_list): + kk, I = get_instance(k_pair) + A[kk] = I + + cost_stack = dict() + marginal_stack = dict() + L_sample = pd.DataFrame(index=["alpha", "group_phase", "K_prime", "K_amp"]) + L_optimize = pd.DataFrame(index=["alpha", "group_phase", "K_prime", "K_amp"]) + L_brute = pd.DataFrame(index=["alpha", "group_phase", "K_prime", "K_amp"]) + + for kk, I in A.items(): + L_sample[kk] = I["L_sample_i"] + L_optimize[kk] = I["L_optimize_i"] + L_brute[kk] = I["L_brute_i"] + + marginal_stack[kk] = I["marginal_stack_i"] + cost_stack[kk] = I["cost"] + + # ## add beam_group dimension + marginal_stack = xr.concat(marginal_stack.values(), dim="k").sortby("k") + L_sample = L_sample.T.sort_values("K_prime") + L_optimize = L_optimize.T.sort_values("K_prime") + L_brute = L_brute.T.sort_values("K_prime") + + print("done with ", group, xi / 1e3) + + # collect + ikey = str(xi) + "_" + "_".join(group) + + marginal_stack.name = "marginals" + marginal_stack = marginal_stack.to_dataset() + marginal_stack["cost"] = (("k"), list(cost_stack.values())) + marginal_stack["weight"] = (("k"), weight_list) + + group_name = str("group" + group[0].split("gt")[1].split("l")[0]) + marginal_stack.coords["beam_group"] = group_name + marginal_stack.coords["x"] = xi + + Marginals[ikey] = marginal_stack.expand_dims(dim="x", axis=0).expand_dims( + dim="beam_group", axis=1 + ) + Marginals[ikey].coords["N_data"] = ( + ("x", "beam_group"), + np.expand_dims(np.expand_dims(N_data, 0), 1), + ) + + L_sample["cost"] = cost_stack + L_sample["weight"] = weight_list + L_collect[group_name, str(int(xi))] = L_sample + + +MM = xr.merge(Marginals.values()) +MM = xr.merge([MM, Prior_smth]) +MM.to_netcdf(save_path + save_name + "_marginals.nc") + +try: + LL = pd.concat(L_collect) + MT.save_pandas_table({"L_sample": LL}, save_name + "_res_table", save_path) +except Exception as e: + print(f"This is a warning: {e}") + +# plot +font_for_print() +F = M.figure_axis_xy(6, 5.5, view_scale=0.7, container=True) + +gs = GridSpec(4, 6, wspace=0.2, hspace=0.8) + +ax0 = F.fig.add_subplot(gs[0:2, -1]) +ax0.tick_params(labelleft=False) + +klims = 0, LL["K_prime"].max() * 1.2 + + +for g in MM.beam_group: + MMi = MM.sel(beam_group=g) + plt.plot( + MMi.weight.T, + MMi.k, + ".", + color=col_dict[str(g.data)], + markersize=3, + linewidth=0.8, + ) + +plt.xlabel("Power") +plt.ylim(klims) + +ax1 = F.fig.add_subplot(gs[0:2, 0:-1]) + +for g in MM.beam_group: + Li = LL.loc[str(g.data)] + + angle_list = np.array(Li["alpha"]) * 180 / np.pi + kk_list = np.array(Li["K_prime"]) + weight_list_i = np.array(Li["weight"]) + + plt.scatter( + angle_list, + kk_list, + s=(weight_list_i * 8e1) ** 2, + c=col_dict[str(g.data)], + label="mode " + str(g.data), + ) + + +dir_best[dir_best > 180] = dir_best[dir_best > 180] - 360 +plt.plot(dir_best, Pwavenumber, ".r", markersize=6) + +dir_interp[dir_interp > 180] = dir_interp[dir_interp > 180] - 360 +plt.plot(dir_interp, Gk.k, "-", color="red", linewidth=0.3, zorder=11) + + +plt.fill_betweenx( + Gk.k, + (dir_interp_smth - spread_smth) * 180 / np.pi, + (dir_interp_smth + spread_smth) * 180 / np.pi, + zorder=1, + color=color_schemes.green1, + alpha=0.2, +) +plt.plot( + dir_interp_smth * 180 / np.pi, Gk.k, ".", markersize=1, color=color_schemes.green1 +) + +ax1.axvline(85, color="gray", linewidth=2) +ax1.axvline(-85, color="gray", linewidth=2) + + +plt.legend() +plt.ylabel("wavenumber (deg)") +plt.xlabel("Angle (deg)") + + +plt.ylim(klims) + +prior_angle_str = str(np.round((prior_sel["alpha"][0]) * 180 / np.pi)) +plt.title(track_name + "\nprior=" + prior_angle_str + "deg", loc="left") + +plt.xlim(min([-90, np.nanmin(dir_best)]), max([np.nanmax(dir_best), 90])) + + +ax3 = F.fig.add_subplot(gs[2, 0:-1]) + +for g in MM.beam_group: + MMi = MM.sel(beam_group=g) + wegihted_margins = (MMi.marginals * MMi.weight).sum(["x", "k"]) / MMi.weight.sum( + ["x", "k"] + ) + plt.plot( + MMi.angle * 180 / np.pi, + wegihted_margins, + ".", + color=col_dict[str(g.data)], + markersize=2, + linewidth=0.8, + ) + +plt.ylabel("Density") +plt.title("weight margins", loc="left") + +plt.xlim(-90, 90) + + +ax3 = F.fig.add_subplot(gs[-1, 0:-1]) + +for g in MM.beam_group: + MMi = MM.sel(beam_group=g) + wegihted_margins = MMi.marginals.mean(["x", "k"]) + plt.plot( + MMi.angle * 180 / np.pi, + wegihted_margins, + ".", + color=col_dict[str(g.data)], + markersize=2, + linewidth=0.8, + ) + +plt.ylabel("Density") +plt.xlabel("Angle (deg)") +plt.title("unweighted margins", loc="left") + +plt.xlim(-90, 90) + +F.save_pup(path=plot_path, name="B04_marginal_distributions") + +MT.json_save( + "B04_success", plot_path, {"time": "time.asctime( time.localtime(time.time()) )"} +)