MNT: introduce pressure unit conversion when using forecast/reanalysis/ensemble data (#955)

* MNT: remove conversion from mbar to Pa in netcdf4 fetched data

* ENH: introducing pressure conversion factor when using forecast or reanalysis

* fix: adding conversion factor argument to ensemble as well

Author: Lucas-Prates

rocketpy/environment/environment.py

@@ -1120,6 +1120,7 @@ def set_atmospheric_model(  # pylint: disable=too-many-statements
         temperature=None,
         wind_u=0,
         wind_v=0,
+        pressure_conversion_factor="Pa",
     ):
         """Define the atmospheric model for this Environment.
 
@@ -1216,6 +1217,12 @@ def set_atmospheric_model(  # pylint: disable=too-many-statements
             m/s). Finally, a callable or function is also accepted. The function
             should take one argument, the height above sea level in meters and
             return a corresponding wind-v in m/s.
+        pressure_conversion_factor : string, int, float
+            This defines the pressure conversion factor to Pa when type is
+            ``forecast`` or ``reanalysis``. The pressure unit from the data may
+            not be in Pascal, so the correction is necessary. Valid strings are
+            ("mbar", "hPa", "Pa"), or a strictly positive number if using a
+            custom pressure unit.
 
         Returns
         -------
@@ -1265,6 +1272,28 @@ def set_atmospheric_model(  # pylint: disable=too-many-statements
             case "windy":
                 self.process_windy_atmosphere(file)
             case "forecast" | "reanalysis" | "ensemble":
+                conversion_factor = 1
+                if not isinstance(pressure_conversion_factor, (float, int, str)):
+                    raise ValueError(
+                        "Argument 'pressure_conversion_factor' must be numeric or a standard pressure unit ('mbar', 'hPa', 'Pa')!"
+                    )
+                if isinstance(pressure_conversion_factor, (float, int)):
+                    if pressure_conversion_factor <= 0:
+                        raise ValueError(
+                            "Argument 'pressure_conversion_factor' must be strictly positive!"
+                        )
+                    else:
+                        conversion_factor = pressure_conversion_factor
+                if isinstance(pressure_conversion_factor, str):
+                    if pressure_conversion_factor.lower() in ("mbar", "hpa"):
+                        conversion_factor = 100
+                    elif pressure_conversion_factor.lower() == "pa":
+                        conversion_factor = 1
+                    else:
+                        raise ValueError(
+                            "Argument 'pressure_conversion_factor' unit must be a standard pressure unit ('mbar', 'hPa', 'Pa')!"
+                        )
+
                 if isinstance(file, str):
                     shortcut_map = self.__atm_type_file_to_function_map.get(type, {})
                     matching_shortcut = next(
@@ -1305,9 +1334,11 @@ def set_atmospheric_model(  # pylint: disable=too-many-statements
                 dataset = fetch_function() if fetch_function is not None else file
 
                 if type in ["forecast", "reanalysis"]:
-                    self.process_forecast_reanalysis(dataset, dictionary)
+                    self.process_forecast_reanalysis(
+                        dataset, dictionary, conversion_factor=conversion_factor
+                    )
                 else:
-                    self.process_ensemble(dataset, dictionary)
+                    self.process_ensemble(dataset, dictionary, conversion_factor)
             case _:  # pragma: no cover
                 raise ValueError(f"Unknown model type '{type}'.")
 
@@ -1686,7 +1717,7 @@ def process_wyoming_sounding(self, file):  # pylint: disable=too-many-statements
         # Save maximum expected height
         self._max_expected_height = data_array[-1, 1]
 
-    def process_forecast_reanalysis(self, file, dictionary):  # pylint: disable=too-many-locals,too-many-statements
+    def process_forecast_reanalysis(self, file, dictionary, conversion_factor):  # pylint: disable=too-many-locals,too-many-statements
         """Import and process atmospheric data from weather forecasts
         and reanalysis given as ``netCDF`` or ``OPeNDAP`` files.
         Sets pressure, temperature, wind-u and wind-v
@@ -1738,6 +1769,9 @@ def process_forecast_reanalysis(self, file, dictionary):  # pylint: disable=too-
                     "u_wind": "ugrdprs",
                     "v_wind": "vgrdprs",
                 }
+        conversion_factor : float, int
+            Specifies the factor by which the pressure will be multiplied
+            in order to transform it to Pascal.
 
         Returns
         -------
@@ -1790,7 +1824,7 @@ def process_forecast_reanalysis(self, file, dictionary):  # pylint: disable=too-
         _, lat_index = find_latitude_index(target_lat, lat_array)
 
         # Get pressure level data from file
-        levels = get_pressure_levels_from_file(data, dictionary)
+        levels = get_pressure_levels_from_file(data, dictionary, conversion_factor)
 
         # Get geopotential data from file
         try:
@@ -1991,7 +2025,7 @@ def process_forecast_reanalysis(self, file, dictionary):  # pylint: disable=too-
         # Close weather data
         data.close()
 
-    def process_ensemble(self, file, dictionary):  # pylint: disable=too-many-locals,too-many-statements
+    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,
         wind-u and wind-v profiles and surface elevation obtained from a weather
@@ -2042,6 +2076,9 @@ def process_ensemble(self, file, dictionary):  # pylint: disable=too-many-locals
                     "u_wind": "ugrdprs",
                     "v_wind": "vgrdprs",
                 }
+        conversion_factor : float, int
+            Specifies the factor by which the pressure will be multiplied
+            in order to transform it to Pascal.
 
         See also
         --------
@@ -2106,7 +2143,7 @@ class for some dictionary examples.
             num_members = 1
 
         # Get pressure level data from file
-        levels = get_pressure_levels_from_file(data, dictionary)
+        levels = get_pressure_levels_from_file(data, dictionary, conversion_factor)
 
         inverse_dictionary = {v: k for k, v in dictionary.items()}
         param_dictionary = {

rocketpy/environment/tools.py

@@ -169,7 +169,7 @@ def geodesic_to_lambert_conformal(lat, lon, projection_variable, x_units="m"):
 ## These functions are meant to be used with netcdf4 datasets
 
 
-def get_pressure_levels_from_file(data, dictionary):
+def get_pressure_levels_from_file(data, dictionary, conversion_factor):
     """Extracts pressure levels from a netCDF4 dataset and converts them to Pa.
 
     Parameters
@@ -178,6 +178,9 @@ def get_pressure_levels_from_file(data, dictionary):
         The netCDF4 dataset containing the pressure level data.
     dictionary : dict
         A dictionary mapping variable names to dataset keys.
+    conversion_factor : float, int
+        Specifies the factor by which the pressure will be multiplied
+        in order to transform it to Pascal.
 
     Returns
     -------
@@ -190,8 +193,7 @@ def get_pressure_levels_from_file(data, dictionary):
         If the pressure levels cannot be read from the file.
     """
     try:
-        # Convert mbar to Pa
-        levels = 100 * data.variables[dictionary["level"]][:]
+        levels = conversion_factor * data.variables[dictionary["level"]][:]
     except KeyError as e:
         raise ValueError(
             "Unable to read pressure levels from file. Check file and dictionary."

tests/acceptance/test_bella_lui_rocket.py

@@ -67,6 +67,7 @@ def test_bella_lui_rocket_data_asserts_acceptance():
         type="Reanalysis",
         file="data/weather/bella_lui_weather_data_ERA5.nc",
         dictionary="ECMWF",
+        pressure_conversion_factor="hPa",
     )
     env.max_expected_height = 2000
 

tests/acceptance/test_ndrt_2020_rocket.py

@@ -83,6 +83,7 @@ def test_ndrt_2020_rocket_data_asserts_acceptance(env_file):
         type="Reanalysis",
         file=env_file,
         dictionary="ECMWF",
+        pressure_conversion_factor="hPa",
     )
     env.max_expected_height = 2000
 

tests/fixtures/environment/environment_fixtures.py

@@ -111,6 +111,7 @@ def environment_spaceport_america_2023():
         type="Reanalysis",
         file="data/weather/spaceport_america_pressure_levels_2023_hourly.nc",
         dictionary="ECMWF",
+        pressure_conversion_factor="hPa",
     )
 
     env.max_expected_height = 6000

tests/integration/environment/test_environment.py

@@ -45,6 +45,7 @@ def test_era5_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=u
         type="Reanalysis",
         file="data/weather/SpaceportAmerica_2018_ERA-5.nc",
         dictionary="ECMWF",
+        pressure_conversion_factor="hPa",
     )
     assert example_spaceport_env.all_info() is None
 

tests/unit/environment/test_environment.py

@@ -577,9 +577,10 @@ def test_set_atmospheric_model_normalizes_shortcut_case_for_forecast(example_pla
 
     called_arguments = {}
 
-    def fake_process_forecast_reanalysis(dataset, dictionary):
+    def fake_process_forecast_reanalysis(dataset, dictionary, conversion_factor):
         called_arguments["dataset"] = dataset
         called_arguments["dictionary"] = dictionary
+        called_arguments["conversion_factr"] = conversion_factor
 
     environment.process_forecast_reanalysis = fake_process_forecast_reanalysis
 
@@ -618,9 +619,10 @@ def test_forecast_shortcut_and_dictionary_are_case_insensitive(
 
     captured = {}
 
-    def fake_process_forecast_reanalysis(file, dictionary):
+    def fake_process_forecast_reanalysis(file, dictionary, conversion_factor):
         captured["file"] = file
         captured["dictionary"] = dictionary
+        captured["conversion_factor"] = conversion_factor
 
     monkeypatch.setattr(
         env, "process_forecast_reanalysis", fake_process_forecast_reanalysis

tests/unit/test_tools.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from rocketpy import Environment
 from rocketpy.tools import (
     calculate_cubic_hermite_coefficients,
     euler313_to_quaternions,
@@ -100,3 +101,39 @@ def test_tuple_handler(input_value, expected_output):
 def test_tuple_handler_exceptions(input_value, expected_exception):
     with pytest.raises(expected_exception):
         tuple_handler(input_value)
+
+
+@pytest.mark.parametrize("pressure_conversion_factor", ["hPa", "mbar", "Pa", 100])
+def test_valid_pressure_conversion_factor(pressure_conversion_factor):
+    env = Environment(
+        gravity=9.81,
+        latitude=47.213476,
+        longitude=9.003336,
+        date=(2020, 2, 22, 13),
+        elevation=407,
+    )
+    env.set_atmospheric_model(
+        type="Reanalysis",
+        file="data/weather/bella_lui_weather_data_ERA5.nc",
+        dictionary="ECMWF",
+        pressure_conversion_factor=pressure_conversion_factor,
+    )
+
+
+@pytest.mark.parametrize("pressure_conversion_factor", [-1, "mPa"])
+def test_invalid_pressure_conversion_factor(pressure_conversion_factor):
+    env = Environment(
+        gravity=9.81,
+        latitude=47.213476,
+        longitude=9.003336,
+        date=(2020, 2, 22, 13),
+        elevation=407,
+    )
+
+    with pytest.raises(ValueError):
+        env.set_atmospheric_model(
+            type="Reanalysis",
+            file="data/weather/bella_lui_weather_data_ERA5.nc",
+            dictionary="ECMWF",
+            pressure_conversion_factor=pressure_conversion_factor,
+        )