moving around usage of dask, attempting to speed up some long running methods

Author: kylemann16

src/silvimetric/commands/scan.py

@@ -55,7 +55,7 @@ def scan(
             count = dask.delayed(data.estimate_count)(extents.bounds).persist()
 
             if filter_empty:
-                chunks = extents.chunk(data, resolution, point_count, depth)
+                chunks = extents.chunk(data, point_count)
                 cell_counts = [ch.cell_count for ch in chunks]
 
             else:

src/silvimetric/commands/shatter.py

@@ -17,6 +17,7 @@
 import dask.array as da
 import dask.bag as db
 from dask.diagnostics import ProgressBar
+from dask import persist, compute
 
 from .. import Extents, Storage, Data, ShatterConfig
 from ..resources.taskgraph import Graph
@@ -159,26 +160,30 @@ def kill_gracefully(signum, frame):
     signal.signal(signal.SIGINT, kill_gracefully)
 
     # leaf_bag: db.Bag = db.from_sequence(leaves)
+    # processes = leaf_bag.map(do_one, config, storage)
     processes = [delayed(do_one)(leaf, config, storage) for leaf in leaves]
 
     ## If dask is distributed, use the futures feature
     dc = get_client()
     consolidate_count = 30
     count = 0
     if dc is not None:
-        pc_futures = futures_of(processes.persist())
+        pc_futures = futures_of(persist(processes))
         for batch in as_completed(pc_futures, with_results=True).batches():
             for _, pack in batch:
                 if isinstance(pack, CancelledError):
                     continue
+                if isinstance(pack, int):
+                    pack = [pack]
                 for pc in pack:
                     if isinstance(pc, BaseException):
-                        print('Error found: ', pc)
-                    else:
+                        config.log.warning('Worker returned exception: ', pc)
+                    if isinstance(pc, int):
                         count += 1
                         if count >= consolidate_count:
                             faf = dc.submit(
-                                storage.consolidate_shatter(config.timestamp)
+                                storage.consolidate_shatter,
+                                timestamp=config.timestamp,
                             )
                             fire_and_forget(faf)
                             count = 0
@@ -192,7 +197,7 @@ def kill_gracefully(signum, frame):
     else:
         # Handle non-distributed dask scenarios
         with ProgressBar():
-            point_count = sum(processes)
+            point_count = sum(*compute(processes))
 
     return point_count
 
@@ -229,11 +234,8 @@ def shatter(config: ShatterConfig) -> int:
     if config.tile_size is not None:
         leaves = extents.get_leaf_children(config.tile_size)
     else:
-        leaves = extents.chunk(
-            data, res_threshold=storage.config.resolution, depth_threshold=30
-        )
-        leaf_count = len(leaves)
-        config.log.debug(f'{leaf_count} tiles to process.')
+        chunks = extents.chunk(data, pc_threshold=600000)
+        leaves = db.from_sequence(chunks).compute()
 
     # Begin main operations
     config.log.debug('Fetching and arranging data...')

src/silvimetric/resources/extents.py

@@ -94,9 +94,7 @@ def disjoint(self, other):
     def chunk(
         self,
         data: Data,
-        res_threshold=100,
         pc_threshold=600000,
-        depth_threshold=6,
     ):
         """
         Split up a dataset into tiles based on the given thresholds. Unlike Scan
@@ -131,34 +129,57 @@ def chunk(
 
         if self.bounds == self.root:
             self.root = chunk.bounds
+        yield from chunk.filter(data, pc_threshold)
 
-        filtered = []
-        curr = db.from_delayed(
-            [
-                dask.delayed(ch.filter)(
-                    data, res_threshold, pc_threshold, depth_threshold, 1
-                )
-                for ch in chunk.split()
-            ]
-        )
-        curr_depth = 1
+    def filter(
+        self,
+        data: Data,
+        pc_threshold=600000,
+        prev_estimate=0,
+    ):
+        """
+        Creates quad tree of chunks for this bounds, runs pdal quickinfo over
+        this to determine if there are any points available. Uses a bottom
+        resolution of 1km.
+
+        :param data: Data object containing point cloud details.
+        :param res_threshold: Resolution threshold., defaults to 100
+        :param pc_threshold: Point count threshold., defaults to 600000
+        :param depth_threshold: Tree depth threshold., defaults to 6
+        :param depth: Current tree depth., defaults to 0
+        :return: Returns a list of Extents.
+        """
 
-        logger = data.storageconfig.log
-        while curr.npartitions > 0:
-            logger.debug(
-                f'Filtering {curr.npartitions} tiles at depth {curr_depth}'
-            )
-            n = curr.compute()
-            to_add = [ne for ne in n if isinstance(ne, Extents)]
-            if to_add:
-                filtered = filtered + to_add
+        pc = data.estimate_count(self.bounds)
 
-            curr = db.from_delayed(
-                [ne for ne in n if not isinstance(ne, Extents)]
-            )
-            curr_depth += 1
+        target_pc = pc_threshold
+        minx, miny, maxx, maxy = self.bounds.get()
 
-        return filtered
+        # is it empty?
+        if not pc:
+            yield self
+        else:
+            # has it hit the threshold yet?
+            area = (maxx - minx) * (maxy - miny)
+            next_split_x = (maxx - minx) / 2
+            next_split_y = (maxy - miny) / 2
+
+            # if the next split would put our area below the resolution, or if
+            # the point count is less than the point threshold then use this
+            # tile as the work unit.
+            if next_split_x < self.resolution or next_split_y < self.resolution:
+                yield self
+            elif pc <= target_pc:
+                yield self
+            elif pc == prev_estimate:
+                yield self
+            else:
+                for ch in self.split():
+                    yield from ch.filter(
+                            data,
+                            pc_threshold,
+                            prev_estimate=pc
+                        )
 
     def split(self):
         """
@@ -202,63 +223,6 @@ def split(self):
         ]
         return exts
 
