Merge pull request #266 from psavery/add-to-dense-function

Add `SparseArray.to_dense()`

python/stempy/io/sparse_array.py

@@ -254,6 +254,29 @@ def write_to_hdf5(self, path):
             # Write out the metadata
             save_dict_to_h5(self.metadata, f.require_group('metadata'))
 
+    def to_dense(self):
+        """Create and return a fully dense version of the sparse array
+
+        If the array shape is large, this may cause the system to
+        run out of memory.
+
+        This is equivalent to `array[:]` if `allow_full_expand` is `True`
+        and `sparse_slicing` is `False`.
+
+        :return: the fully dense array
+        :rtype: np.ndarray
+        """
+        prev_allow_full_expand = self.allow_full_expand
+        prev_sparse_slicing = self.sparse_slicing
+
+        self.allow_full_expand = True
+        self.sparse_slicing = False
+        try:
+            return self[:]
+        finally:
+            self.allow_full_expand = prev_allow_full_expand
+            self.sparse_slicing = prev_sparse_slicing
+
     @property
     def shape(self):
         """The full shape of the data (scan shape + frame shape)

tests/conftest.py

@@ -1,4 +1,3 @@
-import copy
 import io
 
 import numpy as np
@@ -70,6 +69,7 @@ def sparse_array_10x10():
 
     return SparseArray(aa, (2, 2), (11, 11))
 
+
 @pytest.fixture
 def sparse_array_small(electron_data_small):
     kwargs = {
@@ -90,15 +90,7 @@ def full_array_small(sparse_array_small):
     global cached_full_array_small
 
     if cached_full_array_small is None:
-        # Don't allow this fixture to modify the other fixture
-        array = copy.deepcopy(sparse_array_small)
-
-        # Have to change these so we won't return a SparseArray,
-        # and allow it to return a fully expanded array
-        array.sparse_slicing = False
-        array.allow_full_expand = True
-
-        cached_full_array_small = array[:]
+        cached_full_array_small = sparse_array_small.to_dense()
 
     return cached_full_array_small
 

tests/test_sparse_array.py

@@ -131,11 +131,7 @@ def test_bin_scans(sparse_array_small, full_array_small):
                                    shape[1] // factor, factor,
                                    shape[2], shape[3]).sum(axis=(1, 3))
 
-        # Expand our binned sparse array
-        binned.allow_full_expand = True
-        binned.sparse_slicing = False
-
-        assert np.array_equal(binned[:], full_binned)
+        assert np.array_equal(binned.to_dense(), full_binned)
 
         binned_with.append(factor)
 
@@ -166,11 +162,7 @@ def test_bin_frames(sparse_array_small, full_array_small):
                                    factor, shape[3] // factor,
                                    factor).sum(axis=(3, 5))
 
-        # Expand our binned sparse array
-        binned.allow_full_expand = True
-        binned.sparse_slicing = False
-
-        assert np.array_equal(binned[:], full_binned)
+        assert np.array_equal(binned.to_dense(), full_binned)
 
         binned_with.append(factor)