From 2648a67dcdbf14db825ec310fb2c76b1733e0a8b Mon Sep 17 00:00:00 2001
From: Philip Eisenlohr <eisenlohr@egr.msu.edu>
Date: Thu, 5 Dec 2019 12:00:59 -0500
Subject: [PATCH] fixed sort_by to respect updated index (upon subsequent
 adding of new data)

---
 python/damask/table.py     | 60 ++++++++++++++++++++------------------
 python/tests/test_Table.py | 19 ++++++++----
 2 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/python/damask/table.py b/python/damask/table.py
index 9e7b57017..6181fdb1f 100644
--- a/python/damask/table.py
+++ b/python/damask/table.py
@@ -1,4 +1,3 @@
-import random
 import re
 
 import pandas as pd
@@ -24,15 +23,24 @@ class Table():
         self.comments = [] if comments is None else [c for c in comments]
         self.data = pd.DataFrame(data=data)
         self.shapes = shapes
+        self.__label_condensed()
 
+
+    def __label_flat(self):
+        """Label data individually, e.g. v v v ==> 1_v 2_v 3_v."""
+        labels = []
+        for label,shape in self.shapes.items():
+            size = np.prod(shape)
+            labels += ['{}{}'.format('' if size == 1 else '{}_'.format(i+1),label) for i in range(size)]
+        self.data.columns = labels
+        
+        
+    def __label_condensed(self):
+        """Label data condensed, e.g. 1_v 2_v 3_v ==> v v v."""
         labels = []
         for label,shape in self.shapes.items():
             labels += [label] * np.prod(shape)
-
-        if len(labels) != self.data.shape[1]:
-            raise IndexError('Mismatch between shapes and data')
-
-        self.data.rename(columns=dict(zip(range(len(labels)),labels)),inplace=True)
+        self.data.columns = labels
 
 
     def __add_comment(self,label,shape,info):
@@ -87,7 +95,7 @@ class Table():
 
         return Table(data,shapes,comments)
 
-
+    @property
     def labels(self):
         """Return the labels of all columns."""
         return list(self.shapes.keys())
@@ -105,9 +113,11 @@ class Table():
         """
         if re.match(r'[0-9]*?_',label):
             idx,key = label.split('_',1)
-            return self.data[key].to_numpy()[:,int(idx)-1].reshape((-1,1))
+            data = self.data[key].to_numpy()[:,int(idx)-1].reshape((-1,1))
         else: 
-            return self.data[label].to_numpy().reshape((-1,)+self.shapes[label])          # better return shape (N) instead of (N,1), i.e. no reshaping?
+            data = self.data[label].to_numpy().reshape((-1,)+self.shapes[label])
+
+        return data.astype(type(data.flatten()[0]))
 
 
     def set(self,label,data,info=None):
@@ -149,13 +159,14 @@ class Table():
 
         """
         self.__add_comment(label,data.shape[1:],info)
-        
+
         self.shapes[label] = data.shape[1:] if len(data.shape) > 1 else (1,)
-        size = np.prod(data.shape[1:],dtype=int) 
-        self.data = pd.concat([self.data,
-                               pd.DataFrame(data=data.reshape(-1,size),
-                                            columns=[label]*size)],
-                              axis=1)
+        size = np.prod(data.shape[1:],dtype=int)
+        new = pd.DataFrame(data=data.reshape(-1,size),
+                           columns=[label]*size,
+                          )
+        new.index = self.data.index
+        self.data = pd.concat([self.data,new],axis=1)
 
 
     def delete(self,label):
@@ -201,24 +212,15 @@ class Table():
 
         Parameters
         ----------
-        label : list of str or str
+        label : str or list
             Column labels.
-        ascending : bool, optional
+        ascending : bool or list, optional
             Set sort order.
 
         """
-        _temp = []
-        _labels = []
-        for label in labels if isinstance(labels,list) else [labels]:
-            if re.match(r'[0-9]*?_',label):
-                _temp.append(str(random.getrandbits(128)))
-                self.add(_temp[-1],self.get(label))
-                _labels.append(_temp[-1])
-            else: 
-                _labels.append(label)
-
-        self.data.sort_values(_labels,axis=0,inplace=True,ascending=ascending)
-        for t in _temp: self.delete(t)
+        self.__label_flat()
+        self.data.sort_values(labels,axis=0,inplace=True,ascending=ascending)
+        self.__label_condensed()
         self.comments.append('sorted by [{}]'.format(', '.join(labels)))
 
 
diff --git a/python/tests/test_Table.py b/python/tests/test_Table.py
index b17348881..a0dc31975 100644
--- a/python/tests/test_Table.py
+++ b/python/tests/test_Table.py
@@ -9,7 +9,7 @@ from damask import Table
 @pytest.fixture
 def default():
     """Simple Table."""
-    x = np.ones((5,13))
+    x = np.ones((5,13),dtype=float)
     return Table(x,{'F':(3,3),'v':(3,),'s':(1,)},['test data','contains only ones'])
 
 @pytest.fixture
@@ -58,7 +58,7 @@ class TestTable:
         assert np.allclose(d,0.0) and d.shape[1:] == (3,3)
 
     def test_labels(self,default):
-        assert default.labels() == ['F','v','s']
+        assert default.labels == ['F','v','s']
         
     def test_add(self,default):
         d = np.random.random((5,9))
@@ -82,9 +82,9 @@ class TestTable:
             default.get('v')
 
 
-    def test_invalid_initialization(self,default):
-        x = default.get('v')
-        with pytest.raises(IndexError):
+    def test_invalid_initialization(self):
+        x = np.random.random((5,10))
+        with pytest.raises(ValueError):
             Table(x,{'F':(3,3)})
 
     def test_invalid_set(self,default):
@@ -115,7 +115,14 @@ class TestTable:
     def test_sort_revert(self):
         x = np.random.random((5,12))
         t = Table(x,{'F':(3,3),'v':(3,)},['random test data'])
-        t.sort_by('4_F',False)
+        t.sort_by('4_F',ascending=False)
         sort = t.get('4_F')
         assert np.all(np.sort(sort,0)==sort[::-1,:])
 
+    def test_sort(self):
+        t = Table(np.array([[0,1,],[2,1,]]),
+                  {'v':(2,)},
+                  ['test data'])
+        t.add('s',np.array(['b','a']))
+        t.sort_by('s')
+        assert np.all(t.get('1_v') == np.array([2,0]).reshape((2,1)))