Merge pull request #198 from brown-ccv/feat-logger/6-reformat-everything

chore: reformat all the files using black

src/icesat2waves/__init__.py

@@ -1 +1 @@
-"""Tools to analyse data from ICESat-2."""
+"""Tools to analyse data from ICESat-2."""

src/icesat2waves/analysis_db/A02c_IOWAGA_thredds_prior.py

@@ -412,21 +412,13 @@ def draw_range(lon_range, lat_range, *args, **kwargs):
             draw_range(lon_range, lat_range, c="blue", linewidth=0.7, zorder=10)
 
             if fv != "ice":
-                cm = plt.pcolor(
-                    lon, lat, G_beam[fv], vmin=cl[0], vmax=cl[-1], cmap=fc
-                )
+                cm = plt.pcolor(lon, lat, G_beam[fv], vmin=cl[0], vmax=cl[-1], cmap=fc)
                 if G_beam.ice.shape[0] > 1:
-                    plt.contour(
-                        lon, lat, G_beam.ice, colors="black", linewidths=0.6
-                    )
+                    plt.contour(lon, lat, G_beam.ice, colors="black", linewidths=0.6)
             else:
-                cm = plt.pcolor(
-                    lon, lat, G_beam[fv], vmin=cl[0], vmax=cl[-1], cmap=fc
-                )
+                cm = plt.pcolor(lon, lat, G_beam[fv], vmin=cl[0], vmax=cl[-1], cmap=fc)
 
-            plt.title(
-                G_beam[fv].long_name.replace(" ", "\n") + "\n" + fv, loc="left"
-            )
+            plt.title(G_beam[fv].long_name.replace(" ", "\n") + "\n" + fv, loc="left")
             ax1.axis("equal")
 
             ax2 = F.fig.add_subplot(gs[-1, fp])
@@ -471,9 +463,7 @@ def test_nan_frac(imask):
         }
 
         # flatten key_list_pairs.values to a list
-        key_list_pairs2 = [
-            item for pair in key_list_pairs.values() for item in pair
-        ]
+        key_list_pairs2 = [item for pair in key_list_pairs.values() for item in pair]
 
         key_list_scaler = set(key_list) - set(key_list_pairs2)
 
@@ -501,12 +491,8 @@ def test_nan_frac(imask):
             "lontigude",  # TODO: fix typo?
         ]
         Tend["lat"] = [
-            ice_mask_prior.latitude[ice_mask_prior.sum("longitude") == 0]
-            .mean()
-            .data,
-            ice_mask_prior.latitude[ice_mask_prior.sum("longitude") == 0]
-            .std()
-            .data,
+            ice_mask_prior.latitude[ice_mask_prior.sum("longitude") == 0].mean().data,
+            ice_mask_prior.latitude[ice_mask_prior.sum("longitude") == 0].std().data,
             "latitude",
         ]
         Tend = Tend.T

src/icesat2waves/analysis_db/B02_make_spectra_gFT.py

@@ -88,7 +88,6 @@ def run_B02_make_spectra_gFT(
     }
     report_input_parameters(**kargs)
 
-
     workdir, _ = update_paths_mconfig(output_dir, mconfig)
 
     load_path = Path(workdir, batch_key, "B01_regrid")
@@ -111,9 +110,7 @@ def run_B02_make_spectra_gFT(
     nan_fraction = list()
     for k in all_beams:
         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]
-        )
+        nan_fraction.append(np.sum(np.isnan(heights_c_std)) / heights_c_std.shape[0])
 
     del heights_c_std
 
