Coverage for local_installation_linux/mumott/pipelines/utilities/alignment_geometry.py: 91%

1import numpy as np 

2 

3from mumott.data_handling import DataContainer 

4from mumott import Geometry 

5 

6 

7def get_alignment_geometry(data_container: DataContainer) -> tuple[int, np.ndarray[int]]: 

8 """ Define the index of the main rotation axis relative to the data, and the related 

9 tomogram volume. 

10 

11 Parameters 

12 ---------- 

13 data_container 

14 A :class:'DataContainer <mumott.data_handling.DataContainer>' instance. 

15 

16 Returns 

17 ------- 

18 a tuple comprising the index of the main rotation axis (0 or 1) and 

19 the volume of the tomogram related to the data and the main rotation axis 

20 

21 """ 

22 volume_deduced = data_container.geometry.volume_shape 

23 # check that dimensions of the volume match with the projection 

24 # or raise error and stop 

25 projection_shape = data_container.geometry.projection_shape 

26 if not all(element in volume_deduced for element in projection_shape): 26 ↛ 27line 26 didn't jump to line 27, because the condition on line 26 was never true

27 raise ValueError('data_container.geometry.volume_shape and' 

28 ' data_container.geometry.projection_shape must' 

29 ' match in shape') 

30 # abs in case of negative axis 

31 geom_inner_axes_index = np.where( 

32 np.isclose(np.abs(data_container.geometry.inner_axes[0]), 1) 

33 )[0][0] 

34 

35 data_inner_axes_index = np.where( 

36 data_container.diode.shape == 

37 data_container.geometry.volume_shape[geom_inner_axes_index] 

38 )[0][0] - 1 

39 

40 main_rot_axis_deduced = data_inner_axes_index 

41 

42 return main_rot_axis_deduced, volume_deduced 

43 

44 

45def shift_center_of_reconstruction(geometry: Geometry, 

46 shift_vector: tuple[float] = (0., 0., 0.)) -> None: 

47 """ This utility function will shift the ``offsets`` in the 

48 :class:`geometry <mumott.core.geometry.Geometry>` based on a three-dimensional 

49 shift vector. Use this to reposition the reconstruction within the volume. 

50 

51 Parameters 

52 ---------- 

53 geometry 

54 A :class:`Geometry <mumott.core.geometry.Geometry>` instance. 

55 Its `j_offsets` and `k_offsets` are modified in-place. 

56 shift_vector 

57 A ``tuple`` that indicates the direction and magnitude of the 

58 desired shift in ``(x, y, z)``, in units of voxels. 

59 

60 Example 

61 ------- 

62 >>> import numpy as np 

63 >>> from scipy.spatial.transform import Rotation 

64 >>> from mumott.core.geometry import Geometry, GeometryTuple 

65 >>> from mumott.pipelines.utilities.alignment_geometry import shift_center_of_reconstruction 

66 >>> geo = Geometry() 

67 >>> geo.append(GeometryTuple(rotation=np.eye(3), j_offset=0., k_offset=0.)) 

68 >>> geo.append(GeometryTuple(rotation=Rotation.from_euler('y', np.pi/4).as_matrix(), 

69 j_offset=0., 

70 k_offset=0.)) 

71 >>> geo.append(GeometryTuple(rotation=Rotation.from_euler('z', np.pi/4).as_matrix(), 

72 j_offset=0., 

73 k_offset=0.)) 

74 >>> print(geo.j_direction_0) 

75 [1. 0. 0.] 

76 >>> print(geo.k_direction_0) 

77 [0. 0. 1.] 

78 >>> shift_center_of_reconstruction(geo, shift_vector=(-5., 1.7, 3.)) 

79 >>> print(geo[0].j_offset, geo[0].k_offset) 

80 -5.0 3.0 

81 >>> print(geo[1].j_offset, geo[1].k_offset) 

82 -1.4142... 5.6568... 

83 >>> print(geo[2].j_offset, geo[2].k_offset) 

84 -4.737... 3.0 

85 """ 

86 j_vectors = np.einsum( 

87 'kij,i->kj', 

88 geometry.rotations_as_array, 

89 geometry.j_direction_0) 

90 k_vectors = np.einsum( 

91 'kij,i->kj', 

92 geometry.rotations_as_array, 

93 geometry.k_direction_0) 

94 shifts_j = np.einsum('ij, j', j_vectors, shift_vector) 

95 shifts_k = np.einsum('ij, j', k_vectors, shift_vector) 

96 geometry.j_offsets += shifts_j 

97 geometry.k_offsets += shifts_k