From a142185148fd64648e0beca7d67c71853696fc79 Mon Sep 17 00:00:00 2001
From: vhirtham <volker.hirthammer@bam.de>
Date: Fri, 16 Jul 2021 12:45:05 +0200
Subject: [PATCH] 388 timeseries init (#429)

---
 CHANGELOG.md             |  3 ++
 weldx/core.py            | 59 +++++++++++++++++++++++++++-------------
 weldx/tests/test_core.py | 25 +++++++++++++++++
 3 files changed, 68 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ff8f2cb67..28c24587a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,9 @@
 
 ### changes
 
+- `TimeSeries.__init__` accepts `xarray.DataArray` as `data`
+  parameter [[#429]](https://github.com/BAMWelDX/weldx/pull/429)
+
 ### fixes
 
 ### documentation
diff --git a/weldx/core.py b/weldx/core.py
index 595381638..b760ea693 100644
--- a/weldx/core.py
+++ b/weldx/core.py
@@ -295,7 +295,7 @@ def __init__(
         self._units = None
         self._interp_counter = 0
 
-        if isinstance(data, pint.Quantity):
+        if isinstance(data, (pint.Quantity, xr.DataArray)):
             self._initialize_discrete(data, time, interpolation)
         elif isinstance(data, MathematicalExpression):
             self._init_expression(data)
@@ -346,9 +346,25 @@ def __repr__(self):
             )
         return representation + f"Units:\n\t{self.units}\n"
 
+    @staticmethod
+    def _check_data_array(data_array: xr.DataArray):
+        """Raise an exception if the 'DataArray' can't be used as 'self._data'."""
+        try:
+            ut.xr_check_coords(data_array, dict(time={"dtype": ["timedelta64[ns]"]}))
+        except (KeyError, TypeError, ValueError) as e:
+            raise type(e)(
+                "The provided 'DataArray' does not match the required pattern. It "
+                "needs to have a dimension called 'time' with coordinates of type "
+                "'timedelta64[ns]'. The error reported by the comparison function was:"
+                f"\n{e}"
+            )
+
+        if not isinstance(data_array.data, pint.Quantity):
+            raise TypeError("The data of the 'DataArray' must be a 'pint.Quantity'.")
+
     def _initialize_discrete(
         self,
-        data: pint.Quantity,
+        data: Union[pint.Quantity, xr.DataArray],
         time: Union[None, pd.TimedeltaIndex, pint.Quantity],
         interpolation: str,
     ):
@@ -357,24 +373,29 @@ def _initialize_discrete(
         if interpolation is None:
             interpolation = "step"
 
-        # expand dim for scalar input
-        data = Q_(data)
-        if not np.iterable(data):
-            data = np.expand_dims(data, 0)
-
-        # constant value case
-        if time is None:
-            time = pd.TimedeltaIndex([0])
-
-        if isinstance(time, pint.Quantity):
-            time = ut.to_pandas_time_index(time)
-        if not isinstance(time, pd.TimedeltaIndex):
-            raise ValueError(
-                '"time" must be a time quantity or a "pandas.TimedeltaIndex".'
-            )
+        if isinstance(data, xr.DataArray):
+            self._check_data_array(data)
+            data = data.transpose("time", ...)
+            self._data = data
+        else:
+            # expand dim for scalar input
+            data = Q_(data)
+            if not np.iterable(data):
+                data = np.expand_dims(data, 0)
+
+            # constant value case
+            if time is None:
+                time = pd.TimedeltaIndex([0])
+
+            if isinstance(time, pint.Quantity):
+                time = ut.to_pandas_time_index(time)
+            if not isinstance(time, pd.TimedeltaIndex):
+                raise ValueError(
+                    '"time" must be a time quantity or a "pandas.TimedeltaIndex".'
+                )
 
-        dax = xr.DataArray(data=data)
-        self._data = dax.rename({"dim_0": "time"}).assign_coords({"time": time})
+            dax = xr.DataArray(data=data)
+            self._data = dax.rename({"dim_0": "time"}).assign_coords({"time": time})
         self.interpolation = interpolation
 
     def _init_expression(self, data):
diff --git a/weldx/tests/test_core.py b/weldx/tests/test_core.py
index 2c4f7be7f..84fbb1851 100644
--- a/weldx/tests/test_core.py
+++ b/weldx/tests/test_core.py
@@ -5,6 +5,7 @@
 import pandas as pd
 import pint
 import pytest
+import xarray as xr
 
 import weldx.util as ut
 from weldx.constants import WELDX_QUANTITY as Q_
@@ -292,6 +293,30 @@ def test_construction_expression(data, shape_exp, unit_exp):
         assert ts.data_array is None
         assert Q_(1, unit_exp).check(UREG.get_dimensionality(ts.units))
 
+    # test_init_data_array -------------------------------------------------------------
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        "data, dims, coords, exception_type",
+        [
+            (Q_([1, 2, 3], "m"), "time", dict(time=TDI([1, 2, 3])), None),
+            (Q_([1, 2, 3], "m"), "a", dict(a=TDI([1, 2, 3])), KeyError),
+            (Q_([[1, 2]], "m"), ("a", "time"), dict(a=[2], time=TDI([1, 2])), None),
+            (Q_([1, 2, 3], "m"), "time", None, KeyError),
+            (Q_([1, 2, 3], "m"), "time", dict(time=[1, 2, 3]), TypeError),
+            ([1, 2, 3], "time", dict(time=TDI([1, 2, 3])), TypeError),
+        ],
+    )
+    def test_init_data_array(data, dims, coords, exception_type):
+        """Test the `__init__` method with an xarray as data parameter."""
+        da = xr.DataArray(data=data, dims=dims, coords=coords)
+        if exception_type is not None:
+            with pytest.raises(exception_type):
+                TimeSeries(da)
+        else:
+            ts = TimeSeries(da)
+            assert ts.data_array.dims[0] == "time"
+
     # test_construction_exceptions -----------------------------------------------------
 
     values_def = Q_([5, 7, 3, 6, 8], "m")