diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c49c7e..9fe8fe6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ match the `VERSION` file and `v*` git tags. ### Changed +- **DBE's entropy filter runs a native kernel instead of a Python loop of ~4000 `np.histogram` calls.** + Profiling a real `--auto` run (which sets `entropy_bg=True` for most target types) found this filter -- + documented as "cheap" in the code it lives next to -- costing 4.2s on its own, because it runs on every + DBE pass regardless of session size. Fused into one native pass (`patch_entropy_batch`); not bit-exact + (numpy's histogram has a floating-point edge-correction step this doesn't replicate, so entropy values + agree to ~3e-4 absolute, immaterial for the median+MAD threshold this feeds). Cut a real profiled run's + Post-process time by ~2.6s. + - **`--auto`'s robust_pca auto-upgrade now applies to calibration libraries up to 25 frames, not 10.** The threshold was set from a 2026-09-22 benchmark on a synthetic RGB-shaped (2000x3000x3) array that overstated real cost -- this camera's actual calibration frames are raw mono FITS (2048x3056, no x3). diff --git a/CLAUDE.md b/CLAUDE.md index 097d5f0..f565b35 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -223,7 +223,7 @@ with `bayerPattern` to override. - **Online (streaming) sigma-clip** (`src/stacking.py`, used by `--stream`/`src/stream_stack.py`): `online_sigma_clip_seed_burnin` seeds a running Welford `(mean, m2, n_acc)` state from a small MAD-rejected burn-in window, and `online_sigma_clip_fold_frame` folds one further frame in at a time — the per-frame step of a true frame-at-a-time streaming combine, as an alternative to materializing the whole `(N,H,W,C)` aligned stack `sigma_clip_combine` needs. `online_sigma_clip_fold_frame` updates its `(mean, m2, n_acc)` state arrays **in place** (the one exception to this file's alloc-and-return convention) — reallocating and copying three full-resolution float64 arrays on every accepted frame was real allocator churn at full sensor resolution. `online_sigma_clip_combine` (whole-array, benchmark-only) is also available for throughput comparison against the batch kernels. - **Per-frame warp** (`src/registration.py` `apply_transform`): `warp_affine_lanczos3`, a 2D Lanczos-3 affine/shift resample (~5×). CPU path only (GPU unchanged). This is a quality-validated change, not numeric parity: FWHM and flux match scipy order-3, and the mild Lanczos ringing is incoherent across dithered frames (averages to zero in the stack). - **Anisotropic diffusion** (`src/denoising.py` `anisotropic_diffusion`): native Jacobi iteration with periodic boundary (~37×), exact parity. -- **DBE surface fit + patch sampler** (`src/background.py`): `dbe_fit_surface`, Gaussian-weighted local-linear regression with Tukey-biweight IRLS (~2.4× over the scipy RBF it replaced), and `dbe_sample_patches`, the per-patch rejection cascade (~31×). Not just a speedup — the RBF (thin-plate spline + hard outlier-rejection loop) was unbounded and could extrapolate wildly into rejected-sample gaps near bright stars; the local fit is bounded near the sample values by construction. Numpy mirror in `_dbe_fit_surface_numpy`, parity-tested. +- **DBE surface fit + patch sampler** (`src/background.py`): `dbe_fit_surface`, Gaussian-weighted local-linear regression with Tukey-biweight IRLS (~2.4× over the scipy RBF it replaced), and `dbe_sample_patches`, the per-patch rejection cascade (~31×). Not just a speedup — the RBF (thin-plate spline + hard outlier-rejection loop) was unbounded and could extrapolate wildly into rejected-sample gaps near bright stars; the local fit is bounded near the sample values by construction. Numpy mirror in `_dbe_fit_surface_numpy`, parity-tested. **`patch_entropy_batch`**: `dbe_sample_patches`'s own docstring calls the entropy filter "cheap, operates on the small per-patch result" -- true per-call, but under `--auto` (`entropy_bg=True` for most target types) it runs on every DBE pass regardless of session size, and profiling a real run found `_filter_sampled_patches`'s Python loop over up to `Config.DBE_MAX_SAMPLES` (4000) candidate patches, each calling `np.histogram`, costing 4.2s on its own. Fused into one rayon-parallel native pass. Not bit-exact -- numpy's histogram fast path has a floating-point edge-correction pass (compares each value against the actual `linspace` bin edges and nudges the index where rounding put it a bin off) this doesn't replicate, so entropy values agree to ~3e-4 absolute on a real-shaped case, not bit-for-bit; immaterial for what this feeds (a median+MAD outlier threshold over thousands of patches). Cut a real profiled run's Post-process time by ~2.6s. - **Cosmic-ray rejection** (`src/stacking.py` `lacosmic_reject`): `lacosmic_reject_native`, L.A.Cosmic-style Laplacian spike detection + 5×5 median replacement, f32 internally (~2×+ vs numpy; the memory-traffic halving matters more than raw compute under 16-way ProcessPool contention). - **Median filter** (`median_filter_native`): windowed median (reflect boundary) used by lacosmic and the hot-pixel detectors. Interior fast path with contiguous reads; 3×3 uses Paeth's 19-op branchless median network (~13×), 5×5 quickselect (~1.6×). - **Malvar-He-Cutler debayer** (`src/debayer.py` `debayer_malvar`, `--debayer-method malvar`, the default): the true published algorithm (Malvar, He & Cutler 2004), not an approximation — sparse per-pixel tap gather (each output pixel needs at most 2 of the 4 kernels; its own channel is the raw sample) rather than 4 whole-image convolutions, with an interior fast path (no per-tap boundary-index modulo) for all but a 2px border (~2×). Replaces the old cv2 `COLOR_BAYER_*_EA` path, which required requantizing to uint16 first (real precision loss) and only worked when cv2 was installed; this runs on the native float32 data with no dependency. Numpy mirror (`_debayer_malvar_numpy`) validated bit-exact against the `colour-demosaicing` package's reference Malvar2004 implementation across all 4 Bayer patterns (validation-only dependency, not required at runtime — see `tests/test_debayer_malvar.py`). diff --git a/ext/astro_native/Cargo.lock b/ext/astro_native/Cargo.lock index c6fdce9..67de6aa 100644 --- a/ext/astro_native/Cargo.lock +++ b/ext/astro_native/Cargo.lock @@ -49,7 +49,7 @@ checksum = "fb5dfbc6d8d2675589ccbe4d0fd61df2419075625f8c1a62325e718e2b0049f9" [[package]] name = "astro_native" -version = "0.34.0" +version = "0.35.0" dependencies = [ "numpy", "pyo3", diff --git a/ext/astro_native/Cargo.toml b/ext/astro_native/Cargo.toml index 5e0f08c..2b84ac3 100644 --- a/ext/astro_native/Cargo.toml +++ b/ext/astro_native/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "astro_native" -version = "0.34.0" +version = "0.35.0" edition = "2021" description = "Native (Rust) hot-path kernels for OriginStack: stacking combine, etc." diff --git a/ext/astro_native/pyproject.toml b/ext/astro_native/pyproject.toml index d14a42a..f380506 100644 --- a/ext/astro_native/pyproject.toml +++ b/ext/astro_native/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "astro_native" -version = "0.34.0" +version = "0.35.0" description = "Native Rust hot-path kernels for OriginStack" requires-python = ">=3.10" classifiers = ["Programming Language :: Rust"] diff --git a/ext/astro_native/src/lib.rs b/ext/astro_native/src/lib.rs index db59706..121d913 100644 --- a/ext/astro_native/src/lib.rs +++ b/ext/astro_native/src/lib.rs @@ -3328,6 +3328,152 @@ fn dbe_sample_patches<'py>( )) } +/// Shannon entropy of each sampled DBE patch's masked pixel values -- +/// `_filter_sampled_patches`'s entropy-filter loop (`src/background.py`), +/// one native pass over `coords` (rayon-parallel) instead of a Python loop +/// calling `np.histogram` per patch. `dbe_sample_patches` above documents +/// the entropy filter as staying in Python because it's "cheap, operates on +/// the small per-patch result" -- true for the per-patch histogram call +/// itself, but under `--auto` (which sets `entropy_bg=True` for most target +/// types) it runs on every DBE pass regardless of frame count, and a real +/// profiled run found it costing 4.2s on its own: `DBE_MAX_SAMPLES` (4000) +/// candidate patches is enough Python-loop-plus-per-call-numpy-overhead to +/// add up even though each individual histogram is genuinely small. +/// +/// `coords` is `(N, 2)` normalised `(row, col)` fractions, same convention +/// `dbe_sample_patches` returns and `_filter_sampled_patches` consumes: +/// `iy = clip(round(coord[0]*ny_g - 0.5), 0, ny_g-1)` locates the patch's +/// grid cell, whose pixel bounds are then `[round(iy*cell_h), round((iy+1)* +/// cell_h))` -- reproduced here exactly, not re-derived, so a coordinate +/// convention change in the sampler doesn't silently desync the two. +/// Histogram binning uses the same `(value - min) / range * n_bins` +/// uniform-bin formula `np.histogram`'s fast path computes for +/// `range=(mn, mx)` -- but not the edge-correction pass numpy adds after it +/// (comparing each value against the *actual* `linspace` bin edges and +/// nudging the index by one where float rounding put it a bin off), so an +/// occasional value within ~1 ULP of a bin boundary lands one bin over from +/// numpy's answer. Measured against the Python reference on a real-shaped +/// synthetic case: entropy values agree to ~3e-4 absolute (values are +/// O(1), so this is a few parts in 10,000), not bit-exact. Immaterial for +/// what this feeds -- a median+MAD outlier threshold over thousands of +/// patches -- so the edge-correction pass wasn't worth porting. +#[pyfunction] +#[pyo3(signature = (channel, emission_mask, coords, patch_size, n_bins=16))] +fn patch_entropy_batch<'py>( + py: Python<'py>, + channel: PyReadonlyArray2<'py, f32>, + emission_mask: PyReadonlyArray2<'py, f32>, + coords: PyReadonlyArray2<'py, f64>, + patch_size: usize, + n_bins: usize, +) -> PyResult>> { + let ch = channel.as_array(); + let em = emission_mask.as_array(); + let (h, w) = (ch.shape()[0], ch.shape()[1]); + if em.shape() != [h, w] { + return Err(pyo3::exceptions::PyValueError::new_err( + "channel and emission_mask must have the same shape", + )); + } + let ch_flat: Option<&[f32]> = ch.as_slice(); + let em_flat: Option<&[f32]> = em.as_slice(); + let (ch_flat, em_flat) = match (ch_flat, em_flat) { + (Some(c), Some(e)) => (c, e), + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "channel and emission_mask must be C-contiguous", + )) + } + }; + let coords_arr = coords.as_array(); + if coords_arr.shape()[1] != 2 { + return Err(pyo3::exceptions::PyValueError::new_err("coords must be (N, 2)")); + } + let coords_flat: &[f64] = coords_arr + .as_slice() + .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("coords must be C-contiguous"))?; + let n = coords_arr.shape()[0]; + let patch_size = patch_size.max(1); + let ny_g = (h / patch_size).max(1); + let nx_g = (w / patch_size).max(1); + let cell_h = h as f64 / ny_g as f64; + let cell_w = w as f64 / nx_g as f64; + + let mut out = vec![0f64; n]; + py.detach(|| { + out.par_iter_mut().enumerate().for_each(|(i, o)| { + let cy = coords_flat[i * 2]; + let cx = coords_flat[i * 2 + 1]; + let iy = ((cy * ny_g as f64 - 0.5).round() as isize).clamp(0, ny_g as isize - 1) as usize; + let ix = ((cx * nx_g as f64 - 0.5).round() as isize).clamp(0, nx_g as isize - 1) as usize; + let y0 = (iy as f64 * cell_h).round() as usize; + let y1 = (((iy + 1) as f64 * cell_h).round() as usize).min(h); + let x0 = (ix as f64 * cell_w).round() as usize; + let x1 = (((ix + 1) as f64 * cell_w).round() as usize).min(w); + + // f64 throughout the binning arithmetic, matching the numpy + // reference exactly: _patch_entropy's `mn, mx = float(pixels.min()), + // float(pixels.max())` casts to Python float (f64) before calling + // np.histogram, even though `pixels` itself is f32 -- numpy's + // uniform-bin fast path then promotes the per-element bin-index + // arithmetic to f64 too. Using f32 here would still be "close" but + // not actually bit-exact. + let mut mn = f64::INFINITY; + let mut mx = f64::NEG_INFINITY; + let mut count = 0usize; + for y in y0..y1 { + let row_base = y * w; + for x in x0..x1 { + if em_flat[row_base + x] < 0.5 { + let v = ch_flat[row_base + x] as f64; + if v < mn { + mn = v; + } + if v > mx { + mx = v; + } + count += 1; + } + } + } + if count < 4 { + *o = 0.0; + return; + } + let range = mx - mn; + if range < 1e-12 { + *o = 0.0; + return; + } + let mut counts = vec![0u32; n_bins]; + for y in y0..y1 { + let row_base = y * w; + for x in x0..x1 { + if em_flat[row_base + x] < 0.5 { + let v = ch_flat[row_base + x] as f64; + let mut idx = (((v - mn) / range) * n_bins as f64) as usize; + if idx >= n_bins { + idx = n_bins - 1; + } + counts[idx] += 1; + } + } + } + let total = count as f64; + let mut ent = 0.0f64; + for &c in &counts { + if c > 0 { + let p = c as f64 / (total + 1e-12); + ent -= p * p.log2(); + } + } + *o = ent; + }); + }); + + Ok(out.into_pyarray(py)) +} + // ============ Matched-filter star detection ============ // // Mirrors src/star_detect.py::_detect_stars_matched_filter_numpy exactly @@ -7790,6 +7936,7 @@ fn astro_native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(gaussian_filter_native, m)?)?; m.add_function(wrap_pyfunction!(dbe_fit_surface, m)?)?; m.add_function(wrap_pyfunction!(dbe_sample_patches, m)?)?; + m.add_function(wrap_pyfunction!(patch_entropy_batch, m)?)?; m.add_function(wrap_pyfunction!(detect_stars_matched_filter, m)?)?; m.add_function(wrap_pyfunction!(fit_rigid_ransac, m)?)?; m.add_function(wrap_pyfunction!(debayer_malvar, m)?)?; diff --git a/src/background.py b/src/background.py index 14bbb9c..e313e27 100644 --- a/src/background.py +++ b/src/background.py @@ -1135,24 +1135,35 @@ def _filter_sampled_patches(channel: np.ndarray, emission_mask: np.ndarray, # binary emission mask missed. Low-entropy patches (nearly uniform ADU # histogram) are genuine sky background samples. if use_entropy_weights and len(values) >= 8: - H_ch, W_ch = channel.shape - ny_g = max(1, H_ch // patch_size) - nx_g = max(1, W_ch // patch_size) - cell_h_g = H_ch / ny_g - cell_w_g = W_ch / nx_g - entropies = [] - for coord in coords: - iy = int(np.clip(round(coord[0] * ny_g - 0.5), 0, ny_g - 1)) - ix = int(np.clip(round(coord[1] * nx_g - 0.5), 0, nx_g - 1)) - y0 = int(round(iy * cell_h_g)) - y1 = min(int(round((iy + 1) * cell_h_g)), H_ch) - x0 = int(round(ix * cell_w_g)) - x1 = min(int(round((ix + 1) * cell_w_g)), W_ch) - em = emission_mask[y0:y1, x0:x1].ravel() - px = channel[y0:y1, x0:x1].ravel() - px = px[em < 0.5] - entropies.append(_patch_entropy(px)) - entropies = np.array(entropies, dtype=np.float64) + entropies = None + if HAS_NATIVE and hasattr(_native, 'patch_entropy_batch'): + try: + entropies = np.asarray(_native.patch_entropy_batch( + np.ascontiguousarray(channel, dtype=np.float32), + np.ascontiguousarray(emission_mask, dtype=np.float32), + np.ascontiguousarray(coords, dtype=np.float64), + int(patch_size), 16)) + except Exception: + entropies = None + if entropies is None: + H_ch, W_ch = channel.shape + ny_g = max(1, H_ch // patch_size) + nx_g = max(1, W_ch // patch_size) + cell_h_g = H_ch / ny_g + cell_w_g = W_ch / nx_g + entropies = [] + for coord in coords: + iy = int(np.clip(round(coord[0] * ny_g - 0.5), 0, ny_g - 1)) + ix = int(np.clip(round(coord[1] * nx_g - 0.5), 0, nx_g - 1)) + y0 = int(round(iy * cell_h_g)) + y1 = min(int(round((iy + 1) * cell_h_g)), H_ch) + x0 = int(round(ix * cell_w_g)) + x1 = min(int(round((ix + 1) * cell_w_g)), W_ch) + em = emission_mask[y0:y1, x0:x1].ravel() + px = channel[y0:y1, x0:x1].ravel() + px = px[em < 0.5] + entropies.append(_patch_entropy(px)) + entropies = np.array(entropies, dtype=np.float64) med_ent = float(np.median(entropies)) mad_ent = float(np.median(np.abs(entropies - med_ent))) ent_thresh = med_ent + 2.5 * 1.4826 * max(mad_ent, 1e-9) diff --git a/tests/test_native.py b/tests/test_native.py index 1153911..628459d 100644 --- a/tests/test_native.py +++ b/tests/test_native.py @@ -550,6 +550,121 @@ def test_gaussian_filter_native_zero_sigma_is_passthrough(): np.testing.assert_array_equal(got, a) +def _patch_entropy_ref(pixels, n_bins=16): + """The exact Python reference this test pins against -- + src.background._patch_entropy, duplicated here so the test doesn't + depend on background.py's own (possibly-native-routed) call path.""" + if len(pixels) < 4: + return 0.0 + mn, mx = float(pixels.min()), float(pixels.max()) + rng = mx - mn + if rng < 1e-12: + return 0.0 + counts, _ = np.histogram(pixels, bins=n_bins, range=(mn, mx)) + probs = counts / float(counts.sum() + 1e-12) + probs = probs[probs > 0] + return float(-np.sum(probs * np.log2(probs))) + + +def _filter_entropies_ref(channel, emission_mask, patch_size, coords, n_bins=16): + H_ch, W_ch = channel.shape + ny_g = max(1, H_ch // patch_size) + nx_g = max(1, W_ch // patch_size) + cell_h_g = H_ch / ny_g + cell_w_g = W_ch / nx_g + out = [] + for coord in coords: + iy = int(np.clip(round(coord[0] * ny_g - 0.5), 0, ny_g - 1)) + ix = int(np.clip(round(coord[1] * nx_g - 0.5), 0, nx_g - 1)) + y0 = int(round(iy * cell_h_g)) + y1 = min(int(round((iy + 1) * cell_h_g)), H_ch) + x0 = int(round(ix * cell_w_g)) + x1 = min(int(round((ix + 1) * cell_w_g)), W_ch) + em = emission_mask[y0:y1, x0:x1].ravel() + px = channel[y0:y1, x0:x1].ravel() + px = px[em < 0.5] + out.append(_patch_entropy_ref(px, n_bins)) + return np.array(out, dtype=np.float64) + + +def test_patch_entropy_batch_matches_python_reference(): + """patch_entropy_batch replaces _filter_sampled_patches's Python loop + (up to Config.DBE_MAX_SAMPLES=4000 per-patch np.histogram calls) -- + profiling a real --auto run (entropy_bg=True by default for most target + types) found it costing 4.2s on its own. Not bit-exact: numpy's + histogram fast path has a floating-point edge-correction pass this + doesn't replicate, so an occasional value within ~1 ULP of a bin + boundary lands one bin over. Tolerance set from a measured real-shaped + case (~3e-4 absolute on O(1) entropy values), which is immaterial for + what this feeds (a median+MAD outlier threshold).""" + rng = np.random.default_rng(11) + H, W = 800, 1200 + channel = rng.normal(1000.0, 50.0, (H, W)).astype(np.float32) + channel[100:120, 200:230] += 5000.0 # a 'star' to give real mask structure + emission_mask = np.zeros((H, W), dtype=np.float32) + emission_mask[100:120, 200:230] = 1.0 + + patch_size = 64 + ny_g = max(1, H // patch_size) + nx_g = max(1, W // patch_size) + n_coords = 500 + coords = np.column_stack([ + rng.integers(0, ny_g, n_coords) / ny_g + 0.5 / ny_g, + rng.integers(0, nx_g, n_coords) / nx_g + 0.5 / nx_g, + ]).astype(np.float64) + + ref = _filter_entropies_ref(channel, emission_mask, patch_size, coords) + got = np.asarray(native.patch_entropy_batch(channel, emission_mask, coords, patch_size, 16)) + np.testing.assert_allclose(got, ref, atol=1e-3) + + +def test_patch_entropy_batch_rejects_shape_mismatch(): + rng = np.random.default_rng(12) + channel = rng.normal(size=(40, 50)).astype(np.float32) + emission_mask = rng.normal(size=(40, 60)).astype(np.float32) # W mismatch + coords = np.zeros((3, 2), dtype=np.float64) + with pytest.raises(ValueError): + native.patch_entropy_batch(channel, emission_mask, coords, 16, 16) + + +def test_filter_sampled_patches_native_matches_numpy_fallback(): + """End-to-end: _filter_sampled_patches's native and numpy-fallback paths + must agree, not just the isolated kernel.""" + rng = np.random.default_rng(13) + H, W = 300, 400 + channel = rng.normal(1000.0, 50.0, (H, W)).astype(np.float32) + emission_mask = np.zeros((H, W), dtype=np.float32) + patch_size = 32 + ny_g = max(1, H // patch_size) + nx_g = max(1, W // patch_size) + n_coords = 40 + coords = np.column_stack([ + rng.integers(0, ny_g, n_coords) / ny_g + 0.5 / ny_g, + rng.integers(0, nx_g, n_coords) / nx_g + 0.5 / nx_g, + ]).astype(np.float64) + values = rng.normal(1000.0, 5.0, n_coords) + variances = rng.uniform(1.0, 30.0, n_coords) + + assert _background_mod.HAS_NATIVE and hasattr(_background_mod._native, 'patch_entropy_batch') + c_native, v_native = _background_mod._filter_sampled_patches( + channel, emission_mask, patch_size, coords.copy(), values.copy(), variances.copy(), True) + + had = _background_mod.HAS_NATIVE + _background_mod.HAS_NATIVE = False + try: + c_fallback, v_fallback = _background_mod._filter_sampled_patches( + channel, emission_mask, patch_size, coords.copy(), values.copy(), variances.copy(), True) + finally: + _background_mod.HAS_NATIVE = had + + # Not exact-equality: the native kernel's entropy values differ from the + # numpy reference by up to ~3e-4 (see patch_entropy_batch's docstring), + # which can shift which side of the median+MAD threshold a patch whose + # true entropy sits within that tolerance of it falls on. Check the two + # paths agree closely, not bit-for-bit. + assert abs(len(c_native) - len(c_fallback)) <= max(2, int(0.1 * n_coords)) + + def test_gaussian_filter_native_rejects_non_2d_by_signature(): # the numpy binding itself enforces 2D; the Python-side _gaussian_blur # wrapper is what actually gates this in production (see