v3, but pass all arguments to local dendrogram and then skip pruning step

src/dendro/distributed_dendrogram_v3.py

@@ -6,6 +6,9 @@
 from astrodendro.dendrogram import Dendrogram, _make_trunk
 from astrodendro import pruning
 
+from astrodendro.dendrogram import _sorted_by_idx, AnimatedProgressBar
+from astrodendro.structure import Structure as astrodendro_structure
+
 from dendro.distributed_dendrogram import Structure
 
 
@@ -99,7 +102,7 @@ def make_output_astrodendro_compatible(self, is_independent):
     def _compute_single_local_dendrogram(self, local_data, split_dim=0, **kwargs):
 
         t0 = perf_counter()
-        local_dendrogram = Dendrogram.compute(local_data, **kwargs)
+        local_dendrogram = LocalDendrogram.compute(local_data, **kwargs)
         t1 = perf_counter()
         self.time_local_dendrogram = t1 - t0
         self.logger.info(
@@ -554,3 +557,273 @@ def get_adjacent_structures(structure, merged_structures, index_map, peak=False)
             [merged_structures[i].ancestor.idx for i in adjacent_structure_indices]
         )
         return [merged_structures[i] for i in ancestor_indices]
+
+
+class LocalDendrogram(Dendrogram):
+    @staticmethod
+    def compute(
+        data,
+        min_value="min",
+        min_delta=0,
+        min_npix=0,
+        is_independent=None,
+        verbose=False,
+        neighbours=None,
+        wcs=None,
+    ):
+        """
+        Compute a dendrogram from a Numpy array.
+
+        Parameters
+        ----------
+        data : :class:`numpy.ndarray`
+            The n-dimensional array to compute the dendrogram for
+        min_value : float or "min", optional
+            The minimum data value to go down to when computing the
+            dendrogram. Values below this threshold will be ignored. Defaults
+            to the minimum value in the data.
+        min_delta : float, optional
+            The minimum height a leaf has to have in order to be considered an
+            independent entity.
+        min_npix : int, optional
+            The minimum number of pixels/values needed for a leaf to be considered
+            an independent entity.
+        is_independent : function or list of functions, optional
+            A custom function that can be specified that will determine if a
+            leaf can be treated as an independent entity. The signature of the
+            function should be ``func(structure, index=None, value=None)``
+            where ``structure`` is the structure under consideration, and
+            ``index`` and ``value`` are optionally the pixel that is causing
+            the structure to be considered for merging into/attaching to the
+            tree.
+
+            If multiple functions are provided as a list, they
+            are all applied when testing for independence.
+
+        neighbours : function, optional
+            A function that returns the list of neighbours to a given
+            location. Neighbours is called as ``neighbours(dendrogram, idx)``,
+            where ``idx`` is a tuple describing the n-dimensional location
+            of a pixel. It returns a list of N-dimensional locations of
+            neighbours. This function can implement optional adjacency logic.
+
+            .. note:: ``idx`` refers to location in a copy of the input data
+                       that has been padded with one element along each edge.
+
+        wcs : WCS object, optional
+            A WCS object that describes `data`. This is used in the
+            interactive viewer to properly display the data's coordinates
+            on the image axes. (Requires that `wcsaxes` is installed; see
+            http://wcsaxes.readthedocs.org/ for install instructions.)
+
+
+        Examples
+        --------
+
+        The following example demonstrates how to compute a dendrogram from an
+        dataset contained in a FITS file::
+
+            >>> from astropy.io import fits
+            >>> array = fits.getdata('observations.fits')
+            >>> from astrodendro import Dendrogram
+            >>> d = Dendrogram.compute(array)
+
+        Notes
+        -----
+        More information about the above parameters is available from the
+        online documentation at [www.dendrograms.org](www.dendrograms.org).
+        """
+        tests = [pruning.min_delta(min_delta), pruning.min_npix(min_npix)]
+        if is_independent is not None:
+            if hasattr(is_independent, "__iter__"):
+                tests.extend(is_independent)
+            else:
+                tests.append(is_independent)
+        is_independent = pruning.all_true(tests)
+        neighbours = neighbours or Dendrogram.neighbours
+
+        # Default min_val to the minimum in the data
+        if min_value == "min":
+            min_value = np.min(data[np.isfinite(data)]) - 1
+
+        self = Dendrogram()
+        self.data = data
+        self.n_dim = len(data.shape)
+        self.wcs = wcs
+        # For reference, store the parameters used:
+        self.params = dict(min_npix=min_npix, min_value=min_value, min_delta=min_delta)
+
+        # Create a list of all points in the cube above min_value
+        keep = self.data > min_value
+        data_values = self.data[keep]
+        indices = np.vstack(np.where(keep)).transpose()
+
+        if verbose:
+            print(
+                "Generating dendrogram using {:,} of {:,} pixels ({}% of data)".format(
+                    data_values.size,
+                    self.data.size,
+                    (100 * data_values.size / self.data.size),
+                )
+            )
+            progress_bar = AnimatedProgressBar(
+                end=max(data_values.size, 1), width=40, fill="=", blank=" "
+            )
+
+        # Define index array indicating what structure each cell is part of
+        # We expand each dimension by one, so the last value of each
+        # index (accessed with e.g. [nx,#,#] or [-1,#,#]) is always zero
+        # This permits an optimization below when finding adjacent structures
+        self.index_map = -np.ones(np.add(self.data.shape, 1), dtype=np.int32)
+
+        # Dictionary of currently-defined structures:
+        structures = {}
+
+        # Loop from largest to smallest data_value value. Each time, check if
+        # the pixel connects to any existing leaf. Otherwise, create new leaf.
+        count = 0
+
+        for i in np.argsort(data_values)[::-1]:
+
+            def next_idx():
+                return i + 1
+                # Generate IDs index i. We add one to avoid ID 0
+
+            data_value = data_values[i]
+            coord = tuple(indices[i])
+
+            # Print stats
+            count += 1
+            if verbose and (count % 100 == 0):
+                progress_bar + 100
+                progress_bar.show_progress()
+
+            # Check if point is adjacent to any leaf
+            # We don't worry about the edges, because overflow or underflow n
+            # any one dimension will always land on an extra "padding" cell
+            # with value zero added above when index_map was created
+
+            indices_adjacent = neighbours(self, indices[i])
+            adjacent = [
+                self.index_map[c] for c in indices_adjacent if self.index_map[c] > -1
+            ]
+
+            # Replace adjacent elements by its ancestor
+            adjacent = [structures[a].ancestor for a in adjacent]
+
+            # Remove duplicates
+            adjacent = _sorted_by_idx(set(adjacent))
+
+            # What happens next depends on how many unique adjacent structures there are
+
+            if not adjacent:  # No adjacent structures;  Create new leaf:
+                # Set absolute index of the new element
+                idx = next_idx()
+
+                # Create leaf
+                leaf = astrodendro_structure(
+                    coord, data_value, idx=idx, dendrogram=self
+                )
+
+                # Add leaf to overall list
+                structures[idx] = leaf
+
+                # Set absolute index of pixel in index map
+                self.index_map[coord] = idx
+
+            elif len(adjacent) == 1:  # Add to existing leaf or branch
+                # Add point to structure
+                adjacent[0]._add_pixel(coord, data_value)
+
+                # Set absolute index of pixel in index map
+                self.index_map[coord] = adjacent[0].idx
+
+            else:  # Merge leaves
+                # At this stage, the adjacent structures might consist of an
+                # arbitrary number of leaves and branches.
+
+                # Find all leaves that are not important enough to be
+                # kept separate. These leaves will now be treated the
+                # same as the pixel under consideration
+                merge = [
+                    structure
+                    for structure in adjacent
+                    if structure.is_leaf
+                    and (
+                        structure.vmax == data_value
+                        or not is_independent(structure, index=coord, value=data_value)
+                    )
+                ]
+
+                # Remove merges from list of adjacent structures
+                for structure in merge:
+                    adjacent.remove(structure)
+
+                # Now, figure out what object this pixel belongs to
+                # How many significant adjacent structures are left?
+
+                if not adjacent:  # if len(adjacent) == 0:
+                    # There are no separate leaves left (and no branches), so pick the
+                    # first one as the reference and merge all the others onto it
+                    belongs_to = merge.pop()
+                    belongs_to._add_pixel(coord, data_value)
+                elif len(adjacent) == 1:
+                    # There is one significant adjacent leaf/branch left.
+                    belongs_to = adjacent[0]
+                    belongs_to._add_pixel(coord, data_value)
+                else:
+                    # Create a branch
+                    belongs_to = astrodendro_structure(
+                        coord,
+                        data_value,
+                        children=adjacent,
+                        idx=next_idx(),
+                        dendrogram=self,
+                    )
+                    # Add branch to overall list
+                    structures[belongs_to.idx] = belongs_to
+
+                # Set absolute index of pixel in index map
+                self.index_map[coord] = belongs_to.idx
+
+                # Add all insignificant leaves in 'merge' to the same object as this pixel:
+                for m in merge:
+                    # print "Merging leaf %i onto leaf %i" % (i, idx)
+                    # Remove leaf
+                    structures.pop(m.idx)
+                    # Merge the insignificant structure that this pixel now belongs to:
+                    belongs_to._merge(m)
+                    # Update index map
+                    m._fill_footprint(self.index_map, belongs_to.idx)
+
+        if verbose:
+            progress_bar.progress = 100  # Done
+            progress_bar.show_progress()
+            print("")  # newline
+
+        # Create trunk from objects with no ancestors
+        def dummy_is_independent(*args, **kwargs):
+            return True
+
+        _make_trunk(self, structures, dummy_is_independent)
+        # _make_trunk(self, structures, is_independent)
+
+        # Save a list of all structures accessible by ID
+        self._structures_dict = {}
+
+        # Re-assign idx and update index map
+        sorted_structures = sorted(self, key=lambda s: s.smallest_index)
+        for idx, s in enumerate(sorted_structures):
+            s.idx = idx
+            s._fill_footprint(self.index_map, idx, recursive=False)
+            self._structures_dict[idx] = s
+
+        # Remove border from index map
+        s = tuple(slice(0, s, 1) for s in data.shape)
+        self.index_map = self.index_map[s]
+
+        # Add dendrogram index
+        self._index()
+
+        # Return the newly-created dendrogram:
+        return self

