Coverage for src/CSET/cset_workflow/app/fetch_nimrod/bin/fetch_nimrod.py: 80%
87 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-24 15:03 +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 """
97 # Grab the environment variables required for handling Nimrod data.
98 v = _get_needed_environment_variables_nimrod()
100 # Grab the Nimrod handling dictionary.
101 with open(nimrod_met_office, "rt") as fp: # pragma: no cover
102 nimrod_dict = json.load(fp)
104 # Loop over the required Nimrod fields, i.e. 1km 2km or xkm rainfall
105 # accumulation composites or the 5 minute rainfall rate composites.
106 for nimrod_field in v["field"]:
107 if nimrod_field: 107 ↛ 106line 107 didn't jump to line 106 because the condition on line 107 was always true
108 logger.info("Processing Nimrod field: %s", nimrod_field)
110 # Prepare the output directory for the Nimrod field.
111 nimrod_dir = f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['obs_id']}"
112 os.makedirs(nimrod_dir, exist_ok=True)
113 logger.info("Cylc-run Nimrod directory: %s", nimrod_dir)
115 # Prepare the output directory for the Nimrod weights field.
116 if nimrod_dict[nimrod_field]["weights_fname"]: 116 ↛ 126line 116 didn't jump to line 126 because the condition on line 116 was always true
117 nimrod_dir_wei = (
118 f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['wei_id']}"
119 )
120 os.makedirs(nimrod_dir_wei, exist_ok=True)
121 logger.info("Cylc-run Nimrod weights directory: %s", nimrod_dir_wei)
123 # Put +1 hour offset for accumulation radar files as
124 # the time stamps for these files mark the end of the
125 # accumulation period rather than the beginning.
126 radar_offset = timedelta(hours=0.0)
127 if nimrod_field == "Nimrod_comp_xkm": 127 ↛ 129line 127 didn't jump to line 129 because the condition on line 127 was always true
128 radar_offset = timedelta(hours=1.0)
129 if nimrod_field == "Nimrod_comp_1km": 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 radar_offset = timedelta(hours=1.0)
131 if nimrod_field == "Nimrod_comp_2km": 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true
132 radar_offset = timedelta(hours=1.0)
133 date_start_offset = v["data_time"] + radar_offset
135 # Process Nimrod data between the start and end dates.
136 date_use = date_start_offset
137 while date_use < date_start_offset + v["forecast_length"]:
138 # Load the Nimrod data into an Iris cube.
139 nimrod_obs_exist = "False"
140 nimrod_obs = (
141 f"{nimrod_dict[nimrod_field]['basedir']}/"
142 f"{nimrod_dict[nimrod_field]['obs_dir']}/"
143 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
144 f"{nimrod_dict[nimrod_field]['obs_fname']}"
145 )
146 try:
147 nimrod_cube_obs = iris.load_cube(nimrod_obs)
148 nimrod_obs_exist = "True"
149 except OSError:
150 logger.warning("Iris cannot find Nimrod file %s", nimrod_obs)
152 # Load the Nimrod weights into an Iris cube.
153 nimrod_weights_exist = "False"
154 if nimrod_dict[nimrod_field]["weights_fname"] != "": 154 ↛ 173line 154 didn't jump to line 173 because the condition on line 154 was always true
155 nimrod_weights = (
156 f"{nimrod_dict[nimrod_field]['basedir']}/"
157 f"{nimrod_dict[nimrod_field]['weights_dir']}/"
158 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
159 f"{nimrod_dict[nimrod_field]['weights_fname']}"
160 )
161 try:
162 nimrod_cube_weights = iris.load_cube(nimrod_weights)
163 nimrod_weights_exist = "True"
164 except OSError:
165 logger.warning(
166 "Iris cannot find Nimrod weights file %s", nimrod_weights
167 )
169 # Ensure that the weights are unitless.
170 nimrod_cube_weights.units = "1"
172 # QC the the Nimrod observations using the weights field.
173 if nimrod_obs_exist == "True" and nimrod_weights_exist == "True": 173 ↛ 180line 173 didn't jump to line 180 because the condition on line 173 was always true
174 nimrod_cube_obs, nimrod_cube_weights = apply_radar_weights(
175 nimrod_cube_obs, nimrod_cube_weights
176 )
178 # Now that the accumulation weights have been applied if required,
179 # write the Nimrod obs to a NetCDF file.
180 if nimrod_obs_exist == "True": 180 ↛ 190line 180 didn't jump to line 190 because the condition on line 180 was always true
181 filename_obs_nc = (
182 f"{nimrod_dir}/{date_use.strftime('%Y%m%d%H%M')}"
183 f"_{nimrod_dict[nimrod_field]['obs_id']}"
184 )
185 filename_obs_nc = Path(filename_obs_nc).with_suffix(".nc")
186 iris.save(nimrod_cube_obs, filename_obs_nc)
188 # Write the Nimrod weights to a NetCDF file if the switch
189 # for this is set in v["weights"].
190 if v["weights"] == "True" and nimrod_weights_exist == "True": 190 ↛ 200line 190 didn't jump to line 200 because the condition on line 190 was always true
191 filename_wei_nc = (
192 f"{nimrod_dir_wei}/{date_use.strftime('%Y%m%d%H%M')}"
193 f"_{nimrod_dict[nimrod_field]['obs_id']}_weights"
194 )
195 filename_wei_nc = Path(filename_wei_nc).with_suffix(".nc")
196 iris.save(nimrod_cube_weights, filename_wei_nc)
198 # Advance the date/time counter using the time interval
199 # appropriate to the Nimrod field.
200 date_use = date_use + isodate.parse_duration(
201 nimrod_dict[nimrod_field]["freq"]
202 )
205# Run the function that fetches the NImrod radar obs.
206if __name__ == "__main__": 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 retrieve_nimrod()