-    def filter(
-        self,
-        data: Data,
-        res_threshold=100,
-        pc_threshold=600000,
-        depth_threshold=6,
-        depth=0,
-    ):
-        """
-        Creates quad tree of chunks for this bounds, runs pdal quickinfo over
-        this to determine if there are any points available. Uses a bottom
-        resolution of 1km.
-
-        :param data: Data object containing point cloud details.
-        :param res_threshold: Resolution threshold., defaults to 100
-        :param pc_threshold: Point count threshold., defaults to 600000
-        :param depth_threshold: Tree depth threshold., defaults to 6
-        :param depth: Current tree depth., defaults to 0
-        :return: Returns a list of Extents.
-        """
-
-        pc = data.estimate_count(self.bounds)
-        target_pc = pc_threshold
-        minx, miny, maxx, maxy = self.bounds.get()
-
-        # is it empty?
-        if not pc:
-            return []
-        else:
-            # has it hit the threshold yet?
-            area = (maxx - minx) * (maxy - miny)
-            next_split_x = (maxx - minx) / 2
-            next_split_y = (maxy - miny) / 2
-
-            # if the next split would put our area below the resolution, or if
-            # the point count is less than the threshold (600k) then use this
-            # tile as the work unit.
-            if next_split_x < self.resolution or next_split_y < self.resolution:
-                return [self]
-            elif pc < target_pc:
-                return [self]
-            elif area < res_threshold**2 or depth >= depth_threshold:
-                pc_per_cell = pc / (area / self.resolution**2)
-                cell_estimate = ceil(target_pc / pc_per_cell)
-
-                return self.get_leaf_children(cell_estimate)
-            else:
-                return [
-                    dask.delayed(ch.filter)(
-                        data,
-                        res_threshold,
-                        pc_threshold,
-                        depth_threshold,
-                        depth=depth + 1,
-                    )
-                    for ch in self.split()
-                ]
 
     def _find_dims(self, tile_size):
         """
@@ -308,15 +272,13 @@ def get_leaf_children(self, tile_size):
         coords_list = np.array(
             [[*x, *y] for x in dx for y in dy], dtype=np.float64
         )
-        yield from [
-            Extents(
-                Bounds(minx, miny, maxx, maxy),
-                self.resolution,
-                self.alignment,
-                self.root,
-            )
-            for minx, maxx, miny, maxy in coords_list
-        ]
+        for minx, maxx, miny, maxy in coords_list:
+            yield Extents(
+                    Bounds(minx, miny, maxx, maxy),
+                    self.resolution,
+                    self.alignment,
+                    self.root,
+                )
 
     @staticmethod
     def from_storage(tdb_dir: str):

src/silvimetric/resources/metrics/l_moments.py

@@ -13,7 +13,7 @@ def lmom4(data, *args):
     idx = np.arange(n)
 
     # sort in descending order
-    data = np.sort(data.reshape(n))[::-1]
+    data = np.sort(data, kind='quickstort')[::-1]
 
     b0 = data.mean()
     l1: float = b0