reorg and add truncate option for rover data

spf/data_collector.py

@@ -1,4 +1,3 @@
-import concurrent
 import logging
 import multiprocessing
 import queue
@@ -18,6 +17,7 @@
 from spf.dataset.v5_data import v5rx_2xf64_keys, v5rx_f64_keys, v5rx_new_dataset
 from spf.dataset.wall_array_v2_idxs import v2_column_names
 from spf.rf import beamformer_given_steering, get_avg_phase, precompute_steering_vectors
+from spf.scripts.zarr_utils import zarr_shrink
 from spf.sdrpluto.sdr_controller import (
     EmitterConfig,
     PPlus,
@@ -27,7 +27,6 @@
     setup_rx,
     setup_rxtx,
 )
-from spf.utils import zarr_shrink
 
 
 class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor):
@@ -581,7 +580,6 @@ def write_to_record_matrix(self, thread_idx, record_idx, data):
         data.rx_pos_x_mm = rx_pos[0]
         data.rx_pos_y_mm = rx_pos[1]
 
-
         assert data.rx_lo > 1
 
         if not self.yaml_config["dry-run"]:

spf/dataset/benchmark_dataset.py

@@ -1,24 +1,23 @@
+import os
+import pickle
+import random
+
+# ds_fn = "/mnt/4tb_ssd/june_fix/wallarrayv3_2024_06_15_11_44_13_nRX2_bounce.zarr"
+import sys
+import time
 from functools import partial
+
 import numpy as np
 import scipy
-from tqdm import tqdm
+import torch
 from filterpy.monte_carlo import systematic_resample
-
-import pickle
-import os
-
-import time
+from tqdm import tqdm
 
 from spf.dataset.spf_dataset import (
     v5_collate_beamsegnet,
     v5_collate_keys_fast,
     v5spfdataset,
 )
-import torch
-import random
-
-# ds_fn = "/mnt/4tb_ssd/june_fix/wallarrayv3_2024_06_15_11_44_13_nRX2_bounce.zarr"
-import sys
 
 if len(sys.argv) != 2:
     print(f"{sys.argv[0]} input")

spf/dataset/benchmark_dataset_linear.py

@@ -1,21 +1,21 @@
+import os
+import pickle
+import random
+import sys
+import time
 from functools import partial
+
 import numpy as np
 import scipy
-from tqdm import tqdm
+import torch
 from filterpy.monte_carlo import systematic_resample
-
-import pickle
-import os
-import sys
-import time
+from tqdm import tqdm
 
 from spf.dataset.spf_dataset import (
     v5_collate_beamsegnet,
     v5_collate_keys_fast,
     v5spfdataset,
 )
-import torch
-import random
 
 ds_fn = "/mnt/4tb_ssd/june_fix/wallarrayv3_2024_06_15_11_44_13_nRX2_bounce.zarr"
 # ds_fn = "/mnt/4tb_ssd/wallarrayv3_2024_06_15_11_44_13_nRX2_bounce.zarr"

spf/dataset/compare_datasets.py

@@ -1,7 +1,9 @@
 import argparse
 import sys
+
 from spf.dataset.spf_dataset import v5spfdataset
-from spf.utils import compare_and_check, identical_datasets
+from spf.scripts.zarr_utils import compare_and_check
+from spf.utils import identical_datasets
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

spf/dataset/fake_dataset.py

@@ -16,8 +16,9 @@
     create_empirical_p_dist,
     get_empirical_p_dist_parser,
 )
+from spf.scripts.zarr_utils import zarr_open_from_lmdb_store, zarr_shrink
 from spf.sdrpluto.sdr_controller import rx_config_from_receiver_yaml
-from spf.utils import torch_random_signal_matrix, zarr_open_from_lmdb_store, zarr_shrink
+from spf.utils import torch_random_signal_matrix
 
 
 @torch.jit.script

spf/dataset/spf_dataset.py

@@ -68,15 +68,13 @@
     torch_get_phase_diff,
     torch_pi_norm,
 )
