wip, pushing so we can access from pip install

Author: kylemann16

src/silvimetric/cli/cli.py

@@ -256,7 +256,7 @@ def scan_cmd(
     '-a',
     type=AttrParamType(),
     default=[],
-    help='List of attributes to include in Database',
+    help='List of attributes to include in Database, eg. -a Z,Intensity',
 )
 @click.option(
     '--metrics',
@@ -385,10 +385,9 @@ def shatter_cmd(app, pointcloud, bounds, report, tilesize, date, dates):
 @click.option(
     '--attributes',
     '-a',
-    multiple=True,
     type=AttrParamType(),
     default=[],
-    help='List of attributes to include output',
+    help='List of attributes to include output, eg -a Z,Intensity',
 )
 @click.option(
     '--metrics',

src/silvimetric/cli/common.py

@@ -65,6 +65,11 @@ class AttrParamType(click.ParamType):
     # TODO add import similar to metrics
     def convert(self, value, param, ctx) -> list[Attribute]:
         attrs: set[Attribute] = set()
+        if not value:
+            return set(
+                Attributes[a]
+                for a in ['Z', 'ReturnNumber', 'NumberOfReturns', 'Intensity']
+            )
         parsed_values = value.split(',')
         for val in parsed_values:
             try:
@@ -94,7 +99,7 @@ class MetricParamType(click.ParamType):
 
     def convert(self, value, param, ctx) -> list[Metric]:
         if value is None or not value:
-            return list(grid_metrics('Z').values())
+            return list(grid_metrics.get_grid_metrics('Z').values())
         parsed_values = value.split(',')
         metrics: set[Metric] = set()
         for val in parsed_values:
@@ -208,11 +213,6 @@ def dask_handle(
             p.register()
 
     elif scheduler == 'distributed':
-        # raise Exception("Dask distributed scheduler is currently disabled. "
-        #     "Please select another scheduler ('single-threaded' or 'local')")
-
-        # TODO: this is disabled for now. Distributed dask keeps crashing when
-        # used for shatter.
 
         dask_config['scheduler'] = scheduler
         if dasktype == 'processes':

src/silvimetric/resources/storage.py

@@ -1,15 +1,11 @@
-import pathlib
-import contextlib
 import json
-from struct import Struct
-import struct
+import operator
 
 from math import floor
 from datetime import datetime
-from typing_extensions import Generator, Optional
+from typing_extensions import Optional
 
 import tiledb
-from tiledb.metadata import Metadata
 import numpy as np
 import pandas as pd
 
@@ -25,7 +21,6 @@ def ts_overlap(first, second):
         return False
     return True
 
-
 class Storage:
     """Handles storage of shattered data in a TileDB Database."""
 
@@ -35,14 +30,16 @@ def __init__(self, config: StorageConfig):
                 f"Given database directory '{config.tdb_dir}' does not exist"
             )
 
-        self.config = config
-        self.reader: tiledb.SparseArray = None
-        self.writer: tiledb.SparseArray = None
+        self.config: StorageConfig = config
+        self._reader: tiledb.SparseArray = None
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, exc_tb):
+        if self._reader is not None:
+            self._reader.close()
+        self._reader = None
         return
 
     @staticmethod
@@ -87,25 +84,25 @@ def create(config: StorageConfig, ctx: tiledb.Ctx = None):
             name='X',
             domain=(0, xi),
             dtype=np.int32,
-            filters=tiledb.FilterList([ tiledb.ZstdFilter() ]),
+            filters=tiledb.FilterList([tiledb.ZstdFilter()]),
         )
         dim_col = tiledb.Dim(
             name='Y',
             domain=(0, yi),
             dtype=np.int32,
-            filters=tiledb.FilterList([ tiledb.ZstdFilter() ]),
+            filters=tiledb.FilterList([tiledb.ZstdFilter()]),
         )
         domain = tiledb.Domain(dim_row, dim_col)
 
         count_att = tiledb.Attr(
             name='count',
             dtype=np.int32,
-            filters=tiledb.FilterList([ tiledb.ZstdFilter() ]),
+            filters=tiledb.FilterList([tiledb.ZstdFilter()]),
         )
         proc_att = tiledb.Attr(
             name='shatter_process_num',
             dtype=np.uint64,
-            filters=tiledb.FilterList([ tiledb.ZstdFilter() ]),
+            filters=tiledb.FilterList([tiledb.ZstdFilter()]),
         )
         dim_atts = [attr.schema() for attr in config.attrs]
 
