Inference performance: vectorized energy detection, fused per-channel workers, per-cadence streaming, downsample-at-extraction

README.md

@@ -667,6 +667,7 @@ usage: inference [-h] [--data-path DATA_PATH] [--model-path MODEL_PATH]
                  [--detection-window-size DETECTION_WINDOW_SIZE]
                  [--detection-step-size DETECTION_STEP_SIZE]
                  [--stat-threshold STAT_THRESHOLD] [--stamp-width STAMP_WIDTH]
+                 [--store-downsampled-stamps | --no-store-downsampled-stamps]
                  [--overlap-search | --no-overlap-search]
                  [--overlap-fraction OVERLAP_FRACTION]
                  [--preprocess-output-dir PREPROCESS_OUTPUT_DIR]
@@ -733,8 +734,10 @@ options:
                         Number of fine channels per coarse channel (default:
                         1048576)
   --parallel-coarse-chans PARALLEL_COARSE_CHANS
-                        Number of coarse channels to process in parallel per
-                        block (default: 28)
+                        Progress-logging chunk size for energy detection, in
+                        coarse channels per log line (default: the number of
+                        worker processes). Parallelism itself comes from the
+                        persistent worker pool, not this knob.
   --spline-order SPLINE_ORDER
                         Spline order for bandpass fitting (default: 16)
   --detection-window-size DETECTION_WINDOW_SIZE
@@ -749,6 +752,13 @@ options:
   --stamp-width STAMP_WIDTH
                         Width in fine channels of the extracted stamp around
                         each hit (default: 4096; must equal --width-bin)
+  --store-downsampled-stamps, --no-store-downsampled-stamps
+                        Downsample stamps along frequency (by --downsample-
+                        factor) at extraction time, storing stamp_width //
+                        downsample_factor bins per stamp (~8x smaller at
+                        defaults; default: enabled). Pass --no-store-
+                        downsampled-stamps to archive raw-resolution stamps;
+                        loading handles both layouts.
   --overlap-search, --no-overlap-search
                         Additionally extract stamps offset by
                         ±overlap_fraction*stamp_width around each hit. Pass

src/aetherscan/cli.py

@@ -695,7 +695,7 @@ def _add_inference_flags_to(parser):
         "--parallel-coarse-chans",
         type=int,
         default=None,
-        help="Number of coarse channels to process in parallel per block (default: 28)",
+        help="Progress-logging chunk size for energy detection, in coarse channels per log line (default: the number of worker processes). Parallelism itself comes from the persistent worker pool, not this knob.",
     )
     parser.add_argument(
         "--spline-order",
@@ -727,6 +727,12 @@ def _add_inference_flags_to(parser):
         default=None,
         help="Width in fine channels of the extracted stamp around each hit (default: 4096; must equal --width-bin)",
     )
+    parser.add_argument(
+        "--store-downsampled-stamps",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Downsample stamps along frequency (by --downsample-factor) at extraction time, storing stamp_width // downsample_factor bins per stamp (~8x smaller at defaults; default: enabled). Pass --no-store-downsampled-stamps to archive raw-resolution stamps; loading handles both layouts.",
+    )
     parser.add_argument(
         "--overlap-search",
         action=argparse.BooleanOptionalAction,
@@ -1038,6 +1044,11 @@ def apply_args_to_config(args: argparse.Namespace) -> None:
         config.inference.stat_threshold = args.stat_threshold
     if hasattr(args, "stamp_width") and args.stamp_width is not None:
         config.inference.stamp_width = args.stamp_width
+    # store_downsampled_stamps uses argparse.BooleanOptionalAction with default=None so that
+    # the CLI can express "leave the config default" (omit), "force on"
+    # (--store-downsampled-stamps), and "force off" (--no-store-downsampled-stamps)
+    if hasattr(args, "store_downsampled_stamps") and args.store_downsampled_stamps is not None:
+        config.inference.store_downsampled_stamps = args.store_downsampled_stamps
     # overlap_search uses argparse.BooleanOptionalAction with default=None so that
     # the CLI can express "leave the config default" (omit), "force on"
     # (--overlap-search), and "force off" (--no-overlap-search). The `is not None`
@@ -1710,6 +1721,32 @@ def collect_validation_errors(
                         fix_kind="clamp_low",
                     )
                 )
+            # Downsample-at-extraction stores stamp_width // downsample_factor bins per
+            # stamp, so the division must be exact for the stored width to be well-defined
+            store_ds = _resolve(
+                args, "store_downsampled_stamps", config.inference.store_downsampled_stamps
+            )
+            df_inf = _resolve(args, "downsample_factor", config.data.downsample_factor)
+            if (
+                store_ds
+                and sw_inf is not None
+                and df_inf is not None
+                and df_inf > 0
+                and sw_inf % df_inf != 0
+            ):
+                errors.append(
+                    ValidationError(
+                        field="inference.stamp_width",
+                        current=sw_inf,
+                        message=(
+                            f"--stamp-width ({sw_inf}) must be divisible by "
+                            f"--downsample-factor ({df_inf}) when stamps are downsampled at "
+                            "extraction (--store-downsampled-stamps)"
+                        ),
+                        fix_kind="divisibility",
+                        divisor=df_inf,
+                    )
+                )
 
         coarse_channel_width = _resolve(
             args, "coarse_channel_width", config.inference.coarse_channel_width

src/aetherscan/config.py

@@ -291,12 +291,19 @@ class InferenceConfig:
 
     # NOTE: come back to this later (are these params correct?)
     coarse_channel_width: int = 1048576
-    parallel_coarse_chans: int = 28
+    # Progress-logging chunk size for energy detection (coarse channels per log line).
+    # None -> use manager.n_processes. Parallelism itself comes from the persistent worker
+    # pool (one fused task per coarse channel), not from this knob.
+    parallel_coarse_chans: int | None = None
     spline_order: int = 16
     detection_window_size: int = 256
     detection_step_size: int = 128
     stat_threshold: float = 2048.0
     stamp_width: int = 4096
+    # Downsample stamps along frequency at extraction time (by data.downsample_factor), so
+    # the per-cadence .npy stores stamp_width // downsample_factor bins (~8x smaller at
+    # defaults). Disable to archive raw-resolution stamps; loading handles both layouts.
+    store_downsampled_stamps: bool = True
 
     # NOTE: come back to this later (is overlap search implemented correctly? redo analysis from peter forwards search paper)
     overlap_search: bool = True
@@ -571,6 +578,7 @@ def to_dict(self) -> dict:
                 "detection_step_size": self.inference.detection_step_size,
                 "stat_threshold": self.inference.stat_threshold,
                 "stamp_width": self.inference.stamp_width,
+                "store_downsampled_stamps": self.inference.store_downsampled_stamps,
                 "overlap_search": self.inference.overlap_search,
                 "overlap_fraction": self.inference.overlap_fraction,
                 "discard_side_channels": self.inference.discard_side_channels,