fix minor bug in new getitem on dof axis

evalf · Feb 7, 2015 · f8d38f0 · f8d38f0
1 parent e99f1f5
commit f8d38f0
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/nutils/function.py b/nutils/function.py
@@ -671,7 +671,11 @@ def _add( self, other ):
       return align( self.func + other.func, self.axes, self.ndim )
 
   def _take( self, indices, axis ):
-    n = self.axes.index( axis )
+    try:
+      n = self.axes.index( axis )
+    except ValueError:
+      assert isinstance( indices, DofMap )
+      return self
     return align( take( self.func, indices, n ), self.axes, self.ndim )
 
   def _opposite( self ):
@@ -3100,6 +3104,8 @@ def take( arg, index, axis ):
     retval = _call( arg, '_take', index, axis )
     if retval is not None:
       return retval
+    if arg.shape[axis] == 1:
+      return arg
     return DofIndex( arg, axis, index )
 
   if isinstance( index, slice ):