Add plotting and saving functions for data visualization

- Implemented a comprehensive set of plotting functions in `plot_functions.py` for generating time series and map visualizations, including custom legends and nested grid figures.
- Introduced functions for saving figures as PNG and PDF in `save_functions.py`, with options for creating thumbnails and adjusting DPI settings.
- Enhanced the modularity of the plotting functions to accommodate various scenarios and covariates, improving the overall flexibility of the visualization toolkit.

Author: bcreiner

idd-forecast-mbp.Rproj

@@ -1,13 +1,13 @@
-Version: 1.0
-
-RestoreWorkspace: Default
-SaveWorkspace: Default
-AlwaysSaveHistory: Default
-
-EnableCodeIndexing: Yes
-UseSpacesForTab: Yes
-NumSpacesForTab: 2
-Encoding: UTF-8
-
-RnwWeave: Sweave
-LaTeX: pdfLaTeX
+Version: 1.0
+
+RestoreWorkspace: Default
+SaveWorkspace: Default
+AlwaysSaveHistory: Default
+
+EnableCodeIndexing: Yes
+UseSpacesForTab: Yes
+NumSpacesForTab: 2
+Encoding: UTF-8
+
+RnwWeave: Sweave
+LaTeX: pdfLaTeX

src/idd_forecast_mbp/03_modeling/DengueRegression_scam.pdf

@@ -5,7 +5,8 @@
     "smart_format",
     "pretty_bin_labels",
     "add_colorbar",
-    "add_legend_panel_colorbar"
+    "add_legend_panel_colorbar",
+    "clip_data_to_bins"
 ]
 
 import pandas as pd
@@ -31,12 +32,36 @@
 from idd_forecast_mbp.helper_functions import read_income_paths
 from idd_forecast_mbp.parquet_functions import read_parquet_with_integer_ids
 from idd_forecast_mbp.xarray_functions import read_netcdf_with_integer_ids, convert_to_xarray, write_netcdf
+from idd_forecast_mbp.color_functions import (
+    create_change_colormap, 
+    create_outcome_colormap, 
+    create_diverging_colors,
+    get_colors
+)
 
 
-def draw_legend_bins(ax, plot_dict):
-    map_type = plot_dict['map_type']
-    legend_dict = plot_dict['legend_dict']
-    bin_dict = plot_dict['bin_dict']
+def clip_data_to_bins(df, data_column, bins):
+    """Clip data values to ensure they fall within bin boundaries."""
+    df = df.copy()
+    min_bin, max_bin = bins[0], bins[-1]
+    
+    # Show clipping stats
+    original_min, original_max = df[data_column].min(), df[data_column].max()
+    clipped_count = ((df[data_column] < min_bin) | (df[data_column] > max_bin)).sum()
+    
+    if clipped_count > 0:
+        print(f"Clipping {clipped_count} values outside bins [{min_bin}, {max_bin}]")
+        print(f"Original range: [{original_min:.2f}, {original_max:.2f}]")
+    
+    # Clip the data
+    df[data_column] = df[data_column].clip(lower=min_bin, upper=max_bin)
+    
+    return df
+
+def draw_legend_bins(ax, map_plot_dict):
+    map_type = map_plot_dict['map_type']
+    legend_dict = map_plot_dict['legend_dict']
+    bin_dict = map_plot_dict['bin_dict']
     legend_panel = legend_dict['legend_panel']
     legend_bin_spacing = legend_panel['legend_bin_spacing']
     legend_margin = legend_panel['legend_margin']
@@ -49,7 +74,7 @@ def draw_legend_bins(ax, plot_dict):
     bin_height = bin_top - bin_bottom
     bin_label_y = bin_bottom - bin_label_gap
 
-    category_labels = plot_dict['bin_dict']['bin_labels']
+    category_labels = map_plot_dict['bin_dict']['bin_labels']
 
     bin_dict['n_bins'] = len(category_labels)
     n_bins = bin_dict['n_bins']
@@ -72,44 +97,41 @@ def draw_legend_bins(ax, plot_dict):
         rect = Rectangle((bin_left[i], bin_bottom), bin_width, bin_height, facecolor=bin_colors[i], edgecolor='black', linewidth=0.5)
         ax.add_patch(rect)
         ax.text(bin_left[i] + bin_width / 2, bin_label_y, category_labels[i], ha='center', va='top',
-                        fontsize=plot_dict['fontsizes']['legend_label_fontsize'])
+                        fontsize=map_plot_dict['fontsizes']['legend_label_fontsize'])
 
-def add_legend(fig, ax, plot_dict):
-    if plot_dict['have_legend_panel']:
-        if plot_dict['legend_dict']['use_colorbar']:
-            add_legend_panel_colorbar(fig, ax, plot_dict)
+def add_legend(fig, ax, map_plot_dict):
+    if map_plot_dict['have_legend_panel']:
+        if map_plot_dict['legend_dict']['use_colorbar']:
+            add_legend_panel_colorbar(fig, ax, map_plot_dict)
         else:
