From 9f32f28decf3e682e1b0d9eedc7c6bc80fc83c88 Mon Sep 17 00:00:00 2001 From: Carlos Paniagua Date: Wed, 31 Jan 2024 14:09:13 -0500 Subject: [PATCH] feat: cli step 2 -- b03_plot_spectra_ovpy --- .../test-B01_SL_load_single_file.yml | 4 +- .../analysis_db/B03_plot_spectra_ov.py | 903 +++++++++--------- 2 files changed, 455 insertions(+), 452 deletions(-) diff --git a/.github/workflows/test-B01_SL_load_single_file.yml b/.github/workflows/test-B01_SL_load_single_file.yml index 50e88bf9..6c0b9748 100644 --- a/.github/workflows/test-B01_SL_load_single_file.yml +++ b/.github/workflows/test-B01_SL_load_single_file.yml @@ -27,4 +27,6 @@ jobs: run: python src/icesat2_tracks/analysis_db/B01_SL_load_single_file.py --track-name 20190502052058_05180312_005_01 --batch-key SH_testSLsinglefile2 --output-dir ./work - name: second step make_spectra run: python src/icesat2_tracks/analysis_db/B02_make_spectra_gFT.py --track-name SH_20190502_05180312 --batch-key SH_testSLsinglefile2 --output-dir ./work - + - name: third step plot_spectra + run: python src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py --track-name SH_20190502_05180312 --batch-key SH_testSLsinglefile2 --output-dir ./work + \ No newline at end of file diff --git a/src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py b/src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py index ca7984d2..2e4c1721 100644 --- a/src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py +++ b/src/icesat2_tracks/analysis_db/B03_plot_spectra_ov.py @@ -1,134 +1,36 @@ +#!/usr/bin/env python3 """ -This file open a ICEsat2 track applied filters and corections and returns smoothed photon heights on a regular grid in an .nc file. +This file open a ICEsat2 track applied filters and corrections and returns smoothed photon heights on a regular grid in an .nc file. This is python 3 """ -import sys +from pathlib import Path +import matplotlib import numpy as np import xarray as xr from matplotlib.gridspec import GridSpec +import typer + import icesat2_tracks.ICEsat2_SI_tools.io as io import icesat2_tracks.ICEsat2_SI_tools.generalized_FT as gFT import icesat2_tracks.local_modules.m_tools_ph3 as MT from icesat2_tracks.local_modules import m_general_ph3 as M -from icesat2_tracks.config.IceSAT2_startup import mconfig, color_schemes, plt, font_for_print - -track_name, batch_key, test_flag = io.init_from_input( - sys.argv # TODO: Handle via CLI -) # loads standard experiment -hemis, batch = batch_key.split("_") - -load_path = mconfig["paths"]["work"] + batch_key + "/B02_spectra/" -load_file = load_path + "B02_" + track_name -plot_path = ( - mconfig["paths"]["plot"] + "/" + hemis + "/" + batch_key + "/" + track_name + "/" -) # TODO: Update with pathlib -MT.mkdirs_r(plot_path) - -Gk = xr.open_dataset(load_file + "_gFT_k.nc") -Gx = xr.open_dataset(load_file + "_gFT_x.nc") - -Gfft = xr.open_dataset(load_file + "_FFT.nc") - -all_beams = mconfig["beams"]["all_beams"] -high_beams = mconfig["beams"]["high_beams"] -low_beams = mconfig["beams"]["low_beams"] -color_schemes.colormaps2(21) - -col_dict = color_schemes.rels -F = M.figure_axis_xy(9, 3, view_scale=0.5) - -plt.subplot(1, 3, 1) -plt.title(track_name, loc="left") -for k in all_beams: - I = Gk.sel(beam=k) - I2 = Gx.sel(beam=k) - plt.plot(I["lon"], I["lat"], ".", c=col_dict[k], markersize=0.7, linewidth=0.3) - plt.plot(I2["lon"], I2["lat"], "|", c=col_dict[k], markersize=0.7) - - -plt.xlabel("lon") -plt.ylabel("lat") - -plt.subplot(1, 3, 2) - -xscale = 1e3 -for k in all_beams: - I = Gk.sel(beam=k) - plt.plot( - I["x_coord"] / xscale, - I["y_coord"] / xscale, - ".", - c=col_dict[k], - linewidth=0.8, - markersize=0.8, - ) - -plt.xlabel("x_coord (km)") -plt.ylabel("y_coord (km)") - -plt.subplot(1, 3, 3) - -xscale = 1e3 -for k in all_beams: - I = Gk.sel(beam=k) - plt.plot( - I["x_coord"] / xscale, - (I["y_coord"] - I["y_coord"][0]), - ".", - c=col_dict[k], - linewidth=0.8, - markersize=0.8, - ) - -plt.xlabel("x_coord (km)") -plt.ylabel("y_coord deviation (m)") - - -F.save_light(path=plot_path, name="B03_specs_coord_check") - - -def dict_weighted_mean(Gdict, weight_key): - """ - returns the weighted meean of a dict of xarray, data_arrays - weight_key must be in the xr.DataArrays - """ - - akey = list(Gdict.keys())[0] - GSUM = Gdict[akey].copy() - GSUM.data = np.zeros(GSUM.shape) - N_per_stancil = GSUM.N_per_stancil * 0 - N_photons = np.zeros(GSUM.N_per_stancil.size) - - counter = 0 - for I in Gdict.items(): - I = I.squeeze() - if len(I.x) != 0: - GSUM += I.where(~np.isnan(I), 0) * I[weight_key] - N_per_stancil += I[weight_key] - if "N_photons" in GSUM.coords: - N_photons += I["N_photons"] - counter += 1 - - GSUM = GSUM / N_per_stancil - - if "N_photons" in GSUM.coords: - GSUM.coords["N_photons"] = (("x", "beam"), np.expand_dims(N_photons, 1)) - - GSUM["beam"] = ["weighted_mean"] - GSUM.name = "power_spec" - - return GSUM - - -G_gFT_wmean = ( - Gk["gFT_PSD_data"].where(~np.isnan(Gk["gFT_PSD_data"]), 0) * Gk["N_per_stancil"] -).sum("beam") / Gk["N_per_stancil"].sum("beam") -G_gFT_wmean["N_per_stancil"] = Gk["N_per_stancil"].sum("beam") +from icesat2_tracks.config.IceSAT2_startup import ( + mconfig, + color_schemes, + plt, + font_for_print, +) -G_fft_wmean = (Gfft.where(~np.isnan(Gfft), 0) * Gfft["N_per_stancil"]).sum( - "beam" -) / Gfft["N_per_stancil"].sum("beam") -G_fft_wmean["N_per_stancil"] = Gfft["N_per_stancil"].sum("beam") +from icesat2_tracks.clitools import ( + echo, + validate_batch_key, + validate_output_dir, + suppress_stdout, + update_paths_mconfig, + report_input_parameters, + validate_track_name_steps_gt_1, + makeapp, +) def plot_wavenumber_spectrogram(ax, Gi, clev, title=None, plot_photon_density=True): @@ -162,402 +64,501 @@ def plot_wavenumber_spectrogram(ax, Gi, clev, title=None, plot_photon_density=Tr plt.title(title, loc="left") -Gmean = G_gFT_wmean.rolling(k=5, center=True).mean() +def plot_data_eta(D, offset=0, **kargs): + eta_1 = D.eta # + D.x + y_data = D.y_model + offset + plt.plot(eta_1, y_data, **kargs) + return eta_1 -try: - k_max = Gmean.k[Gmean.isel(x=slice(0, 5)).mean("x").argmax().data].data -except: - k_max = Gmean.k[Gmean.isel(x=slice(0, 20)).mean("x").argmax().data].data -k_max_range = (k_max * 0.75, k_max * 1, k_max * 1.25) -font_for_print() -F = M.figure_axis_xy(6.5, 5.6, container=True, view_scale=1) -Lmeters = Gk.L.data[0] +def plot_model_eta(D, ax, offset=0, **kargs): + eta = D.eta # + D.x + y_data = D.y_model + offset + plt.plot(eta, y_data, **kargs) -plt.suptitle("gFT Slope Spectrograms\n" + track_name, y=0.98) -gs = GridSpec(3, 3, wspace=0.2, hspace=0.5) + ax.axvline(eta[0].data, linewidth=0.1, color=kargs["color"], alpha=0.5) + ax.axvline(eta[-1].data, linewidth=0.1, color=kargs["color"], alpha=0.5) -Gplot = ( - G_gFT_wmean.squeeze() - .rolling(k=10, min_periods=1, center=True) - .median() - .rolling(x=3, min_periods=1, center=True) - .median() -) -dd = 10 * np.log10(Gplot) -dd = dd.where(~np.isinf(dd), np.nan) -clev_log = M.clevels([dd.quantile(0.01).data, dd.quantile(0.98).data * 1.2], 31) * 1 -xlims = Gmean.x[0] / 1e3, Gmean.x[-1] / 1e3 +matplotlib.use("Agg") # prevent plot windows from opening -k = high_beams[0] -for pos, k, pflag in zip( - [gs[0, 0], gs[0, 1], gs[0, 2]], high_beams, [True, False, False] -): - ax0 = F.fig.add_subplot(pos) - Gplot = Gk.sel(beam=k).gFT_PSD_data.squeeze() - dd2 = 10 * np.log10(Gplot) - dd2 = dd2.where(~np.isinf(dd2), np.nan) - plot_wavenumber_spectrogram( - ax0, dd2, clev_log, title=k + " unsmoothed", plot_photon_density=True - ) - plt.xlim(xlims) - if pflag: - plt.ylabel("Wave length\n(meters)") - plt.legend() -for pos, k, pflag in zip( - [gs[1, 0], gs[1, 1], gs[1, 2]], low_beams, [True, False, False] +def run_B03_plot_spectra_ov( + track_name: str = typer.Option(..., callback=validate_track_name_steps_gt_1), + batch_key: str = typer.Option(..., callback=validate_batch_key), + ID_flag: bool = True, + output_dir: str = typer.Option(None, callback=validate_output_dir), ): - ax0 = F.fig.add_subplot(pos) - Gplot = Gk.sel(beam=k).gFT_PSD_data.squeeze() - dd2 = 10 * np.log10(Gplot) - dd2 = dd2.where(~np.isinf(dd2), np.nan) - plot_wavenumber_spectrogram( - ax0, dd2, clev_log, title=k + " unsmoothed", plot_photon_density=True - ) - plt.xlim(xlims) - if pflag: - plt.ylabel("Wave length\n(meters)") - plt.legend() - -ax0 = F.fig.add_subplot(gs[2, 0]) + """ + TODO: add docstring + """ + with suppress_stdout(): + track_name, batch_key, _ = io.init_from_input( + [ + None, + track_name, + batch_key, + ID_flag, + ] # init_from_input expects sys.argv with 4 elements + ) -plot_wavenumber_spectrogram( - ax0, - dd, - clev_log, - title="smoothed weighted mean \n10 $\log_{10}( (m/m)^2 m )$", - plot_photon_density=True, -) -plt.xlim(xlims) + kargs = { + "track_name": track_name, + "batch_key": batch_key, + "ID_flag": ID_flag, + "output_dir": output_dir, + } + report_input_parameters(**kargs) -line_styles = ["--", "-", "--"] -for k_max, style in zip(k_max_range, line_styles): - ax0.axhline(2 * np.pi / k_max, color="red", linestyle=style, linewidth=0.5) + hemis, batch = batch_key.split("_") -if pflag: - plt.ylabel("Wave length\n(meters)") - plt.legend() + workdir, plotsdir = update_paths_mconfig(output_dir, mconfig) -pos = gs[2, 1] -ax0 = F.fig.add_subplot(pos) -plt.title("Photons density ($m^{-1}$)", loc="left") - -for k in all_beams: - I = Gk.sel(beam=k)["gFT_PSD_data"] - plt.plot(Gplot.x / 1e3, I.N_photons / I.L.data, label=k, linewidth=0.8) -plt.plot( - Gplot.x / 1e3, - G_gFT_wmean.N_per_stancil / 3 / I.L.data, - c="black", - label="ave Photons", - linewidth=0.8, -) -plt.xlim(xlims) -plt.xlabel("Distance from the Ice Edge (km)") - -pos = gs[2, 2] - -ax0 = F.fig.add_subplot(pos) -ax0.set_yscale("log") - -plt.title("Peak Spectal Power", loc="left") - -x0 = Gk.x[0].data -for k in all_beams: - I = Gk.sel(beam=k)["gFT_PSD_data"] - plt.scatter( - I.x.data / 1e3, - I.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k").data, - s=0.5, - marker=".", - color="red", - alpha=0.3, - ) - I = Gfft.sel(beam=k) - plt.scatter( - (x0 + I.x.data) / 1e3, - I.power_spec.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), - s=0.5, - marker=".", - c="blue", - alpha=0.3, - ) + load_path = Path(workdir, batch_key, "B02_spectra") + load_file = str(load_path / ("B02_" + track_name)) + plot_path = Path(plotsdir, hemis, batch_key, track_name) + MT.mkdirs_r(plot_path) + Gk = xr.open_dataset(load_file + "_gFT_k.nc") + Gx = xr.open_dataset(load_file + "_gFT_x.nc") -Gplot = G_fft_wmean.squeeze() -Gplot = Gplot.power_spec[:, Gplot.N_per_stancil >= Gplot.N_per_stancil.max().data * 0.9] -plt.plot( - (x0 + Gplot.x) / 1e3, - Gplot.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), - ".", - markersize=1.5, - c="blue", - label="FFT", -) + Gfft = xr.open_dataset(load_file + "_FFT.nc") -Gplot = G_gFT_wmean.squeeze() -plt.plot( - Gplot.x / 1e3, - Gplot.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), - ".", - markersize=1.5, - c="red", - label="gFT", -) + all_beams = mconfig["beams"]["all_beams"] + high_beams = mconfig["beams"]["high_beams"] + low_beams = mconfig["beams"]["low_beams"] + color_schemes.colormaps2(21) -plt.ylabel("1e-3 $(m)^2~m$") -plt.legend() + col_dict = color_schemes.rels + F = M.figure_axis_xy(9, 3, view_scale=0.5) -F.save_light(path=plot_path, name="B03_specs_L" + str(Lmeters)) + plt.subplot(1, 3, 1) + plt.title(track_name, loc="left") + for k in all_beams: + I = Gk.sel(beam=k) + I2 = Gx.sel(beam=k) + plt.plot(I["lon"], I["lat"], ".", c=col_dict[k], markersize=0.7, linewidth=0.3) + plt.plot(I2["lon"], I2["lat"], "|", c=col_dict[k], markersize=0.7) -Gk.sel(beam=k).gFT_PSD_data.plot() + plt.xlabel("lon") + plt.ylabel("lat") + plt.subplot(1, 3, 2) -def plot_model_eta(D, ax, offset=0, xscale=1e3, **kargs): - eta = D.eta + D.x - y_data = D.y_model + offset - plt.plot(eta / xscale, y_data, **kargs) - - ax.axvline(eta[0].data / xscale, linewidth=2, color=kargs["color"], alpha=0.5) - ax.axvline(eta[-1].data / xscale, linewidth=2, color=kargs["color"], alpha=0.5) - - -def add_info(D, Dk, ylims): - eta = D.eta + D.x - N_per_stancil, ksize = Dk.N_per_stancil.data, Dk.k.size - plt.text( - eta[0].data, - ylims[-1], - " N=" - + numtostr(N_per_stancil) - + " N/2M= " - + fltostr(N_per_stancil / 2 / ksize, 1), - ) + xscale = 1e3 + for k in all_beams: + I = Gk.sel(beam=k) + plt.plot( + I["x_coord"] / xscale, + I["y_coord"] / xscale, + ".", + c=col_dict[k], + linewidth=0.8, + markersize=0.8, + ) + plt.xlabel("x_coord (km)") + plt.ylabel("y_coord (km)") -def plot_data_eta(D, offset=0, xscale=1e3, **kargs): - eta_1 = D.eta + D.x - y_data = D.y_model + offset - plt.plot(eta_1 / xscale, y_data, **kargs) - return eta_1 + plt.subplot(1, 3, 3) + xscale = 1e3 + for k in all_beams: + I = Gk.sel(beam=k) + plt.plot( + I["x_coord"] / xscale, + (I["y_coord"] - I["y_coord"][0]), + ".", + c=col_dict[k], + linewidth=0.8, + markersize=0.8, + ) -def plot_data_eta(D, offset=0, **kargs): - eta_1 = D.eta # + D.x - y_data = D.y_model + offset - plt.plot(eta_1, y_data, **kargs) - return eta_1 + plt.xlabel("x_coord (km)") + plt.ylabel("y_coord deviation (m)") + + F.save_light(path=plot_path, name="B03_specs_coord_check") + + G_gFT_wmean = ( + Gk["gFT_PSD_data"].where(~np.isnan(Gk["gFT_PSD_data"]), 0) * Gk["N_per_stancil"] + ).sum("beam") / Gk["N_per_stancil"].sum("beam") + G_gFT_wmean["N_per_stancil"] = Gk["N_per_stancil"].sum("beam") + + G_fft_wmean = (Gfft.where(~np.isnan(Gfft), 0) * Gfft["N_per_stancil"]).sum( + "beam" + ) / Gfft["N_per_stancil"].sum("beam") + G_fft_wmean["N_per_stancil"] = Gfft["N_per_stancil"].sum("beam") + Gmean = G_gFT_wmean.rolling(k=5, center=True).mean() + + try: + k_max = Gmean.k[Gmean.isel(x=slice(0, 5)).mean("x").argmax().data].data + except Exception: + k_max = Gmean.k[Gmean.isel(x=slice(0, 20)).mean("x").argmax().data].data + + k_max_range = (k_max * 0.75, k_max * 1, k_max * 1.25) + font_for_print() + F = M.figure_axis_xy(6.5, 5.6, container=True, view_scale=1) + Lmeters = Gk.L.data[0] + + plt.suptitle("gFT Slope Spectrograms\n" + track_name, y=0.98) + gs = GridSpec(3, 3, wspace=0.2, hspace=0.5) + + Gplot = ( + G_gFT_wmean.squeeze() + .rolling(k=10, min_periods=1, center=True) + .median() + .rolling(x=3, min_periods=1, center=True) + .median() + ) + dd = 10 * np.log10(Gplot) + dd = dd.where(~np.isinf(dd), np.nan) + clev_log = M.clevels([dd.quantile(0.01).data, dd.quantile(0.98).data * 1.2], 31) * 1 + xlims = Gmean.x[0] / 1e3, Gmean.x[-1] / 1e3 -def plot_model_eta(D, ax, offset=0, **kargs): - eta = D.eta # + D.x - y_data = D.y_model + offset - plt.plot(eta, y_data, **kargs) + k = high_beams[0] + for pos, k, pflag in zip( + [gs[0, 0], gs[0, 1], gs[0, 2]], high_beams, [True, False, False] + ): + ax0 = F.fig.add_subplot(pos) + Gplot = Gk.sel(beam=k).gFT_PSD_data.squeeze() + dd2 = 10 * np.log10(Gplot) + dd2 = dd2.where(~np.isinf(dd2), np.nan) + plot_wavenumber_spectrogram( + ax0, dd2, clev_log, title=k + " unsmoothed", plot_photon_density=True + ) + plt.xlim(xlims) + if pflag: + plt.ylabel("Wave length\n(meters)") + plt.legend() - ax.axvline(eta[0].data, linewidth=0.1, color=kargs["color"], alpha=0.5) - ax.axvline(eta[-1].data, linewidth=0.1, color=kargs["color"], alpha=0.5) + for pos, k, pflag in zip( + [gs[1, 0], gs[1, 1], gs[1, 2]], low_beams, [True, False, False] + ): + ax0 = F.fig.add_subplot(pos) + Gplot = Gk.sel(beam=k).gFT_PSD_data.squeeze() + dd2 = 10 * np.log10(Gplot) + dd2 = dd2.where(~np.isinf(dd2), np.nan) + plot_wavenumber_spectrogram( + ax0, dd2, clev_log, title=k + " unsmoothed", plot_photon_density=True + ) + plt.xlim(xlims) + if pflag: + plt.ylabel("Wave length\n(meters)") + plt.legend() + ax0 = F.fig.add_subplot(gs[2, 0]) -if "y_data" in Gx.sel(beam="gt3r").keys(): - print("ydata is ", ("y_data" in Gx.sel(beam="gt3r").keys())) -else: - print("ydata is ", ("y_data" in Gx.sel(beam="gt3r").keys())) - MT.json_save("B03_fail", plot_path, {"reason": "no y_data"}) - print("failed, exit") - exit() + plot_wavenumber_spectrogram( + ax0, + dd, + clev_log, + title="smoothed weighted mean \n10 $\log_{10}( (m/m)^2 m )$", + plot_photon_density=True, + ) + plt.xlim(xlims) -fltostr, numtostr = MT.float_to_str, MT.num_to_str + line_styles = ["--", "-", "--"] + for k_max, style in zip(k_max_range, line_styles): + ax0.axhline(2 * np.pi / k_max, color="red", linestyle=style, linewidth=0.5) -font_for_print() + if pflag: + plt.ylabel("Wave length\n(meters)") + plt.legend() -MT.mkdirs_r(plot_path + "B03_spectra/") + pos = gs[2, 1] + ax0 = F.fig.add_subplot(pos) + plt.title("Photons density ($m^{-1}$)", loc="left") -x_pos_sel = np.arange(Gk.x.size)[~np.isnan(Gk.mean("beam").mean("k").gFT_PSD_data.data)] -x_pos_max = ( - Gk.mean("beam") - .mean("k") - .gFT_PSD_data[~np.isnan(Gk.mean("beam").mean("k").gFT_PSD_data)] - .argmax() - .data -) -xpp = x_pos_sel[[int(i) for i in np.round(np.linspace(0, x_pos_sel.size - 1, 4))]] -xpp = np.insert(xpp, 0, x_pos_max) + for k in all_beams: + I = Gk.sel(beam=k)["gFT_PSD_data"] + plt.plot(Gplot.x / 1e3, I.N_photons / I.L.data, label=k, linewidth=0.8) + plt.plot( + Gplot.x / 1e3, + G_gFT_wmean.N_per_stancil / 3 / I.L.data, + c="black", + label="ave Photons", + linewidth=0.8, + ) + plt.xlim(xlims) + plt.xlabel("Distance from the Ice Edge (km)") -for i in xpp: - F = M.figure_axis_xy(6, 8, container=True, view_scale=0.8) + pos = gs[2, 2] - plt.suptitle( - "gFT Model and Spectrograms | x=" + str(Gk.x[i].data) + " \n" + track_name, - y=0.95, - ) - gs = GridSpec(5, 6, wspace=0.2, hspace=0.7) + ax0 = F.fig.add_subplot(pos) + ax0.set_yscale("log") - ax0 = F.fig.add_subplot(gs[0:2, :]) - col_d = color_schemes.__dict__["rels"] + plt.title("Peak Spectral Power", loc="left") - neven = True - offs = 0 + x0 = Gk.x[0].data for k in all_beams: - Gx_1 = Gx.isel(x=i).sel(beam=k) - Gk_1 = Gk.isel(x=i).sel(beam=k) - - plot_model_eta( - Gx_1, - ax0, - offset=offs, - linestyle="-", - color=col_d[k], - linewidth=0.4, - alpha=1, - zorder=12, + I = Gk.sel(beam=k)["gFT_PSD_data"] + plt.scatter( + I.x.data / 1e3, + I.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k").data, + s=0.5, + marker=".", + color="red", + alpha=0.3, ) - ylims = -np.nanstd(Gx_1.y_data) * 3, np.nanstd(Gx_1.y_data) * 3 - - # oringial data - eta_1 = plot_data_eta( - Gx_1, offset=offs, linestyle="-", c="k", linewidth=1, alpha=0.5, zorder=11 + I = Gfft.sel(beam=k) + plt.scatter( + (x0 + I.x.data) / 1e3, + I.power_spec.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), + s=0.5, + marker=".", + c="blue", + alpha=0.3, ) - # reconstruct in gaps - FT = gFT.generalized_Fourier(Gx_1.eta + Gx_1.x, None, Gk_1.k) - _ = FT.get_H() - FT.p_hat = np.concatenate([Gk_1.gFT_cos_coeff, Gk_1.gFT_sin_coeff]) - plt.plot( - Gx_1.eta, - FT.model() + offs, - "-", - c="orange", - linewidth=0.3, - alpha=1, - zorder=2, - ) + Gplot = G_fft_wmean.squeeze() + Gplot = Gplot.power_spec[ + :, Gplot.N_per_stancil >= Gplot.N_per_stancil.max().data * 0.9 + ] + plt.plot( + (x0 + Gplot.x) / 1e3, + Gplot.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), + ".", + markersize=1.5, + c="blue", + label="FFT", + ) + + Gplot = G_gFT_wmean.squeeze() + plt.plot( + Gplot.x / 1e3, + Gplot.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"), + ".", + markersize=1.5, + c="red", + label="gFT", + ) + + plt.ylabel("1e-3 $(m)^2~m$") + plt.legend() - if neven: - neven = False - offs += 0.3 - else: - neven = True - offs += 0.6 + F.save_light(path=plot_path, name="B03_specs_L" + str(Lmeters)) - dx = eta_1.diff("eta").mean().data + Gk.sel(beam=k).gFT_PSD_data.plot() - eta_ticks = np.linspace(Gx_1.eta.data[0], Gx_1.eta.data[-1], 11) + if "y_data" in Gx.sel(beam="gt3r").keys(): + print("ydata is ", ("y_data" in Gx.sel(beam="gt3r").keys())) + else: + print("ydata is ", ("y_data" in Gx.sel(beam="gt3r").keys())) + MT.json_save("B03_fail", plot_path, {"reason": "no y_data"}) + echo("failed, exit", "red") + exit() - ax0.set_xticks(eta_ticks) - ax0.set_xticklabels(eta_ticks / 1e3) - plt.xlim(eta_1[0].data - 40 * dx, eta_1[-1].data + 40 * dx) - plt.title("Model reconst.", loc="left") + fltostr, _ = MT.float_to_str, MT.num_to_str - plt.ylabel("relative slope (m/m)") - plt.xlabel( - "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km" + font_for_print() + + MT.mkdirs_r(plot_path / "B03_spectra/") + + x_pos_sel = np.arange(Gk.x.size)[ + ~np.isnan(Gk.mean("beam").mean("k").gFT_PSD_data.data) + ] + x_pos_max = ( + Gk.mean("beam") + .mean("k") + .gFT_PSD_data[~np.isnan(Gk.mean("beam").mean("k").gFT_PSD_data)] + .argmax() + .data ) + xpp = x_pos_sel[[int(i) for i in np.round(np.linspace(0, x_pos_sel.size - 1, 4))]] + xpp = np.insert(xpp, 0, x_pos_max) - # spectra - # define threshold - k_thresh = 0.085 - ax1_list = list() - dd_max = list() - for pos, kgroup, lflag in zip( - [gs[2, 0:2], gs[2, 2:4], gs[2, 4:]], - [["gt1l", "gt1r"], ["gt2l", "gt2r"], ["gt3l", "gt3r"]], - [True, False, False], - ): - ax11 = F.fig.add_subplot(pos) - ax11.tick_params(labelleft=lflag) - ax1_list.append(ax11) - for k in kgroup: - Gx_1 = Gx.isel(x=i).sel(beam=k) - Gk_1 = Gk.isel(x=i).sel(beam=k) + for i in xpp: + F = M.figure_axis_xy(6, 8, container=True, view_scale=0.8) - klim = Gk_1.k[0], Gk_1.k[-1] + plt.suptitle( + "gFT Model and Spectrograms | x=" + str(Gk.x[i].data) + " \n" + track_name, + y=0.95, + ) + gs = GridSpec(5, 6, wspace=0.2, hspace=0.7) - if "l" in k: - dd = Gk_1.gFT_PSD_data - plt.plot(Gk_1.k, dd, color="gray", linewidth=0.5, alpha=0.5) + ax0 = F.fig.add_subplot(gs[0:2, :]) + col_d = color_schemes.__dict__["rels"] - dd = Gk_1.gFT_PSD_data.rolling(k=10, min_periods=1, center=True).mean() - plt.plot(Gk_1.k, dd, color=col_d[k], linewidth=0.8) - # handle the 'All-NaN slice encountered' warning - if np.all(np.isnan(dd.data)): - dd_max.append(np.nan) + neven = True + offs = 0 + for k in all_beams: + Gx_1 = Gx.isel(x=i).sel(beam=k) + Gk_1 = Gk.isel(x=i).sel(beam=k) + + plot_model_eta( + Gx_1, + ax0, + offset=offs, + linestyle="-", + color=col_d[k], + linewidth=0.4, + alpha=1, + zorder=12, + ) + + # original data + eta_1 = plot_data_eta( + Gx_1, + offset=offs, + linestyle="-", + c="k", + linewidth=1, + alpha=0.5, + zorder=11, + ) + + # reconstruct in gaps + FT = gFT.generalized_Fourier(Gx_1.eta + Gx_1.x, None, Gk_1.k) + _ = FT.get_H() + FT.p_hat = np.concatenate([Gk_1.gFT_cos_coeff, Gk_1.gFT_sin_coeff]) + plt.plot( + Gx_1.eta, + FT.model() + offs, + "-", + c="orange", + linewidth=0.3, + alpha=1, + zorder=2, + ) + + if neven: + neven = False + offs += 0.3 else: - dd_max.append(np.nanmax(dd.data)) - - plt.xlim(klim) - if lflag: - plt.ylabel("$(m/m)^2/k$") - plt.title("Energy Spectra", loc="left") - - plt.xlabel("wavenumber k (2$\pi$ m$^{-1}$)") - - ax11.axvline(k_thresh, linewidth=1, color="gray", alpha=1) - ax11.axvspan(k_thresh, klim[-1], color="gray", alpha=0.5, zorder=12) - - if not np.all(np.isnan(dd_max)): - max_vale = np.nanmax(dd_max) - for ax in ax1_list: - ax.set_ylim(0,max_vale * 1.1) + neven = True + offs += 0.6 - ax0 = F.fig.add_subplot(gs[-2:, :]) + dx = eta_1.diff("eta").mean().data - neven = True - offs = 0 - for k in all_beams: - Gx_1 = Gx.isel(x=i).sel(beam=k) - Gk_1 = Gk.isel(x=i).sel(beam=k) + eta_ticks = np.linspace(Gx_1.eta.data[0], Gx_1.eta.data[-1], 11) - ylims = -np.nanstd(Gx_1.y_data) * 3, np.nanstd(Gx_1.y_data) * 3 + ax0.set_xticks(eta_ticks) + ax0.set_xticklabels(eta_ticks / 1e3) + plt.xlim(eta_1[0].data - 40 * dx, eta_1[-1].data + 40 * dx) + plt.title("Model reconst.", loc="left") - # oringial data - eta_1 = plot_data_eta( - Gx_1, offset=offs, linestyle="-", c="k", linewidth=1.5, alpha=0.5, zorder=11 + plt.ylabel("relative slope (m/m)") + plt.xlabel( + "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km" ) - # reconstruct in gaps - FT = gFT.generalized_Fourier(Gx_1.eta + Gx_1.x, None, Gk_1.k) - _ = FT.get_H() - FT.p_hat = np.concatenate([Gk_1.gFT_cos_coeff, Gk_1.gFT_sin_coeff]) + # spectra + # define threshold + k_thresh = 0.085 + ax1_list = list() + dd_max = list() + for pos, kgroup, lflag in zip( + [gs[2, 0:2], gs[2, 2:4], gs[2, 4:]], + [["gt1l", "gt1r"], ["gt2l", "gt2r"], ["gt3l", "gt3r"]], + [True, False, False], + ): + ax11 = F.fig.add_subplot(pos) + ax11.tick_params(labelleft=lflag) + ax1_list.append(ax11) + for k in kgroup: + Gx_1 = Gx.isel(x=i).sel(beam=k) + Gk_1 = Gk.isel(x=i).sel(beam=k) + + klim = Gk_1.k[0], Gk_1.k[-1] + + if "l" in k: + dd = Gk_1.gFT_PSD_data + plt.plot(Gk_1.k, dd, color="gray", linewidth=0.5, alpha=0.5) + + dd = Gk_1.gFT_PSD_data.rolling(k=10, min_periods=1, center=True).mean() + plt.plot(Gk_1.k, dd, color=col_d[k], linewidth=0.8) + # handle the 'All-NaN slice encountered' warning + if np.all(np.isnan(dd.data)): + dd_max.append(np.nan) + else: + dd_max.append(np.nanmax(dd.data)) + + plt.xlim(klim) + if lflag: + plt.ylabel("$(m/m)^2/k$") + plt.title("Energy Spectra", loc="left") + + plt.xlabel("wavenumber k (2$\pi$ m$^{-1}$)") + + ax11.axvline(k_thresh, linewidth=1, color="gray", alpha=1) + ax11.axvspan(k_thresh, klim[-1], color="gray", alpha=0.5, zorder=12) + + if not np.all(np.isnan(dd_max)): + max_vale = np.nanmax(dd_max) + for ax in ax1_list: + ax.set_ylim(0, max_vale * 1.1) - p_hat_k = np.concatenate([Gk_1.k, Gk_1.k]) - k_mask = p_hat_k < k_thresh - FT.p_hat[~k_mask] = 0 + ax0 = F.fig.add_subplot(gs[-2:, :]) - plt.plot( - Gx_1.eta, - FT.model() + offs, - "-", - c=col_d[k], - linewidth=0.8, - alpha=1, - zorder=12, - ) + neven = True + offs = 0 + for k in all_beams: + Gx_1 = Gx.isel(x=i).sel(beam=k) + Gk_1 = Gk.isel(x=i).sel(beam=k) + + # original data + eta_1 = plot_data_eta( + Gx_1, + offset=offs, + linestyle="-", + c="k", + linewidth=1.5, + alpha=0.5, + zorder=11, + ) + + # reconstruct in gaps + FT = gFT.generalized_Fourier(Gx_1.eta + Gx_1.x, None, Gk_1.k) + _ = FT.get_H() + FT.p_hat = np.concatenate([Gk_1.gFT_cos_coeff, Gk_1.gFT_sin_coeff]) + + p_hat_k = np.concatenate([Gk_1.k, Gk_1.k]) + k_mask = p_hat_k < k_thresh + FT.p_hat[~k_mask] = 0 + + plt.plot( + Gx_1.eta, + FT.model() + offs, + "-", + c=col_d[k], + linewidth=0.8, + alpha=1, + zorder=12, + ) + + if neven: + neven = False + offs += 0.3 + else: + neven = True + offs += 0.6 - if neven: - neven = False - offs += 0.3 - else: - neven = True - offs += 0.6 + dx = eta_1.diff("eta").mean().data - dx = eta_1.diff("eta").mean().data + eta_ticks = np.linspace(Gx_1.eta.data[0], Gx_1.eta.data[-1], 11) - eta_ticks = np.linspace(Gx_1.eta.data[0], Gx_1.eta.data[-1], 11) + ax0.set_xticks(eta_ticks) + ax0.set_xticklabels(eta_ticks / 1e3) + plt.xlim(eta_1[1000].data - 40 * dx, eta_1[-1000].data + 40 * dx) + plt.title("Low-Wavenumber Model reconst.", loc="left") - ax0.set_xticks(eta_ticks) - ax0.set_xticklabels(eta_ticks / 1e3) - plt.xlim(eta_1[1000].data - 40 * dx, eta_1[-1000].data + 40 * dx) - plt.title("Low-Wavenumber Model reconst.", loc="left") + plt.ylabel("relative slope (m/m)") + plt.xlabel( + "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km" + ) - plt.ylabel("relative slope (m/m)") - plt.xlabel( - "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km" + F.save_pup( + path=str(plot_path) + "B03_spectra/", name="B03_freq_reconst_x" + str(i) + ) + + MT.json_save( + "B03_success", + plot_path, + {"time": "time.asctime( time.localtime(time.time()) )"}, ) - F.save_pup(path=plot_path + "B03_spectra/", name="B03_freq_reconst_x" + str(i)) -MT.json_save( - "B03_success", plot_path, {"time": "time.asctime( time.localtime(time.time()) )"} -) +if __name__ == "__main__": + step3app = makeapp(run_B03_plot_spectra_ov, name="plotspectra") + step3app()