Add population read and diagnostic

JGCRI · Jul 27, 2023 · e77dc42 · e77dc42
1 parent 91202bd
commit e77dc42
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 11 deletions.
diff --git a/bed/diagnostics.py b/bed/diagnostics.py
@@ -3,9 +3,10 @@
 import pandas as pd
 import xarray as xr
 import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
 
 
-def diagnostics(fake_param: int = 1, data=None, name='figure'):
+def diagnostics(data=None, name='figure'):
     """Fake function to remove.
 
     :param fake_param:                          A fake integer
@@ -24,19 +25,53 @@ def diagnostics(fake_param: int = 1, data=None, name='figure'):
 
     if data != None:
         ds = data.temperature
-        ds.data_vars
-        ds.T2.XTIME
+        # ds.data_vars
+        # ds.T2.XTIME
         t2_1 = ds.T2.isel(Time=1)
         t2_2 = ds.T2.isel(Time=12)
         t2 = xr.concat([t2_1, t2_2], pd.Index([t2_1.XTIME.values.astype('str')[0:19], t2_2.XTIME.values.astype('str')[0:19]], name="T"))
-        t2.values.shape
-        t2.coords
-        plot_t2 = t2.plot(x="west_east", y="south_north", col="T")
-        plt.show()
+        # t2.values.shape
+        # t2.coords
+
+        t2.plot(x="west_east", y="south_north", col="T")
+        # plt.show()
+        # Save the plot into the diagnostics folder created by read_data()
+        plt.savefig(fname=os.path.join(data.dir_diagnostics, 'diagnostic_temperature.png'))
+        plt.close()
+
+    logging.info(f"Diagnostic plots saved to: {os.path.join(data.dir_diagnostics, 'diagnostic_temperature.png')}")
+    logging.info('Plotting diagnostics for temperature data complete.')
+
+    # Plot Population Data
+    logging.info('Plotting diagnostics for temperature data...')
+
+    if data != None:
+        ds = data.population
+
+        ticks = [0, 1000, 5000, 10000, 50000, 100000, 500000, 1000000, 5000000]
+        subplot_kws = dict(projection=ccrs.PlateCarree(),
+                           facecolor='white')
+
+        plt.figure(figsize=[10, 5])
+        p = ds.ssp2_2020.plot(x='lon', y='lat',
+                              vmin=0, vmax=5000000,
+                              size=4,
+                              aspect=ds.dims['lon']/ds.dims['lat'],
+                              cmap='viridis',
+                              subplot_kws=subplot_kws,
+                              transform=ccrs.PlateCarree(),
+                              add_labels=False,
+                              add_colorbar=False,
+                              levels=ticks)
+
+        cb = plt.colorbar(p, shrink=0.7, ticks=ticks)
+        cb.ax.tick_params(labelsize=14)
+        # plt.show()
         # Save the plot into the diagnostics folder created by read_data()
-        plt.savefig(fname=os.path.join(data.dir_diagnostics, name + '.png'))
+        plt.savefig(fname=os.path.join(data.dir_diagnostics, 'diagnostic_population.png'))
+        plt.close()
 
-    logging.info(f"Diagnostic plots saved to: {os.path.join(data.dir_diagnostics, name + '.png')}")
+    logging.info(f"Diagnostic plots saved to: {os.path.join(data.dir_diagnostics, 'diagnostic_population.png')}")
     logging.info('Plotting diagnostics for temperature data complete.')
 
 

diff --git a/bed/read_data.py b/bed/read_data.py
@@ -59,4 +59,9 @@ def __init__(self, config_file=''):
             else: # If user gives full path
                 self.temperature = xr.open_dataset(self.config['path_temperature_ncdf'])
 
+            if os.path.exists(os.path.abspath(os.path.join(self.dir_root, self.config['path_population_ncdf']))):
+                self.population = xr.open_dataset(os.path.abspath(os.path.join(self.dir_root, self.config['path_population_ncdf'])))
+            else: # If user gives full path
+                self.population = xr.open_dataset(self.config['path_population_ncdf'])
+
         logging.info('Class Data inside module read_data completed.')
diff --git a/extras/dev_tests.py b/extras/dev_tests.py
@@ -6,8 +6,6 @@
 # pytest
 import os
 import bed
-
-
 # a1 = bed.Bed(config_file= os.path.join(data_folder,"example_config.yml")) # Coming from Model.py
 
 # Get example data
@@ -20,6 +18,7 @@
 data = bed.Data(config_file=os.path.join(data_folder, "example_config.yml"))
 data.example_dataset
 data.temperature
+data.population
 
 # Run Diagnsotics
 bed.diagnostics(data=data)