Coverage for src/CSET/cset_workflow/app/fetch_nimrod/bin/fetch_nimrod.py: 80%
85 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 14:18 +0000
1#! /usr/bin/env python3
3"""Retrieve UK Nimrod radar observations. Specific to the Met Office."""
5import json
6import logging
7import os
8from datetime import datetime, timedelta
9from pathlib import Path
11import iris
12import isodate
13import numpy as np
15logger = logging.getLogger(__name__)
17iris.FUTURE.save_split_attrs = True
18iris.FUTURE.date_microseconds = True
20# Define the file to use when retrieving the sensitive Nimrod handling data.
21nimrod_met_office = "restricted_nimrod_met_office.json"
24def _get_needed_environment_variables_nimrod() -> dict:
25 """Load the needed variables from the environment to retrieve UK Nimrod data."""
26 radar_sources = []
27 if os.environ["NIMROD_COMP_XKM"] == "True": 27 ↛ 29line 27 didn't jump to line 29 because the condition on line 27 was always true
28 radar_sources.append("Nimrod_comp_xkm")
29 if os.environ["NIMROD_COMP_1KM"] == "True": 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true
30 radar_sources.append("Nimrod_comp_1km")
31 if os.environ["NIMROD_COMP_2KM"] == "True": 31 ↛ 32line 31 didn't jump to line 32 because the condition on line 31 was never true
32 radar_sources.append("Nimrod_comp_2km")
33 if os.environ["NIMROD_COMP_5MIN"] == "True": 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true
34 radar_sources.append("Nimrod_comp_5min")
35 variables = {
36 "field": radar_sources,
37 "weights": os.environ["NIMROD_WEIGHTS"],
38 "date_type": "initiation",
39 "data_time": datetime.fromisoformat(os.environ["CYLC_TASK_CYCLE_POINT"]),
40 "forecast_length": isodate.parse_duration(os.environ["ANALYSIS_LENGTH"]),
41 "rose_datac": os.environ["ROSE_DATAC"],
42 }
43 logger.debug("Environment variables loaded for Nimrod: %s", variables)
44 return variables
47def apply_radar_weights(cube_obs: iris.cube.Cube, cube_wei: iris.cube.Cube):
48 """Apply the Nimrod weights to the radar hourly rainfall accumulation data.
50 Parameters
51 ----------
52 cube_obs: Cube or
53 2 dimensional Cube of the radar rainfall accumulation data.
55 cube_wei: Cube or
56 2 dimensional Cube of the Nimrod rainfall accumulation weights.
58 """
59 # The weights are the number of 5 minute rainfall rates used to
60 # compute the hourly rainfall accumulation. The weight should be
61 # in the range 0 to 13.
62 # Define the minimum weight to accept as good data.
63 weight_min = 11
65 # Define the value to use for duff radar data
66 duff_value = 0.0
68 # Check if the weights are packed as weights / 32
69 # and if they are, unpack them.
70 # Note: if the weights are packed, then the maximum value
71 # found will be 13 / 32 = 0.40625, i.e. all values < 1.
72 weights = cube_wei.data
73 if weights.max() < 1.0:
74 weights = (weights * 32).round().astype(int)
75 cube_wei.data = weights
76 logger.info("Unpacked Nimrod weights file.")
78 # Apply the weights.
79 cube_obs_weighted = cube_obs
80 cube_obs_weighted.data = np.where(
81 weights < weight_min, duff_value, cube_obs_weighted.data
82 )
84 # Return the QC'd Nimrod data and the weights used.
85 return cube_obs_weighted, cube_wei
88def retrieve_nimrod():
89 """Fetch the observations corresponding to a model run.
91 The following environment variables need to be set:
92 * NIMROD_COMP_XKM, NIMROD_COMP_1KM, NIMROD_COMP_2KM, NIMROD_COMP_5MIN
93 * CYLC_TASK_CYCLE_POINT
94 * ANALYSIS_LENGTH
95 * ROSE_DATAC
96 """
98 # Grab the environment variables required for handling Nimrod data.
99 def _load_nimrod(path: Path) -> iris.cube.Cube | None:
100 """Load a Nimrod file, or return None if it's missing."""
101 if not path.exists(): 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 logger.warning("Nimrod file not found: %s", path)
103 return None
104 return iris.load_cube(path)
106 v = _get_needed_environment_variables_nimrod()
108 # Grab the Nimrod handling dictionary.
109 with open(nimrod_met_office, "rt") as fp: # pragma: no cover
110 nimrod_dict = json.load(fp)
112 # Loop over the required Nimrod fields, i.e. 1km 2km or xkm rainfall
113 # accumulation composites or the 5 minute rainfall rate composites.
114 for nimrod_field in v["field"]:
115 if nimrod_field: 115 ↛ 114line 115 didn't jump to line 114 because the condition on line 115 was always true
116 logger.info("Processing Nimrod field: %s", nimrod_field)
117 nimrod_cfg = nimrod_dict[nimrod_field]
118 # Prepare the output directory for the Nimrod field.
119 nimrod_dir = Path(v["rose_datac"]) / "data" / nimrod_cfg["obs_id"]
120 os.makedirs(nimrod_dir, exist_ok=True)
121 logger.info("Cylc-run Nimrod directory: %s", nimrod_dir)
123 # Prepare the output directory for the Nimrod weights field.
124 if nimrod_dict[nimrod_field]["weights_fname"]: 124 ↛ 132line 124 didn't jump to line 132 because the condition on line 124 was always true
125 nimrod_dir_wei = Path(v["rose_datac"]) / "data" / nimrod_cfg["wei_id"]
126 os.makedirs(nimrod_dir_wei, exist_ok=True)
127 logger.info("Cylc-run Nimrod weights directory: %s", nimrod_dir_wei)
129 # Put +1 hour offset for accumulation radar files as
130 # the time stamps for these files mark the end of the
131 # accumulation period rather than the beginning.
132 radar_offset = timedelta(hours=0.0)
133 if nimrod_field == "Nimrod_comp_xkm": 133 ↛ 135line 133 didn't jump to line 135 because the condition on line 133 was always true
134 radar_offset = timedelta(hours=1.0)
135 if nimrod_field == "Nimrod_comp_1km": 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true
136 radar_offset = timedelta(hours=1.0)
137 if nimrod_field == "Nimrod_comp_2km": 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 radar_offset = timedelta(hours=1.0)
139 date_start_offset = v["data_time"] + radar_offset
141 # Process Nimrod data between the start and end dates.
142 date_use = date_start_offset
143 while date_use < date_start_offset + v["forecast_length"]:
144 # Load the Nimrod data into an Iris cube.
146 nimrod_obs_path = (
147 Path(nimrod_cfg["basedir"])
148 / nimrod_cfg["obs_dir"]
149 / str(date_use.year)
150 / f"{date_use:%Y%m%d%H%M}{nimrod_cfg['obs_fname']}"
151 )
153 nimrod_cube_obs = _load_nimrod(nimrod_obs_path)
155 # Load the Nimrod weights into an Iris cube.
156 nimrod_cube_weights = None
157 if nimrod_cfg["weights_fname"]: 157 ↛ 172line 157 didn't jump to line 172 because the condition on line 157 was always true
158 nimrod_weights_path = (
159 Path(nimrod_cfg["basedir"])
160 / nimrod_cfg["weights_dir"]
161 / str(date_use.year)
162 / f"{date_use:%Y%m%d%H%M}{nimrod_cfg['weights_fname']}"
163 )
165 nimrod_cube_weights = _load_nimrod(nimrod_weights_path)
167 if nimrod_cube_weights is not None: 167 ↛ 172line 167 didn't jump to line 172 because the condition on line 167 was always true
168 # Ensure that the weights are unitless.
169 nimrod_cube_weights.units = "1"
171 # QC the the Nimrod observations using the weights field.
172 if nimrod_cube_obs is not None and nimrod_cube_weights is not None: 172 ↛ 179line 172 didn't jump to line 179 because the condition on line 172 was always true
173 nimrod_cube_obs, nimrod_cube_weights = apply_radar_weights(
174 nimrod_cube_obs, nimrod_cube_weights
175 )
177 # Now that the accumulation weights have been applied if required,
178 # write the Nimrod obs to a NetCDF file.
179 if nimrod_cube_obs is not None: 179 ↛ 189line 179 didn't jump to line 189 because the condition on line 179 was always true
180 filename_obs_nc = (
181 f"{nimrod_dir}/{date_use.strftime('%Y%m%d%H%M')}"
182 f"_{nimrod_dict[nimrod_field]['obs_id']}"
183 )
184 filename_obs_nc = Path(filename_obs_nc).with_suffix(".nc")
185 iris.save(nimrod_cube_obs, filename_obs_nc)
187 # Write the Nimrod weights to a NetCDF file if the switch
188 # for this is set in v["weights"].
189 if v["weights"] == "True" and nimrod_cube_weights is not None: 189 ↛ 199line 189 didn't jump to line 199 because the condition on line 189 was always true
190 filename_wei_nc = (
191 f"{nimrod_dir_wei}/{date_use.strftime('%Y%m%d%H%M')}"
192 f"_{nimrod_dict[nimrod_field]['obs_id']}_weights"
193 )
194 filename_wei_nc = Path(filename_wei_nc).with_suffix(".nc")
195 iris.save(nimrod_cube_weights, filename_wei_nc)
197 # Advance the date/time counter using the time interval
198 # appropriate to the Nimrod field.
199 date_use = date_use + isodate.parse_duration(
200 nimrod_dict[nimrod_field]["freq"]
201 )
204# Run the function that fetches the Nimrod radar obs.
205if __name__ == "__main__": 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 retrieve_nimrod()