From 1327fca5d7d6fb0dddd707699faa32cb2ac12cb2 Mon Sep 17 00:00:00 2001 From: Ting-Hong Shieh <32212900+ting-hong-shieh@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:39:24 +0800 Subject: [PATCH 1/3] ENH: create ensembles from user-defined profiles --- .../environment/1-atm-models/ensemble.rst | 74 +++- rocketpy/environment/environment.py | 343 ++++++++++++++++++ rocketpy/environment/tools.py | 6 +- tests/unit/environment/test_environment.py | 126 +++++++ 4 files changed, 547 insertions(+), 2 deletions(-) diff --git a/docs/user/environment/1-atm-models/ensemble.rst b/docs/user/environment/1-atm-models/ensemble.rst index a2c75b118..b7cdfd6be 100644 --- a/docs/user/environment/1-atm-models/ensemble.rst +++ b/docs/user/environment/1-atm-models/ensemble.rst @@ -20,6 +20,78 @@ forecast and obtain a range of possible outcomes. Ensemble Forecast ----------------- +Creating a Custom Ensemble +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Use :meth:`rocketpy.Environment.create_ensemble` to combine two or more +atmospheric profiles for one location. Each member is a mapping with +``pressure``, ``temperature``, ``wind_u`` and ``wind_v`` profiles. Each profile +is a two-column array: the first column is geometric height above sea level in +meters, and the second column uses Pa for pressure, K for temperature and m/s +for either wind component. + +The pressure profiles determine a common isobaric grid. Temperature and wind +are interpolated onto that grid, then written to a GEFS-compatible NetCDF file. +Every pressure profile must overlap the others and decrease with height. + +.. code-block:: python + + import numpy as np + + from rocketpy import Environment + + env = Environment( + date=(2026, 9, 1, 12), + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + ) + + pressure = np.array([85000, 70000, 50000]) # Pa + height_0 = np.array([1500, 3000, 5500]) # m ASL + height_1 = np.array([1550, 3100, 5650]) # m ASL + + profiles = [ + { + "pressure": np.column_stack((height_0, pressure)), + "temperature": np.column_stack((height_0, [278, 268, 250])), + "wind_u": np.column_stack((height_0, [2, 5, 9])), + "wind_v": np.column_stack((height_0, [-1, 1, 4])), + }, + { + "pressure": np.column_stack((height_1, pressure)), + "temperature": np.column_stack((height_1, [280, 269, 251])), + "wind_u": np.column_stack((height_1, [4, 7, 12])), + "wind_v": np.column_stack((height_1, [0, 2, 6])), + }, + ] + + ensemble_file = env.create_ensemble(profiles, "my_ensemble.nc") + +The method activates member 0 after writing the file. Select another member +with the same method used for forecast ensembles: + +.. code-block:: python + + env.select_ensemble_member(1) + +Another Environment can load the returned file with the ``GEFS`` mapping: + +.. code-block:: python + + saved_env = Environment( + date=(2026, 9, 1, 12), + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + ) + saved_env.set_atmospheric_model( + type="Ensemble", + file=ensemble_file, + dictionary="GEFS", + ) + + Global Ensemble Forecast System (GEFS) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -106,4 +178,4 @@ Ensemble Reanalysis ------------------- Ensemble reanalyses are also possible with RocketPy. See the -:ref:`reanalysis_ensemble` section for more information. \ No newline at end of file +:ref:`reanalysis_ensemble` section for more information. diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index edf3a342c..cf523cbd7 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -5,6 +5,7 @@ import os import re import warnings +from collections.abc import Mapping from collections import namedtuple from datetime import datetime @@ -2803,6 +2804,348 @@ def process_forecast_reanalysis(self, file, dictionary, conversion_factor): # p # Close weather data data.close() + @staticmethod + def _prepare_ensemble_profile_source(source, variable, member): + """Validate and normalize a user-defined atmospheric profile.""" + if isinstance(source, Function): + if not source.is_array_source(): + raise TypeError( + f"Member {member} '{variable}' must be an array-backed " + "Function or a two-column array." + ) + source = source.source + + try: + profile = np.asarray(source, dtype=float) + except (TypeError, ValueError) as exc: + raise TypeError( + f"Member {member} '{variable}' must be a two-column numeric array." + ) from exc + + if profile.ndim != 2 or profile.shape[1] != 2 or len(profile) < 2: + raise ValueError( + f"Member {member} '{variable}' must contain at least two " + "[height, value] rows." + ) + if not np.all(np.isfinite(profile)): + raise ValueError( + f"Member {member} '{variable}' contains non-finite values." + ) + + profile = profile[np.argsort(profile[:, 0])] + if np.any(np.diff(profile[:, 0]) <= 0): + raise ValueError(f"Member {member} '{variable}' heights must be unique.") + return profile + + def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-many-statements + self, + profiles, + file_name="custom_ensemble.nc", + pressure_levels=None, + overwrite=False, + ): + """Create and activate an ensemble from user-defined profiles. + + RocketPy writes the profiles with GEFS-compatible variable names. + Another Environment can load the returned file by passing + ``type="Ensemble"`` and ``dictionary="GEFS"`` to + :meth:`Environment.set_atmospheric_model`. + + Parameters + ---------- + profiles : sequence of mappings + Atmospheric profiles for each ensemble member. Every mapping must + define ``pressure``, ``temperature``, ``wind_u`` and ``wind_v``. + Each value must be a two-column array whose first column is + geometric height above sea level in meters. The second column uses + Pa for pressure, K for temperature and m/s for either wind + component. Array-backed :class:`rocketpy.Function` objects are also + accepted. Pressure must decrease strictly with increasing height. + file_name : str or os.PathLike, optional + Path of the NetCDF file to create. The ``.nc`` suffix is appended + when omitted. Default is ``"custom_ensemble.nc"``. + pressure_levels : array-like, optional + Common pressure levels in Pa. By default, the union of sampled + pressure levels inside the range shared by every member is used. + overwrite : bool, optional + Whether an existing file may be replaced. Default is ``False``. + + Returns + ------- + str + Absolute path of the created NetCDF file. + + Raises + ------ + TypeError + If profiles or profile values have invalid types. + ValueError + If fewer than two members are supplied, required variables are + missing, profiles are invalid, or the members have no usable common + pressure range. + FileExistsError + If the output exists and ``overwrite`` is ``False``. + + Notes + ----- + The first member is activated after the file is created. Use + :meth:`Environment.select_ensemble_member` to activate another member. + """ + self.__validate_datetime() + + if isinstance(profiles, (str, bytes, Mapping)): + raise TypeError("'profiles' must be a sequence of member mappings.") + try: + profiles = list(profiles) + except TypeError as exc: + raise TypeError( + "'profiles' must be a sequence of member mappings." + ) from exc + if len(profiles) < 2: + raise ValueError("At least two atmospheric profiles are required.") + + required_variables = ("pressure", "temperature", "wind_u", "wind_v") + members = [] + for member_index, member in enumerate(profiles): + if not isinstance(member, Mapping): + raise TypeError( + f"Member {member_index} must be a mapping of profile names " + "to two-column arrays." + ) + missing = [name for name in required_variables if name not in member] + if missing: + raise ValueError( + f"Member {member_index} is missing required profile(s): " + f"{', '.join(missing)}." + ) + + prepared = { + variable: self._prepare_ensemble_profile_source( + member[variable], variable, member_index + ) + for variable in required_variables + } + pressure = prepared["pressure"][:, 1] + if np.any(pressure <= 0): + raise ValueError( + f"Member {member_index} pressure values must be positive." + ) + if np.any(np.diff(pressure) >= 0): + raise ValueError( + f"Member {member_index} pressure must decrease strictly " + "with increasing height." + ) + if np.any(prepared["temperature"][:, 1] <= 0): + raise ValueError( + f"Member {member_index} temperature values must be positive." + ) + members.append(prepared) + + common_min_pressure = max(member["pressure"][-1, 1] for member in members) + common_max_pressure = min(member["pressure"][0, 1] for member in members) + if common_min_pressure >= common_max_pressure: + raise ValueError("Ensemble members have no common pressure range.") + + if pressure_levels is None: + common_levels = np.concatenate( + [member["pressure"][:, 1] for member in members] + ) + common_levels = common_levels[ + (common_levels >= common_min_pressure) + & (common_levels <= common_max_pressure) + ] + pressure_levels = np.unique(common_levels)[::-1] + else: + try: + pressure_levels = np.asarray(pressure_levels, dtype=float) + except (TypeError, ValueError) as exc: + raise TypeError("'pressure_levels' must be a numeric array.") from exc + if pressure_levels.ndim != 1: + raise ValueError("'pressure_levels' must be one-dimensional.") + if not np.all(np.isfinite(pressure_levels)) or np.any(pressure_levels <= 0): + raise ValueError( + "'pressure_levels' must contain only finite, positive values." + ) + if len(np.unique(pressure_levels)) != len(pressure_levels): + raise ValueError("'pressure_levels' must not contain duplicates.") + pressure_levels = np.sort(pressure_levels)[::-1] + + if len(pressure_levels) < 2: + raise ValueError( + "At least two pressure levels inside the common range are required." + ) + if ( + pressure_levels[-1] < common_min_pressure + or pressure_levels[0] > common_max_pressure + ): + raise ValueError( + "'pressure_levels' must stay inside the pressure range shared " + "by every member." + ) + + member_heights = [] + member_values = {name: [] for name in required_variables[1:]} + for member_index, member in enumerate(members): + pressure_profile = member["pressure"] + heights = np.interp( + pressure_levels, + pressure_profile[::-1, 1], + pressure_profile[::-1, 0], + ) + member_heights.append(heights) + + for variable in required_variables[1:]: + profile = member[variable] + if heights[0] < profile[0, 0] or heights[-1] > profile[-1, 0]: + raise ValueError( + f"Member {member_index} '{variable}' does not cover all " + "heights in the common pressure range." + ) + member_values[variable].append( + np.interp(heights, profile[:, 0], profile[:, 1]) + ) + + geometric_heights = np.asarray(member_heights) + if np.any(geometric_heights <= -self.earth_radius): + raise ValueError("Profile heights must be greater than -Earth's radius.") + geopotential_heights = ( + self.earth_radius + * geometric_heights + / (self.earth_radius + geometric_heights) + ) + + try: + file_path = os.fspath(file_name) + except TypeError as exc: + raise TypeError( + "'file_name' must be a string or path-like object." + ) from exc + if not file_path.lower().endswith(".nc"): + file_path += ".nc" + file_path = os.path.abspath(file_path) + if os.path.exists(file_path) and not overwrite: + raise FileExistsError( + f"'{file_path}' already exists. Pass overwrite=True to replace it." + ) + + latitude_bounds = np.array( + [max(-90, self.latitude - 0.01), min(90, self.latitude + 0.01)] + ) + grid_longitude = 0 if self.longitude == 360 else self.longitude + longitude_bounds = np.array( + [ + max(-180, grid_longitude - 0.01), + min(360, grid_longitude + 0.01), + ] + ) + data_shape = ( + 1, + len(members), + len(pressure_levels), + len(latitude_bounds), + len(longitude_bounds), + ) + + with netCDF4.Dataset(file_path, "w", format="NETCDF4") as dataset: + dataset.Conventions = "CF-1.8" + dataset.title = "RocketPy user-defined atmospheric ensemble" + dataset.source = "RocketPy Environment.create_ensemble" + dataset.history = ( + f"Created {datetime.now(tz=pytz.UTC).isoformat()} by RocketPy" + ) + dataset.comment = ( + "Profiles are spatially constant across the 2 x 2 grid " + "surrounding the launch coordinates." + ) + dataset.launch_latitude = self.latitude + dataset.launch_longitude = self.longitude + + dataset.createDimension("time", 1) + dataset.createDimension("ens", len(members)) + dataset.createDimension("lev", len(pressure_levels)) + dataset.createDimension("lat", len(latitude_bounds)) + dataset.createDimension("lon", len(longitude_bounds)) + + time = dataset.createVariable("time", "f8", ("time",)) + time.long_name = "profile valid time" + time.standard_name = "time" + time.units = ( + f"hours since {self.datetime_date.strftime('%Y-%m-%d %H:%M:%S')} UTC" + ) + time.calendar = "gregorian" + time.axis = "T" + time[:] = [0] + + ensemble = dataset.createVariable("ens", "i4", ("ens",)) + ensemble.long_name = "ensemble member" + ensemble.units = "1" + ensemble[:] = np.arange(len(members)) + + level = dataset.createVariable("lev", "f8", ("lev",)) + level.long_name = "pressure level" + level.standard_name = "air_pressure" + level.units = "hPa" + level.positive = "down" + level.axis = "Z" + level[:] = pressure_levels / 100 + + latitude = dataset.createVariable("lat", "f8", ("lat",)) + latitude.long_name = "latitude" + latitude.standard_name = "latitude" + latitude.units = "degrees_north" + latitude.axis = "Y" + latitude[:] = latitude_bounds + + longitude = dataset.createVariable("lon", "f8", ("lon",)) + longitude.long_name = "longitude" + longitude.standard_name = "longitude" + longitude.units = "degrees_east" + longitude.axis = "X" + longitude[:] = longitude_bounds + + dimensions = ("time", "ens", "lev", "lat", "lon") + variables = { + "hgtprs": ( + geopotential_heights, + "geopotential height", + "geopotential_height", + "m", + ), + "tmpprs": ( + np.asarray(member_values["temperature"]), + "air temperature", + "air_temperature", + "K", + ), + "ugrdprs": ( + np.asarray(member_values["wind_u"]), + "eastward wind", + "eastward_wind", + "m s-1", + ), + "vgrdprs": ( + np.asarray(member_values["wind_v"]), + "northward wind", + "northward_wind", + "m s-1", + ), + } + for name, (values, long_name, standard_name, units) in variables.items(): + variable = dataset.createVariable( + name, "f8", dimensions, zlib=True, complevel=4 + ) + variable.long_name = long_name + variable.standard_name = standard_name + variable.units = units + variable.coordinates = "time ens lev lat lon" + variable[:] = np.broadcast_to( + values[None, :, :, None, None], data_shape + ) + + self.set_atmospheric_model(type="Ensemble", file=file_path, dictionary="GEFS") + logger.info("Atmospheric ensemble saved at '%s'.", file_path) + return file_path + def process_ensemble(self, file, dictionary, conversion_factor): # pylint: disable=too-many-locals,too-many-statements """Import and process atmospheric data from weather ensembles given as ``netCDF`` or ``OPeNDAP`` files. Sets pressure, temperature, diff --git a/rocketpy/environment/tools.py b/rocketpy/environment/tools.py index cb0f4d5ad..c681d4d39 100644 --- a/rocketpy/environment/tools.py +++ b/rocketpy/environment/tools.py @@ -730,8 +730,12 @@ def get_interval_date_from_time_array(time_array, units=None): Returns ------- int - The interval in hours between two times in the time array. + The interval in hours between times in the array, or 0 when the array + contains a single time. """ + if len(time_array) < 2: + return 0 + units = units or time_array.units return netCDF4.num2date( (time_array[-1] - time_array[0]) / (len(time_array) - 1), diff --git a/tests/unit/environment/test_environment.py b/tests/unit/environment/test_environment.py index bee3decf1..4db6a13f2 100644 --- a/tests/unit/environment/test_environment.py +++ b/tests/unit/environment/test_environment.py @@ -4,6 +4,7 @@ import numpy as np import numpy.testing as npt +import netCDF4 import pytest import pytz @@ -14,6 +15,7 @@ geodesic_to_utm, get_final_date_from_time_array, get_initial_date_from_time_array, + get_interval_date_from_time_array, get_pressure_levels_from_file, pressure_unit_to_factor, utm_to_geodesic, @@ -22,6 +24,130 @@ from rocketpy.tools import geopotential_height_to_geometric_height +def _user_defined_ensemble_profiles(): + """Return two members with a shared isobaric grid.""" + pressure = np.array([101325.0, 90000.0, 80000.0]) + member_0_height = np.array([0.0, 1000.0, 2000.0]) + member_1_height = np.array([100.0, 1100.0, 2100.0]) + return [ + { + "pressure": np.column_stack((member_0_height, pressure)), + "temperature": np.column_stack((member_0_height, [288.0, 281.0, 275.0])), + "wind_u": np.column_stack((member_0_height, [1.0, 2.0, 3.0])), + "wind_v": np.column_stack((member_0_height, [-1.0, -2.0, -3.0])), + }, + { + "pressure": np.column_stack((member_1_height, pressure)), + "temperature": np.column_stack((member_1_height, [290.0, 283.0, 277.0])), + "wind_u": np.column_stack((member_1_height, [4.0, 5.0, 6.0])), + "wind_v": np.column_stack((member_1_height, [-4.0, -5.0, -6.0])), + }, + ] + + +def test_time_array_interval_helper_accepts_a_single_time(): + """A static user ensemble has no forecast interval.""" + + class SingleTimeArray: + """Minimal single-value NetCDF-like time coordinate.""" + + units = "hours since 2025-06-01 12:00:00" + + def __len__(self): + return 1 + + assert get_interval_date_from_time_array(SingleTimeArray()) == 0 + + +def test_create_ensemble_exports_and_activates_profiles(tmp_path): + """Export user profiles and expose each member through Environment.""" + # Arrange + env = Environment( + date=(2025, 6, 1, 12), + latitude=32.99, + longitude=-106.97, + elevation=0, + ) + output = tmp_path / "test_ensemble" + + # Act + file_path = env.create_ensemble(_user_defined_ensemble_profiles(), file_name=output) + + # Assert + assert file_path == str(output) + ".nc" + assert env.atmospheric_model_type == "Ensemble" + assert env.num_ensemble_members == 2 + assert env.ensemble_member == 0 + assert env.pressure(1000) == pytest.approx(90000) + assert env.temperature(1000) == pytest.approx(281) + assert env.wind_velocity_x(1000) == pytest.approx(2) + + env.select_ensemble_member(1) + assert env.pressure(1100) == pytest.approx(90000) + assert env.temperature(1100) == pytest.approx(283) + assert env.wind_velocity_x(1100) == pytest.approx(5) + assert env.wind_velocity_y(1100) == pytest.approx(-5) + + with netCDF4.Dataset(file_path) as dataset: + assert dataset.Conventions == "CF-1.8" + assert dataset.source == "RocketPy Environment.create_ensemble" + assert dataset.variables["time"].long_name == "profile valid time" + assert { + name: len(dataset.dimensions[name]) for name in ("ens", "lev", "time") + } == {"ens": 2, "lev": 3, "time": 1} + assert dataset.variables["lev"].units == "hPa" + assert dataset.variables["tmpprs"].standard_name == "air_temperature" + npt.assert_allclose(dataset.variables["lev"][:], [1013.25, 900, 800]) + + +def test_create_ensemble_file_round_trip(tmp_path): + """Reload the exported file using the existing GEFS ensemble mapping.""" + # Arrange + source_env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + file_path = source_env.create_ensemble( + _user_defined_ensemble_profiles(), file_name=tmp_path / "round_trip.nc" + ) + loaded_env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act + loaded_env.set_atmospheric_model(type="Ensemble", file=file_path, dictionary="GEFS") + loaded_env.select_ensemble_member(1) + + # Assert + assert loaded_env.num_ensemble_members == 2 + assert loaded_env.pressure(1100) == pytest.approx(90000) + assert loaded_env.temperature(1100) == pytest.approx(283) + assert loaded_env.wind_velocity_x(1100) == pytest.approx(5) + assert loaded_env.wind_velocity_y(1100) == pytest.approx(-5) + + +def test_create_ensemble_rejects_non_overlapping_pressure_profiles(tmp_path): + """Reject members that cannot be sampled on a common pressure grid.""" + # Arrange + profiles = _user_defined_ensemble_profiles() + heights = profiles[1]["pressure"][:, 0] + profiles[1]["pressure"] = np.column_stack((heights, [70000.0, 60000.0, 50000.0])) + + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(ValueError, match="no common pressure range"): + env.create_ensemble(profiles, file_name=tmp_path / "invalid.nc") + + +def test_create_ensemble_does_not_overwrite_by_default(tmp_path): + """Preserve an existing ensemble file unless overwrite is explicit.""" + # Arrange + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + file_path = env.create_ensemble( + _user_defined_ensemble_profiles(), file_name=tmp_path / "existing.nc" + ) + + # Act / Assert + with pytest.raises(FileExistsError, match="overwrite=True"): + env.create_ensemble(_user_defined_ensemble_profiles(), file_name=file_path) + + class DummyLambertProjection: """Minimal projection metadata container for unit tests.""" From 5f5d26ee39668898ee268d524dd5d6c07ba9316a Mon Sep 17 00:00:00 2001 From: Ting-Hong Shieh <32212900+ting-hong-shieh@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:07:45 +0800 Subject: [PATCH 2/3] MNT: split custom ensemble creation into helpers --- rocketpy/environment/environment.py | 343 ++++++++++++++++------------ 1 file changed, 203 insertions(+), 140 deletions(-) diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index cf523cbd7..460f0bc89 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -2837,62 +2837,8 @@ def _prepare_ensemble_profile_source(source, variable, member): raise ValueError(f"Member {member} '{variable}' heights must be unique.") return profile - def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-many-statements - self, - profiles, - file_name="custom_ensemble.nc", - pressure_levels=None, - overwrite=False, - ): - """Create and activate an ensemble from user-defined profiles. - - RocketPy writes the profiles with GEFS-compatible variable names. - Another Environment can load the returned file by passing - ``type="Ensemble"`` and ``dictionary="GEFS"`` to - :meth:`Environment.set_atmospheric_model`. - - Parameters - ---------- - profiles : sequence of mappings - Atmospheric profiles for each ensemble member. Every mapping must - define ``pressure``, ``temperature``, ``wind_u`` and ``wind_v``. - Each value must be a two-column array whose first column is - geometric height above sea level in meters. The second column uses - Pa for pressure, K for temperature and m/s for either wind - component. Array-backed :class:`rocketpy.Function` objects are also - accepted. Pressure must decrease strictly with increasing height. - file_name : str or os.PathLike, optional - Path of the NetCDF file to create. The ``.nc`` suffix is appended - when omitted. Default is ``"custom_ensemble.nc"``. - pressure_levels : array-like, optional - Common pressure levels in Pa. By default, the union of sampled - pressure levels inside the range shared by every member is used. - overwrite : bool, optional - Whether an existing file may be replaced. Default is ``False``. - - Returns - ------- - str - Absolute path of the created NetCDF file. - - Raises - ------ - TypeError - If profiles or profile values have invalid types. - ValueError - If fewer than two members are supplied, required variables are - missing, profiles are invalid, or the members have no usable common - pressure range. - FileExistsError - If the output exists and ``overwrite`` is ``False``. - - Notes - ----- - The first member is activated after the file is created. Use - :meth:`Environment.select_ensemble_member` to activate another member. - """ - self.__validate_datetime() - + def _prepare_ensemble_profiles(self, profiles): + """Validate and normalize every user-defined ensemble member.""" if isinstance(profiles, (str, bytes, Mapping)): raise TypeError("'profiles' must be a sequence of member mappings.") try: @@ -2940,7 +2886,11 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma f"Member {member_index} temperature values must be positive." ) members.append(prepared) + return members + @staticmethod + def _prepare_ensemble_pressure_levels(members, pressure_levels): + """Return a valid pressure grid shared by every ensemble member.""" common_min_pressure = max(member["pressure"][-1, 1] for member in members) common_max_pressure = min(member["pressure"][0, 1] for member in members) if common_min_pressure >= common_max_pressure: @@ -2982,9 +2932,13 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma "'pressure_levels' must stay inside the pressure range shared " "by every member." ) + return pressure_levels + def _interpolate_ensemble_profiles(self, members, pressure_levels): + """Interpolate ensemble members onto their common pressure grid.""" + value_variables = ("temperature", "wind_u", "wind_v") member_heights = [] - member_values = {name: [] for name in required_variables[1:]} + member_values = {name: [] for name in value_variables} for member_index, member in enumerate(members): pressure_profile = member["pressure"] heights = np.interp( @@ -2994,7 +2948,7 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma ) member_heights.append(heights) - for variable in required_variables[1:]: + for variable in value_variables: profile = member[variable] if heights[0] < profile[0, 0] or heights[-1] > profile[-1, 0]: raise ValueError( @@ -3013,7 +2967,11 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma * geometric_heights / (self.earth_radius + geometric_heights) ) + return geopotential_heights, member_values + @staticmethod + def _prepare_ensemble_file_path(file_name, overwrite): + """Normalize the output path and protect existing files.""" try: file_path = os.fspath(file_name) except TypeError as exc: @@ -3027,7 +2985,59 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma raise FileExistsError( f"'{file_path}' already exists. Pass overwrite=True to replace it." ) + return file_path + + def _create_ensemble_time_coordinate(self, dataset): + """Create the valid-time coordinate for an ensemble dataset.""" + time = dataset.createVariable("time", "f8", ("time",)) + time.long_name = "profile valid time" + time.standard_name = "time" + time.units = ( + f"hours since {self.datetime_date.strftime('%Y-%m-%d %H:%M:%S')} UTC" + ) + time.calendar = "gregorian" + time.axis = "T" + time[:] = [0] + + @staticmethod + def _create_ensemble_member_and_level_coordinates( + dataset, member_count, pressure_levels + ): + """Create the ensemble-member and pressure-level coordinates.""" + ensemble = dataset.createVariable("ens", "i4", ("ens",)) + ensemble.long_name = "ensemble member" + ensemble.units = "1" + ensemble[:] = np.arange(member_count) + + level = dataset.createVariable("lev", "f8", ("lev",)) + level.long_name = "pressure level" + level.standard_name = "air_pressure" + level.units = "hPa" + level.positive = "down" + level.axis = "Z" + level[:] = pressure_levels / 100 + @staticmethod + def _create_ensemble_spatial_coordinates( + dataset, latitude_bounds, longitude_bounds + ): + """Create latitude and longitude coordinates for an ensemble dataset.""" + latitude = dataset.createVariable("lat", "f8", ("lat",)) + latitude.long_name = "latitude" + latitude.standard_name = "latitude" + latitude.units = "degrees_north" + latitude.axis = "Y" + latitude[:] = latitude_bounds + + longitude = dataset.createVariable("lon", "f8", ("lon",)) + longitude.long_name = "longitude" + longitude.standard_name = "longitude" + longitude.units = "degrees_east" + longitude.axis = "X" + longitude[:] = longitude_bounds + + def _create_ensemble_coordinates(self, dataset, member_count, pressure_levels): + """Create dimensions and coordinate variables for an ensemble dataset.""" latitude_bounds = np.array( [max(-90, self.latitude - 0.01), min(90, self.latitude + 0.01)] ) @@ -3038,14 +3048,75 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma min(360, grid_longitude + 0.01), ] ) - data_shape = ( + + dataset.createDimension("time", 1) + dataset.createDimension("ens", member_count) + dataset.createDimension("lev", len(pressure_levels)) + dataset.createDimension("lat", len(latitude_bounds)) + dataset.createDimension("lon", len(longitude_bounds)) + + self._create_ensemble_time_coordinate(dataset) + self._create_ensemble_member_and_level_coordinates( + dataset, member_count, pressure_levels + ) + self._create_ensemble_spatial_coordinates( + dataset, latitude_bounds, longitude_bounds + ) + + return ( 1, - len(members), + member_count, len(pressure_levels), len(latitude_bounds), len(longitude_bounds), ) + @staticmethod + def _create_ensemble_data_variables( + dataset, data_shape, geopotential_heights, member_values + ): + """Create and populate the atmospheric variables in an ensemble dataset.""" + variables = { + "hgtprs": ( + geopotential_heights, + "geopotential height", + "geopotential_height", + "m", + ), + "tmpprs": ( + np.asarray(member_values["temperature"]), + "air temperature", + "air_temperature", + "K", + ), + "ugrdprs": ( + np.asarray(member_values["wind_u"]), + "eastward wind", + "eastward_wind", + "m s-1", + ), + "vgrdprs": ( + np.asarray(member_values["wind_v"]), + "northward wind", + "northward_wind", + "m s-1", + ), + } + dimensions = ("time", "ens", "lev", "lat", "lon") + for name, (values, long_name, standard_name, units) in variables.items(): + variable = dataset.createVariable( + name, "f8", dimensions, zlib=True, complevel=4 + ) + variable.long_name = long_name + variable.standard_name = standard_name + variable.units = units + variable.coordinates = "time ens lev lat lon" + variable[:] = np.broadcast_to(values[None, :, :, None, None], data_shape) + + def _write_ensemble_file( + self, file_path, pressure_levels, geopotential_heights, member_values + ): + """Write prepared ensemble data to a GEFS-compatible NetCDF file.""" with netCDF4.Dataset(file_path, "w", format="NETCDF4") as dataset: dataset.Conventions = "CF-1.8" dataset.title = "RocketPy user-defined atmospheric ensemble" @@ -3060,87 +3131,79 @@ def create_ensemble( # pylint: disable=too-many-branches,too-many-locals,too-ma dataset.launch_latitude = self.latitude dataset.launch_longitude = self.longitude - dataset.createDimension("time", 1) - dataset.createDimension("ens", len(members)) - dataset.createDimension("lev", len(pressure_levels)) - dataset.createDimension("lat", len(latitude_bounds)) - dataset.createDimension("lon", len(longitude_bounds)) - - time = dataset.createVariable("time", "f8", ("time",)) - time.long_name = "profile valid time" - time.standard_name = "time" - time.units = ( - f"hours since {self.datetime_date.strftime('%Y-%m-%d %H:%M:%S')} UTC" + data_shape = self._create_ensemble_coordinates( + dataset, len(geopotential_heights), pressure_levels ) - time.calendar = "gregorian" - time.axis = "T" - time[:] = [0] - - ensemble = dataset.createVariable("ens", "i4", ("ens",)) - ensemble.long_name = "ensemble member" - ensemble.units = "1" - ensemble[:] = np.arange(len(members)) - - level = dataset.createVariable("lev", "f8", ("lev",)) - level.long_name = "pressure level" - level.standard_name = "air_pressure" - level.units = "hPa" - level.positive = "down" - level.axis = "Z" - level[:] = pressure_levels / 100 - - latitude = dataset.createVariable("lat", "f8", ("lat",)) - latitude.long_name = "latitude" - latitude.standard_name = "latitude" - latitude.units = "degrees_north" - latitude.axis = "Y" - latitude[:] = latitude_bounds - - longitude = dataset.createVariable("lon", "f8", ("lon",)) - longitude.long_name = "longitude" - longitude.standard_name = "longitude" - longitude.units = "degrees_east" - longitude.axis = "X" - longitude[:] = longitude_bounds - - dimensions = ("time", "ens", "lev", "lat", "lon") - variables = { - "hgtprs": ( - geopotential_heights, - "geopotential height", - "geopotential_height", - "m", - ), - "tmpprs": ( - np.asarray(member_values["temperature"]), - "air temperature", - "air_temperature", - "K", - ), - "ugrdprs": ( - np.asarray(member_values["wind_u"]), - "eastward wind", - "eastward_wind", - "m s-1", - ), - "vgrdprs": ( - np.asarray(member_values["wind_v"]), - "northward wind", - "northward_wind", - "m s-1", - ), - } - for name, (values, long_name, standard_name, units) in variables.items(): - variable = dataset.createVariable( - name, "f8", dimensions, zlib=True, complevel=4 - ) - variable.long_name = long_name - variable.standard_name = standard_name - variable.units = units - variable.coordinates = "time ens lev lat lon" - variable[:] = np.broadcast_to( - values[None, :, :, None, None], data_shape - ) + self._create_ensemble_data_variables( + dataset, data_shape, geopotential_heights, member_values + ) + + def create_ensemble( + self, + profiles, + file_name="custom_ensemble.nc", + pressure_levels=None, + overwrite=False, + ): + """Create and activate an ensemble from user-defined profiles. + + RocketPy writes the profiles with GEFS-compatible variable names. + Another Environment can load the returned file by passing + ``type="Ensemble"`` and ``dictionary="GEFS"`` to + :meth:`Environment.set_atmospheric_model`. + + Parameters + ---------- + profiles : sequence of mappings + Atmospheric profiles for each ensemble member. Every mapping must + define ``pressure``, ``temperature``, ``wind_u`` and ``wind_v``. + Each value must be a two-column array whose first column is + geometric height above sea level in meters. The second column uses + Pa for pressure, K for temperature and m/s for either wind + component. Array-backed :class:`rocketpy.Function` objects are also + accepted. Pressure must decrease strictly with increasing height. + file_name : str or os.PathLike, optional + Path of the NetCDF file to create. The ``.nc`` suffix is appended + when omitted. Default is ``"custom_ensemble.nc"``. + pressure_levels : array-like, optional + Common pressure levels in Pa. By default, the union of sampled + pressure levels inside the range shared by every member is used. + overwrite : bool, optional + Whether an existing file may be replaced. Default is ``False``. + + Returns + ------- + str + Absolute path of the created NetCDF file. + + Raises + ------ + TypeError + If profiles or profile values have invalid types. + ValueError + If fewer than two members are supplied, required variables are + missing, profiles are invalid, or the members have no usable common + pressure range. + FileExistsError + If the output exists and ``overwrite`` is ``False``. + + Notes + ----- + The first member is activated after the file is created. Use + :meth:`Environment.select_ensemble_member` to activate another member. + """ + self.__validate_datetime() + members = self._prepare_ensemble_profiles(profiles) + pressure_levels = self._prepare_ensemble_pressure_levels( + members, pressure_levels + ) + geopotential_heights, member_values = self._interpolate_ensemble_profiles( + members, pressure_levels + ) + file_path = self._prepare_ensemble_file_path(file_name, overwrite) + self._write_ensemble_file( + file_path, pressure_levels, geopotential_heights, member_values + ) self.set_atmospheric_model(type="Ensemble", file=file_path, dictionary="GEFS") logger.info("Atmospheric ensemble saved at '%s'.", file_path) From 41d5d53c448d9cf19ca6b5376d08e29e2b4ed166 Mon Sep 17 00:00:00 2001 From: Ting-Hong Shieh <32212900+ting-hong-shieh@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:32:46 +0800 Subject: [PATCH 3/3] TST: cover custom ensemble validation --- tests/unit/environment/test_environment.py | 162 ++++++++++++++++++++- 1 file changed, 161 insertions(+), 1 deletion(-) diff --git a/tests/unit/environment/test_environment.py b/tests/unit/environment/test_environment.py index 4db6a13f2..61d6c3ff4 100644 --- a/tests/unit/environment/test_environment.py +++ b/tests/unit/environment/test_environment.py @@ -8,7 +8,7 @@ import pytest import pytz -from rocketpy import Environment +from rocketpy import Environment, Function from rocketpy.environment.tools import ( find_longitude_index, geodesic_to_lambert_conformal, @@ -148,6 +148,166 @@ def test_create_ensemble_does_not_overwrite_by_default(tmp_path): env.create_ensemble(_user_defined_ensemble_profiles(), file_name=file_path) +def test_create_ensemble_accepts_array_functions_and_explicit_levels(tmp_path): + """Accept array-backed Functions and sort explicit pressure levels.""" + # Arrange + profiles = _user_defined_ensemble_profiles() + for member in profiles: + for variable, source in member.items(): + member[variable] = Function(source) + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act + file_path = env.create_ensemble( + profiles, + file_name=tmp_path / "function_profiles.nc", + pressure_levels=[80000, 101325, 90000], + ) + + # Assert + with netCDF4.Dataset(file_path) as dataset: + npt.assert_allclose(dataset.variables["lev"][:], [1013.25, 900, 800]) + + +@pytest.mark.parametrize( + "source, error, match", + [ + (Function(lambda height: height), TypeError, "array-backed Function"), + (object(), TypeError, "two-column numeric array"), + (np.array([0.0, 1.0]), ValueError, "at least two"), + (np.array([[0.0, 1.0], [1.0, np.inf]]), ValueError, "non-finite"), + (np.array([[0.0, 1.0], [0.0, 2.0]]), ValueError, "heights must be unique"), + ], +) +def test_create_ensemble_rejects_invalid_profile_sources( + tmp_path, source, error, match +): + """Reject profile sources that cannot define a finite height-value curve.""" + # Arrange + profiles = _user_defined_ensemble_profiles() + profiles[0]["wind_u"] = source + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(error, match=match): + env.create_ensemble(profiles, file_name=tmp_path / "invalid_source.nc") + + +def test_create_ensemble_rejects_invalid_profile_collections(tmp_path): + """Reject invalid ensemble containers and incomplete members.""" + # Arrange + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + profiles = _user_defined_ensemble_profiles() + output = tmp_path / "invalid_collection.nc" + + # Act / Assert + with pytest.raises(TypeError, match="sequence of member mappings"): + env.create_ensemble(profiles[0], file_name=output) + with pytest.raises(TypeError, match="sequence of member mappings"): + env.create_ensemble(1, file_name=output) + with pytest.raises(ValueError, match="At least two"): + env.create_ensemble(profiles[:1], file_name=output) + with pytest.raises(TypeError, match="Member 1 must be a mapping"): + env.create_ensemble([profiles[0], None], file_name=output) + + incomplete_profiles = _user_defined_ensemble_profiles() + incomplete_profiles[1].pop("wind_v") + with pytest.raises(ValueError, match="missing required profile.*wind_v"): + env.create_ensemble(incomplete_profiles, file_name=output) + + +@pytest.mark.parametrize( + "variable, values, match", + [ + ("pressure", [101325.0, 0.0, 80000.0], "pressure values must be positive"), + ( + "pressure", + [101325.0, 80000.0, 90000.0], + "pressure must decrease strictly", + ), + ("temperature", [288.0, 0.0, 275.0], "temperature values must be positive"), + ], +) +def test_create_ensemble_rejects_invalid_profile_values( + tmp_path, variable, values, match +): + """Reject nonphysical pressure and temperature profile values.""" + # Arrange + profiles = _user_defined_ensemble_profiles() + profiles[0][variable][:, 1] = values + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(ValueError, match=match): + env.create_ensemble(profiles, file_name=tmp_path / "invalid_values.nc") + + +@pytest.mark.parametrize( + "pressure_levels, error, match", + [ + (["invalid", "values"], TypeError, "numeric array"), + ([[101325.0, 90000.0]], ValueError, "one-dimensional"), + ([101325.0, np.nan], ValueError, "finite, positive"), + ([101325.0, 101325.0], ValueError, "duplicates"), + ([90000.0], ValueError, "At least two pressure levels"), + ([110000.0, 90000.0], ValueError, "inside the pressure range"), + ], +) +def test_create_ensemble_rejects_invalid_pressure_levels( + tmp_path, pressure_levels, error, match +): + """Reject explicit pressure grids that cannot be shared by all members.""" + # Arrange + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(error, match=match): + env.create_ensemble( + _user_defined_ensemble_profiles(), + file_name=tmp_path / "invalid_levels.nc", + pressure_levels=pressure_levels, + ) + + +def test_create_ensemble_rejects_profiles_without_height_coverage(tmp_path): + """Require every variable to span the common pressure-grid heights.""" + # Arrange + profiles = _user_defined_ensemble_profiles() + profiles[0]["temperature"] = profiles[0]["temperature"][1:] + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(ValueError, match="temperature.*does not cover all heights"): + env.create_ensemble(profiles, file_name=tmp_path / "incomplete_height.nc") + + +def test_create_ensemble_rejects_heights_below_earth_center(tmp_path): + """Reject geometric heights at or below the coordinate singularity.""" + # Arrange + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + profiles = _user_defined_ensemble_profiles() + invalid_heights = np.array( + [-env.earth_radius - 2000, -env.earth_radius - 1000, -env.earth_radius - 1] + ) + for member in profiles: + for source in member.values(): + source[:, 0] = invalid_heights + + # Act / Assert + with pytest.raises(ValueError, match="greater than -Earth's radius"): + env.create_ensemble(profiles, file_name=tmp_path / "invalid_height.nc") + + +def test_create_ensemble_rejects_invalid_file_name(): + """Require the NetCDF output name to implement the path protocol.""" + # Arrange + env = Environment(date=(2025, 6, 1, 12), latitude=32.99, longitude=-106.97) + + # Act / Assert + with pytest.raises(TypeError, match="string or path-like"): + env.create_ensemble(_user_defined_ensemble_profiles(), file_name=object()) + + class DummyLambertProjection: """Minimal projection metadata container for unit tests."""