Ready for initial release: bigfixes, more strict linter

Author: thetianshuhuang

docs/tss/index.md

@@ -44,7 +44,7 @@ The neural radar development kit ships with a small statistical module (`nrdk.ts
     ```
     where `rho_t` is the `t`-lag autocorrelation. The sum of autocorrelations is empirically estimated from the performance metrics, where the sum is calculated up to `N / 2` or the first negative autocorrelation, whichever is first.
     
-    For details about how we calculate this, see [`effective_sample_size`][tss.stats.effective_sample_size].
+    For details about how we calculate this, see [`effective_sample_size`][nrdk.tss.stats.effective_sample_size].
 
     ??? question "Why cut off when `rho_t` is negative?"
 
@@ -62,7 +62,7 @@ The neural radar development kit ships with a small statistical module (`nrdk.ts
 
 ## General Usage
 
-In addition to using the [low level API][tss.stats], we provide a [CLI](#cli) and a [high level API][tss] which can be used to index results and evaluations, then load the evaluations and calculate statistics.
+In addition to using the [low level API][nrdk.tss.stats], we provide a [CLI](#cli) and a [high level API][nrdk.tss] which can be used to index results and evaluations, then load the evaluations and calculate statistics.
 
 ### File Format
 
@@ -160,7 +160,7 @@ The file path to each evaluation, relative to some base path, should contain inf
 
 ## High Level API
 
-**Index evaluations**: using [`index`][tss.index], provide a base path where the evaluations are stored, and a regex pattern for finding evaluation files and extracting their `experiment` and `trace` names.
+**Index evaluations**: using [`index`][nrdk.tss.index], provide a base path where the evaluations are stored, and a regex pattern for finding evaluation files and extracting their `experiment` and `trace` names.
 
 ```python
 import tss
@@ -178,7 +178,7 @@ index = tss.index_results(path, pattern)
     a directory with 20k total files on a SMB share). You may want to
     cache the index or save them to disk somewhere!
 
-**Compute Statistics**: we provide a all-inclusive [`dataframe_from_index`][tss.dataframe_from_index] function which returns a dataframe containing summary statistics for the specified index, given a key of interest and baseline method.
+**Compute Statistics**: we provide a all-inclusive [`dataframe_from_index`][nrdk.tss.dataframe_from_index] function which returns a dataframe containing summary statistics for the specified index, given a key of interest and baseline method.
 
 ```python
 experiments = ["small/p10", "small/p20", "small/p50", "small/base"]

grt/config/model/decoder/lidar2d.yaml

@@ -14,7 +14,7 @@ occ2d:
   d_model: ${globals.d_model}
   num_layers: ${globals.dec_layers}
   shape: [1, 1, 1024, 256]
-  pos_scale: [1.0, 1.0, 1.0, 1.0]
-  global_scale: 16.0
+  scale: [1.0, 1.0, 1.0, 1.0]
+  w_min: 0.2
   patch: [1, 1, 16, 16]
   out_dim: 0

grt/config/model/decoder/lidar3d.yaml

@@ -14,7 +14,7 @@ occ3d:
   d_model: ${globals.d_model}
   num_layers: 4
   shape: [1, 64, 128, 64]
-  pos_scale: [1.0, 1.0, 1.0, 1.0]
-  global_scale: 16.0
+  scale: [1.0, 1.0, 1.0, 1.0]
+  w_min: 0.2
   patch: [1, 8, 8, 8]
   out_dim: 0

grt/config/model/decoder/semseg.yaml

@@ -14,7 +14,7 @@ semseg:
   d_model: ${globals.d_model}
   num_layers: 4
   shape: [1, 160, 160, 1]
-  pos_scale: [1.0, 1.0, 0.666, 1.0]
-  global_scale: 16.0
+  scale: [1.0, 1.0, 0.666, 1.0]
+  w_min: 0.2
   patch: [1, 5, 5, 1]
   out_dim: 8

grt/config/model/tokenizer/grt.yaml

@@ -3,6 +3,6 @@ d_model: ${globals.d_model}
 patch: [1, 2, 2, 8, 4]  # (time, doppler, elevation, azimuth, range)
 squeeze: [2, 3]  # elevation + azimuth
 n_channels: 3 
-scale: [1.0, 1.0, 1.0, 1.0, 1.0]
+scale: [1.0, 1.0, 1.0]
 w_min: 0.2
 positions: nd

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nrdk"
-version = "0.0.0"
+version = "0.1.0"
 authors = [
   { name="Tianshu Huang", email="tianshu2@andrew.cmu.edu" },
 ]
@@ -79,13 +79,16 @@ indent-width = 4
 
 [tool.ruff.lint]
 select = [
-    "W", "N", "I", "D", "NPY"
+    "E", "F", "W", "N", "I", "D", "NPY"
 ]
 ignore = [
     "D102",  # Inheriting docstrings is allowed
     "D107",  # Initializer goes in the class docstring
     "D401",  # "Imperative mood" NLP doesn't always make sense
     "D413",  # No blank lines after section is preferred by mkdocstrings
+    "F722",  # Jaxtyping annotations confuse ruff
+    "F821",
+    "N806",  # We do math, so capitalization should be allowed
 
     # Must pick one...
     "D213",  # Pick 'multi-line-summary-first-line'.

src/nrdk/config/config.py

@@ -23,10 +23,11 @@ def expand(path: str | None = None, **nested: Nested) -> list[str]:
     """
     def _expand(nested: Nested, base: str | None = None):
 
-        if base is not None:
-            _join = lambda p: os.path.join(base, p)
-        else:
-            _join = lambda p: p
+        def _join(p: str) -> str:
+            if base is not None:
+                return os.path.join(base, p)
+            else:
+                return p
 
         if isinstance(nested, Sequence):
             return [_join(item) for item in nested]

src/nrdk/modules/position.py

@@ -109,7 +109,7 @@ def forward(
         """
         if positions is None:
             positions = [
-                torch.linspace(-1.0, 1.0, steps=n, device=x.device)
+                torch.linspace(-1.0, 1.0, steps=n, device=x.device)[None, :]
                 for n in x.shape[1:-1]]
 
         if len(positions) != len(x.shape) - 2:

src/nrdk/objectives/__init__.py

@@ -70,10 +70,11 @@
 
     !!! info
 
-        Each `*Data` type is described using protocols such that any classes which
-        have the same attributes as the protocol can be used as inputs. For
-        example, `SemsegData` can be replaced by any object (e.g., a dataclass)
-        with a `semseg` attribute with dtype `UInt8` and shape `batch t h w`.
+        Each `*Data` type is described using protocols such that any classes
+        which have the same attributes as the protocol can be used as inputs.
+        For example, `SemsegData` can be replaced by any object (e.g., a
+        dataclass) with a `semseg` attribute with dtype `UInt8` and shape
+        `batch t h w`.
         ```python
         class SemsegData(Protocol, Generic[TArray]):
             semseg: UInt8[TArray, "batch t h w"]

src/nrdk/tss/__init__.py

@@ -44,5 +44,5 @@
     "NestedValues",
     "dataframe_from_index", "dataframe_from_stats",
     "experiments_from_index", "stats_from_experiments",
-    "index", "stats", "utils"
+    "index", "stats", "utils", "_cli_main"
 ]