combine table.add and .set

python/damask/_table.py

@@ -363,14 +363,14 @@ class Table:
             data: np.ndarray,
             info: str = None) -> 'Table':
         """
-        Set column data.
+        Add new or replace existing column data.
 
         Parameters
         ----------
         label : str
             Column label.
         data : numpy.ndarray
-            Replacement data.
+            Column data.
         info : str, optional
             Human-readable information about the modified data.
 
@@ -382,49 +382,32 @@ class Table:
         """
         dup = self.copy()
         dup._add_comment(label, data.shape[1:], info)
+
         if m := re.match(r'(.*)\[((\d+,)*(\d+))\]',label):
             key = m.group(1)
-            idx = np.ravel_multi_index(tuple(map(int,m.group(2).split(","))),
-                                       self.shapes[key])
-            iloc = dup.data.columns.get_loc(key).tolist().index(True) + idx
-            dup.data.iloc[:,iloc] = data
         else:
-            dup.data[label]       = data.reshape(dup.data[label].shape)
-        return dup
+            key = label
 
+        if key in dup.shapes:
 
-    def add(self,
-            label: str,
-            data: np.ndarray,
-            info: str = None) -> 'Table':
-        """
-        Add column data.
+            if m:
+                idx = np.ravel_multi_index(tuple(map(int,m.group(2).split(","))),
+                                           self.shapes[key])
+                iloc = dup.data.columns.get_loc(key).tolist().index(True) + idx
+                dup.data.iloc[:,iloc] = data
+            else:
+                dup.data[label]       = data.reshape(dup.data[label].shape)
 
-        Parameters
-        ----------
-        label : str
-            Column label.
-        data : numpy.ndarray
-            New data.
-        info : str, optional
-            Human-readable information about the new data.
+        else:
 
-        Returns
-        -------
-        updated : damask.Table
-            Updated table.
+            dup.shapes[label] = data.shape[1:] if len(data.shape) > 1 else (1,)
+            size = np.prod(data.shape[1:],dtype=int)
+            new = pd.DataFrame(data=data.reshape(-1,size),
+                               columns=[label]*size,
+                              )
+            new.index = dup.data.index
+            dup.data = pd.concat([dup.data,new],axis=1)
 
-        """
-        dup = self.copy()
-        dup._add_comment(label,data.shape[1:],info)
-
-        dup.shapes[label] = data.shape[1:] if len(data.shape) > 1 else (1,)
-        size = np.prod(data.shape[1:],dtype=int)
-        new = pd.DataFrame(data=data.reshape(-1,size),
-                           columns=[label]*size,
-                          )
-        new.index = dup.data.index
-        dup.data = pd.concat([dup.data,new],axis=1)
         return dup
 
 

python/tests/test_Grid.py

@@ -441,7 +441,7 @@ class TestGrid:
         z = np.ones(cells.prod())
         z[cells[:2].prod()*int(cells[2]/2):] = 0
         t = Table({'coords':3,'z':1},np.column_stack((coords,z)))
-        t = t.add('indicator',t.get('coords')[:,0])
+        t = t.set('indicator',t.get('coords')[:,0])
         g = Grid.from_table(t,'coords',['indicator','z'])
         assert g.N_materials == g.cells[0]*2 and (g.material[:,:,-1]-g.material[:,:,0] == cells[0]).all()
 

python/tests/test_Orientation.py

@@ -169,7 +169,7 @@ class TestOrientation:
     def test_from_directions(self,kwargs):
         for a,b in np.random.random((10,2,3)):
             c = np.cross(b,a)
-            if np.all(np.isclose(c,0)): continue
+            if np.allclose(c,0): continue
             o = Orientation.from_directions(uvw=a,hkl=c,**kwargs)
             x = o.to_pole(uvw=a)
             z = o.to_pole(hkl=c)

python/tests/test_Table.py

@@ -51,7 +51,7 @@ class TestTable:
 
     def test_add(self,default):
         d = np.random.random((5,9))
-        assert np.allclose(d,default.add('nine',d,'random data').get('nine'))
+        assert np.allclose(d,default.set('nine',d,'random data').get('nine'))
 
     def test_isclose(self,default):
         assert default.isclose(default).all()
@@ -200,6 +200,6 @@ class TestTable:
         t = Table({'v':(2,)},
                   np.array([[0,1,],[2,1,]]),
                   ['test data'])\
-            .add('s',np.array(['b','a']))\
+            .set('s',np.array(['b','a']))\
             .sort_by('s')
         assert np.all(t.get('v')[:,0] == np.array([2,0]))