@@ -185,7 +182,9 @@ def run_B02_make_spectra_gFT(
     for k in all_beams:
         dist_i = io.get_beam_var_hdf_store(Gd[k], "x")
         x_mask = (dist_i > xlims[0]) & (dist_i < xlims[1])
-        _logger.debug("k: %s, sum/range: %s", k, sum(x_mask["x"]) / (xlims[1] - xlims[0]))
+        _logger.debug(
+            "k: %s, sum/range: %s", k, sum(x_mask["x"]) / (xlims[1] - xlims[0])
+        )
 
     _logger.debug("-reduced frequency resolution")
     kk = kk[::2]

src/icesat2waves/analysis_db/B03_plot_spectra_ov.py

@@ -139,9 +139,7 @@ def run_B03_plot_spectra_ov(
     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(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")
@@ -185,8 +183,7 @@ def run_B03_plot_spectra_ov(
 
     # TODO: refactor to make more readable. CP
     G_gFT_wmean = (
-        Gk["gFT_PSD_data"].where(~np.isnan(Gk["gFT_PSD_data"]), 0)
-        * Gk["N_per_stancil"]
+        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")
 
@@ -219,9 +216,7 @@ def run_B03_plot_spectra_ov(
     )
     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
-    )
+    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
 
@@ -313,9 +308,7 @@ def run_B03_plot_spectra_ov(
         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"
-            ),
+            I.power_spec.sel(k=slice(k_max_range[0], k_max_range[2])).integrate("k"),
             s=0.5,
             marker=".",
             c="blue",
@@ -376,19 +369,14 @@ def run_B03_plot_spectra_ov(
         .argmax()
         .data
     )
-    xpp = x_pos_sel[
-        [int(i) for i in np.round(np.linspace(0, x_pos_sel.size - 1, 4))]
-    ]
+    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 i in xpp:
         F = M.FigureAxisXY(6, 8, container=True, view_scale=0.8)
 
         plt.suptitle(
-            "gFT Model and Spectrograms | x="
-            + str(Gk.x[i].data)
-            + " \n"
-            + track_name,
+            "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)
@@ -457,9 +445,7 @@ def run_B03_plot_spectra_ov(
         plt.ylabel("relative slope (m/m)")
         # TODO: compute xlabel as fstring. CP
         plt.xlabel(
-            "segment distance $\eta$ (km) @ x="
-            + fltostr(Gx_1.x.data / 1e3, 2)
-            + "km"
+            "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km"
         )
 
         # spectra
@@ -485,9 +471,7 @@ def run_B03_plot_spectra_ov(
                     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()
+                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)):
@@ -567,14 +551,10 @@ def run_B03_plot_spectra_ov(
         plt.ylabel("relative slope (m/m)")
         # TODO: compute xlabel as fstring. CP
         plt.xlabel(
-            "segment distance $\eta$ (km) @ x="
-            + fltostr(Gx_1.x.data / 1e3, 2)
-            + "km"
+            "segment distance $\eta$ (km) @ x=" + fltostr(Gx_1.x.data / 1e3, 2) + "km"
         )
 
-        F.save_pup(
-            path=str(plot_path / "B03_spectra"), name=f"B03_freq_reconst_x{i}"
-        )
+        F.save_pup(path=str(plot_path / "B03_spectra"), name=f"B03_freq_reconst_x{i}")
 
     MT.json_save(
         "B03_success",

src/icesat2waves/analysis_db/B04_angle.py

@@ -142,12 +142,10 @@ def run_B04_angle(
 
     #### 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"
-    ]
+    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))]
@@ -266,7 +264,9 @@ def run_B04_angle(
     prior_angle = Prior_smth.Prior_direction * 180 / np.pi
     if (abs(prior_angle) > 80).all():
         _logger.critical(
-            "Prior angle is %s %s. Quit.", prior_angle.min().data, prior_angle.max().data,
+            "Prior angle is %s %s. Quit.",
+            prior_angle.min().data,
+            prior_angle.max().data,
         )
         dd_save = {
             "time": time.asctime(time.localtime(time.time())),
@@ -452,9 +452,7 @@ def plot_instance(
                 "-",
                 c=next(col_list),
             )
-            plt.xlim(
-                x_concat[y_concat == y_pos][0], x_concat[y_concat == y_pos][-1]
-            )
+            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])
@@ -465,9 +463,7 @@ def plot_instance(
             )
 
         if optimze is True:
-            SM.plot_optimze(
-                color="r", markersize=10, zorder=12, label="Dual Annealing"
-            )
+            SM.plot_optimze(color="r", markersize=10, zorder=12, label="Dual Annealing")
 
         if sample is True:
             SM.plot_sample(
@@ -624,9 +620,7 @@ def plot_instance(
         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_direction.data,
                 Prior_smth.sel(k=k_prime_max, method="nearest").Prior_spread.data,
             )
         }
@@ -657,9 +651,7 @@ def get_instance(k_pair):
                     Prior_smth.sel(
                         k=k_prime_max, method="nearest"
                     ).Prior_direction.data,
-                    Prior_smth.sel(
-                        k=k_prime_max, method="nearest"
-                    ).Prior_spread.data,
+                    Prior_smth.sel(k=k_prime_max, method="nearest").Prior_spread.data,
                 )
             }
 
@@ -676,9 +668,7 @@ def get_instance(k_pair):
                 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
-            )
+            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
@@ -714,9 +704,7 @@ def get_instance(k_pair):
                 "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
-            )
+            z_model = SM.objective_func(fitter.params, *fitting_args, test_flag=True)
             cost = (fitter.residual**2).sum() / (z_concat**2).sum()
 
             if plot_flag:
