Pair Aligner Preparation

.dockerignore

@@ -1,15 +1,16 @@
 specs/**
-
+*.gz
 # Traing data logging
 **/wandb/**
 **/.neptune/**
 **/tensorboard/**
-
+task_dashboard
+tests
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
 *$py.class
-
+.git/
 # C extensions
 *.so
 

zetta_utils/builder/constants.py

@@ -1,4 +1,4 @@
 from typing import Final
 
-DEFAULT_VERSION: Final = "0.0.0"
+DEFAULT_VERSION: Final = "0.4"
 DEFAULT_VERSION_SPEC: Final = ">=0.0.0"

zetta_utils/geometry/bbox.py

@@ -515,7 +515,7 @@ def snapped(
         self,
         grid_offset: Sequence[float],
         grid_size: Sequence[float],
-        mode: Literal["shrink", "expand"],
+        mode: Literal["shrink", "expand", "round_down"],
     ) -> BBox3D:
         """Returns a BoundingBox snapped to a grid with the given offset and size.
 
@@ -536,8 +536,7 @@ def snapped(
                 floor(round((b[1] - o) / s, VEC3D_PRECISION)) * s + o
                 for b, o, s in zip(self.bounds, grid_offset, grid_size)
             )
-        else:
-            assert mode == "expand", "Typechecking error"
+        elif mode == "expand":
             start_final = tuple(
                 floor(round((b[0] - o) / s, VEC3D_PRECISION)) * s + o
                 for b, o, s in zip(self.bounds, grid_offset, grid_size)
@@ -546,6 +545,17 @@ def snapped(
                 floor(round((b[1] - o) / s + 1, VEC3D_PRECISION) - EPS) * s + o
                 for b, o, s in zip(self.bounds, grid_offset, grid_size)
             )
+        else:
+            assert mode == "round_down", "Typechecking error"
+            start_final = tuple(
+                floor(round((b[0] - o) / s, VEC3D_PRECISION)) * s + o
+                for b, o, s in zip(self.bounds, grid_offset, grid_size)
+            )
+            end_final = tuple(
+                floor(round((b[1] - o) / s, VEC3D_PRECISION)) * s + o
+                for b, o, s in zip(self.bounds, grid_offset, grid_size)
+            )
+
         return BBox3D.from_coords(
             start_coord=cast(tuple[float, float, float], start_final),
             end_coord=cast(tuple[float, float, float], end_final),

zetta_utils/layer/precomputed.py

@@ -24,7 +24,7 @@
 # wrapper to cache using absolute paths with '/info'.
 # invalidates the cached infofile if the infofile is local and has since been deleted.
 def get_info(path: str) -> dict[str, Any]:
-    info_path = _to_info_path(path)
+    info_path = to_info_path(path)
     if is_local(info_path):
         if not CloudFile(info_path).exists():
             _info_cache.pop(_info_hash_key(info_path), None)
@@ -126,8 +126,9 @@ def from_optional_reference(
                     missing_params.append("bbox")
 
                 raise ValueError(
-                    f"When 'reference_path' is None, the following parameters must be "
-                    f"specified and cannot be None: {', '.join(missing_params)}"
+                    f"When not giving a reference info file through 'reference_path', "
+                    f"the following parameters must be specified and cannot be None: "
+                    f"{', '.join(missing_params)}"
                 )
             if inherit_all_params:
                 raise ValueError(
@@ -281,6 +282,13 @@ def __attrs_post_init__(self):
         ):
             raise ValueError("Exactly one of `info_path`/`info_spec_params` must be provided")
 
+        if self.info_path is not None and not CloudFile(to_info_path(self.info_path)).exists():
+            raise FileNotFoundError(
+                f"The infofile at '{self.info_path}' does not exist. If you are trying to create "
+                "a new layer, you must provide 'info_scales' to specify the scales for the info "
+                "file. If you are trying to use an existing layer, make sure the path is correct."
+            )
+
     def make_info(self) -> dict:
         if self.info_path is not None:
             return get_info(self.info_path)
@@ -326,7 +334,7 @@ def update_info(self, path: str, overwrite: bool, keep_existing_scales: bool) ->
         """
         try:
             existing_info = get_info(path)
-        except FileNotFoundError:
+        except (FileNotFoundError, KeyError):
             existing_info = None
 
         new_info = self.make_info()
@@ -370,7 +378,7 @@ def update_info(self, path: str, overwrite: bool, keep_existing_scales: bool) ->
 
             if existing_info != new_info:
                 _write_info(new_info, path)
-                _info_cache[_info_hash_key(_to_info_path(path))] = new_info
+                _info_cache[_info_hash_key(to_info_path(path))] = new_info
                 return True
 
         return False
@@ -406,11 +414,11 @@ def _get_info_from_info_path(info_path: str) -> dict[str, Any]:
 
 
 def _write_info(info: dict[str, Any], path: str) -> None:  # pragma: no cover
-    info_path = _to_info_path(path)
+    info_path = to_info_path(path)
     CloudFile(info_path).put_json(info, cache_control="no-cache")
 
 
-def _to_info_path(path: str) -> str:
+def to_info_path(path: str) -> str:
     if not path.endswith("/info"):
         path = os.path.join(path, "info")
     return abspath(path)

zetta_utils/layer/volumetric/cloudvol/build.py

@@ -135,6 +135,7 @@ def build_cv_layer(  # pylint: disable=too-many-locals
                 "specified path is assumed to exist and will be used without modifications. "
                 "Therefore, all other 'info_*' parameters must be None as they won't be used."
             )
+
         info_spec = PrecomputedInfoSpec(info_path=path)
 
     backend = CVBackend(

zetta_utils/training/datasets/collection_dataset.py

@@ -78,7 +78,7 @@ def build_collection_dataset(  # pylint: disable=too-many-locals
         this_resolution = [resolution[0], resolution[1], z_resolution]
         if isinstance(annotation.ng_annotation, AxisAlignedBoundingBoxAnnotation):
             bbox = BBox3D.from_ng_bbox(annotation.ng_annotation, (1, 1, 1)).snapped(
-                (0, 0, 0), this_resolution, "shrink"
+                (0, 0, 0), this_resolution, "round_down"
             )
 
             this_dset = LayerDataset(
@@ -94,7 +94,8 @@ def build_collection_dataset(  # pylint: disable=too-many-locals
             if len(this_dset) == 0:
                 raise RuntimeError(
                     f"The following annotation indicates bounding box {bbox} "
-                    f"which is smaller than the chunk size {chunk_size}"
+                    f"which is smaller than the chunk size {chunk_size} at "
+                    f"resolution {this_resolution} ({bbox.to_slices(this_resolution)})"
                 )
             datasets[str(i)] = this_dset
     dset = JointDataset(mode="vertical", datasets=datasets)

zetta_utils/training/lightning/train.py

@@ -116,9 +116,9 @@ def lightning_train(
         return
 
     if image is None:
-        raise ValueError("Must provide a container image for remote training.")
+        raise ValueError("Must provide a `image` when local_run=False.")
     if resource_limits is None:
-        raise ValueError("Must provide resource limits for remote training.")
+        raise ValueError("Must provide `resource_limits` when local_run=False.")
 
     assert resource_allocation.gcloud.check_image_exists(image), image
 
@@ -144,7 +144,7 @@ def lightning_train(
         if isinstance(v, dict):
             # argument given as spec, use it directly
             train_args[k] = v
-        else:
+        elif v is not None:
             arg_spec = builder.get_initial_builder_spec(v)
             if arg_spec is None:
                 raise RuntimeError(

zetta_utils/viz/rendering.py

@@ -129,7 +129,7 @@ def render_fld(  # pylint: disable=too-many-locals,too-many-arguments
     fld: Tensor,
     figsize: tuple[int, int] = (8, 8),
     dpi: int = 80,
-    grid_size: int = 50,
+    grid_size: int = 20,
     alpha: float = 0.6,
     linewidth: float = 0.5,
 ) -> npt.NDArray:  # pragma: no cover