From 3ee94aa9ca58bc0bcf027862c0391034c59921cc Mon Sep 17 00:00:00 2001
From: phillip <pc494@cam.ac.uk>
Date: Thu, 3 Sep 2020 12:23:49 +0100
Subject: [PATCH] Updates and clean up of gridding functionality

---
 .../generators/rotation_list_generators.py    | 28 +++++++++++--------
 .../test_rotation_list_generator.py           | 14 ++++------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/diffsims/generators/rotation_list_generators.py b/diffsims/generators/rotation_list_generators.py
index 1f8407ee..071774be 100644
--- a/diffsims/generators/rotation_list_generators.py
+++ b/diffsims/generators/rotation_list_generators.py
@@ -25,6 +25,9 @@
 from itertools import product
 
 from orix.sampling.sample_generators import get_sample_fundamental, get_sample_local
+from orix.quaternion.rotation import Rotation
+from orix.vector.neo_euler import AxAngle
+from orix.vector import Vector3d
 
 from transforms3d.euler import euler2axangle, axangle2euler
 from transforms3d.euler import axangle2euler, euler2axangle, euler2mat
@@ -65,7 +68,7 @@ def get_list_from_orix(grid, rounding=2):
     rotation_list = z.data.tolist()
     i = 0
     while i < len(rotation_list):
-        rotation_list[i] = tuple(np.round(rotation_list[i], decimals=rounding))
+        rotation_list[i] = tuple(np.round(np.rad2deg(rotation_list[i]), decimals=rounding))
         i += 1
 
     return rotation_list
@@ -103,18 +106,19 @@ def get_local_grid(resolution=2, center=None, grid_width=10):
     ----------
     resolution : float, optional
         The characteristic distance between a rotation and its neighbour (degrees)
-    center : orix.quaternion.rotation.Rotation, optional
+    center : euler angle tuple or orix.quaternion.rotation.Rotation, optional
         The rotation at which the grid is centered. If None (default) uses the identity
     grid_width : float, optional
         The largest angle of rotation away from center that is acceptable (degrees)
 
-    See Also
-    --------
-
     Returns
     -------
     rotation_list : list of tuples
     """
+    if isinstance(center,tuple):
+        z = np.deg2rad(np.asarray(center))
+        center = Rotation.from_euler(z,convention="bunge",direction="crystal2lab")
+
     orix_grid = get_sample_local(
         resolution=resolution, center=center, grid_width=grid_width
     )
@@ -147,13 +151,15 @@ def get_grid_around_beam_direction(beam_rotation, resolution, angular_range=(0,
     >>> beam_rotation = get_rotation_from_z_to_direction(structure,[1,1,1])
     >>> grid = get_grid_around_beam_direction(beam_rotation,1)
     """
+    z = np.deg2rad(np.asarray(beam_rotation))
+    beam_rotation = Rotation.from_euler(z,convention="bunge",direction="crystal2lab")
 
-    beam_rotation = np.deg2rad(beam_rotation)
-    axangle = euler2axangle(
-        beam_rotation[0], beam_rotation[1], beam_rotation[2], "rzxz"
-    )
-    rotation_list = None
-    raise NotImplementedError("This functionality will be (re)added in future")
+    angles = np.deg2rad(np.arange(start=angular_range[0],stop=angular_range[1],step=resolution))
+    axes   = np.repeat([[0,0,1]],angles.shape[0],axis=0)
+    in_plane_rotation = Rotation.from_neo_euler(AxAngle.from_axes_angles(axes,angles))
+
+    orix_grid = beam_rotation * in_plane_rotation
+    rotation_list = get_list_from_orix(orix_grid,rounding=2)
     return rotation_list
 
 
diff --git a/diffsims/tests/test_generators/test_rotation_list_generator.py b/diffsims/tests/test_generators/test_rotation_list_generator.py
index 3af67b4f..b3704dee 100644
--- a/diffsims/tests/test_generators/test_rotation_list_generator.py
+++ b/diffsims/tests/test_generators/test_rotation_list_generator.py
@@ -30,7 +30,7 @@
     "grid",
     [
         pytest.param(
-            get_local_grid(resolution=30, center=None, grid_width=35),
+            get_local_grid(resolution=30, center=(0,1,0), grid_width=35),
             marks=pytest.mark.xfail(reason="Downstream bug"),
         ),
         get_fundamental_zone_grid(space_group=20, resolution=20),
@@ -41,14 +41,12 @@ def test_get_grid(grid):
     assert len(grid) > 0
     assert isinstance(grid[0], tuple)
 
-
-@pytest.mark.xfail(reason="Functionality removed")
 def test_get_grid_around_beam_direction():
-    grid_simple = get_grid_around_beam_direction([1, 1, 1], 1, (0, 360))
-    assert isinstance(grid_simple, list)
-    assert isinstance(grid_simple[0], tuple)
-    assert len(grid_simple) == 360
-
+    grid = get_grid_around_beam_direction((0,90,0),resolution=2,angular_range=(0,9))
+    assert isinstance(grid, list)
+    assert isinstance(grid[0], tuple)
+    assert len(grid) == 5 # should have 0,2,4,6 and 8
+    assert np.allclose([x[1] for x in grid],90) #taking z to y
 
 @pytest.mark.parametrize(
     "crystal_system",