diff --git a/README.md b/README.md
index 7751c77..e36299b 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 # smap-drydown
 This repository contains code for analyzing soil moisture drydowns from [the Soil Moisture Active Passive (SMAP)](https://smap.jpl.nasa.gov/data/) as detailed in the corresponding manuscript: 
 
-    Araki, R., Morgan, B.E., McMillan, H.K., Caylor, K.K. Nonlinear Soil Moisture Loss Function Reveals Vegetation Responses to Water Availability. Geophysical Research Letters (in review)
+    Araki, R., Morgan, B.E., McMillan, H.K., Caylor, K.K. Nonlinear Soil Moisture Loss Function Reveals Vegetation Responses to Water Availability. Submitted to Geophysical Research Letters (in review)
 
 ## Getting started
 1. Clone the repository
@@ -9,30 +9,64 @@ This repository contains code for analyzing soil moisture drydowns from [the Soi
 $ git clone git@github.com:RY4GIT/smap-drydown.git
 ```
 
-2. Create a virtual environment (select appropriate yml file according to your OS). 
+2. Create a virtual environment (select an appropriate yml file according to your OS). 
 ```bash
 $ cd smap-drydown
 $ conda env create -f environment_linux.yml
 $ conda activate SMAP
 ```
 
-3. Download the SMAP and ancillary data from appropriate sources.
+3. Download the SMAP and ancillary data from appropriate sources using scripts in `data_mng`
 
 4. In the `analysis` directory, create `config.ini`, based on `config_example.ini`
 
 5. Run `analysis\__main__.py`
 
-## Contents
+6. Visualize the results using scripts in `notebooks`. The results file is large (~130 MB) and is therefore available upon request.
 
-### `data_mng` 
-Contains scripts to retrieve & curate input data.
+## Contents
 
 ### `analysis`
 Contains scripts to implement the drydown analysis and model fits. 
-This code has been further refactored in https://github.com/ecohydro/drydowns: checkout if you are interested in. 
+
+The functions for loss calculations and drydown models are contained in `DrydownModel.py`. 
+This code has been further refactored in https://github.com/ecohydro/drydowns; check it out if you are interested.
+
+
+### `data_mng` 
+Contains scripts to retrieve and curate input data. 
+
+All data are pre-curated in "datarods" format, which stores the data as a long time series at a single SMAP grid.
+
+- SMAP soil moisture data
+    - Download data using `retrieve_NSIDC_Data_SPL3SMP.ipynb`
+    - Preprocess data using `create_datarods_SPL3SMP.ipynb`
+- SMAP precipitation data
+    - Download data using `retrieve_NSIDC_Data_SPL4SMGP.ipynb`
+    - Preprocess data using `create_datarods_SPL4SMGP.py`
+- dPET (Singer et al., 2020) data
+    - Download daily data from [the website](http://doi.org/10.5523/bris.qb8ujazzda0s2aykkv0oq0ctp)
+    - Preprocess data using `create_datarods_PET.py`
+- SMAP ancillary data
+    - Download data from [the website](https://doi.org/10.5067/HB8BPJ13TDQJ)
+    - Preprocess data using `read_ancillary_landcover_data.ipynb`
+    - After obtaining precipitation and PET data, run `calc_aridityindex.py`
+- Rangeland data
+    - Download data using `retrieve_rangeland_data.sh`
+    - Preprocess data using `read_rangeland_data.py`
+- Other utilities
+    - `retrieve_NSIDC_Data_datacheck.ipynb`: check if all the data are downloaded from NSIDC
+    - `create_datarods_datacheck.ipynb`: check if all the data are preprocessed
+    - `identify_unusable_grid.ipynb`: identify grids located on open water
+
 
 ### `notebooks`
-Contains scripts that are used to test functions or visualise results
+Contains scripts used to test functions or visualize the models and results
+- `figs_stats_datapreprocess.py`: Preprocess result files to reduce execution time
+- `figs_method.py` & `figs_method_tau.py`: Visualize the loss functions and drydown models
+- `figs_stats.py` and `figs_stats_rev.py`: Visualize the results
+- `figs_drydown.py`: Plot observed and modeled drydown curves
+
 
 ## Contact
-Ryoko Araki, raraki8150 (at) sdsu.edu
\ No newline at end of file
+Ryoko Araki, raraki8159 (at) sdsu.edu
\ No newline at end of file
diff --git a/data_mng/read_MODISvegetation.ipynb b/data_mng/archive/read_MODISvegetation.ipynb
similarity index 100%
rename from data_mng/read_MODISvegetation.ipynb
rename to data_mng/archive/read_MODISvegetation.ipynb
diff --git a/data_mng/read_bassiouni_data.ipynb b/data_mng/archive/read_bassiouni_data.ipynb
similarity index 100%
rename from data_mng/read_bassiouni_data.ipynb
rename to data_mng/archive/read_bassiouni_data.ipynb
diff --git a/data_mng/retrieve_MODIS44B.sh b/data_mng/archive/retrieve_MODIS44B.sh
similarity index 100%
rename from data_mng/retrieve_MODIS44B.sh
rename to data_mng/archive/retrieve_MODIS44B.sh
diff --git a/data_mng/retrieve_MODIS_LAI_from_Appeears_pt_sample.py b/data_mng/archive/retrieve_MODIS_LAI_from_Appeears_pt_sample.py
similarity index 100%
rename from data_mng/retrieve_MODIS_LAI_from_Appeears_pt_sample.py
rename to data_mng/archive/retrieve_MODIS_LAI_from_Appeears_pt_sample.py
diff --git a/data_mng/retrieve_MODIS_NDVI_from_Appears_pt_sample.ipynb b/data_mng/archive/retrieve_MODIS_NDVI_from_Appears_pt_sample.ipynb
similarity index 100%
rename from data_mng/retrieve_MODIS_NDVI_from_Appears_pt_sample.ipynb
rename to data_mng/archive/retrieve_MODIS_NDVI_from_Appears_pt_sample.ipynb
diff --git a/data_mng/retrieve_NSIDC_Data_SPL3SMP_E.ipynb b/data_mng/archive/retrieve_NSIDC_Data_SPL3SMP_E.ipynb
similarity index 100%
rename from data_mng/retrieve_NSIDC_Data_SPL3SMP_E.ipynb
rename to data_mng/archive/retrieve_NSIDC_Data_SPL3SMP_E.ipynb
diff --git a/data_mng/retrieve_appeears_smap_data.py b/data_mng/archive/retrieve_appeears_smap_data.py
similarity index 100%
rename from data_mng/retrieve_appeears_smap_data.py
rename to data_mng/archive/retrieve_appeears_smap_data.py
diff --git a/notebooks/check_ancillary_info.ipynb b/notebooks/archive/check_ancillary_info.ipynb
similarity index 99%
rename from notebooks/check_ancillary_info.ipynb
rename to notebooks/archive/check_ancillary_info.ipynb
index 5fba677..40db03f 100644
--- a/notebooks/check_ancillary_info.ipynb
+++ b/notebooks/archive/check_ancillary_info.ipynb
@@ -1960,7 +1960,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "SMAP",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -1974,9 +1974,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }
diff --git a/notebooks/figs_stats_agu24.py b/notebooks/archive/figs_stats_agu24.py
similarity index 100%
rename from notebooks/figs_stats_agu24.py
rename to notebooks/archive/figs_stats_agu24.py
diff --git a/notebooks/test_covmat.py b/notebooks/archive/test_covmat.py
similarity index 100%
rename from notebooks/test_covmat.py
rename to notebooks/archive/test_covmat.py
diff --git a/notebooks/test_event_separation_w_rainfall_v2.ipynb b/notebooks/archive/test_event_separation_w_rainfall_v2.ipynb
similarity index 100%
rename from notebooks/test_event_separation_w_rainfall_v2.ipynb
rename to notebooks/archive/test_event_separation_w_rainfall_v2.ipynb
diff --git a/notebooks/toy_drydown_models.ipynb b/notebooks/archive/toy_drydown_models.ipynb
similarity index 100%
rename from notebooks/toy_drydown_models.ipynb
rename to notebooks/archive/toy_drydown_models.ipynb
diff --git a/notebooks/figs_drydown.py b/notebooks/figs_drydown.py
index 256432d..ad8acb7 100644
--- a/notebooks/figs_drydown.py
+++ b/notebooks/figs_drydown.py
@@ -22,8 +22,10 @@
 # %% Plot config
 
 ############ CHANGE HERE FOR CHECKING DIFFERENT RESULTS ###################
-dir_name = f"raraki_2024-05-13_global_piecewise"  # f"raraki_2024-04-26"
+dir_name = f"raraki_2024-12-03_revision"  # f"raraki_2024-04-26"
 ###########################################################################
+# f"raraki_2024-05-13_global_piecewise" was used for the 1st version of the manuscript
+# f"raraki_2024-12-03_revision" was used for the revised version of the manuscript
 
 ################ CHANGE HERE FOR PLOT VISUAL CONFIG #########################
 
@@ -95,7 +97,7 @@ def get_dataframe(varname, event):
         EASE_row_index=event.EASE_row_index,
         EASE_column_index=event.EASE_column_index,
     )
-    _df = pd.read_csv(os.path.join(data_dir, datarod_dir, varname, fn))
+    _df = pd.read_csv(os.path.join(data_dir, datarods_dir, varname, fn))
 
     # Set time index and crop
     df = set_time_index(_df, index_name="time")
@@ -598,21 +600,55 @@ def plot_drydown(
 print(f"Next to try (in df): {df_filt.sample(n=1).index}")
 print(f"Next to try (not in df): {not_in_filt_indices.to_series().sample(n=1).index}")
 # %%
-# # check_1ts_range(df.loc[event_id], verbose=True)
-# # %%
-# plt.scatter(df["event_length"], df["q_q"])
-# plt.scatter(df_filt["event_length"], df_filt["q_q"])
-
-# # %%
-# plt.scatter(df["large_q_criteria"], df["q_q"])
-# plt.scatter(df_filt["large_q_criteria"], df_filt["q_q"])
-# # %%
-# # %%
-# df_filt["q_r_squared"].hist()
-# plt.xlim([0, 1])
-# # %%
-# plt.scatter(df_filt["q_q"], df_filt["q_ETmax"])
-# # %%
-# ax = df["q_q"].hist(vmin=0, vmax=4)
-# ax.set_xlim([0, 3])
+
+# %%
+##########################################################
+# Select plot for revision
+
+# # CONUS
+# lat_min, lat_max = 24.396308, 49.384358
+# lon_min, lon_max = -125.000000, -66.934570
+# Define bounding box coordinates
+lon_min, lon_max = 67.507970, 89.653185
+lat_min, lat_max = 9.279354, 31.365056
+
+df_filt = df[
+    (df["q_r_squared"] < df["tauexp_r_squared"])
+    & (df["q_r_squared"] < 0.95)
+    & (df["q_r_squared"] > 0.8)
+    & (df["sm_range"] > 0.20)
+    & (df["large_q_criteria"] < 0.8)
+    & (df["first3_avail2"])
+    & (df["q_q"] > 1.0e-04)
+    & (df["event_length"] >= 7)
+    & (df["AI"] < 0.5)
+    # & (df["latitude"] >= lat_min)
+    # & (df["latitude"] <= lat_max)
+    # & (df["longitude"] >= lon_min)
+    # & (df["longitude"] <= lon_max)
+    # & (df["q_q"] > 1.7)
+    # & (df["q_q"] < 2.0)
+    & (df["sand_fraction"] > 0.7)
+]
+# print(df_filt)
+print(f"Try (in df): {df_filt.sample(n=5).index}")
+
+# %%
+df_filt.columns
+# %%|
+################################################
+event_id = 304485
+################################################
+
+# Arid example: 462921, 127783, 602222, 218603: 816131, 304485, 320094
+# Cropland example: 506447, 484442,
+
+save = True
+plot_drydown(
+    df=df, event_id=event_id, plot_precip=False, save=save, days_after_to_plot=13
+)
+
+# print(df.loc[event_id])
+print(f"Next to try (in df): {df_filt.sample(n=1).index}")
+print(f"Next to try (not in df): {not_in_filt_indices.to_series().sample(n=1).index}")
 # %%
diff --git a/notebooks/figs_stats_datapreprocess.py b/notebooks/figs_stats_datapreprocess.py
index 35de718..d02a797 100644
--- a/notebooks/figs_stats_datapreprocess.py
+++ b/notebooks/figs_stats_datapreprocess.py
@@ -37,7 +37,7 @@
 ############################|###############################################
 # f"raraki_2024-05-13_global_piecewise" was used for the 1st version of the manuscript
 
-# Ryoko Araki, Bryn Morgan, Hilary K McMillan, Kelly K Caylor. 
+# Ryoko Araki, Bryn Morgan, Hilary K McMillan, Kelly K Caylor.
 # Nonlinear Soil Moisture Loss Function Reveals Vegetation Responses to Water Availability. ESS Open Archive . August 01, 2024.
 # DOI: 10.22541/essoar.172251989.99347091/v1
 
@@ -137,6 +137,7 @@
 df = df.assign(diff_bic_q_tauexp=df["q_bic"] - df["tauexp_bic"])
 df = df.assign(diff_bic_q_exp=df["q_bic"] - df["exp_bic"])
 
+
 def df_availability(row):
     # Check df point availability in the first 3 time steps of observation
     # Define a helper function to convert string to list
diff --git a/notebooks/figs_stats_rev.py b/notebooks/figs_stats_rev.py
index a67c32e..e2005cc 100644
--- a/notebooks/figs_stats_rev.py
+++ b/notebooks/figs_stats_rev.py
@@ -52,6 +52,7 @@
 dir_name = f"raraki_2024-12-03_revision"  # "raraki_2024-02-02"  # f"raraki_2023-11-25_global_95asmax"
 ###########################################################################
 # f"raraki_2024-05-13_global_piecewise" was used for the 1st version of the manuscript
+# f"raraki_2024-12-03_revision" was used for the revised version of the manuscript
 
 # Ryoko Araki, Bryn Morgan, Hilary K McMillan, Kelly K Caylor.
 # Nonlinear Soil Moisture Loss Function Reveals Vegetation Responses to Water Availability. ESS Open Archive . August 01, 2024.
@@ -103,10 +104,11 @@
 # DATA IMPORT
 
 df = pd.read_csv(os.path.join(output_dir, dir_name, results_file))
-print("Loaded results file")
+print("Loaded results file\n")
 
 coord_info = pd.read_csv(os.path.join(data_dir, datarods_dir, coord_info_file))
 
+# %%
 # Get bins for ancillary data
 sand_bin_list = [i * 0.1 for i in range(11)]
 sand_bin_list = sand_bin_list[1:]
@@ -193,8 +195,13 @@ def print_model_success(message, df):
 print_model_success("q model fit successful:", df_filt_q)
 print_model_success("exp model fit successful:", df_filt_exp)
 print_model_success("tau-exp model fit successful:", df_filt_tauexp)
-print_model_success("both q and exp model fit successful:", df_filt_q_and_exp)
-print_model_success("both q and tau-exp model fit successful:", df_filt_q_and_tauexp)
+print_model_success(
+    "both q and exp model fit successful (used for comparison):", df_filt_q_and_exp
+)
+print_model_success(
+    "both q and tau-exp model fit successful (used for comparison):",
+    df_filt_q_and_tauexp,
+)
 # print_model_success("either q or tau-exp:", df_filt_q_or_tauexp)
 # print_model_success("either q or exp model fit successful:", df_filt_q_or_exp)
 print_model_success(
@@ -210,17 +217,31 @@ def print_performance_comparison(df, model1, model2):
     )
 
 
+# %% In terms of AIC, BIC, and p-values
+print("============ EVENT-WISE COMPARISON ==============")
+print(r"In terms of $R^2$")
+print("tau-linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_tauexp, "q_r_squared", "tauexp_r_squared")
+print("linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_exp, "q_r_squared", "exp_r_squared")
 
-# %% In terms of AIC, BIC, and p-values
+
+print("\n")
+print(r"In terms of $AIC$")
+print("tau-linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_tauexp, "q_aic", "tauexp_aic")
+print("linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_exp, "q_aic", "exp_aic")
 
+print("\n")
+print(r"In terms of $AIC$")
+print("tau-linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_tauexp, "q_bic", "tauexp_bic")
+print("linear vs nonlinear")
 print_performance_comparison(df_filt_q_and_exp, "q_bic", "exp_bic")
 
 
+# %%
 def print_p_value(df, varname, thresh):
     n_better = sum(df[varname] < thresh)
     percentage_better = n_better / len(df) * 100
@@ -229,7 +250,11 @@ def print_p_value(df, varname, thresh):
     )
 
 
+print("\n")
+print(r"In terms of $p$-value")
+print("Pre-filtered q values")
 print_p_value(df_filt_q, "q_eq_1_p", 0.05)
+print("All q values")
 print_p_value(df_filt_allq, "q_eq_1_p", 0.05)
 
 
@@ -276,17 +301,24 @@ def count_model_performance(df, varname):
     # sns.histplot(grouped["count"], binwidth=0.5, color="#2c7fb8", fill=False, linewidth=3)
 
 
-print("\nIn terms of $R^2$")
-count_model_performance(df_filt_q_or_exp, "diff_R2_q_tauexp")
-count_model_performance(df_filt_q_or_exp, "diff_R2_q_exp")
+print("============ GRID-WISE COMPARISON ==============")
+print(r"In terms of $R^2$")
+print("tau-linear vs nonlinear")
+count_model_performance(df_filt_q_and_tauexp, "diff_R2_q_tauexp")
+print("linear vs nonlinear")
+count_model_performance(df_filt_q_and_exp, "diff_R2_q_exp")
 
-print("\nIn terms of $AIC$")
-count_model_performance(df_filt_q_or_exp, "diff_aic_q_tauexp")
-count_model_performance(df_filt_q_or_exp, "diff_aic_q_exp")
+print(r"In terms of $AIC$")
+print("tau-linear vs nonlinear")
+count_model_performance(df_filt_q_and_tauexp, "diff_aic_q_tauexp")
+print("linear vs nonlinear")
+count_model_performance(df_filt_q_and_exp, "diff_aic_q_exp")
 
-print("\nIn terms of $BIC$")
-count_model_performance(df_filt_q_or_exp, "diff_bic_q_tauexp")
-count_model_performance(df_filt_q_or_exp, "diff_bic_q_exp")
+print(r"In terms of $BIC$")
+print("tau-linear vs nonlinear")
+count_model_performance(df_filt_q_and_tauexp, "diff_bic_q_tauexp")
+print("linear vs nonlinear")
+count_model_performance(df_filt_q_and_exp, "diff_bic_q_exp")
 
 
 # %%
@@ -655,6 +687,43 @@ def print_global_stats(df, diff_var, model_desc):
 
 save_figure(fig_map_tau, fig_dir, f"tau_map_{stat_type}", "png", 900)
 # save_figure(fig_map_q, fig_dir, f"q_map_{stat_type}", "pdf", 1200)
+
+
+# %%
+def plot_hist(df, var_key):
+    fig, ax = plt.subplots(figsize=(5.5, 5))
+
+    # Create the histogram with a bin width of 1
+    sns.histplot(
+        df[var_key], binwidth=0.2, color="tab:blue", fill=False, linewidth=3, ax=ax
+    )
+
+    # Calculate median and mean
+    median_value = df[var_key].median()
+    mean_value = df[var_key].mean()
+
+    # Add median and mean as vertical lines
+    ax.axvline(
+        median_value, color="tab:grey", linestyle="--", linewidth=3, label=f"Median"
+    )
+    ax.axvline(mean_value, color="tab:grey", linestyle=":", linewidth=3, label=f"Mean")
+
+    # Setting the x limit
+    ax.set_xlim(0, 10)
+
+    # Adding title and labels
+    ax.set_xlabel(var_dict[var_key]["symbol"])
+    ax.set_ylabel("Frequency")
+    fig.legend(loc="upper right", bbox_to_anchor=(0.93, 0.9), fontsize="small")
+
+    return fig, ax
+
+
+plt.rcParams.update({"font.size": 30})
+fig_q_hist, _ = plot_hist(df=df_filt_q, var_key="tauexp_tau")
+if save:
+    save_figure(fig_q_hist, fig_dir, f"tau_hist", "png", 900)
+    save_figure(fig_q_hist, fig_dir, f"tau_hist", "pdf", 1200)
 # %%
 print(f"Global median q: {df_filt_q['tauexp_tau'].median()}")
 print(f"Global mean q: {df_filt_q['tauexp_tau'].mean()}")
@@ -735,4 +804,207 @@ def plot_map_counts(ax, df, coord_info, cmap, norm, title=""):
 if save:
     save_figure(fig_map_events, fig_dir, f"map_eventcounts", "png", 900)
 
+
+# %%
+import matplotlib.cm as cm
+
+
+def plot_boxplots(df, x_var, y_var, cmap_name):
+    plt.rcParams.update({"font.size": 12})
+    fig, ax = plt.subplots(figsize=(6, 4))
+
+    # Extract unique categories for the x variable
+    categories = df[x_var["column_name"]].dropna().unique()
+    categories = sorted(
+        categories, key=lambda x: x.left
+    )  # Sort intervals by their left edge
+
+    # Generate a colormap for the categories
+    cmap = cm.get_cmap(cmap_name, len(categories))
+    colors = [cmap(i) for i in range(len(categories))]
+    palette = dict(zip(categories, colors))  # Map categories to colors
+
+    sns.boxplot(
+        x=x_var["column_name"],
+        y=y_var["column_name"],
+        data=df,
+        palette=cmap_name,
+        # boxprops=dict(facecolor="lightgray"),
+        ax=ax,
+    )
+
+    ax.set_xticklabels(ax.get_xticklabels(), rotation=45)
+    ax.set_xlabel(f'{x_var["label"]} {x_var["unit"]}')
+    ax.set_ylabel(f'{y_var["label"]} {y_var["unit"]}')
+    ax.set_ylim(y_var["lim"][0] * 5, y_var["lim"][1] * 5)
+    fig.tight_layout()
+
+    return fig, ax
+
+
+# %% sand
+fig_box_sand, _ = plot_boxplots(
+    df_filt_q_and_tauexp,
+    var_dict["sand_bins"],
+    var_dict["diff_R2_tauexp"],
+    cmap_name=sand_cmap,
+)
+fig_box_sand.savefig(
+    os.path.join(fig_dir, f"box_diff_R2_tauexp_sand.png"),
+    dpi=600,
+    bbox_inches="tight",
+)
+
+# %% Aridity index
+# %% sand
+fig_box_ai, _ = plot_boxplots(
+    df_filt_q_and_tauexp,
+    var_dict["ai_bins"],
+    var_dict["diff_R2_tauexp"],
+    cmap_name=ai_cmap,
+)
+fig_box_ai.savefig(
+    os.path.join(fig_dir, f"box_diff_R2_tauexp_ai.png"),
+    dpi=600,
+    bbox_inches="tight",
+)
+# %% sand
+fig_box_sand, _ = plot_boxplots(
+    df_filt_q_and_exp,
+    var_dict["sand_bins"],
+    var_dict["diff_aic_q_exp"],
+    cmap_name=sand_cmap,
+)
+fig_box_sand.savefig(
+    os.path.join(fig_dir, f"box_diff_aic_exp_sand.png"),
+    dpi=600,
+    bbox_inches="tight",
+)
+
+# %% Aridity index
+fig_box_ai, _ = plot_boxplots(
+    df_filt_q_and_exp,
+    var_dict["ai_bins"],
+    var_dict["diff_aic_q_exp"],
+    cmap_name=ai_cmap,
+)
+fig_box_ai.savefig(
+    os.path.join(fig_dir, f"box_diff_aic_exp_ai.png"),
+    dpi=600,
+    bbox_inches="tight",
+)
+
+
+# %% Vegatation
+def wrap_at_space(text, max_width):
+    parts = text.split(" ")
+    wrapped_parts = [wrap(part, max_width) for part in parts]
+    return "\n".join([" ".join(wrapped_part) for wrapped_part in wrapped_parts])
+
+
+def plot_boxplots_categorical(df, x_var, y_var, categories, colors):
+    # Create the figure and axes
+    fig, ax = plt.subplots(figsize=(6, 4))
+
+    # Plot the boxplot with specified colors and increased alpha
+    sns.boxplot(
+        x=x_var["column_name"],
+        y=y_var["column_name"],
+        data=df,
+        # hue=x_var['column_name'],
+        legend=False,
+        order=categories,
+        palette=colors,
+        ax=ax,
+    )
+
+    for patch in ax.artists:
+        r, g, b, a = patch.get_facecolor()
+        patch.set_facecolor(mcolors.to_rgba((r, g, b), alpha=0.5))
+
+    # Optionally, adjust layout
+    plt.tight_layout()
+    ax.set_xlabel(f'{x_var["label"]}')
+    max_label_width = 20
+    ax.set_xticklabels(
+        [
+            wrap_at_space(label.get_text(), max_label_width)
+            for label in ax.get_xticklabels()
+        ]
+    )
+    plt.setp(ax.get_xticklabels(), rotation=45)
+    ax.set_ylabel(f'{y_var["label"]} {y_var["unit"]}')
+    # Show the plot
+    ax.set_ylim(y_var["lim"][0] * 5, y_var["lim"][1] * 5)
+    plt.tight_layout()
+    plt.show()
+
+    return fig, ax
+
+
+# %%
+fig_box_veg, _ = plot_boxplots_categorical(
+    df_filt_q_and_tauexp,
+    var_dict["veg_class"],
+    var_dict["diff_R2_tauexp"],
+    categories=vegetation_color_dict.keys(),
+    colors=list(vegetation_color_dict.values()),
+)
+fig_box_veg.savefig(
+    os.path.join(fig_dir, f"box_diff_R2_tauexp_veg.png"), dpi=600, bbox_inches="tight"
+)
+
+# %%
+fig_box_veg, _ = plot_boxplots_categorical(
+    df_filt_q_and_exp,
+    var_dict["veg_class"],
+    var_dict["diff_aic_q_exp"],
+    categories=vegetation_color_dict.keys(),
+    colors=list(vegetation_color_dict.values()),
+)
+fig_box_veg.savefig(
+    os.path.join(fig_dir, f"box_diff_aic_exp_veg.png"), dpi=600, bbox_inches="tight"
+)
+
+
+# def plot_box_ai_veg(df):
+#     plt.rcParams.update({"font.size": 26})  # Adjust the font size as needed
+
+#     fig, ax = plt.subplots(figsize=(20, 8))
+#     for i, category in enumerate(vegetation_color_dict.keys()):
+#         subset = df[df["name"] == category]
+#         sns.boxplot(
+#             x="name",
+#             y="AI",
+#             df=subset,
+#             color=vegetation_color_dict[category],
+#             ax=ax,
+#             linewidth=2,
+#         )
+
+#     # ax = sns.violinplot(x='abbreviation', y='q_q', df=filtered_df, order=vegetation_orders, palette=palette_dict) # boxprops=dict(facecolor='lightgray'),
+#     max_label_width = 20
+#     ax.set_xticklabels(
+#         [
+#             wrap_at_space(label.get_text(), max_label_width)
+#             for label in ax.get_xticklabels()
+#         ]
+#     )
+#     plt.setp(ax.get_xticklabels(), rotation=45)
+
+#     # ax.set_xticklabels([textwrap.fill(t.get_text(), 10) for t in ax.get_xticklabels()])
+#     ax.set_ylabel("Aridity index [MAP/MAE]")
+#     ax.set_xlabel("IGBP Landcover Class")
+#     ax.set_ylim(0, 2.0)
+#     ax.set_title("(a)", loc="left")
+#     plt.tight_layout()
+
+#     return fig, ax
+
+
+# fig_box_ai_veg, _ = plot_box_ai_veg(df_filt_q)
+# fig_box_ai_veg.savefig(
+#     os.path.join(fig_dir, f"sup_box_ai_veg.png"), dpi=1200, bbox_inches="tight"
+# )
+
 # %%