added keepdims=False option to project_stereographic

python/damask/util.py

@@ -193,7 +193,7 @@ def scale_to_coprime(v):
     return m
 
 
-def project_stereographic(vector,normalize=False,direction='z'):
+def project_stereographic(vector,direction='z',normalize=True,keepdims=False):
     """
     Apply stereographic projection to vector.
 
@@ -201,23 +201,37 @@ def project_stereographic(vector,normalize=False,direction='z'):
     ----------
     vector : numpy.ndarray of shape (...,3)
         Vector coordinates to be projected.
-    normalize : bool
-        Ensure unit length for vector. Defaults to False.
     direction : str or int
-        Projection direction 'x'|0, 'y'|1, or 'z'|2.
+        Projection direction 'x' | 0, 'y' | 1, or 'z' | 2.
         Defaults to 'z'.
+    normalize : bool
+        Ensure unit length of input vector. Defaults to True.
+    keepdims : bool
+        Maintain three-dimensional output coordinates.
+        Default two-dimensional output uses right-handed frame spanned by
+        the next and next-next axis relative to the projection direction,
+        e.g. x-y when projecting along z and z-x when projecting along y.
 
     Returns
     -------
-    coordinates : numpy.ndarray of shape (...,3)
+    coordinates : numpy.ndarray of shape (...,2 | 3)
         Projected coordinates.
 
+    Examples
+    --------
+    >>> project_stereographic(np.ones(3))
+        [0.3660254, 0.3660254]
+    >>> project_stereographic(np.ones(3),direction='x',normalize=False,keepdims=True)
+        [0, 0.5, 0.5]
+    >>> project_stereographic([0,1,1],direction='y',normalize=True,keepdims=False)
+        [0.41421356, 0]
+
     """
     shift = 2-('xyzXYZ'.index(direction)%3 if isinstance(direction,str) else int(direction))
     v_ = np.roll(vector/np.linalg.norm(vector,axis=-1,keepdims=True) if normalize else vector,
                  shift,axis=-1)
     return np.roll(np.block([v_[...,:2]/(1+np.abs(v_[...,2:3])),np.zeros_like(v_[...,2:3])]),
-                   -shift,axis=-1)
+                   -shift if keepdims else 0,axis=-1)[...,:3 if keepdims else 2]
 
 
 def execution_stamp(class_name,function_name=None):

python/tests/test_util.py

@@ -49,17 +49,18 @@ class TestUtil:
         dist_sampled = np.histogram(centers[selected],bins)[0]/N_samples*np.sum(dist)
         assert np.sqrt(((dist - dist_sampled) ** 2).mean()) < .025 and selected.shape[0]==N_samples
 
-    @pytest.mark.parametrize('point,normalize,direction,answer',
+    @pytest.mark.parametrize('point,direction,normalize,keepdims,answer',
                              [
-                              ([1,0,0],False,'z',[1,0,0]),
-                              ([1,0,0],True, 'z',[1,0,0]),
-                              ([0,1,1],False,'z',[0,0.5,0]),
-                              ([0,1,1],True, 'y',[0,0,0.41421356]),
-                              ([1,1,1],False,'x',[0,0.5,0.5]),
-                              ([1,1,1],True, 'y',[0.3660254, 0,0.3660254]),
+                              ([1,0,0],'z',False,True, [1,0,0]),
+                              ([1,0,0],'z',True, False,[1,0]),
+                              ([0,1,1],'z',False,True, [0,0.5,0]),
+                              ([0,1,1],'y',True, False,[0.41421356,0]),
+                              ([1,1,0],'x',False,False,[0.5,0]),
+                              ([1,1,1],'y',True, True, [0.3660254, 0,0.3660254]),
                              ])
-    def test_project_stereographic(self,point,normalize,direction,answer):
-        assert np.allclose(util.project_stereographic(np.array(point),normalize=normalize,direction=direction),answer)
+    def test_project_stereographic(self,point,direction,normalize,keepdims,answer):
+        assert np.allclose(util.project_stereographic(np.array(point),direction=direction,
+                                                      normalize=normalize,keepdims=keepdims),answer)
 
     @pytest.mark.parametrize('fro,to,mode,answer',
                              [