@@ -822,9 +810,7 @@ def get_instance(k_pair):
         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_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():

src/icesat2waves/analysis_db/B05_define_angle.py

@@ -325,9 +325,7 @@ def define_angle(
             ax_list[group] = ax0
 
             data = corrected_marginals.isel(x=xi).sel(beam_group=group)
-            weights = derive_weights(
-                Marginals.weight.isel(x=xi).sel(beam_group=group)
-            )
+            weights = derive_weights(Marginals.weight.isel(x=xi).sel(beam_group=group))
             weights = weights**2
 
             # derive angle axis
@@ -349,9 +347,9 @@ def define_angle(
             data_collect[group] = data_wmean
 
         data_collect = xr.concat(data_collect.values(), dim="beam_group")
-        final_data = (group_weight * data_collect).sum(
+        final_data = (group_weight * data_collect).sum("beam_group") / group_weight.sum(
             "beam_group"
-        ) / group_weight.sum("beam_group").data
+        ).data
 
         plt.sca(ax_sum)
         plt.stairs(final_data, x_angle, color="k", alpha=1, linewidth=0.8)
@@ -378,9 +376,7 @@ def define_angle(
 
         priors_k = Marginals.Prior_direction[~np.isnan(k_mask.isel(x=xi))]
         for pk in priors_k:
-            ax_final.axvline(
-                pk, color=color_schemes.cascade2, linewidth=1, alpha=0.7
-            )
+            ax_final.axvline(pk, color=color_schemes.cascade2, linewidth=1, alpha=0.7)
 
         plt.stairs(final_data, x_angle, color="k", alpha=0.5, linewidth=0.8)
 
@@ -466,9 +462,7 @@ def define_angle(
     _title = f"{track_name}\nPower Spectra (m/m)$^2$ k$^{{-1}}$"
     plt.title(_title, loc="left")
 
-    cbar = plt.colorbar(
-        fraction=0.018, pad=0.01, orientation="vertical", label="Power"
-    )
+    cbar = plt.colorbar(fraction=0.018, pad=0.01, orientation="vertical", label="Power")
     cbar.outline.set_visible(False)
     clev_ticks = np.round(clev_spec[::3], 0)
     cbar.set_ticks(clev_ticks)
@@ -529,9 +523,9 @@ def define_angle(
         i_spec = weighted_spec.sel(x=slice(x_range[0], x_range[-1]))
         i_dir = corrected_marginals.sel(x=slice(x_range[0], x_range[-1]))
 
-        dir_data = (i_dir * i_dir.N_data).sum(
+        dir_data = (i_dir * i_dir.N_data).sum(["beam_group", "x"]) / i_dir.N_data.sum(
             ["beam_group", "x"]
-        ) / i_dir.N_data.sum(["beam_group", "x"])
+        )
         lims = get_first_and_last_nonzero_data(dir_data)
 
         plot_data = build_plot_data(dir_data, i_spec, lims)