condensing code

python/damask/table.py

@@ -21,27 +21,27 @@ class Table():
             Additional, human-readable information.
         
         """
+        self.comments = [] if comments is None else [c for c in comments]
         self.data = pd.DataFrame(data=data)
-
-        labels = {}
-        i = 0
-        for label in shapes.keys():
-            for components in range(np.prod(shapes[label])):
-                labels[i] = label
-                i+=1
-
-        if i != self.data.shape[1]:
-            raise IndexError('Shape mismatch between shapes and data')
-
-        self.data.rename(columns=labels,inplace=True)
-
-        if comments is None:
-            self.comments = []
-        else:
-            self.comments = [c for c in comments]
-
         self.shapes = shapes
 
+        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)
+
+
+    def __add_comment(self,label,shape,info):
+        if info is not None:
+            self.comments.append('{}{}: {}'.format(label,
+                                                   ' '+str(shape) if np.prod(shape,dtype=int) > 1 else '',
+                                                   info))
+
+    
     @staticmethod
     def from_ASCII(fname):
         """
@@ -67,7 +67,7 @@ class Table():
             header = int(header)
         else:
             raise Exception
-        comments = [f.readline()[:-1] for i in range(header-1)]
+        comments = [f.readline()[:-1] for i in range(1,header)]
         labels   = f.readline().split()
        
         shapes = {}
@@ -83,7 +83,7 @@ class Table():
                 else:
                     shapes[label] = (1,)
         
-        data = pd.read_csv(f,names=[i for i in range(len(labels))],sep=r'\s+').to_numpy()
+        data = pd.read_csv(f,names=list(range(len(labels))),sep=r'\s+').to_numpy()
 
         return Table(data,shapes,comments)
 
@@ -107,7 +107,8 @@ class Table():
             idx,key = label.split('_',1)
             return self.data[key].to_numpy()[:,int(idx)-1].reshape((-1,1))
         else: 
-            return self.data[label].to_numpy().reshape((-1,)+self.shapes[label])
+            return self.data[label].to_numpy().reshape((-1,)+self.shapes[label])          # better return shape (N) instead of (N,1), i.e. no reshaping?
+
 
     def set(self,label,data,info=None):
         """
@@ -123,11 +124,7 @@ class Table():
             Human-readable information about the new data.
 
         """
-        if info is not None:
-            if np.prod(data.shape[1:],dtype=int) == 1: 
-                self.comments.append('{}: {}'.format(label,info))
-            else:
-                self.comments.append('{} {}: {}'.format(label,data.shape[1:],info))
+        self.__add_comment(label,data.shape[1:],info)
 
         if re.match(r'[0-9]*?_',label):
             idx,key = label.split('_',1)
@@ -136,6 +133,7 @@ class Table():
         else: 
             self.data[label]       = data.reshape(self.data[label].shape)
 
+
     def add(self,label,data,info=None):
         """
         Add column data.
@@ -150,17 +148,15 @@ class Table():
             Human-readable information about the modified data.
 
         """
-        if info is not None:
-           if np.prod(data.shape[1:],dtype=int) == 1: 
-               self.comments.append('{}: {}'.format(label,info))
-           else:
-               self.comments.append('{} {}: {}'.format(label,data.shape[1:],info))
+        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) 
-        new_data = pd.DataFrame(data=data.reshape(-1,size),
-                                columns=[label for l in range(size)])
-        self.data = pd.concat([self.data,new_data],axis=1)
+        self.data = pd.concat([self.data,
+                               pd.DataFrame(data=data.reshape(-1,size),
+                                            columns=[label]*size)],
+                              axis=1)
+
 
     def delete(self,label):
         """
@@ -176,6 +172,7 @@ class Table():
 
         del self.shapes[label]
 
+
     def rename(self,label_old,label_new,info=None):
         """
         Rename column data.
@@ -190,9 +187,10 @@ class Table():
         """
         self.data.rename(columns={label_old:label_new},inplace=True)
 
-        comments = '{} => {}'.format(label_old,label_new)
-        comments += ': {}'.format(info) if info is not None else ''
-        self.comments.append(comments)
+        self.comments.append('{} => {}{}'.format(label_old,
+                                                 label_new,
+                                                 '' if info is None else ': {}'.format(info),
+                                                 ))
 
         self.shapes[label_new] = self.shapes.pop(label_old)
 
@@ -223,6 +221,7 @@ class Table():
         for t in _temp: self.delete(t)
         self.comments.append('sorted by [{}]'.format(', '.join(labels)))
 
+
     def to_ASCII(self,fname):
         """
         Store as plain text file.