From 6916160ed1977979bf89b8b8c7182cb0b236172b Mon Sep 17 00:00:00 2001
From: Philip Eisenlohr <eisenlohr@egr.msu.edu>
Date: Tue, 15 Feb 2022 16:38:02 -0500
Subject: [PATCH] implemented VTK.add(damask.Table)

---
 python/damask/_table.py  |  2 +-
 python/damask/_vtk.py    | 53 +++++++++++++++++++++++++++-------------
 python/tests/test_VTK.py | 18 ++++++++++++++
 3 files changed, 55 insertions(+), 18 deletions(-)

diff --git a/python/damask/_table.py b/python/damask/_table.py
index 1572c4f76..e592a5c7c 100644
--- a/python/damask/_table.py
+++ b/python/damask/_table.py
@@ -326,7 +326,7 @@ class Table:
 
 
     @property
-    def labels(self) -> List[Tuple[int, ...]]:
+    def labels(self) -> List[str]:
         return list(self.shapes)
 
 
diff --git a/python/damask/_vtk.py b/python/damask/_vtk.py
index 16b72e98d..56ad9f35c 100644
--- a/python/damask/_vtk.py
+++ b/python/damask/_vtk.py
@@ -38,6 +38,18 @@ class VTK:
         self.vtk_data = vtk_data
 
 
+    @property
+    def N_points(self) -> int:
+        """Number of points in vtkdata."""
+        return self.vtk_data.GetNumberOfPoints()
+
+
+    @property
+    def N_cells(self) -> int:
+        """Number of cells in vtkdata."""
+        return self.vtk_data.GetNumberOfCells()
+
+
     @staticmethod
     def from_image_data(cells: IntSequence,
                         size: FloatSequence,
@@ -295,7 +307,7 @@ class VTK:
     # Check https://blog.kitware.com/ghost-and-blanking-visibility-changes/ for missing data
     # Needs support for damask.Table
     def add(self,
-            data: Union[np.ndarray, np.ma.MaskedArray],
+            data: Union[np.ndarray, np.ma.MaskedArray, 'Table'],
             label: str = None):
         """
         Add data to either cells or points.
@@ -309,20 +321,17 @@ class VTK:
             Data label.
 
         """
-        N_points = self.vtk_data.GetNumberOfPoints()
-        N_cells  = self.vtk_data.GetNumberOfCells()
 
-        if isinstance(data,np.ndarray):
-            if label is None:
-                raise ValueError('No label defined for numpy.ndarray')
+        def _add_array(self,
+                       data: np.ndarray,
+                       label: str):
 
-            N_data = data.shape[0]
-            data_ = (data if not isinstance(data,np.ma.MaskedArray) else
-                     np.where(data.mask,data.fill_value,data)).reshape(N_data,-1)
+            N_data   = data.shape[0]
 
-            if data_.dtype in [np.double,np.longdouble]:
-                d = np_to_vtk(data_.astype(np.single),deep=True)                                    # avoid large files
-            elif data_.dtype.type is np.str_:
+            data_ = data.reshape(N_data,-1)\
+                        .astype(np.single if data.dtype in [np.double,np.longdouble] else data.dtype)
+
+            if data.dtype.type is np.str_:
                 d = vtk.vtkStringArray()
                 for s in np.squeeze(data_):
                     d.InsertNextValue(s)
@@ -331,14 +340,24 @@ class VTK:
 
             d.SetName(label)
 
-            if   N_data == N_points:
+            if   N_data == self.N_points:
                 self.vtk_data.GetPointData().AddArray(d)
-            elif N_data == N_cells:
+            elif N_data == self.N_cells:
                 self.vtk_data.GetCellData().AddArray(d)
             else:
-                raise ValueError(f'Cell / point count ({N_cells} / {N_points}) differs from data ({N_data}).')
+                raise ValueError(f'Data count mismatch ({N_data} ≠ {self.N_points} & {self.N_cells})')
+
+
+        if isinstance(data,np.ndarray):
+            if label is not None:
+                _add_array(self,
+                           np.where(data.mask,data.fill_value,data) if isinstance(data,np.ma.MaskedArray) else data,
+                           label)
+            else:
+                raise ValueError('No label defined for numpy.ndarray')
         elif isinstance(data,Table):
-            raise NotImplementedError('damask.Table')
+            for l in data.labels:
+                _add_array(self,data.get(l),l)
         else:
             raise TypeError
 
@@ -383,7 +402,7 @@ class VTK:
             # string array
             return np.array([vtk_array.GetValue(i) for i in range(vtk_array.GetNumberOfValues())]).astype(str)
         except UnboundLocalError:
-            raise ValueError(f'Array "{label}" not found.')
+            raise ValueError(f'Array "{label}" not found')
 
 
     @property
diff --git a/python/tests/test_VTK.py b/python/tests/test_VTK.py
index 8f1dc4d3d..40bfbef23 100644
--- a/python/tests/test_VTK.py
+++ b/python/tests/test_VTK.py
@@ -7,7 +7,9 @@ import numpy as np
 import numpy.ma as ma
 import vtk
 
+from collections import defaultdict
 from damask import VTK
+from damask import Table
 from damask import grid_filters
 
 @pytest.fixture
@@ -152,6 +154,22 @@ class TestVTK:
         assert (np.squeeze(data.reshape(N_values,-1)) == default.get('data')).all()
 
 
+    @pytest.mark.parametrize('shapes',[{'scalar':(1,),'vector':(3,),'tensor':(3,3)},
+                                       {'vector':(6,),'tensor':(3,3)},
+                                       {'tensor':(3,3),'scalar':(1,)}])
+    def test_add_table(self,default,shapes):
+        N = np.random.choice([default.N_points,default.N_cells])
+        d = defaultdict(dict)
+        for k,s in shapes.items():
+            d[k]['shape'] = s
+            d[k]['data'] = np.random.random(N*np.prod(s)).reshape((N,-1))
+        t = Table(np.column_stack([d[k]['data'] for k in shapes.keys()]),shapes)
+        default.add(t)
+        for k,s in shapes.items():
+            assert np.allclose(default.get(k).reshape((N,-1)),d[k]['data'],
+                               rtol=1e-7)
+
+
     def test_add_masked(self,default):
         data = np.random.rand(5*6*7,3)
         masked = ma.MaskedArray(data,mask=data<.4,fill_value=42.)