Skip to content

Commit a54a0ba

Browse files
MateusStanoMateusStano
committed
BUG: migrate to THREDDS
1 parent 50e057a commit a54a0ba

3 files changed

Lines changed: 438 additions & 130 deletions

File tree

‎rocketpy/environment/environment.py‎

Lines changed: 187 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
find_latitude_index,
2828
find_longitude_index,
2929
find_time_index,
30+
geodesic_to_lambert_conformal,
3031
geodesic_to_utm,
3132
get_elevation_data_from_dataset,
3233
get_final_date_from_time_array,
@@ -605,6 +606,61 @@ def __set_earth_rotation_vector(self):
605606

606607
# Validators (used to verify an attribute is being set correctly.)
607608

609+
@staticmethod
610+
def __dictionary_matches_dataset(dictionary, dataset):
611+
"""Check whether a mapping dictionary is compatible with a dataset."""
612+
variables = dataset.variables
613+
required_keys = (
614+
"time",
615+
"latitude",
616+
"longitude",
617+
"level",
618+
"temperature",
619+
"u_wind",
620+
"v_wind",
621+
)
622+
623+
for key in required_keys:
624+
variable_name = dictionary.get(key)
625+
if variable_name is None or variable_name not in variables:
626+
return False
627+
628+
projection_name = dictionary.get("projection")
629+
if projection_name is not None and projection_name not in variables:
630+
return False
631+
632+
geopotential_height_name = dictionary.get("geopotential_height")
633+
geopotential_name = dictionary.get("geopotential")
634+
has_geopotential_height = (
635+
geopotential_height_name is not None
636+
and geopotential_height_name in variables
637+
)
638+
has_geopotential = (
639+
geopotential_name is not None and geopotential_name in variables
640+
)
641+
642+
return has_geopotential_height or has_geopotential
643+
644+
def __resolve_dictionary_for_dataset(self, dictionary, dataset):
645+
"""Resolve a compatible mapping dictionary for the loaded dataset.
646+
647+
If the provided mapping is incompatible with the dataset variables,
648+
this method tries built-in mappings and falls back to the first
649+
compatible one.
650+
"""
651+
if self.__dictionary_matches_dataset(dictionary, dataset):
652+
return dictionary
653+
654+
for model_name, candidate in self.__weather_model_map.all_dictionaries.items():
655+
if self.__dictionary_matches_dataset(candidate, dataset):
656+
warnings.warn(
657+
"Provided weather mapping does not match dataset variables. "
658+
f"Falling back to built-in mapping '{model_name}'."
659+
)
660+
return candidate
661+
662+
return dictionary
663+
608664
def __validate_dictionary(self, file, dictionary):
609665
# removed CMC until it is fixed.
610666
available_models = [
@@ -1200,6 +1256,36 @@ def set_atmospheric_model( # pylint: disable=too-many-statements
12001256
case "windy":
12011257
self.process_windy_atmosphere(file)
12021258
case "forecast" | "reanalysis" | "ensemble":
1259+
if isinstance(file, str):
1260+
shortcut_map = self.__atm_type_file_to_function_map.get(type, {})
1261+
matching_shortcut = next(
1262+
(
1263+
shortcut
1264+
for shortcut in shortcut_map
1265+
if shortcut.lower() == file.lower()
1266+
),
1267+
None,
1268+
)
1269+
if matching_shortcut is not None:
1270+
file = matching_shortcut
1271+
1272+
if isinstance(file, str):
1273+
file_upper = file.upper()
1274+
if type == "forecast" and file_upper == "HIRESW":
1275+
raise ValueError(
1276+
"The HIRESW latest-model shortcut is currently "
1277+
"unavailable because NOMADS OPeNDAP is deactivated. "
1278+
"Please use another forecast source or provide a "
1279+
"compatible dataset path/URL explicitly."
1280+
)
1281+
if type == "ensemble" and file_upper == "GEFS":
1282+
raise ValueError(
1283+
"The GEFS latest-model shortcut is currently "
1284+
"unavailable because NOMADS OPeNDAP is deactivated. "
1285+
"Please use another ensemble source or provide a "
1286+
"compatible dataset path/URL explicitly."
1287+
)
1288+
12031289
dictionary = self.__validate_dictionary(file, dictionary)
12041290
try:
12051291
fetch_function = self.__atm_type_file_to_function_map[type][file]
@@ -1661,20 +1747,34 @@ def process_forecast_reanalysis(self, file, dictionary): # pylint: disable=too-
16611747
# Read weather file
16621748
if isinstance(file, str):
16631749
data = netCDF4.Dataset(file)
1664-
if dictionary["time"] not in data.variables.keys():
1665-
dictionary = self.__weather_model_map.get("ECMWF_v0")
16661750
else:
16671751
data = file
16681752

1753+
dictionary = self.__resolve_dictionary_for_dataset(dictionary, data)
1754+
16691755
# Get time, latitude and longitude data from file
16701756
time_array = data.variables[dictionary["time"]]
1671-
lon_list = data.variables[dictionary["longitude"]][:].tolist()
1672-
lat_list = data.variables[dictionary["latitude"]][:].tolist()
1757+
lon_array = data.variables[dictionary["longitude"]]
1758+
lat_array = data.variables[dictionary["latitude"]]
1759+
1760+
# Some THREDDS datasets use projected x/y coordinates.
1761+
if dictionary.get("projection") is not None:
1762+
projection_variable = data.variables[dictionary["projection"]]
1763+
x_units = getattr(lon_array, "units", "m")
1764+
target_lon, target_lat = geodesic_to_lambert_conformal(
1765+
self.latitude,
1766+
self.longitude,
1767+
projection_variable,
1768+
x_units=x_units,
1769+
)
1770+
else:
1771+
target_lon = self.longitude
1772+
target_lat = self.latitude
16731773

16741774
# Find time, latitude and longitude indexes
16751775
time_index = find_time_index(self.datetime_date, time_array)
1676-
lon, lon_index = find_longitude_index(self.longitude, lon_list)
1677-
_, lat_index = find_latitude_index(self.latitude, lat_list)
1776+
lon, lon_index = find_longitude_index(target_lon, lon_array)
1777+
_, lat_index = find_latitude_index(target_lat, lat_array)
16781778

16791779
# Get pressure level data from file
16801780
levels = get_pressure_levels_from_file(data, dictionary)
@@ -1732,9 +1832,9 @@ def process_forecast_reanalysis(self, file, dictionary): # pylint: disable=too-
17321832
) from e
17331833

17341834
# Prepare for bilinear interpolation
1735-
x, y = self.latitude, lon
1736-
x1, y1 = lat_list[lat_index - 1], lon_list[lon_index - 1]
1737-
x2, y2 = lat_list[lat_index], lon_list[lon_index]
1835+
x, y = target_lat, lon
1836+
x1, y1 = float(lat_array[lat_index - 1]), float(lon_array[lon_index - 1])
1837+
x2, y2 = float(lat_array[lat_index]), float(lon_array[lon_index])
17381838

17391839
# Determine properties in lat, lon
17401840
height = bilinear_interpolation(
@@ -1786,6 +1886,17 @@ def process_forecast_reanalysis(self, file, dictionary): # pylint: disable=too-
17861886
wind_vs[:, 1, 1],
17871887
)
17881888

1889+
# Some datasets expose different level counts between fields
1890+
# (e.g., temperature on isobaric1 and geopotential on isobaric).
1891+
min_profile_length = min(
1892+
len(levels), len(height), len(temper), len(wind_u), len(wind_v)
1893+
)
1894+
levels = levels[:min_profile_length]
1895+
height = height[:min_profile_length]
1896+
temper = temper[:min_profile_length]
1897+
wind_u = wind_u[:min_profile_length]
1898+
wind_v = wind_v[:min_profile_length]
1899+
17891900
# Determine wind speed, heading and direction
17901901
wind_speed = calculate_wind_speed(wind_u, wind_v)
17911902
wind_heading = calculate_wind_heading(wind_u, wind_v)
@@ -1843,22 +1954,25 @@ def process_forecast_reanalysis(self, file, dictionary): # pylint: disable=too-
18431954
)
18441955
else:
18451956
self.atmospheric_model_interval = 0
1846-
self.atmospheric_model_init_lat = lat_list[0]
1847-
self.atmospheric_model_end_lat = lat_list[-1]
1848-
self.atmospheric_model_init_lon = lon_list[0]
1849-
self.atmospheric_model_end_lon = lon_list[-1]
1957+
self.atmospheric_model_init_lat = float(lat_array[0])
1958+
self.atmospheric_model_end_lat = float(lat_array[len(lat_array) - 1])
1959+
self.atmospheric_model_init_lon = float(lon_array[0])
1960+
self.atmospheric_model_end_lon = float(lon_array[len(lon_array) - 1])
18501961