-            draw_legend_bins(ax, plot_dict)
+            draw_legend_bins(ax, map_plot_dict)
     else:
-        add_colorbar(fig, ax, plot_dict)
+        add_colorbar(fig, ax, map_plot_dict)
 
-def get_bin_info(plot_dict, valid_data, map_data_masked):
-    map_type = plot_dict['map_type']
-    bin_dict = plot_dict['bin_dict']
-    # Categorization logic (same as before)
-    bins = bin_dict['bins']
-    n_bins = bin_dict['n_bins']
-    base_colormap = plt.colormaps[plot_dict['colors_dict']['base_cmap']]
-    if bins[0] == 0:
-        white_rgba = to_rgba('white')
-        bin_colors = np.vstack([white_rgba, base_colormap(np.linspace(0.1, 0.9, n_bins - 1))])
+def get_bin_info(map_plot_dict, plot_data):
+    map_type = map_plot_dict['map_type']
 
+    if map_type == 'outcome':
+        # Outcome map
+        create_outcome_colormap(map_plot_dict)            
     else:
-        bin_colors = base_colormap(np.linspace(0.1, 0.9, n_bins))
+        create_change_colormap(map_plot_dict)
 
-    cmap = ListedColormap(bin_colors)
-    categorical_data = np.full_like(map_data_masked, np.nan)
+    map_plot_dict['bin_dict']['bin_labels'] = pretty_bin_labels(map_plot_dict)
+    bin_dict = map_plot_dict['bin_dict']
+    bins = bin_dict['bins']
+    n_bins = bin_dict['n_bins']
+    plot_data_values = plot_data.values if hasattr(plot_data, 'values') else plot_data
+    categorical_data = np.full_like(plot_data_values, np.nan)
     for i in range(n_bins):
-        if i == 0:
-            mask = map_data_masked <= bins[i+1]
+        if i == 0.0:
+            mask = plot_data_values <= bins[i+1]
         elif i == n_bins - 1:
-            mask = map_data_masked > bins[i]
+            mask = plot_data_values > bins[i]
         else:
-            mask = (map_data_masked > bins[i]) & (map_data_masked <= bins[i+1])
+            mask = (plot_data_values > bins[i]) & (plot_data_values <= bins[i+1])
         categorical_data[mask] = i
 
-    bin_dict['bin_colors'] = bin_colors
-    bin_dict['cmap'] = cmap
     bin_dict['categorical_data'] = categorical_data
 
 def smart_format(val):
@@ -120,14 +142,14 @@ def smart_format(val):
         s = f"{val:,.2f}".rstrip('0').rstrip('.')
         return s
 
-def pretty_bin_labels(plot_dict):
-    bins = plot_dict['bin_dict']['bins']
-    le = plot_dict['bin_dict']['le']
-    ge = plot_dict['bin_dict']['ge']
-    zero_bin = plot_dict['bin_dict']['zero_bin']
-    prefix_units = plot_dict['bin_dict']['prefix_units']
-    suffix_units = plot_dict['bin_dict']['suffix_units']
-    abbreviate_labels = plot_dict['bin_dict']['abbreviate_labels']
+def pretty_bin_labels(map_plot_dict):
+    bins = map_plot_dict['bin_dict']['bins']
+    le = map_plot_dict['bin_dict']['le']
+    ge = map_plot_dict['bin_dict']['ge']
+    zero_bin = map_plot_dict['bin_dict']['zero_bin']
+    prefix_units = map_plot_dict['bin_dict']['prefix_units']
+    suffix_units = map_plot_dict['bin_dict']['suffix_units']
+    abbreviate_labels = map_plot_dict['bin_dict']['abbreviate_labels']
     # bins: array-like of bin edges
     # fmt: format for numbers (default: 2 significant digits)
 
@@ -188,12 +210,12 @@ def pretty_bin_labels(plot_dict):
 
     return labels
 
-def add_colorbar(fig, ax, plot_dict):
-    figure_dict = plot_dict['figure_dict']
-    legend_dict = plot_dict['legend_dict']
-    bins = plot_dict['bin_dict']['bins']
+def add_colorbar(fig, ax, map_plot_dict):
+    figure_dict = map_plot_dict['figure_dict']
+    legend_dict = map_plot_dict['legend_dict']
+    bins = map_plot_dict['bin_dict']['bins']
     bin_centers = [(bins[i] + bins[i+1]) / 2 for i in range(len(bins)-1)]
