Updated read from HEALPix catalog file.

1. Updated these codes to read HEALPix data from catalog file.
- pyflextrkr/idfeature_driver.py
- pyflextrkr/idclouds_tbpf.py
- pyflextrkr/remap_healpix.py
2. Updated the example config file: config/config_mcs_tbpf_scream_healpix9.yml
3. Added intake libraries in the environment list:
- environment.yml
- requirements.txt

Author: feng045

config/config_mcs_tbpf_scream_healpix9.yml

@@ -26,19 +26,28 @@ timeout: 3600  # [seconds] Dask timeout limit
 startdate: '20190901.0000'
 enddate: '20200901.0000'
 
+# HEALPix parameters
+catalog_file: '/global/cfs/cdirs/m4549/scream-cess-healpix/scream_catalog.yaml'
+catalog_source: 'scream2D_hrly'
+# Catalog parameters, can have multiple entries
+catalog_params:
+  zoom: 9
+#  time: 'PT1H'
+olr_varname: 'rlut'  # OLR variable name
+pcp_varname: 'pr'   # Precipitation variable name
+pcp_convert_factor: 3600000.  # Convert precipitation flux from [m/s] to [mm/h]
+input_format: 'zarr'  # 'zarr' or 'netcdf'
+
 # Specify tracking input data date/time string format
 # This is the preprocessed file that contains Tb & rainrate
 # E.g., databasename20181101.011503.nc --> yyyymodd.hhmmss
 # E.g., databasename2018-11-01_01:15:00 --> yyyy-mo-dd_hh:mm:ss
 time_format: 'yyyymoddhh'
-databasename:  'scream2D_hrly_rlut_hp9_v6'
-precipdata_basename: 'scream2D_hrly_pr_hp9_v6'
-hp_zoomlev: '9'  # HEALPix zoom level
-input_format: 'zarr'  # 'zarr' or 'netcdf'
+databasename: ''  # This is not used for HEALPix data
 
 # Input files directory
-clouddata_path: '/pscratch/sd/w/wcmca1/scream-cess-healpix/'
-precipdata_path: '/pscratch/sd/w/wcmca1/scream-cess-healpix/'
+clouddata_path: ''  # Not needed for catalog data
+# precipdata_path: '/pscratch/sd/w/wcmca1/scream-cess-healpix/' # Not needed for catalog data
 # Working directory for the tracking data
 root_path: '/pscratch/sd/w/wcmca1/scream-cess-healpix/mcs_tracking_hp9/'
 # Working sub-directory names
@@ -60,9 +69,6 @@ pixel_radius:  10.0  # [km] Spatial resolution of the input data
 datatimeresolution: 1.0  # [hour] Temporal resolution of the input data
 # Variable names in the input data
 olr2tb: True
-olr_varname: 'rlut'
-pcp_varname: 'pr'
-pcp_convert_factor: 3600000.  # Convert precipitation flux from [m/s] to [mm/h]
 clouddatasource: 'model'
 time_dimname: 'time'
 x_dimname: 'lon'

environment.yml

@@ -10,6 +10,8 @@ dependencies:
   - dask
   - ffmpeg
   - healpix
+  - intake==0.7.0
+  - intake-xarray==0.7.0
   - ipython
   - joblib
   - matplotlib

pyflextrkr/idclouds_tbpf.py

