From 6c182a52d7ed59e7ce7e9fef299a11cde2f0a7ef Mon Sep 17 00:00:00 2001
From: Mostafa Farrag <moah.farag@gmail.com>
Date: Mon, 20 Jan 2025 12:36:26 +0100
Subject: [PATCH] add `get_units` method to the `TimModel`

---
 hydrolib/core/dflowfm/tim/models.py | 56 +++++++++++++++++++++++++++++
 tests/dflowfm/test_tim.py           | 16 +++++++++
 2 files changed, 72 insertions(+)

diff --git a/hydrolib/core/dflowfm/tim/models.py b/hydrolib/core/dflowfm/tim/models.py
index 42ad4e6bb..c7d9617f1 100644
--- a/hydrolib/core/dflowfm/tim/models.py
+++ b/hydrolib/core/dflowfm/tim/models.py
@@ -270,3 +270,59 @@ def as_dataframe(self, columns: List[Any] = None) -> DataFrame:
         if not columns:
             columns = self.quantities_names
         return DataFrame(time_series, index=index, columns=columns)
+
+    def get_units(self):
+        """Return the units for each quantity in the timeseries.
+
+        Returns:
+            List[str]: A list of units for each quantity in the timeseries.
+
+        Examples:
+            Create a `TimModel` object from a .tim file:
+                >>> from hydrolib.core.dflowfm.tim.models import TimModel
+                >>> tim_model = TimModel(filepath="tests/data/input/source-sink/tim-5-columns.tim")
+                >>> tim_model.quantities_names = ["discharge", "waterlevel", "temperature", "salinity", "initialtracer"]
+                >>> print(tim_model.get_units())
+                ['m3/s', 'm', 'C', 'ppt', 'Unknown']
+        """
+        if self.quantities_names is None:
+            return None
+        return self.map_to_units(self.quantities_names)
+
+    @staticmethod
+    def map_to_units(quantities_names: List[str]) -> List[str]:
+        """
+        Maps each quantity in the input list to a specific unit based on its content.
+
+        Args:
+            quantities_names (list of str): A list of strings to be checked for specific keywords.
+
+        Returns:
+            list of str: A list of corresponding units for each input string.
+
+        Examples:
+            >>> from hydrolib.core.dflowfm.tim.models import TimModel
+            >>> quantities_names = ["discharge", "waterlevel", "salinity", "temperature"]
+            >>> TimModel.map_to_units(quantities_names)
+            ['m3/s', 'm', 'ppt', 'C']
+        """
+        # Define the mapping of keywords to units
+        unit_mapping = {
+            "discharge": "m3/s",
+            "waterlevel": "m",
+            "salinity": "ppt",
+            "temperature": "C",
+        }
+
+        # Generate the list of units based on the mapping
+        units = []
+        for string in quantities_names:
+            for keyword, unit in unit_mapping.items():
+                if keyword in string.lower():
+                    units.append(unit)
+                    break
+            else:
+                # Append "Unknown" if no keywords match
+                units.append("Unknown")
+
+        return units
diff --git a/tests/dflowfm/test_tim.py b/tests/dflowfm/test_tim.py
index 14ab7e976..4cc2bee53 100644
--- a/tests/dflowfm/test_tim.py
+++ b/tests/dflowfm/test_tim.py
@@ -118,6 +118,22 @@ def test_initialize_with_quantities_names(self, input_files_dir: Path):
         assert model.quantities_names == ["a"]
         assert len(model.timeseries) == 13
 
+    def test_get_units(self, input_files_dir: Path):
+        """
+        Test the `get_units` method. The method should return the units of the quantities in the timeseries.
+        """
+        path = input_files_dir / "tim/single_data_for_timeseries.tim"
+        model = TimModel(path, quantities_names=["discharge"])
+        assert model.get_units() == ["m3/s"]
+        model.quantities_names = ["waterlevel"]
+        assert model.get_units() == ["m"]
+        model.quantities_names = ["temperature"]
+        assert model.get_units() == ["C"]
+        model.quantities_names = ["salinity"]
+        assert model.get_units() == ["ppt"]
+        model.quantities_names = ["initialtracer-anyname"]
+        assert model.get_units() == ["Unknown"]
+
     def test_as_dataframe(self):
         model = TimModel(
             timeseries=self.single_data_for_timeseries_floats,