rename _prepare_4d_data to _prepare_soil_data, use Union for typing

Author: SarahAlidoost

PyStemmusScope/save.py

@@ -24,8 +24,7 @@
 
 import logging
 from pathlib import Path
-from typing import Dict
-from typing import List
+from typing import Dict, List, Union
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -56,15 +55,15 @@ def _select_forcing_variables(forcing_dict: Dict, forcing_var: str, alma_var: st
     return data_array
 
 
-def _shorten_data_array(data: xr.DataArray, time_steps: str)-> xr.DataArray:
+def _shorten_data_array(data: Union[xr.DataArray, xr.Dataset], time_steps: str)-> Union[xr.DataArray, xr.Dataset]:
     """Shorten data based on time_steps.
 
     Args:
-        data(xr.DataArray): data to be shortend.
+        data(xr.DataArray or xr.Dataset): data to be shortend.
         time_steps(str): number of time steps to shorten.
 
     Returns:
-        xr.DataArray: subset of data with the lenght of time equal to time_steps.
+        xr.DataArray or xr.Dataset: subset of data with the lenght of time equal to time_steps.
     """
 
     if time_steps != "NA":
@@ -74,8 +73,8 @@ def _shorten_data_array(data: xr.DataArray, time_steps: str)-> xr.DataArray:
     return data
 
 
-def _prepare_4d_data(file_name: str, var_name: str, time: List) -> xr.DataArray:
-    """Reshape a `pandas.DataFrame` to include z dimension.
+def _prepare_soil_data(file_name: str, var_name: str, time: List) -> xr.DataArray:
+    """Return simulated soil temperature and soil moisture as `xr.DataArray`.
 
     Args:
         file_name(str): csv file name generated by Stemmus_Scope model.
@@ -304,7 +303,7 @@ def to_netcdf(config: Dict, cf_filename: str) -> str:
 
         # create data array
         elif alma_name in {"SoilTemp", "SoilMoist"}:
-            data_array = _prepare_4d_data(file_name, alma_name, time.values)
+            data_array = _prepare_soil_data(file_name, alma_name, time.values)
         else:
             data_array = _prepare_3d_data(
                 file_name, df["Variable name in STEMMUS-SCOPE"], alma_name, time.values

tests/test_save.py

@@ -167,7 +167,7 @@ def test_save_to_netcdf(self, cf_convention, _make_csv_file, model_with_setup):
         assert "z" not in dataset
 
 
-class Test4dData:
+class TestSoilData:
     @pytest.fixture
     def cf_convention(self, tmp_path):
         convention = [