Coverage for src/CSET/cset_workflow/app/fetch_nimrod/bin/fetch_nimrod.py: 83%
81 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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
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 # Form the Nimrod start and end dates.
105 date_start = v["data_time"]
106 date_end = v["data_time"] + v["forecast_length"]
108 # Loop over the required Nimrod fields, i.e. 1km 2km or xkm rainfall
109 # accumulation composites or the 5 minute rainfall rate composites.
110 for nimrod_field in v["field"]:
111 if nimrod_field: 111 ↛ 110line 111 didn't jump to line 110 because the condition on line 111 was always true
112 logger.info("Processing Nimrod field: %s", nimrod_field)
114 # Prepare the output directory for the Nimrod field.
115 nimrod_dir = f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['obs_id']}"
116 os.makedirs(nimrod_dir, exist_ok=True)
117 logger.info("Cylc-run Nimrod directory: %s", nimrod_dir)
119 # Prepare the output directory for the Nimrod weights field.
120 if nimrod_dict[nimrod_field]["weights_fname"]: 120 ↛ 128line 120 didn't jump to line 128 because the condition on line 120 was always true
121 nimrod_dir_wei = (
122 f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['wei_id']}"
123 )
124 os.makedirs(nimrod_dir_wei, exist_ok=True)
125 logger.info("Cylc-run Nimrod weights directory: %s", nimrod_dir_wei)
127 # Process Nimrod data between the start and end dates.
128 date_use = date_start
129 while date_use <= date_end:
130 # Load the Nimrod data into an Iris cube.
131 nimrod_obs_exist = "False"
132 nimrod_obs = (
133 f"{nimrod_dict[nimrod_field]['basedir']}/"
134 f"{nimrod_dict[nimrod_field]['obs_dir']}/"
135 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
136 f"{nimrod_dict[nimrod_field]['obs_fname']}"
137 )
138 try:
139 nimrod_cube_obs = iris.load_cube(nimrod_obs)
140 nimrod_obs_exist = "True"
141 except OSError:
142 logger.warning("Iris cannot find Nimrod file %s", nimrod_obs)
144 # Load the Nimrod weights into an Iris cube.
145 nimrod_weights_exist = "False"
146 if nimrod_dict[nimrod_field]["weights_fname"] != "": 146 ↛ 165line 146 didn't jump to line 165 because the condition on line 146 was always true
147 nimrod_weights = (
148 f"{nimrod_dict[nimrod_field]['basedir']}/"
149 f"{nimrod_dict[nimrod_field]['weights_dir']}/"
150 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
151 f"{nimrod_dict[nimrod_field]['weights_fname']}"
152 )
153 try:
154 nimrod_cube_weights = iris.load_cube(nimrod_weights)
155 nimrod_weights_exist = "True"
156 except OSError:
157 logger.warning(
158 "Iris cannot find Nimrod weights file %s", nimrod_weights
159 )
161 # Ensure that the weights are unitless.
162 nimrod_cube_weights.units = "1"
164 # QC the the Nimrod observations using the weights field.
165 if nimrod_obs_exist == "True" and nimrod_weights_exist == "True": 165 ↛ 172line 165 didn't jump to line 172 because the condition on line 165 was always true
166 nimrod_cube_obs, nimrod_cube_weights = apply_radar_weights(
167 nimrod_cube_obs, nimrod_cube_weights
168 )
170 # Now that the accumulation weights have been applied if required,
171 # write the Nimrod obs to a NetCDF file.
172 if nimrod_obs_exist == "True": 172 ↛ 182line 172 didn't jump to line 182 because the condition on line 172 was always true
173 filename_obs_nc = (
174 f"{nimrod_dir}/{date_use.strftime('%Y%m%d%H%M')}"
175 f"_{nimrod_dict[nimrod_field]['obs_id']}"
176 )
177 filename_obs_nc = Path(filename_obs_nc).with_suffix(".nc")
178 iris.save(nimrod_cube_obs, filename_obs_nc)
180 # Write the Nimrod weights to a NetCDF file if the switch
181 # for this is set in v["weights"].
182 if v["weights"] == "True" and nimrod_weights_exist == "True": 182 ↛ 192line 182 didn't jump to line 192 because the condition on line 182 was always true
183 filename_wei_nc = (
184 f"{nimrod_dir_wei}/{date_use.strftime('%Y%m%d%H%M')}"
185 f"_{nimrod_dict[nimrod_field]['obs_id']}_weights"
186 )
187 filename_wei_nc = Path(filename_wei_nc).with_suffix(".nc")
188 iris.save(nimrod_cube_weights, filename_wei_nc)
190 # Advance the date/time counter using the time interval
191 # appropriate to the Nimrod field.
192 date_use = date_use + isodate.parse_duration(
193 nimrod_dict[nimrod_field]["freq"]
194 )
197# Run the function that fetches the NImrod radar obs.
198if __name__ == "__main__": 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true
199 retrieve_nimrod()