@@ -152,11 +149,10 @@ def create(config: StorageConfig, ctx: tiledb.Ctx = None):
         schema.check()
 
         tiledb.SparseArray.create(config.tdb_dir, schema)
-        writer = tiledb.SparseArray(config.tdb_dir, 'w')
-        writer.meta['config'] = str(config)
+        with tiledb.SparseArray(config.tdb_dir, 'w') as writer:
+            writer.meta['config'] = str(config)
 
         s = Storage(config)
-        s.writer = writer
         s.save_config()
 
         return s
@@ -182,8 +178,7 @@ def from_db(tdb_dir: str, ctx: tiledb.Ctx = None):
         storage = Storage(config)
 
         # set the metadata for storage object so we don't have to query again
-        storage._meta = metadata
-        storage.reader = reader
+        storage._reader = reader
         storage.save_config()
 
         return storage
@@ -196,8 +191,8 @@ def save_config(self) -> None:
         # key later
         with self.open('w') as w:
             w.meta['config'] = str(self.config)
-        with self.open('r') as r:
-            self._meta = r.meta
+        if self._reader is not None:
+            self._reader.reopen()
 
     def get_config(self) -> StorageConfig:
         """
@@ -236,20 +231,13 @@ def get_metadata(self, key: str) -> str:
         :return: Metadata value found in storage.
         """
         # if meta hasn't been set up, do so
-        if self._meta is None:
-            with self.open('r') as tdb:
-                self._meta = tdb.meta
-                # this should be latest metadata, handle a keyerror higher up
-                return self._meta[key]
-
-        # if it was already set up and the key doesn't exist,
-        # reload metadata and check there
-        if key not in self._meta.keys():
-            with self.open('r') as tdb:
-                self._meta = tdb.meta
-
-        return self._meta['key']
-
+        reader = self.open('r')
+        try:
+            return reader.meta[key]
+        except KeyError:
+            reader.reopen()
+            self._reader = reader
+            return reader.meta[key]
 
     def save_metadata(self, key: str, data: str) -> None:
         """
@@ -259,12 +247,11 @@ def save_metadata(self, key: str, data: str) -> None:
         :param data: Data to save to metadata.
         """
         # if writer isn't set up, do it now
-        if self._writer is None:
-            self._writer = self.open('w')
-
         # propogate the key-value to both tiledb and the local copy
-        self._meta[key] = data
-        self._writer.meta[key] = data
+        with self.open('w') as w:
+            w.meta[key] = data
+        if self._reader is not None:
+            self._reader.reopen()
 
     def get_tdb_context(self):
         cfg = tiledb.Config()
@@ -299,9 +286,7 @@ def get_derived_names(self) -> list[str]:
             if not m.attributes or a.name in [ma.name for ma in m.attributes]
         ]
 
