matrix-fma-and-fancy-indexing (#41)

A `Matrix` capability release. The dense matrix type gains NumPy-style fancy
indexing, comparison masks and lexicographic comparison operators, a
``where`` selector, a single-rounding ``fma`` with vector broadcasting, and
``sqrt``. Two older method names move to their NumPy spellings: ``select``
becomes ``take`` and ``clip`` adopts ``min`` / ``max`` keyword bounds.

**New Features**

- **Fancy indexing** — ``m[[r0, r1]]`` / ``m[[r0, r1], :]`` gather rows and
  ``m[:, [c0, c1]]`` gathers columns, returning a :class:`Matrix`; the
  matching assignment forms scatter into rows or columns with last-write-wins
  duplicates and all-or-nothing validation. New :meth:`Matrix.take` and
  :meth:`Matrix.put` expose the same gather/scatter as methods, ``put`` with
  an ``accumulate=True`` mode that folds duplicate indices.
- **Comparison masks** — :meth:`Matrix.less`, ``less_equal``, ``greater``,
  ``greater_equal``, ``equal``, and ``not_equal`` return a ``1.0`` / ``0.0``
  mask matrix, accepting a same-shape matrix, a scalar (including ``bool``), a
  ``1x1`` matrix, a broadcasting row/column vector, or a list/tuple of
  numbers. Distinct from the comparison operators, which return a single
  bool.
- **Lexicographic comparison operators** — ``<`` ``<=`` ``>`` ``>=`` ``==``
  ``!=`` compare element by element in row-major order and return a single
  :class:`bool`. ``==`` / ``!=`` are total: a shape mismatch or an
  uncoercible list/tuple yields ``False`` / ``True`` rather than raising, so
  ``matrix in some_list`` works. A ``NaN`` never decides the comparison, so an
  all-``NaN`` matrix compares ``==`` equal to itself. Defining value equality
  makes :class:`Matrix` unhashable.
- **`Matrix.where(mask, a, b)`** — a NumPy-style selector taking *a* where the
  mask is non-zero (``NaN`` counts as non-zero) and *b* elsewhere; *a* and *b*
  may each be a scalar, a same-shape matrix, or a list/tuple of numbers.
- **`Matrix.fma(b, c)`** — fused multiply-add computing single-rounding
  ``self * b + c``; *b* and *c* may be a same-shape matrix, a ``1x1`` matrix,
  a scalar, or a row / column vector that broadcasts against ``self``. The
  contraction kernel is preserved so hardware FMA still applies. Use it as an
  accuracy primitive — compare results with :meth:`Matrix.allclose`, never
  ``==``.
- **`Matrix.sqrt()`** — element-wise square root (negative inputs map to
  ``NaN``), with an ``in_place=True`` form.

**Breaking Changes**

- **`Matrix.select` renamed to `Matrix.take`.** The gather method is now
  spelled :meth:`Matrix.take` to match NumPy and pair with the new
  :meth:`Matrix.put`. Replace ``m.select(indices, axis)`` with
  ``m.take(indices, axis)``; the signature and semantics are otherwise
  unchanged.
- **`Matrix.clip` bounds are now `min` / `max` keywords.** The signature
  changes from ``clip(min_or_maxval, maxval=None)`` to
  ``clip(min=None, max=None)``, matching :func:`numpy.clip`. Either bound may
  be omitted to leave that side unbounded: ``m.clip(min=0.0)`` clamps only
  below, ``m.clip(max=255.0)`` only above.

**Documentation**

- Expanded the :doc:`api` matrix surface for the new indexing, masking,
  comparison, ``where``, ``fma``, and ``sqrt`` methods via the
  ``__init__.pyi`` stub docstrings, including the totality, ``NaN``, and
  broadcasting rules.

**Tests**

- Extensive `test_matrix.py` additions covering fancy-index gather/scatter,
  ``take`` / ``put`` (including accumulate and all-or-nothing validation),
  the comparison masks, lexicographic operators (totality, ``NaN``,
  reflected-scalar, and list/tuple/bool coercion edge cases), ``where``
  selection and value propagation, and ``fma`` row/column broadcasting.

**Internal**

- New `bench_fma` and `bench_take` micro-benchmarks in
  `scripts/bench_matrix.py`. The `examples/boids.py` demo migrates from
  ``select`` to ``take``.

Signed-off-by: Matthew A Johnson <matthew@matthewajohnson.org>
Signed-off-by: Matthew A Johnson <matjoh@microsoft.com>

Author: Matthew A Johnson

.flake8

@@ -17,7 +17,8 @@ extend-exclude = scripts/_vendored_warehouse_wheel.py
 per-file-ignores =
     test/*:D102,D103,D403
     # .pyi stubs intentionally:
+    #   A002: arguments shadowing Python builtins (Matrix.clip min/max bounds)
     #   A003: methods shadowing Python builtins (Matrix.sum/min/max/abs/round)
     #   D402: Sphinx-style docstrings start with the function signature
     #   N802: NumPy-style ``Matrix.T`` transpose property is upper-case
-    *.pyi:A003,D402,N802
+    *.pyi:A003,D402,N802,A002

CHANGELOG.md

@@ -1,3 +1,77 @@
+## 2026-06-17 - Version 0.12.0
+A `Matrix` capability release. The dense matrix type gains NumPy-style fancy
+indexing, comparison masks and lexicographic comparison operators, a
+``where`` selector, a single-rounding ``fma`` with vector broadcasting, and
+``sqrt``. Two older method names move to their NumPy spellings: ``select``
+becomes ``take`` and ``clip`` adopts ``min`` / ``max`` keyword bounds.
+
+**New Features**
+
+- **Fancy indexing** — ``m[[r0, r1]]`` / ``m[[r0, r1], :]`` gather rows and
+  ``m[:, [c0, c1]]`` gathers columns, returning a :class:`Matrix`; the
+  matching assignment forms scatter into rows or columns with last-write-wins
+  duplicates and all-or-nothing validation. New :meth:`Matrix.take` and
+  :meth:`Matrix.put` expose the same gather/scatter as methods, ``put`` with
+  an ``accumulate=True`` mode that folds duplicate indices.
+- **Comparison masks** — :meth:`Matrix.less`, ``less_equal``, ``greater``,
+  ``greater_equal``, ``equal``, and ``not_equal`` return a ``1.0`` / ``0.0``
+  mask matrix, accepting a same-shape matrix, a scalar (including ``bool``), a
+  ``1x1`` matrix, a broadcasting row/column vector, or a list/tuple of
+  numbers. Distinct from the comparison operators, which return a single
+  bool.
+- **Lexicographic comparison operators** — ``<`` ``<=`` ``>`` ``>=`` ``==``
+  ``!=`` compare element by element in row-major order and return a single
+  :class:`bool`. ``==`` / ``!=`` are total: a shape mismatch or an
+  uncoercible list/tuple yields ``False`` / ``True`` rather than raising, so
+  ``matrix in some_list`` works. A ``NaN`` never decides the comparison, so an
+  all-``NaN`` matrix compares ``==`` equal to itself. Defining value equality
+  makes :class:`Matrix` unhashable.
+- **`Matrix.where(mask, a, b)`** — a NumPy-style selector taking *a* where the
+  mask is non-zero (``NaN`` counts as non-zero) and *b* elsewhere; *a* and *b*
+  may each be a scalar, a same-shape matrix, or a list/tuple of numbers.
+- **`Matrix.fma(b, c)`** — fused multiply-add computing single-rounding
+  ``self * b + c``; *b* and *c* may be a same-shape matrix, a ``1x1`` matrix,
+  a scalar, or a row / column vector that broadcasts against ``self``. The
+  contraction kernel is preserved so hardware FMA still applies. Use it as an
+  accuracy primitive — compare results with :meth:`Matrix.allclose`, never
+  ``==``.
+- **`Matrix.sqrt()`** — element-wise square root (negative inputs map to
+  ``NaN``), with an ``in_place=True`` form.
+
+**Breaking Changes**
+
+- **`Matrix.select` renamed to `Matrix.take`.** The gather method is now
+  spelled :meth:`Matrix.take` to match NumPy and pair with the new
+  :meth:`Matrix.put`. Replace ``m.select(indices, axis)`` with
+  ``m.take(indices, axis)``; the signature and semantics are otherwise
+  unchanged.
+- **`Matrix.clip` bounds are now `min` / `max` keywords.** The signature
+  changes from ``clip(min_or_maxval, maxval=None)`` to
+  ``clip(min=None, max=None)``, matching :func:`numpy.clip`. Either bound may
+  be omitted to leave that side unbounded: ``m.clip(min=0.0)`` clamps only
+  below, ``m.clip(max=255.0)`` only above.
+
+**Documentation**
+
+- Expanded the :doc:`api` matrix surface for the new indexing, masking,
+  comparison, ``where``, ``fma``, and ``sqrt`` methods via the
+  ``__init__.pyi`` stub docstrings, including the totality, ``NaN``, and
+  broadcasting rules.
+
+**Tests**
+
+- Extensive `test_matrix.py` additions covering fancy-index gather/scatter,
+  ``take`` / ``put`` (including accumulate and all-or-nothing validation),
+  the comparison masks, lexicographic operators (totality, ``NaN``,
+  reflected-scalar, and list/tuple/bool coercion edge cases), ``where``
+  selection and value propagation, and ``fma`` row/column broadcasting.
+
+**Internal**
+
+- New `bench_fma` and `bench_take` micro-benchmarks in
+  `scripts/bench_matrix.py`. The `examples/boids.py` demo migrates from
+  ``select`` to ``take``.
+
 ## 2026-06-14 - Version 0.11.0
 A behavior-dispatch release. `@when` becomes a **runtime decorator backed by
 a content-addressed marshalled-code registry** instead of a transpile-time

CITATION.cff

@@ -5,6 +5,6 @@ authors:
   given-names: "Matthew Alastair"
   orcid: "https://orcid.org/0000-0002-1019-8036"
 title: "bocpy"
-version: 0.11.0
-date-released: 2026-06-14
+version: 0.12.0
+date-released: 2026-06-17
 url: "https://github.com/microsoft/bocpy"

examples/boids.py

@@ -313,8 +313,8 @@ def build_cell_data(self, positions: Matrix, velocities: Matrix, row: int, colum
 
         assert len(boids) > 0, "Invalid grid cell"
 
-        positions = positions.select(boids)
-        velocities = velocities.select(boids)
+        positions = positions.take(boids)
+        velocities = velocities.take(boids)
         return CellData(Cell(row, column), tuple(boids), Cown(positions), Cown(velocities))
 
     def step(self, width: int, height: int):

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bocpy"
-version = "0.11.0"
+version = "0.12.0"
 authors = [
     {name = "bocpy Team", email="bocpy@microsoft.com"}
 ]

scripts/bench_matrix.py

@@ -298,6 +298,48 @@ def bench_vecdot(results: list[dict[str, float]]) -> None:
     _print_row(results[-1])
 
 
+def bench_fma(results: list[dict[str, float]]) -> None:
+    """Bench fma(b, c) vs a*b + c.
+
+    On the shipped x86-64 SSE2 build libc ``fma()`` is an unvectorisable
+    libcall, so ``a*b + c`` is expected to win on throughput there -- fma's
+    headline is accuracy (single rounding), not speed. A hardware-FMA host
+    (arm64, or x86-64 built with ``-mfma``) is where any throughput claim
+    must be measured.
+    """
+    _print_section("fma(b, c) vs a*b + c (accuracy primitive; libcall on SSE2)")
+
+    for rows, cols in [(1000, 3), (256, 256)]:
+        a = _make_matrix(rows, cols, seed=40)
+        b = _make_matrix(rows, cols, seed=41)
+        c = _make_matrix(rows, cols, seed=42)
+        label = f"{rows}x{cols}"
+        results.append(measure(
+            f"fma(b, c)              matrix b,c {label}",
+            lambda a=a, b=b, c=c: a.fma(b, c),
+        ))
+        _print_row(results[-1])
+        results.append(measure(
+            f"a*b + c                matrix b,c {label}",
+            lambda a=a, b=b, c=c: a * b + c,
+        ))
+        _print_row(results[-1])
+
+    a = _make_matrix(256, 256, seed=40)
+    b = _make_matrix(256, 256, seed=41)
+    c = _make_matrix(256, 256, seed=42)
+    results.append(measure(
+        "fma(b, c, in_place)    matrix b,c 256x256",
+        lambda a=a, b=b, c=c: a.fma(b, c, in_place=True),
+    ))
+    _print_row(results[-1])
+    results.append(measure(
+        "fma(2.0, 1.0)          scalar b,c 256x256",
+        lambda a=a: a.fma(2.0, 1.0),
+    ))
+    _print_row(results[-1])
+
+
 def bench_cross(results: list[dict[str, float]]) -> None:
     """Bench cross across scalar and batch shapes."""
     _print_section("cross at 1x3, 1000x2 (2D row batch), 1000x3 (3D row batch), 3x1000 (3D col batch)")
@@ -624,22 +666,79 @@ def bench_transpose(results: list[dict[str, float]]) -> None:
     _print_row(results[-1])
 
 
-def bench_select(results: list[dict[str, float]]) -> None:
-    """Bench row and column gather across small and large index lists."""
-    _print_section("select (row / column gather)")
+def bench_take(results: list[dict[str, float]]) -> None:
+    """Bench row and column gather across wide, narrow, and column cases.
+
+    The shared ``gather_axis`` helper copies each selected row with
+    ``memcpy`` (source and destination rows are both contiguous); the
+    wide row case (1000x100) is where that path dominates. The narrow-C
+    row case (1000x3) isolates the per-index unbox + bounds-check cost,
+    a larger fraction of the work when each copied row is short. Column
+    gather stays element-wise (inherently strided) and confirms no
+    regression there.
+    """
+    _print_section("take (row / column gather)")
 
     shape = (1000, 100)
     m = _make_matrix(*shape, seed=108)
     row_idx = list(range(0, 1000, 10))
     col_idx = list(range(0, 100, 2))
     results.append(measure(
-        f"select rows (100/1000)  shape={shape}",
-        lambda m=m, i=row_idx: m.select(i, 0),
+        f"take rows (100/1000)    shape={shape}",
+        lambda m=m, i=row_idx: m.take(i, 0),
+    ))
+    _print_row(results[-1])
+    results.append(measure(
+        f"take cols (50/100)      shape={shape}",
+        lambda m=m, i=col_idx: m.take(i, 1),
+    ))
+    _print_row(results[-1])
+
+    narrow = _make_matrix(1000, 3, seed=51)
+    results.append(measure(
+        "take rows (narrow C)    shape=(1000, 3) -> 100 rows",
+        lambda m=narrow, i=row_idx: m.take(i, 0),
+    ))
+    _print_row(results[-1])
+
+
+def bench_scatter(results: list[dict[str, float]]) -> None:
+    """Bench list-key scatter assignment (rows vs columns, fill vs matrix).
+
+    Row scatter writes each selected destination row with a contiguous
+    ``memcpy`` (the headline write-side path); column scatter is strided.
+    The augmented ``m[[...]] += v`` case measures the gather -> in-place op
+    -> scatter triple pass against the plain single scatter; if that ratio
+    is large *and* the form is hot, that is the signal to revisit the
+    deferred fused single-pass kernel.
+    """
+    _print_section("scatter (row / column assignment)")
+
+    shape = (1000, 100)
+    row_idx = list(range(0, 1000, 10))
+    col_idx = list(range(0, 100, 2))
+    rows_rhs = _make_matrix(len(row_idx), 100, seed=110)
+    cols_rhs = _make_matrix(1000, len(col_idx), seed=111)
+
+    m = _make_matrix(*shape, seed=112)
+    results.append(measure(
+        f"scatter rows = scalar    shape={shape}",
+        lambda m=m, i=row_idx: m.__setitem__(i, 0.5),
+    ))
+    _print_row(results[-1])
+    results.append(measure(
+        f"scatter rows = matrix    shape={shape} (memcpy)",
+        lambda m=m, i=row_idx, v=rows_rhs: m.__setitem__(i, v),
+    ))
+    _print_row(results[-1])
+    results.append(measure(
+        f"scatter cols = matrix    shape={shape} (strided)",
+        lambda m=m, i=col_idx, v=cols_rhs: m.__setitem__((slice(None), i), v),
     ))
     _print_row(results[-1])
     results.append(measure(
-        f"select cols (50/100)    shape={shape}",
-        lambda m=m, i=col_idx: m.select(i, 1),
+        f"scatter rows += scalar   shape={shape} (triple pass)",
+        lambda m=m, i=row_idx: m.__setitem__(i, m[i] + 0.5),
     ))
     _print_row(results[-1])
 
@@ -827,7 +926,8 @@ def main(argv: list[str] | None = None) -> int:
     bench_binary_arithmetic(results)
     bench_matmul(results)
     bench_transpose(results)
-    bench_select(results)
+    bench_take(results)
+    bench_scatter(results)
     bench_copy_clip_allclose(results)
     bench_construction(results)
     bench_factories(results)
@@ -836,6 +936,7 @@ def main(argv: list[str] | None = None) -> int:
     bench_magnitude_squared(results)
     bench_negate(results)
     bench_vecdot(results)
+    bench_fma(results)
     bench_cross(results)
     bench_normalize(results)
     bench_perpendicular(results)

sphinx/source/api.rst

@@ -125,7 +125,7 @@ Math
 .. autoclass:: Matrix
     :members:
     :undoc-members:
-    :special-members: __init__
+    :special-members: __init__, __eq__, __ne__, __lt__, __le__, __gt__, __ge__, __getitem__, __setitem__
 
 
 Messaging

sphinx/source/conf.py

@@ -14,7 +14,7 @@
 project = 'bocpy'
 copyright = '2026, Microsoft'
 author = 'Microsoft'
-release = '0.11.0'
+release = '0.12.0'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration