ENH: Compatibility with MERRA-2 atmosphere reanalysis files #825 (#877)

* created new branch

* ENH: Add MERRA-2 compatibility with unit conversion and documentation.

* ENH: Finalize MERRA-2 support (Resolve conflicts, Update Docs)

* Style: Fix ruff formatting errors

* REV: Address all code review comments

* Apply suggestion from @Gui-FernandesBR

---------

Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com>

Author: Monta120

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 Attention: The newest changes should be on top -->
 
 ### Added
+ENH: Compatibility with MERRA-2 atmosphere reanalysis files [#825](https://github.com/RocketPy-Team/RocketPy/pull/825)
 
 - ENH: Enable only radial burning [#815](https://github.com/RocketPy-Team/RocketPy/pull/815)
 - ENH: Add thrustcurve api integration to retrieve motor eng data [#870](https://github.com/RocketPy-Team/RocketPy/pull/870)

docs/user/environment/1-atm-models/reanalysis.rst

@@ -46,6 +46,27 @@ ERA5 data can be downloaded from the
     processed by RocketPy. It is recommended that you download only the \
     necessary data.
 
+MERRA-2
+-------
+
+The Modern-Era Retrospective analysis for Research and Applications, Version 2 (MERRA-2) is a NASA atmospheric reanalysis for the satellite era using the Goddard Earth Observing System, Version 5 (GEOS-5) with its Atmospheric Data Assimilation System (ADAS).
+
+You can download these files from the `NASA GES DISC <https://disc.gsfc.nasa.gov/>`_.
+
+To use MERRA-2 data in RocketPy, you generally need the **Assimilated Meteorological Fields** collection (specifically the 3D Pressure Level data, usually named ``inst3_3d_asm_Np``). Note that MERRA-2 files typically use the ``.nc4`` extension (NetCDF-4), which is fully supported by RocketPy.
+
+You can load these files using the ``dictionary="MERRA2"`` argument:
+
+.. code-block:: python
+
+    env.set_atmospheric_model(
+        type="Reanalysis",
+        file="MERRA2_400.inst3_3d_asm_Np.20230620.nc4",
+        dictionary="MERRA2"
+    )
+
+RocketPy automatically handles the unit conversion for MERRA-2's surface geopotential (energy) to geometric height (meters).
+
 
 Setting the Environment
 ^^^^^^^^^^^^^^^^^^^^^^^

rocketpy/environment/environment.py

@@ -607,7 +607,7 @@ def __set_earth_rotation_vector(self):
 
     def __validate_dictionary(self, file, dictionary):
         # removed CMC until it is fixed.
-        available_models = ["GFS", "NAM", "RAP", "HIRESW", "GEFS", "ERA5"]
+        available_models = ["GFS", "NAM", "RAP", "HIRESW", "GEFS", "ERA5", "MERRA2"]
         if isinstance(dictionary, str):
             dictionary = self.__weather_model_map.get(dictionary)
         elif file in available_models:
@@ -1186,7 +1186,7 @@ def set_atmospheric_model(  # pylint: disable=too-many-statements
             ``Reanalysis`` or ``Ensemble``. It specifies the dictionary to be
             used when reading ``netCDF`` and ``OPeNDAP`` files, allowing the
             correct retrieval of data. Acceptable values include ``ECMWF``,
-            ``NOAA`` and ``UCAR`` for default dictionaries which can generally
+            ``NOAA``, ``UCAR`` and ``MERRA2`` for default dictionaries which can generally
             be used to read datasets from these institutes. Alternatively, a
             dictionary structure can also be given, specifying the short names
             used for time, latitude, longitude, pressure levels, temperature
@@ -1893,10 +1893,20 @@ def process_forecast_reanalysis(self, file, dictionary):  # pylint: disable=too-
         self._max_expected_height = max(height[0], height[-1])
 
         # Get elevation data from file
-        if dictionary["surface_geopotential_height"] is not None:
+        if dictionary.get("surface_geopotential_height") is not None:
             self.elevation = get_elevation_data_from_dataset(
                 dictionary, data, time_index, lat_index, lon_index, x, y, x1, x2, y1, y2
             )
+        # 2. If not found, try Geopotential (m^2/s^2) and convert
+        elif dictionary.get("surface_geopotential") is not None:
+            temp_dict = dictionary.copy()
+            temp_dict["surface_geopotential_height"] = dictionary[
+                "surface_geopotential"
+            ]
+            surface_geopotential_value = get_elevation_data_from_dataset(
+                temp_dict, data, time_index, lat_index, lon_index, x, y, x1, x2, y1, y2
+            )
+            self.elevation = surface_geopotential_value / self.standard_g
 
         # Compute info data
         self.atmospheric_model_init_date = get_initial_date_from_time_array(time_array)

rocketpy/environment/weather_model_mapping.py

@@ -114,6 +114,19 @@ class WeatherModelMapping:
         "u_wind": "ugrdprs",
         "v_wind": "vgrdprs",
     }
+    MERRA2 = {
+        "time": "time",
+        "latitude": "lat",
+        "longitude": "lon",
+        "level": "lev",
+        "temperature": "T",
+        "surface_geopotential_height": None,
+        "surface_geopotential": "PHIS",  # special key for Geopotential (m^2/s^2)
+        "geopotential_height": "H",
+        "geopotential": None,
+        "u_wind": "U",
+        "v_wind": "V",
+    }
 
     def __init__(self):
         """Initialize the class, creates a dictionary with all the weather models
@@ -129,6 +142,7 @@ def __init__(self):
             "CMC": self.CMC,
             "GEFS": self.GEFS,
             "HIRESW": self.HIRESW,
+            "MERRA2": self.MERRA2,
         }
 
     def get(self, model):

tests/conftest.py

@@ -1,3 +1,5 @@
+import netCDF4
+import numpy as np
 import matplotlib
 import pytest
 
@@ -72,3 +74,67 @@ def pytest_collection_modifyitems(config, items):
     for item in items:
         if "slow" in item.keywords:
             item.add_marker(skip_slow)
+
+
+@pytest.fixture
+def merra2_file_path(tmp_path):  # pylint: disable=too-many-statements
+    """
+    Generates a temporary NetCDF file that STRICTLY mimics the structure of a
+    NASA MERRA-2 'inst3_3d_asm_Np' file (Assimilated Meteorological Fields)
+    because MERRA-2 files are too large.
+
+    """
+    file_path = tmp_path / "MERRA2_300.inst3_3d_asm_Np.20230620.nc4"
+
+    with netCDF4.Dataset(file_path, "w", format="NETCDF4") as nc:
+        # Define Dimensions
+        nc.createDimension("lon", 5)
+        nc.createDimension("lat", 5)
+        nc.createDimension("lev", 5)
+        nc.createDimension("time", None)
+
+        # Define Coordinates
+        lon = nc.createVariable("lon", "f8", ("lon",))
+        lon.units = "degrees_east"
+        lon[:] = np.linspace(-180, 180, 5)
+
+        lat = nc.createVariable("lat", "f8", ("lat",))
+        lat.units = "degrees_north"
+        lat[:] = np.linspace(-90, 90, 5)
+
+        lev = nc.createVariable("lev", "f8", ("lev",))
+        lev.units = "hPa"
+        lev[:] = np.linspace(1000, 100, 5)
+
+        time = nc.createVariable("time", "i4", ("time",))
+        time.units = "minutes since 2023-06-20 00:00:00"
+        time[:] = [720]
+
+        # Define Data Variables
+        # NetCDF names are uppercase ("T") to match NASA specs.
+
+        t_var = nc.createVariable("T", "f4", ("time", "lev", "lat", "lon"))
+        t_var.units = "K"
+        t_var[:] = np.full((1, 5, 5, 5), 300.0)
+
+        u_var = nc.createVariable("U", "f4", ("time", "lev", "lat", "lon"))
+        u_var.units = "m s-1"
+        u_var[:] = np.full((1, 5, 5, 5), 10.0)
+
+        v_var = nc.createVariable("V", "f4", ("time", "lev", "lat", "lon"))
+        v_var.units = "m s-1"
+        v_var[:] = np.full((1, 5, 5, 5), 5.0)
+
+        h_var = nc.createVariable("H", "f4", ("time", "lev", "lat", "lon"))
+        h_var.units = "m"
+        h_var[:] = np.linspace(0, 10000, 5).reshape(1, 5, 1, 1) * np.ones((1, 5, 5, 5))
+
+        # PHIS: Surface Geopotential Height [m2 s-2]
+        phis = nc.createVariable("PHIS", "f4", ("time", "lat", "lon"))
+        phis.units = "m2 s-2"
+
+        # We set PHIS to 9806.65 (Energy).
+        # We expect the code to divide by ~9.8 and get ~1000.0 (Height).
+        phis[:] = np.full((1, 5, 5), 9806.65)
+
+    return str(file_path)

tests/integration/environment/test_environment.py

@@ -258,3 +258,33 @@ def test_cmc_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=un
     """
     example_spaceport_env.set_atmospheric_model(type="Ensemble", file="CMC")
     assert example_spaceport_env.all_info() is None
+
+
+def test_merra2_full_specification_compliance(merra2_file_path, example_plain_env):
+    """
+    Tests that RocketPy loads a file complying with NASA MERRA-2 file specs.
+    """
+    # 1. Initialize Environment
+    env = example_plain_env
+    env.set_date((2023, 6, 20, 12))
+    env.set_location(latitude=0, longitude=0)
+
+    # 2. Force standard gravity to a known constant for precise math checking
+    env.standard_g = 9.80665
+
+    # 3. Load the Atmospheric Model (Using the file generated above)
+    env.set_atmospheric_model(
+        type="Reanalysis", file=merra2_file_path, dictionary="MERRA2"
+    )
+
+    # 4. Verify Unit Conversion (Energy -> Height)
+    # Input: 9806.65 m2/s2
+    # Expected: 1000.0 m
+    print(f"Calculated Elevation: {env.elevation} m")
+    assert abs(env.elevation - 1000.0) < 1e-6, (
+        f"Failed to convert PHIS (m2/s2) to meters. Got {env.elevation}, expected 1000.0"
+    )
+
+    # 5. Verify Variable Mapping
+    assert abs(env.temperature(0) - 300.0) < 1e-6
+    assert env.wind_speed(0) > 0