tests/test_distributed_dendrogram_v3.py

@@ -218,7 +218,10 @@ def test_example_pseudo_parallel_local_params(ntasks):
 
     d_ref = DistributedDendrogramV3.compute_pseudo_parallel(data, ntasks, **kwargs)
     d = DistributedDendrogramV3.compute_pseudo_parallel(
-        data, ntasks, **kwargs, min_npix_loc=min_npix, min_delta_loc=min_delta
+        data,
+        ntasks,
+        **kwargs,
+        min_npix_loc=min_npix,  # , min_delta_loc=min_delta
     )
     assert d._iterations < d_ref._iterations  # we were faster with the local parameters
     if ntasks < 4:
@@ -237,7 +240,7 @@ def test_example_pseudo_parallel_local_params(ntasks):
     if ht.comm.rank == 0:
         logging.basicConfig(level=logging.INFO)
 
-    # test_1D_v3_pseudo_parallel(2, 32, 6, 0.0, 0.0, show=True)
+    # test_1D_v3_pseudo_parallel(2, 64, 6, 0.1, 0.0, show=True)
     # test_1D_v3_pseudo_parallel(2, 33, 0, 0.0, 0.0, show=True)
     # test_1D_v3(2, 33, 0, 0.0, 0.0)
     # test_2D_v3_pseudo_parallel(4, 32, 3, 0, 0, 0)