-from spf.sdrpluto.sdr_controller import rx_config_from_receiver_yaml
-from spf.utils import (
-    SEGMENTATION_VERSION,
+from spf.scripts.zarr_utils import (
     new_yarr_dataset,
-    rx_spacing_to_str,
-    to_bin,
     zarr_open_from_lmdb_store,
     zarr_shrink,
 )
+from spf.sdrpluto.sdr_controller import rx_config_from_receiver_yaml
+from spf.utils import SEGMENTATION_VERSION, rx_spacing_to_str, to_bin
 
 
 # from Stackoverflow

spf/dataset/v4_data.py

@@ -1,4 +1,4 @@
-from spf.utils import zarr_new_dataset
+from spf.scripts.zarr_utils import zarr_new_dataset
 
 v4rx_f64_keys = [
     "system_timestamp",

spf/dataset/v5_data.py

@@ -1,4 +1,11 @@
-from spf.utils import zarr_new_dataset
+import os
+import shutil
+
+from spf.scripts.zarr_utils import (
+    compare_and_copy,
+    zarr_new_dataset,
+    zarr_open_from_lmdb_store,
+)
 
 v5rx_f64_keys = [
     "system_timestamp",

spf/rf.py

@@ -5,14 +5,14 @@
 import torch
 from numba import njit
 
+from spf.scripts.zarr_utils import zarr_open_from_lmdb_store_cm
+
 try:
     import cupy as cp
 except:
     pass
 from scipy.stats import trim_mean
 
-from spf.utils import zarr_open_from_lmdb_store_cm
-
 # numba = False
 
 """

spf/scripts/truncate_zarr.py

@@ -0,0 +1,17 @@
+import argparse
+import sys
+
+from spf.dataset.v4_data import v4rx_2xf64_keys, v4rx_f64_keys
+from spf.scripts.zarr_utils import truncate_zarr
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("zarr_filename_in", type=str, help="input zarr")
+    args = parser.parse_args()
+
+    # this only works for version 4
+
+    return_code = truncate_zarr(
+        args.zarr_filename_in, f64_keys=v4rx_f64_keys, f64x2_keys=v4rx_2xf64_keys
+    )
+    sys.exit(return_code)

spf/dataset/v4_tx_rx_to_v5.py → spf/scripts/v4_tx_rx_to_v5.py

@@ -10,8 +10,9 @@
 from pyproj import Proj
 
 from spf.dataset.v5_data import v5rx_2xf64_keys, v5rx_f64_keys
-from spf.gps.boundaries import franklin_safe  # crissy_boundary_convex
-from spf.utils import zarr_new_dataset, zarr_open_from_lmdb_store, zarr_shrink
+from spf.gps.boundaries import franklin_safe
+from spf.scripts.zarr_utils import zarr_new_dataset  # crissy_boundary_convex
+from spf.scripts.zarr_utils import zarr_open_from_lmdb_store, zarr_shrink
 
 
 def lat_lon_to_xy(lat, lon, center_lat, center_lon):

spf/dataset/yarr_rechunk.py → spf/scripts/yarr_rechunk.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-from spf.utils import new_yarr_dataset, zarr_open_from_lmdb_store
+from spf.scripts.zarr_utils import new_yarr_dataset, zarr_open_from_lmdb_store
 
 if __name__ == "__main__":
     if len(sys.argv) != 3:

spf/dataset/zarr_rechunk.py → spf/scripts/zarr_rechunk.py

@@ -4,13 +4,11 @@
 import sys
 
 import yaml
-import zarr
 
 from spf.dataset.v5_data import v5rx_2xf64_keys, v5rx_f64_keys
-from spf.utils import (
+from spf.scripts.zarr_utils import (
     compare_and_check,
     compare_and_copy,
-    new_yarr_dataset,
     zarr_new_dataset,
     zarr_open_from_lmdb_store,
     zarr_shrink,

spf/dataset/zarr_shrink.py → spf/scripts/zarr_shrink.py

@@ -1,6 +1,6 @@
 import sys
 
-from spf.utils import zarr_shrink
+from spf.scripts.zarr_utils import zarr_shrink
 
 if __name__ == "__main__":
     if len(sys.argv) != 2: