diff --git a/python/xy/_native.py b/python/xy/_native.py index e0298da5..b9c07c41 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -24,7 +24,7 @@ from .config import MAX_CONTOUR_WORK, MAX_SCREEN_DIM -ABI_VERSION = 40 +ABI_VERSION = 41 # Rust reports invalid arguments (and, via the ffi_guard panic shield, any # internal panic) by returning `usize::MAX` from size-returning entry points. @@ -546,6 +546,18 @@ def _load() -> ctypes.CDLL: ctypes.c_double, ctypes.c_void_p, ] + lib.xy_polygon_select.restype = ctypes.c_size_t + lib.xy_polygon_select.argtypes = [ + ctypes.c_void_p, # x + ctypes.c_void_p, # y + ctypes.c_size_t, # len + ctypes.c_void_p, # candidate row ids + ctypes.c_size_t, # candidate count + ctypes.c_void_p, # polygon x + ctypes.c_void_p, # polygon y + ctypes.c_size_t, # polygon vertices + ctypes.c_void_p, # output row IDs + ] lib.xy_sample_mask.restype = ctypes.c_int32 lib.xy_sample_mask.argtypes = [ ctypes.c_void_p, @@ -2510,6 +2522,47 @@ def range_indices( return out if written == len(out) else out[:written].copy() +def polygon_select( + x: npt.NDArray[np.float64], + y: npt.NDArray[np.float64], + rows: npt.NDArray[np.uint32], + poly_x: npt.NDArray[np.float64], + poly_y: npt.NDArray[np.float64], +) -> npt.NDArray[np.uint32]: + """Those of `rows` whose canonical point is inside the lasso polygon. + + Even-odd ray casting, order preserved. Non-finite coordinates are never + inside, matching `range_indices`.""" + x = _as_f64(x, "x") + y = _as_f64(y, "y") + if len(x) != len(y): + raise ValueError("x and y must have equal length") + poly_x = _as_f64(poly_x, "polygon x") + poly_y = _as_f64(poly_y, "polygon y") + if len(poly_x) != len(poly_y): + raise ValueError("polygon x and y must have equal length") + rows = np.ascontiguousarray(rows, dtype=np.uint32) + if rows.ndim != 1: + raise ValueError("rows must be 1-D") + out = np.empty(len(rows), dtype=np.uint32) + if len(rows) == 0: + return out + written = _lib.xy_polygon_select( + _ptr_f64(x), + _ptr_f64(y), + len(x), + rows.ctypes.data, + len(rows), + _ptr_f64(poly_x), + _ptr_f64(poly_y), + len(poly_x), + out.ctypes.data, + ) + if written == _USIZE_MAX: + raise ValueError("invalid polygon_select arguments") + return out if written == len(out) else out[:written].copy() + + def sample_mask( ids: npt.NDArray[np.uint64], seed: int, diff --git a/python/xy/interaction.py b/python/xy/interaction.py index dc43875e..89057c51 100644 --- a/python/xy/interaction.py +++ b/python/xy/interaction.py @@ -225,6 +225,14 @@ def select_polygon( The polygon's bounding box first reuses the zone-pruned range predicate; ray casting then runs only on those candidates rather than every row. + + The cast itself is native (`kernels.polygon_select`). Vectorizing it over + NumPy meant one pass per polygon EDGE, each materializing full-length + crossing and edge-intersection temporaries: a 64-vertex lasso over 160k + candidates moved roughly half a gigabyte through memory to answer a + question that fits in registers. Crossing parity is order-independent, so + the native kernel walks edges inside the point loop and returns the same + rows. """ polygon = np.asarray(points, dtype=np.float64) if polygon.ndim != 2 or polygon.shape[1:] != (2,) or not 3 <= len(polygon) <= 2048: @@ -239,25 +247,15 @@ def select_polygon( float(polygon[:, 1].max()), trace_id, ) + poly_x = np.ascontiguousarray(polygon[:, 0]) + poly_y = np.ascontiguousarray(polygon[:, 1]) out: dict[int, np.ndarray] = {} for tid, rows in candidates.items(): if len(rows) == 0: out[tid] = rows continue trace = fig.traces[tid] - x = trace.x.values[rows] - y = trace.y.values[rows] - inside = np.zeros(len(rows), dtype=np.bool_) - j = len(polygon) - 1 - for i in range(len(polygon)): - xi, yi = polygon[i] - xj, yj = polygon[j] - crosses = (yi > y) != (yj > y) - with np.errstate(divide="ignore", invalid="ignore"): - edge_x = (xj - xi) * (y - yi) / (yj - yi) + xi - inside ^= crosses & (x < edge_x) - j = i - out[tid] = rows[inside] + out[tid] = kernels.polygon_select(trace.x.values, trace.y.values, rows, poly_x, poly_y) return out diff --git a/python/xy/kernels.py b/python/xy/kernels.py index e53cedf3..520663f2 100644 --- a/python/xy/kernels.py +++ b/python/xy/kernels.py @@ -62,6 +62,7 @@ valid_indices_f64 = _impl.valid_indices_f64 remap_u8 = _impl.remap_u8 range_indices = _impl.range_indices +polygon_select = _impl.polygon_select sample_mask = _impl.sample_mask sample_range_indices = _impl.sample_range_indices stratified_sample_range_u8 = _impl.stratified_sample_range_u8 @@ -121,6 +122,7 @@ "marching_triangles", "min_max", "normalize_f32", + "polygon_select", "polygon_triangles", "pyramid_append", "pyramid_build", diff --git a/spec/design/rust-engine.md b/spec/design/rust-engine.md index b8a2e3b5..8b5641ef 100644 --- a/spec/design/rust-engine.md +++ b/spec/design/rust-engine.md @@ -25,7 +25,7 @@ clear ImportError when it can't load, with no pure-Python fallback. | Concern | Today | Verdict | |---|---|---| -| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, local_log_density | Rust (ABI v40) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | +| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, polygon (lasso) selection, local_log_density | Rust (ABI v41) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; lasso ray casting (`xy_polygon_select`, §34) walks edges inside the point loop instead of one NumPy pass per edge, and buckets edges by y so a point tests only those spanning its row — the answer is the same crossing parity, without the per-edge full-length temporaries that made a 2048-vertex lasso cost ~370 ms over 160k candidates. Bucketing is declined below 16 vertices (the build outweighs the scans it saves) and for a non-finite polygon (slab bounds stop meaning anything); above 16 the slab count tracks the vertex count, capped at 256 and shed further so the CSR never exceeds 2²⁰ entries, which bounds both the allocation and its u32 cursors independently of what a caller passes; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | | fixed-width string/bytes/bool factorization | Rust (ABI v36) | correct — compact palettes use a bounded L1-resident codebook with full-record collision checks and emit exact counts; U1 uses a direct Unicode-scalar table with endian support; ≥512k rows probe a prefix then encode disjoint chunks in parallel, merging late labels by canonical first-row order before any retry; Python sees only unique labels and retains display-label ordering policy | | static display-list raster, row-banded polyline/point/segment paint, batched fill+stroke triangle meshes, affine scatter projection plus typed color/size resolution, density/heatmap colormap and sampling | Rust (ABI v36) | correct — commands borrow f32/u8 payload or canonical spans synchronously; compact stratified sampling reuses factorization counts; batched/banded output is byte-identical | | signal processing: `xy_rfft`, `xy_welch_spectra`, `xy_spectrogram` | Rust (ABI v36) | correct — O(N) transforms over sample columns; window/segment policy stays in Python | diff --git a/src/kernels.rs b/src/kernels.rs index 339e99e8..d5991b19 100644 --- a/src/kernels.rs +++ b/src/kernels.rs @@ -4778,6 +4778,194 @@ pub fn range_indices( range_indices_impl(x, y, lo_x, hi_x, lo_y, hi_y, par_threads(x.len()), out) } +/// Which of `rows` land inside the lasso polygon, by even-odd ray casting +/// (§34 selection). Writes the surviving canonical row ids to `out` (capacity +/// `rows.len()`, order preserved) and returns how many. +/// +/// The predicate is edge-for-edge the one the Python selection path used: for +/// each edge from vertex `j` to vertex `i`, the horizontal ray is crossed when +/// `(y_i > y) != (y_j > y)` and `x < (x_j - x_i) * (y - y_i) / (y_j - y_i) + +/// x_i`, and membership is the parity of the crossings. `y_j - y_i` cannot be +/// zero where the endpoints straddle the ray, so the guarded division here +/// covers exactly the lanes the vectorized form computed unconditionally and +/// then masked its inf/NaN away. `x_j - x_i` and `y_j - y_i` are hoisted per +/// edge — exact subtractions, so the arithmetic is unchanged; the division is +/// NOT turned into a reciprocal multiply, which would not be. +/// +/// Parity is order-independent, so walking edges inside the point loop rather +/// than points inside the edge loop gives the same answer while holding the +/// whole polygon in registers: no per-edge full-length temporaries (a +/// 64-vertex lasso over 160k candidates moved ~0.5 GB through memory). +/// +/// A non-finite coordinate fails every comparison and so is never inside, +/// matching both the vectorized predicate and `range_indices`. +pub fn polygon_select( + x: &[f64], + y: &[f64], + rows: &[u32], + poly_x: &[f64], + poly_y: &[f64], + out: &mut [u32], +) -> usize { + assert_eq!(x.len(), y.len()); + assert_eq!(poly_x.len(), poly_y.len()); + assert!(out.len() >= rows.len()); + let n = poly_x.len(); + if n < 3 { + return 0; + } + // (x_i, y_i, y_j, x_j - x_i, y_j - y_i) per edge, hoisted so the inner + // loop is a bounds-check-free walk over a couple of KB. + let edges: Vec<(f64, f64, f64, f64, f64)> = (0..n) + .map(|i| { + let j = if i == 0 { n - 1 } else { i - 1 }; + ( + poly_x[i], + poly_y[i], + poly_y[j], + poly_x[j] - poly_x[i], + poly_y[j] - poly_y[i], + ) + }) + .collect(); + let index = PolygonSlabs::build(poly_y, &edges); + let mut write = 0; + for &r in rows { + let (xv, yv) = (x[r as usize], y[r as usize]); + let mut inside = false; + // An edge can only flip parity for a point inside its own y-extent, + // so the slab index hands back a superset of the edges that can + // cross this row and the rest are provably no-ops — same crossings, + // same parity, but a freehand lasso's hundreds of vertices stop + // being scanned in full for every candidate point. + for &e in index.candidates(yv) { + let (xi, yi, yj, dx, dy) = edges[e as usize]; + if ((yi > yv) != (yj > yv)) && xv < dx * (yv - yi) / dy + xi { + inside = !inside; + } + } + if inside { + out[write] = r; + write += 1; + } + } + write +} + +/// Ceiling on `PolygonSlabs` CSR entries (4 MB of u32 at the worst case), so +/// neither the allocation nor the u32 cursors depend on caller-supplied +/// vertex counts. +const POLYGON_INDEX_RECORD_BUDGET: usize = 1 << 20; + +/// Edges bucketed by y, so a point tests only those spanning its row. +struct PolygonSlabs { + y0: f64, + y1: f64, + inv_h: f64, + n: usize, + /// CSR over slabs into `items`; empty when the index was declined. + starts: Vec, + items: Vec, + /// Every edge, in order — the answer when the index was declined. + all: Vec, +} + +impl PolygonSlabs { + /// Bucketing is declined for a non-finite polygon (slab bounds stop + /// meaning anything), for a vertex count too small to pay for the build + /// (< 16), and for one so large that `POLYGON_INDEX_RECORD_BUDGET` leaves + /// fewer than 8 slabs to sort into. Slab count otherwise tracks the vertex + /// count, capped at 256. + fn build(poly_y: &[f64], edges: &[(f64, f64, f64, f64, f64)]) -> Self { + let n_edges = edges.len(); + let unindexed = || Self { + y0: 0.0, + y1: 0.0, + inv_h: 0.0, + n: 0, + starts: Vec::new(), + items: Vec::new(), + all: (0..n_edges as u32).collect(), + }; + if n_edges < 16 || !poly_y.iter().all(|v| v.is_finite()) { + return unindexed(); + } + let mut y0 = f64::INFINITY; + let mut y1 = f64::NEG_INFINITY; + for &v in poly_y { + y0 = y0.min(v); + y1 = y1.max(v); + } + let height = y1 - y0; + if height <= 0.0 { + return unindexed(); + } + // An edge is filed once per slab its y-extent spans, so the index + // holds at most `edges x slabs` entries. Capping that product bounds + // both the allocation and the u32 cursors, and shedding slabs rather + // than refusing outright keeps a huge polygon indexed — the fallback + // scan is O(edges) for every point, which is exactly what hurts most + // at high vertex counts. `select_polygon` caps callers at 2048 + // vertices, well inside the full 256 slabs; this guards the kernel's + // own contract, which has no such ceiling. + let n = n_edges + .clamp(16, 256) + .min((POLYGON_INDEX_RECORD_BUDGET / n_edges).max(1)); + if n < 8 { + return unindexed(); + } + let inv_h = n as f64 / height; + // Half-open [lo, hi) per edge, matching the crossing test's own + // convention, so an edge lands in exactly the slabs it can cross. + let slab_of = |v: f64| (((v - y0) * inv_h) as usize).min(n - 1); + let span = |e: &(f64, f64, f64, f64, f64)| { + let (lo, hi) = if e.1 <= e.2 { (e.1, e.2) } else { (e.2, e.1) }; + (slab_of(lo), slab_of(hi)) + }; + let mut counts = vec![0u32; n + 1]; + for e in edges { + let (a, b) = span(e); + for c in counts.iter_mut().take(b + 2).skip(a + 1) { + *c += 1; + } + } + for s in 0..n { + counts[s + 1] += counts[s]; + } + let mut fill = counts.clone(); + let mut items = vec![0u32; counts[n] as usize]; + for (i, e) in edges.iter().enumerate() { + let (a, b) = span(e); + for slot in fill.iter_mut().take(b + 1).skip(a) { + items[*slot as usize] = i as u32; + *slot += 1; + } + } + Self { + y0, + y1, + inv_h, + n, + starts: counts, + items, + all: Vec::new(), + } + } + + /// Edges that could cross the ray at `yv`. A y outside the polygon's own + /// extent (or NaN) crosses nothing. + fn candidates(&self, yv: f64) -> &[u32] { + if self.n == 0 { + return &self.all; + } + if !(yv >= self.y0 && yv < self.y1) { + return &[]; + } + let s = (((yv - self.y0) * self.inv_h) as usize).min(self.n - 1); + &self.items[self.starts[s] as usize..self.starts[s + 1] as usize] + } +} + /// Append the global indices (`base + i`) of in-window rows to `out`, /// returning the match count. NaN fails every comparison → skipped, matching /// the historical behavior. Dispatches to the AVX2 clone when available. @@ -5080,6 +5268,158 @@ fn min_max_impl(data: &[f64], threads: usize) -> Option<(f64, f64)> { mod tests { use super::*; + /// The vectorized formulation `polygon_select` replaced, kept as the + /// executable definition of the predicate: one pass per edge, crossing + /// parity accumulated over every candidate. + fn polygon_reference( + x: &[f64], + y: &[f64], + rows: &[u32], + poly_x: &[f64], + poly_y: &[f64], + ) -> Vec { + let n = poly_x.len(); + let mut inside = vec![false; rows.len()]; + let mut j = n - 1; + for i in 0..n { + let (xi, yi) = (poly_x[i], poly_y[i]); + let (xj, yj) = (poly_x[j], poly_y[j]); + for (k, &r) in rows.iter().enumerate() { + let (xv, yv) = (x[r as usize], y[r as usize]); + let crosses = (yi > yv) != (yj > yv); + let edge_x = (xj - xi) * (yv - yi) / (yj - yi) + xi; + inside[k] ^= crosses && xv < edge_x; + } + j = i; + } + rows.iter() + .zip(&inside) + .filter(|(_, &b)| b) + .map(|(&r, _)| r) + .collect() + } + + fn polygon_case(poly_x: &[f64], poly_y: &[f64], x: &[f64], y: &[f64]) { + let rows: Vec = (0..x.len() as u32).collect(); + let mut out = vec![0u32; rows.len()]; + let n = polygon_select(x, y, &rows, poly_x, poly_y, &mut out); + assert_eq!( + &out[..n], + &polygon_reference(x, y, &rows, poly_x, poly_y)[..], + "polygon with {} vertices", + poly_x.len() + ); + } + + /// The y-slab index must only ever restrict the candidate edge set, never + /// the answer: indexed (>= 16 vertices) and unindexed paths agree with the + /// per-edge reference on convex, concave, self-intersecting, degenerate + /// and axis-aligned polygons, and on non-finite data. + #[test] + fn polygon_select_matches_the_per_edge_reference_indexed_and_not() { + let mut x = Vec::new(); + let mut y = Vec::new(); + // Deterministic lattice plus points landing exactly on vertices/edges. + for i in 0..60 { + for j in 0..60 { + x.push(i as f64 * 2.0); + y.push(j as f64 * 2.0); + } + } + x.extend_from_slice(&[f64::NAN, 0.0, f64::INFINITY, -f64::INFINITY, 50.0]); + y.extend_from_slice(&[50.0, f64::NAN, 50.0, 50.0, f64::INFINITY]); + + // Triangle and axis-aligned box: below the index threshold, and the + // box's horizontal edges are the zero-height ones. + polygon_case(&[10.0, 90.0, 50.0], &[10.0, 20.0, 80.0], &x, &y); + polygon_case(&[20.0, 80.0, 80.0, 20.0], &[20.0, 20.0, 80.0, 80.0], &x, &y); + // Bowtie: self-intersecting, so even-odd parity is load-bearing. + polygon_case(&[10.0, 90.0, 90.0, 10.0], &[10.0, 90.0, 10.0, 90.0], &x, &y); + // Collinear and zero-height polygons decline the index. + polygon_case(&[0.0, 50.0, 100.0], &[0.0, 50.0, 100.0], &x, &y); + polygon_case(&[0.0, 50.0, 100.0], &[30.0, 30.0, 30.0], &x, &y); + + // Indexed paths: circles and a concave star past the 16-vertex gate. + for n in [16usize, 17, 64, 257] { + let (px, py): (Vec, Vec) = (0..n) + .map(|k| { + let t = std::f64::consts::TAU * k as f64 / n as f64; + (50.0 + 30.0 * t.cos(), 50.0 + 30.0 * t.sin()) + }) + .unzip(); + polygon_case(&px, &py, &x, &y); + } + let (px, py): (Vec, Vec) = (0..24) + .map(|k| { + let t = std::f64::consts::TAU * k as f64 / 24.0; + let r = if k % 2 == 0 { 40.0 } else { 12.0 }; + (50.0 + r * t.cos(), 50.0 + r * t.sin()) + }) + .unzip(); + polygon_case(&px, &py, &x, &y); + + // A non-finite vertex declines the index; the scan must still agree. + polygon_case( + &[10.0, 90.0, 90.0, 10.0, 50.0], + &[10.0, 10.0, 90.0, 90.0, f64::NAN], + &x, + &y, + ); + + // Fewer than three vertices selects nothing. + let rows: Vec = (0..x.len() as u32).collect(); + let mut out = vec![0u32; rows.len()]; + assert_eq!( + polygon_select(&x, &y, &rows, &[0.0, 9.0], &[0.0, 9.0], &mut out), + 0 + ); + assert_eq!( + polygon_select(&x, &y, &[], &[0.0, 9.0, 5.0], &[0.0, 9.0, 5.0], &mut out), + 0 + ); + } + + /// `polygon_select` is a public export with no polygon-size ceiling of its + /// own — `Figure.select_polygon` caps callers at 2048 vertices, the kernel + /// does not — so the slab index must bound its own CSR rather than trust + /// the vertex count it is handed. + #[test] + fn polygon_slab_index_stays_inside_its_record_budget() { + for n in [16usize, 2048, 65_536, 200_000] { + // Worst case for the CSR: every edge spans the full height, so + // each is filed into every slab. + let poly_x: Vec = (0..n).map(|i| i as f64).collect(); + let poly_y: Vec = (0..n) + .map(|i| if i % 2 == 0 { 0.0 } else { 100.0 }) + .collect(); + let edges: Vec<(f64, f64, f64, f64, f64)> = (0..n) + .map(|i| { + let j = if i == 0 { n - 1 } else { i - 1 }; + ( + poly_x[i], + poly_y[i], + poly_y[j], + poly_x[j] - poly_x[i], + poly_y[j] - poly_y[i], + ) + }) + .collect(); + let index = PolygonSlabs::build(&poly_y, &edges); + assert!( + index.items.len() <= POLYGON_INDEX_RECORD_BUDGET, + "{n} vertices produced {} CSR entries", + index.items.len() + ); + // Either the index was declined outright or it kept enough slabs + // to be worth consulting. + assert!( + index.n == 0 || index.n >= 8, + "{n} vertices -> {} slabs", + index.n + ); + } + } + #[test] fn factorize_fixed_preserves_first_seen_codes_and_full_record_identity() { let rows = [b"ab\0", b"xy\0", b"ab\0", b"abx", b"xy\0"]; diff --git a/src/lib.rs b/src/lib.rs index d568d998..0d893374 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -82,7 +82,7 @@ unsafe fn borrowed_byte_spans<'a>( /// ABI version — bumped on any signature change. The Python wrapper checks this /// at load time and refuses a mismatched library loudly (§33 comm-versioning /// rule, applied to the in-process boundary). -pub const ABI_VERSION: u32 = 40; +pub const ABI_VERSION: u32 = 41; const FACTORIZE_CAPACITY_EXCEEDED: usize = usize::MAX - 1; #[no_mangle] @@ -2797,6 +2797,60 @@ pub unsafe extern "C" fn xy_range_indices( }) } +/// Canonical row ids from `rows` that fall inside the lasso polygon, by +/// even-odd ray casting. Returns the count written; `out` must hold +/// `n_rows` u32s. A polygon of fewer than 3 vertices selects nothing. +/// +/// # Safety +/// `x`/`y` must point to `len` readable f64s, `rows` to `n_rows` readable +/// u32s, `poly_x`/`poly_y` to `n_poly` readable f64s, and `out` to `n_rows` +/// writable u32s. +#[no_mangle] +#[allow(clippy::too_many_arguments)] +pub unsafe extern "C" fn xy_polygon_select( + x: *const f64, + y: *const f64, + len: usize, + rows: *const u32, + n_rows: usize, + poly_x: *const f64, + poly_y: *const f64, + n_poly: usize, + out: *mut u32, +) -> usize { + // u32 index ceiling — see xy_m4_indices. + if len > u32::MAX as usize { + return usize::MAX; + } + if n_rows == 0 { + return 0; + } + // Fewer than three vertices encloses nothing. Answer before building any + // slice: `from_raw_parts` requires a non-null, aligned pointer even at + // length zero, so a caller passing null for an empty polygon must not + // reach the constructions below. + if n_poly < 3 { + return 0; + } + if x.is_null() || y.is_null() || rows.is_null() || out.is_null() { + return usize::MAX; + } + if poly_x.is_null() || poly_y.is_null() { + return usize::MAX; + } + let x = std::slice::from_raw_parts(x, len); + let y = std::slice::from_raw_parts(y, len); + let rows = std::slice::from_raw_parts(rows, n_rows); + let poly_x = std::slice::from_raw_parts(poly_x, n_poly); + let poly_y = std::slice::from_raw_parts(poly_y, n_poly); + let out = std::slice::from_raw_parts_mut(out, n_rows); + // An out-of-range row id panics the bounds check inside; ffi_guard turns + // that into the error sentinel rather than unwinding across the ABI. + ffi_guard(usize::MAX, || { + kernels::polygon_select(x, y, rows, poly_x, poly_y, out) + }) +} + /// Per-point local log density for a subset. Returns 1 on success, 0 on invalid /// grid/window arguments. /// diff --git a/tests/test_select_polygon.py b/tests/test_select_polygon.py new file mode 100644 index 00000000..08ea099d --- /dev/null +++ b/tests/test_select_polygon.py @@ -0,0 +1,146 @@ +"""Lasso selection: the native ray cast must answer exactly what the +vectorized per-edge predicate answered (design dossier §34).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from xy import kernels +from xy._figure import Figure + + +def _per_edge_reference( + x: np.ndarray, y: np.ndarray, rows: np.ndarray, polygon: np.ndarray +) -> np.ndarray: + """One pass per polygon edge, crossing parity accumulated across rows. + + This is the definition `kernels.polygon_select` implements; keeping it + here makes the contract executable rather than a comment. + """ + xs, ys = x[rows], y[rows] + inside = np.zeros(len(rows), dtype=np.bool_) + j = len(polygon) - 1 + for i in range(len(polygon)): + xi, yi = polygon[i] + xj, yj = polygon[j] + crosses = (yi > ys) != (yj > ys) + with np.errstate(divide="ignore", invalid="ignore"): + edge_x = (xj - xi) * (ys - yi) / (yj - yi) + xi + inside ^= crosses & (xs < edge_x) + j = i + return rows[inside] + + +def _ring(n: int, *, r: float = 30.0, cx: float = 50.0, cy: float = 50.0) -> np.ndarray: + t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + return np.stack([cx + r * np.cos(t), cy + r * np.sin(t)], axis=1) + + +def _star(n: int = 24) -> np.ndarray: + t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + r = np.where(np.arange(n) % 2 == 0, 40.0, 12.0) + return np.stack([50.0 + r * np.cos(t), 50.0 + r * np.sin(t)], axis=1) + + +POLYGONS = { + "triangle": np.array([[10.0, 10.0], [90.0, 20.0], [50.0, 80.0]]), + # Horizontal edges are the zero-height ones the vectorized form divided + # by and masked away. + "axis_aligned_box": np.array([[20.0, 20.0], [80.0, 20.0], [80.0, 80.0], [20.0, 80.0]]), + # Self-intersecting, so even-odd parity is load-bearing. + "bowtie": np.array([[10.0, 10.0], [90.0, 90.0], [90.0, 10.0], [10.0, 90.0]]), + "collinear": np.array([[0.0, 0.0], [50.0, 50.0], [100.0, 100.0]]), + "repeated_vertex": np.array([[10.0, 10.0], [10.0, 10.0], [90.0, 90.0], [10.0, 90.0]]), + "flat": np.array([[0.0, 30.0], [50.0, 30.0], [100.0, 30.0]]), + # Straddles the >= 16-vertex gate where edges get bucketed by y. + "ring_15": _ring(15), + "ring_16": _ring(16), + "ring_64": _ring(64), + "ring_257": _ring(257), + # The API ceiling `Figure.select_polygon` enforces. + "ring_2048": _ring(2048), + "star_concave": _star(), + "covers_everything": _ring(32, r=500.0), + "encloses_nothing": _ring(32, r=1e-6), +} + + +@pytest.fixture(scope="module") +def sample() -> tuple[np.ndarray, np.ndarray]: + """A lattice (points land exactly on vertices and edges) plus the + non-finite rows a canonical column is allowed to carry (§19).""" + gx, gy = np.meshgrid(np.arange(0.0, 100.0, 2.0), np.arange(0.0, 100.0, 2.0)) + x = np.concatenate([gx.ravel(), [np.nan, 0.0, np.inf, -np.inf, 50.0, 20.0]]) + y = np.concatenate([gy.ravel(), [50.0, np.nan, 50.0, 50.0, np.inf, 20.0]]) + return np.ascontiguousarray(x), np.ascontiguousarray(y) + + +@pytest.mark.parametrize("name", sorted(POLYGONS)) +def test_native_cast_matches_the_per_edge_predicate(name, sample): + x, y = sample + polygon = POLYGONS[name] + rows = np.arange(len(x), dtype=np.uint32) + got = kernels.polygon_select( + x, y, rows, np.ascontiguousarray(polygon[:, 0]), np.ascontiguousarray(polygon[:, 1]) + ) + assert np.array_equal(got, _per_edge_reference(x, y, rows, polygon)) + assert got.dtype == np.uint32 + + +@pytest.mark.parametrize( + "rows", + [ + np.empty(0, dtype=np.uint32), + np.array([0], dtype=np.uint32), + np.arange(0, 2400, 7, dtype=np.uint32), + ], + ids=["empty", "single", "strided"], +) +def test_candidate_subsets_preserve_order_and_membership(rows, sample): + """The real call shape: `rows` is the zone-pruned bbox subset, and the + reply must stay in canonical ascending order.""" + x, y = sample + polygon = POLYGONS["ring_64"] + got = kernels.polygon_select( + x, y, rows, np.ascontiguousarray(polygon[:, 0]), np.ascontiguousarray(polygon[:, 1]) + ) + assert np.array_equal(got, _per_edge_reference(x, y, rows, polygon)) + assert np.all(np.diff(got.astype(np.int64)) > 0) + assert set(got.tolist()) <= set(rows.tolist()) + + +def test_figure_lasso_selects_the_disc_and_nothing_outside_it(): + rng = np.random.default_rng(67) + x = rng.uniform(0.0, 100.0, 20_000) + y = rng.uniform(0.0, 100.0, 20_000) + fig = Figure() + fig.scatter(x, y) + polygon = _ring(64, r=20.0).tolist() + + selected = fig.select_polygon(polygon)[0] + inside_radius = np.hypot(x[selected] - 50.0, y[selected] - 50.0) + # Every selected point is within the circumscribed radius; the inscribed + # radius bounds what the 64-gon must have caught. + assert inside_radius.max() <= 20.0 + caught = np.hypot(x - 50.0, y - 50.0) <= 20.0 * np.cos(np.pi / 64) + assert set(np.flatnonzero(caught).tolist()) <= set(selected.tolist()) + + +def test_polygon_vertex_and_row_validation(): + x = np.array([1.0, 2.0, 3.0]) + y = np.array([1.0, 2.0, 3.0]) + rows = np.array([0, 1, 2], dtype=np.uint32) + square_x = np.array([0.0, 4.0, 4.0, 0.0]) + square_y = np.array([0.0, 0.0, 4.0, 4.0]) + + # Fewer than three vertices encloses nothing rather than erroring. + assert len(kernels.polygon_select(x, y, rows, square_x[:2], square_y[:2])) == 0 + + with pytest.raises(ValueError): + kernels.polygon_select(x, y[:2], rows, square_x, square_y) + with pytest.raises(ValueError): + kernels.polygon_select(x, y, rows, square_x, square_y[:3]) + # A row id past the end of the column is a caller bug, not a silent drop. + with pytest.raises(ValueError): + kernels.polygon_select(x, y, np.array([99], dtype=np.uint32), square_x, square_y)