Mimic parallel read in unstable

Author: jinmannwong

src/pproc/thermal_indices.py

@@ -34,125 +34,13 @@
 )
 from pproc.thermo import helpers
 from pproc.thermo.indices import ComputeIndices
-from pproc.thermo.wrappers import ArrayFieldList
+from pproc.thermo.wrappers import ArrayFieldList, GribMetadata
 
 logging.getLogger("pproc").setLevel(os.environ.get("PPROC_LOGLEVEL", "INFO").upper())
 logging.basicConfig(format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
 logger = logging.getLogger(__name__)
 
-
-@metered("Process step", out=logger.info)
-def process_step(args, config, window_id, fields, accum: Accumulator, accum_metadata: list[GribFieldMetadata], recovery: Recovery,
-):
-    fields = load_input("inst", config, step)
-    if len(accum["step"].coords) > 1:
-        # Set step range for de-accumulated fields
-        step_range = "-".join(map(str, accum["step"].coords))
-        accum_fields = earthkit.data.FieldList.from_array(
-            accum.values,
-            [x.override(stepType="diff", stepRange=step_range) for x in accum_metadata],
-        )
-        helpers.write(config.outputs.accum, accum_fields)
-        fields += accum_fields
-    
-    helpers.check_field_sizes(fields)
-    basetime, validtime = helpers.get_datetime(fields)
-    step = helpers.get_step(fields)
-    time = basetime.hour
-
-    logger.info(
-        f"Compute indices step {step}, validtime {validtime.isoformat()} - "
-        + f"basetime {basetime.date().isoformat()}, time {time}"
-    )
-    logger.debug(f"Inputs \n {fields.ls(namespace='mars')}")
-    indices = ComputeIndices(config.out_keys)
-    params = fields.indices()["param"]
-
-    # Windspeed - shortName ws
-    if config.is_target_param(step, {"10si", 207}):
-        ws = indices.calc_field("10si", indices.calc_ws, fields)
-        config.write(step, ws)
-
-    # Cosine of Solar Zenith Angle - shortName cossza - ECMWF product
-    # TODO: 214001 only exists for GRIB1 -- but here we use it for GRIB2 (waiting for WMO)
-    if config.is_target_param(step, {"cossza", 214001}):
-        cossza = indices.calc_field("cossza", indices.calc_cossza_int, fields)
-        config.write(step, cossza)
-
-    # direct solar radiation - shortName dsrp - ECMWF product
-    if config.is_target_param(step, {"dsrp", 47}) and "dsrp" not in params:
-        dsrp = indices.calc_field("dsrp", indices.calc_dsrp, fields)
-        config.write(step, dsrp)
-
-    # Mean Radiant Temperature - shortName mrt - ECMWF product
-    if config.is_target_param(step, {"mrt", 261002}):
-        mrt = indices.calc_field("mrt", indices.calc_mrt, fields)
-        config.write(step, mrt)
-
-    # Univeral Thermal Climate Index - shortName utci - ECMWF product
-    if config.is_target_param(step, {"utci", 261001}):
-        utci = indices.calc_field(
-            "utci",
-            indices.calc_utci,
-            fields,
-            print_misses=args.utci_misses,
-            validate=args.validateutci,
-        )
-        config.write(step, utci)
-
-    # Heat Index (adjusted) - shortName heatx - ECMWF product
-    if config.is_target_param(step, {"heatx", 260004}):
-        heatx = indices.calc_field("heatx", indices.calc_heatx, fields)
-        config.write(step, heatx)
-
-    # Wind Chill factor - shortName wcf - ECMWF product
-    if config.is_target_param(step, {"wcf", 260005}):
-        wcf = indices.calc_field("wcf", indices.calc_wcf, fields)
-        config.write(step, wcf)
-
-    # Apparent Temperature - shortName aptmp - ECMWF product
-    if config.is_target_param(step, {"aptmp", 260255}):
-        aptmp = indices.calc_field("aptmp", indices.calc_aptmp, fields)
-        config.write(step, aptmp)
-
-    # Relative humidity percent at 2m - shortName 2r - ECMWF product
-    if config.is_target_param(step, {"2r", 260242}):
-        rhp = indices.calc_field("2r", indices.calc_rhp, fields)
-        config.write(step, rhp)
-
-    # Humidex - shortName hmdx
-    if config.is_target_param(step, {"hmdx", 261016}):
-        hmdx = indices.calc_field("hmdx", indices.calc_hmdx, fields)
-        config.write(step, hmdx)
-
-    # Normal Effective Temperature - shortName nefft
-    if config.is_target_param(step, {"nefft", 261018}):
-        nefft = indices.calc_field("nefft", indices.calc_nefft, fields)
-        config.write(step, nefft)
-
-    # Globe Temperature - shortName gt
-    if config.is_target_param(step, {"gt", 261015}):
-        gt = indices.calc_field("gt", indices.calc_gt, fields)
-        config.write(step, gt)
-
-    # Wet-bulb potential temperature - shortName wbpt
-    if config.is_target_param(step, {"wbpt", 261022}):
-        wbpt = indices.calc_field("wbpt", indices.calc_wbpt, fields)
-        config.write(step, wbpt)
-
-    # Wet Bulb Globe Temperature - shortName wbgt
-    if config.is_target_param(step, {"wbgt", 261014}):  #
-        wbgt = indices.calc_field("wbgt", indices.calc_wbgt, fields)
-        config.write(step, wbgt)
-
-    # effective temperature 261017
-    # standard effective temperature 261019
-
-    config.flush_targets()
-    recovery.add_checkpoint(window_id)
-
-    if args.usage:
-        print_usage()
+__version__ = "2.0.0"
 
 
 class ThermoTarget:
@@ -254,13 +142,13 @@ def flush_targets(self):
 
 
 def load_input(source: str, config: ThermoConfig, step: int):
-    logger.debug(f"Retrieve step {step}: source {source}")
     req = config.sources.get(source, {}).copy()
     if len(req) == 0:
         return None
     req.update(config.override_input)
     req["step"] = [step]
 
+    logger.debug(f"Retrieve step {step}: source {source}")
     src = req.pop("source")
     if ":" in src:
         src, loc = src.split(":")
@@ -281,6 +169,127 @@ def load_input(source: str, config: ThermoConfig, step: int):
     return earthkit.data.FieldList.from_array(ds.values, ds.metadata())
 
 
+@metered("Process step", out=logger.info)
+def process_step(
+    args,
+    config,
+    window_id,
+    step,
+    accum,
+    accum_metadata,
+    recovery,
+):
+    fields = load_input("inst", config, step)
+    if len(accum["step"].coords) > 1:
+        # Set step range for de-accumulated fields
+        step_range = "-".join(map(str, accum["step"].coords))
+        accum_fields = earthkit.data.FieldList.from_array(
+            accum.values,
+            [x.override(stepType="diff", stepRange=step_range) for x in accum_metadata],
+        )
+        config.write(step, accum_fields)
+        fields += accum_fields
+
+    helpers.check_field_sizes(fields)
+    basetime, validtime = helpers.get_datetime(fields)
+    step = helpers.get_step(fields)
+    time = basetime.hour
+
+    logger.info(
+        f"Compute indices step {step}, validtime {validtime.isoformat()} - "
+        + f"basetime {basetime.date().isoformat()}, time {time}"
+    )
+    logger.debug(f"Inputs \n {fields.ls(namespace='mars')}")
+    indices = ComputeIndices(config.out_keys)
+    params = fields.indices()["param"]
+
+    # Windspeed - shortName ws
+    if config.is_target_param(step, {"10si", 207}):
+        ws = indices.calc_field("10si", indices.calc_ws, fields)
+        config.write(step, ws)
+
+    # Cosine of Solar Zenith Angle - shortName cossza - ECMWF product
+    # TODO: 214001 only exists for GRIB1 -- but here we use it for GRIB2 (waiting for WMO)
+    if config.is_target_param(step, {"cossza", 214001}):
+        cossza = indices.calc_field("cossza", indices.calc_cossza_int, fields)
+        config.write(step, cossza)
+
+    # direct solar radiation - shortName dsrp - ECMWF product
+    if config.is_target_param(step, {"dsrp", 47}) and "dsrp" not in params:
+        dsrp = indices.calc_field("dsrp", indices.calc_dsrp, fields)
+        config.write(step, dsrp)
+
+    # Mean Radiant Temperature - shortName mrt - ECMWF product
+    if config.is_target_param(step, {"mrt", 261002}):
+        mrt = indices.calc_field("mrt", indices.calc_mrt, fields)
+        config.write(step, mrt)
+
+    # Univeral Thermal Climate Index - shortName utci - ECMWF product
+    if config.is_target_param(step, {"utci", 261001}):
+        utci = indices.calc_field(
+            "utci",
+            indices.calc_utci,
+            fields,
+            print_misses=args.utci_misses,
+            validate=args.validateutci,
+        )
+        config.write(step, utci)
+
+    # Heat Index (adjusted) - shortName heatx - ECMWF product
+    if config.is_target_param(step, {"heatx", 260004}):
+        heatx = indices.calc_field("heatx", indices.calc_heatx, fields)
+        config.write(step, heatx)
+
+    # Wind Chill factor - shortName wcf - ECMWF product
+    if config.is_target_param(step, {"wcf", 260005}):
+        wcf = indices.calc_field("wcf", indices.calc_wcf, fields)
+        config.write(step, wcf)
+
+    # Apparent Temperature - shortName aptmp - ECMWF product
+    if config.is_target_param(step, {"aptmp", 260255}):
+        aptmp = indices.calc_field("aptmp", indices.calc_aptmp, fields)
+        config.write(step, aptmp)
+
+    # Relative humidity percent at 2m - shortName 2r - ECMWF product
+    if config.is_target_param(step, {"2r", 260242}):
+        rhp = indices.calc_field("2r", indices.calc_rhp, fields)
+        config.write(step, rhp)
+
+    # Humidex - shortName hmdx
+    if config.is_target_param(step, {"hmdx", 261016}):
+        hmdx = indices.calc_field("hmdx", indices.calc_hmdx, fields)
+        config.write(step, hmdx)
+
+    # Normal Effective Temperature - shortName nefft
+    if config.is_target_param(step, {"nefft", 261018}):
+        nefft = indices.calc_field("nefft", indices.calc_nefft, fields)
+        config.write(step, nefft)
+
+    # Globe Temperature - shortName gt
+    if config.is_target_param(step, {"gt", 261015}):
+        gt = indices.calc_field("gt", indices.calc_gt, fields)
+        config.write(step, gt)
+
+    # Wet-bulb potential temperature - shortName wbpt
+    if config.is_target_param(step, {"wbpt", 261022}):
+        wbpt = indices.calc_field("wbpt", indices.calc_wbpt, fields)
+        config.write(step, wbpt)
+
+    # Wet Bulb Globe Temperature - shortName wbgt
+    if config.is_target_param(step, {"wbgt", 261014}):  #
+        wbgt = indices.calc_field("wbgt", indices.calc_wbgt, fields)
+        config.write(step, wbgt)
+
+    # effective temperature 261017
+    # standard effective temperature 261019
+
+    config.flush_targets()
+    recovery.add_checkpoint(window_id)
+
+    if args.usage:
+        print_usage()
+
+
 def get_parser():
     parser = default_parser("Compute thermal indices")
 
@@ -385,8 +394,9 @@ def main(args: List[str] = sys.argv[1:]):
                 executor.submit(
                     thermo_partial,
                     window_id,
-                    accum, 
-                    accum_data.metadata(),
+                    step,
+                    accum,
+                    [GribMetadata(x._handle) for x in accum_data.metadata()],
                 )
         executor.wait()
     recovery.clean_file()