-    def open(
-        self, mode: str = 'r', timestamp=None
-    ) -> Generator[tiledb.SparseArray, None, None]:
+    def open(self, mode: str = 'r', timestamp=None) -> tiledb.SparseArray:
         """
         Open stream for TileDB database in given mode and at given timestamp.
 
@@ -318,22 +303,18 @@ def open(
         # other threads present
         ctx = self.get_tdb_context()
 
-        if tiledb.object_type(self.config.tdb_dir) == 'array':
-            if mode in ['w', 'r', 'd', 'm']:
-                tdb = tiledb.open(
-                    self.config.tdb_dir, mode, timestamp=timestamp, ctx=ctx
-                )
-            else:
-                raise Exception(f"Given open mode '{mode}' is not valid")
-        elif pathlib.Path(self.config.tdb_dir).exists():
-            raise Exception(
-                f'Path {self.config.tdb_dir} already exists and is not'
-                ' initialized for TileDB access.'
+        # non-timestamped reader and writer are stored as member variables to
+        # avoid opening and closing too many io objects.
+        if timestamp is not None or mode != 'r':
+            return tiledb.open(
+                self.config.tdb_dir, mode, timestamp=timestamp, ctx=ctx
             )
-        else:
-            raise Exception(f'Path {self.config.tdb_dir} does not exist')
+        else:  # no timestamp and mode is 'r'
+            if self._reader is None or not self._reader.isopen:
+                self._reader = tiledb.open(self.config.tdb_dir, 'r')
 
-        return tdb
+            self._reader.reopen()
+            return self._reader
 
     def write(self, data_in: pd.DataFrame, timestamp):
         data_in = data_in.rename(columns={'xi': 'X', 'yi': 'Y'})
@@ -375,17 +356,17 @@ def reserve_time_slot(self) -> int:
         shatter process.
 
         :param config: Shatter config will be written as metadata to reserve
-            time slot.
+        time slot.
 
         :return: Time slot.
         """
-        latest = self.get_metadata('config')
-
-        time = latest['next_time_slot']
-        latest['next_time_slot'] = time + 1
-        self.save_metadata('config', json.dumps(latest))
+        # make sure we're dealing with the latest config
+        cfg = self.get_config()
+        self.config = cfg
+        time = self.config.next_time_slot
+        self.config.next_time_slot = time + 1
+        self.save_config()
 
-        self.config.next_time_slot = latest['next_time_slot']
         return time
 
     def get_history(
@@ -472,14 +453,16 @@ def delete(self, time_slot: int) -> ShatterConfig:
         """
 
         self.config.log.debug(f'Deleting time slot {time_slot}...')
-        with self.open('r') as r:
-            sh_cfg = ShatterConfig.from_string(r.meta[f'shatter_{time_slot}'])
-            sh_cfg.mbr = ()
-            sh_cfg.finished = False
+
+        r = self.open('r')
+        sh_cfg = ShatterConfig.from_string(r.meta[f'shatter_{time_slot}'])
+        sh_cfg.mbr = ()
+        sh_cfg.finished = False
 
         self.config.log.debug('Deleting fragments...')
-        with self.open('d') as d:
-            d.query(cond=f'shatter_process_num=={time_slot}').submit()
+        d = self.open('d')
+        d.query(cond=f'shatter_process_num=={time_slot}').submit()
+        d.close()
 
         self.config.log.debug('Rewriting config.')
         with self.open('w') as w:

tests/test_cli.py

@@ -211,11 +211,9 @@ def test_cli_extract(
     ):
         atts = []
         for a in extract_config.attrs:
-            atts.append('-a')
             atts.append(a.name)
         ms = []
         for m in extract_config.metrics:
-            ms.append('-m')
             ms.append(m.name)
         out_dir = extract_config.out_dir
         tdb_dir = extract_config.tdb_dir
@@ -228,8 +226,10 @@ def test_cli_extract(
                 '--scheduler',
                 'single-threaded',
                 'extract',
-                *atts,
-                *ms,
+                '-a',
+                ','.join(atts),
+                '-m',
+                ','.join(ms),
                 '--outdir',
                 out_dir,
             ],

tests/test_shatter.py

@@ -5,6 +5,7 @@
 
 import pytest
 import numpy as np
+import pandas as pd
 import dask
 
 from silvimetric import Extents, Log, info, shatter, Storage
@@ -47,9 +48,15 @@ def confirm_one_entry(storage, maxy, base, pointcount, num_entries=1):
 
         for xi in range(xdom):
             for yi in range(ydom):
-                assert bool(
-                    np.all(vals.loc[xi, yi].Z[0] == (val_const - yi - 1))
-                )
+                try:
+                    z = vals.loc[xi,yi].Z
+                    if isinstance(z, np.ndarray):
+                        assert np.all(z == (val_const - yi - 1))
+                    elif isinstance(z, pd.Series):
+                        for z1 in z.values:
+                            np.all(z1 == (val_const - yi - 1))
+                except Exception as e:
+                    print('asfasdf')
 
 
 class Test_Shatter(object):  # noqa: D101