@@ -107,8 +107,8 @@ def idclouds_tbpf(
         # Read landmask file to get target lat/lon grid
         landmask_filename = config.get('landmask_filename', None)
         dslm = xr.open_dataset(landmask_filename)
-        lon = dslm.lon.data
-        lat = dslm.lat.data
+        lon = dslm.lon.values
+        lat = dslm.lat.values
         dslm.close()
 
         # Find the HEALPix pixels that are closest to the target grid points
@@ -120,11 +120,14 @@ def idclouds_tbpf(
         # Note this would change the calendar type of the original time coordinate
         time_coord = input_data[time_coordname]
         time_pd = pd.to_datetime(time_coord.dt.strftime("%Y-%m-%dT%H:%M:%S").item())
-        # Regrid variables, expand time dimension so the variable has dimensions [time, y, x]
-        olr = input_data[olr_varname].isel(cell=pix).expand_dims({time_dimname:[time_pd]})
-        pcp = input_data[pcp_varname].isel(cell=pix).expand_dims({time_dimname:[time_pd]})
-        # Combine DataArrays into a single Dataset
-        rawdata = xr.Dataset({olr_varname: olr, pcp_varname: pcp})
+        # Remap DataSet to lat/lon grid, expand time dimension so it has dimensions [time, y, x]
+        rawdata = input_data.isel(cell=pix).expand_dims({time_dimname:[time_pd]})
+
+        # # Remap variables, expand time dimension so the variable has dimensions [time, y, x]
+        # olr = input_data[olr_varname].isel(cell=pix).expand_dims({time_dimname:[time_pd]})
+        # pcp = input_data[pcp_varname].isel(cell=pix).expand_dims({time_dimname:[time_pd]})
+        # # Combine DataArrays into a single Dataset
+        # rawdata = xr.Dataset({olr_varname: olr, pcp_varname: pcp})
 
     # NetCDF format
     elif input_format.lower() == "netcdf":
@@ -164,8 +167,8 @@ def idclouds_tbpf(
             logger.debug(f'Added Timestamp: {file_timestamp} calculated from filename to the input data')
 
     # Get data coordinates
-    lat = rawdata[y_coordname].data
-    lon = rawdata[x_coordname].data
+    lat = rawdata[y_coordname].values
+    lon = rawdata[x_coordname].values
     time_decode = rawdata[time_coordname]
 
     # Check coordinate dimensions
@@ -210,8 +213,8 @@ def idclouds_tbpf(
     # Subset dataset
     rawdata = rawdata[subset_dict]
     # Get lat/lon coordinates again
-    lat = rawdata[y_coordname].data
-    lon = rawdata[x_coordname].data
+    lat = rawdata[y_coordname].values
+    lon = rawdata[x_coordname].values
     # Check coordinate dimensions
     if (lat.ndim == 1) | (lon.ndim == 1):
         # Mesh 1D coordinate into 2D
@@ -223,13 +226,12 @@ def idclouds_tbpf(
 
     # Convert OLR to Tb if olr2tb flag is set
     if olr2tb is True:
-        olr = rawdata[olr_varname].data
+        olr = rawdata[olr_varname].values
         original_ir = olr_to_tb(olr)
     else:
         # Read Tb from data
-        original_ir = rawdata[tb_varname].data
-    rawdata.close()
-
+        original_ir = rawdata[tb_varname].values
+    # rawdata.close()
 
     # Loop over each time
     ntimes = get_length(time_decode)
@@ -337,7 +339,7 @@ def idclouds_tbpf(
                         if final_nclouds > 0:
 
                             # Convert precipitation factor to unit [mm/hour]
-                            pcp = rawdata[pcp_varname].data * pcp_convert_factor
+                            pcp = rawdata[pcp_varname].values * pcp_convert_factor
 
                             # For 'gpmirimerg', precipitation is averaged to 1-hourly
                             # and put in first time dimension
@@ -359,8 +361,9 @@ def idclouds_tbpf(
                             # Replace values <=0 with 0 before smoothing
                             pcp_linkpf[pcp_linkpf <= 0] = 0
 
+                            # Check piriodic boundary conditions
                             if pbc_direction != 'none':
-                                # Step 2: Extend and pad data
+                                # Extend and pad data
                                 pcp_linkpf_orig = np.copy(pcp_linkpf)
                                 pcp_linkpf, padded_x, padded_y = pad_and_extend(pcp_linkpf, config)
                                 # Smooth pcp_linkpf using convolve filter (handles NaN)

pyflextrkr/idfeature_driver.py

@@ -22,10 +22,9 @@ def idfeature_driver(config):
     logger.info('Identifying features from raw data')
 
     clouddata_path = config["clouddata_path"]
-    databasename = config["databasename"]
+    databasename = config.get("databasename", "")
     start_basetime = config.get("start_basetime", None)
     end_basetime = config.get("end_basetime", None)
-    # time_format = config["time_format"]
     run_parallel = config["run_parallel"]
     feature_type = config["feature_type"]
     input_format = config.get("input_format", "netcdf")
@@ -46,50 +45,54 @@ def idfeature_driver(config):
 
     if input_format.lower() == "zarr":
 
-        # Get precipitation data info from config
-        precipdata_path = config["precipdata_path"]
-        precipdata_basename = config["precipdata_basename"]
+        import intake     # For catalogs
+
+        # Get catalog info from config
+        catalog_file = config["catalog_file"]
+        catalog_source = config["catalog_source"]
+        catalog_params = config.get("catalog_params", {})
+        olr_varname = config['olr_varname']
+        pcp_varname = config['pcp_varname']
         start_date = config["startdate"]
         end_date = config["enddate"]
 
-        # OLR Zarr filename
-        fn_olr = f"{clouddata_path}{databasename}.zarr"
-        fn_pr = f"{precipdata_path}{precipdata_basename}.zarr"
-        # Read HEALPix zarr file
-        ds_olr = xr.open_dataset(fn_olr)
-        ds_pr = xr.open_dataset(fn_pr)
+        # Load the catalog
+        in_catalog = intake.open_catalog(catalog_file)
+        # Get the DataSet from the catalog
+        ds = in_catalog[catalog_source](**catalog_params).to_dask()
+
+        # Subset to keep only the required variables
+        all_vars = list(ds.data_vars)
+        keep_vars = [olr_varname, pcp_varname]
+        drop_vars = [var for var in all_vars if var not in keep_vars]
+        ds = ds.drop_vars(drop_vars)
+
         # Check the calendar type of the time coordinate
-        calendar = ds_olr['time'].dt.calendar
-        # Add coordinates (lat and lon)
-        # ds_olr = ds_olr.pipe(egh.attach_coords)
-        # ds_pr = ds_pr.pipe(egh.attach_coords)
+        calendar = ds['time'].dt.calendar
         # Convert start_date and end_date to pandas.Timestamp
         start_datetime = pd.to_datetime(start_date, format='%Y%m%d.%H%M')
         end_datetime = pd.to_datetime(end_date, format='%Y%m%d.%H%M')
         # Convert pandas.Timestamp to cftime objects based on the calendar type
         start_datetime_cftime = convert_to_cftime(start_datetime, calendar)
         end_datetime_cftime = convert_to_cftime(end_datetime, calendar)
         # Subset the Dataset using the cftime objects
-        ds_olr = ds_olr.sel(time=slice(start_datetime_cftime, end_datetime_cftime))
-        ds_pr = ds_pr.sel(time=slice(start_datetime_cftime, end_datetime_cftime))
+        ds = ds.sel(time=slice(start_datetime_cftime, end_datetime_cftime))
 
         # Get the number of time steps
-        nfiles = ds_olr.sizes['time']
+        nfiles = ds.sizes['time']
         logger.info(f"Total number of time steps to process: {nfiles}")
 
         # Serial
         if run_parallel == 0:
             for ifile in range(0, nfiles):
                 # Subset one time from the DataSets and combine them
-                ds = xr.merge([ds_olr.isel(time=ifile), ds_pr.isel(time=ifile)])
-                id_feature(ds, config)
+                id_feature(ds.isel(time=ifile), config)
         # Parallel
         elif run_parallel >= 1:
             results = []
             for ifile in range(0, nfiles):
                 # Subset one time from the DataSets and combine them
-                ds = xr.merge([ds_olr.isel(time=ifile), ds_pr.isel(time=ifile)])
-                result = dask.delayed(id_feature)(ds, config)
+                result = dask.delayed(id_feature)(ds.isel(time=ifile), config)
                 results.append(result)
             final_result = dask.compute(*results)
             wait(final_result)

pyflextrkr/remap_healpix.py

@@ -5,6 +5,7 @@
 import logging
 import dask.array as da
 import healpix as hp
+import intake
 from pyflextrkr.ft_utilities import setup_logging   
 
 def remap_mask_to_healpix(config):
@@ -30,18 +31,17 @@ def remap_mask_to_healpix(config):
     try:
         from dask.distributed import get_client
         client = get_client()
-        parallel = True
         logger.info(f"Using existing Dask client with {len(client.scheduler_info()['workers'])} workers")
     except ValueError:
         logger.warning("No Dask client found, continuing without explicit client")
         client = None
-        parallel = False
 
     # Get config parameters
     pixeltracking_outpath = config.get("pixeltracking_outpath")
-    clouddata_path = config.get("clouddata_path")
-    databasename = config.get("databasename")
-    hp_zoomlev = config.get("hp_zoomlev")
+    catalog_file = config.get("catalog_file")
+    catalog_source = config.get("catalog_source")
+    catalog_params = config.get("catalog_params", {})
+    hp_zoomlev = catalog_params.get("zoom")
     startdate = config.get("startdate")
     enddate = config.get("enddate")
     outpath = os.path.dirname(os.path.normpath(pixeltracking_outpath)) + "/"
@@ -50,8 +50,6 @@ def remap_mask_to_healpix(config):
     in_mask_filebase = presets.get("mask", {}).get("out_filebase", "mcs_mask_latlon_")
     # Input mask Zarr store
     in_mask_dir = f"{outpath}{in_mask_filebase}{startdate}_{enddate}.zarr"
-    # Input HEALPix Zarr store
-    in_hp_dir = f"{clouddata_path}{databasename}.zarr"
 
     # Build output filename
     out_mask_filebase = presets.get("healpix", {}).get("out_filebase", "mcs_mask_hp")
@@ -67,8 +65,8 @@ def remap_mask_to_healpix(config):
     if os.path.exists(in_mask_dir) is False:
         logger.error(f"Input mask file {in_mask_dir} does not exist. Skipping remap.")
         return out_zarr
-    if os.path.exists(in_hp_dir) is False:
-        logger.error(f"Input HEALPix file {in_hp_dir} does not exist. Skipping remap.")
+    if os.path.isfile(catalog_file) is False:
+        logger.error(f"Catalog file {catalog_file} does not exist. Skipping remap.")
         return out_zarr
 
     # Get chunk size from config
@@ -82,8 +80,10 @@ def remap_mask_to_healpix(config):
     # Modify mask grid for remapping
     ds_mask = prepare_grid_for_analysis(ds_mask)
 
-    # Read HEALPix zarr file
-    ds_hp = xr.open_dataset(in_hp_dir)
+    # Load the catalog
+    in_catalog = intake.open_catalog(catalog_file)
+    # Get the DataSet from the catalog
+    ds_hp = in_catalog[catalog_source](**catalog_params).to_dask()
 
     # Make remap lat/lon for HEALPix
     remap_lons, remap_lats = hp.pix2ang(

requirements.txt

@@ -5,6 +5,8 @@ colormath
 dask
 ffmpeg
 healpix
+intake==0.7.0
+intake-xarray==0.7.0
 ipython
 joblib
 matplotlib