Coverage for src/CSET/cset_workflow/app/fetch_nimrod/bin/fetch_nimrod.py: 82%
80 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-14 13:50 +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
15iris.FUTURE.save_split_attrs = True
16iris.FUTURE.date_microseconds = True
18# Define the file to use when retrieving the sensitive Nimrod handling data.
19nimrod_met_office = "restricted_nimrod_met_office.json"
22def _get_needed_environment_variables_nimrod() -> dict:
23 """Load the needed variables from the environment to retrieve UK Nimrod data."""
24 radar_sources = []
25 if os.environ["NIMROD_COMP_XKM"] == "True": 25 ↛ 27line 25 didn't jump to line 27 because the condition on line 25 was always true
26 radar_sources.append("Nimrod_comp_xkm")
27 if os.environ["NIMROD_COMP_1KM"] == "True": 27 ↛ 28line 27 didn't jump to line 28 because the condition on line 27 was never true
28 radar_sources.append("Nimrod_comp_1km")
29 if os.environ["NIMROD_COMP_2KM"] == "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_2km")
31 if os.environ["NIMROD_COMP_5MIN"] == "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_5min")
33 variables = {
34 "field": radar_sources,
35 "weights": os.environ["NIMROD_WEIGHTS"],
36 "date_type": "initiation",
37 "data_time": datetime.fromisoformat(os.environ["CYLC_TASK_CYCLE_POINT"]),
38 "forecast_length": isodate.parse_duration(os.environ["ANALYSIS_LENGTH"]),
39 "rose_datac": os.environ["ROSE_DATAC"],
40 }
41 logging.debug("Environment variables loaded for Nimrod: %s", variables)
42 return variables
45def apply_radar_weights(cube_obs: iris.cube.Cube, cube_wei: iris.cube.Cube):
46 """Apply the Nimrod weights to the radar hourly rainfall accumulation data.
48 Parameters
49 ----------
50 cube_obs: Cube or
51 2 dimensional Cube of the radar rainfall accumulation data.
53 cube_wei: Cube or
54 2 dimensional Cube of the Nimrod rainfall accumulation weights.
56 """
57 # The weights are the number of 5 minute rainfall rates used to
58 # compute the hourly rainfall accumulation. The weight should be
59 # in the range 0 to 13.
60 # Define the minimum weight to accept as good data.
61 weight_min = 11
63 # Define the value to use for duff radar data
64 duff_value = 0.0
66 # Check if the weights are packed as weights / 32
67 # and if they are, unpack them.
68 # Note: if the weights are packed, then the maximum value
69 # found will be 13 / 32 = 0.40625, i.e. all values < 1.
70 weights = cube_wei.data
71 if weights.max() < 1.0:
72 weights = (weights * 32).round().astype(int)
73 cube_wei.data = weights
74 logging.info("Unpacked Nimrod weights file.")
76 # Apply the weights.
77 cube_obs_weighted = cube_obs
78 cube_obs_weighted.data = np.where(
79 weights < weight_min, duff_value, cube_obs_weighted.data
80 )
82 # Return the QC'd Nimrod data and the weights used.
83 return cube_obs_weighted, cube_wei
86def retrieve_nimrod():
87 """Fetch the observations corresponding to a model run.
89 The following environment variables need to be set:
90 * NIMROD_COMP_XKM, NIMROD_COMP_1KM, NIMROD_COMP_2KM, NIMROD_COMP_5MIN
91 * CYLC_TASK_CYCLE_POINT
92 * ANALYSIS_LENGTH
93 * ROSE_DATAC
94 """
95 # Grab the environment variables required for handling Nimrod data.
96 v = _get_needed_environment_variables_nimrod()
98 # Grab the Nimrod handling dictionary.
99 with open(nimrod_met_office, "rt") as fp: # pragma: no cover
100 nimrod_dict = json.load(fp)
102 # Form the Nimrod start and end dates.
103 date_start = v["data_time"]
104 date_end = v["data_time"] + v["forecast_length"]
106 # Loop over the required Nimrod fields, i.e. 1km 2km or xkm rainfall
107 # accumulation composites or the 5 minute rainfall rate composites.
108 for nimrod_field in v["field"]:
109 if nimrod_field: 109 ↛ 108line 109 didn't jump to line 108 because the condition on line 109 was always true
110 logging.info("Processing Nimrod field: %s", nimrod_field)
112 # Prepare the output directory for the Nimrod field.
113 nimrod_dir = f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['obs_id']}"
114 os.makedirs(nimrod_dir, exist_ok=True)
115 logging.info("Cylc-run Nimrod directory: %s", nimrod_dir)
117 # Prepare the output directory for the Nimrod weights field.
118 if nimrod_dict[nimrod_field]["weights_fname"]: 118 ↛ 126line 118 didn't jump to line 126 because the condition on line 118 was always true
119 nimrod_dir_wei = (
120 f"{v['rose_datac']}/data/{nimrod_dict[nimrod_field]['wei_id']}"
121 )
122 os.makedirs(nimrod_dir_wei, exist_ok=True)
123 logging.info("Cylc-run Nimrod weights directory: %s", nimrod_dir_wei)
125 # Process Nimrod data between the start and end dates.
126 date_use = date_start
127 while date_use <= date_end:
128 # Load the Nimrod data into an Iris cube.
129 nimrod_obs_exist = "False"
130 nimrod_obs = (
131 f"{nimrod_dict[nimrod_field]['basedir']}/"
132 f"{nimrod_dict[nimrod_field]['obs_dir']}/"
133 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
134 f"{nimrod_dict[nimrod_field]['obs_fname']}"
135 )
136 try:
137 nimrod_cube_obs = iris.load_cube(nimrod_obs)
138 nimrod_obs_exist = "True"
139 except OSError:
140 logging.warning("Iris cannot find Nimrod file %s", nimrod_obs)
142 # Load the Nimrod weights into an Iris cube.
143 nimrod_weights_exist = "False"
144 if nimrod_dict[nimrod_field]["weights_fname"] != "": 144 ↛ 163line 144 didn't jump to line 163 because the condition on line 144 was always true
145 nimrod_weights = (
146 f"{nimrod_dict[nimrod_field]['basedir']}/"
147 f"{nimrod_dict[nimrod_field]['weights_dir']}/"
148 f"{date_use.year}/{date_use.strftime('%Y%m%d%H%M')}"
149 f"{nimrod_dict[nimrod_field]['weights_fname']}"
150 )
151 try:
152 nimrod_cube_weights = iris.load_cube(nimrod_weights)
153 nimrod_weights_exist = "True"
154 except OSError:
155 logging.warning(
156 "Iris cannot find Nimrod weights file %s", nimrod_weights
157 )
159 # Ensure that the weights are unitless.
160 nimrod_cube_weights.units = "1"
162 # QC the the Nimrod observations using the weights field.
163 if nimrod_obs_exist == "True" and nimrod_weights_exist == "True": 163 ↛ 170line 163 didn't jump to line 170 because the condition on line 163 was always true
164 nimrod_cube_obs, nimrod_cube_weights = apply_radar_weights(
165 nimrod_cube_obs, nimrod_cube_weights
166 )
168 # Now that the accumulation weights have been applied if required,
169 # write the Nimrod obs to a NetCDF file.
170 if nimrod_obs_exist == "True": 170 ↛ 180line 170 didn't jump to line 180 because the condition on line 170 was always true
171 filename_obs_nc = (
172 f"{nimrod_dir}/{date_use.strftime('%Y%m%d%H%M')}"
173 f"_{nimrod_dict[nimrod_field]['obs_id']}"
174 )
175 filename_obs_nc = Path(filename_obs_nc).with_suffix(".nc")
176 iris.save(nimrod_cube_obs, filename_obs_nc)
178 # Write the Nimrod weights to a NetCDF file if the switch
179 # for this is set in v["weights"].
180 if v["weights"] == "True" and nimrod_weights_exist == "True": 180 ↛ 190line 180 didn't jump to line 190 because the condition on line 180 was always true
181 filename_wei_nc = (
182 f"{nimrod_dir_wei}/{date_use.strftime('%Y%m%d%H%M')}"
183 f"_{nimrod_dict[nimrod_field]['obs_id']}_weights"
184 )
185 filename_wei_nc = Path(filename_wei_nc).with_suffix(".nc")
186 iris.save(nimrod_cube_weights, filename_wei_nc)
188 # Advance the date/time counter using the time interval
189 # appropriate to the Nimrod field.
190 date_use = date_use + isodate.parse_duration(
191 nimrod_dict[nimrod_field]["freq"]
192 )
195# Run the function that fetches the NImrod radar obs.
196if __name__ == "__main__": 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true
197 retrieve_nimrod()