-    sm = plt.cm.ScalarMappable(cmap=plot_dict['bin_dict']['cmap'], norm=plot_dict['bin_dict']['norm'])
+    sm = plt.cm.ScalarMappable(cmap=map_plot_dict['bin_dict']['cmap'], norm=map_plot_dict['bin_dict']['norm'])
     sm.set_array([])
     cbar = fig.colorbar(sm, ax=ax, orientation='horizontal', 
                         shrink=legend_dict['color_bar_dict']['shrink'], 
@@ -204,22 +226,22 @@ def add_colorbar(fig, ax, plot_dict):
     cbar.set_ticklabels(figure_dict['bin_labels'], fontsize=figure_dict['tick_font_size'])
     cbar.set_label(figure_dict['colorbar_label'], fontsize=figure_dict['colorbar_title_font_size'])
 
-def add_legend_panel_colorbar(fig, ax_legend, plot_dict):
+def add_legend_panel_colorbar(fig, ax_legend, map_plot_dict):
     """
     Fixed version that uses the calculated parameters properly
     """
     ax_legend.clear()
     ax_legend.axis('off')
 
     # Get colorbar data
-    bin_dict = plot_dict['bin_dict']
+    bin_dict = map_plot_dict['bin_dict']
     cmap = bin_dict['cmap'] 
     norm = bin_dict['norm']
     bins = bin_dict.get('bins', None)
     bin_labels = bin_dict.get('bin_labels', None)
-    colorbar_label = plot_dict['full_outcome_label']
+    colorbar_label = map_plot_dict['full_outcome_label']
 
-    color_bar_dict = plot_dict['legend_dict']['color_bar_dict']
+    color_bar_dict = map_plot_dict['legend_dict']['color_bar_dict']
 
     # Create ScalarMappable
     sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
@@ -242,7 +264,7 @@ def add_legend_panel_colorbar(fig, ax_legend, plot_dict):
 
     cbar_ax = fig.add_axes([cbar_left, cbar_bottom, cbar_width, cbar_height])
 
-    extend_option = 'max' if plot_dict.get('extend_colorbar', False) else 'neither'
+    extend_option = 'max' if map_plot_dict.get('extend_colorbar', False) else 'neither'
     cbar = fig.colorbar(sm, cax=cbar_ax, orientation='horizontal', extend=extend_option)
 
     if bins is not None and bin_labels is not None:

src/idd_forecast_mbp/03_modeling/Rplot.pdf

@@ -13,6 +13,7 @@
 FIGURES_PATH = MODEL_ROOT / "07-figures"
 GBD_DATA_PATH = f"{RAW_DATA_PATH}/gbd"
 LSAE_INPUT_PATH = PROCESSED_DATA_PATH / "lsae_1209"
+MANUSCRIPT_PATH = MODEL_ROOT / "08-manuscript_material"
 
 FHS_RESULTS_PATH = '/mnt/share/forecasting/data/9/future'
 
@@ -262,21 +263,21 @@
 
 ssp_scenario_map = {
     "ssp126": {
-        "name": "RCP 2.6",
+        "name": "RCP2.6",
         "rcp_scenario": 2.6,
         "color": "#046C9A",
         "dhs_scenario": 66,
         "dhs_vbd_scenario": 75
     },
     "ssp245": {
-        "name": "RCP 4.5",
+        "name": "RCP4.5",
         "rcp_scenario": 4.5,
         "color": "#E58601",
         "dhs_scenario": 0,
         "dhs_vbd_scenario": 0
     },
     "ssp585": {
-        "name": "RCP 8.5",
+        "name": "RCP8.5",
         "rcp_scenario": 8.5,
         "color": "#A42820",
         "dhs_scenario": 54,
@@ -286,19 +287,19 @@
 
 ssp_scenarios = {
     "ssp126": {
-        "name": "RCP 2.6",
+        "name": "RCP2.6",
         "rcp_scenario": 2.6,
         "color": "#046C9A",
         "dhs_scenario": 66
     },
     "ssp245": {
-        "name": "RCP 4.5",
+        "name": "RCP4.5",
         "rcp_scenario": 4.5,
         "color": "#E58601",
         "dhs_scenario": 0
     },
     "ssp585": {
-        "name": "RCP 8.5",
+        "name": "RCP8.5",
         "rcp_scenario": 8.5,
         "color": "#A42820",
         "dhs_scenario": 54
@@ -367,3 +368,32 @@
         }        
     }
 }
+
+
+covariate_map = {
+    'suitability' : {
+        'var': 'suitability',
+        'ylabel': 'Suitability (days)',
+        'title' : 'Temperature Suitability'
+    },
+    'gdppc_mean': {
+        'var': 'gdppc_mean',
+        'ylabel': 'GDP per Capita',
+        'title' : 'GDP per Capita (in 2020 USD)'
+    },
+    'dah_pc': {
+        'var': 'dah_pc',
+        'ylabel': 'DAH per Capita',
+        'title' : 'DAH per Capita (in 2020 USD)'
+    },
+    'flooding_pc': {
+        'var': 'flooding_pc',
+        'ylabel': 'Flood days per capita',
+        'title' : 'flood days per capita'
+    },
+    'urbanization': {
+        'var': 'urbanization',
+        'ylabel': 'Urbanization (%)',
+        'title' : 'Urbanization'
+    }
+}