18511962
# Save debugging data
1852-
self.lat_array = lat_list
1853-
self.lon_array = lon_list
1963+
self.lat_array = [x1, x2]
1964+
self.lon_array = [y1, y2]
18541965
self.lon_index = lon_index
18551966
self.lat_index = lat_index
18561967
self.geopotentials = geopotentials
18571968
self.wind_us = wind_us
18581969
self.wind_vs = wind_vs
18591970
self.levels = levels
18601971
self.temperatures = temperatures
1861-
self.time_array = time_array[:].tolist()
1972+
self.time_array = [
1973+
float(time_array[0]),
1974+
float(time_array[time_array.shape[0] - 1]),
1975+
]
18621976
self.height = height
18631977

18641978
# Close weather data
@@ -1937,23 +2051,40 @@ def process_ensemble(self, file, dictionary): # pylint: disable=too-many-locals
19372051
else:
19382052
data = file
19392053

2054+
dictionary = self.__resolve_dictionary_for_dataset(dictionary, data)
2055+
19402056
# Get time, latitude and longitude data from file
19412057
time_array = data.variables[dictionary["time"]]
1942-
lon_list = data.variables[dictionary["longitude"]][:].tolist()
1943-
lat_list = data.variables[dictionary["latitude"]][:].tolist()
2058+
lon_array = data.variables[dictionary["longitude"]]
2059+
lat_array = data.variables[dictionary["latitude"]]
2060+
2061+
# Some THREDDS datasets use projected x/y coordinates.
2062+
#TODO CHECK THIS I AM NOT SURE?????
2063+
if dictionary.get("projection") is not None:
2064+
projection_variable = data.variables[dictionary["projection"]]
2065+
x_units = getattr(lon_array, "units", "m")
2066+
target_lon, target_lat = geodesic_to_lambert_conformal(
2067+
self.latitude,
2068+
self.longitude,
2069+
projection_variable,
2070+
x_units=x_units,
2071+
)
2072+
else:
2073+
target_lon = self.longitude
2074+
target_lat = self.latitude
19442075

19452076
# Find time, latitude and longitude indexes
19462077
time_index = find_time_index(self.datetime_date, time_array)
1947-
lon, lon_index = find_longitude_index(self.longitude, lon_list)
1948-
_, lat_index = find_latitude_index(self.latitude, lat_list)
2078+
lon, lon_index = find_longitude_index(target_lon, lon_array)
2079+
_, lat_index = find_latitude_index(target_lat, lat_array)
19492080

19502081
# Get ensemble data from file
2082+
has_ensemble_dimension = True
19512083
try:
19522084
num_members = len(data.variables[dictionary["ensemble"]][:])
1953-
except KeyError as e:
1954-
raise ValueError(
1955-
"Unable to read ensemble data from file. Check file and dictionary."
1956-
) from e
2085+
except KeyError:
2086+
has_ensemble_dimension = False
2087+
num_members = 1
19572088

19582089
# Get pressure level data from file
19592090
levels = get_pressure_levels_from_file(data, dictionary)
@@ -2012,10 +2143,16 @@ def process_ensemble(self, file, dictionary): # pylint: disable=too-many-locals
20122143
"Unable to read wind-v component. Check file and dictionary."
20132144
) from e
20142145

2146+
if not has_ensemble_dimension:
2147+
geopotentials = np.expand_dims(geopotentials, axis=0)
2148+
temperatures = np.expand_dims(temperatures, axis=0)
2149+
wind_us = np.expand_dims(wind_us, axis=0)
2150+
wind_vs = np.expand_dims(wind_vs, axis=0)
2151+
20152152
# Prepare for bilinear interpolation
2016-
x, y = self.latitude, lon
2017-
x1, y1 = lat_list[lat_index - 1], lon_list[lon_index - 1]
2018-
x2, y2 = lat_list[lat_index], lon_list[lon_index]
2153+
x, y = target_lat, lon
2154+
x1, y1 = float(lat_array[lat_index - 1]), float(lon_array[lon_index - 1])
2155+
x2, y2 = float(lat_array[lat_index]), float(lon_array[lon_index])
20192156

