diff --git a/mortie/__init__.py b/mortie/__init__.py index aa6bfc1..9f14531 100644 --- a/mortie/__init__.py +++ b/mortie/__init__.py @@ -53,6 +53,7 @@ heal_norm, # Inverse functions infer_order_from_morton, + is_point, mort2bbox, mort2geo, mort2healpix, @@ -63,6 +64,7 @@ norm2mort, norm2uniq, order2res, + orders_of, res2display, uniq2geo, unique2parent, @@ -75,6 +77,8 @@ 'mort2bbox', 'mort2polygon', 'infer_order_from_morton', + 'orders_of', + 'is_point', 'validate_morton', 'mort2norm', 'norm2uniq', diff --git a/mortie/tests/test_coverage.py b/mortie/tests/test_coverage.py index 58a6f85..7b4de9d 100644 --- a/mortie/tests/test_coverage.py +++ b/mortie/tests/test_coverage.py @@ -710,7 +710,8 @@ def test_moc_high_order(self, order): sr_lons = [-76.55, -76.54997, -76.54997, -76.55, -76.55] moc = mortie.morton_coverage_moc(sr_lats, sr_lons, order=order) assert len(moc) > 0 - orders = mortie.infer_order_from_morton(moc) + # a MOC is mixed-order by construction: per-element orders (issue #116) + orders = mortie.orders_of(moc) assert int(np.max(orders)) <= order # Lossless densify: the MOC expands to exactly the flat order cover. flat = set(int(x) for x in mortie.morton_coverage(sr_lats, sr_lons, order=order)) @@ -722,7 +723,7 @@ def test_zagg_child_order_plus_three(self): (the old cap rejected order 22) and reach order 22 on the boundary.""" moc = mortie.morton_coverage_moc(self.RING_LATS, self.RING_LONS, order=22) assert len(moc) > 0 - assert int(np.max(mortie.infer_order_from_morton(moc))) == 22 + assert int(np.max(mortie.orders_of(moc))) == 22 def test_moc_densifies_to_flat_order_19(self): """MOC at order 19 densifies losslessly to the flat cover.""" diff --git a/mortie/tests/test_mixed_order.py b/mortie/tests/test_mixed_order.py new file mode 100644 index 0000000..66a8a2d --- /dev/null +++ b/mortie/tests/test_mixed_order.py @@ -0,0 +1,260 @@ +"""Mixed-order support in the geo kernels (issue #116). + +Covers the per-element introspection surface ``orders_of`` / ``is_point``: +pure-numpy suffix decodes of the packed word per the spec page's §1 suffix +table (``docs/specification.md``), golden-pinned against the table itself — +including the 28/47 preorder band boundaries and both hemispheres — and +cross-checked against the Rust kernel's depth decode so the two cannot drift. +""" + +import numpy as np +import pytest + +from mortie import ( + _rustie, + geo2mort, + infer_order_from_morton, + is_point, + mort2bbox, + mort2geo, + mort2polygon, + order2res, + orders_of, +) + +# One northern and one southern location: base cells 7-11 set bit 63, so the +# southern column exercises the large-unsigned half of the word space. +_LATS = np.array([45.0, -45.0]) +_LONS = np.array([45.0, -120.0]) + + +def _area(order): + return geo2mort(_LATS, _LONS, order=order) + + +def _point(): + return geo2mort(_LATS, _LONS, points=True) + + +class TestOrdersOf: + """Golden pins for the spec §1 suffix table.""" + + def test_suffix_0_to_27_order_is_suffix(self): + """Orders 0-27 store the order as the suffix value itself (spec §1).""" + for order in range(28): + words = _area(order) + assert np.all((words & np.uint64(0x3F)) == order) + np.testing.assert_array_equal(orders_of(words), order) + + def test_suffix_28_to_47_preorder_band(self): + """The 28..=47 band is parent-first preorder r = t28*5 + (t29+1 | 0): + each t28 owns a 5-block — the order-28 parent slot, then its four + order-29 children (spec §1).""" + # Craft every band suffix onto a real order-29 prefix+body (the body + # is fully populated for suffixes >= 28, so every value is canonical). + for base_word in _area(29): + stem = np.uint64(base_word) & ~np.uint64(0x3F) + for suffix in range(28, 48): + word = stem | np.uint64(suffix) + r = suffix - 28 + expected = 28 if r % 5 == 0 else 29 + assert orders_of(word)[0] == expected, ( + f"suffix {suffix}: expected order {expected}" + ) + assert not is_point(word)[0] + + def test_suffix_48_to_63_point_band(self): + """Suffixes 48..=63 are order-29 points, max-encoded (spec §1).""" + for base_word in _area(29): + stem = np.uint64(base_word) & ~np.uint64(0x3F) + for suffix in range(48, 64): + word = stem | np.uint64(suffix) + assert orders_of(word)[0] == 29 + assert is_point(word)[0] + + def test_matches_kernel_depth_decode(self): + """orders_of agrees with the Rust kernel's depth decode elementwise — + the numpy suffix table and the kernel cannot drift apart.""" + words = np.concatenate( + [_area(o) for o in (0, 1, 6, 13, 19, 24, 27, 28, 29)] + [_point()] + ) + _, depths = _rustie.rust_mort2nested(np.ascontiguousarray(words)) + np.testing.assert_array_equal(orders_of(words), depths) + + def test_shape_family_and_dtype(self): + """Scalar in -> length-1 uint8 ndarray (the geo2mort shape family).""" + got = orders_of(int(_area(6)[0])) + assert got.shape == (1,) and got.dtype == np.uint8 and got[0] == 6 + assert orders_of(np.array([], dtype=np.uint64)).shape == (0,) + + +class TestIsPoint: + def test_point_words_are_points(self): + np.testing.assert_array_equal(is_point(_point()), True) + + def test_area_words_are_not_points(self): + """Area words at every order — order-29 areas included — are not + points: kind rides the suffix band, not the order (spec §4).""" + for order in (0, 6, 28, 29): + np.testing.assert_array_equal(is_point(_area(order)), False) + + def test_shape_family(self): + got = is_point(int(_point()[0])) + assert got.shape == (1,) and got.dtype == np.bool_ and got[0] + + +class TestInferOrderMixedRaise: + """infer_order_from_morton returns one scalar order; mixed input raises + (issue #116 — it previously returned the first element's order silently).""" + + def test_uniform_scalar_and_array_unchanged(self): + for order in (0, 6, 29): + words = _area(order) + assert infer_order_from_morton(int(words[0])) == order + assert infer_order_from_morton(words) == order + assert infer_order_from_morton(_point()) == 29 + + def test_mixed_raises_naming_orders(self): + mixed = np.concatenate([_area(6), _area(19)]) + with pytest.raises(ValueError, match=r"Mixed orders.*\[6, 19\].*orders_of"): + infer_order_from_morton(mixed) + + def test_mixed_area_point_is_uniform_29(self): + """Order-29 areas and points share order 29 — kind is not order, so + this is NOT a mixed-order array (spec §4).""" + assert infer_order_from_morton(np.concatenate([_area(29), _point()])) == 29 + + +_MIXED_ORDERS = (0, 6, 19, 24, 29) + + +def _mixed_area(rng_seed=42): + """Interleaved mixed-order area words spanning both hemispheres.""" + words = np.concatenate([_area(o) for o in _MIXED_ORDERS]) + np.random.default_rng(rng_seed).shuffle(words) + return words + + +class TestGroupDispatchOracle: + """The correctness oracle (issue #116): group-by-order dispatch must equal + the per-order uniform kernel calls, scattered back to input positions.""" + + def test_mort2geo_matches_uniform(self): + words = _mixed_area() + orders = orders_of(words) + lat, lon = mort2geo(words) + assert lat.shape == lon.shape == words.shape + for order in np.unique(orders): + mask = orders == order + ulat, ulon = mort2geo(np.ascontiguousarray(words[mask])) + np.testing.assert_array_equal(lat[mask], ulat) + np.testing.assert_array_equal(lon[mask], ulon) + + def test_mort2geo_points_group_with_29(self): + """Point words dispatch through the order-29 group: locations match + the uniform all-point call (a point's location IS its mort2geo).""" + words = np.concatenate([_area(6), _point(), _area(19)]) + lat, lon = mort2geo(words) + plat, plon = mort2geo(_point()) + np.testing.assert_array_equal(lat[2:4], plat) + np.testing.assert_array_equal(lon[2:4], plon) + + def test_mort2bbox_matches_uniform(self): + words = _mixed_area() + orders = orders_of(words) + bboxes = mort2bbox(words) + assert len(bboxes) == words.size + for order in np.unique(orders): + (idx,) = np.nonzero(orders == order) + uniform = mort2bbox(np.ascontiguousarray(words[idx])) + assert [bboxes[i] for i in idx] == uniform + + def test_mort2polygon_matches_uniform(self): + for step in (1, 2): + words = _mixed_area() + orders = orders_of(words) + polygons = mort2polygon(words, step=step) + assert len(polygons) == words.size + # Rings are 4*step + 1 vertices at every order (closure included). + assert all(len(p) == 4 * step + 1 for p in polygons) + for order in np.unique(orders): + (idx,) = np.nonzero(orders == order) + uniform = mort2polygon(np.ascontiguousarray(words[idx]), step=step) + assert [polygons[i] for i in idx] == uniform + + def test_singleton_group_scatter(self): + """A group of exactly one element exercises the bare-dict/bare-ring + return of the length-1 uniform call; scatter must still place it.""" + words = np.concatenate([_area(6), _area(19)[:1]]) + bboxes = mort2bbox(words) + assert bboxes[2] == mort2bbox(int(words[2])) + polygons = mort2polygon(words) + assert polygons[2] == mort2polygon(int(words[2])) + lat, lon = mort2geo(words) + slat, slon = mort2geo(int(words[2])) + assert lat[2] == slat[0] and lon[2] == slon[0] + + def test_uniform_input_unchanged(self): + """Uniform arrays never enter the dispatch branch: results are the + pre-#116 uniform kernel outputs (scalar/array postures alike).""" + words = _area(6) + lat, lon = mort2geo(words) + slat, slon = mort2geo(int(words[0])) + assert lat[0] == slat[0] and lon[0] == slon[0] + assert mort2bbox(words)[0] == mort2bbox(int(words[0])) + assert mort2polygon(words)[0] == mort2polygon(int(words[0])) + + +class TestPointPostures: + """mort2polygon/mort2bbox on point words yield the geometry of the point's + containing order-29 cell (issue #116, espg review): a point still has a + well-defined containing box/polygon — the cell that contains it — and a + group of points covers a well-defined area element by element. The oracle + is the same location's order-29 AREA word: a point and the order-29 area + word at that location decode to the identical HEALPix cell (spec §4).""" + + def test_bbox_points_yield_containing_cell(self): + """A point's bbox equals the bbox of its containing order-29 cell (the + order-29 area word at the same location).""" + assert mort2bbox(_point()) == mort2bbox(_area(29)) + + def test_polygon_points_yield_containing_cell(self): + assert mort2polygon(_point()) == mort2polygon(_area(29)) + + def test_scalar_point_yields_containing_cell(self): + point_word = int(_point()[0]) + area_word = int(_area(29)[0]) + assert mort2bbox(point_word) == mort2bbox(area_word) + assert mort2polygon(point_word) == mort2polygon(area_word) + + def test_mixed_area_point_scatters(self): + """Points mixed with area words scatter correctly through the + group-by-order dispatch: each point element yields its containing + order-29 cell geometry at its input position.""" + words = np.concatenate([_area(6), _point()]) + bboxes = mort2bbox(words) + assert len(bboxes) == words.size + # The order-6 area words keep their uniform bboxes... + assert bboxes[:2] == mort2bbox(_area(6)) + # ...and the trailing points scatter to their containing-cell bboxes. + assert bboxes[2:] == mort2bbox(_area(29)) + + polygons = mort2polygon(words) + assert len(polygons) == words.size + assert polygons[:2] == mort2polygon(_area(6)) + assert polygons[2:] == mort2polygon(_area(29)) + + def test_order29_area_words_still_work(self): + """Order-29 AREA words are genuine cells (spec §4) — points now group + with them and share their geometry.""" + assert len(mort2bbox(_area(29))) == 2 + assert len(mort2polygon(_area(29))) == 2 + + +class TestResolutionCrossCheck: + def test_orders_of_indexes_order2res(self): + """The per-element order indexes the resolution ladder exactly as the + encoding order does (spec §3).""" + words = np.concatenate([_area(o) for o in _MIXED_ORDERS]) + expected = np.repeat([order2res(o) for o in _MIXED_ORDERS], len(_LATS)) + np.testing.assert_allclose(order2res(orders_of(words)), expected, rtol=1e-12) diff --git a/mortie/tools.py b/mortie/tools.py index 3c80971..c409157 100644 --- a/mortie/tools.py +++ b/mortie/tools.py @@ -194,24 +194,105 @@ def geo2mort(lats, lons, order=None, points=None): return np.ascontiguousarray(np.atleast_1d(result), dtype=np.uint64) +def orders_of(morton): + """Per-element HEALPix order of packed morton words. + + Vectorized numpy decode of the 6-bit suffix (bits 5-0) per the spec page's + suffix table (``docs/specification.md`` §1): + + * suffix ``0..=27`` — variable-length area element; the order *is* the + suffix value (``0`` = base-cell-only). + * suffix ``28..=47`` — order-28/29 area cells in parent-first preorder + ``suffix = 28 + t28*5 + (t29 present ? t29 + 1 : 0)``: each ``t28`` owns + a 5-block (the order-28 parent, then its four order-29 children), so + ``(suffix - 28) % 5 == 0`` is order 28 and everything else is order 29. + * suffix ``48..=63`` — order-29 **point** (max-encoded, no area claim — + spec §4); points are order 29 by definition. + + Pure bit arithmetic — words are not validated (the empty sentinel ``0`` + decodes as order 0; use :func:`validate_morton` to reject malformed + words). This is the per-element, mixed-order-native counterpart of + :func:`infer_order_from_morton`. + + Parameters + ---------- + morton : int or array-like + Packed morton word(s) (``uint64``). + + Returns + ------- + ndarray + ``uint8`` order per element, 0-29 (scalar in -> length-1 ndarray, + matching :func:`geo2mort`). + """ + m = np.atleast_1d(np.asarray(morton, dtype=np.uint64)) + suffix = (m & np.uint64(0x3F)).astype(np.uint8) + # 0..=27: order == suffix. 28..=47: order-28 on the 5-block parent slots, + # order 29 otherwise. 48..=63: order-29 point. + orders = suffix.copy() + band = (suffix >= 28) & (suffix <= 47) + orders[band] = np.where((suffix[band] - 28) % 5 == 0, 28, 29) + orders[suffix >= 48] = 29 + return orders + + +def is_point(morton): + """Per-element point-kind predicate for packed morton words. + + Kind is carried by the encoding itself (spec §4): suffix ``0..=47`` + decodes as an **area** word, suffix ``48..=63`` as an order-29 **point** + (a location with no area claim — ``docs/specification.md`` §1 suffix + table). Pure bit arithmetic; words are not validated (see + :func:`validate_morton`). + + Parameters + ---------- + morton : int or array-like + Packed morton word(s) (``uint64``). + + Returns + ------- + ndarray + ``bool`` per element, True for point words (scalar in -> length-1 + ndarray, matching :func:`geo2mort`). + """ + m = np.atleast_1d(np.asarray(morton, dtype=np.uint64)) + return (m & np.uint64(0x3F)) >= np.uint64(48) + + def infer_order_from_morton(morton): - """Infer the HEALPix order of a packed morton word. + """Infer the single HEALPix order of packed morton word(s). Decodes through the packed-u64 kernel (issue #48): the order is carried in - the word's suffix, not in any decimal-digit count. + the word's suffix, not in any decimal-digit count. The return is one + scalar order, so array input must be uniform-order; mixed-order input + raises, naming the distinct orders (issue #116 — previously the first + element's order was returned silently). For per-element orders of a mixed + array use :func:`orders_of`. Parameters ---------- - morton : int - Packed morton word. + morton : int or array-like + Packed morton word(s), all at one order. Returns ------- int The HEALPix order. + + Raises + ------ + ValueError + If the words are at mixed orders. """ m = np.atleast_1d(np.asarray(morton, dtype=np.uint64)) _, depths = _rust_mort2nested(np.ascontiguousarray(m)) + distinct = np.unique(depths) + if distinct.size > 1: + raise ValueError( + f"Mixed orders in morton array: {[int(d) for d in distinct]}; " + "use orders_of for per-element orders" + ) return int(depths[0]) @@ -281,10 +362,15 @@ def mort2norm(morton): return empty, empty.copy(), 0 # The packed-u64 kernel decodes each word to (nested, depth); the depth is - # the HEALPix order (no decimal-digit scan). Reject mixed orders. + # the HEALPix order (no decimal-digit scan). Reject mixed orders: the + # return contract is a single scalar order (the geo kernels above this + # dispatch group-by-order and never hit this — issue #116). nested, depths = _rust_mort2nested(np.ascontiguousarray(morton)) if np.any(depths != depths[0]): - raise ValueError(f"Mixed orders in morton array: {set(int(d) for d in depths)}") + raise ValueError( + f"Mixed orders in morton array: {sorted(set(int(d) for d in depths))}; " + "use orders_of for per-element orders" + ) order = int(depths[0]) # nested ids are HEALPix cell ids (<< 2^58 for order <= 29), so int64 is safe @@ -362,10 +448,16 @@ def mort2geo(morton): This is the inverse of geo2mort, returning the center coordinates of the HEALPix cell identified by the morton index. + Mixed-order arrays are supported (issue #116): elements are grouped by + order (:func:`orders_of`), each group runs the uniform kernel, and the + results scatter back to input positions. Point words (spec §4) are order + 29 by definition and group with order 29 — a point's location is exactly + what mort2geo returns. + Parameters ---------- morton : int or array-like - Morton index + Morton index (mixed orders allowed) Returns ------- @@ -377,6 +469,19 @@ def mort2geo(morton): # Handle scalar vs array input to match geo2mort behavior input_is_scalar = np.isscalar(morton) + # Group-by-order dispatch for mixed-order input (issue #116). + if not input_is_scalar: + words = np.atleast_1d(np.asarray(morton, dtype=np.uint64)) + orders = orders_of(words) + unique_orders = np.unique(orders) + if unique_orders.size > 1: + lat = np.empty(words.size, dtype=np.float64) + lon = np.empty(words.size, dtype=np.float64) + for order in unique_orders: + mask = orders == order + lat[mask], lon[mask] = mort2geo(words[mask]) + return lat, lon + # Decode morton to normalized address and parent normed, parent, order = mort2norm(morton) @@ -399,10 +504,18 @@ def mort2bbox(morton): normalized to use consistent representation based on hemisphere voting, preventing bbox misinterpretation as spanning the entire globe. + Mixed-order arrays are supported (issue #116): elements are grouped by + order (:func:`orders_of`), each group runs the uniform kernel, and the + results scatter back to input positions. Point words (spec §4) are order + 29 by definition and group with order 29 — a point yields the bounding box + of its containing order-29 cell (the cell that contains the point), which + is exactly the bbox of the order-29 **area** word at the same location. A + group of points therefore covers a well-defined area, element by element. + Parameters ---------- morton : int or array-like - Morton index + Morton index (mixed orders allowed) Returns ------- @@ -413,6 +526,22 @@ def mort2bbox(morton): morton = np.atleast_1d(morton) is_scalar = len(morton) == 1 + words = np.asarray(morton, dtype=np.uint64) + # Group-by-order dispatch for mixed-order input (issue #116). + orders = orders_of(words) + unique_orders = np.unique(orders) + if unique_orders.size > 1: + bboxes = [None] * words.size + for order in unique_orders: + (idx,) = np.nonzero(orders == order) + group = mort2bbox(words[idx]) + if idx.size == 1: + bboxes[idx[0]] = group # length-1 call returns the bare dict + else: + for i, bbox in zip(idx, group): + bboxes[i] = bbox + return bboxes + # First get the pixel center normed, parent, order = mort2norm(morton) uniq = norm2uniq(normed, parent, order) @@ -577,10 +706,34 @@ def mort2polygon(morton, step=1): normalized to use consistent longitude representation (-180 or +180) based on which hemisphere contains the majority of vertices. This prevents spatial libraries from misinterpreting touching polygons as crossing polygons. + + Mixed-order arrays are supported (issue #116): elements are grouped by + order (:func:`orders_of`), each group runs the uniform kernel, and the + results scatter back to input positions (rings are 4*step+1 vertices at + every order, so mixed orders do not change the output shape). Point words + (spec §4) are order 29 by definition and group with order 29 — a point + yields the polygon ring of its containing order-29 cell, exactly the ring + of the order-29 **area** word at the same location. """ morton = np.atleast_1d(morton) is_scalar = len(morton) == 1 + words = np.asarray(morton, dtype=np.uint64) + # Group-by-order dispatch for mixed-order input (issue #116). + orders = orders_of(words) + unique_orders = np.unique(orders) + if unique_orders.size > 1: + polygons = [None] * words.size + for order in unique_orders: + (idx,) = np.nonzero(orders == order) + group = mort2polygon(words[idx], step=step) + if idx.size == 1: + polygons[idx[0]] = group # length-1 call returns the bare ring + else: + for i, polygon in zip(idx, group): + polygons[i] = polygon + return polygons + # Get pixel information normed, parent, order = mort2norm(morton) uniq = norm2uniq(normed, parent, order)