Adding production WRF pipeline

wildfire_occurrence/model/analysis/wrf_analysis.py

@@ -0,0 +1,484 @@
+import os
+import xwrf
+import netCDF4
+import logging
+import xarray as xr
+from glob import glob
+from rasterio.crs import CRS
+from wrf import getvar, interplevel
+
+__all__ = ["xwrf"]
+
+
+class WRFAnalysis(object):
+
+    def __init__(self, input_filename_regex):
+
+        # wrf input_filename_regex
+        self.wrf_filename_regex = input_filename_regex
+
+        # get list of wrf input filenames
+        self.wrf_filenames = sorted(glob(self.wrf_filename_regex))
+
+        # get dataset into xr using xwrf format
+        self.wrf_dataset = xr.open_mfdataset(
+            self.wrf_filenames,
+            engine="netcdf4",
+            parallel=True,
+            concat_dim="Time",
+            combine="nested",
+            chunks=None,
+            decode_times=False,
+            decode_coords="all",
+        ).xwrf.postprocess(drop_diagnostic_variable_components=False)
+
+        # get CRS from xwrf consolidation
+        self.crs = CRS.from_string(
+            str(self.wrf_dataset['wrf_projection'].values))
+
+        # assign crs to override crs=None and to be compliant with rioxarray
+        self.wrf_dataset.rio.write_crs(self.crs, inplace=True)
+
+        # get netCDF objects compatible with wrf-python
+        # this is needed since wrf-python does not accept xwrf input
+        self.wrf_python_dataset = [
+            netCDF4.Dataset(f) for f in self.wrf_filenames]
+
+        # get list of variables, and remove grid_mapping attribute
+        # this is needed to be compliant with rioxarray
+        vars_list = list(self.wrf_dataset.data_vars)
+        for var in vars_list:
+            if 'grid_mapping' in self.wrf_dataset[var].attrs:
+                del self.wrf_dataset[var].attrs['grid_mapping']
+
+    def compute_all_and_write(
+                self,
+                timeidx=0,
+                output_variables=["LANDMASK"],
+                output_filename=None
+            ):
+        """
+        We use this function to compute lightning specific variables
+        and to store them in the same dataset.
+        """
+
+        """
+        ['Times', 'LU_INDEX', 'ZS', 'DZS', 'VAR_SSO', 'BATHYMETRY_FLAG',
+        'U', 'V', 'W','PH', 'PHB', 'T', 'THM', 'HFX_FORCE', 'LH_FORCE',
+        'TSK_FORCE', 'HFX_FORCE_TEND','LH_FORCE_TEND', 'TSK_FORCE_TEND',
+        'MU', 'MUB', 'NEST_POS', 'P', 'PB', 'FNM', 'FNP','RDNW', 'RDN',
+        'DNW', 'DN', 'CFN', 'CFN1', 'THIS_IS_AN_IDEAL_RUN', 'P_HYD', 'Q2',
+        'T2', 'TH2', 'PSFC', 'U10', 'V10', 'LPI', 'RDX', 'RDY', 'AREA2D',
+        'DX2D', 'RESM','ZETATOP', 'CF1', 'CF2', 'CF3', 'ITIMESTEP', 'QVAPOR',
+        'QCLOUD', 'QRAIN', 'QICE','QSNOW', 'QGRAUP', 'QNICE', 'QNRAIN',
+        'SHDMAX','SHDMIN', 'SNOALB', 'TSLB', 'SMOIS','SH2O', 'SMCREL',
+        'SEAICE', 'XICEM', 'SFROFF', 'UDROFF', 'IVGTYP', 'ISLTYP', 'VEGFRA',
+        'GRDFLX', 'ACGRDFLX', 'ACSNOM', 'SNOW', 'SNOWH', 'CANWAT', 'SSTSK',
+        'WATER_DEPTH', 'COSZEN', 'LAI', 'U10E', 'V10E', 'DTAUX3D', 'DTAUY3D',
+        'DUSFCG', 'DVSFCG', 'VAR', 'CON', 'OA1', 'OA2', 'OA3', 'OA4', 'OL1',
+        'OL2', 'OL3', 'OL4', 'TKE_PBL', 'EL_PBL', 'O3_GFS_DU', 'MAPFAC_M',
+        'MAPFAC_U', 'MAPFAC_V', 'MAPFAC_MX', 'MAPFAC_MY', 'MAPFAC_UX',
+        'MAPFAC_UY', 'MAPFAC_VX', 'MF_VX_INV', 'MAPFAC_VY', 'F', 'E',
+        'SINALPHA','COSALPHA', 'HGT', 'TSK', 'P_TOP', 'GOT_VAR_SSO', 'T00',
+        'P00', 'TLP','TISO', 'TLP_STRAT', 'P_STRAT', 'MAX_MSFTX', 'MAX_MSFTY',
+        'RAINC','RAINSH', 'RAINNC', 'SNOWNC', 'GRAUPELNC', 'HAILNC',
+        'REFL_10CM','CLDFRA','SWDOWN', 'GLW', 'SWNORM', 'ACSWUPT', 'ACSWUPTC',
+        'ACSWDNT', 'ACSWDNTC','ACSWUPB', 'ACSWUPBC', 'ACSWDNB', 'ACSWDNBC',
+        'ACLWUPT', 'ACLWUPTC', 'ACLWDNT','ACLWDNTC', 'ACLWUPB', 'ACLWUPBC',
+        'ACLWDNB', 'ACLWDNBC', 'SWUPT', 'SWUPTC', 'SWDNT', 'SWDNTC', 'SWUPB',
+        'SWUPBC', 'SWDNB', 'SWDNBC', 'LWUPT', 'LWUPTC', 'LWDNT', 'LWDNTC',
+        'LWUPB', 'LWUPBC', 'LWDNB', 'LWDNBC', 'OLR', 'ALBEDO', 'ALBBCK',
+        'EMISS', 'NOAHRES', 'TMN', 'XLAND', 'UST', 'PBLH', 'HFX', 'QFX',
+        'LH', 'ACHFX', 'ACLHF', 'SNOWC', 'SR', 'SAVE_TOPO_FROM_REAL',
+        'REFD_MAX', 'ISEEDARR_SPPT', 'ISEEDARR_SKEBS', 'ISEEDARR_RAND_PERTURB',
+        'ISEEDARRAY_SPP_CONV', 'ISEEDARRAY_SPP_PBL', 'ISEEDARRAY_SPP_LSM',
+        'C1H', 'C2H', 'C1F', 'C2F', 'C3H', 'C4H', 'C3F', 'C4F', 'PCB', 'PC',
+        'LANDMASK', 'LAKEMASK', 'SST', 'SST_INPUT', 'air_potential_temperature'
+        'air_pressure', 'geopotential', 'geopotential_height', 'wind_east',
+        'wind_north']
+        """
+
+        # create a copy of the dataset with a single time step
+        wrf_dataset_single_time = self.wrf_dataset.isel(Time=timeidx)
+
+        # compute LPI - LPI is already computed by xwrf
+
+        # compute Helicity
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            Helicity=self.compute_var('helicity', timeidx))
+
+        # compute LCL
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            LCL=self.compute_var('lcl', timeidx))
+
+        # compute PW
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            PW=self.compute_var('pw', timeidx))
+
+        # compute SLP
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            SLP=self.compute_var('slp', timeidx))
+
+        # compute GPZ levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            GPZ500=self.compute_gpz(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            GPZ700=self.compute_gpz(700, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            GPZ750=self.compute_gpz(750, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            GPZ850=self.compute_gpz(850, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            GPZ1000=self.compute_gpz(1000, timeidx))
+
+        # compute DZ levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            DZ500_1000=self.compute_dz(
+                wrf_dataset_single_time['GPZ500'],
+                wrf_dataset_single_time['GPZ1000']
+            )
+        )
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            DZ850_1000=self.compute_dz(
+                wrf_dataset_single_time['GPZ850'],
+                wrf_dataset_single_time['GPZ1000']
+            )
+        )
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            DZ700_850=self.compute_dz(
+                wrf_dataset_single_time['GPZ700'],
+                wrf_dataset_single_time['GPZ850']
+            )
+        )
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            DZ700_1000=self.compute_dz(
+                wrf_dataset_single_time['GPZ700'],
+                wrf_dataset_single_time['GPZ1000']
+            )
+        )
+
+        # compute RH2
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RH2=self.compute_var('rh2', timeidx))
+
+        # compute RH levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RH500=self.compute_rh(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RH700=self.compute_rh(700, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RH800=self.compute_rh(800, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RH850=self.compute_rh(850, timeidx))
+
+        # compute T2 - T2 is already computed by xwrf
+
+        # compute Td2
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TD2=self.compute_var('td2', timeidx))
+
+        # compute TD levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TD500=self.compute_td(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TD700=self.compute_td(700, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TD850=self.compute_td(850, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TD1000=self.compute_td(1000, timeidx))
+
+        # compute TC levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TC500=self.compute_tc(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TC700=self.compute_tc(700, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TC850=self.compute_tc(850, timeidx))
+
+        # compute TP levels, double-check equation for this one
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TP500=self.compute_tp(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TP850=self.compute_tp(850, timeidx))
+
+        # compute SHOW
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            SHOW=self.compute_show(
+                wrf_dataset_single_time['TC500'],
+                wrf_dataset_single_time['TP850']
+            )
+        )
+
+        # compute TT
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            TT=self.compute_tt(
+                wrf_dataset_single_time['TC850'],
+                wrf_dataset_single_time['TD850'],
+                wrf_dataset_single_time['TC500']
+            )
+        )
+
+        # compute Rain
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            RAINTotal=self.compute_rain(
+                wrf_dataset_single_time['RAINNC'],
+                wrf_dataset_single_time['RAINC']
+            )
+        )
+
+        # compute W levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            W500=self.compute_w(500, timeidx))
+
+        # compute WA levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            WA500=self.compute_wa(500, timeidx))
+
+        # compute cloud frac levels
+        cloud_frac_variables = self.compute_cloudfrac(timeidx)
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            CFLow=cloud_frac_variables[0])
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            CFMed=cloud_frac_variables[1])
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            CFHigh=cloud_frac_variables[2])
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            CFTotal=cloud_frac_variables[3])
+
+        # compute T levels
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            T500=self.compute_t(500, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            T750=self.compute_t(750, timeidx))
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            T850=self.compute_t(850, timeidx))
+
+        # PLI IS INCORRECT, WE NEED TO FIX THIS
+        # compute PLI, might be incorrectly computed, need to double check
+        # https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.lifted_index.html
+        wrf_dataset_single_time = wrf_dataset_single_time.assign(
+            PLI=self.compute_pli(
+                wrf_dataset_single_time['T500'],
+                wrf_dataset_single_time['TP500']
+            )
+        )
+
+        wrf_dataset_single_time[output_variables].rio.write_crs(
+            self.crs, inplace=True)
+
+        wrf_dataset_single_time = wrf_dataset_single_time[
+            output_variables].to_array()
+        wrf_dataset_single_time.rio.write_nodata(-10001, inplace=True)
+        wrf_dataset_single_time.attrs['long_name'] = output_variables
+
+        if output_filename is not None:
+            wrf_dataset_single_time.rio.to_raster(
+                output_filename,
+                BIGTIFF='IF_SAFER',
+                compress='LZW',
+                driver='GTiff',
+                dtype='float32',
+                recalc_transform=False
+            )
+
+        return wrf_dataset_single_time
+
+    def compute_var(self, var_name: str, timeidx: int = 0):
+        var_output = getvar(
+            self.wrf_python_dataset, var_name, timeidx=timeidx)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_gpz(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        gpz = getvar(
+            self.wrf_python_dataset, 'geopt', timeidx=timeidx) / 9.81
+        var_output = interplevel(gpz, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_rh(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        rh = getvar(
+            self.wrf_python_dataset, 'rh', timeidx=timeidx)
+        var_output = interplevel(rh, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_td(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        td = getvar(
+            self.wrf_python_dataset, 'td', timeidx=timeidx)
+        var_output = interplevel(td, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_t(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        t = getvar(
+            self.wrf_python_dataset, 'T', timeidx=timeidx)
+        var_output = interplevel(t, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_tc(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        tc = getvar(
+            self.wrf_python_dataset, 'tc', timeidx=timeidx)
+        var_output = interplevel(tc, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_tp(self, pressure_level: int = 500, timeidx: int = 0):
+        # consider removing this extract calculation, try to find it
+        # if its already computed
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        tc = getvar(
+            self.wrf_python_dataset, 'tc', timeidx=timeidx)
+        tc_interpolated = interplevel(tc, pressure, pressure_level)
+        var_output = \
+            (tc_interpolated + 273.15) * \
+            ((500. / pressure_level)**0.286) - 273.15
+        # (tc_850 + 273.15)*((500/850)^0.286) - 273.15
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_dz(self, gpz1, gpz2):
+        return gpz1 - gpz2
+
+    def compute_show(self, tc_500, tp_850):
+        return tc_500 - tp_850
+
+    def compute_tt(self, tc_850, td_850, tc_500):
+        return tc_850 + td_850 - 2 * tc_500
+
+    def compute_rain(self, rain_exp, rain_con):
+        return rain_exp + rain_con
+
+    def compute_w(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        w = getvar(
+            self.wrf_python_dataset, 'W', timeidx=timeidx)
+        w = w[:pressure.shape[0], :, :]
+        var_output = interplevel(w, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_wa(self, pressure_level: int = 500, timeidx: int = 0):
+        pressure = getvar(
+            self.wrf_python_dataset, 'pressure', timeidx=timeidx)
+        wa = getvar(
+            self.wrf_python_dataset, 'wa', timeidx=timeidx)
+        var_output = interplevel(wa, pressure, pressure_level)
+        var_output = var_output.rio.write_crs(self.crs, inplace=True)
+        return var_output.rename({'south_north': 'y', 'west_east': 'x'})
+
+    def compute_cloudfrac(self, timeidx: int = 0):
+        cloudfrac = getvar(
+            self.wrf_python_dataset, "cloudfrac", timeidx=timeidx)
+        cloudfrac = cloudfrac.rename({'south_north': 'y', 'west_east': 'x'})
+        cloudfrac = cloudfrac.rio.write_crs(self.crs, inplace=True)
+        low_cloudfrac = cloudfrac[0, :, :]
+        mid_cloudfrac = cloudfrac[1, :, :]
+        high_cloudfrac = cloudfrac[2, :, :]
+        total_cloudfrac = (low_cloudfrac + mid_cloudfrac + high_cloudfrac) / 3
+        return low_cloudfrac, mid_cloudfrac, high_cloudfrac, total_cloudfrac
+
+    def compute_pli(self, t_500, tp_500):
+        return t_500 - tp_500
+
+
+# -----------------------------------------------------------------------------
+# Invoke the main
+# -----------------------------------------------------------------------------
+if __name__ == "__main__":
+
+    filename_regex = \
+        '/explore/nobackup/projects/ilab/projects/LobodaTFO/data/WRF_Data/' + \
+        'WRF_Simulations/2022-07-03/wrfout_d02*'
+
+    filename_regex = \
+        '/explore/nobackup/projects/ilab/projects/LobodaTFO/operations/' + \
+        '2023-06-24_2023-07-04/output/wrfout_d02*'
+    data_filenames = glob(filename_regex)
+
+    for filename in data_filenames:
+
+        # /explore/nobackup/projects/ilab/projects/LobodaTFO/data/WRF_Data/WRF_Simulations/*/wrfout_d02*
+        # create WRFAnalysis object, stores wrf_dataset
+        wrf_analysis = WRFAnalysis(filename)
+
+        # output variables for the lightning model
+        # these are the variables of importance between both 24 and 48 models
+        # output_variables = [
+        #    'PLI', 'GPZ500', 'GPZ700', 'TD500', 'CFTotal',
+        #    'RH500', 'SLP', 'W500', 'RH700', 'CFLow', 'TD2',
+        #    'TT', 'Helicity', 'GPZ850', 'SHOW', 'LCL',
+        #    'RH2', 'T850', 'RH850', 'Rain', 'T2', 'DZ700_850',
+        #    'RH800', 'T500', 'PW', 'T750'
+        # ] # BT missing
+
+        output_variables = [
+            'CFTotal', 'CFLow', 'CFMed', 'CFHigh',
+            'DZ700_850',
+            'GPZ500', 'GPZ700', 'GPZ750', 'GPZ850',
+            'Helicity',
+            'LCL',
+            'PLI', 'PW',
+            'RAINTotal',
+            'RH2', 'RH500', 'RH700', 'RH800', 'RH850',
+            'SHOW',
+            'SLP',
+            'TD2', 'TD500',
+            'TT', 'T2', 'T500', 'T750', 'T850',
+            'W500', 'WA500'
+        ]  # BT missing
+
+        # looks good - Helicity, SLP, 'GPZ500', 'TD500', 'RH500',
+        # 'TD2', 'LCL', 'PW', 'RH2', 'RAINTotal'
+        # TT, TC500
+        # wrong - PLI
+        # maybe
+        #   CFTotoal (wrong, local is 0, ours is higher)
+        #   CFlow (wrong, local is 0, ours is higher)
+        #   CFMed  (wrong, local is 0, ours is higher)
+        #   CFHigh (wrong, local is 0, ours is higher)
+        #   'GPZ750' no-data problems, kind of similar
+        #   'GPZ700' no local data to compare to
+        #   'RH700' no local data to compare to
+        #   'GPZ850' no local data to compare to
+        #   'SHOW' wrong because of no data
+        #   'T500' no local data to compare to
+        #    'RH800' wrong because of no data
+        #    'T2' numbers look far away - our T2 is in Kelvin
+        #    'RH850' no local data to compare to
+        #    'T850' no local data to compare to
+        #   'W500'
+        #   'WA500' -looks good
+        #  'DZ700_850' - looks good
+        # 'BT' missing
+
+        output_dir = 'output'  # os.path.dirname(os.path.abspath(filename))
+
+        # TODO: make this for loop parallel later
+        for t_idx, delta_time in enumerate(wrf_analysis.wrf_dataset.Times.values):
+
+            logging.info(f'Processing t_idx: {t_idx}, timestamp: {delta_time}')
+            output_filename = os.path.join(
+                output_dir,
+                f"d02_{delta_time.astype(str).replace(':', '-')}.tif")
+
+            if not os.path.isfile(output_filename):
+
+                wrf_analysis.compute_all_and_write(
+                    timeidx=t_idx,
+                    output_variables=output_variables,
+                    output_filename=output_filename
+                )

wildfire_occurrence/model/config.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional
 from dataclasses import dataclass, field
 
 
@@ -34,3 +34,25 @@ class Config:
     wrf_config: Optional[dict] = field(
         default_factory=lambda: {
             'interval_seconds': 10800, 'num_metgrid_levels': 27})
+
+    # Output filename from WRF to extract variables from
+    wrf_output_filename: Optional[str] = 'wrfout_d02_*_00:00:00'
+
+    # List for posprocessing of variables
+    wrf_output_variables: Optional[List[str]] = field(
+        default_factory=lambda: [
+            'CFTotal', 'CFLow', 'CFMed', 'CFHigh',
+            'DZ700_850',
+            'GPZ500', 'GPZ700', 'GPZ750', 'GPZ850',
+            'Helicity',
+            'LCL',
+            'PLI', 'PW',
+            'RAINTotal',
+            'RH2', 'RH500', 'RH700', 'RH800', 'RH850',
+            'SHOW',
+            'SLP',
+            'TD2', 'TD500',
+            'TT', 'T2', 'T500', 'T750', 'T850',
+            'W500', 'WA500'
+        ]
+    )

wildfire_occurrence/model/data_download/ncep_fnl.py

@@ -9,7 +9,15 @@ from datetime import date
 from typing import List, Literal
 from multiprocessing import Pool, cpu_count
 
-__data_source__ = 'https://rda.ucar.edu/datasets/ds083.2'
+__past_data_source__ = 'https://rda.ucar.edu/datasets/ds083.2'
+__future_data_source__ = 'https://rda.ucar.edu/datasets/ds084.1'
+__projection_data_source__ = 'https://rda.ucar.edu/datasets/ds316-1'
+
+DATASET_URL = {
+    'prod': 'https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod',
+    'ds084.1': 'https://data.rda.ucar.edu/ds084.1',  # future-short
+    'ds083.2': 'https://stratus.rda.ucar.edu/ds083.2',  # past
+}
 
 
 class NCEP_FNL(object):
@@ -20,6 +28,8 @@ class NCEP_FNL(object):
                 start_date: str = date.today(),
                 end_date: str = date.today(),
                 hour_intervals: List = ['00', '06', '12', '18'],
+                dataset: str = None,
+                resolution: str = '1p00',  # 1p00, 0p50, 0p25
                 n_procs: int = cpu_count()
             ):
 
@@ -37,22 +47,37 @@ class NCEP_FNL(object):
         if isinstance(end_date, str):
             self.end_date = datetime.datetime.strptime(
                 end_date, '%Y-%m-%d').date()
+        elif isinstance(end_date, datetime.datetime):
+            self.end_date = end_date.date()
         else:
             self.end_date = end_date
 
         # define hour intervals
         self.hour_intervals = hour_intervals
 
-        # TODO: IF WE ARE DOWNLOADING INTO THE FUTURE
-        # THEN WE NEED TO SPECIFY THIS IS FROM THE OTHER
-        # DATASET AND NOT FROM THE CURRENT GFS
+        # define resolution to download
+        self.resolution = resolution
+
+        # dataset to download, select based on past vs future
+        if dataset is not None:
+            # this means the user specified the dataset manually
+            self.dataset = dataset
+        else:
+            # automatically select future dataset
+            if self.end_date > datetime.datetime.now().date():
+
+                # specify NOAA production GFS dataset
+                self.dataset = 'prod'
+
+                # modify the hour interval to match end date
+                # 384 is the longest time interval produced by NOAA
+                self.hour_intervals = [
+                    f'{interval:03}' for interval in range(0, 385, 3)]
+
+            # automatically select past archive dataset
+            else:
+                self.dataset = 'ds083.2'
 
-        # make sure we do not download data into the future
-        # if self.end_date > datetime.datetime.now():
-        #    self.end_date = datetime.datetime.now()
-        #    self.hour_intervals = [
-        #        d for d in self.hour_intervals
-        #        if int(d) <= self.end_date.hour - 6]
         logging.info(
             f'Downloading data from {self.start_date} to {self.end_date}')
 
@@ -84,18 +109,23 @@ class NCEP_FNL(object):
         }
 
         # define data url
-        self.data_url = 'https://rda.ucar.edu'
+        self.set_data_url(self.dataset)
 
+        # setup grib format
         if self.start_date.year < 2008:
             self.grib_format = 'grib1'
         else:
             self.grib_format = 'grib2'
 
-        self.dataset_path = f'/data/OS/ds083.2/{self.grib_format}'
-
         # nnumber of processors to use
         self.n_procs = n_procs
 
+    def set_data_url(self, dataset: str):
+        try:
+            self.data_url = DATASET_URL[dataset]
+        except KeyError:
+            sys.exit(f'{dataset} dataset not supported')
+
     def _authenticate(self, action: Literal["auth", "cleanup"] = "auth"):
 
         if action == "cleanup":
@@ -167,18 +197,34 @@ class NCEP_FNL(object):
         return
 
     def _get_filenames(self):
+        # list to store filenames
         filenames_list = []
-        daterange = pd.date_range(self.start_date, self.end_date)
-        for single_date in daterange:
-            year = single_date.strftime("%Y")
+
+        # dataset path for ds083.2, past archive data
+        if self.dataset == 'ds083.2':
+            daterange = pd.date_range(self.start_date, self.end_date)
+            for single_date in daterange:
+                year = single_date.strftime("%Y")
+                for hour in self.hour_intervals:
+                    filename = os.path.join(
+                        f'/{self.grib_format}/',
+                        f'{year}/{single_date.strftime("%Y.%m")}',
+                        f'fnl_{single_date.strftime("%Y%m%d")}_' +
+                        f'{hour}_00.{self.grib_format}'
+                    )
+                    filenames_list.append(filename)
+
+        # dataset path for production
+        # https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod/gfs.20230623/00/atmos/gfs.t00z.pgrb2.1p00.f000
+        elif self.dataset == 'prod':
             for hour in self.hour_intervals:
                 filename = os.path.join(
-                    self.dataset_path,
-                    f'{year}/{single_date.strftime("%Y.%m")}',
-                    f'fnl_{single_date.strftime("%Y%m%d")}_' +
-                    f'{hour}_00.{self.grib_format}'
+                    f'/gfs.{self.start_date.strftime("%Y%m%d")}',
+                    '00/atmos',
+                    f'gfs.t00z.pgrb2.{self.resolution}.f{hour}'
                 )
                 filenames_list.append(filename)
+
         return filenames_list
 
 
@@ -188,15 +234,16 @@ class NCEP_FNL(object):
 if __name__ == "__main__":
 
     dates = [
-        '2003-06-23',
-        '2005-06-11',
-        '2023-06-04'
+        #'2003-06-23',
+        #'2005-06-11',
+        #'2023-06-04'
+        '2023-06-23'
     ]
 
     for init_date in dates:
 
         start_date = datetime.datetime.strptime(init_date, "%Y-%m-%d")
-        end_date = (start_date + datetime.timedelta(days=10))
+        end_date = (start_date + datetime.timedelta(days=2))
 
         downloader = NCEP_FNL(
             output_dir='output/NCEP_FNL',

wildfire_occurrence/model/pipelines/wrf_pipeline.py

@@ -5,11 +5,15 @@ import logging
 import datetime
 from glob import glob
 from pathlib import Path
+from itertools import repeat
+from omegaconf import OmegaConf
+from multiprocessing import Pool, cpu_count
 from jinja2 import Environment, PackageLoader, select_autoescape
 
 from wildfire_occurrence.model.config import Config
 from wildfire_occurrence.model.common import read_config
 from wildfire_occurrence.model.data_download.ncep_fnl import NCEP_FNL
+from wildfire_occurrence.model.analysis.wrf_analysis import WRFAnalysis
 
 
 class WRFPipeline(object):
@@ -19,6 +23,7 @@ class WRFPipeline(object):
                 config_filename: str,
                 start_date: str,
                 forecast_lenght: str,
+                multi_node: bool = False
             ):
 
         # Configuration file intialization
@@ -42,7 +47,7 @@ class WRFPipeline(object):
             f'{self.end_date.strftime("%Y-%m-%d")}'
         )
         os.makedirs(self.simulation_dir, exist_ok=True)
-        logging.info(f'Created output directory {self.simulation_dir}')
+        logging.info(f'Created model output directory {self.simulation_dir}')
 
         # Setup data_dir
         self.data_dir = os.path.join(self.simulation_dir, 'data')
@@ -55,6 +60,8 @@ class WRFPipeline(object):
         # Setup wps directory
         self.local_wps_path = os.path.join(self.simulation_dir, 'WPS')
         self.local_wrf_path = os.path.join(self.simulation_dir, 'em_real')
+        self.local_wrf_output = os.path.join(self.simulation_dir, 'output')
+        self.local_wrf_output_vars = os.path.join(self.simulation_dir, 'variables')
 
         # Setup configuration filenames
         self.wps_conf_filename = os.path.join(self.conf_dir, 'namelist.wps')
@@ -65,6 +72,9 @@ class WRFPipeline(object):
         self.local_wrf_conf = os.path.join(
             self.local_wrf_path, 'namelist.input')
 
+        # setup multi_node variable
+        self.conf.multi_node = multi_node
+
     # -------------------------------------------------------------------------
     # setup
     # -------------------------------------------------------------------------
@@ -120,15 +130,15 @@ class WRFPipeline(object):
         if not self.conf.multi_node:
             geodrid_cmd = \
                 'singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
-                'mpirun -np 40 --oversubscribe ./geogrid.exe'
+                f'mpirun -np {cpu_count()} --oversubscribe ./geogrid.exe'
         else:
-            geodrid_cmd = \
-                'srun --mpi=pmix -N 2 -n 80 singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
-                f'{self.conf.container_path} ' + \
-                './geogrid.exe'
+            geodrid_cmd = 'bash /panfs/ccds02/nobackup/projects/ilab/projects/LobodaTFO/operations/2015-07-23_2015-08-02/WPS/run_geogrid.sh'
+            #'mpirun -np 40 --host gpu016 --oversubscribe singularity exec -B /explore/nobackup/projects/ilab,' + \
+            #'$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+            #f'{self.conf.container_path} ' + \
+            #'bash /panfs/ccds02/nobackup/projects/ilab/projects/LobodaTFO/operations/2015-07-23_2015-08-02/WPS/run_geogrid.sh'
 
         # run geogrid command
         os.system(geodrid_cmd)
@@ -160,22 +170,24 @@ class WRFPipeline(object):
         os.chdir(self.local_wps_path)
         logging.info(f'Changed working directory to {self.local_wps_path}')
 
+        # find all files in directory, and extract common prefix
+        common_prefix = os.path.commonprefix(
+            glob(f'{self.data_dir}/{str(self.start_date.year)}/*'))
+
         # run link_grib
-        os.system(
-            f'./link_grib.csh {self.data_dir}/{str(self.start_date.year)}/' +
-            f'fnl_{str(self.start_date.year)}')
+        os.system(f'./link_grib.csh {common_prefix}')
         logging.info('Done with link_grib.csh')
 
         # setup ungrib command
         if not self.conf.multi_node:
             ungrib_cmd = \
                 'singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
-                f'{self.conf.container_path} mpirun ./ungrib.exe'
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
+                f'{self.conf.container_path} ./ungrib.exe'
         else:
             ungrib_cmd = \
                 'srun --mpi=pmix -N 1 -n 1 singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
                 './ungrib.exe'
 
@@ -203,13 +215,13 @@ class WRFPipeline(object):
         if not self.conf.multi_node:
             metgrid_cmd = \
                 'singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
-                'mpirun -np 40 --oversubscribe ./metgrid.exe'
+                f'mpirun -np {cpu_count()} --oversubscribe ./metgrid.exe'
         else:
             metgrid_cmd = \
-                'srun --mpi=pmix -N 1 -n 40 singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                f'srun --mpi=pmix -N 1 -n {cpu_count()} singularity exec -B /explore/nobackup/projects/ilab,' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
                 './metgrid.exe'
 
@@ -250,13 +262,13 @@ class WRFPipeline(object):
         if not self.conf.multi_node:
             real_cmd = \
                 'singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
-                'mpirun -np 40 --oversubscribe ./real.exe'
+                f'mpirun -np {cpu_count()} --oversubscribe ./real.exe'
         else:
             real_cmd = \
                 'srun --mpi=pmix -N 2 -n 80 singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
                 './real.exe'
 
@@ -284,13 +296,13 @@ class WRFPipeline(object):
         if not self.conf.multi_node:
             wrf_cmd = \
                 'singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
-                'mpirun -np 40 --oversubscribe ./wrf.exe'
+                f'mpirun -np {cpu_count()} --oversubscribe ./wrf.exe'
         else:
             wrf_cmd = \
                 'srun --mpi=pmix -N 2 -n 80 singularity exec -B /explore/nobackup/projects/ilab,' + \
-                '$NOBACKUP,/lscratch,/panfs/ccds02/nobackup/projects/ilab ' + \
+                '$NOBACKUP,/panfs/ccds02/nobackup/projects/ilab ' + \
                 f'{self.conf.container_path} ' + \
                 './wrf.exe'
 
@@ -302,6 +314,114 @@ class WRFPipeline(object):
 
         return
 
+    # -------------------------------------------------------------------------
+    # postprocess
+    # -------------------------------------------------------------------------
+    def postprocess(self) -> None:
+
+        logging.info('Preparing to postprocess and extract variables')
+
+        # create output directory
+        os.makedirs(self.local_wrf_output, exist_ok=True)
+        logging.info(f'Created WRF output directory {self.local_wrf_output}')
+
+        # if the files have not been moved, move them to the output dir
+        if len(os.listdir(self.local_wrf_output)) == 0:
+
+            # get filenames, make sure they exist
+            wrf_output_filenames = \
+                glob(os.path.join(self.local_wrf_path, 'auxhist24_d0*')) + \
+                glob(os.path.join(self.local_wrf_path, 'wrfout_d0*'))
+            assert len(wrf_output_filenames) > 0, \
+                'WRF output (auxhist24_d0*, wrfout_d0*) not found. Re-run WRF.'
+
+            # move output to clean directory
+            for filename in wrf_output_filenames:
+                shutil.move(filename, self.local_wrf_output)
+            logging.info(f'Moved WRF output to {self.local_wrf_output}')
+
+        # Get WRF output filename
+        wrf_output_filename = glob(
+            os.path.join(self.local_wrf_output, self.conf.wrf_output_filename))
+        assert len(wrf_output_filename) == 1, \
+            f'WRF output filename not found under {self.local_wrf_output}.'
+
+        # Get first item from the list
+        wrf_output_filename = wrf_output_filename[0]
+        logging.info(f'Loading {wrf_output_filename}')
+
+        # create WRFAnalysis object, stores wrf_dataset
+        wrf_analysis = WRFAnalysis(wrf_output_filename)
+
+        # variables output dir
+        os.makedirs(self.local_wrf_output_vars, exist_ok=True)
+        logging.info(
+            f'Created WRF vars output directory {self.local_wrf_output_vars}')
+
+        """
+        # parallel extraction of variables
+        p = Pool(processes=cpu_count())
+        p.starmap(
+            self._compute_variables,
+            zip(
+                range(len(wrf_analysis.wrf_dataset.Times.values)),
+                wrf_analysis.wrf_dataset.Times.values,
+                repeat(wrf_output_filename)
+            )
+        )
+        """
+
+        # serial extraction of variables
+        for t_idx, delta_time in \
+                enumerate(wrf_analysis.wrf_dataset.Times.values):
+
+            logging.info(f'Processing t_idx: {t_idx}, timestamp: {delta_time}')
+
+            # setup output filename
+            output_filename = os.path.join(
+                    self.local_wrf_output_vars,
+                    f"d02_{delta_time.astype(str).replace(':', '-')}.tif")
+
+            # if the imagery does not exist
+            if not os.path.isfile(output_filename):
+
+                # compute WRF variables and output to disk
+                wrf_analysis.compute_all_and_write(
+                    timeidx=t_idx,
+                    output_variables=OmegaConf.to_object(
+                        self.conf.wrf_output_variables),
+                    output_filename=output_filename
+                )
+
+    # -------------------------------------------------------------------------
+    # _compute_variables
+    # -------------------------------------------------------------------------
+    def _compute_variables(self, t_idx, delta_time, wrf_output_filename):
+
+        logging.info(f'Processing t_idx: {t_idx}, timestamp: {delta_time}')
+
+        # setup output filename
+        output_filename = os.path.join(
+                self.local_wrf_output_vars,
+                f"d02_{delta_time.astype(str).replace(':', '-')}.tif")
+
+        # if the imagery does not exist
+        if not os.path.isfile(output_filename):
+
+            # unfortunately, netCDF object is not pickable, thus we need
+            # to redifine it on every single process, hopefully there
+            # will be an implementation on it at some point by them
+            wrf_analysis = WRFAnalysis(wrf_output_filename)
+
+            # compute WRF variables and output to disk
+            wrf_analysis.compute_all_and_write(
+                timeidx=t_idx,
+                output_variables=OmegaConf.to_object(
+                        self.conf.wrf_output_variables),
+                output_filename=output_filename
+            )
+        return
+
     # -------------------------------------------------------------------------
     # setup_wps_config
     # -------------------------------------------------------------------------

wildfire_occurrence/templates/config.yaml

@@ -12,13 +12,46 @@ wps_config:
 
 wrf_config:
   interval_seconds: 10800
-  num_metgrid_levels: 27
+  num_metgrid_levels: 34
 
 # Directories to mount inside the container
 container_mounts:
   - /explore/nobackup/projects/ilab'
   - '$NOBACKUP'
-  - '/lscratch'
   - '/panfs/ccds02/nobackup/projects/ilab'
 
-multi_node: False
+multi_node: False
+
+wrf_output_filename: 'wrfout_d02_*_00:00:00'
+
+wrf_output_variables:
+  - 'CFTotal'
+#  - 'CFLow'
+#  - 'CFMed'
+#  - 'CFHigh'
+#  - 'DZ700_850'
+#  - 'GPZ500'
+#  - 'GPZ700'
+#  - 'GPZ750'
+#  - 'GPZ850'
+#  - 'Helicity'
+#  - 'LCL'
+#  - 'PLI'
+#  - 'PW'
+#  - 'RAINTotal'
+#  - 'RH2'
+#  - 'RH500'
+#  - 'RH700'
+#  - 'RH800'
+#  - 'RH850'
+#  - 'SHOW'
+#  - 'SLP'
+#  - 'TD2'
+#  - 'TD500'
+#  - 'TT'
+#  - 'T2'
+#  - 'T500'
+#  - 'T750'
+#  - 'T850'
+#  - 'W500'
+#  - 'WA500'

wildfire_occurrence/view/wrf_pipeline_cli.py

@@ -42,6 +42,14 @@ def main():
                         dest='forecast_lenght',
                         help='Lenght of WRF forecast')
 
+    parser.add_argument('-mn',
+                        '--multi-node',
+                        type=bool,
+                        required=False,
+                        default=False,
+                        dest='multi_node',
+                        help='Multiple nodes for WRF forecast')
+
     parser.add_argument(
                         '-s',
                         '--pipeline-step',
@@ -52,10 +60,10 @@ def main():
                         help='Pipeline step to perform',
                         default=[
                             'setup', 'geogrid', 'ungrib', 'metgrid',
-                            'real', 'wrf', 'all'],
+                            'real', 'wrf', 'postprocess', 'all'],
                         choices=[
                             'setup', 'geogrid', 'ungrib', 'metgrid',
-                            'real', 'wrf', 'all'])
+                            'real', 'wrf', 'postprocess', 'all'])
 
     args = parser.parse_args()
 
@@ -75,7 +83,9 @@ def main():
 
     # Initialize pipeline object
     pipeline = WRFPipeline(
-        args.config_file, args.start_date, args.forecast_lenght)
+        args.config_file, args.start_date,
+        args.forecast_lenght, args.multi_node
+    )
 
     # WRF pipeline steps
     if "setup" in args.pipeline_step or "all" in args.pipeline_step:
@@ -90,6 +100,8 @@ def main():
         pipeline.real()
     if "wrf" in args.pipeline_step or "all" in args.pipeline_step:
         pipeline.wrf()
+    if "postprocess" in args.pipeline_step or "all" in args.pipeline_step:
+        pipeline.postprocess()
 
     logging.info(f'Took {(time.time()-timer)/60.0:.2f} min.')