fix Dirac basis evaluate_t

Author: chris-langfield

src/aspire/basis/dirac.py

@@ -67,14 +67,17 @@ def _evaluate_t(self, x):
              `self.count` and whose remaining dimensions correspond to
              higher dimensions of `v`.
         """
-        x = np.squeeze(x.asnumpy())
+        import pdb
+
+        pdb.set_trace()
+        x = m_reshape(x.asnumpy(), new_shape=self.ndim * (self.nres,) + (x.shape[0],))
         x, sz_roll = unroll_dim(x, self.ndim + 1)
         x = m_reshape(x, new_shape=(self._sz_prod,) + x.shape[self.ndim :])
         v = np.zeros(shape=(self.count,) + x.shape[1:], dtype=self.dtype)
         v = x[self._mask, ...]
         v = roll_dim(v, sz_roll)
 
-        return v
+        return np.squeeze(v)
 
     def expand(self, x):
         return self.evaluate_t(x)

tests/test_Diracbasis.py

@@ -4,7 +4,7 @@
 
 from aspire.basis import DiracBasis
 from aspire.image import Image
-from aspire.utils.matlab_compat import m_flatten
+from aspire.utils.matlab_compat import m_flatten, m_reshape
 
 
 class DiracBasisTestCase(TestCase):
@@ -192,13 +192,23 @@ def testDiracEvaluate_t(self):
                 ],
             ]
         )
+        # First test single image
         result = self.basis.evaluate_t(x)
-
         # evaluate_t should return a NumPy array
         self.assertTrue(isinstance(result, np.ndarray))
-
+        # the result should be a flattened array of the values of x
+        # in particular, for one image, its shape should be (size*size,)
+        # not (size*size, 1)
         self.assertTrue(np.allclose(result, m_flatten(x)))
 
+        # Now test a stack of images
+        stack = np.array([x] * 10)
+        result_stack = self.basis.evaluate_t(stack)
+        # the result should be of the shape (size*size, 10)
+        flat_x = m_flatten(x)
+        compare_array = m_reshape(np.array([flat_x] * 10), (self.basis.nres**2, 10))
+        self.assertTrue(np.allclose(result_stack, compare_array))
+
     def testInitWithIntSize(self):
         # make sure we can instantiate with just an int as a shortcut
         self.assertEqual((8, 8), DiracBasis(8).sz)