20202157
# Determine properties in lat, lon
20212158
height = bilinear_interpolation(
@@ -2067,6 +2204,19 @@ def process_ensemble(self, file, dictionary): # pylint: disable=too-many-locals
20672204
wind_vs[:, :, 1, 1],
20682205
)
20692206

2207+
min_profile_length = min(
2208+
len(levels),
2209+
height.shape[1],
2210+
temper.shape[1],
2211+
wind_u.shape[1],
2212+
wind_v.shape[1],
2213+
)
2214+
levels = levels[:min_profile_length]
2215+
height = height[:, :min_profile_length]
2216+
temper = temper[:, :min_profile_length]
2217+
wind_u = wind_u[:, :min_profile_length]
2218+
wind_v = wind_v[:, :min_profile_length]
2219+
20702220
# Determine wind speed, heading and direction
20712221
wind_speed = calculate_wind_speed(wind_u, wind_v)
20722222
wind_heading = calculate_wind_heading(wind_u, wind_v)
@@ -2099,22 +2249,25 @@ def process_ensemble(self, file, dictionary): # pylint: disable=too-many-locals
20992249
self.atmospheric_model_init_date = get_initial_date_from_time_array(time_array)
21002250
self.atmospheric_model_end_date = get_final_date_from_time_array(time_array)
21012251
self.atmospheric_model_interval = get_interval_date_from_time_array(time_array)
2102-
self.atmospheric_model_init_lat = lat_list[0]
2103-
self.atmospheric_model_end_lat = lat_list[-1]
2104-
self.atmospheric_model_init_lon = lon_list[0]
2105-
self.atmospheric_model_end_lon = lon_list[-1]
2252+
self.atmospheric_model_init_lat = float(lat_array[0])
2253+
self.atmospheric_model_end_lat = float(lat_array[len(lat_array) - 1])
2254+
self.atmospheric_model_init_lon = float(lon_array[0])
2255+
self.atmospheric_model_end_lon = float(lon_array[len(lon_array) - 1])
21062256

21072257
# Save debugging data
2108-
self.lat_array = lat_list
2109-
self.lon_array = lon_list
2258+
self.lat_array = [x1, x2]
2259+
self.lon_array = [y1, y2]
21102260
self.lon_index = lon_index
21112261
self.lat_index = lat_index
21122262
self.geopotentials = geopotentials
21132263
self.wind_us = wind_us
21142264
self.wind_vs = wind_vs
21152265
self.levels = levels
21162266
self.temperatures = temperatures
2117-
self.time_array = time_array[:].tolist()
2267+
self.time_array = [
2268+
float(time_array[0]),
2269+
float(time_array[time_array.shape[0] - 1]),
2270+
]
21182271
self.height = height
21192272

21202273
# Close weather data

0 commit comments

Comments
 (0)