Bugfix: Small Bugfix for polar flatten simulation

Author: CSSFrancis

diffsims/simulations/simulation2d.py

@@ -327,14 +327,27 @@ def polar_flatten_simulations(self, radial_axes=None, azimuthal_axes=None):
         intensities_templates = np.zeros((len(flattened_vectors), max_num_spots))
         for i, v in enumerate(flattened_vectors):
             r, t = v.to_flat_polar()
+            inten = v.intensity
             if radial_axes is not None and azimuthal_axes is not None:
-                r = get_closest(radial_axes, r)
-                t = get_closest(azimuthal_axes, t)
-                r = r[r < len(radial_axes)]
-                t = t[t < len(azimuthal_axes)]
-            r_templates[i, : len(r)] = r
-            theta_templates[i, : len(t)] = t
-            intensities_templates[i, : len(v.intensity)] = v.intensity
+                r = get_closest(
+                    radial_axes, r
+                )  # convert from real to pixel coordinates
+                t = get_closest(
+                    azimuthal_axes, t
+                )  # convert from real to pixel coordinates
+                mask = r < len(radial_axes) - 1  # if out of bounds, ignore
+                r = r[mask]  # if out of bounds, ignore
+                t = t[mask]  # if out of bounds, ignore
+                inten = inten[mask]  # if out of bounds, ignore
+                mask = t < len(azimuthal_axes) - 1  # if out of bounds, ignore
+                t = t[mask]  # if out of bounds, ignore
+                r = r[mask]  # if out of bounds, ignore
+                inten = inten[mask]  # if out of bounds, ignore
+            r_templates[i, : len(r)] = (
+                r  # set the r coordinates (len r and t should be the same)
+            )
+            theta_templates[i, : len(r)] = t  # set the theta coordinates
+            intensities_templates[i, : len(inten)] = inten
         if radial_axes is not None and azimuthal_axes is not None:
             r_templates = np.array(r_templates, dtype=int)
             theta_templates = np.array(theta_templates, dtype=int)

diffsims/tests/simulations/test_simulations2d.py

@@ -46,7 +46,20 @@ class TestSingleSimulation:
     def single_simulation(self, al_phase):
         gen = SimulationGenerator(accelerating_voltage=200)
         rot = Rotation.from_axes_angles([1, 0, 0], 45, degrees=True)
-        coords = DiffractingVector(phase=al_phase, xyz=[[1, 0, 0]])
+        coords = DiffractingVector(
+            phase=al_phase,
+            xyz=[
+                [1, 0, 0],
+                [2, 0, 0],
+                [3, 3, 0],
+                [-4, 0, 0],
+                [-5, 0, 0],
+                [-6, 0, 0],
+                [-7, 0, 0],
+                [-8, 0, 0],
+            ],
+            intensity=[1, 2, 3, 4, 5, 6, 7, 8],
+        )
         sim = Simulation2D(
             phases=al_phase, simulation_generator=gen, coordinates=coords, rotations=rot
         )
@@ -91,12 +104,12 @@ def test_polar_flatten(self, single_simulation):
             theta_templates,
             intensities_templates,
         ) = single_simulation.polar_flatten_simulations()
-        assert r_templates.shape == (1, 1)
-        assert theta_templates.shape == (1, 1)
-        assert intensities_templates.shape == (1, 1)
+        assert r_templates.shape == (1, 8)
+        assert theta_templates.shape == (1, 8)
+        assert intensities_templates.shape == (1, 8)
 
     def test_polar_flatten_axes(self, single_simulation):
-        radial_axes = np.linspace(0, 1, 10)
+        radial_axes = np.linspace(0, 7, 5)
         theta_axes = np.linspace(0, 2 * np.pi, 10)
         (
             r_templates,
@@ -105,9 +118,13 @@ def test_polar_flatten_axes(self, single_simulation):
         ) = single_simulation.polar_flatten_simulations(
             radial_axes=radial_axes, azimuthal_axes=theta_axes
         )
-        assert r_templates.shape == (1, 1)
-        assert theta_templates.shape == (1, 1)
-        assert intensities_templates.shape == (1, 1)
+        assert r_templates.shape == (1, 8)
+        assert theta_templates.shape == (1, 8)
+        assert intensities_templates.shape == (1, 8)
+        # The last 2 elements should be zero
+        np.testing.assert_array_equal(r_templates[:, 6:], 0)
+        np.testing.assert_array_equal(theta_templates[:, 6:], 0)
+        np.testing.assert_array_equal(intensities_templates[:, 6:], 0)
 
     def test_deepcopy(self, single_simulation):
         copied = single_simulation.deepcopy()