From 29e8e722be6f29ef350790d03d5ea73e3bb9d358 Mon Sep 17 00:00:00 2001 From: Matthew A Johnson Date: Sat, 4 Jul 2026 15:55:01 +0100 Subject: [PATCH] matrix-improvements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A `Matrix` gather/reduce release. The dense matrix type gains top-k selection, along-axis gather/scatter, interleaved repetition, a lazy element iterator, masked reductions, and trigonometric/sign ufuncs, while the random-number factories move to an independent per-interpreter generator. A crash in the noticeboard mutator thread when reconstructing Matrix-valued entries is fixed, a family of allocation-overflow and cross-interpreter acquisition guards harden the C core, and passing an empty cown-list group to ``@when`` no longer crashes or drops a behavior parameter. It also adds a zero-copy `MatrixView` — read/write row and column windows that alias a matrix's storage, with ``.x``/``.y``/``.z``/``.w`` accessors and per-row/column iterators — plus value binning via `Matrix.bin` and `Matrix.digitize`. One breaking change: the ``in_place`` flag on the unary element-wise methods is now keyword-only. **New Features** - **`Matrix.topk(k, axis=None, largest=True, where=None, as_matrix=False)`** — the *k* extreme elements per reduction group, in sorted order, returning a ``(values, indices)`` tuple. NumPy tie-break (first occurrence wins, NaN sorts last); a masked group shorter than *k* pads with ``NaN`` values and ``-1`` indices. Indices come back as Python lists by default (a flat ``list[int]`` for ``axis=None``, a list of per-group lists for an axis) so they feed straight into fancy indexing; ``as_matrix=True`` returns a same-shape index :class:`Matrix` instead. - **`Matrix.take_along_axis` / `Matrix.put_along_axis`** — the ``np.take_along_axis`` / ``np.put_along_axis`` gather and scatter, one index per row (``axis=1``) or per column (``axis=0``). ``take_along_axis`` accepts a keyword-only ``out=`` target; ``put_along_axis`` has an ``accumulate=True`` mode. Pairs directly with :meth:`Matrix.argmin` / :meth:`Matrix.argmax`. - **`Matrix.repeat_interleave(repeats, axis=None)`** — the ``np.repeat`` / ``torch.repeat_interleave`` interleaved (not tiled) copy of each element, row, or column. - **`Matrix.values()`** — a lazy row-major iterator over every element, holding a strong reference to its source and re-checking cown acquisition on each step. - **`Matrix.full(size, value)`** — a constant-filled constructor alongside :meth:`Matrix.zeros` / :meth:`Matrix.ones`. - **`Matrix.sign` / `Matrix.cos` / `Matrix.sin`** — element-wise ufuncs (``sign`` maps ±0.0 and NaN to ``0``), each with keyword-only ``in_place=`` / ``out=``. - **Masked reductions** — ``sum`` / ``mean`` / ``min`` / ``max`` / ``magnitude`` / ``magnitude_squared`` and ``argmin`` / ``argmax`` accept a ``where=`` same-shape mask matrix; only cells whose mask is non-zero (NaN counts as included) are considered. An all-excluded group publishes an op-specific sentinel: ``0`` for the additive ops, ``NaN`` for ``mean`` / ``min`` / ``max`` (matching NumPy's empty-slice semantics), and ``-1`` for ``argmin`` / ``argmax``. - **Axis-wise `argmin` / `argmax`** — ``axis=0`` / ``axis=1`` return the per-column / per-row extreme positions as a Python ``list[int]`` (directly usable in fancy indexing), or a :class:`Matrix` vector with ``as_matrix=True``. ``axis=None`` still returns a single ``int``. - **`MatrixView` — zero-copy row/column windows** — :meth:`Matrix.row_view` and :meth:`Matrix.column_view` (and the ``m.view[...]`` slice-grammar accessor) return a read/write 1-D view that *aliases* the matrix's storage: reading or writing an element (``rv[i]``, ``rv[i] = x``, ``rv[:] = ...``, or the ``.x`` / ``.y`` / ``.z`` / ``.w`` accessors) touches the underlying buffer directly. A view pins its whole backing buffer, materialises to a fresh contiguous :class:`Matrix` for arithmetic, and is deliberately not shippable across interpreters (call :meth:`MatrixView.copy` to detach an owned copy). :meth:`MatrixView.values` lazily iterates the (possibly strided) elements. - **`Matrix.row_views()` / `Matrix.column_views()`** — iterators that yield a fresh :class:`MatrixView` over each row / column in turn, a cleaner spelling of ``for i in range(m.rows): m.row_view(i)``. - **`MatrixView.x` / `.y` / `.z` / `.w`** — get / set the first four elements of a view, honouring its offset and stride so a column or reversed view reads and writes the correct cells (mirroring the :class:`Matrix` accessors). - **`Matrix.bin(bins, *, bounds=None, right=False, in_place=False, out=None)`** — map every element to its equal-width bin index over a range. ``bounds`` is an optional ``(min, max)`` tuple of finite numbers that fixes the range and skips the ``O(size)`` min/max scan; when omitted the range is the matrix's own min/max (ignoring NaN). Indices clamp to ``[0, bins-1]`` (so a value at the max lands in the last bin) and a ``NaN`` element maps to ``-1``. Clip-family output routing: ``in_place`` / keyword-only ``out`` are mutually exclusive. - **`Matrix.digitize(edges, *, right=False, in_place=False, out=None)`** — the ``numpy.digitize`` map: each element to the index of the interval bounded by a strictly increasing *edges* vector (a :class:`Matrix` vector or any coercible sequence). ``right`` selects the boundary side; a ``NaN`` element maps to ``-1``. Same output routing as ``bin``. **Improvements** - **Allocation-free `out=` on gathers** — :meth:`Matrix.take` and :meth:`Matrix.take_along_axis` accept a keyword-only ``out=`` :class:`Matrix` to write into, rejected if it aliases the source (a reordering gather would read cells it had already overwritten). - **`Matrix.clip` gains `in_place=` / `out=`** — mutate in place or write into a caller-supplied matrix, mutually exclusive. - **Slice subscripts always return a `Matrix`** — a slice anywhere in the key (even a length-1 one, e.g. ``m[0:1, 0:1]``) keeps the result a :class:`Matrix`; only an all-integer key collapses to a ``float`` scalar. - **Reproducible, race-free Matrix RNG** — :meth:`Matrix.normal` / :meth:`Matrix.uniform` / :meth:`Matrix.seed` now draw from a per-interpreter splitmix64 generator held in the module state instead of the process-global C ``rand()`` / ``srand()``. Each worker sub-interpreter gets an independent stream, so ``seed()`` is reproducible within an interpreter and concurrent draws are no longer a data race on non-glibc or free-threaded builds. **Bug Fixes** - **Noticeboard Matrix-entry segfault** — the noticeboard mutator thread (a second OS thread in the primary interpreter that never ran the ``_math`` module init) crashed dereferencing a NULL cached module-state pointer while reconstructing a Matrix-valued noticeboard entry during a ``notice_update`` snapshot. The main interpreter's module state is now published at module exec (``MAIN_MATH_STATE``) and resolved on that thread; a regression test exercises both the XIData-reconstruction and pickle paths on the mutator thread. - **`repeat_interleave` integer-overflow → out-of-bounds heap write** — the overflow guard bounded only the repeated dimension, so a large ``repeats`` with a small repeated axis could wrap the ``rows * columns * repeats`` product and under-allocate. The guard now bounds the total element count, and ``impl_new`` re-checks the ``rows * columns`` product as a backstop for every constructor; both raise ``OverflowError`` instead of corrupting the heap. - **`Matrix.allclose` missing acquisition guard** — the only data-touching method that dereferenced both operands' buffers without an ``impl_check_acquired`` check now raises ``RuntimeError`` on a released or cross-interpreter operand instead of reading memory it does not own. - **NULL-deref under memory pressure** — the ``Matrix`` subscript / item paths now check the ``impl_new`` result before use. - **Empty cown-list `@when` group crashed or dropped a parameter** — passing an empty list to :func:`when` (e.g. ``@when(a, [])`` or ``@when([], b)``) dropped the corresponding behavior parameter: a trailing empty group raised ``TypeError`` (argument tuple too small) and a leading or middle empty group left a ``NULL`` argument slot that segfaulted the call. Empty groups now reconstruct as ``[]`` at their correct parameter slot, and the low-level ``BehaviorCapsule`` constructor validates the ``group_id`` / slot-count relationship so a malformed capsule can no longer drive an out-of-bounds argument-tuple write. **Breaking Changes** - **`in_place` is keyword-only on the unary element-wise methods** — ``ceil`` / ``floor`` / ``round`` / ``negate`` / ``abs`` / ``sqrt`` / ``sign`` / ``cos`` / ``sin`` no longer accept ``in_place`` positionally, matching :meth:`Matrix.clip`. Replace ``m.negate(True)`` with ``m.negate(in_place=True)``. **Documentation** - Expanded the :doc:`api` matrix surface via the ``__init__.pyi`` stub docstrings for the new methods and parameters, including ``@overload`` signatures for the polymorphic ``topk`` / ``argmin`` / ``argmax`` return types. The :class:`Matrix` class docstring now states the output-routing convention (which methods offer ``out=`` vs ``in_place=``) and the masked-reduction empty-group sentinel table in one place. - Documented the :class:`MatrixView` surface (construction, aliasing, not-shippable semantics, ``.x``/``.y``/``.z``/``.w`` accessors), the ``row_views`` / ``column_views`` iterators, and ``Matrix.bin`` / ``Matrix.digitize`` via the ``__init__.pyi`` stubs and the :doc:`api` page, and switched the boids example to zero-copy ``row_view`` / ``row_views`` for its per-boid coordinate reads. **Tests** - Added golden and fuzzed coverage for ``topk`` (list and matrix index forms), ``take_along_axis`` / ``put_along_axis``, ``repeat_interleave`` (including the overflow regression), masked reductions and arg-reductions, the ``values()`` iterator, ``sign`` / ``cos`` / ``sin``, slice-forces-matrix, and the ``allclose`` acquisition guard. Added ``TestFactoriesOnWorker`` to exercise the per-interpreter RNG from inside ``@when`` behaviors, and a ``TestNoticeboardMatrixEntry`` regression class for the mutator-thread fix. - Added construction, read/write, aliasing, arithmetic-materialisation, ownership-guard, not-shippable, strided-accessor, and fuzzed coverage for :class:`MatrixView` and the ``row_views`` / ``column_views`` iterators, and golden + fuzzed (vs a pure-Python oracle) coverage for ``bin`` (data and fixed ``bounds``, non-finite-bounds rejection, ``NaN`` → ``-1``) and ``digitize`` (``right`` sides, monotonic-edge validation, in-place edge aliasing). **Internal** - Reworked the RNG onto a per-interpreter splitmix64 state seeded at module exec, and stopped caching the published main-interpreter state into a cold thread's thread-local so a torn-down module surfaces a clean ``RuntimeError`` rather than a dangling pointer. - De-duplicated the three ``topk`` axis branches into a single group loop behind ``topk_gather_group`` / ``topk_write_group`` / ``topk_pack_result``, and extracted the shared gather output-allocation + alias-check block (``gather_output``) used by the whole-axis and along-axis gathers. Signed-off-by: Matthew A Johnson --- CHANGELOG.md | 191 + CITATION.cff | 4 +- examples/boids.py | 9 +- pyproject.toml | 2 +- sphinx/source/api.rst | 52 + sphinx/source/conf.py | 2 +- src/bocpy/__init__.py | 5 +- src/bocpy/__init__.pyi | 763 ++- src/bocpy/_core.c | 116 +- src/bocpy/_math.c | 6967 ++++++++++++++++------- src/bocpy/behaviors.py | 4 +- templates/c_abi_consumer/pyproject.toml | 4 +- test/test_boc.py | 161 +- test/test_matrix.py | 3398 ++++++++++- test/test_noticeboard.py | 67 +- test/test_scheduling_stress.py | 2 + 16 files changed, 9367 insertions(+), 2380 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4461a2c..dd963b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,194 @@ +## 2026-07-02 - Version 0.14.0 +A `Matrix` gather/reduce release. The dense matrix type gains top-k selection, +along-axis gather/scatter, interleaved repetition, a lazy element iterator, +masked reductions, and trigonometric/sign ufuncs, while the random-number +factories move to an independent per-interpreter generator. A crash in the +noticeboard mutator thread when reconstructing Matrix-valued entries is fixed, +a family of allocation-overflow and cross-interpreter acquisition guards +harden the C core, and passing an empty cown-list group to ``@when`` no longer +crashes or drops a behavior parameter. It also adds a zero-copy `MatrixView` — +read/write row and column windows that alias a matrix's storage, with +``.x``/``.y``/``.z``/``.w`` accessors and per-row/column iterators — plus value +binning via `Matrix.bin` and `Matrix.digitize`. One breaking change: the +``in_place`` flag on the unary element-wise methods is now keyword-only. + +**New Features** + +- **`Matrix.topk(k, axis=None, largest=True, where=None, as_matrix=False)`** — + the *k* extreme elements per reduction group, in sorted order, returning a + ``(values, indices)`` tuple. NumPy tie-break (first occurrence wins, NaN + sorts last); a masked group shorter than *k* pads with ``NaN`` values and + ``-1`` indices. Indices come back as Python lists by default (a flat + ``list[int]`` for ``axis=None``, a list of per-group lists for an axis) so + they feed straight into fancy indexing; ``as_matrix=True`` returns a + same-shape index :class:`Matrix` instead. +- **`Matrix.take_along_axis` / `Matrix.put_along_axis`** — the + ``np.take_along_axis`` / ``np.put_along_axis`` gather and scatter, one index + per row (``axis=1``) or per column (``axis=0``). ``take_along_axis`` accepts + a keyword-only ``out=`` target; ``put_along_axis`` has an ``accumulate=True`` + mode. Pairs directly with :meth:`Matrix.argmin` / :meth:`Matrix.argmax`. +- **`Matrix.repeat_interleave(repeats, axis=None)`** — the + ``np.repeat`` / ``torch.repeat_interleave`` interleaved (not tiled) copy of + each element, row, or column. +- **`Matrix.values()`** — a lazy row-major iterator over every element, + holding a strong reference to its source and re-checking cown acquisition on + each step. +- **`Matrix.full(size, value)`** — a constant-filled constructor alongside + :meth:`Matrix.zeros` / :meth:`Matrix.ones`. +- **`Matrix.sign` / `Matrix.cos` / `Matrix.sin`** — element-wise ufuncs + (``sign`` maps ±0.0 and NaN to ``0``), each with keyword-only + ``in_place=`` / ``out=``. +- **`Matrix.reciprocal`** — the ``np.reciprocal`` element-wise ``1 / x`` ufunc + (zeros yield ±``inf``), with keyword-only ``in_place=`` / ``out=``. +- **Masked reductions** — ``sum`` / ``mean`` / ``min`` / ``max`` / + ``magnitude`` / ``magnitude_squared`` and ``argmin`` / ``argmax`` accept a + ``where=`` same-shape mask matrix; only cells whose mask is non-zero (NaN + counts as included) are considered. An all-excluded group publishes an + op-specific sentinel: ``0`` for the additive ops, ``NaN`` for ``mean`` / + ``min`` / ``max`` (matching NumPy's empty-slice semantics), and ``-1`` for + ``argmin`` / ``argmax``. +- **Axis-wise `argmin` / `argmax`** — ``axis=0`` / ``axis=1`` return the + per-column / per-row extreme positions as a Python ``list[int]`` (directly + usable in fancy indexing), or a :class:`Matrix` vector with + ``as_matrix=True``. ``axis=None`` still returns a single ``int``. +- **`MatrixView` — zero-copy row/column windows** — :meth:`Matrix.row_view` + and :meth:`Matrix.column_view` (and the ``m.view[...]`` slice-grammar + accessor) return a read/write 1-D view that *aliases* the matrix's storage: + reading or writing an element (``rv[i]``, ``rv[i] = x``, ``rv[:] = ...``, or + the ``.x`` / ``.y`` / ``.z`` / ``.w`` accessors) touches the underlying + buffer directly. A view pins its whole backing buffer, materialises to a + fresh contiguous :class:`Matrix` for arithmetic, and is deliberately not + shippable across interpreters (call :meth:`MatrixView.copy` to detach an + owned copy). A view is directly iterable (``for x in rv``) and + :meth:`MatrixView.values` lazily iterates its (possibly strided) elements. + Pass ``read_only=True`` to :meth:`Matrix.row_view` / :meth:`Matrix.column_view` + for an immutable view whose writes (item / slice assignment and the + ``.x``/``.y``/``.z``/``.w`` setters) raise ``TypeError``; ``.T`` and + sub-slices inherit the flag, and the ``.read_only`` property reports it. +- **`Matrix.row_views()` / `Matrix.column_views()`** — iterators that yield a + fresh :class:`MatrixView` over each row / column in turn, a cleaner spelling + of ``for i in range(m.rows): m.row_view(i)``. Pass ``read_only=True`` for + immutable per-row / per-column iteration. +- **`MatrixView.x` / `.y` / `.z` / `.w`** — get / set the first four elements + of a view, honouring its offset and stride so a column or reversed view + reads and writes the correct cells (mirroring the :class:`Matrix` + accessors). +- **`Matrix.bin(bins, *, bounds=None, right=False, in_place=False, out=None)`** + — map every element to its equal-width bin index over a range. ``bounds`` is + an optional ``(min, max)`` tuple of finite numbers that fixes the range and + skips the ``O(size)`` min/max scan; when omitted the range is the matrix's + own min/max (ignoring NaN). Indices clamp to ``[0, bins-1]`` (so a value at + the max lands in the last bin) and a ``NaN`` element maps to ``-1``. + Clip-family output routing: ``in_place`` / keyword-only ``out`` are mutually + exclusive. +- **`Matrix.digitize(edges, *, right=False, in_place=False, out=None)`** — the + ``numpy.digitize`` map: each element to the index of the interval bounded by + a strictly increasing *edges* vector (a :class:`Matrix` vector or any + coercible sequence). ``right`` selects the boundary side; a ``NaN`` element + maps to ``-1``. Same output routing as ``bin``. + +**Improvements** + +- **Allocation-free `out=` on gathers** — :meth:`Matrix.take` and + :meth:`Matrix.take_along_axis` accept a keyword-only ``out=`` :class:`Matrix` + to write into, rejected if it aliases the source (a reordering gather would + read cells it had already overwritten). +- **`Matrix.clip` gains `in_place=` / `out=`** — mutate in place or write into + a caller-supplied matrix, mutually exclusive. +- **Slice subscripts always return a `Matrix`** — a slice anywhere in the key + (even a length-1 one, e.g. ``m[0:1, 0:1]``) keeps the result a + :class:`Matrix`; only an all-integer key collapses to a ``float`` scalar. +- **Reproducible, race-free Matrix RNG** — :meth:`Matrix.normal` / + :meth:`Matrix.uniform` / :meth:`Matrix.seed` now draw from a per-interpreter + splitmix64 generator held in the module state instead of the process-global + C ``rand()`` / ``srand()``. Each worker sub-interpreter gets an independent + stream, so ``seed()`` is reproducible within an interpreter and concurrent + draws are no longer a data race on non-glibc or free-threaded builds. + +**Bug Fixes** + +- **Noticeboard Matrix-entry segfault** — the noticeboard mutator thread (a + second OS thread in the primary interpreter that never ran the ``_math`` + module init) crashed dereferencing a NULL cached module-state pointer while + reconstructing a Matrix-valued noticeboard entry during a ``notice_update`` + snapshot. The main interpreter's module state is now published at module + exec (``MAIN_MATH_STATE``) and resolved on that thread; a regression test + exercises both the XIData-reconstruction and pickle paths on the mutator + thread. +- **`repeat_interleave` integer-overflow → out-of-bounds heap write** — the + overflow guard bounded only the repeated dimension, so a large ``repeats`` + with a small repeated axis could wrap the ``rows * columns * repeats`` + product and under-allocate. The guard now bounds the total element count, + and ``impl_new`` re-checks the ``rows * columns`` product as a backstop for + every constructor; both raise ``OverflowError`` instead of corrupting the + heap. +- **`Matrix.allclose` missing acquisition guard** — the only data-touching + method that dereferenced both operands' buffers without an + ``impl_check_acquired`` check now raises ``RuntimeError`` on a released or + cross-interpreter operand instead of reading memory it does not own. +- **NULL-deref under memory pressure** — the ``Matrix`` subscript / item paths + now check the ``impl_new`` result before use. +- **Empty cown-list `@when` group crashed or dropped a parameter** — passing an + empty list to :func:`when` (e.g. ``@when(a, [])`` or ``@when([], b)``) dropped + the corresponding behavior parameter: a trailing empty group raised + ``TypeError`` (argument tuple too small) and a leading or middle empty group + left a ``NULL`` argument slot that segfaulted the call. Empty groups now + reconstruct as ``[]`` at their correct parameter slot, and the low-level + ``BehaviorCapsule`` constructor validates the ``group_id`` / slot-count + relationship so a malformed capsule can no longer drive an out-of-bounds + argument-tuple write. + +**Breaking Changes** + +- **`in_place` is keyword-only on the unary element-wise methods** — + ``ceil`` / ``floor`` / ``round`` / ``negate`` / ``abs`` / ``sqrt`` / + ``sign`` / ``cos`` / ``sin`` no longer accept ``in_place`` positionally, + matching :meth:`Matrix.clip`. Replace ``m.negate(True)`` with + ``m.negate(in_place=True)``. + +**Documentation** + +- Expanded the :doc:`api` matrix surface via the ``__init__.pyi`` stub + docstrings for the new methods and parameters, including ``@overload`` + signatures for the polymorphic ``topk`` / ``argmin`` / ``argmax`` return + types. The :class:`Matrix` class docstring now states the output-routing + convention (which methods offer ``out=`` vs ``in_place=``) and the + masked-reduction empty-group sentinel table in one place. +- Documented the :class:`MatrixView` surface (construction, aliasing, + not-shippable semantics, ``.x``/``.y``/``.z``/``.w`` accessors), the + ``row_views`` / ``column_views`` iterators, and ``Matrix.bin`` / + ``Matrix.digitize`` via the ``__init__.pyi`` stubs and the :doc:`api` page, + and switched the boids example to zero-copy ``row_view`` / ``row_views`` for + its per-boid coordinate reads. + +**Tests** + +- Added golden and fuzzed coverage for ``topk`` (list and matrix index forms), + ``take_along_axis`` / ``put_along_axis``, ``repeat_interleave`` (including + the overflow regression), masked reductions and arg-reductions, the + ``values()`` iterator, ``sign`` / ``cos`` / ``sin``, slice-forces-matrix, + and the ``allclose`` acquisition guard. Added ``TestFactoriesOnWorker`` to + exercise the per-interpreter RNG from inside ``@when`` behaviors, and a + ``TestNoticeboardMatrixEntry`` regression class for the mutator-thread fix. +- Added construction, read/write, aliasing, arithmetic-materialisation, + ownership-guard, not-shippable, strided-accessor, and fuzzed coverage for + :class:`MatrixView` and the ``row_views`` / ``column_views`` iterators, and + golden + fuzzed (vs a pure-Python oracle) coverage for ``bin`` (data and + fixed ``bounds``, non-finite-bounds rejection, ``NaN`` → ``-1``) and + ``digitize`` (``right`` sides, monotonic-edge validation, in-place edge + aliasing). + +**Internal** + +- Reworked the RNG onto a per-interpreter splitmix64 state seeded at module + exec, and stopped caching the published main-interpreter state into a + cold thread's thread-local so a torn-down module surfaces a clean + ``RuntimeError`` rather than a dangling pointer. +- De-duplicated the three ``topk`` axis branches into a single group loop + behind ``topk_gather_group`` / ``topk_write_group`` / ``topk_pack_result``, + and extracted the shared gather output-allocation + alias-check block + (``gather_output``) used by the whole-axis and along-axis gathers. + ## 2026-06-24 - Version 0.13.0 A `Matrix` ergonomics release. The dense matrix type gains named method forms of the arithmetic operators, a numpy-style ``out=`` target on every diff --git a/CITATION.cff b/CITATION.cff index 91ac82e..bf5b9d9 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -5,6 +5,6 @@ authors: given-names: "Matthew Alastair" orcid: "https://orcid.org/0000-0002-1019-8036" title: "bocpy" -version: 0.13.0 -date-released: 2026-06-24 +version: 0.14.0 +date-released: 2026-07-02 url: "https://github.com/microsoft/bocpy" \ No newline at end of file diff --git a/examples/boids.py b/examples/boids.py index 6e0e830..c32cda5 100644 --- a/examples/boids.py +++ b/examples/boids.py @@ -92,7 +92,7 @@ def avoid_others(neighbors: Matrix, pos: Matrix, bounds = BoundingBox(left, top, right, bottom) move = Matrix.vector([0, 0]) - for npos in neighbors: + for npos in neighbors.row_views(): if bounds.is_outside(npos.x, npos.y): continue @@ -263,7 +263,7 @@ def spatial_hashing(self, positions: Matrix): self.grid_cells.clear() - for i, pos in enumerate(positions): + for i, pos in enumerate(positions.row_views()): r = int_coord(pos.y, self.spacing) c = int_coord(pos.x, self.spacing) self.grid_cells.add(Cell(r, c)) @@ -303,9 +303,8 @@ def build_cell_data(self, positions: Matrix, velocities: Matrix, row: int, colum start = self.cell_start[h] end = self.cell_start[h + 1] boids = [] - for i in range(start, end): - b = self.cell_entries[i] - pos = positions[b] + for b in self.cell_entries[start:end]: + pos = positions.row_view(b) if box.is_outside(pos.x, pos.y): continue diff --git a/pyproject.toml b/pyproject.toml index 71c6835..2869df4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "bocpy" -version = "0.13.0" +version = "0.14.0" authors = [ {name = "bocpy Team", email="bocpy@microsoft.com"} ] diff --git a/sphinx/source/api.rst b/sphinx/source/api.rst index 9487d72..75e57b1 100644 --- a/sphinx/source/api.rst +++ b/sphinx/source/api.rst @@ -128,6 +128,58 @@ Math :special-members: __init__, __eq__, __ne__, __lt__, __le__, __gt__, __ge__, __getitem__, __setitem__ +Matrix views +^^^^^^^^^^^^ + +:meth:`Matrix.row_view`, :meth:`Matrix.column_view`, and the ``m.view[...]`` +accessor return a :class:`MatrixView` — a read/write, 1-D window that +**shares** the matrix's storage instead of copying it. This is the +borrow-by-view counterpart to plain subscripting (``m[i]`` / ``m[i, j]`` / +``m[rows, cols]``), which always **copies**:: + + m = Matrix(3, 4, [float(i) for i in range(12)]) + + row = m.row_view(1) # borrow: aliases m's storage + row[0] = 99.0 # writes THROUGH to m[1, 0] + col = m.view[:, 2] # slice-notation borrow of column 2 + snap = m[1] # copy: an independent 1 x 4 Matrix + +Key semantics: + +* **Aliasing / write-through.** Reading or writing a view element touches the + source matrix directly (``view[i] = x``, ``view[:] = scalar``, + ``view[:] = sequence``). Sub-views (``view[1:3]``) and the free transpose + (``view.T``, an O(1) orientation flip) alias the same storage. +* **Copy vs borrow.** ``m[...]`` copies; ``m.row_view`` / ``m.column_view`` / + ``m.view[...]`` borrow. Use :meth:`MatrixView.copy` (or + :meth:`Matrix.from_view`) to materialise an independent owned + :class:`Matrix`. +* **Iterating views.** :meth:`Matrix.row_views` / :meth:`Matrix.column_views` + yield a fresh borrow over each row / column in turn — a cleaner spelling of + ``for i in range(m.rows): m.row_view(i)``. Each yielded view is a distinct + object (never a reused cursor), so ``list(m.row_views())`` behaves like any + other iterator. Plain ``for row in m`` still **copies** each row. +* **Whole-buffer lifetime pin.** A view keeps the *entire* backing matrix + alive for as long as the view exists, even a one-element span. +* **Not shippable.** A view cannot be pickled, sent, or placed in a + :class:`Cown`; doing so raises :class:`TypeError` naming the ``.copy()`` + remedy. ``.copy()`` yields an owned, shippable matrix. +* **Arithmetic materialises.** ``+``, ``-``, ``*``, ``/``, ``@``, unary ``-``, + and ``abs()`` on a view return a fresh owned :class:`Matrix`; augmented + ``view += x`` therefore rebinds the name to a new matrix rather than writing + through. +* **Column views are strided.** A column view materialises (for arithmetic or + ``.copy()``) by gathering strided elements, so it costs a little more than a + contiguous row view. +* **Ownership.** Every value access requires the current interpreter to own + the source matrix; a released base raises :class:`RuntimeError`. + +.. autoclass:: MatrixView + :members: + :undoc-members: + :special-members: __len__, __getitem__, __setitem__ + + Messaging --------- diff --git a/sphinx/source/conf.py b/sphinx/source/conf.py index 8ca1227..a1dfaa5 100644 --- a/sphinx/source/conf.py +++ b/sphinx/source/conf.py @@ -14,7 +14,7 @@ project = 'bocpy' copyright = '2026, Microsoft' author = 'Microsoft' -release = '0.13.0' +release = '0.14.0' # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/src/bocpy/__init__.py b/src/bocpy/__init__.py index 002ee72..f13bd67 100644 --- a/src/bocpy/__init__.py +++ b/src/bocpy/__init__.py @@ -6,7 +6,7 @@ import sys from ._core import drain, receive, send, set_tags, TIMEOUT -from ._math import Matrix +from ._math import Matrix, MatrixView from .behaviors import (Behaviors, Cown, notice_delete, notice_read, notice_seed, notice_update, notice_write, noticeboard, PinnedCown, pump, PumpResult, quiesce, @@ -67,7 +67,8 @@ def get_sources() -> list[str]: return [] -__all__ = ["Behaviors", "Cown", "Matrix", "PinnedCown", "PumpResult", +__all__ = ["Behaviors", "Cown", "Matrix", "MatrixView", "PinnedCown", + "PumpResult", "REMOVED", "TIMEOUT", "WORKER_COUNT", "__version__", "drain", "get_include", "get_sources", diff --git a/src/bocpy/__init__.pyi b/src/bocpy/__init__.pyi index 7319a54..510fae6 100644 --- a/src/bocpy/__init__.pyi +++ b/src/bocpy/__init__.pyi @@ -88,6 +88,25 @@ class Matrix: Supports element-wise arithmetic (``+``, ``-``, ``*``, ``/``), matrix multiplication (``@``), in-place variants (``+=``, ``-=``, ``*=``, ``/=``), unary ``-`` and ``abs()``, and subscript indexing with integers or slices. + + **Output-routing convention.** Methods that produce a same-shape result + (the unary element-wise ops ``ceil``/``floor``/``round``/``negate``/ + ``abs``/``sqrt``/``reciprocal``/``sign``/``cos``/``sin``, plus ``clip``) + accept a keyword-only ``in_place=`` to mutate ``self`` and a keyword-only + ``out=`` to write into a caller-supplied matrix; the two are mutually + exclusive. + Methods with a natural in-place meaning but no fresh-allocation form + (``normalize``, ``perpendicular``, ``transpose``) accept only ``in_place=``. + Reductions and gathers (``add``/``subtract``/``multiply``/``divide``, the + six comparisons, ``take``, ``take_along_axis``, ``where``) accept only + ``out=``. Scatter-style writers (``put``, ``put_along_axis``) always mutate + and return ``self``. + + **Masked-reduction empty-group sentinels.** When a ``where=`` mask + excludes every element of a reduction group, the published value depends on + the op: additive ops (``sum``, ``magnitude``, ``magnitude_squared``) yield + ``0``; ``mean``, ``min``, and ``max`` yield ``NaN`` (matching NumPy's + empty-slice semantics); ``argmin``/``argmax`` yield ``-1``. """ def __init__(self, rows: int, columns: int, @@ -171,6 +190,17 @@ class Matrix: def shape(self) -> tuple[int, int]: """The ``(rows, columns)`` shape of the matrix.""" + @property + def view(self) -> "MatrixViewIndexer": + """Slice-notation constructor for read/write row and column views. + + ``m.view[i]`` and ``m.view[i, cols]`` return a row + :class:`MatrixView`; ``m.view[rows, j]`` returns a column + :class:`MatrixView`. Unlike ``m[...]`` (which copies), the result + **aliases** ``m``'s storage. A two-axis key must resolve to a single + row or column; a sub-matrix selection raises :class:`TypeError`. + """ + def transpose(self, in_place: bool = False) -> "Matrix": """Return a transposed matrix. @@ -179,36 +209,55 @@ class Matrix: return a new transposed :class:`Matrix`. """ - def sum(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def sum(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None) -> Union[float, "Matrix"]: """Sum of matrix elements. :param axis: If ``None``, return the total sum as a float. If ``0``, return a 1 x *columns* row vector of column sums. If ``1``, return a *rows* x 1 column vector of row sums. + :param where: Optional same-shape mask matrix. An element is included + only where its mask cell is non-zero (NaN counts as included). A + group with no included element sums to ``0``. """ - def mean(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def mean(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None) -> Union[float, "Matrix"]: """Arithmetic mean of matrix elements. :param axis: If ``None``, return the overall mean as a float. If ``0``, return a 1 x *columns* row vector of column means. If ``1``, return a *rows* x 1 column vector of row means. + :param where: Optional same-shape mask matrix. The mean is taken over + only the elements whose mask cell is non-zero (NaN counts as + included). A group with no included element yields ``NaN`` + (matching NumPy's mean of an empty slice), distinguishing it from + a genuine zero mean. """ - def magnitude(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def magnitude(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None) -> Union[float, "Matrix"]: """Euclidean magnitude (L2 norm) of matrix elements. :param axis: If ``None``, return the total magnitude as a float. If ``0``, return a 1 x *columns* row vector of column magnitudes. If ``1``, return a *rows* x 1 column vector of row magnitudes. + :param where: Optional same-shape mask matrix. Only elements whose + mask cell is non-zero contribute (NaN counts as included). A + group with no included element yields ``0``. """ - def magnitude_squared(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def magnitude_squared(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None + ) -> Union[float, "Matrix"]: """Sum of squared elements (squared L2 norm), avoiding the square-root step. :param axis: If ``None``, return the total squared magnitude as a float. If ``0``, return a 1 x *columns* row vector of column squared magnitudes. If ``1``, return a *rows* x 1 column vector of row squared magnitudes. + :param where: Optional same-shape mask matrix. Only elements whose + mask cell is non-zero contribute (NaN counts as included). A + group with no included element yields ``0``. """ def vecdot(self, other: "Matrix", @@ -445,51 +494,103 @@ class Matrix: or a ``Nx2`` / ``2xN`` batch. """ - def min(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def min(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None) -> Union[float, "Matrix"]: """Minimum of matrix elements. :param axis: If ``None``, return the overall minimum as a float. If ``0``, return a 1 x *columns* row vector of column minima. If ``1``, return a *rows* x 1 column vector of row minima. + :param where: Optional same-shape mask matrix. Only elements whose + mask cell is non-zero are considered (NaN counts as included). A + group with no included element yields ``NaN``. """ - def max(self, axis: Optional[int] = None) -> Union[float, "Matrix"]: + def max(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None) -> Union[float, "Matrix"]: """Maximum of matrix elements. :param axis: If ``None``, return the overall maximum as a float. If ``0``, return a 1 x *columns* row vector of column maxima. If ``1``, return a *rows* x 1 column vector of row maxima. + :param where: Optional same-shape mask matrix. Only elements whose + mask cell is non-zero are considered (NaN counts as included). A + group with no included element yields ``NaN``. """ - def argmin(self, axis: Optional[int] = None) -> Union[int, "Matrix"]: + @overload + def argmin(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, + as_matrix: Literal[False] = False) -> Union[int, list[int]]: ... + + @overload + def argmin(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, *, + as_matrix: Literal[True]) -> Union[int, "Matrix"]: ... + + def argmin(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, + as_matrix: bool = False + ) -> Union[int, list[int], "Matrix"]: """Index of the minimum element (first occurrence on ties). :param axis: If ``None``, return the flat row-major index of the - overall minimum as an ``int``. If ``0``, return a 1 x *columns* - row vector of per-column row indices. If ``1``, return a - *rows* x 1 column vector of per-row column indices. + overall minimum as an ``int``. If ``0``, reduce down the rows; + if ``1``, reduce across the columns. + :param where: Optional same-shape mask matrix. An element is + considered only where its mask cell is non-zero (NaN counts as + included). A group with no included element yields ``-1``, the + integer analog of a masked minimum returning NaN. + :param as_matrix: Controls the axis-wise index form. With ``False`` + (the default) the per-group indices are returned as a Python + ``list[int]``, directly usable in fancy indexing (e.g. fed to + :meth:`take_along_axis`). With ``True`` they are instead a + :class:`Matrix` vector of ``float`` index positions. Ignored for + ``axis=None`` (always a single ``int``). .. note:: - NaN elements are skipped unless the running extreme starts at - NaN (element 0 along the reduced axis), which pins the result - to that position. This differs from NumPy, which propagates NaN. + NaN elements are skipped unless the first *included* element + along the reduced axis is NaN, which pins the result to that + position. This differs from NumPy, which propagates NaN. """ - def argmax(self, axis: Optional[int] = None) -> Union[int, "Matrix"]: + @overload + def argmax(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, + as_matrix: Literal[False] = False) -> Union[int, list[int]]: ... + + @overload + def argmax(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, *, + as_matrix: Literal[True]) -> Union[int, "Matrix"]: ... + + def argmax(self, axis: Optional[int] = None, + where: Optional["Matrix"] = None, + as_matrix: bool = False + ) -> Union[int, list[int], "Matrix"]: """Index of the maximum element (first occurrence on ties). :param axis: If ``None``, return the flat row-major index of the - overall maximum as an ``int``. If ``0``, return a 1 x *columns* - row vector of per-column row indices. If ``1``, return a - *rows* x 1 column vector of per-row column indices. + overall maximum as an ``int``. If ``0``, reduce down the rows; + if ``1``, reduce across the columns. + :param where: Optional same-shape mask matrix. An element is + considered only where its mask cell is non-zero (NaN counts as + included). A group with no included element yields ``-1``, the + integer analog of a masked maximum returning NaN. + :param as_matrix: Controls the axis-wise index form. With ``False`` + (the default) the per-group indices are returned as a Python + ``list[int]``, directly usable in fancy indexing (e.g. fed to + :meth:`take_along_axis`). With ``True`` they are instead a + :class:`Matrix` vector of ``float`` index positions. Ignored for + ``axis=None`` (always a single ``int``). .. note:: - NaN elements are skipped unless the running extreme starts at - NaN (element 0 along the reduced axis), which pins the result - to that position. This differs from NumPy, which propagates NaN. + NaN elements are skipped unless the first *included* element + along the reduced axis is NaN, which pins the result to that + position. This differs from NumPy, which propagates NaN. """ - def ceil(self, in_place: bool = False, *, + def ceil(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Round each element up to the nearest integer. @@ -499,7 +600,7 @@ class Matrix: exclusive with ``in_place``. """ - def floor(self, in_place: bool = False, *, + def floor(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Round each element down to the nearest integer. @@ -509,7 +610,7 @@ class Matrix: exclusive with ``in_place``. """ - def round(self, in_place: bool = False, *, + def round(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Round each element to the nearest integer (banker's rounding). @@ -519,7 +620,7 @@ class Matrix: exclusive with ``in_place``. """ - def negate(self, in_place: bool = False, *, + def negate(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Negate every element. @@ -529,7 +630,7 @@ class Matrix: exclusive with ``in_place``. """ - def abs(self, in_place: bool = False, *, + def abs(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Take the absolute value of every element. @@ -539,7 +640,7 @@ class Matrix: exclusive with ``in_place``. """ - def sqrt(self, in_place: bool = False, *, + def sqrt(self, *, in_place: bool = False, out: Optional["Matrix"] = None) -> "Matrix": """Take the square root of every element. @@ -552,8 +653,52 @@ class Matrix: exclusive with ``in_place``. """ + def reciprocal(self, *, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Take the reciprocal (``1 / x``) of every element. + + Zero elements yield +/-``inf`` (no exception is raised), matching + :func:`numpy.reciprocal`. + + :param in_place: When ``True``, mutate ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + """ + + def sign(self, *, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Return -1, 0, or 1 for each element by sign (``NaN`` maps to 0). + + :param in_place: When ``True``, mutate ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + """ + + def cos(self, *, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Take the cosine (radians) of every element. + + :param in_place: When ``True``, mutate ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + """ + + def sin(self, *, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Take the sine (radians) of every element. + + :param in_place: When ``True``, mutate ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + """ + def less(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self < other`` as a 0/1 mask matrix. Distinct from the ``<`` operator, which returns a single @@ -562,6 +707,8 @@ class Matrix: :param other: A same-shape matrix, a scalar (including ``bool``), a ``1x1`` matrix, a row/column vector that broadcasts (same rules as arithmetic), or a list/tuple of numbers. + :param out: A same-shape :class:`Matrix` to write the mask into + (allocation-free); returned in place of a fresh matrix. :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons yield ``0.0``. :raises ValueError: on a non-broadcastable shape or an empty @@ -571,7 +718,8 @@ class Matrix: """ def less_equal(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self <= other`` as a 0/1 mask matrix. Distinct from the ``<=`` operator, which returns a single @@ -589,7 +737,8 @@ class Matrix: """ def greater(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self > other`` as a 0/1 mask matrix. Distinct from the ``>`` operator, which returns a single @@ -607,7 +756,8 @@ class Matrix: """ def greater_equal(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self >= other`` as a 0/1 mask matrix. Distinct from the ``>=`` operator, which returns a single @@ -625,7 +775,8 @@ class Matrix: """ def equal(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self == other`` as a 0/1 mask matrix. Distinct from the ``==`` operator, which returns a single @@ -643,7 +794,8 @@ class Matrix: """ def not_equal(self, other: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Element-wise ``self != other`` as a 0/1 mask matrix. Distinct from the ``!=`` operator, which returns a single @@ -661,7 +813,8 @@ class Matrix: """ def clip(self, min: Optional[float] = None, - max: Optional[float] = None) -> "Matrix": + max: Optional[float] = None, *, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": """Clamp every element to ``[min, max]``. The first argument is the lower bound and the second the upper @@ -671,15 +824,165 @@ class Matrix: :param min: Lower clipping bound, or ``None`` for no lower bound. :param max: Upper clipping bound, or ``None`` for no upper bound. - :return: A new clipped :class:`Matrix`. - :raises ValueError: if both *min* and *max* are ``None``. + :param in_place: When ``True``, clamp ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + :return: A new clipped :class:`Matrix`, or ``self``/``out`` when + ``in_place``/``out`` is used. + :raises ValueError: if both *min* and *max* are ``None``, or if + ``in_place`` and ``out`` are both given. :raises AssertionError: if both bounds are given and *max* < *min*. - :raises TypeError: if a given bound is not a real number. + :raises TypeError: if a given bound is not a real number, or if + *out* is not a :class:`Matrix`. + """ + + def bin(self, bins: int, *, + bounds: Optional[tuple[float, float]] = None, + right: bool = False, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Map each element to its equal-width bin index over a range. + + The range is split into *bins* equal-width intervals and each element + is replaced by the integer index (stored as a ``float``) of the + interval it falls in. By default the range is taken from the matrix's + own elements (both bounds computed ignoring ``NaN``); pass *bounds* to + fix it and skip the min/max scan. + + :param bins: Number of equal-width bins; must be ``>= 1``. + :param bounds: An optional ``(min, max)`` tuple fixing the value + range. When given, the ``O(size)`` data scan is skipped and a + value outside the range clamps into the first or last bin. When + ``None`` (the default), the range is the matrix's own + ``(min, max)``. + :param right: When ``False`` (default), a value on an interior + boundary falls into the upper bin; when ``True``, into the lower + bin. A value equal to ``max`` always lands in the last bin + (index ``bins - 1``). + :param in_place: When ``True``, bin ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + :return: A new :class:`Matrix` of bin indices, or ``self``/``out`` + when ``in_place``/``out`` is used. + :raises ValueError: if *bins* < 1, if *bounds* is not a length-2 + ``(min, max)`` pair of finite numbers with ``max >= min``, or if + ``in_place`` and ``out`` are both given. + :raises TypeError: if a *bounds* value is not a real number, or if + *out* is not a :class:`Matrix`. + + .. note:: + A ``NaN`` element maps to ``-1`` (the invalid-index sentinel). A + degenerate range (``max == min``) sends every element to bin 0. + Because the equal-width edges are implicit, boundary placement may + differ by at most one bin from :meth:`digitize` fed edges derived + from the same range. + """ + + def digitize(self, edges: Union["Matrix", Sequence[float]], *, + right: bool = False, in_place: bool = False, + out: Optional["Matrix"] = None) -> "Matrix": + """Map each element to the index of the edge interval it falls in. + + Mirrors :func:`numpy.digitize`. Each element is replaced by the + index (stored as a ``float``) of the bin bounded by *edges*. + + :param edges: A strictly increasing vector of interior boundaries, + given as a :class:`Matrix` vector (1 x *n* or *n* x 1) or any + coercible sequence of numbers. + :param right: When ``False`` (default), the index is the count of + edges ``<=`` the value (bin ``i`` means + ``edges[i-1] <= value < edges[i]``); when ``True``, the count of + edges ``<`` the value (``edges[i-1] < value <= edges[i]``). + Indices span ``[0, len(edges)]``. + :param in_place: When ``True``, digitize ``self`` and return it. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. Mutually + exclusive with ``in_place``. + :return: A new :class:`Matrix` of bin indices, or ``self``/``out`` + when ``in_place``/``out`` is used. + :raises ValueError: if *edges* is not a vector, is empty, is not + strictly increasing or contains ``NaN``, or if ``in_place`` and + ``out`` are both given. + :raises TypeError: if *out* is not a :class:`Matrix`. + + .. note:: + A ``NaN`` element maps to ``-1`` (the invalid-index sentinel). """ def copy(self) -> "Matrix": """Return a deep copy of this matrix.""" + def values(self) -> Iterator[float]: + """Yield every element as a ``float`` in row-major order. + + Lazy: one ``float`` is boxed per step, so streaming a large + matrix never materialises an intermediate list. The current + interpreter must own the matrix for the duration of iteration. + """ + + def row_view(self, row: int, + columns: Optional[Union[int, slice]] = None, *, + read_only: bool = False) -> "MatrixView": + """Return a 1-D view onto a single row (no copy). + + The view **aliases** this matrix's storage: reading or writing an + element (``rv[i]``, ``rv[i] = x``, ``rv[:] = ...``) touches the + underlying matrix directly. *columns* selects a sub-span of the row — + an ``int`` (a single cell), a ``slice`` (including non-unit and + negative steps), or ``None`` (the whole row). The view pins the whole + backing buffer alive and cannot be pickled or sent across + interpreters; call :meth:`MatrixView.copy` to detach an owned + :class:`Matrix`. + + :param row: The row index (negative counts from the end). + :param columns: A column selector, or ``None`` for the whole row. + :param read_only: When ``True``, the view — and any view derived from + it via ``.T`` or slicing — rejects writes with ``TypeError``. + :return: A :class:`MatrixView` aliasing the selected row. + """ + + def column_view(self, column: int, + rows: Optional[Union[int, slice]] = None, *, + read_only: bool = False) -> "MatrixView": + """Return a 1-D view onto a single column (no copy). + + The column-oriented twin of :meth:`row_view`; see it for the aliasing, + whole-buffer-pin, and not-shippable semantics. *rows* selects a + sub-span of the column. + + :param column: The column index (negative counts from the end). + :param rows: A row selector, or ``None`` for the whole column. + :param read_only: When ``True``, the view (and any view derived from + it) rejects writes with ``TypeError``. + :return: A :class:`MatrixView` aliasing the selected column. + """ + + def row_views(self, *, + read_only: bool = False) -> Iterator["MatrixView"]: + """Yield a :class:`MatrixView` over each row in turn. + + Each yielded view is a distinct object aliasing this matrix's + storage (a ``1 x N`` row view per row). A cleaner spelling of + ``for i in range(m.rows): m.row_view(i)``. The current interpreter + must own the matrix for the duration of iteration. + + :param read_only: When ``True``, every yielded view is immutable + (writes raise ``TypeError``) — a convenient way to iterate a + matrix's rows without risking accidental mutation. + """ + + def column_views(self, *, + read_only: bool = False) -> Iterator["MatrixView"]: + """Yield a :class:`MatrixView` over each column in turn. + + The column-oriented twin of :meth:`row_views`; each yielded view is a + distinct ``M x 1`` column view aliasing this matrix's storage. + + :param read_only: When ``True``, every yielded view is immutable + (writes raise ``TypeError``). + """ + def __reduce__(self) -> tuple: """Support pickling and :func:`copy.deepcopy`. @@ -688,7 +991,8 @@ class Matrix: object overhead. The current interpreter must own the matrix. """ - def take(self, indices: Union[list[int], tuple[int]], axis=0) -> "Matrix": + def take(self, indices: Union[list[int], tuple[int]], axis=0, *, + out: Optional["Matrix"] = None) -> "Matrix": """Return a new matrix containing only the selected rows or columns. *indices* is a 1-D list or tuple of ints. Negative indices count @@ -699,10 +1003,19 @@ class Matrix: :param indices: The row or column indices to take. Must be a non-empty list or tuple of ints. :param axis: ``0`` to take rows, ``1`` to take columns. + :param out: A pre-allocated :class:`Matrix` of the selection shape + (``len(indices)`` rows by this matrix's columns for ``axis=0``; + this matrix's rows by ``len(indices)`` columns for ``axis=1``) + to write the result into and return, avoiding a fresh + allocation. All indices are validated before any write, so a + rejected call leaves *out* untouched. *out* must not alias + ``self``. :raises IndexError: if an index is out of range, or if *indices* is empty. :raises KeyError: if *axis* is not ``0`` or ``1``. - :raises TypeError: if an index is not an int. + :raises TypeError: if an index is not an int, or if *out* is not a + :class:`Matrix`. + :raises ValueError: if *out* has the wrong shape or aliases ``self``. :raises OverflowError: if an index exceeds the platform word size. """ @@ -748,6 +1061,157 @@ class Matrix: :raises OverflowError: if an index exceeds the platform word size. """ + def take_along_axis(self, indices: Union[list[int], tuple[int]], + axis=0, *, + out: Optional["Matrix"] = None) -> "Matrix": + """Gather one element per row or column along an axis. + + The ``np.take_along_axis`` counterpart of :meth:`take` (which + selects whole rows or columns). With ``axis=1`` *indices* gives one + column index per row, so ``len(indices)`` must equal the row count, + and the result is a ``rows x 1`` column vector with + ``out[r] == self[r][indices[r]]``. With ``axis=0`` *indices* gives + one row index per column, so ``len(indices)`` must equal the column + count, and the result is a ``1 x columns`` row vector with + ``out[c] == self[indices[c]][c]``. + + This pairs directly with :meth:`argmin`/:meth:`argmax` along the + same axis: the index list they return feeds straight back in to + gather the reduced values. Negative indices count from the end. A + ``bool`` element is treated as the integer ``0``/``1``. + + :param indices: One index per row (``axis=1``) or per column + (``axis=0``). Must be a list or tuple of ints whose length + matches that axis. + :param axis: ``1`` to gather a column index per row, ``0`` to gather + a row index per column. + :param out: A pre-allocated :class:`Matrix` of the result shape + (``1 x columns`` for ``axis=0``; ``rows x 1`` for ``axis=1``) + to write the result into and return, avoiding a fresh + allocation. All indices are validated before any write, so a + rejected call leaves *out* untouched. *out* must not alias + ``self``. + :raises IndexError: if an index is out of range. + :raises ValueError: if ``len(indices)`` does not match the axis, or + if *out* has the wrong shape or aliases ``self``. + :raises KeyError: if *axis* is not ``0`` or ``1``. + :raises TypeError: if an index is not an int, or if *out* is not a + :class:`Matrix`. + :raises OverflowError: if an index exceeds the platform word size. + """ + + def put_along_axis(self, indices: Union[list[int], tuple[int]], + value: Union[int, float, "Matrix"], axis=0, + accumulate: bool = False) -> "Matrix": + """Assign one element per row or column along an axis in place. + + The write-side counterpart of :meth:`take_along_axis`. With + ``axis=1`` element ``i`` lands in ``self[i][indices[i]]`` and + ``len(indices)`` must equal the row count; with ``axis=0`` it lands + in ``self[indices[i]][i]`` and ``len(indices)`` must equal the column + count. *value* may be a scalar (a real number or a ``1x1`` matrix, + broadcast over the selection) or a vector matching the selection + shape (``rows x 1`` for ``axis=1``, ``1 x columns`` for ``axis=0``). + + All indices and the *value* shape are validated before any element + is written, so a rejected call leaves the matrix unchanged. Negative + indices count from the end. A ``bool`` index element is treated as + the integer ``0``/``1``. + + With ``accumulate=False`` (the default) duplicate indices follow + last-write-wins. With ``accumulate=True`` the values are *added* + into the selection, so duplicate indices fold additively + (scatter-add). + + :param indices: One index per row (``axis=1``) or per column + (``axis=0``). Must be a list or tuple of ints whose length + matches that axis. + :param value: The scalar or vector to assign into the selection. + :param axis: ``1`` to index a column per row, ``0`` to index a row + per column. + :param accumulate: When ``True``, add into the selection instead of + overwriting, so duplicate indices accumulate. + :return: ``self`` (to allow chaining). + :raises IndexError: if an index is out of range. + :raises ValueError: if ``len(indices)`` does not match the axis, or + if a matrix *value* shape does not match the selection shape. + :raises KeyError: if *axis* is not ``0`` or ``1``. + :raises TypeError: if *value* is neither a real number nor a matrix, + or if an index is not an int. + :raises OverflowError: if an index exceeds the platform word size. + """ + + def repeat_interleave(self, repeats: int, axis=None) -> "Matrix": + """Repeat each element, row, or column consecutively. + + Interleaved like ``np.repeat`` / ``torch.repeat_interleave`` (not + tiled): ``[a, b]`` with ``repeats=2`` becomes ``[a, a, b, b]``. With + ``axis=0`` each row is repeated into a ``(rows*repeats) x columns`` + matrix; with ``axis=1`` each column is repeated into + ``rows x (columns*repeats)``; with ``axis=None`` (the default) the + row-major buffer is flattened to a ``1 x (size*repeats)`` row vector. + + :param repeats: The number of consecutive copies of each element, + row, or column. Must be a positive integer. + :param axis: ``0`` to repeat rows, ``1`` to repeat columns, or + ``None`` to flatten and repeat each element. + :return: A new matrix; the receiver is unchanged. + :raises ValueError: if *repeats* is not positive, or if *axis* is not + ``-2``, ``-1``, ``0``, ``1``, or ``None``. + :raises OverflowError: if the result would exceed the platform word + size. + """ + + @overload + def topk(self, k: int, axis: Optional[int] = None, largest: bool = True, + where: Optional["Matrix"] = None, + as_matrix: Literal[False] = False + ) -> tuple["Matrix", Union[list[int], list[list[int]]]]: ... + + @overload + def topk(self, k: int, axis: Optional[int] = None, largest: bool = True, + where: Optional["Matrix"] = None, *, + as_matrix: Literal[True] + ) -> tuple["Matrix", "Matrix"]: ... + + def topk(self, k: int, axis: Optional[int] = None, largest: bool = True, + where: Optional["Matrix"] = None, as_matrix: bool = False + ) -> tuple["Matrix", Union[list[int], list[list[int]], "Matrix"]]: + """The *k* extreme elements per reduction group, in sorted order. + + Returns a ``(values, indices)`` tuple. ``largest=True`` (the default) + selects the *k* greatest in descending order; ``largest=False`` + selects the *k* smallest in ascending order. Ties keep the first + occurrence (NumPy tie-break) and any NaN sorts last. + + :param k: The number of elements to select per group. Must be a + positive integer no larger than the reduced axis length (every + cell counts, masked or not). + :param axis: With ``None`` (the default) the whole row-major buffer + is one group and *values* is a ``1 x k`` row vector. With ``0`` + each column is reduced down the rows and *values* is + ``k x columns``. With ``1`` each row is reduced across the + columns and *values* is ``rows x k``. + :param largest: Select the largest *k* (descending) when ``True``, + else the smallest *k* (ascending). + :param where: Optional same-shape mask matrix. An element is + considered only where its mask cell is non-zero (NaN counts as + included). A group with fewer than *k* included elements fills its + leading slots and pads the rest with ``NaN`` values and ``-1`` + indices. + :param as_matrix: Controls the *indices* form. With ``False`` (the + default) *indices* is Python ints: a flat ``list[int]`` for + ``axis=None``, or a list of *columns*/*rows* lists (each *k* + indices) for ``axis=0``/``axis=1``. With ``True`` *indices* is + instead a :class:`Matrix` of the same shape as *values*, each cell + the source index (as a ``float``) of the value beside it, and the + ``-1`` pad becomes ``-1.0``. + :return: A ``(values, indices)`` tuple as described above. + :raises ValueError: if *k* is not positive, exceeds the reduced axis + length, or if *axis* is not ``-2``, ``-1``, ``0``, ``1``, or + ``None``. + """ + def __add__(self, other: Union["Matrix", int, float]) -> "Matrix": """Element-wise addition.""" @@ -793,6 +1257,12 @@ class Matrix: raises :class:`IndexError` and an empty list raises :class:`IndexError`. + A result collapses to a Python ``float`` only when every selector is + an integer and the selection is a single cell. A slice anywhere in + the key keeps the result a :class:`Matrix`, even when it is 1x1: so + ``m[0:1, 0:1]``, ``m[i, 0:1]``, and ``m[0:1, j]`` all return a 1x1 + :class:`Matrix`, while ``m[i, j]`` returns a ``float``. + Column gather requires the bare ``:`` row selector: ``m[0:R, [c]]`` (a full *range* rather than ``:``) raises :class:`IndexError`. Paired lists (``m[[r], [c]]``), @@ -947,7 +1417,8 @@ class Matrix: def where(cls, mask: "Matrix", a: Union["Matrix", int, float, Sequence[Union[int, float]]], b: Union["Matrix", int, float, - Sequence[Union[int, float]]]) -> "Matrix": + Sequence[Union[int, float]]], *, + out: Optional["Matrix"] = None) -> "Matrix": """Select element-wise from *a* or *b* on a truthy mask. Returns a fresh matrix taking *a* where the corresponding *mask* @@ -960,6 +1431,9 @@ class Matrix: *mask*'s shape. :param b: A scalar (including ``bool``), a list/tuple of numbers, or a :class:`Matrix` matching *mask*'s shape. + :param out: A same-shape :class:`Matrix` to write the result into + (allocation-free); returned in place of a fresh matrix. May alias + *mask*, *a*, or *b*. :return: A new :class:`Matrix` with *mask*'s shape. :raises TypeError: if *a* or *b* is neither a matrix, a scalar, nor a list/tuple of numbers, or is a list/tuple holding a non-number. @@ -985,6 +1459,15 @@ class Matrix: :return: A new :class:`Matrix` with every element set to ``1.0``. """ + @classmethod + def full(cls, size: tuple[int, int], value: float) -> "Matrix": + """Create a matrix filled with a constant value. + + :param size: A ``(rows, columns)`` tuple specifying the shape. + :param value: The value every element is set to. + :return: A new :class:`Matrix` with every element set to ``value``. + """ + @classmethod def normal(cls, mean: Optional[float], stddev: Optional[float], size: Optional[tuple[int, int]] = None) -> Union[float, "Matrix"]: @@ -1014,11 +1497,10 @@ class Matrix: :param value: The seed value. .. note:: - The generator is the process-global C library PRNG shared by - every sub-interpreter, so a seed only makes subsequent draws - reproducible when random generation stays on a single thread; - concurrent draws interleave on the shared state. The sequence - is also not portable across platforms. + Each interpreter owns an independent splitmix64 stream, so a + seed makes that interpreter's subsequent draws reproducible; + parallel workers seed their own streams independently. The + sequence is not portable across platforms. """ @classmethod @@ -1042,6 +1524,199 @@ class Matrix: :return: A new :class:`Matrix` containing the concatenated data. """ + @classmethod + def from_view(cls, view: "MatrixView") -> "Matrix": + """Materialise a :class:`MatrixView` into a new owned :class:`Matrix`. + + Equivalent to :meth:`MatrixView.copy`, in the constructor-family + spelling next to :meth:`zeros` / :meth:`ones` / :meth:`vector`. The + result shares no storage with the view's base. + + :param view: The :class:`MatrixView` to copy. + :return: A new :class:`Matrix` holding the view's elements. + :raises TypeError: if *view* is not a :class:`MatrixView`. + """ + + +class MatrixView: + """A read/write, 1-D window onto a single row or column of a :class:`Matrix`. + + A view **aliases** its source matrix's storage rather than copying it: + reading an element reads the matrix, and writing one (``v[i] = x``, + ``v[:] = scalar``, ``v[:] = sequence``) writes straight through to the + matrix. A view **pins the whole backing buffer** alive for its lifetime and + is **not shippable** — it cannot be pickled or sent across interpreters; + call :meth:`copy` to detach an independent, shippable :class:`Matrix`. + + Construct a view via :meth:`Matrix.row_view`, :meth:`Matrix.column_view`, + or the ``m.view[...]`` accessor; direct construction raises + :class:`TypeError`. Arithmetic (``+``, ``-``, ``*``, ``/``, ``@``, unary + ``-``, ``abs()``) materialises the view and returns a fresh owned + :class:`Matrix`; augmented ``v += x`` therefore **rebinds** the name to a + new :class:`Matrix` and does not write through. Every value access requires + the current interpreter to own the source matrix. + """ + + @property + def rows(self) -> int: + """The row count (``1`` for a row view, else the element count).""" + + @property + def columns(self) -> int: + """The column count (``1`` for a column view, else the element count).""" + + @property + def size(self) -> int: + """The number of elements in the view.""" + + @property + def shape(self) -> tuple[int, int]: + """The ``(rows, columns)`` shape (``1 x N`` row or ``N x 1`` column).""" + + @property + def is_row(self) -> bool: + """``True`` for a row view, ``False`` for a column view.""" + + @property + def read_only(self) -> bool: + """``True`` if writes through this view raise ``TypeError``. + + Set at construction via the ``read_only=`` flag on + :meth:`Matrix.row_view` / :meth:`Matrix.column_view` / + :meth:`Matrix.row_views` / :meth:`Matrix.column_views`, and inherited + by views derived through ``.T`` or slicing. + """ + + @property + def T(self) -> "MatrixView": + """A free (no-copy) transposed view: same storage, flipped orientation.""" + + @property + def x(self) -> float: + """The first view element; writes through to the source matrix.""" + + @x.setter + def x(self, value): + """Set the first view element (writes through).""" + + @property + def y(self) -> float: + """The second view element; writes through (``IndexError`` if absent).""" + + @y.setter + def y(self, value): + """Set the second view element (writes through).""" + + @property + def z(self) -> float: + """The third view element; writes through (``IndexError`` if absent).""" + + @z.setter + def z(self, value): + """Set the third view element (writes through).""" + + @property + def w(self) -> float: + """The fourth view element; writes through (``IndexError`` if absent).""" + + @w.setter + def w(self, value): + """Set the fourth view element (writes through).""" + + def __len__(self) -> int: + """Return the number of elements in the view.""" + + def __getitem__(self, key: Union[int, slice] + ) -> Union[float, "MatrixView"]: + """Read an element (``int`` key) or a sub-view (``slice`` key). + + A sub-view aliases the same storage; slices support non-unit and + negative steps. A tuple key raises :class:`TypeError` (a view is 1-D). + """ + + def __setitem__(self, key: Union[int, slice], + value: Union[int, float, "Matrix", "MatrixView", + Sequence[Union[int, float]]]): + """Write through to the source matrix. + + An ``int`` key assigns a single scalar. A ``slice`` key assigns a + scalar (broadcast to every selected element) or a sequence / + :class:`Matrix` / :class:`MatrixView` whose element **count** matches + the slice length (matching is length-only — a view is 1-D, so RHS + orientation is irrelevant). + + :raises IndexError: if an ``int`` key is out of range. + :raises ValueError: if a sequence RHS length differs from the slice. + :raises TypeError: on a tuple key, on ``del v[i]`` (a view cannot + resize), or on a non-numeric scalar RHS. + """ + + def values(self) -> Iterator[float]: + """Yield every element as a ``float`` in view order (lazy, guarded).""" + + def __iter__(self) -> Iterator[float]: + """Iterate the view's elements as ``float`` s. + + A view is 1-D, so ``for x in view`` yields scalars (like iterating a + 1-D array), exactly equivalent to :meth:`values`. + """ + + def copy(self) -> "Matrix": + """Return an independent owned :class:`Matrix` of the view's elements. + + The escape hatch for shipping: the result shares no storage with the + source, so it can be pickled, sent, or placed in a :class:`Cown`. + """ + + def __add__(self, other: Union["Matrix", "MatrixView", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise addition; returns a fresh owned :class:`Matrix`.""" + + def __radd__(self, other) -> "Matrix": + """Reflected element-wise addition.""" + + def __sub__(self, other) -> "Matrix": + """Element-wise subtraction; returns a fresh owned :class:`Matrix`.""" + + def __rsub__(self, other) -> "Matrix": + """Reflected element-wise subtraction.""" + + def __mul__(self, other) -> "Matrix": + """Element-wise multiplication; returns a fresh owned :class:`Matrix`.""" + + def __rmul__(self, other) -> "Matrix": + """Reflected element-wise multiplication.""" + + def __truediv__(self, other) -> "Matrix": + """Element-wise division; returns a fresh owned :class:`Matrix`.""" + + def __rtruediv__(self, other) -> "Matrix": + """Reflected element-wise division.""" + + def __matmul__(self, other) -> "Matrix": + """Matrix multiplication; returns a fresh owned :class:`Matrix`.""" + + def __rmatmul__(self, other) -> "Matrix": + """Reflected matrix multiplication.""" + + def __neg__(self) -> "Matrix": + """Element-wise negation (``-v``); returns a fresh :class:`Matrix`.""" + + def __abs__(self) -> "Matrix": + """Element-wise absolute value (``abs(v)``); returns a :class:`Matrix`.""" + + +class MatrixViewIndexer: + """Private accessor returned by :attr:`Matrix.view` for ``m.view[...]``. + + Not part of the public API — obtained only via :attr:`Matrix.view`. Its + ``__getitem__`` builds a :class:`MatrixView` from native subscript grammar + (``m.view[i]``, ``m.view[i, cols]``, ``m.view[rows, j]``). + """ + + def __getitem__(self, key: Union[int, tuple]) -> "MatrixView": + """Build a row/column :class:`MatrixView` from ``m.view[...]`` grammar.""" + T = TypeVar("T") diff --git a/src/bocpy/_core.c b/src/bocpy/_core.c index cf0131a..88352b8 100644 --- a/src/bocpy/_core.c +++ b/src/bocpy/_core.c @@ -9,6 +9,7 @@ #include "boc_terminator.h" #include #include +#include typedef struct boc_queue BOCQueue; @@ -3319,6 +3320,12 @@ typedef struct behavior_s { BOCCown *result; /// @brief Grouping identifiers, used for reassembling lists of cowns int *group_ids; + /// @brief Number of top-level @when argument slots (cown parameters). + /// @details May exceed the number of distinct groups present in + /// @c group_ids: an empty list group contributes no @c args entry, so this + /// count -- supplied by @c whencall -- is the only record that the slot + /// exists. @ref behavior_execute_impl reconstructs each empty slot as @c []. + Py_ssize_t num_arg_slots; /// @brief The args buffer BOCCown **args; /// @brief The number of args @@ -3394,6 +3401,7 @@ BOCBehavior *behavior_new() { behavior->result = NULL; behavior->rc = 0; behavior->group_ids = NULL; + behavior->num_arg_slots = 0; behavior->args_size = 0; behavior->args = NULL; behavior->captures_size = 0; @@ -3569,10 +3577,16 @@ static int BehaviorCapsule_init(PyObject *op, PyObject *args, PyObject *result = NULL; PyObject *cowns_list = NULL; PyObject *captures = NULL; + Py_ssize_t num_arg_slots = 0; - if (!PyArg_ParseTuple(args, "O!O!OO", &PyUnicode_Type, &thunk, + if (!PyArg_ParseTuple(args, "O!O!OOn", &PyUnicode_Type, &thunk, BOC_STATE->cown_capsule_type, &result, &cowns_list, - &captures)) { + &captures, &num_arg_slots)) { + return -1; + } + + if (num_arg_slots < 0) { + PyErr_SetString(PyExc_ValueError, "num_arg_slots must be non-negative"); return -1; } @@ -3596,6 +3610,8 @@ static int BehaviorCapsule_init(PyObject *op, PyObject *args, self->behavior = behavior; BEHAVIOR_INCREF(behavior); + behavior->num_arg_slots = num_arg_slots; + behavior->thunk = tag_from_PyUnicode(thunk, NULL); if (behavior->thunk == NULL) { return -1; @@ -3637,6 +3653,18 @@ static int BehaviorCapsule_init(PyObject *op, PyObject *args, PyObject *cown; if (!PyArg_ParseTuple(item, "iO!", &group_id, BOC_STATE->cown_capsule_type, &cown)) { + Py_DECREF(cowns); + Py_DECREF(cowns_list_fast); + return -1; + } + + // Bound group_id (exec uses abs(group_id)-1): reject 0, INT_MIN, > slots. + if (group_id == 0 || group_id == INT_MIN || + (Py_ssize_t)abs(group_id) > num_arg_slots) { + PyErr_Format(PyExc_ValueError, + "cown group_id %d out of range for %zd argument slot(s)", + group_id, num_arg_slots); + Py_DECREF(cowns); Py_DECREF(cowns_list_fast); return -1; } @@ -3664,6 +3692,12 @@ static int BehaviorCapsule_init(PyObject *op, PyObject *args, return -1; } + // Guard the caller-controlled tuple-size sum against Py_ssize_t overflow. + if (behavior->num_arg_slots > PY_SSIZE_T_MAX - behavior->captures_size) { + PyErr_SetString(PyExc_OverflowError, "too many behavior argument slots"); + return -1; + } + // +2 over the cown args: one hold for the result cown, one so dispatch waits // for 2PL phase-2 to finish before the thunk's release calls can race it. behavior->count = (int_least64_t)(behavior->args_size + 2); @@ -4102,73 +4136,67 @@ static PyObject *BehaviorCapsule_release(PyObject *op, /// failure static PyObject *behavior_execute_impl(BOCBehavior *behavior, PyObject *boc_export) { - size_t num_groups = 0; - if (behavior->args_size > 0) { - num_groups = abs(behavior->group_ids[behavior->args_size - 1]); - } - - num_groups += behavior->captures_size; - - PyObject *thunk_args = PyTuple_New(num_groups); + // Cown i -> slot abs(group_id)-1; untouched slots become empty groups []. + size_t num_arg_slots = (size_t)behavior->num_arg_slots; + size_t num_args = num_arg_slots + (size_t)behavior->captures_size; + PyObject *thunk_args = PyTuple_New((Py_ssize_t)num_args); if (thunk_args == NULL) { return NULL; } - Py_ssize_t arg_idx = 0; BOCCown **ptr = behavior->args; - - PyObject *group_list = NULL; - int current_group_id = 0; for (Py_ssize_t i = 0; i < behavior->args_size; ++i, ++ptr) { - PyObject *capsule = cown_capsule_wrap(*ptr, false); int group_id = behavior->group_ids[i]; + Py_ssize_t slot = (Py_ssize_t)(abs(group_id) - 1); - if (group_id == current_group_id) { - if (PyList_Append(group_list, capsule) < 0) { - Py_DECREF(thunk_args); - Py_DECREF(group_list); - return NULL; - } - - Py_DECREF(capsule); - continue; - } - - if (group_list != NULL) { - PyTuple_SET_ITEM(thunk_args, arg_idx, group_list); - arg_idx += 1; - group_list = NULL; - current_group_id = 0; + PyObject *capsule = cown_capsule_wrap(*ptr, false); + if (capsule == NULL) { + Py_DECREF(thunk_args); + return NULL; } if (group_id > 0) { - PyTuple_SET_ITEM(thunk_args, arg_idx, capsule); - arg_idx += 1; + // Singleton argument: the slot holds the capsule directly. + PyTuple_SET_ITEM(thunk_args, slot, capsule); continue; } - group_list = PyList_New(1); + // Group member: append to the slot's list, creating it on first sight. + PyObject *group_list = PyTuple_GET_ITEM(thunk_args, slot); if (group_list == NULL) { + group_list = PyList_New(0); + if (group_list == NULL) { + Py_DECREF(capsule); + Py_DECREF(thunk_args); + return NULL; + } + PyTuple_SET_ITEM(thunk_args, slot, group_list); + } + if (PyList_Append(group_list, capsule) < 0) { + Py_DECREF(capsule); Py_DECREF(thunk_args); return NULL; } - - current_group_id = group_id; - PyList_SET_ITEM(group_list, 0, capsule); + Py_DECREF(capsule); } - if (group_list != NULL) { - PyTuple_SET_ITEM(thunk_args, arg_idx, group_list); - arg_idx += 1; - group_list = NULL; - current_group_id = 0; + // Any cown slot still unset is an empty list group -> []. + for (Py_ssize_t slot = 0; slot < (Py_ssize_t)num_arg_slots; ++slot) { + if (PyTuple_GET_ITEM(thunk_args, slot) == NULL) { + PyObject *empty = PyList_New(0); + if (empty == NULL) { + Py_DECREF(thunk_args); + return NULL; + } + PyTuple_SET_ITEM(thunk_args, slot, empty); + } } ptr = behavior->captures; - for (Py_ssize_t i = 0; i < behavior->captures_size; ++i, ++arg_idx, ++ptr) { + for (Py_ssize_t i = 0; i < behavior->captures_size; ++i, ++ptr) { PyObject *value = Py_NewRef((*ptr)->value); - PyTuple_SET_ITEM(thunk_args, arg_idx, value); + PyTuple_SET_ITEM(thunk_args, (Py_ssize_t)num_arg_slots + i, value); } PyObject *thunk = PyObject_GetAttrString(boc_export, behavior->thunk->str); diff --git a/src/bocpy/_math.c b/src/bocpy/_math.c index ccb6727..139e27c 100644 --- a/src/bocpy/_math.c +++ b/src/bocpy/_math.c @@ -9,6 +9,7 @@ #include #include #include +#include #include @@ -132,6 +133,17 @@ static int update_row_ptrs(matrix_impl *matrix) { static matrix_impl *impl_new(size_t rows, size_t columns) { assert(rows > 0 && columns > 0); + // Guard the rows * columns product itself: a wrapped size_t would + // under-allocate `data` below (the calloc byte-product check only sees + // the already-truncated size), leaving a matrix that advertises more + // cells than its backing store holds -> out-of-bounds access. Every + // constructor (zeros/ones/full/normal/uniform, repeat_interleave, ...) + // routes through here, so this single check covers them all. + if (columns != 0 && rows > SIZE_MAX / columns) { + PyErr_SetString(PyExc_OverflowError, "Matrix dimensions are too large"); + return NULL; + } + matrix_impl *matrix = (matrix_impl *)PyMem_RawMalloc(sizeof(matrix_impl)); if (matrix == NULL) { PyErr_NoMemory(); @@ -773,6 +785,148 @@ static void dispatch_agg_columnwise(matrix_impl *m, enum AggregateOps op, } } +/* -------------------------------------------------------------------------- + Masked aggregate kernels (the ``where=`` path). + + Stamped from BOC_AGG_OPS so a new aggregate automatically gains masked + support. An excluded element (mask cell == 0.0; a NaN mask cell counts as + included, matching where()'s truthiness) is replaced by the op's INIT, + which is the neutral element for that op's STEP and so leaves the + accumulator unchanged, and it is not counted toward ``cnt``. When a group + has no included element the published value is BOC_AGG_MASKED_EMPTY_: + the additive ops (Sum/Magnitude/MagnitudeSquared) collapse to 0 (their + additive identity), while Mean and min/max yield NaN -- there is no element + to average or choose, and NaN matches NumPy's empty-slice mean. + + These run only when a caller passes ``where=`` (the slow path), so they use + a single accumulator rather than the LANES=4 unrolling of the unmasked + kernels -- simplicity over SIMD. Rows/columns are addressed through + row_ptrs so the mask and matrix stay aligned cell-for-cell. + -------------------------------------------------------------------------- */ +#define BOC_AGG_MASKED_EMPTY_Sum 0.0 +#define BOC_AGG_MASKED_EMPTY_Mean NAN +#define BOC_AGG_MASKED_EMPTY_Magnitude 0.0 +#define BOC_AGG_MASKED_EMPTY_MagnitudeSquared 0.0 +#define BOC_AGG_MASKED_EMPTY_Minimum NAN +#define BOC_AGG_MASKED_EMPTY_Maximum NAN + +#define DEFINE_AGG_MASKED_EWISE(ENUM, STAMP, INIT, STEP, MERGE, FINAL, LANES) \ + static double impl_##STAMP##_masked_ewise(const matrix_impl *m, \ + const matrix_impl *mask) { \ + double agg = (INIT); \ + size_t cnt = 0; \ + const double *sp = m->data; \ + const double *kp = mask->data; \ + const size_t n = m->size; \ + for (size_t i = 0; i < n; ++i) { \ + const int inc = (kp[i] != 0.0); \ + const double value = inc ? sp[i] : (INIT); \ + cnt += (size_t)inc; \ + agg = (STEP); \ + } \ + return (cnt > 0) ? (FINAL) : (BOC_AGG_MASKED_EMPTY_##ENUM); \ + } + +#define DEFINE_AGG_MASKED_ROWWISE(ENUM, STAMP, INIT, STEP, MERGE, FINAL, \ + LANES) \ + static void impl_##STAMP##_masked_rowwise( \ + const matrix_impl *m, const matrix_impl *mask, matrix_impl *vec) { \ + const size_t M = m->rows; \ + const size_t N = m->columns; \ + assert(vec->rows == M && vec->columns == 1); \ + for (size_t r = 0; r < M; ++r) { \ + const double *sp = m->row_ptrs[r]; \ + const double *kp = mask->row_ptrs[r]; \ + double agg = (INIT); \ + size_t cnt = 0; \ + for (size_t c = 0; c < N; ++c) { \ + const int inc = (kp[c] != 0.0); \ + const double value = inc ? sp[c] : (INIT); \ + cnt += (size_t)inc; \ + agg = (STEP); \ + } \ + vec->data[r] = (cnt > 0) ? (FINAL) : (BOC_AGG_MASKED_EMPTY_##ENUM); \ + } \ + } + +#define DEFINE_AGG_MASKED_COLUMNWISE(ENUM, STAMP, INIT, STEP, MERGE, FINAL, \ + LANES) \ + static void impl_##STAMP##_masked_columnwise( \ + const matrix_impl *m, const matrix_impl *mask, matrix_impl *vec) { \ + const size_t M = m->rows; \ + const size_t N = m->columns; \ + assert(vec->rows == 1 && vec->columns == N); \ + for (size_t c = 0; c < N; ++c) { \ + double agg = (INIT); \ + size_t cnt = 0; \ + for (size_t r = 0; r < M; ++r) { \ + const int inc = (mask->row_ptrs[r][c] != 0.0); \ + const double value = inc ? m->row_ptrs[r][c] : (INIT); \ + cnt += (size_t)inc; \ + agg = (STEP); \ + } \ + vec->data[c] = (cnt > 0) ? (FINAL) : (BOC_AGG_MASKED_EMPTY_##ENUM); \ + } \ + } + +#define X(E, S, I, ST, MG, F, L) DEFINE_AGG_MASKED_EWISE(E, S, I, ST, MG, F, L) +BOC_AGG_OPS(X) +#undef X + +#define X(E, S, I, ST, MG, F, L) \ + DEFINE_AGG_MASKED_ROWWISE(E, S, I, ST, MG, F, L) +BOC_AGG_OPS(X) +#undef X + +#define X(E, S, I, ST, MG, F, L) \ + DEFINE_AGG_MASKED_COLUMNWISE(E, S, I, ST, MG, F, L) +BOC_AGG_OPS(X) +#undef X + +static double dispatch_agg_masked_ewise(matrix_impl *m, matrix_impl *mask, + enum AggregateOps op) { + switch (op) { +#define X(ENUM, STAMP, ...) \ + case ENUM: \ + return impl_##STAMP##_masked_ewise(m, mask); + BOC_AGG_OPS(X) +#undef X + default: + fprintf(stderr, "Unknown aggregate op\n"); + return nan(""); + } +} + +static void dispatch_agg_masked_rowwise(matrix_impl *m, matrix_impl *mask, + enum AggregateOps op, + matrix_impl *vec) { + switch (op) { +#define X(ENUM, STAMP, ...) \ + case ENUM: \ + impl_##STAMP##_masked_rowwise(m, mask, vec); \ + return; + BOC_AGG_OPS(X) +#undef X + default: + fprintf(stderr, "Unknown aggregate op\n"); + } +} + +static void dispatch_agg_masked_columnwise(matrix_impl *m, matrix_impl *mask, + enum AggregateOps op, + matrix_impl *vec) { + switch (op) { +#define X(ENUM, STAMP, ...) \ + case ENUM: \ + impl_##STAMP##_masked_columnwise(m, mask, vec); \ + return; + BOC_AGG_OPS(X) +#undef X + default: + fprintf(stderr, "Unknown aggregate op\n"); + } +} + /* -------------------------------------------------------------------------- Unary family X-macro template. @@ -819,7 +973,11 @@ static void dispatch_agg_columnwise(matrix_impl *m, enum AggregateOps op, X(Round, round, nearbyint(v)) \ X(Negate, negate, -v) \ X(Abs, abs, fabs(v)) \ - X(Sqrt, sqrt, sqrt(v)) + X(Sqrt, sqrt, sqrt(v)) \ + X(Reciprocal, reciprocal, 1.0 / v) \ + X(Sign, sign, (double)((v > 0.0) - (v < 0.0))) \ + X(Cos, cos, cos(v)) \ + X(Sin, sin, sin(v)) #define DEFINE_UNARY(ENUM, STAMP, EXPR) \ BOC_CANARY_NOINLINE \ @@ -891,6 +1049,7 @@ typedef struct range_s { Py_ssize_t stop; Py_ssize_t step; size_t count; + bool scalar; ///< True iff the key was a bare integer (collapses to a float). } range; /// @brief This processes the arguments to __get__ to produce the actual @@ -907,10 +1066,12 @@ int range_read(range *range, PyObject *key, size_t length) { } stop = start + 1; step = 1; + range->scalar = true; } else if (PySlice_Check(key)) { if (PySlice_Unpack(key, &start, &stop, &step) < 0) { return -1; } + range->scalar = false; } else { PyErr_SetString(PyExc_TypeError, "Key must be a long or a slice"); return -1; @@ -1072,10 +1233,34 @@ static bool impl_check_acquired(matrix_impl *matrix, bool set_error) { return true; } +/// @brief A read/write, storage-sharing 1-D window into a Matrix. +/// @details Holds a strong reference (one IMPL_INCREF) to the ROOT ``base`` +/// matrix and describes a single row or column span as an affine +/// range over ``base->data``: element ``i`` lives at +/// ``base->data[offset + i * stride]``. ``offset`` and ``stride`` are +/// signed so a negative-step slice (e.g. ``view[::-1]``) round-trips. +/// A view is read/write by default (or read-only when constructed +/// with ``read_only=True``), pins the whole backing buffer for its +/// lifetime, and is never cross-interpreter shareable. +typedef struct boc_matrix_view { + PyObject_HEAD matrix_impl *base; ///< strong IMPL_INCREF to the ROOT storage + Py_ssize_t offset; ///< first element index into base->data + size_t count; ///< element count (>= 1) + Py_ssize_t stride; ///< signed element stride + bool is_row; ///< true => 1 x count, false => count x 1 + bool read_only; ///< true => writes through the view raise +} MatrixViewObject; + typedef struct { int_least64_t interpid; PyTypeObject *matrix_type; + PyTypeObject *values_iter_type; + PyTypeObject *view_type; + PyTypeObject *view_values_iter_type; + PyTypeObject *view_indexer_type; + PyTypeObject *views_iter_type; PyObject *matrix_unpickle; + uint64_t prng_state; ///< Per-interpreter PRNG state (splitmix64). } _math_module_state; static thread_local _math_module_state *LOCAL_STATE; @@ -1085,6 +1270,62 @@ static thread_local _math_module_state *LOCAL_STATE; LOCAL_STATE = (_math_module_state *)PyModule_GetState(m); \ } while (0) +/// @brief Resolve this thread's @ref _math_module_state, warming the cache. +/// @details Fast path returns the thread-local set at module exec. A handful +/// of paths (Matrix pickle reduce/unpickle and the XIData +/// reconstructor) can also be reached from a *second* thread in the +/// same interpreter that never ran module init -- concretely the +/// primary interpreter's noticeboard mutator thread, which +/// reconstructs Matrix-valued noticeboard entries during +/// @c notice_update snapshots. On that cold thread @ref LOCAL_STATE +/// is NULL, so we resolve *this interpreter's* already-imported +/// module from @c sys.modules and read its state. Module state is +/// per-interpreter, so this is correct on any interpreter without a +/// published global. @c PyImport_GetModule is a plain @c sys.modules +/// lookup: it returns the module @ref _math_module_exec already +/// created and never re-runs init, so no duplicate state is built. +/// +/// The lookup is deliberately *not* cached in @ref LOCAL_STATE: it +/// runs only on the cold noticeboard path, while the fast path is +/// reserved for the threads that ran module init. +/// +/// Invariant relied on by the many hot paths that dereference +/// @ref LOCAL_STATE *directly* (the number-protocol ops, +/// @c unwrap_matrix, @c use_out_target, @c unwrap_mask, the values +/// iterator, ...): those run only on threads that executed module +/// init, so @ref LOCAL_STATE is warm. The one cold thread that +/// reaches @c _math (the noticeboard mutator) enters solely through +/// @ref Matrix_reduce, which routes through this helper and never +/// chains a @ref LOCAL_STATE -direct call. If a future noticeboard +/// path executes a Matrix *method* on that thread, re-audit: it must +/// go through @ref math_local_state, not a bare @ref LOCAL_STATE. +/// @return the module state, or NULL with a RuntimeError set if the module is +/// no longer present in @c sys.modules for the current interpreter +/// (does not happen in bocpy's threading model, where the noticeboard +/// thread is stopped before module teardown; guarded defensively +/// rather than crashing). +static _math_module_state *math_local_state(void) { + if (LOCAL_STATE != NULL) { + return LOCAL_STATE; + } + PyObject *name = PyUnicode_FromString("bocpy._math"); + if (name == NULL) { + return NULL; + } + PyObject *module = PyImport_GetModule(name); + Py_DECREF(name); + if (module == NULL) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_RuntimeError, + "bocpy._math module state is unavailable on this thread"); + } + return NULL; + } + _math_module_state *state = (_math_module_state *)PyModule_GetState(module); + Py_DECREF(module); + return state; +} + typedef struct matrix_object { PyObject_HEAD matrix_impl *impl; } MatrixObject; @@ -1208,7 +1449,12 @@ PyObject *wrap_matrix(PyTypeObject *type, matrix_impl *impl) { } static PyObject *wrap_impl_or_free(matrix_impl *impl) { - PyObject *matrix = wrap_matrix(LOCAL_STATE->matrix_type, impl); + _math_module_state *state = math_local_state(); + if (state == NULL) { + impl_free(impl); + return NULL; + } + PyObject *matrix = wrap_matrix(state->matrix_type, impl); if (matrix == NULL) { impl_free(impl); } @@ -1216,6 +1462,81 @@ static PyObject *wrap_impl_or_free(matrix_impl *impl) { return matrix; } +// --------------------------------------------------------------------------- +// MatrixView: a read/write, storage-sharing 1-D window (row or column) into a +// Matrix. A view holds one strong reference to the ROOT matrix_impl and reads +// through base->data[offset + i*stride]; it is read/write, pins the whole +// backing buffer for its lifetime, and is never cross-interpreter shareable. +// +// view_new and view_materialize sit ahead of unwrap_matrix (which materialises +// a view operand); the rest of the MatrixView surface follows below. +// --------------------------------------------------------------------------- + +/// @brief The single guarded constructor for every MatrixView. +/// @details Rejects an empty span (``count == 0``) and validates that both +/// element endpoints (``offset`` and ``offset + (count-1)*stride``) +/// fall inside ``[0, base->size)`` -- the load-bearing guard against +/// an out-of-bounds read. Every caller (row_view / column_view / +/// sub-view / the m.view[...] accessor) builds the span from an +/// in-bounds range and leaves ``stride`` unused when ``count == 1``, +/// so ``(count-1)*stride`` stays within the buffer and cannot +/// overflow. Takes exactly one ``IMPL_INCREF`` on the root ``base``. +/// ``offset``/``stride`` are signed so negative-step slices work. +/// When ``read_only`` is set, writes through the view (item / slice +/// assignment and the ``.x``/``.y``/``.z``/``.w`` setters) raise +/// ``TypeError``; derived views (``.T``, sub-slices) inherit it. +static PyObject *view_new(matrix_impl *base, Py_ssize_t offset, size_t count, + Py_ssize_t stride, bool is_row, bool read_only) { + if (count == 0) { + PyErr_SetString(PyExc_IndexError, "cannot create an empty view"); + return NULL; + } + Py_ssize_t last = offset + (Py_ssize_t)(count - 1) * stride; + Py_ssize_t size = (Py_ssize_t)base->size; + if (offset < 0 || offset >= size || last < 0 || last >= size) { + PyErr_SetString(PyExc_IndexError, "view span out of range"); + return NULL; + } + + PyTypeObject *type = LOCAL_STATE->view_type; + MatrixViewObject *view = (MatrixViewObject *)type->tp_alloc(type, 0); + if (view == NULL) { + return NULL; + } + view->base = base; + view->offset = offset; + view->count = count; + view->stride = stride; + view->is_row = is_row; + view->read_only = read_only; + IMPL_INCREF(base); + return (PyObject *)view; +} + +/// @brief Materialise a view into a fresh, contiguous, owned matrix_impl. +/// @details Guards ownership, then gathers ``count`` strided elements into a +/// new ``1 x count`` (row) or ``count x 1`` (column) buffer. Returns +/// the impl at refcount 0 (the caller wraps or INCREFs it), or NULL +/// with an exception set. +static matrix_impl *view_materialize(MatrixViewObject *v) { + if (!impl_check_acquired(v->base, true)) { + return NULL; + } + size_t rows = v->is_row ? 1 : v->count; + size_t cols = v->is_row ? v->count : 1; + matrix_impl *out = impl_new(rows, cols); + if (out == NULL) { + return NULL; + } + const double *src = v->base->data; + double *dst = out->data; + Py_ssize_t idx = v->offset; + for (size_t i = 0; i < v->count; ++i, idx += v->stride) { + dst[i] = src[idx]; + } + return out; +} + matrix_impl *unwrap_matrix(PyObject *op, bool seq_as_column) { PyTypeObject *type = LOCAL_STATE->matrix_type; MatrixObject *matrix; @@ -1232,6 +1553,18 @@ matrix_impl *unwrap_matrix(PyObject *op, bool seq_as_column) { return impl; } + // A MatrixView materialises to a fresh contiguous impl so every arithmetic + // kernel keeps its unit-stride fast path; return it INCREF'd to match the + // Matrix branch's caller-DECREF contract. + if (Py_TYPE(op) == LOCAL_STATE->view_type) { + matrix_impl *out = view_materialize((MatrixViewObject *)op); + if (out == NULL) { + return NULL; + } + IMPL_INCREF(out); + return out; + } + impl = impl_new_from_sequence(op, seq_as_column); if (impl == NULL) { return NULL; @@ -1241,1348 +1574,2353 @@ matrix_impl *unwrap_matrix(PyObject *op, bool seq_as_column) { return impl; } -static PyObject *Matrix_transpose(PyObject *op, PyObject *args, - PyObject *kwds) { - MatrixObject *matrix = (MatrixObject *)op; - matrix_impl *impl = matrix->impl; +// MatrixView value-access paths (using the view_new / view_materialize helpers +// above). Invariant: every path that reads or writes base storage MUST call +// impl_check_acquired(base, true) before touching data. The pure-metadata +// getters is_row and .T are exempt (no data touch) -- a released view still has +// a valid orientation. A new getter/setter that reads or writes data without +// the guard is a bug. - int in_place = 0; - static char *kwlist[] = {"in_place", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|p", kwlist, &in_place)) { +static PyObject *MatrixView_get_rows(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return NULL; } + return PyLong_FromSize_t(v->is_row ? 1 : v->count); +} - if (!impl_check_acquired(impl, true)) { +static PyObject *MatrixView_get_columns(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return NULL; } + return PyLong_FromSize_t(v->is_row ? v->count : 1); +} - if (in_place) { - if (impl_transpose_in_place(impl) < 0) { - return NULL; - } - return Py_NewRef(op); - } - - matrix_impl *transpose = impl_transpose(impl); - if (transpose == NULL) { +static PyObject *MatrixView_get_size(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return NULL; } - - return wrap_impl_or_free(transpose); + return PyLong_FromSize_t(v->count); } -/// @brief Validate a caller-supplied ``out=`` target and alias its buffer. -/// @details Implements the numpy-style ``out=`` convention for the -/// allocation-free elementwise paths: the result is written into an -/// existing matrix the caller owns rather than a fresh allocation. -/// The target must be a Matrix owned by the current interpreter whose -/// shape exactly matches the operation's result shape; the shape is -/// validated here, before any kernel writes, so a rejected target -/// leaves every buffer untouched (the validate-then-write contract -/// shared with ``put``). The target may alias the primary -/// same-shape operand because the kernel reads and writes that -/// operand at the same index; aliasing a lower-rank broadcast -/// operand is impossible because the result-shape check rejects any -/// target whose shape differs from the result. -/// @param out_target The ``out=`` Matrix object (never NULL / Py_None here). -/// @param out_op Receives a new reference to ``out_target`` on success. -/// @param want_rows Required row count of the result. -/// @param want_cols Required column count of the result. -/// @return The target's impl (no extra C refcount taken — the live Python -/// reference owns it), or NULL with an exception set on failure. -static matrix_impl *use_out_target(PyObject *out_target, PyObject **out_op, - size_t want_rows, size_t want_cols) { - if (Py_TYPE(out_target) != LOCAL_STATE->matrix_type) { - PyErr_SetString(PyExc_TypeError, "out must be a Matrix"); - return NULL; - } - matrix_impl *impl = ((MatrixObject *)out_target)->impl; - if (!impl_check_acquired(impl, true)) { - return NULL; - } - if (impl->rows != want_rows || impl->columns != want_cols) { - PyErr_Format(PyExc_ValueError, - "out shape %zux%zu does not match result %zux%zu", impl->rows, - impl->columns, want_rows, want_cols); +static PyObject *MatrixView_get_shape(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return NULL; } - *out_op = Py_NewRef(out_target); - return impl; + Py_ssize_t rows = (Py_ssize_t)(v->is_row ? 1 : v->count); + Py_ssize_t cols = (Py_ssize_t)(v->is_row ? v->count : 1); + return Py_BuildValue("(nn)", rows, cols); } -/// @brief Sets the output of an arithmetic operation. -/// @details The output of an arithmetic operation is one of three things: a -/// caller-supplied ``out=`` target (when ``out_target`` is a non-None Matrix), -/// the left-hand side itself (in-place operations), or a freshly allocated -/// matrix of the same dimensions as the left-hand side. ``out_target`` and -/// ``inplace`` are mutually exclusive (asserted here; callers reject the -/// combination before reaching this point). -/// -/// The result shape and wrap type come from the already-unwrapped ``lhs`` -/// impl and the interpreter's canonical Matrix type -- never by casting -/// ``lhs_op``. ``lhs_op`` may be a sequence that ``unwrap_matrix`` coerced -/// into ``lhs`` (e.g. ``[1, 2, 3] + matrix``), so it is not guaranteed to be -/// a Matrix; it is dereferenced only as the in-place return value, a path -/// reachable solely with a genuine Matrix left operand. -/// @param lhs The unwrapped impl of the left-hand operand (gives result shape) -/// @param lhs_op The PyObject of the left-hand operand (in-place return only) -/// @param out_op A pointer to the output pointer of the equation -/// @param out_target Optional ``out=`` Matrix (NULL or Py_None when unused) -/// @param inplace Whether this is an inplace operation -/// @return The matrix wrapped by out_op, or NULL in the case of an error -static matrix_impl *set_output(matrix_impl *lhs, PyObject *lhs_op, - PyObject **out_op, PyObject *out_target, - bool inplace) { - matrix_impl *out; - if (out_target != NULL && out_target != Py_None) { - assert(!inplace); - return use_out_target(out_target, out_op, lhs->rows, lhs->columns); +// is_row is pure orientation metadata; no ownership guard (a released view +// still has a well-defined orientation). +static PyObject *MatrixView_get_is_row(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (v->is_row) { + Py_RETURN_TRUE; } - if (inplace) { - *out_op = Py_NewRef(lhs_op); - return lhs; + Py_RETURN_FALSE; +} + +// read_only is pure metadata (no data touch); no ownership guard, matching +// is_row / .T. +static PyObject *MatrixView_get_read_only(PyObject *op, + void *Py_UNUSED(dummy)) { + if (((MatrixViewObject *)op)->read_only) { + Py_RETURN_TRUE; } + Py_RETURN_FALSE; +} - out = impl_new(lhs->rows, lhs->columns); - if (out == NULL) { +// .T is a free (O(1), no-copy) transpose: a new view over the same storage +// with is_row flipped. No data guard -- view_new only reads base->size. +static PyObject *MatrixView_get_T(PyObject *op, void *Py_UNUSED(dummy)) { + MatrixViewObject *v = (MatrixViewObject *)op; + return view_new(v->base, v->offset, v->count, v->stride, !v->is_row, + v->read_only); +} + +// Shared .x/.y/.z/.w component access, matching Matrix's convenience +// accessors: component k is the view's k-th element, at +// base->data[offset + k*stride]. Read/write like the rest of the view: the +// getter guards ownership and the setter writes through to the base. Absent +// components (k >= count) raise IndexError, exactly as on a Matrix. +static PyObject *view_get_component(PyObject *op, size_t k) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return NULL; } - - *out_op = wrap_matrix(LOCAL_STATE->matrix_type, out); - if (*out_op == NULL) { - impl_free(out); + if (v->count <= k) { + PyErr_SetNone(PyExc_IndexError); return NULL; } - - return out; + return PyFloat_FromDouble( + v->base->data[v->offset + (Py_ssize_t)k * v->stride]); } -/// @brief Tri-state representation of an optional ``axis`` kwarg. -/// @details ``has_axis`` is true iff the caller passed a non-None axis; -/// ``axis`` is the validated int value (and is undefined when -/// ``has_axis`` is false). Replaces the historical NO_AXIS=-1000 -/// sentinel which collided with the integer -1000. -typedef struct { - bool has_axis; - int axis; -} AxisArg; - -/// @brief Decode an optional ``axis`` keyword argument into an AxisArg. -/// @details Accepts ``NULL`` or ``Py_None`` (no axis), an ``int`` in -/// ``INT_MIN..INT_MAX``. Rejects ``bool`` (subclass of int) and -/// overflows. Returns 0 on success and writes through ``*out``; -/// returns -1 with TypeError / OverflowError set on failure. -static int decode_axis_kwarg(PyObject *axis_obj, AxisArg *out) { - if (axis_obj == NULL || axis_obj == Py_None) { - out->has_axis = false; - out->axis = 0; - return 0; - } - if (PyBool_Check(axis_obj)) { - PyErr_SetString(PyExc_TypeError, "axis must be an int or None, not bool"); +static int view_set_component(PyObject *op, PyObject *arg, size_t k) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return -1; } - if (!PyLong_Check(axis_obj)) { - PyErr_SetString(PyExc_TypeError, "axis must be an int or None"); + if (v->read_only) { + PyErr_SetString(PyExc_TypeError, "cannot write to a read-only MatrixView"); return -1; } - long ax = PyLong_AsLong(axis_obj); - if (ax == -1 && PyErr_Occurred()) { + if (v->count <= k) { + PyErr_SetNone(PyExc_IndexError); return -1; } - if (ax < INT_MIN || ax > INT_MAX) { - PyErr_Format(PyExc_OverflowError, "axis %ld out of int range", ax); + // unwrap_double reports a non-number by returning false without setting an + // exception (it is also used as a type probe elsewhere), so name the error + // here rather than returning -1 with no exception pending. + if (!unwrap_double(arg, + &v->base->data[v->offset + (Py_ssize_t)k * v->stride])) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "expected a real number"); + } return -1; } - out->has_axis = true; - out->axis = (int)ax; return 0; } -/// @brief Decode, validate (-2/-1/0/1) and normalise (-1->1, -2->0) axis. -/// @details Single entry point for the cross / perpendicular / angle / -/// normalize / aggregate methods that accept only the four -/// standard axis values. After this call, ``out->axis`` is -/// either 0 (column-wise) or 1 (row-wise) when ``has_axis`` is -/// true. Returns -1 with the appropriate exception set on any -/// decode or range error. -static int parse_validate_normalise_axis(PyObject *axis_obj, AxisArg *out) { - if (decode_axis_kwarg(axis_obj, out) < 0) { - return -1; +static PyObject *MatrixView_get_x(PyObject *op, void *Py_UNUSED(dummy)) { + return view_get_component(op, 0); +} +static int MatrixView_set_x(PyObject *op, PyObject *arg, void *Py_UNUSED(d)) { + return view_set_component(op, arg, 0); +} +static PyObject *MatrixView_get_y(PyObject *op, void *Py_UNUSED(dummy)) { + return view_get_component(op, 1); +} +static int MatrixView_set_y(PyObject *op, PyObject *arg, void *Py_UNUSED(d)) { + return view_set_component(op, arg, 1); +} +static PyObject *MatrixView_get_z(PyObject *op, void *Py_UNUSED(dummy)) { + return view_get_component(op, 2); +} +static int MatrixView_set_z(PyObject *op, PyObject *arg, void *Py_UNUSED(d)) { + return view_set_component(op, arg, 2); +} +static PyObject *MatrixView_get_w(PyObject *op, void *Py_UNUSED(dummy)) { + return view_get_component(op, 3); +} +static int MatrixView_set_w(PyObject *op, PyObject *arg, void *Py_UNUSED(d)) { + return view_set_component(op, arg, 3); +} + +static Py_ssize_t MatrixView_length(PyObject *op) { + return (Py_ssize_t)((MatrixViewObject *)op)->count; +} + +// A view is 1-D: an int key yields a float, a slice yields a sub-view (composed +// against the root), a tuple / anything else is rejected. Sub-view offset and +// stride are computed with overflow-safe arithmetic (stride may be negative). +static PyObject *MatrixView_subscript(PyObject *op, PyObject *key) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { + return NULL; } - if (!out->has_axis) { - return 0; + if (PyLong_Check(key)) { + Py_ssize_t i = PyLong_AsSsize_t(key); + if (i == -1 && PyErr_Occurred()) { + return NULL; + } + if (i < 0) { + i += (Py_ssize_t)v->count; + } + if (i < 0 || (size_t)i >= v->count) { + PyErr_SetString(PyExc_IndexError, "view index out of range"); + return NULL; + } + return PyFloat_FromDouble(v->base->data[v->offset + i * v->stride]); } - int ax = out->axis; - if (ax != 0 && ax != 1 && ax != -1 && ax != -2) { - PyErr_SetString(PyExc_ValueError, "axis must be -2, -1, 0, or 1"); + if (PySlice_Check(key)) { + range r; + if (range_read(&r, key, v->count) < 0) { + return NULL; + } + // The stride only matters for a multi-element sub-view; for count <= 1 it + // is never dereferenced, so avoid the stride*step multiply that a giant + // slice step could overflow. r.start*stride reaches a valid parent element + // and stays within the buffer, so it cannot overflow. + Py_ssize_t new_stride = r.count > 1 ? v->stride * r.step : v->stride; + Py_ssize_t new_offset = v->offset + r.start * v->stride; + return view_new(v->base, new_offset, r.count, new_stride, v->is_row, + v->read_only); + } + PyErr_SetString(PyExc_TypeError, + "MatrixView is 1-D; index with a single int or slice"); + return NULL; +} + +// Item / slice assignment writes THROUGH the shared storage of a read/write +// view (a read-only view rejects the write with TypeError up front). An int +// key writes one scalar; a slice key broadcasts a +// scalar or copies a length-matched sequence / Matrix / view. Matching is +// length-only -- a view is 1-D, so RHS orientation is irrelevant. A view cannot +// resize, so deletion (value == NULL) is rejected. The ownership guard is the +// same one the reads use, so writing an unowned view raises RuntimeError +// symmetrically. The step*stride multiply is guarded behind count > 1 (unused +// for a single element), matching the read sub-view path. +static int MatrixView_ass_subscript(PyObject *op, PyObject *key, + PyObject *value) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { return -1; } - if (ax == -1) { - out->axis = 1; - } else if (ax == -2) { - out->axis = 0; + if (v->read_only) { + PyErr_SetString(PyExc_TypeError, "cannot write to a read-only MatrixView"); + return -1; } - return 0; -} - -static int Matrix_aggregate(PyObject *matrix_op, AxisArg axis, - PyObject **out_op, enum AggregateOps agg) { - MatrixObject *matrix = (MatrixObject *)matrix_op; - matrix_impl *impl = matrix->impl; - - if (!impl_check_acquired(impl, true)) { + if (value == NULL) { + PyErr_SetString(PyExc_TypeError, "cannot delete elements of a MatrixView"); return -1; } - if (!axis.has_axis) { - *out_op = PyFloat_FromDouble(dispatch_agg_ewise(impl, agg)); + if (PyLong_Check(key)) { + Py_ssize_t i = PyLong_AsSsize_t(key); + if (i == -1 && PyErr_Occurred()) { + return -1; + } + if (i < 0) { + i += (Py_ssize_t)v->count; + } + if (i < 0 || (size_t)i >= v->count) { + PyErr_SetString(PyExc_IndexError, "view index out of range"); + return -1; + } + double scalar = PyFloat_AsDouble(value); + if (scalar == -1.0 && PyErr_Occurred()) { + return -1; + } + v->base->data[v->offset + i * v->stride] = scalar; return 0; } - if (axis.axis == 0) { - matrix_impl *vector = impl_new(1, impl->columns); - if (vector == NULL) { + if (PySlice_Check(key)) { + range r; + if (range_read(&r, key, v->count) < 0) { return -1; } + Py_ssize_t step = r.count > 1 ? v->stride * r.step : v->stride; + Py_ssize_t start = v->offset + r.start * v->stride; - dispatch_agg_columnwise(impl, agg, vector); - *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); - if (*out_op == NULL) { - impl_free(vector); - return -1; + // A scalar RHS broadcasts to every selected element. + if (PyLong_Check(value) || PyFloat_Check(value)) { + double scalar = PyFloat_AsDouble(value); + if (scalar == -1.0 && PyErr_Occurred()) { + return -1; + } + Py_ssize_t idx = start; + for (size_t k = 0; k < r.count; ++k, idx += step) { + v->base->data[idx] = scalar; + } + return 0; } + // Otherwise the RHS is a sequence / Matrix / view: length-only match. + matrix_impl *rhs = unwrap_matrix(value, false); + if (rhs == NULL) { + return -1; + } + if (rhs->size != r.count) { + PyErr_Format(PyExc_ValueError, + "cannot assign %zu values to a view of length %zu", + rhs->size, r.count); + IMPL_DECREF(rhs); + return -1; + } + const double *src = rhs->data; + Py_ssize_t idx = start; + for (size_t k = 0; k < r.count; ++k, idx += step) { + v->base->data[idx] = src[k]; + } + IMPL_DECREF(rhs); return 0; } - // Fall-through is axis == 1 (row-wise): parse_validate_normalise_axis - // restricts axis to {0, 1}. - matrix_impl *vector = impl_new(impl->rows, 1); - if (vector == NULL) { - return -1; - } + PyErr_SetString(PyExc_TypeError, + "MatrixView is 1-D; index with a single int or slice"); + return -1; +} - dispatch_agg_rowwise(impl, agg, vector); - *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); - if (*out_op == NULL) { - impl_free(vector); - return -1; +// repr/str delegate to a materialised temporary Matrix (via runtime str()), +// prefixed so a human sees it is a borrow. Degrade (never raise) on a released +// base, matching Matrix's "" contract. +static PyObject *view_render(PyObject *op) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, false)) { + return PyUnicode_FromString("MatrixView()"); + } + matrix_impl *out = view_materialize(v); + if (out == NULL) { + return NULL; + } + PyObject *tmp = wrap_impl_or_free(out); + if (tmp == NULL) { + return NULL; + } + PyObject *inner = PyObject_Str(tmp); + Py_DECREF(tmp); + if (inner == NULL) { + return NULL; } + PyObject *res = PyUnicode_FromFormat("MatrixView(%U)", inner); + Py_DECREF(inner); + return res; +} - return 0; +static PyObject *MatrixView_repr(PyObject *op) { return view_render(op); } +static PyObject *MatrixView_str(PyObject *op) { return view_render(op); } + +// Lazy strided iterator returned by MatrixView.values(): a strong ref to the +// view plus a cursor. Re-checks ownership every step (the base may be released +// between yields), mirroring MatrixValuesIter. +typedef struct { + PyObject_HEAD MatrixViewObject *view; + size_t cursor; +} MatrixViewValuesIterObject; + +static void MatrixViewValuesIter_dealloc(PyObject *op) { + MatrixViewValuesIterObject *it = (MatrixViewValuesIterObject *)op; + PyTypeObject *type = Py_TYPE(op); + Py_XDECREF(it->view); + type->tp_free(op); + Py_DECREF(type); } -#define MATRIX_AGGREGATE(agg) \ - static PyObject *Matrix_##agg##_method(PyObject *op, PyObject *args, \ - PyObject *kwds) { \ - PyObject *out = NULL; \ - PyObject *axis_obj = NULL; \ - static char *kwlist[] = {"axis", NULL}; \ - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &axis_obj)) { \ - return NULL; \ - } \ - AxisArg axis; \ - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { \ - return NULL; \ - } \ - if (Matrix_aggregate(op, axis, &out, agg) < 0) { \ - return NULL; \ - } \ - return out; \ +static PyObject *MatrixViewValuesIter_next(PyObject *op) { + MatrixViewValuesIterObject *it = (MatrixViewValuesIterObject *)op; + MatrixViewObject *v = it->view; + if (!impl_check_acquired(v->base, true)) { + return NULL; } + if (it->cursor >= v->count) { + return NULL; // StopIteration + } + double value = v->base->data[v->offset + (Py_ssize_t)it->cursor * v->stride]; + it->cursor++; + return PyFloat_FromDouble(value); +} -MATRIX_AGGREGATE(Sum) -MATRIX_AGGREGATE(Mean) -MATRIX_AGGREGATE(Magnitude) -MATRIX_AGGREGATE(Minimum) -MATRIX_AGGREGATE(Maximum) -MATRIX_AGGREGATE(MagnitudeSquared) +static PyObject *MatrixViewValuesIter_self(PyObject *op) { + Py_INCREF(op); + return op; +} -/* -------------------------------------------------------------------------- - Arg-reduction (argmin / argmax) kernels. +// A view iterator is only ever born through MatrixView.values(); direct +// construction would leave view == NULL and NULL-deref on first step. +static PyObject *MatrixViewValuesIter_new(PyTypeObject *Py_UNUSED(type), + PyObject *Py_UNUSED(args), + PyObject *Py_UNUSED(kwds)) { + PyErr_SetString(PyExc_TypeError, + "cannot create 'bocpy._math.MatrixViewValuesIter' instances " + "directly; use MatrixView.values()"); + return NULL; +} - These do not fit the BOC_AGG_OPS X-macro: that table accumulates a - single double, whereas an arg-reduction must also carry the index of - the running extreme. Comparisons are strict so the first occurrence of - a tied extreme wins, matching NumPy. Indices are published as doubles - in the result matrix (the Matrix type stores only doubles). - -------------------------------------------------------------------------- */ -// Arg-reduction (argmin/argmax) kernels: kept out of the BOC_AGG_OPS X-macro -// because they must carry the running index, not just a double accumulator. -// Strict comparisons make the first tied extreme win (NumPy tie-break); -// indices are published as doubles (Matrix stores only doubles). -static Py_ssize_t argextreme_ewise(matrix_impl *m, bool want_max) { - const double *p = m->data; - double best = p[0]; - Py_ssize_t best_i = 0; - for (size_t i = 1; i < m->size; ++i) { - const double v = p[i]; - if (want_max ? (v > best) : (v < best)) { - best = v; - best_i = (Py_ssize_t)i; - } +static PyType_Slot MatrixViewValuesIter_slots[] = { + {Py_tp_new, MatrixViewValuesIter_new}, + {Py_tp_dealloc, MatrixViewValuesIter_dealloc}, + {Py_tp_iter, MatrixViewValuesIter_self}, + {Py_tp_iternext, MatrixViewValuesIter_next}, + {0, NULL}, +}; + +static PyType_Spec MatrixViewValuesIter_Spec = { + .name = "bocpy._math.MatrixViewValuesIter", + .basicsize = sizeof(MatrixViewValuesIterObject), + .itemsize = 0, + .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE, + .slots = MatrixViewValuesIter_slots}; + +static PyObject *MatrixView_values(PyObject *op, PyObject *Py_UNUSED(ignored)) { + MatrixViewObject *v = (MatrixViewObject *)op; + if (!impl_check_acquired(v->base, true)) { + return NULL; } - return best_i; + PyTypeObject *type = LOCAL_STATE->view_values_iter_type; + MatrixViewValuesIterObject *it = + (MatrixViewValuesIterObject *)type->tp_alloc(type, 0); + if (it == NULL) { + return NULL; + } + Py_INCREF(op); + it->view = v; + it->cursor = 0; + return (PyObject *)it; } -static void argextreme_columnwise(matrix_impl *m, bool want_max, - matrix_impl *out) { - const size_t M = m->rows; - const size_t N = m->columns; - for (size_t c = 0; c < N; ++c) { - double best = m->data[c]; - size_t best_r = 0; - for (size_t r = 1; r < M; ++r) { - const double v = m->data[r * N + c]; - if (want_max ? (v > best) : (v < best)) { - best = v; - best_r = r; - } - } - out->data[c] = (double)best_r; - } +// A view is 1-D, so iterating it yields its elements (like a 1-D array), not +// sub-views: `for x in view` is exactly `for x in view.values()`. Delegates to +// MatrixView_values so both share one lazy, ownership-re-checking iterator. +static PyObject *MatrixView_iter(PyObject *op) { + return MatrixView_values(op, NULL); } -static void argextreme_rowwise(matrix_impl *m, bool want_max, - matrix_impl *out) { - const size_t M = m->rows; - const size_t N = m->columns; - for (size_t r = 0; r < M; ++r) { - const double *row = m->data + r * N; - double best = row[0]; - size_t best_c = 0; - for (size_t c = 1; c < N; ++c) { - const double v = row[c]; - if (want_max ? (v > best) : (v < best)) { - best = v; - best_c = c; - } - } - out->data[r] = (double)best_c; +// Materialise the view into an independent, owned Matrix (canonical type). The +// result shares no storage with the base, so mutating either leaves the other +// unchanged -- the escape hatch for shipping a view across interpreters. +static PyObject *MatrixView_copy(PyObject *op, PyObject *Py_UNUSED(ignored)) { + matrix_impl *out = view_materialize((MatrixViewObject *)op); + if (out == NULL) { + return NULL; } + return (PyObject *)wrap_impl_or_free(out); } -static int Matrix_argextreme(PyObject *matrix_op, AxisArg axis, - PyObject **out_op, bool want_max) { - MatrixObject *matrix = (MatrixObject *)matrix_op; - matrix_impl *impl = matrix->impl; +// A view aliases its base's storage and pins the whole backing buffer, so it +// is deliberately not shippable. Refuse pickling eagerly (at dumps time) with +// a message that names the .copy() remedy -- overriding __reduce__ takes +// precedence over object.__reduce_ex__, so this fires for pickle.dumps and for +// every bocpy path that serialises (Cown payloads, send()). +static PyObject *MatrixView_reduce(PyObject *Py_UNUSED(op), + PyObject *Py_UNUSED(ignored)) { + PyErr_SetString(PyExc_TypeError, + "cannot pickle or send a MatrixView across interpreters; " + "call .copy() to get an owned Matrix"); + return NULL; +} - if (!impl_check_acquired(impl, true)) { - return -1; +static PyMethodDef MatrixView_methods[] = { + {"values", MatrixView_values, METH_NOARGS, + "values($self, /)\n--\n\n" + "Return a lazy iterator over the view's elements."}, + {"copy", MatrixView_copy, METH_NOARGS, + "copy($self, /)\n--\n\n" + "Return an independent Matrix holding a copy of the view's elements."}, + {"__reduce__", MatrixView_reduce, METH_NOARGS, + "__reduce__($self, /)\n--\n\n" + "Refuse pickling; a view aliases its base and cannot be shipped. " + "Call .copy() to get an owned, shippable Matrix."}, + {NULL} /* Sentinel */ +}; + +static PyGetSetDef MatrixView_getset[] = { + {"rows", (getter)MatrixView_get_rows, NULL, NULL, NULL}, + {"columns", (getter)MatrixView_get_columns, NULL, NULL, NULL}, + {"size", (getter)MatrixView_get_size, NULL, NULL, NULL}, + {"shape", (getter)MatrixView_get_shape, NULL, NULL, NULL}, + {"is_row", (getter)MatrixView_get_is_row, NULL, NULL, NULL}, + {"read_only", (getter)MatrixView_get_read_only, NULL, NULL, NULL}, + {"T", (getter)MatrixView_get_T, NULL, NULL, NULL}, + {"x", (getter)MatrixView_get_x, (setter)MatrixView_set_x, NULL, NULL}, + {"y", (getter)MatrixView_get_y, (setter)MatrixView_set_y, NULL, NULL}, + {"z", (getter)MatrixView_get_z, (setter)MatrixView_set_z, NULL, NULL}, + {"w", (getter)MatrixView_get_w, (setter)MatrixView_set_w, NULL, NULL}, + {NULL} /* Sentinel */ +}; + +// A MatrixView is only ever born through view_new (via Matrix.row_view / +// column_view / view[...]); direct construction (and pickle __newobj__ +// reconstruction) would leave base == NULL and NULL-deref on first use. +static PyObject *MatrixView_new(PyTypeObject *Py_UNUSED(type), + PyObject *Py_UNUSED(args), + PyObject *Py_UNUSED(kwds)) { + PyErr_SetString(PyExc_TypeError, + "cannot create 'bocpy._math.MatrixView' instances directly; " + "use Matrix.row_view() / Matrix.column_view()"); + return NULL; +} + +static void MatrixView_dealloc(PyObject *op) { + MatrixViewObject *self = (MatrixViewObject *)op; + PyTypeObject *type = Py_TYPE(op); + IMPL_DECREF(self->base); + type->tp_free(op); + Py_DECREF(type); +} + +// A view participates in arithmetic by materialising at the unwrap_matrix +// boundary, so the binary and matmul number slots reuse Matrix's own +// operand-symmetric kernels verbatim (they never dereference a left operand's +// impl -- both operands go through unwrap_matrix). Forward-declared here; the +// macro-stamped definitions live further down with the Matrix number protocol. +static PyObject *Matrix_Add_op(PyObject *lhs, PyObject *rhs); +static PyObject *Matrix_Subtract_op(PyObject *lhs, PyObject *rhs); +static PyObject *Matrix_Multiply_op(PyObject *lhs, PyObject *rhs); +static PyObject *Matrix_Divide_op(PyObject *lhs, PyObject *rhs); +static PyObject *Matrix_matmul(PyObject *lhs_op, PyObject *rhs_op); + +// Unary ops cannot reuse Matrix_Negate_op / Matrix_Abs_op: those blind-cast +// their operand to MatrixObject and read self->impl. Instead materialise the +// view to an owned Matrix and delegate to its number protocol. +static PyObject *view_unary(PyObject *op, PyObject *(*fn)(PyObject *)) { + matrix_impl *out = view_materialize((MatrixViewObject *)op); + if (out == NULL) { + return NULL; + } + PyObject *tmp = wrap_impl_or_free(out); + if (tmp == NULL) { + return NULL; } + PyObject *result = fn(tmp); + Py_DECREF(tmp); + return result; +} - // Defensive: the public constructors reject zero-size matrices, but guard - // each axis so a future empty-capable path can't read p[0] out of bounds. - const char *empty_error = "arg-reduction of an empty matrix is undefined"; +static PyObject *MatrixView_negative(PyObject *op) { + return view_unary(op, PyNumber_Negative); +} - if (!axis.has_axis) { - if (impl->size == 0) { - PyErr_SetString(PyExc_ValueError, empty_error); - return -1; - } - *out_op = PyLong_FromSsize_t(argextreme_ewise(impl, want_max)); - return *out_op == NULL ? -1 : 0; +static PyObject *MatrixView_absolute(PyObject *op) { + return view_unary(op, PyNumber_Absolute); +} + +static PyType_Slot MatrixView_slots[] = { + {Py_tp_new, MatrixView_new}, + {Py_tp_dealloc, MatrixView_dealloc}, + {Py_tp_getset, MatrixView_getset}, + {Py_tp_methods, MatrixView_methods}, + {Py_mp_subscript, MatrixView_subscript}, + {Py_mp_ass_subscript, MatrixView_ass_subscript}, + {Py_mp_length, MatrixView_length}, + {Py_sq_length, MatrixView_length}, + {Py_tp_iter, MatrixView_iter}, + {Py_tp_repr, MatrixView_repr}, + {Py_tp_str, MatrixView_str}, + // Binary / matmul reuse Matrix's operand-symmetric kernels; the view is + // materialised at unwrap_matrix. Item/slice assignment writes through + // (Py_mp_ass_subscript); there are deliberately no Py_nb_inplace_* slots, + // so ``view += x`` rebinds the name to a fresh Matrix rather than mutating. + {Py_nb_add, Matrix_Add_op}, + {Py_nb_subtract, Matrix_Subtract_op}, + {Py_nb_multiply, Matrix_Multiply_op}, + {Py_nb_true_divide, Matrix_Divide_op}, + {Py_nb_matrix_multiply, Matrix_matmul}, + {Py_nb_negative, MatrixView_negative}, + {Py_nb_absolute, MatrixView_absolute}, + {0, NULL}, +}; + +static PyType_Spec MatrixView_Spec = {.name = "bocpy._math.MatrixView", + .basicsize = sizeof(MatrixViewObject), + .itemsize = 0, + .flags = Py_TPFLAGS_DEFAULT | + Py_TPFLAGS_IMMUTABLETYPE, + .slots = MatrixView_slots}; + +// MatrixViewIndexer: the object returned by Matrix.view, giving m.view[...] +// slice-notation construction of MatrixViews. Holds a strong ref to the source +// Matrix; guards ownership at __getitem__ time (so m.view itself is cheap). +typedef struct { + PyObject_HEAD MatrixObject *matrix; +} MatrixViewIndexerObject; + +static PyObject *MatrixViewIndexer_new(PyTypeObject *Py_UNUSED(type), + PyObject *Py_UNUSED(args), + PyObject *Py_UNUSED(kwds)) { + PyErr_SetString(PyExc_TypeError, + "cannot create 'bocpy._math.MatrixViewIndexer' instances " + "directly; use Matrix.view"); + return NULL; +} + +static void MatrixViewIndexer_dealloc(PyObject *op) { + MatrixViewIndexerObject *ix = (MatrixViewIndexerObject *)op; + PyTypeObject *type = Py_TYPE(op); + Py_XDECREF(ix->matrix); + type->tp_free(op); + Py_DECREF(type); +} + +static PyObject *MatrixViewIndexer_subscript(PyObject *op, PyObject *key) { + MatrixViewIndexerObject *ix = (MatrixViewIndexerObject *)op; + matrix_impl *impl = ix->matrix->impl; + if (!impl_check_acquired(impl, true)) { + return NULL; } + const Py_ssize_t columns = (Py_ssize_t)impl->columns; - if (axis.axis == 0) { - if (impl->rows == 0) { - PyErr_SetString(PyExc_ValueError, empty_error); - return -1; + // 1-D int key -> the whole row as a row view (== m.view[i, :]). + if (PyLong_Check(key)) { + Py_ssize_t r = PyLong_AsSsize_t(key); + if (r == -1 && PyErr_Occurred()) { + return NULL; } - matrix_impl *vector = impl_new(1, impl->columns); - if (vector == NULL) { - return -1; + if (r < 0) { + r += (Py_ssize_t)impl->rows; } - argextreme_columnwise(impl, want_max, vector); - *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); - if (*out_op == NULL) { - impl_free(vector); - return -1; + if (r < 0 || (size_t)r >= impl->rows) { + PyErr_SetString(PyExc_IndexError, "row index out of range"); + return NULL; } - return 0; + return view_new(impl, r * columns, impl->columns, 1, true, false); } - if (impl->columns == 0) { - PyErr_SetString(PyExc_ValueError, empty_error); - return -1; + // 2-tuple (row_key, col_key): orientation from which axis is a bare int. + if (PyTuple_Check(key) && PyTuple_GET_SIZE(key) == 2) { + range rr; + range cc; + if (range_read(&rr, PyTuple_GET_ITEM(key, 0), impl->rows) < 0) { + return NULL; + } + if (range_read(&cc, PyTuple_GET_ITEM(key, 1), impl->columns) < 0) { + return NULL; + } + // Base offset of the first selected cell is always a valid flat index + // in [0, size), so it cannot overflow. + Py_ssize_t base_off = rr.start * columns + cc.start; + if (rr.scalar && cc.scalar) { + return view_new(impl, base_off, 1, 1, true, false); // 1-element row view + } + if (rr.scalar) { + return view_new(impl, base_off, cc.count, cc.step, true, + false); // row view + } + if (cc.scalar) { + // Column view: stride matters only for count > 1 (where |rr.step| < rows, + // so columns*rr.step < size and cannot overflow); for count <= 1 the + // stride is unused, so skip the multiply a giant step could overflow. + Py_ssize_t stride = rr.count > 1 ? columns * rr.step : columns; + return view_new(impl, base_off, rr.count, stride, false, + false); // column view + } + PyErr_SetString(PyExc_TypeError, + "a view must be a single row or column; use m[...] to copy " + "a submatrix"); + return NULL; } - matrix_impl *vector = impl_new(impl->rows, 1); - if (vector == NULL) { - return -1; + + PyErr_SetString(PyExc_TypeError, + "a view must be a single row or column; use m[...] to copy a " + "submatrix"); + return NULL; +} + +static PyType_Slot MatrixViewIndexer_slots[] = { + {Py_tp_new, MatrixViewIndexer_new}, + {Py_tp_dealloc, MatrixViewIndexer_dealloc}, + {Py_mp_subscript, MatrixViewIndexer_subscript}, + {0, NULL}, +}; + +static PyType_Spec MatrixViewIndexer_Spec = { + .name = "bocpy._math.MatrixViewIndexer", + .basicsize = sizeof(MatrixViewIndexerObject), + .itemsize = 0, + .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE, + .slots = MatrixViewIndexer_slots}; + +// Decode an optional slice/None into a full range over `length`. None (axis +// omitted) yields the whole axis; a slice decodes via range_read (steps and +// negative steps honoured). Returns 0 on success, -1 with an exception set. +static int view_axis_range(PyObject *key, size_t length, range *out) { + if (key == NULL || key == Py_None) { + out->start = 0; + out->stop = (Py_ssize_t)length; + out->step = 1; + out->count = length; + out->scalar = false; + return 0; } - argextreme_rowwise(impl, want_max, vector); - *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); - if (*out_op == NULL) { - impl_free(vector); - return -1; + return range_read(out, key, length); +} + +static PyObject *Matrix_row_view(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + Py_ssize_t r = 0; + PyObject *cols = NULL; + int read_only = 0; + static char *kwlist[] = {"row", "columns", "read_only", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "n|O$p", kwlist, &r, &cols, + &read_only)) { + return NULL; } - return 0; + if (!impl_check_acquired(impl, true)) { + return NULL; + } + if (r < 0) { + r += (Py_ssize_t)impl->rows; + } + if (r < 0 || (size_t)r >= impl->rows) { + PyErr_SetString(PyExc_IndexError, "row index out of range"); + return NULL; + } + range c; + if (view_axis_range(cols, impl->columns, &c) < 0) { + return NULL; + } + Py_ssize_t offset = r * (Py_ssize_t)impl->columns + c.start; + return view_new(impl, offset, c.count, c.step, true, (bool)read_only); } -#define MATRIX_ARGEXTREME(name, want_max_val) \ - static PyObject *Matrix_##name##_method(PyObject *op, PyObject *args, \ - PyObject *kwds) { \ - PyObject *out = NULL; \ - PyObject *axis_obj = NULL; \ - static char *kwlist[] = {"axis", NULL}; \ - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &axis_obj)) { \ - return NULL; \ - } \ - AxisArg axis; \ - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { \ - return NULL; \ - } \ - if (Matrix_argextreme(op, axis, &out, want_max_val) < 0) { \ - return NULL; \ - } \ - return out; \ +static PyObject *Matrix_column_view(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + Py_ssize_t col = 0; + PyObject *rows = NULL; + int read_only = 0; + static char *kwlist[] = {"column", "rows", "read_only", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "n|O$p", kwlist, &col, &rows, + &read_only)) { + return NULL; + } + if (!impl_check_acquired(impl, true)) { + return NULL; + } + if (col < 0) { + col += (Py_ssize_t)impl->columns; + } + if (col < 0 || (size_t)col >= impl->columns) { + PyErr_SetString(PyExc_IndexError, "column index out of range"); + return NULL; + } + range rr; + if (view_axis_range(rows, impl->rows, &rr) < 0) { + return NULL; } + Py_ssize_t offset = rr.start * (Py_ssize_t)impl->columns + col; + // Column view: stride matters only for count > 1 (where |rr.step| < rows, so + // columns*rr.step < size and cannot overflow); for count <= 1 it is unused, + // so skip the multiply that a giant slice step could otherwise overflow. + Py_ssize_t stride = rr.count > 1 ? (Py_ssize_t)impl->columns * rr.step + : (Py_ssize_t)impl->columns; + return view_new(impl, offset, rr.count, stride, false, (bool)read_only); +} -MATRIX_ARGEXTREME(argmin, false) -MATRIX_ARGEXTREME(argmax, true) +static PyObject *Matrix_get_view(PyObject *op, void *Py_UNUSED(dummy)) { + // Cheap: no ownership guard here -- a stale accessor fails at __getitem__. + PyTypeObject *type = LOCAL_STATE->view_indexer_type; + MatrixViewIndexerObject *ix = + (MatrixViewIndexerObject *)type->tp_alloc(type, 0); + if (ix == NULL) { + return NULL; + } + Py_INCREF(op); + ix->matrix = (MatrixObject *)op; + return (PyObject *)ix; +} -enum BroadcastShape { BCAST_NONE = 0, BCAST_ROW, BCAST_COL }; +// Iterator returned by Matrix.row_views() / column_views(): yields a FRESH +// MatrixView for each row (is_row) or column, one at a time. Holds a strong +// ref to the source matrix and re-checks ownership every step. Each yielded +// view is a distinct object -- never a reused cursor -- so list(m.row_views()) +// and comprehensions behave like any other iterator. +typedef struct { + PyObject_HEAD MatrixObject *matrix; + size_t cursor; + bool is_row; + bool read_only; +} MatrixViewsIterObject; + +static void MatrixViewsIter_dealloc(PyObject *op) { + MatrixViewsIterObject *it = (MatrixViewsIterObject *)op; + PyTypeObject *type = Py_TYPE(op); + Py_XDECREF(it->matrix); + type->tp_free(op); + Py_DECREF(type); +} -/* -------------------------------------------------------------------------- - Two-operand aggregate family X-macro template. +static PyObject *MatrixViewsIter_next(PyObject *op) { + MatrixViewsIterObject *it = (MatrixViewsIterObject *)op; + matrix_impl *impl = it->matrix->impl; + if (!impl_check_acquired(impl, true)) { + return NULL; + } + size_t limit = it->is_row ? impl->rows : impl->columns; + if (it->cursor >= limit) { + return NULL; // StopIteration + } + Py_ssize_t columns = (Py_ssize_t)impl->columns; + size_t i = it->cursor++; + if (it->is_row) { + // Row i occupies the contiguous span [i*columns, i*columns + columns). + return view_new(impl, (Py_ssize_t)i * columns, impl->columns, 1, true, + it->read_only); + } + // Column i is strided by `columns` down the rows starting at flat index i. + return view_new(impl, (Py_ssize_t)i, impl->rows, columns, false, + it->read_only); +} - BOC_2AGG_OPS is the single source of truth for the two-operand aggregate - op set. It is stamped in three places: - 1. impl__total — flat traversal returning a scalar. - 2. impl__rowwise — per-row reduction writing Mx1 output. - 3. impl__columnwise — per-column accumulation writing 1xN output. - REQUIRES caller-zeroed output buffer - (dispatcher uses impl_new -> PyMem_RawCalloc). +static PyObject *MatrixViewsIter_self(PyObject *op) { + Py_INCREF(op); + return op; +} - Each walker carries the per-shape switch at the TOP of the body, with - three specialised inner loops (NONE / ROW / COL). Moving the switch - inside the inner loop would defeat contraction; the per-shape pointer - arithmetic and broadcast behaviour are intentionally different. - The dispatcher (Matrix_vecdot) canonicalises operands so the matrix is - always the LHS and the vector is always the RHS before calling the - helpers. +// A views iterator is only ever born through Matrix.row_views() / +// column_views(); direct construction would leave matrix == NULL and +// NULL-deref on first step. +static PyObject *MatrixViewsIter_new(PyTypeObject *Py_UNUSED(type), + PyObject *Py_UNUSED(args), + PyObject *Py_UNUSED(kwds)) { + PyErr_SetString(PyExc_TypeError, + "cannot create 'bocpy._math.MatrixViewsIter' instances " + "directly; use Matrix.row_views() / Matrix.column_views()"); + return NULL; +} - Names in scope inside STEP: - lhs (double) — current left-hand-side value - rhs (double) — current right-hand-side value (in BCAST_COL: the - per-row scalar, hoisted out of the inner loop) - agg (double) — accumulator (for total/rowwise it is a local; - for columnwise the per-iteration cell of the output - vector is loaded into `agg`, STEP runs, then the - cell is written back) +static PyType_Slot MatrixViewsIter_slots[] = { + {Py_tp_new, MatrixViewsIter_new}, + {Py_tp_dealloc, MatrixViewsIter_dealloc}, + {Py_tp_iter, MatrixViewsIter_self}, + {Py_tp_iternext, MatrixViewsIter_next}, + {0, NULL}, +}; - Add an op: append one X(NAME, INIT, STEP) row to BOC_2AGG_OPS and the - three impl__* kernels are stamped automatically. Wire it into a - Python entry point separately (mirrors Matrix_vecdot). +static PyType_Spec MatrixViewsIter_Spec = { + .name = "bocpy._math.MatrixViewsIter", + .basicsize = sizeof(MatrixViewsIterObject), + .itemsize = 0, + .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE, + .slots = MatrixViewsIter_slots}; - See `.github/skills/commenting-c-and-python/SKILL.md` for the - "Exception: X-macro descriptor tables" convention. - -------------------------------------------------------------------------- */ -#define BOC_2AGG_OPS(X) \ - /* name init step (lhs, rhs in scope; agg accumulator) */ \ - X(vecdot, 0.0, agg += lhs * rhs) +static PyObject *make_views_iter(PyObject *op, bool is_row, bool read_only) { + MatrixObject *self = (MatrixObject *)op; + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + PyTypeObject *type = LOCAL_STATE->views_iter_type; + MatrixViewsIterObject *it = (MatrixViewsIterObject *)type->tp_alloc(type, 0); + if (it == NULL) { + return NULL; + } + Py_INCREF(op); + it->matrix = self; + it->cursor = 0; + it->is_row = is_row; + it->read_only = read_only; + return (PyObject *)it; +} -#define DEFINE_2AGG_TOTAL(NAME, INIT, STEP) \ - BOC_CANARY_NOINLINE \ - static double impl_##NAME##_total(const matrix_impl *lm, \ - const matrix_impl *rm, \ - enum BroadcastShape shape) { \ - double agg = (INIT); \ - switch (shape) { \ - case BCAST_NONE: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - for (size_t i = 0; i < lm->size; ++i, ++lp, ++rp) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - STEP; \ - } \ - break; \ - } \ - case BCAST_ROW: { \ - const double *lp = lm->data; \ - const size_t M = lm->rows; \ - const size_t N = lm->columns; \ - for (size_t r = 0; r < M; ++r) { \ - const double *rp = rm->data; \ - for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - STEP; \ - } \ - } \ - break; \ - } \ - case BCAST_COL: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - const size_t M = lm->rows; \ - const size_t N = lm->columns; \ - for (size_t r = 0; r < M; ++r) { \ - const double rhs = *rp++; \ - for (size_t c = 0; c < N; ++c, ++lp) { \ - const double lhs = *lp; \ - STEP; \ - } \ - } \ - break; \ - } \ - } \ - return agg; \ +static PyObject *Matrix_row_views(PyObject *op, PyObject *args, + PyObject *kwds) { + int read_only = 0; + static char *kwlist[] = {"read_only", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|$p", kwlist, &read_only)) { + return NULL; } + return make_views_iter(op, true, (bool)read_only); +} -#define DEFINE_2AGG_ROWWISE(NAME, INIT, STEP) \ - BOC_CANARY_NOINLINE \ - static void impl_##NAME##_rowwise( \ - const matrix_impl *lm, const matrix_impl *rm, matrix_impl *out_Mx1, \ - enum BroadcastShape shape) { \ - const size_t M = lm->rows; \ - const size_t N = lm->columns; \ - double *out_ptr = out_Mx1->data; \ - switch (shape) { \ - case BCAST_NONE: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - for (size_t r = 0; r < M; ++r, ++out_ptr) { \ - double agg = (INIT); \ - for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - STEP; \ - } \ - *out_ptr = agg; \ - } \ - break; \ - } \ - case BCAST_ROW: { \ - const double *lp = lm->data; \ - for (size_t r = 0; r < M; ++r, ++out_ptr) { \ - const double *rp = rm->data; \ - double agg = (INIT); \ - for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - STEP; \ - } \ - *out_ptr = agg; \ - } \ - break; \ - } \ - case BCAST_COL: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - for (size_t r = 0; r < M; ++r, ++out_ptr) { \ - const double rhs = *rp++; \ - double agg = (INIT); \ - for (size_t c = 0; c < N; ++c, ++lp) { \ - const double lhs = *lp; \ - STEP; \ - } \ - *out_ptr = agg; \ - } \ - break; \ - } \ - } \ +static PyObject *Matrix_column_views(PyObject *op, PyObject *args, + PyObject *kwds) { + int read_only = 0; + static char *kwlist[] = {"read_only", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|$p", kwlist, &read_only)) { + return NULL; } + return make_views_iter(op, false, (bool)read_only); +} -/* Columnwise: the per-output-cell accumulator IS the output vector slot - itself. Loading `*out_ptr` into a local `agg`, running STEP, and - writing back keeps STEP uniform across all three walkers; the compiler - collapses the load/store pair to a single += against memory. The - caller must hand in a zero-initialised buffer (impl_new -> calloc); - a recycled non-zero buffer would silently produce wrong sums because - STEP accumulates with `+=`. */ -#define DEFINE_2AGG_COLUMNWISE(NAME, INIT, STEP) \ - BOC_CANARY_NOINLINE \ - static void impl_##NAME##_columnwise( \ - const matrix_impl *lm, const matrix_impl *rm, matrix_impl *out_1xN, \ - enum BroadcastShape shape) { \ - const size_t M = lm->rows; \ - const size_t N = lm->columns; \ - (void)(INIT); \ - switch (shape) { \ - case BCAST_NONE: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - for (size_t r = 0; r < M; ++r) { \ - double *out_ptr = out_1xN->data; \ - for (size_t c = 0; c < N; ++c, ++lp, ++rp, ++out_ptr) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - double agg = *out_ptr; \ - STEP; \ - *out_ptr = agg; \ - } \ - } \ - break; \ - } \ - case BCAST_ROW: { \ - const double *lp = lm->data; \ - for (size_t r = 0; r < M; ++r) { \ - const double *rp = rm->data; \ - double *out_ptr = out_1xN->data; \ - for (size_t c = 0; c < N; ++c, ++lp, ++rp, ++out_ptr) { \ - const double lhs = *lp; \ - const double rhs = *rp; \ - double agg = *out_ptr; \ - STEP; \ - *out_ptr = agg; \ - } \ - } \ - break; \ - } \ - case BCAST_COL: { \ - const double *lp = lm->data; \ - const double *rp = rm->data; \ - for (size_t r = 0; r < M; ++r) { \ - const double rhs = *rp++; \ - double *out_ptr = out_1xN->data; \ - for (size_t c = 0; c < N; ++c, ++lp, ++out_ptr) { \ - const double lhs = *lp; \ - double agg = *out_ptr; \ - STEP; \ - *out_ptr = agg; \ - } \ - } \ - break; \ - } \ - } \ +static PyObject *Matrix_transpose(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *matrix = (MatrixObject *)op; + matrix_impl *impl = matrix->impl; + + int in_place = 0; + static char *kwlist[] = {"in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|p", kwlist, &in_place)) { + return NULL; } -#define X(N, I, S) DEFINE_2AGG_TOTAL(N, I, S) -BOC_2AGG_OPS(X) -#undef X + if (!impl_check_acquired(impl, true)) { + return NULL; + } -#define X(N, I, S) DEFINE_2AGG_ROWWISE(N, I, S) -BOC_2AGG_OPS(X) -#undef X + if (in_place) { + if (impl_transpose_in_place(impl) < 0) { + return NULL; + } + return Py_NewRef(op); + } -#define X(N, I, S) DEFINE_2AGG_COLUMNWISE(N, I, S) -BOC_2AGG_OPS(X) -#undef X + matrix_impl *transpose = impl_transpose(impl); + if (transpose == NULL) { + return NULL; + } -/// @brief Axis-aware inner product: sum of element-wise products. -/// @details The canonicalisation swap rearranges the dispatch-argument -/// pointers ``mat_arg`` / ``vec_arg`` so the helpers always see -/// ``(matrix, vector)``. The refcounted ``rhs`` from -/// ``unwrap_matrix`` is preserved unchanged so the IMPL_DECREF at -/// ``done`` always matches the single INCREF. ``self->impl`` is -/// NOT refcount-paired here (mirrors Matrix_transpose). -static PyObject *Matrix_vecdot(PyObject *op, PyObject *args, PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - PyObject *other = NULL; - PyObject *axis = NULL; - PyObject *result = NULL; - matrix_impl *rhs = NULL; + return wrap_impl_or_free(transpose); +} - static char *kwlist[] = {"", "axis", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", kwlist, &other, &axis)) { +/// @brief Validate a caller-supplied ``out=`` target and alias its buffer. +/// @details Implements the numpy-style ``out=`` convention for the +/// allocation-free elementwise paths: the result is written into an +/// existing matrix the caller owns rather than a fresh allocation. +/// The target must be a Matrix owned by the current interpreter whose +/// shape exactly matches the operation's result shape; the shape is +/// validated here, before any kernel writes, so a rejected target +/// leaves every buffer untouched (the validate-then-write contract +/// shared with ``put``). The target may alias the primary +/// same-shape operand because the kernel reads and writes that +/// operand at the same index; aliasing a lower-rank broadcast +/// operand is impossible because the result-shape check rejects any +/// target whose shape differs from the result. +/// @param out_target The ``out=`` Matrix object (never NULL / Py_None here). +/// @param out_op Receives a new reference to ``out_target`` on success. +/// @param want_rows Required row count of the result. +/// @param want_cols Required column count of the result. +/// @return The target's impl (no extra C refcount taken — the live Python +/// reference owns it), or NULL with an exception set on failure. +static matrix_impl *use_out_target(PyObject *out_target, PyObject **out_op, + size_t want_rows, size_t want_cols) { + if (Py_TYPE(out_target) != LOCAL_STATE->matrix_type) { + PyErr_SetString(PyExc_TypeError, "out must be a Matrix"); return NULL; } - if (!impl_check_acquired(self->impl, true)) { + matrix_impl *impl = ((MatrixObject *)out_target)->impl; + if (!impl_check_acquired(impl, true)) { return NULL; } - - rhs = unwrap_matrix(other, false); - if (rhs == NULL) { - goto done; - } - - matrix_impl *lhs = self->impl; - matrix_impl *mat_arg = lhs; - matrix_impl *vec_arg = rhs; - enum BroadcastShape shape; - - if (lhs->rows == rhs->rows && lhs->columns == rhs->columns) { - shape = BCAST_NONE; - } else if (lhs->rows == rhs->rows && - (lhs->columns == 1 || rhs->columns == 1)) { - shape = BCAST_COL; - if (lhs->columns == 1) { - mat_arg = rhs; - vec_arg = lhs; - } - } else if (lhs->columns == rhs->columns && - (lhs->rows == 1 || rhs->rows == 1)) { - shape = BCAST_ROW; - if (lhs->rows == 1) { - mat_arg = rhs; - vec_arg = lhs; - } - } else if ((lhs->rows == 1 || lhs->columns == 1) && - (rhs->rows == 1 || rhs->columns == 1) && lhs->size == rhs->size) { - shape = BCAST_NONE; - } else { + if (impl->rows != want_rows || impl->columns != want_cols) { PyErr_Format(PyExc_ValueError, - "vecdot: lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, - lhs->columns, rhs->rows, rhs->columns); - goto done; + "out shape %zux%zu does not match result %zux%zu", impl->rows, + impl->columns, want_rows, want_cols); + return NULL; } + *out_op = Py_NewRef(out_target); + return impl; +} - AxisArg axis_arg; - if (parse_validate_normalise_axis(axis, &axis_arg) < 0) { - goto done; +/// @brief Sets the output of an arithmetic operation. +/// @details The output of an arithmetic operation is one of three things: a +/// caller-supplied ``out=`` target (when ``out_target`` is a non-None Matrix), +/// the left-hand side itself (in-place operations), or a freshly allocated +/// matrix of the same dimensions as the left-hand side. ``out_target`` and +/// ``inplace`` are mutually exclusive (asserted here; callers reject the +/// combination before reaching this point). +/// +/// The result shape and wrap type come from the already-unwrapped ``lhs`` +/// impl and the interpreter's canonical Matrix type -- never by casting +/// ``lhs_op``. ``lhs_op`` may be a sequence that ``unwrap_matrix`` coerced +/// into ``lhs`` (e.g. ``[1, 2, 3] + matrix``), so it is not guaranteed to be +/// a Matrix; it is dereferenced only as the in-place return value, a path +/// reachable solely with a genuine Matrix left operand. +/// @param lhs The unwrapped impl of the left-hand operand (gives result shape) +/// @param lhs_op The PyObject of the left-hand operand (in-place return only) +/// @param out_op A pointer to the output pointer of the equation +/// @param out_target Optional ``out=`` Matrix (NULL or Py_None when unused) +/// @param inplace Whether this is an inplace operation +/// @return The matrix wrapped by out_op, or NULL in the case of an error +static matrix_impl *set_output(matrix_impl *lhs, PyObject *lhs_op, + PyObject **out_op, PyObject *out_target, + bool inplace) { + matrix_impl *out; + if (out_target != NULL && out_target != Py_None) { + assert(!inplace); + return use_out_target(out_target, out_op, lhs->rows, lhs->columns); } - - if (!axis_arg.has_axis) { - result = PyFloat_FromDouble(impl_vecdot_total(mat_arg, vec_arg, shape)); - } else if (axis_arg.axis == 0) { - matrix_impl *out = impl_new(1, mat_arg->columns); - if (out != NULL) { - impl_vecdot_columnwise(mat_arg, vec_arg, out, shape); - result = wrap_impl_or_free(out); - } - } else { - // axis == 1 (row-wise): parse_validate_normalise_axis restricts axis - // to {0, 1}. - matrix_impl *out = impl_new(mat_arg->rows, 1); - if (out != NULL) { - impl_vecdot_rowwise(mat_arg, vec_arg, out, shape); - result = wrap_impl_or_free(out); - } + if (inplace) { + *out_op = Py_NewRef(lhs_op); + return lhs; } -done: - IMPL_DECREF(rhs); - return result; -} - -/// @brief A classified broadcast operand: scalar or same-shape matrix. -/// @details ``full`` is an INCREF'd matrix with the same shape as the -/// receiver, and the caller must IMPL_DECREF it; when ``full`` is -/// NULL the operand is the scalar in ``scalar``. Shared by ``fma`` -/// and ``scaled_add`` to classify a multiplier / scale operand. -typedef struct { - double scalar; - matrix_impl *full; -} broadcast_operand; - -/// @brief Materialise a row- or column-vector broadcast into a full buffer. -/// @details Expands a ``1xN`` row vector (``is_row``) or an ``Mx1`` column -/// vector into a fresh contiguous ``rows`` x ``columns`` matrix so -/// the contiguous fma kernel keeps its unit stride (and hardware -/// FMA). The returned impl has refcount 0; the caller INCREFs it. -static matrix_impl *impl_broadcast_vector(const matrix_impl *vec, size_t rows, - size_t columns, bool is_row) { - matrix_impl *out = impl_new(rows, columns); + out = impl_new(lhs->rows, lhs->columns); if (out == NULL) { return NULL; } - double *dst = out->data; - const double *v = vec->data; - if (is_row) { - for (size_t r = 0; r < rows; ++r) { - for (size_t c = 0; c < columns; ++c) { - *dst++ = v[c]; - } - } - } else { - for (size_t r = 0; r < rows; ++r) { - for (size_t c = 0; c < columns; ++c) { - *dst++ = v[r]; - } - } + + *out_op = wrap_matrix(LOCAL_STATE->matrix_type, out); + if (*out_op == NULL) { + impl_free(out); + return NULL; } + return out; } -/// @brief Classify a broadcast operand as a scalar or a same-shape matrix. -/// @details A Matrix whose size is 1 is treated as a scalar so a ``1x1`` -/// broadcasts like a number (the in-house scalar rule used -/// elsewhere); a same-shape Matrix is INCREF'd into ``out->full``. -/// A ``1xN`` row vector (matching self's columns) or an ``Mx1`` -/// column vector (matching self's rows) is materialised into a fresh -/// same-shape buffer stored in ``out->full`` so the contiguous kernel -/// keeps its unit stride; any other matrix shape raises ``ValueError`` -/// naming both shapes. Any real number is accepted as a scalar. -/// @param op_name Method name used as the error-message prefix (e.g. ``fma``). -/// @param operand The operand object to classify. -/// @param self_impl The receiver's impl, used for the shape check. -/// @param what Operand name used in error messages. -/// @param out Output operand descriptor (zeroed before classification). -/// @return 0 on success; -1 with an exception set on failure. -static int classify_broadcast_operand(const char *op_name, PyObject *operand, - const matrix_impl *self_impl, - const char *what, - broadcast_operand *out) { - out->scalar = 0.0; - out->full = NULL; - - if (Py_TYPE(operand) == LOCAL_STATE->matrix_type) { - matrix_impl *impl = ((MatrixObject *)operand)->impl; - if (!impl_check_acquired(impl, true)) { - return -1; - } - if (impl->size == 1) { - out->scalar = impl->data[0]; - return 0; - } - if (impl->rows == self_impl->rows && impl->columns == self_impl->columns) { - IMPL_INCREF(impl); - out->full = impl; - return 0; - } - if (impl->rows == 1 && impl->columns == self_impl->columns) { - matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, - self_impl->columns, true); - if (full == NULL) { - return -1; - } - IMPL_INCREF(full); - out->full = full; - return 0; - } - if (impl->columns == 1 && impl->rows == self_impl->rows) { - matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, - self_impl->columns, false); - if (full == NULL) { - return -1; - } - IMPL_INCREF(full); - out->full = full; - return 0; - } - PyErr_Format(PyExc_ValueError, - "%s: %s shape %zux%zu incompatible with self %zux%zu", op_name, - what, impl->rows, impl->columns, self_impl->rows, - self_impl->columns); - return -1; - } +/// @brief Tri-state representation of an optional ``axis`` kwarg. +/// @details ``has_axis`` is true iff the caller passed a non-None axis; +/// ``axis`` is the validated int value (and is undefined when +/// ``has_axis`` is false). Replaces the historical NO_AXIS=-1000 +/// sentinel which collided with the integer -1000. +typedef struct { + bool has_axis; + int axis; +} AxisArg; - if (unwrap_double(operand, &out->scalar)) { +/// @brief Decode an optional ``axis`` keyword argument into an AxisArg. +/// @details Accepts ``NULL`` or ``Py_None`` (no axis), an ``int`` in +/// ``INT_MIN..INT_MAX``. Rejects ``bool`` (subclass of int) and +/// overflows. Returns 0 on success and writes through ``*out``; +/// returns -1 with TypeError / OverflowError set on failure. +static int decode_axis_kwarg(PyObject *axis_obj, AxisArg *out) { + if (axis_obj == NULL || axis_obj == Py_None) { + out->has_axis = false; + out->axis = 0; return 0; } - if (PyErr_Occurred()) { + if (PyBool_Check(axis_obj)) { + PyErr_SetString(PyExc_TypeError, "axis must be an int or None, not bool"); return -1; } - - PyErr_Format(PyExc_TypeError, "%s: %s must be a Matrix or a real number", - op_name, what); - return -1; -} - -/// @brief Fused multiply-add kernel: ``out = a*b + c`` with one rounding. -/// @details ``b_step`` / ``c_step`` are 1 for a full same-shape operand and -/// 0 to repeat a scalar (the pointer then aims at a single local -/// ``double``). Each output cell reads only its own index, so the -/// in-place case (``out == a``) and ``b`` / ``c`` aliasing ``a`` -/// are all safe. ``fma()`` rounds the product once, unlike -/// ``a*b + c`` which rounds twice under ``-ffp-contract=off``. -/// BOC_FMA_MULTIVERSION clones this kernel for hardware FMA on -/// glibc x86-64 (see the macro definition for why glibc-only). -BOC_CANARY_NOINLINE -BOC_FMA_MULTIVERSION -static void impl_fma(const double *a, const double *b, size_t b_step, - const double *c, size_t c_step, double *out, size_t n) { - for (size_t i = 0; i < n; ++i, ++a, b += b_step, c += c_step, ++out) { - *out = fma(*a, *b, *c); + if (!PyLong_Check(axis_obj)) { + PyErr_SetString(PyExc_TypeError, "axis must be an int or None"); + return -1; } -} - -/// @brief Fused multiply-add: single-rounding ``self*b + c``. -/// @details ``b`` and ``c`` may each be a same-shape matrix, a ``1x1`` -/// matrix, a ``1xN`` row vector or ``Mx1`` column vector that -/// broadcasts against ``self``, or a scalar; any other matrix shape -/// raises ``ValueError``. Both operands are validated before any -/// allocation, so a rejected operand leaves ``self`` untouched. -/// With ``in_place=True`` the result is written into ``self`` and -/// ``self`` is returned. -static PyObject *Matrix_fma(PyObject *op, PyObject *args, PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - PyObject *b_op = NULL; - PyObject *c_op = NULL; - int in_place = 0; - PyObject *out_op = NULL; - broadcast_operand bop = {0.0, NULL}; - broadcast_operand cop = {0.0, NULL}; - - static char *kwlist[] = {"", "", "in_place", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &b_op, &c_op, - &in_place)) { - return NULL; + long ax = PyLong_AsLong(axis_obj); + if (ax == -1 && PyErr_Occurred()) { + return -1; } - - if (!impl_check_acquired(self->impl, true)) { - return NULL; + if (ax < INT_MIN || ax > INT_MAX) { + PyErr_Format(PyExc_OverflowError, "axis %ld out of int range", ax); + return -1; } + out->has_axis = true; + out->axis = (int)ax; + return 0; +} - if (classify_broadcast_operand("fma", b_op, self->impl, "b", &bop) < 0) { - goto done; +/// @brief Decode, validate (-2/-1/0/1) and normalise (-1->1, -2->0) axis. +/// @details Single entry point for the cross / perpendicular / angle / +/// normalize / aggregate methods that accept only the four +/// standard axis values. After this call, ``out->axis`` is +/// either 0 (column-wise) or 1 (row-wise) when ``has_axis`` is +/// true. Returns -1 with the appropriate exception set on any +/// decode or range error. +static int parse_validate_normalise_axis(PyObject *axis_obj, AxisArg *out) { + if (decode_axis_kwarg(axis_obj, out) < 0) { + return -1; } - if (classify_broadcast_operand("fma", c_op, self->impl, "c", &cop) < 0) { - goto done; + if (!out->has_axis) { + return 0; } - - matrix_impl *out = set_output(self->impl, op, &out_op, NULL, in_place); - if (out == NULL) { - out_op = NULL; - goto done; + int ax = out->axis; + if (ax != 0 && ax != 1 && ax != -1 && ax != -2) { + PyErr_SetString(PyExc_ValueError, "axis must be -2, -1, 0, or 1"); + return -1; } - - double b_scalar = bop.scalar; - double c_scalar = cop.scalar; - const double *b_ptr = bop.full != NULL ? bop.full->data : &b_scalar; - const double *c_ptr = cop.full != NULL ? cop.full->data : &c_scalar; - size_t b_step = bop.full != NULL ? 1 : 0; - size_t c_step = cop.full != NULL ? 1 : 0; - - impl_fma(self->impl->data, b_ptr, b_step, c_ptr, c_step, out->data, - self->impl->size); - -done: - IMPL_DECREF(bop.full); - IMPL_DECREF(cop.full); - return out_op; -} - -/// @brief Two-rounding scaled-add kernel: ``out = y + s * x``. -/// @details Deliberately NOT fused: the product ``s[i] * x[i]`` is rounded to -/// double first, then the sum with ``y[i]`` is rounded again, so the -/// result is bit-for-bit identical to the naive ``y[i] + s * x[i]`` -/// expression in IEEE-754 double (two roundings). This is the -/// two-rounding complement to ``impl_fma``'s single rounding; the -/// ``-ffp-contract=off`` build flag guarantees the two statements -/// are never contracted back into an fma. ``s_step`` is 1 for a full -/// same-shape scale buffer and 0 to repeat a broadcast scalar (the -/// pointer then aims at a single local ``double``); ``x`` and ``y`` -/// are the same shape (unit stride). Each output cell reads only its -/// own index, so the in-place case (``out == y``) and ``x`` / ``s`` -/// aliasing ``y`` are all safe. -BOC_CANARY_NOINLINE -static void impl_scaled_add(const double *y, const double *s, size_t s_step, - const double *x, double *out, size_t n) { - for (size_t i = 0; i < n; ++i, ++y, s += s_step, ++x, ++out) { - const double prod = (*s) * (*x); - *out = *y + prod; + if (ax == -1) { + out->axis = 1; + } else if (ax == -2) { + out->axis = 0; } + return 0; } -/// @brief Scaled add: ``self + s * x`` with two roundings. -/// @details The two-rounding sibling of ``fma``: the product ``s * x`` is -/// rounded to double, then the sum with ``self`` is rounded again, so -/// the result is bit-for-bit identical to ``self + s * x`` and -/// distinct from the single-rounded ``fma``. ``s`` is the scale and -/// may be a scalar, a ``1x1`` matrix, a ``1xN`` row vector or ``Mx1`` -/// column vector that broadcasts against ``self``, or a same-shape -/// matrix (the same operand rules as ``fma``'s multiplier); ``x`` is -/// a same-shape matrix. Both operands are validated before any -/// allocation, so a rejected operand leaves ``self`` untouched (the -/// validate-then-write contract shared with ``put``). With -/// ``in_place=True`` the result is written into ``self``'s existing -/// buffer (allocating nothing) and ``self`` is returned; otherwise a -/// fresh matrix is returned. ``s`` and ``x`` may alias ``self``. -static PyObject *Matrix_scaled_add(PyObject *op, PyObject *args, - PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - PyObject *s_op = NULL; - PyObject *x_op = NULL; - int in_place = 0; - PyObject *out_op = NULL; - broadcast_operand sop = {0.0, NULL}; - matrix_impl *x = NULL; - - static char *kwlist[] = {"", "", "in_place", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &s_op, &x_op, - &in_place)) { +// Validate and unwrap a where= mask: it must be a Matrix of the same shape +// as the matrix being aggregated and owned by the current interpreter. A +// mask cell == 0.0 excludes that element; NaN counts as included (matching +// where()'s truthiness). Returns a borrowed impl (the caller holds the +// Python object alive for the synchronous call) or NULL with an exception. +static matrix_impl *unwrap_mask(PyObject *where_op, matrix_impl *impl) { + if (Py_TYPE(where_op) != LOCAL_STATE->matrix_type) { + PyErr_SetString(PyExc_TypeError, "where must be a Matrix mask"); return NULL; } - if (!impl_check_acquired(self->impl, true)) { + matrix_impl *mask = ((MatrixObject *)where_op)->impl; + if (!impl_check_acquired(mask, true)) { return NULL; } - if (classify_broadcast_operand("scaled_add", s_op, self->impl, "s", &sop) < - 0) { + if (mask->rows != impl->rows || mask->columns != impl->columns) { + PyErr_Format(PyExc_ValueError, + "where mask shape %zux%zu does not match matrix shape " + "%zux%zu", + mask->rows, mask->columns, impl->rows, impl->columns); return NULL; } - x = unwrap_matrix(x_op, false); - if (x == NULL) { - goto done; + return mask; +} + +static int Matrix_aggregate(PyObject *matrix_op, AxisArg axis, + PyObject *where_op, PyObject **out_op, + enum AggregateOps agg) { + MatrixObject *matrix = (MatrixObject *)matrix_op; + matrix_impl *impl = matrix->impl; + + if (!impl_check_acquired(impl, true)) { + return -1; } - matrix_impl *y = self->impl; - if (x->rows != y->rows || x->columns != y->columns) { - PyErr_Format(PyExc_ValueError, - "scaled_add: x shape %zux%zu does not match self %zux%zu", - x->rows, x->columns, y->rows, y->columns); - goto done; + matrix_impl *mask = NULL; + if (where_op != NULL && where_op != Py_None) { + mask = unwrap_mask(where_op, impl); + if (mask == NULL) { + return -1; + } } - matrix_impl *out = set_output(y, op, &out_op, NULL, in_place); - if (out == NULL) { - out_op = NULL; - goto done; + if (!axis.has_axis) { + double value = mask ? dispatch_agg_masked_ewise(impl, mask, agg) + : dispatch_agg_ewise(impl, agg); + *out_op = PyFloat_FromDouble(value); + return 0; } - double s_scalar = sop.scalar; - const double *s_ptr = sop.full != NULL ? sop.full->data : &s_scalar; - size_t s_step = sop.full != NULL ? 1 : 0; + if (axis.axis == 0) { + matrix_impl *vector = impl_new(1, impl->columns); + if (vector == NULL) { + return -1; + } - impl_scaled_add(y->data, s_ptr, s_step, x->data, out->data, y->size); + if (mask) { + dispatch_agg_masked_columnwise(impl, mask, agg, vector); + } else { + dispatch_agg_columnwise(impl, agg, vector); + } + *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); + if (*out_op == NULL) { + impl_free(vector); + return -1; + } -done: - IMPL_DECREF(sop.full); - IMPL_DECREF(x); - return out_op; -} + return 0; + } -/// @details ``has_axis`` / ``explicit_axis`` carry an optional caller- -/// supplied axis (already normalised to 0 or 1 by -/// ``parse_validate_normalise_axis``). For the doubly-valid -/// ``2x2`` / ``3x3`` shapes ``explicit_axis`` picks the -/// orientation (``0`` -> columns, default -> rows). For all -/// other shapes only one orientation is valid; supplying an -/// ``explicit_axis`` that contradicts that orientation returns -/// ``CROSS_INVALID`` so the caller raises rather than running -/// the wrong kernel. Returns ``CROSS_INVALID`` for any shape -/// that has no valid cross-product interpretation. -enum CrossAxis { - CROSS_SCALAR_2D_1x2, - CROSS_SCALAR_2D_2x1, - CROSS_ROWS_2D_Nx2, - CROSS_COLS_2D_2xN, - CROSS_SCALAR_3D_1x3, - CROSS_SCALAR_3D_3x1, - CROSS_ROWS_3D_Nx3, - CROSS_COLS_3D_3xN, - CROSS_INVALID -}; + // Fall-through is axis == 1 (row-wise): parse_validate_normalise_axis + // restricts axis to {0, 1}. + matrix_impl *vector = impl_new(impl->rows, 1); + if (vector == NULL) { + return -1; + } -static enum CrossAxis classify_cross_axis(const matrix_impl *impl, - bool has_axis, int explicit_axis) { - const size_t M = impl->rows; - const size_t N = impl->columns; - if (M == 2 && N == 2) { - return (has_axis && explicit_axis == 0) ? CROSS_COLS_2D_2xN - : CROSS_ROWS_2D_Nx2; + if (mask) { + dispatch_agg_masked_rowwise(impl, mask, agg, vector); + } else { + dispatch_agg_rowwise(impl, agg, vector); } - if (M == 3 && N == 3) { - return (has_axis && explicit_axis == 0) ? CROSS_COLS_3D_3xN - : CROSS_ROWS_3D_Nx3; + *out_op = wrap_matrix(Py_TYPE(matrix_op), vector); + if (*out_op == NULL) { + impl_free(vector); + return -1; } - if (M == 1 && N == 2) { - if (has_axis && explicit_axis == 0) { - return CROSS_INVALID; - } - return CROSS_SCALAR_2D_1x2; + + return 0; +} + +#define MATRIX_AGGREGATE(agg) \ + static PyObject *Matrix_##agg##_method(PyObject *op, PyObject *args, \ + PyObject *kwds) { \ + PyObject *out = NULL; \ + PyObject *axis_obj = NULL; \ + PyObject *where_obj = NULL; \ + static char *kwlist[] = {"axis", "where", NULL}; \ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO", kwlist, &axis_obj, \ + &where_obj)) { \ + return NULL; \ + } \ + AxisArg axis; \ + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { \ + return NULL; \ + } \ + if (Matrix_aggregate(op, axis, where_obj, &out, agg) < 0) { \ + return NULL; \ + } \ + return out; \ } - if (M == 1 && N == 3) { - if (has_axis && explicit_axis == 0) { - return CROSS_INVALID; + +MATRIX_AGGREGATE(Sum) +MATRIX_AGGREGATE(Mean) +MATRIX_AGGREGATE(Magnitude) +MATRIX_AGGREGATE(Minimum) +MATRIX_AGGREGATE(Maximum) +MATRIX_AGGREGATE(MagnitudeSquared) + +/* -------------------------------------------------------------------------- + Arg-reduction (argmin / argmax) kernels. + + These do not fit the BOC_AGG_OPS X-macro: that table accumulates a + single double, whereas an arg-reduction must also carry the index of + the running extreme. Comparisons are strict so the first occurrence of + a tied extreme wins, matching NumPy. Indices are published as doubles + in the result matrix (the Matrix type stores only doubles). + -------------------------------------------------------------------------- */ +// Arg-reduction (argmin/argmax) kernels: kept out of the BOC_AGG_OPS X-macro +// because they must carry the running index, not just a double accumulator. +// Strict comparisons make the first tied extreme win (NumPy tie-break); +// indices are published as doubles (Matrix stores only doubles). +static Py_ssize_t argextreme_ewise(matrix_impl *m, bool want_max) { + const double *p = m->data; + double best = p[0]; + Py_ssize_t best_i = 0; + for (size_t i = 1; i < m->size; ++i) { + const double v = p[i]; + if (want_max ? (v > best) : (v < best)) { + best = v; + best_i = (Py_ssize_t)i; } - return CROSS_SCALAR_3D_1x3; } - if (M == 2 && N == 1) { - if (has_axis && explicit_axis == 1) { - return CROSS_INVALID; + return best_i; +} + +static void argextreme_columnwise(matrix_impl *m, bool want_max, + matrix_impl *out) { + const size_t M = m->rows; + const size_t N = m->columns; + for (size_t c = 0; c < N; ++c) { + double best = m->data[c]; + size_t best_r = 0; + for (size_t r = 1; r < M; ++r) { + const double v = m->data[r * N + c]; + if (want_max ? (v > best) : (v < best)) { + best = v; + best_r = r; + } } - return CROSS_SCALAR_2D_2x1; + out->data[c] = (double)best_r; } - if (M == 3 && N == 1) { - if (has_axis && explicit_axis == 1) { - return CROSS_INVALID; +} + +static void argextreme_rowwise(matrix_impl *m, bool want_max, + matrix_impl *out) { + const size_t M = m->rows; + const size_t N = m->columns; + for (size_t r = 0; r < M; ++r) { + const double *row = m->data + r * N; + double best = row[0]; + size_t best_c = 0; + for (size_t c = 1; c < N; ++c) { + const double v = row[c]; + if (want_max ? (v > best) : (v < best)) { + best = v; + best_c = c; + } } - return CROSS_SCALAR_3D_3x1; + out->data[r] = (double)best_c; } - if (N == 2 && M != 3) { - if (has_axis && explicit_axis == 0) { - return CROSS_INVALID; +} + +/* Masked arg-reduction kernels: among the included elements (mask cell != + 0.0; NaN counts as included, matching where()'s truthiness) find the index + of the first strict extreme. The first included element seeds the running + extreme (best_i < 0 guard); strict comparisons then keep the first + occurrence on a tie and -- exactly like the unmasked kernels -- a NaN never + displaces a real running extreme, while a NaN seed pins the result to its + position. An all-excluded group has no argument and yields + ARGEXTREME_MASKED_EMPTY (-1), the integer analog of the masked aggregate's + NaN. Rows/columns are addressed through row_ptrs so mask and matrix stay + aligned cell-for-cell. */ +#define ARGEXTREME_MASKED_EMPTY (-1) + +static Py_ssize_t argextreme_masked_ewise(matrix_impl *m, matrix_impl *mask, + bool want_max) { + const double *p = m->data; + const double *kp = mask->data; + double best = 0.0; + Py_ssize_t best_i = ARGEXTREME_MASKED_EMPTY; + for (size_t i = 0; i < m->size; ++i) { + if (kp[i] == 0.0) { + continue; } - return CROSS_ROWS_2D_Nx2; - } - if (M == 2 && N != 3) { - if (has_axis && explicit_axis == 1) { - return CROSS_INVALID; + const double v = p[i]; + if (best_i < 0 || (want_max ? (v > best) : (v < best))) { + best = v; + best_i = (Py_ssize_t)i; } - return CROSS_COLS_2D_2xN; } - if (N == 3 && M != 2) { - if (has_axis && explicit_axis == 0) { - return CROSS_INVALID; + return best_i; +} + +static void argextreme_masked_columnwise(matrix_impl *m, matrix_impl *mask, + bool want_max, matrix_impl *out) { + const size_t M = m->rows; + const size_t N = m->columns; + for (size_t c = 0; c < N; ++c) { + double best = 0.0; + Py_ssize_t best_r = ARGEXTREME_MASKED_EMPTY; + for (size_t r = 0; r < M; ++r) { + if (mask->row_ptrs[r][c] == 0.0) { + continue; + } + const double v = m->row_ptrs[r][c]; + if (best_r < 0 || (want_max ? (v > best) : (v < best))) { + best = v; + best_r = (Py_ssize_t)r; + } } - return CROSS_ROWS_3D_Nx3; + out->data[c] = (double)best_r; } - if (M == 3 && N != 2) { - if (has_axis && explicit_axis == 1) { - return CROSS_INVALID; +} + +static void argextreme_masked_rowwise(matrix_impl *m, matrix_impl *mask, + bool want_max, matrix_impl *out) { + const size_t M = m->rows; + const size_t N = m->columns; + for (size_t r = 0; r < M; ++r) { + const double *row = m->row_ptrs[r]; + const double *krow = mask->row_ptrs[r]; + double best = 0.0; + Py_ssize_t best_c = ARGEXTREME_MASKED_EMPTY; + for (size_t c = 0; c < N; ++c) { + if (krow[c] == 0.0) { + continue; + } + const double v = row[c]; + if (best_c < 0 || (want_max ? (v > best) : (v < best))) { + best = v; + best_c = (Py_ssize_t)c; + } } - return CROSS_COLS_3D_3xN; - } - if (N == 2) { - return CROSS_ROWS_2D_Nx2; - } - if (M == 2) { - return CROSS_COLS_2D_2xN; - } - if (N == 3) { - return CROSS_ROWS_3D_Nx3; - } - if (M == 3) { - return CROSS_COLS_3D_3xN; + out->data[r] = (double)best_c; } - return CROSS_INVALID; } -/// @brief 2D / 3D cross product against another vector or batch. -/// @details Five paths share one dispatcher. For 1x2 / 2x1 inputs the -/// result is the scalar z-component -/// ``self.x * other.y - self.y * other.x`` as a Python float; -/// for 1x3 / 3x1 the result is a same-shape Matrix preserving -/// ``self``'s orientation. For Nx2 / 2xN batches the result is -/// a per-vector scalar collected in a Mx1 (rows) or 1xN (cols) -/// Matrix; for Nx3 / 3xN batches the result is a same-shape -/// Matrix of per-vector cross products. The ``axis`` keyword -/// disambiguates 2x2 / 3x3 squares (default: rows; ``axis=0`` -/// forces columns). For scalar inputs ``other``'s orientation -/// is irrelevant; for batch inputs ``other`` must have the same -/// shape as ``self``. -static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - PyObject *other_op = NULL; - PyObject *axis_obj = NULL; - PyObject *result = NULL; - matrix_impl *rhs = NULL; - - static char *kwlist[] = {"", "axis", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", kwlist, &other_op, - &axis_obj)) { +// Convert an Mx1 or 1xN vector of index positions (stored as doubles by the +// argextreme kernels) into a flat Python list[int] — directly usable in fancy +// indexing without a float-matrix conversion. +static PyObject *argextreme_indices_to_list(const matrix_impl *vector) { + PyObject *list = PyList_New((Py_ssize_t)vector->size); + if (list == NULL) { return NULL; } - if (!impl_check_acquired(self->impl, true)) { - return NULL; + for (size_t i = 0; i < vector->size; ++i) { + PyObject *idx = PyLong_FromSsize_t((Py_ssize_t)vector->data[i]); + if (idx == NULL) { + Py_DECREF(list); + return NULL; + } + PyList_SET_ITEM(list, (Py_ssize_t)i, idx); } + return list; +} - AxisArg axis; - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { - return NULL; - } +static int Matrix_argextreme(PyObject *matrix_op, AxisArg axis, + PyObject *where_op, PyObject **out_op, + bool want_max, bool as_matrix) { + MatrixObject *matrix = (MatrixObject *)matrix_op; + matrix_impl *impl = matrix->impl; - rhs = unwrap_matrix(other_op, false); - if (rhs == NULL) { - goto done; - } - - matrix_impl *lhs = self->impl; - enum CrossAxis flavor = classify_cross_axis(lhs, axis.has_axis, axis.axis); - if (flavor == CROSS_INVALID) { - PyErr_SetString( - PyExc_ValueError, - "cross requires a 2D or 3D vector or Nx2 or 2xN or Nx3 or 3xN matrix"); - goto done; + if (!impl_check_acquired(impl, true)) { + return -1; } - if (flavor == CROSS_SCALAR_2D_1x2 || flavor == CROSS_SCALAR_2D_2x1) { - if (rhs->size != 2) { - PyErr_Format(PyExc_ValueError, - "cross: 2D vector lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; + matrix_impl *mask = NULL; + if (where_op != NULL && where_op != Py_None) { + mask = unwrap_mask(where_op, impl); + if (mask == NULL) { + return -1; } - const double *a = lhs->data; - const double *b = rhs->data; - result = PyFloat_FromDouble(a[0] * b[1] - a[1] * b[0]); - goto done; } - if (flavor == CROSS_SCALAR_3D_1x3 || flavor == CROSS_SCALAR_3D_3x1) { - if (rhs->size != 3) { - PyErr_Format(PyExc_ValueError, - "cross: 3D vector lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; - } - matrix_impl *out = impl_new(lhs->rows, lhs->columns); - if (out == NULL) { - goto done; + + // Defensive: the public constructors reject zero-size matrices, but guard + // each axis so a future empty-capable path can't read p[0] out of bounds. + // (The masked kernels never index p[0] unconditionally, so an all-excluded + // group is fine -- they return the -1 "no argument" sentinel instead.) + const char *empty_error = "arg-reduction of an empty matrix is undefined"; + + if (!axis.has_axis) { + if (impl->size == 0) { + PyErr_SetString(PyExc_ValueError, empty_error); + return -1; } - const double *a = lhs->data; - const double *b = rhs->data; - out->data[0] = a[1] * b[2] - a[2] * b[1]; - out->data[1] = a[2] * b[0] - a[0] * b[2]; - out->data[2] = a[0] * b[1] - a[1] * b[0]; - result = wrap_impl_or_free(out); - goto done; + Py_ssize_t idx = mask ? argextreme_masked_ewise(impl, mask, want_max) + : argextreme_ewise(impl, want_max); + *out_op = PyLong_FromSsize_t(idx); + return *out_op == NULL ? -1 : 0; } - if (flavor == CROSS_ROWS_2D_Nx2) { - const size_t N = lhs->rows; - const bool same_shape = - (lhs->rows == rhs->rows && lhs->columns == rhs->columns); - const bool broadcast = - (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); - if (!same_shape && !broadcast) { - PyErr_Format(PyExc_ValueError, - "cross: Nx2 batch lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; + if (axis.axis == 0) { + if (impl->rows == 0) { + PyErr_SetString(PyExc_ValueError, empty_error); + return -1; } - matrix_impl *out = impl_new(N, 1); - if (out == NULL) { - goto done; + matrix_impl *vector = impl_new(1, impl->columns); + if (vector == NULL) { + return -1; } - const double *a = lhs->data; - double *dst = out->data; - if (same_shape) { - const double *b = rhs->data; - for (size_t i = 0; i < N; ++i) { - double ax = *a++; - double ay = *a++; - double bx = *b++; - double by = *b++; - *dst++ = ax * by - ay * bx; - } + if (mask) { + argextreme_masked_columnwise(impl, mask, want_max, vector); } else { - const double bx = rhs->data[0]; - const double by = rhs->data[1]; - for (size_t i = 0; i < N; ++i) { - double ax = *a++; - double ay = *a++; - *dst++ = ax * by - ay * bx; - } - } - result = wrap_impl_or_free(out); - goto done; - } - - if (flavor == CROSS_COLS_2D_2xN) { - const size_t N = lhs->columns; - const bool same_shape = - (lhs->rows == rhs->rows && lhs->columns == rhs->columns); - const bool broadcast = - (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); - if (!same_shape && !broadcast) { - PyErr_Format(PyExc_ValueError, - "cross: 2xN batch lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; - } - matrix_impl *out = impl_new(1, N); - if (out == NULL) { - goto done; + argextreme_columnwise(impl, want_max, vector); } - const double *ax_row = lhs->data; - const double *ay_row = lhs->data + N; - double *dst = out->data; - if (same_shape) { - const double *bx_row = rhs->data; - const double *by_row = rhs->data + N; - for (size_t j = 0; j < N; ++j) { - *dst++ = ax_row[j] * by_row[j] - ay_row[j] * bx_row[j]; - } + if (as_matrix) { + *out_op = (PyObject *)wrap_impl_or_free(vector); } else { - const double bx = rhs->data[0]; - const double by = rhs->data[1]; - for (size_t j = 0; j < N; ++j) { - *dst++ = ax_row[j] * by - ay_row[j] * bx; - } + *out_op = argextreme_indices_to_list(vector); + impl_free(vector); } - result = wrap_impl_or_free(out); - goto done; + return *out_op == NULL ? -1 : 0; } - if (flavor == CROSS_ROWS_3D_Nx3) { - const size_t N = lhs->rows; - const bool same_shape = - (lhs->rows == rhs->rows && lhs->columns == rhs->columns); - const bool broadcast = - (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); - if (!same_shape && !broadcast) { - PyErr_Format(PyExc_ValueError, - "cross: Nx3 batch lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; - } - matrix_impl *out = impl_new(N, 3); - if (out == NULL) { - goto done; - } - const double *a = lhs->data; - double *dst = out->data; - if (same_shape) { - const double *b = rhs->data; - for (size_t i = 0; i < N; ++i) { - double ax = a[0], ay = a[1], az = a[2]; - double bx = b[0], by = b[1], bz = b[2]; - dst[0] = ay * bz - az * by; - dst[1] = az * bx - ax * bz; - dst[2] = ax * by - ay * bx; - a += 3; - b += 3; - dst += 3; - } - } else { - const double bx = rhs->data[0]; - const double by = rhs->data[1]; - const double bz = rhs->data[2]; - for (size_t i = 0; i < N; ++i) { - double ax = a[0], ay = a[1], az = a[2]; - dst[0] = ay * bz - az * by; - dst[1] = az * bx - ax * bz; - dst[2] = ax * by - ay * bx; - a += 3; - dst += 3; - } - } - result = wrap_impl_or_free(out); - goto done; + if (impl->columns == 0) { + PyErr_SetString(PyExc_ValueError, empty_error); + return -1; + } + matrix_impl *vector = impl_new(impl->rows, 1); + if (vector == NULL) { + return -1; + } + if (mask) { + argextreme_masked_rowwise(impl, mask, want_max, vector); + } else { + argextreme_rowwise(impl, want_max, vector); + } + if (as_matrix) { + *out_op = (PyObject *)wrap_impl_or_free(vector); + } else { + *out_op = argextreme_indices_to_list(vector); + impl_free(vector); } + return *out_op == NULL ? -1 : 0; +} - if (flavor == CROSS_COLS_3D_3xN) { - const size_t N = lhs->columns; - const bool same_shape = - (lhs->rows == rhs->rows && lhs->columns == rhs->columns); - const bool broadcast = - (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); - if (!same_shape && !broadcast) { - PyErr_Format(PyExc_ValueError, - "cross: 3xN batch lhs %zux%zu incompatible with rhs %zux%zu", - lhs->rows, lhs->columns, rhs->rows, rhs->columns); - goto done; - } - matrix_impl *out = impl_new(3, N); - if (out == NULL) { - goto done; - } - const double *ax_row = lhs->data; - const double *ay_row = lhs->data + N; - const double *az_row = lhs->data + 2 * N; - double *dx_row = out->data; - double *dy_row = out->data + N; - double *dz_row = out->data + 2 * N; - if (same_shape) { - const double *bx_row = rhs->data; - const double *by_row = rhs->data + N; - const double *bz_row = rhs->data + 2 * N; - for (size_t j = 0; j < N; ++j) { - double ax = ax_row[j], ay = ay_row[j], az = az_row[j]; - double bx = bx_row[j], by = by_row[j], bz = bz_row[j]; - dx_row[j] = ay * bz - az * by; - dy_row[j] = az * bx - ax * bz; - dz_row[j] = ax * by - ay * bx; - } - } else { - const double bx = rhs->data[0]; - const double by = rhs->data[1]; - const double bz = rhs->data[2]; - for (size_t j = 0; j < N; ++j) { - double ax = ax_row[j], ay = ay_row[j], az = az_row[j]; - dx_row[j] = ay * bz - az * by; - dy_row[j] = az * bx - ax * bz; - dz_row[j] = ax * by - ay * bx; - } - } - result = wrap_impl_or_free(out); - goto done; +#define MATRIX_ARGEXTREME(name, want_max_val) \ + static PyObject *Matrix_##name##_method(PyObject *op, PyObject *args, \ + PyObject *kwds) { \ + PyObject *out = NULL; \ + PyObject *axis_obj = NULL; \ + PyObject *where_obj = NULL; \ + int as_matrix = 0; \ + static char *kwlist[] = {"axis", "where", "as_matrix", NULL}; \ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOp", kwlist, &axis_obj, \ + &where_obj, &as_matrix)) { \ + return NULL; \ + } \ + AxisArg axis; \ + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { \ + return NULL; \ + } \ + if (Matrix_argextreme(op, axis, where_obj, &out, want_max_val, \ + (bool)as_matrix) < 0) { \ + return NULL; \ + } \ + return out; \ } - PyErr_SetString(PyExc_RuntimeError, - "internal: unhandled CrossAxis in Matrix_cross"); +MATRIX_ARGEXTREME(argmin, false) +MATRIX_ARGEXTREME(argmax, true) -done: - IMPL_DECREF(rhs); - return result; -} +enum BroadcastShape { BCAST_NONE = 0, BCAST_ROW, BCAST_COL }; + +/* -------------------------------------------------------------------------- + Two-operand aggregate family X-macro template. + + BOC_2AGG_OPS is the single source of truth for the two-operand aggregate + op set. It is stamped in three places: + 1. impl__total — flat traversal returning a scalar. + 2. impl__rowwise — per-row reduction writing Mx1 output. + 3. impl__columnwise — per-column accumulation writing 1xN output. + REQUIRES caller-zeroed output buffer + (dispatcher uses impl_new -> PyMem_RawCalloc). + + Each walker carries the per-shape switch at the TOP of the body, with + three specialised inner loops (NONE / ROW / COL). Moving the switch + inside the inner loop would defeat contraction; the per-shape pointer + arithmetic and broadcast behaviour are intentionally different. + The dispatcher (Matrix_vecdot) canonicalises operands so the matrix is + always the LHS and the vector is always the RHS before calling the + helpers. + + Names in scope inside STEP: + lhs (double) — current left-hand-side value + rhs (double) — current right-hand-side value (in BCAST_COL: the + per-row scalar, hoisted out of the inner loop) + agg (double) — accumulator (for total/rowwise it is a local; + for columnwise the per-iteration cell of the output + vector is loaded into `agg`, STEP runs, then the + cell is written back) + + Add an op: append one X(NAME, INIT, STEP) row to BOC_2AGG_OPS and the + three impl__* kernels are stamped automatically. Wire it into a + Python entry point separately (mirrors Matrix_vecdot). + + See `.github/skills/commenting-c-and-python/SKILL.md` for the + "Exception: X-macro descriptor tables" convention. + -------------------------------------------------------------------------- */ +#define BOC_2AGG_OPS(X) \ + /* name init step (lhs, rhs in scope; agg accumulator) */ \ + X(vecdot, 0.0, agg += lhs * rhs) + +#define DEFINE_2AGG_TOTAL(NAME, INIT, STEP) \ + BOC_CANARY_NOINLINE \ + static double impl_##NAME##_total(const matrix_impl *lm, \ + const matrix_impl *rm, \ + enum BroadcastShape shape) { \ + double agg = (INIT); \ + switch (shape) { \ + case BCAST_NONE: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + for (size_t i = 0; i < lm->size; ++i, ++lp, ++rp) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + STEP; \ + } \ + break; \ + } \ + case BCAST_ROW: { \ + const double *lp = lm->data; \ + const size_t M = lm->rows; \ + const size_t N = lm->columns; \ + for (size_t r = 0; r < M; ++r) { \ + const double *rp = rm->data; \ + for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + STEP; \ + } \ + } \ + break; \ + } \ + case BCAST_COL: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + const size_t M = lm->rows; \ + const size_t N = lm->columns; \ + for (size_t r = 0; r < M; ++r) { \ + const double rhs = *rp++; \ + for (size_t c = 0; c < N; ++c, ++lp) { \ + const double lhs = *lp; \ + STEP; \ + } \ + } \ + break; \ + } \ + } \ + return agg; \ + } + +#define DEFINE_2AGG_ROWWISE(NAME, INIT, STEP) \ + BOC_CANARY_NOINLINE \ + static void impl_##NAME##_rowwise( \ + const matrix_impl *lm, const matrix_impl *rm, matrix_impl *out_Mx1, \ + enum BroadcastShape shape) { \ + const size_t M = lm->rows; \ + const size_t N = lm->columns; \ + double *out_ptr = out_Mx1->data; \ + switch (shape) { \ + case BCAST_NONE: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + for (size_t r = 0; r < M; ++r, ++out_ptr) { \ + double agg = (INIT); \ + for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + STEP; \ + } \ + *out_ptr = agg; \ + } \ + break; \ + } \ + case BCAST_ROW: { \ + const double *lp = lm->data; \ + for (size_t r = 0; r < M; ++r, ++out_ptr) { \ + const double *rp = rm->data; \ + double agg = (INIT); \ + for (size_t c = 0; c < N; ++c, ++lp, ++rp) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + STEP; \ + } \ + *out_ptr = agg; \ + } \ + break; \ + } \ + case BCAST_COL: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + for (size_t r = 0; r < M; ++r, ++out_ptr) { \ + const double rhs = *rp++; \ + double agg = (INIT); \ + for (size_t c = 0; c < N; ++c, ++lp) { \ + const double lhs = *lp; \ + STEP; \ + } \ + *out_ptr = agg; \ + } \ + break; \ + } \ + } \ + } + +/* Columnwise: the per-output-cell accumulator IS the output vector slot + itself. Loading `*out_ptr` into a local `agg`, running STEP, and + writing back keeps STEP uniform across all three walkers; the compiler + collapses the load/store pair to a single += against memory. The + caller must hand in a zero-initialised buffer (impl_new -> calloc); + a recycled non-zero buffer would silently produce wrong sums because + STEP accumulates with `+=`. */ +#define DEFINE_2AGG_COLUMNWISE(NAME, INIT, STEP) \ + BOC_CANARY_NOINLINE \ + static void impl_##NAME##_columnwise( \ + const matrix_impl *lm, const matrix_impl *rm, matrix_impl *out_1xN, \ + enum BroadcastShape shape) { \ + const size_t M = lm->rows; \ + const size_t N = lm->columns; \ + (void)(INIT); \ + switch (shape) { \ + case BCAST_NONE: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + for (size_t r = 0; r < M; ++r) { \ + double *out_ptr = out_1xN->data; \ + for (size_t c = 0; c < N; ++c, ++lp, ++rp, ++out_ptr) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + double agg = *out_ptr; \ + STEP; \ + *out_ptr = agg; \ + } \ + } \ + break; \ + } \ + case BCAST_ROW: { \ + const double *lp = lm->data; \ + for (size_t r = 0; r < M; ++r) { \ + const double *rp = rm->data; \ + double *out_ptr = out_1xN->data; \ + for (size_t c = 0; c < N; ++c, ++lp, ++rp, ++out_ptr) { \ + const double lhs = *lp; \ + const double rhs = *rp; \ + double agg = *out_ptr; \ + STEP; \ + *out_ptr = agg; \ + } \ + } \ + break; \ + } \ + case BCAST_COL: { \ + const double *lp = lm->data; \ + const double *rp = rm->data; \ + for (size_t r = 0; r < M; ++r) { \ + const double rhs = *rp++; \ + double *out_ptr = out_1xN->data; \ + for (size_t c = 0; c < N; ++c, ++lp, ++out_ptr) { \ + const double lhs = *lp; \ + double agg = *out_ptr; \ + STEP; \ + *out_ptr = agg; \ + } \ + } \ + break; \ + } \ + } \ + } + +#define X(N, I, S) DEFINE_2AGG_TOTAL(N, I, S) +BOC_2AGG_OPS(X) +#undef X + +#define X(N, I, S) DEFINE_2AGG_ROWWISE(N, I, S) +BOC_2AGG_OPS(X) +#undef X + +#define X(N, I, S) DEFINE_2AGG_COLUMNWISE(N, I, S) +BOC_2AGG_OPS(X) +#undef X + +/// @brief Axis-aware inner product: sum of element-wise products. +/// @details The canonicalisation swap rearranges the dispatch-argument +/// pointers ``mat_arg`` / ``vec_arg`` so the helpers always see +/// ``(matrix, vector)``. The refcounted ``rhs`` from +/// ``unwrap_matrix`` is preserved unchanged so the IMPL_DECREF at +/// ``done`` always matches the single INCREF. ``self->impl`` is +/// NOT refcount-paired here (mirrors Matrix_transpose). +static PyObject *Matrix_vecdot(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *other = NULL; + PyObject *axis = NULL; + PyObject *result = NULL; + matrix_impl *rhs = NULL; + + static char *kwlist[] = {"", "axis", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", kwlist, &other, &axis)) { + return NULL; + } + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + rhs = unwrap_matrix(other, false); + if (rhs == NULL) { + goto done; + } + + matrix_impl *lhs = self->impl; + matrix_impl *mat_arg = lhs; + matrix_impl *vec_arg = rhs; + enum BroadcastShape shape; + + if (lhs->rows == rhs->rows && lhs->columns == rhs->columns) { + shape = BCAST_NONE; + } else if (lhs->rows == rhs->rows && + (lhs->columns == 1 || rhs->columns == 1)) { + shape = BCAST_COL; + if (lhs->columns == 1) { + mat_arg = rhs; + vec_arg = lhs; + } + } else if (lhs->columns == rhs->columns && + (lhs->rows == 1 || rhs->rows == 1)) { + shape = BCAST_ROW; + if (lhs->rows == 1) { + mat_arg = rhs; + vec_arg = lhs; + } + } else if ((lhs->rows == 1 || lhs->columns == 1) && + (rhs->rows == 1 || rhs->columns == 1) && lhs->size == rhs->size) { + shape = BCAST_NONE; + } else { + PyErr_Format(PyExc_ValueError, + "vecdot: lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, + lhs->columns, rhs->rows, rhs->columns); + goto done; + } + + AxisArg axis_arg; + if (parse_validate_normalise_axis(axis, &axis_arg) < 0) { + goto done; + } + + if (!axis_arg.has_axis) { + result = PyFloat_FromDouble(impl_vecdot_total(mat_arg, vec_arg, shape)); + } else if (axis_arg.axis == 0) { + matrix_impl *out = impl_new(1, mat_arg->columns); + if (out != NULL) { + impl_vecdot_columnwise(mat_arg, vec_arg, out, shape); + result = wrap_impl_or_free(out); + } + } else { + // axis == 1 (row-wise): parse_validate_normalise_axis restricts axis + // to {0, 1}. + matrix_impl *out = impl_new(mat_arg->rows, 1); + if (out != NULL) { + impl_vecdot_rowwise(mat_arg, vec_arg, out, shape); + result = wrap_impl_or_free(out); + } + } + +done: + IMPL_DECREF(rhs); + return result; +} + +/// @brief A classified broadcast operand: scalar or same-shape matrix. +/// @details ``full`` is an INCREF'd matrix with the same shape as the +/// receiver, and the caller must IMPL_DECREF it; when ``full`` is +/// NULL the operand is the scalar in ``scalar``. Shared by ``fma`` +/// and ``scaled_add`` to classify a multiplier / scale operand. +typedef struct { + double scalar; + matrix_impl *full; +} broadcast_operand; + +/// @brief Materialise a row- or column-vector broadcast into a full buffer. +/// @details Expands a ``1xN`` row vector (``is_row``) or an ``Mx1`` column +/// vector into a fresh contiguous ``rows`` x ``columns`` matrix so +/// the contiguous fma kernel keeps its unit stride (and hardware +/// FMA). The returned impl has refcount 0; the caller INCREFs it. +static matrix_impl *impl_broadcast_vector(const matrix_impl *vec, size_t rows, + size_t columns, bool is_row) { + matrix_impl *out = impl_new(rows, columns); + if (out == NULL) { + return NULL; + } + double *dst = out->data; + const double *v = vec->data; + if (is_row) { + for (size_t r = 0; r < rows; ++r) { + for (size_t c = 0; c < columns; ++c) { + *dst++ = v[c]; + } + } + } else { + for (size_t r = 0; r < rows; ++r) { + for (size_t c = 0; c < columns; ++c) { + *dst++ = v[r]; + } + } + } + return out; +} + +/// @brief Classify a broadcast operand as a scalar or a same-shape matrix. +/// @details A Matrix whose size is 1 is treated as a scalar so a ``1x1`` +/// broadcasts like a number (the in-house scalar rule used +/// elsewhere); a same-shape Matrix is INCREF'd into ``out->full``. +/// A ``1xN`` row vector (matching self's columns) or an ``Mx1`` +/// column vector (matching self's rows) is materialised into a fresh +/// same-shape buffer stored in ``out->full`` so the contiguous kernel +/// keeps its unit stride; any other matrix shape raises ``ValueError`` +/// naming both shapes. Any real number is accepted as a scalar. +/// @param op_name Method name used as the error-message prefix (e.g. ``fma``). +/// @param operand The operand object to classify. +/// @param self_impl The receiver's impl, used for the shape check. +/// @param what Operand name used in error messages. +/// @param out Output operand descriptor (zeroed before classification). +/// @return 0 on success; -1 with an exception set on failure. +static int classify_broadcast_operand(const char *op_name, PyObject *operand, + const matrix_impl *self_impl, + const char *what, + broadcast_operand *out) { + out->scalar = 0.0; + out->full = NULL; + + if (Py_TYPE(operand) == LOCAL_STATE->matrix_type) { + matrix_impl *impl = ((MatrixObject *)operand)->impl; + if (!impl_check_acquired(impl, true)) { + return -1; + } + if (impl->size == 1) { + out->scalar = impl->data[0]; + return 0; + } + if (impl->rows == self_impl->rows && impl->columns == self_impl->columns) { + IMPL_INCREF(impl); + out->full = impl; + return 0; + } + if (impl->rows == 1 && impl->columns == self_impl->columns) { + matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, + self_impl->columns, true); + if (full == NULL) { + return -1; + } + IMPL_INCREF(full); + out->full = full; + return 0; + } + if (impl->columns == 1 && impl->rows == self_impl->rows) { + matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, + self_impl->columns, false); + if (full == NULL) { + return -1; + } + IMPL_INCREF(full); + out->full = full; + return 0; + } + PyErr_Format(PyExc_ValueError, + "%s: %s shape %zux%zu incompatible with self %zux%zu", op_name, + what, impl->rows, impl->columns, self_impl->rows, + self_impl->columns); + return -1; + } + + if (unwrap_double(operand, &out->scalar)) { + return 0; + } + if (PyErr_Occurred()) { + return -1; + } + + PyErr_Format(PyExc_TypeError, "%s: %s must be a Matrix or a real number", + op_name, what); + return -1; +} + +/// @brief Fused multiply-add kernel: ``out = a*b + c`` with one rounding. +/// @details ``b_step`` / ``c_step`` are 1 for a full same-shape operand and +/// 0 to repeat a scalar (the pointer then aims at a single local +/// ``double``). Each output cell reads only its own index, so the +/// in-place case (``out == a``) and ``b`` / ``c`` aliasing ``a`` +/// are all safe. ``fma()`` rounds the product once, unlike +/// ``a*b + c`` which rounds twice under ``-ffp-contract=off``. +/// BOC_FMA_MULTIVERSION clones this kernel for hardware FMA on +/// glibc x86-64 (see the macro definition for why glibc-only). +BOC_CANARY_NOINLINE +BOC_FMA_MULTIVERSION +static void impl_fma(const double *a, const double *b, size_t b_step, + const double *c, size_t c_step, double *out, size_t n) { + for (size_t i = 0; i < n; ++i, ++a, b += b_step, c += c_step, ++out) { + *out = fma(*a, *b, *c); + } +} + +/// @brief Fused multiply-add: single-rounding ``self*b + c``. +/// @details ``b`` and ``c`` may each be a same-shape matrix, a ``1x1`` +/// matrix, a ``1xN`` row vector or ``Mx1`` column vector that +/// broadcasts against ``self``, or a scalar; any other matrix shape +/// raises ``ValueError``. Both operands are validated before any +/// allocation, so a rejected operand leaves ``self`` untouched. +/// With ``in_place=True`` the result is written into ``self`` and +/// ``self`` is returned. +static PyObject *Matrix_fma(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *b_op = NULL; + PyObject *c_op = NULL; + int in_place = 0; + PyObject *out_op = NULL; + broadcast_operand bop = {0.0, NULL}; + broadcast_operand cop = {0.0, NULL}; + + static char *kwlist[] = {"", "", "in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &b_op, &c_op, + &in_place)) { + return NULL; + } + + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + if (classify_broadcast_operand("fma", b_op, self->impl, "b", &bop) < 0) { + goto done; + } + if (classify_broadcast_operand("fma", c_op, self->impl, "c", &cop) < 0) { + goto done; + } + + matrix_impl *out = set_output(self->impl, op, &out_op, NULL, in_place); + if (out == NULL) { + out_op = NULL; + goto done; + } + + double b_scalar = bop.scalar; + double c_scalar = cop.scalar; + const double *b_ptr = bop.full != NULL ? bop.full->data : &b_scalar; + const double *c_ptr = cop.full != NULL ? cop.full->data : &c_scalar; + size_t b_step = bop.full != NULL ? 1 : 0; + size_t c_step = cop.full != NULL ? 1 : 0; + + impl_fma(self->impl->data, b_ptr, b_step, c_ptr, c_step, out->data, + self->impl->size); + +done: + IMPL_DECREF(bop.full); + IMPL_DECREF(cop.full); + return out_op; +} + +/// @brief Two-rounding scaled-add kernel: ``out = y + s * x``. +/// @details Deliberately NOT fused: the product ``s[i] * x[i]`` is rounded to +/// double first, then the sum with ``y[i]`` is rounded again, so the +/// result is bit-for-bit identical to the naive ``y[i] + s * x[i]`` +/// expression in IEEE-754 double (two roundings). This is the +/// two-rounding complement to ``impl_fma``'s single rounding; the +/// ``-ffp-contract=off`` build flag guarantees the two statements +/// are never contracted back into an fma. ``s_step`` is 1 for a full +/// same-shape scale buffer and 0 to repeat a broadcast scalar (the +/// pointer then aims at a single local ``double``); ``x`` and ``y`` +/// are the same shape (unit stride). Each output cell reads only its +/// own index, so the in-place case (``out == y``) and ``x`` / ``s`` +/// aliasing ``y`` are all safe. +BOC_CANARY_NOINLINE +static void impl_scaled_add(const double *y, const double *s, size_t s_step, + const double *x, double *out, size_t n) { + for (size_t i = 0; i < n; ++i, ++y, s += s_step, ++x, ++out) { + const double prod = (*s) * (*x); + *out = *y + prod; + } +} + +/// @brief Scaled add: ``self + s * x`` with two roundings. +/// @details The two-rounding sibling of ``fma``: the product ``s * x`` is +/// rounded to double, then the sum with ``self`` is rounded again, so +/// the result is bit-for-bit identical to ``self + s * x`` and +/// distinct from the single-rounded ``fma``. ``s`` is the scale and +/// may be a scalar, a ``1x1`` matrix, a ``1xN`` row vector or ``Mx1`` +/// column vector that broadcasts against ``self``, or a same-shape +/// matrix (the same operand rules as ``fma``'s multiplier); ``x`` is +/// a same-shape matrix. Both operands are validated before any +/// allocation, so a rejected operand leaves ``self`` untouched (the +/// validate-then-write contract shared with ``put``). With +/// ``in_place=True`` the result is written into ``self``'s existing +/// buffer (allocating nothing) and ``self`` is returned; otherwise a +/// fresh matrix is returned. ``s`` and ``x`` may alias ``self``. +static PyObject *Matrix_scaled_add(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *s_op = NULL; + PyObject *x_op = NULL; + int in_place = 0; + PyObject *out_op = NULL; + broadcast_operand sop = {0.0, NULL}; + matrix_impl *x = NULL; + + static char *kwlist[] = {"", "", "in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &s_op, &x_op, + &in_place)) { + return NULL; + } + + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + if (classify_broadcast_operand("scaled_add", s_op, self->impl, "s", &sop) < + 0) { + return NULL; + } + + x = unwrap_matrix(x_op, false); + if (x == NULL) { + goto done; + } + + matrix_impl *y = self->impl; + if (x->rows != y->rows || x->columns != y->columns) { + PyErr_Format(PyExc_ValueError, + "scaled_add: x shape %zux%zu does not match self %zux%zu", + x->rows, x->columns, y->rows, y->columns); + goto done; + } + + matrix_impl *out = set_output(y, op, &out_op, NULL, in_place); + if (out == NULL) { + out_op = NULL; + goto done; + } + + double s_scalar = sop.scalar; + const double *s_ptr = sop.full != NULL ? sop.full->data : &s_scalar; + size_t s_step = sop.full != NULL ? 1 : 0; + + impl_scaled_add(y->data, s_ptr, s_step, x->data, out->data, y->size); + +done: + IMPL_DECREF(sop.full); + IMPL_DECREF(x); + return out_op; +} + +/// @details ``has_axis`` / ``explicit_axis`` carry an optional caller- +/// supplied axis (already normalised to 0 or 1 by +/// ``parse_validate_normalise_axis``). For the doubly-valid +/// ``2x2`` / ``3x3`` shapes ``explicit_axis`` picks the +/// orientation (``0`` -> columns, default -> rows). For all +/// other shapes only one orientation is valid; supplying an +/// ``explicit_axis`` that contradicts that orientation returns +/// ``CROSS_INVALID`` so the caller raises rather than running +/// the wrong kernel. Returns ``CROSS_INVALID`` for any shape +/// that has no valid cross-product interpretation. +enum CrossAxis { + CROSS_SCALAR_2D_1x2, + CROSS_SCALAR_2D_2x1, + CROSS_ROWS_2D_Nx2, + CROSS_COLS_2D_2xN, + CROSS_SCALAR_3D_1x3, + CROSS_SCALAR_3D_3x1, + CROSS_ROWS_3D_Nx3, + CROSS_COLS_3D_3xN, + CROSS_INVALID +}; + +static enum CrossAxis classify_cross_axis(const matrix_impl *impl, + bool has_axis, int explicit_axis) { + const size_t M = impl->rows; + const size_t N = impl->columns; + if (M == 2 && N == 2) { + return (has_axis && explicit_axis == 0) ? CROSS_COLS_2D_2xN + : CROSS_ROWS_2D_Nx2; + } + if (M == 3 && N == 3) { + return (has_axis && explicit_axis == 0) ? CROSS_COLS_3D_3xN + : CROSS_ROWS_3D_Nx3; + } + if (M == 1 && N == 2) { + if (has_axis && explicit_axis == 0) { + return CROSS_INVALID; + } + return CROSS_SCALAR_2D_1x2; + } + if (M == 1 && N == 3) { + if (has_axis && explicit_axis == 0) { + return CROSS_INVALID; + } + return CROSS_SCALAR_3D_1x3; + } + if (M == 2 && N == 1) { + if (has_axis && explicit_axis == 1) { + return CROSS_INVALID; + } + return CROSS_SCALAR_2D_2x1; + } + if (M == 3 && N == 1) { + if (has_axis && explicit_axis == 1) { + return CROSS_INVALID; + } + return CROSS_SCALAR_3D_3x1; + } + if (N == 2 && M != 3) { + if (has_axis && explicit_axis == 0) { + return CROSS_INVALID; + } + return CROSS_ROWS_2D_Nx2; + } + if (M == 2 && N != 3) { + if (has_axis && explicit_axis == 1) { + return CROSS_INVALID; + } + return CROSS_COLS_2D_2xN; + } + if (N == 3 && M != 2) { + if (has_axis && explicit_axis == 0) { + return CROSS_INVALID; + } + return CROSS_ROWS_3D_Nx3; + } + if (M == 3 && N != 2) { + if (has_axis && explicit_axis == 1) { + return CROSS_INVALID; + } + return CROSS_COLS_3D_3xN; + } + if (N == 2) { + return CROSS_ROWS_2D_Nx2; + } + if (M == 2) { + return CROSS_COLS_2D_2xN; + } + if (N == 3) { + return CROSS_ROWS_3D_Nx3; + } + if (M == 3) { + return CROSS_COLS_3D_3xN; + } + return CROSS_INVALID; +} + +/// @brief 2D / 3D cross product against another vector or batch. +/// @details Five paths share one dispatcher. For 1x2 / 2x1 inputs the +/// result is the scalar z-component +/// ``self.x * other.y - self.y * other.x`` as a Python float; +/// for 1x3 / 3x1 the result is a same-shape Matrix preserving +/// ``self``'s orientation. For Nx2 / 2xN batches the result is +/// a per-vector scalar collected in a Mx1 (rows) or 1xN (cols) +/// Matrix; for Nx3 / 3xN batches the result is a same-shape +/// Matrix of per-vector cross products. The ``axis`` keyword +/// disambiguates 2x2 / 3x3 squares (default: rows; ``axis=0`` +/// forces columns). For scalar inputs ``other``'s orientation +/// is irrelevant; for batch inputs ``other`` must have the same +/// shape as ``self``. +static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *other_op = NULL; + PyObject *axis_obj = NULL; + PyObject *result = NULL; + matrix_impl *rhs = NULL; + + static char *kwlist[] = {"", "axis", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", kwlist, &other_op, + &axis_obj)) { + return NULL; + } + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + return NULL; + } + + rhs = unwrap_matrix(other_op, false); + if (rhs == NULL) { + goto done; + } + + matrix_impl *lhs = self->impl; + enum CrossAxis flavor = classify_cross_axis(lhs, axis.has_axis, axis.axis); + if (flavor == CROSS_INVALID) { + PyErr_SetString( + PyExc_ValueError, + "cross requires a 2D or 3D vector or Nx2 or 2xN or Nx3 or 3xN matrix"); + goto done; + } + + if (flavor == CROSS_SCALAR_2D_1x2 || flavor == CROSS_SCALAR_2D_2x1) { + if (rhs->size != 2) { + PyErr_Format(PyExc_ValueError, + "cross: 2D vector lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + const double *a = lhs->data; + const double *b = rhs->data; + result = PyFloat_FromDouble(a[0] * b[1] - a[1] * b[0]); + goto done; + } + if (flavor == CROSS_SCALAR_3D_1x3 || flavor == CROSS_SCALAR_3D_3x1) { + if (rhs->size != 3) { + PyErr_Format(PyExc_ValueError, + "cross: 3D vector lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + matrix_impl *out = impl_new(lhs->rows, lhs->columns); + if (out == NULL) { + goto done; + } + const double *a = lhs->data; + const double *b = rhs->data; + out->data[0] = a[1] * b[2] - a[2] * b[1]; + out->data[1] = a[2] * b[0] - a[0] * b[2]; + out->data[2] = a[0] * b[1] - a[1] * b[0]; + result = wrap_impl_or_free(out); + goto done; + } + + if (flavor == CROSS_ROWS_2D_Nx2) { + const size_t N = lhs->rows; + const bool same_shape = + (lhs->rows == rhs->rows && lhs->columns == rhs->columns); + const bool broadcast = + (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); + if (!same_shape && !broadcast) { + PyErr_Format(PyExc_ValueError, + "cross: Nx2 batch lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + matrix_impl *out = impl_new(N, 1); + if (out == NULL) { + goto done; + } + const double *a = lhs->data; + double *dst = out->data; + if (same_shape) { + const double *b = rhs->data; + for (size_t i = 0; i < N; ++i) { + double ax = *a++; + double ay = *a++; + double bx = *b++; + double by = *b++; + *dst++ = ax * by - ay * bx; + } + } else { + const double bx = rhs->data[0]; + const double by = rhs->data[1]; + for (size_t i = 0; i < N; ++i) { + double ax = *a++; + double ay = *a++; + *dst++ = ax * by - ay * bx; + } + } + result = wrap_impl_or_free(out); + goto done; + } + + if (flavor == CROSS_COLS_2D_2xN) { + const size_t N = lhs->columns; + const bool same_shape = + (lhs->rows == rhs->rows && lhs->columns == rhs->columns); + const bool broadcast = + (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); + if (!same_shape && !broadcast) { + PyErr_Format(PyExc_ValueError, + "cross: 2xN batch lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + matrix_impl *out = impl_new(1, N); + if (out == NULL) { + goto done; + } + const double *ax_row = lhs->data; + const double *ay_row = lhs->data + N; + double *dst = out->data; + if (same_shape) { + const double *bx_row = rhs->data; + const double *by_row = rhs->data + N; + for (size_t j = 0; j < N; ++j) { + *dst++ = ax_row[j] * by_row[j] - ay_row[j] * bx_row[j]; + } + } else { + const double bx = rhs->data[0]; + const double by = rhs->data[1]; + for (size_t j = 0; j < N; ++j) { + *dst++ = ax_row[j] * by - ay_row[j] * bx; + } + } + result = wrap_impl_or_free(out); + goto done; + } + + if (flavor == CROSS_ROWS_3D_Nx3) { + const size_t N = lhs->rows; + const bool same_shape = + (lhs->rows == rhs->rows && lhs->columns == rhs->columns); + const bool broadcast = + (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); + if (!same_shape && !broadcast) { + PyErr_Format(PyExc_ValueError, + "cross: Nx3 batch lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + matrix_impl *out = impl_new(N, 3); + if (out == NULL) { + goto done; + } + const double *a = lhs->data; + double *dst = out->data; + if (same_shape) { + const double *b = rhs->data; + for (size_t i = 0; i < N; ++i) { + double ax = a[0], ay = a[1], az = a[2]; + double bx = b[0], by = b[1], bz = b[2]; + dst[0] = ay * bz - az * by; + dst[1] = az * bx - ax * bz; + dst[2] = ax * by - ay * bx; + a += 3; + b += 3; + dst += 3; + } + } else { + const double bx = rhs->data[0]; + const double by = rhs->data[1]; + const double bz = rhs->data[2]; + for (size_t i = 0; i < N; ++i) { + double ax = a[0], ay = a[1], az = a[2]; + dst[0] = ay * bz - az * by; + dst[1] = az * bx - ax * bz; + dst[2] = ax * by - ay * bx; + a += 3; + dst += 3; + } + } + result = wrap_impl_or_free(out); + goto done; + } + + if (flavor == CROSS_COLS_3D_3xN) { + const size_t N = lhs->columns; + const bool same_shape = + (lhs->rows == rhs->rows && lhs->columns == rhs->columns); + const bool broadcast = + (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); + if (!same_shape && !broadcast) { + PyErr_Format(PyExc_ValueError, + "cross: 3xN batch lhs %zux%zu incompatible with rhs %zux%zu", + lhs->rows, lhs->columns, rhs->rows, rhs->columns); + goto done; + } + matrix_impl *out = impl_new(3, N); + if (out == NULL) { + goto done; + } + const double *ax_row = lhs->data; + const double *ay_row = lhs->data + N; + const double *az_row = lhs->data + 2 * N; + double *dx_row = out->data; + double *dy_row = out->data + N; + double *dz_row = out->data + 2 * N; + if (same_shape) { + const double *bx_row = rhs->data; + const double *by_row = rhs->data + N; + const double *bz_row = rhs->data + 2 * N; + for (size_t j = 0; j < N; ++j) { + double ax = ax_row[j], ay = ay_row[j], az = az_row[j]; + double bx = bx_row[j], by = by_row[j], bz = bz_row[j]; + dx_row[j] = ay * bz - az * by; + dy_row[j] = az * bx - ax * bz; + dz_row[j] = ax * by - ay * bx; + } + } else { + const double bx = rhs->data[0]; + const double by = rhs->data[1]; + const double bz = rhs->data[2]; + for (size_t j = 0; j < N; ++j) { + double ax = ax_row[j], ay = ay_row[j], az = az_row[j]; + dx_row[j] = ay * bz - az * by; + dy_row[j] = az * bx - ax * bz; + dz_row[j] = ax * by - ay * bx; + } + } + result = wrap_impl_or_free(out); + goto done; + } + + PyErr_SetString(PyExc_RuntimeError, + "internal: unhandled CrossAxis in Matrix_cross"); + +done: + IMPL_DECREF(rhs); + return result; +} /// @brief Replace every zero entry in @p vector with 1.0. /// @note Contract: @p vector must be a freshly-computed magnitude vector @@ -2600,345 +3938,1065 @@ static void sanitize_divisor(matrix_impl *vector) { } } -/// @brief Normalize @p impl into @p out along the given axis. -/// @details ``axis.has_axis == false`` divides every element by the matrix's -/// total magnitude; ``axis.axis == 0`` divides each column by its -/// own magnitude; ``axis.axis == 1`` divides each row by its own -/// magnitude. The all-zero input case is preserved (see -/// ``sanitize_divisor``). Self-aliasing is supported — pass -/// ``out == impl`` for in-place operation. -static int do_normalize(matrix_impl *impl, AxisArg axis, matrix_impl *out) { - if (!axis.has_axis) { - double m = dispatch_agg_ewise(impl, Magnitude); - if (m == 0.0) { - if (out != impl) { - memcpy(out->data, impl->data, impl->size * sizeof(double)); - } - return 0; +/// @brief Normalize @p impl into @p out along the given axis. +/// @details ``axis.has_axis == false`` divides every element by the matrix's +/// total magnitude; ``axis.axis == 0`` divides each column by its +/// own magnitude; ``axis.axis == 1`` divides each row by its own +/// magnitude. The all-zero input case is preserved (see +/// ``sanitize_divisor``). Self-aliasing is supported — pass +/// ``out == impl`` for in-place operation. +static int do_normalize(matrix_impl *impl, AxisArg axis, matrix_impl *out) { + if (!axis.has_axis) { + double m = dispatch_agg_ewise(impl, Magnitude); + if (m == 0.0) { + if (out != impl) { + memcpy(out->data, impl->data, impl->size * sizeof(double)); + } + return 0; + } + dispatch_bin_scalar(impl, m, out, Divide); + return 0; + } + + if (axis.axis == 0) { + matrix_impl *divisor = impl_new(1, impl->columns); + if (divisor == NULL) { + return -1; + } + dispatch_agg_columnwise(impl, Magnitude, divisor); + sanitize_divisor(divisor); + dispatch_bin_rowwise(impl, divisor, out, Divide); + impl_free(divisor); + return 0; + } + + // axis == 1 (row-wise): parse_validate_normalise_axis restricts axis + // to {0, 1}. + matrix_impl *divisor = impl_new(impl->rows, 1); + if (divisor == NULL) { + return -1; + } + dispatch_agg_rowwise(impl, Magnitude, divisor); + sanitize_divisor(divisor); + dispatch_bin_columnwise(impl, divisor, out, Divide); + impl_free(divisor); + return 0; +} + +static PyObject *Matrix_normalize(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *axis_obj = NULL; + int in_place = 0; + PyObject *out_op = NULL; + + static char *kwlist[] = {"axis", "in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Op", kwlist, &axis_obj, + &in_place)) { + return NULL; + } + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + return NULL; + } + + matrix_impl *out = set_output(self->impl, op, &out_op, NULL, in_place); + if (out == NULL) { + return NULL; + } + if (do_normalize(self->impl, axis, out) < 0) { + Py_DECREF(out_op); + return NULL; + } + return out_op; +} + +enum Vec2Axis { + VEC2_SCALAR_1x2, + VEC2_SCALAR_2x1, + VEC2_ROWS_Nx2, + VEC2_COLS_2xN, + VEC2_INVALID +}; + +/// @brief Classify a matrix as a 2D vector or batch of 2D vectors. +/// @details ``has_axis`` / ``explicit_axis`` carry an optional caller- +/// supplied axis (already normalised to 0 or 1). The ``2x2`` +/// shape is doubly-valid and ``explicit_axis`` picks the +/// orientation (``0`` -> columns, default -> rows). For all +/// other shapes only one orientation is valid; supplying an +/// ``explicit_axis`` that contradicts that orientation returns +/// ``VEC2_INVALID``. Returns ``VEC2_INVALID`` for any shape +/// that is not a 2D vector or Nx2 / 2xN batch. +static enum Vec2Axis classify_vec2_axis(const matrix_impl *impl, bool has_axis, + int explicit_axis) { + const size_t M = impl->rows; + const size_t N = impl->columns; + if (M == 1 && N == 2) { + if (has_axis && explicit_axis == 0) { + return VEC2_INVALID; + } + return VEC2_SCALAR_1x2; + } + if (M == 2 && N == 1) { + if (has_axis && explicit_axis == 1) { + return VEC2_INVALID; + } + return VEC2_SCALAR_2x1; + } + if (M == 2 && N == 2) { + return (has_axis && explicit_axis == 0) ? VEC2_COLS_2xN : VEC2_ROWS_Nx2; + } + if (N == 2) { + if (has_axis && explicit_axis == 0) { + return VEC2_INVALID; + } + return VEC2_ROWS_Nx2; + } + if (M == 2) { + if (has_axis && explicit_axis == 1) { + return VEC2_INVALID; + } + return VEC2_COLS_2xN; + } + return VEC2_INVALID; +} + +/// @brief Fill @p out with the 2D perpendicular of every vector in @p impl. +/// @details Row-batch and ``1x2`` scalar share one pointer walk; column- +/// batch and ``2x1`` scalar share another. Self-aliasing is NOT +/// supported here \u2014 callers needing in-place must use the +/// dedicated in-place helper. +static void impl_perpendicular_out_of_place(const matrix_impl *impl, + matrix_impl *out, + enum Vec2Axis flavor) { + const size_t M = impl->rows; + const size_t N = impl->columns; + if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_ROWS_Nx2) { + const double *src = impl->data; + double *dst = out->data; + for (size_t r = 0; r < M; ++r) { + const double sx = *src++; + const double sy = *src++; + *dst++ = -sy; + *dst++ = sx; + } + return; + } + const double *src_x = impl->data; + const double *src_y = impl->data + N; + double *dst_x = out->data; + double *dst_y = out->data + N; + for (size_t c = 0; c < N; ++c, ++src_x, ++src_y, ++dst_x, ++dst_y) { + *dst_x = -*src_y; + *dst_y = *src_x; + } +} + +/// @brief In-place 2D perpendicular: swap each (x, y) pair to (-y, x). +static void impl_perpendicular_in_place(matrix_impl *impl, + enum Vec2Axis flavor) { + const size_t M = impl->rows; + const size_t N = impl->columns; + if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_ROWS_Nx2) { + double *p = impl->data; + for (size_t r = 0; r < M; ++r, p += 2) { + const double temp = p[0]; + p[0] = -p[1]; + p[1] = temp; + } + return; + } + double *p = impl->data; + for (size_t c = 0; c < N; ++c, ++p) { + const double temp = p[0]; + p[0] = -p[N]; + p[N] = temp; + } +} + +static PyObject *Matrix_perpendicular(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *axis_obj = NULL; + int in_place = 0; + PyObject *out_op = NULL; + + static char *kwlist[] = {"axis", "in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Op", kwlist, &axis_obj, + &in_place)) { + return NULL; + } + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + return NULL; + } + + enum Vec2Axis flavor = + classify_vec2_axis(self->impl, axis.has_axis, axis.axis); + if (flavor == VEC2_INVALID) { + PyErr_SetString(PyExc_ValueError, + "perpendicular requires a 2D vector or Nx2 or 2xN matrix"); + return NULL; + } + + if (in_place) { + impl_perpendicular_in_place(self->impl, flavor); + return Py_NewRef(op); + } + + matrix_impl *out = set_output(self->impl, op, &out_op, NULL, false); + if (out == NULL) { + return NULL; + } + impl_perpendicular_out_of_place(self->impl, out, flavor); + return out_op; +} + +/// @brief Angle of every 2D vector in @p impl, computed via ``atan2``. +/// @details Returns a Python float for a single vector input, an ``M\xc3\x971`` +/// column matrix for an ``Nx2`` row batch, or a ``1\xc3\x97N`` row +/// matrix for a ``2xN`` column batch. The ``2x2`` ambiguous shape +/// defaults to per-row; pass ``axis=0`` to force per-column. +static PyObject *Matrix_angle(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *axis_obj = NULL; + + static char *kwlist[] = {"axis", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &axis_obj)) { + return NULL; + } + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + return NULL; + } + + enum Vec2Axis flavor = + classify_vec2_axis(self->impl, axis.has_axis, axis.axis); + if (flavor == VEC2_INVALID) { + PyErr_SetString(PyExc_ValueError, + "angle requires a 2D vector or Nx2 or 2xN matrix"); + return NULL; + } + + matrix_impl *impl = self->impl; + if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_SCALAR_2x1) { + return PyFloat_FromDouble(atan2(impl->data[1], impl->data[0])); + } + + if (flavor == VEC2_ROWS_Nx2) { + const size_t M = impl->rows; + matrix_impl *out = impl_new(M, 1); + if (out == NULL) { + return NULL; + } + const double *src = impl->data; + double *dst = out->data; + for (size_t r = 0; r < M; ++r) { + *dst++ = atan2(src[1], src[0]); + src += 2; + } + return wrap_impl_or_free(out); + } + + // Fall-through is the 2xN column-batch case: classify_vec2_axis already + // rejected every shape other than VEC2_ROWS_Nx2 and this one. + const size_t N = impl->columns; + matrix_impl *out = impl_new(1, N); + if (out == NULL) { + return NULL; + } + const double *xp = impl->data; + const double *yp = impl->data + N; + double *dst = out->data; + for (size_t c = 0; c < N; ++c, ++xp, ++yp, ++dst) { + *dst = atan2(*yp, *xp); + } + return wrap_impl_or_free(out); +} + +/* MATRIX_UNARY_METHOD stamps a Python METH_VARARGS|METH_KEYWORDS wrapper + that calls the per-op kernel impl__ewise directly — no runtime + dispatch. The ``in_place`` kwarg routes the output through + ``set_output``: when true, the kernel aliases its input and output + buffers and the method returns ``self`` (refcount-incremented); when + false, a fresh matrix is allocated and returned. The keyword-only + ``out`` target writes the result into a caller-supplied same-shape + matrix (allocation-free, numpy convention) and returns it; ``out`` and + ``in_place`` are mutually exclusive. See the BOC_UNARY_OPS top-of-family + block comment for the full template. */ +#define MATRIX_UNARY_METHOD(ENUM, STAMP) \ + static PyObject *Matrix_##ENUM##_method(PyObject *op, PyObject *args, \ + PyObject *kwds) { \ + MatrixObject *self = (MatrixObject *)op; \ + matrix_impl *impl = self->impl; \ + int in_place = 0; \ + PyObject *out_target = NULL; \ + static char *kwlist[] = {"in_place", "out", NULL}; \ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|$pO", kwlist, &in_place, \ + &out_target)) { \ + return NULL; \ + } \ + if (!impl_check_acquired(impl, true)) { \ + return NULL; \ + } \ + if (in_place && out_target != NULL && out_target != Py_None) { \ + PyErr_SetString(PyExc_ValueError, \ + "out and in_place are mutually exclusive"); \ + return NULL; \ + } \ + PyObject *out_op = NULL; \ + matrix_impl *out = set_output(impl, op, &out_op, out_target, in_place); \ + if (out == NULL) { \ + return NULL; \ + } \ + impl_##STAMP##_ewise(impl, out); \ + return out_op; \ + } + +#define X(E, S, EX) MATRIX_UNARY_METHOD(E, S) +BOC_UNARY_OPS(X) +#undef X + +/// @brief Clamp every element to ``[min, max]``; either bound may be omitted +/// @details The first positional argument is the lower bound, the second the +/// upper. An omitted bound (``None``) is realised as +/// ``-INFINITY`` / ``+INFINITY`` so the inner loop stays branch-free; +/// NaN elements pass through unchanged. Raises ``ValueError`` if both +/// bounds are omitted and ``AssertionError`` if ``max < min``. +/// Output routing matches the unary elementwise family via +/// ``set_output``: ``in_place=True`` clamps ``self`` in place and +/// returns it, the keyword-only ``out`` writes into a caller-supplied +/// same-shape matrix, and the two are mutually exclusive. +static PyObject *Matrix_clip(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + + if (!impl_check_acquired(impl, true)) { + return NULL; + } + + PyObject *minval_op = Py_None; + PyObject *maxval_op = Py_None; + int in_place = 0; + PyObject *out_target = NULL; + static char *kwlist[] = {"min", "max", "in_place", "out", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO$pO", kwlist, &minval_op, + &maxval_op, &in_place, &out_target)) { + return NULL; + } + + if (in_place && out_target != NULL && out_target != Py_None) { + PyErr_SetString(PyExc_ValueError, + "out and in_place are mutually exclusive"); + return NULL; + } + + bool has_min = minval_op != Py_None; + bool has_max = maxval_op != Py_None; + if (!has_min && !has_max) { + PyErr_SetString(PyExc_ValueError, "clip: must provide min and/or max"); + return NULL; + } + + double minval = -INFINITY; + double maxval = INFINITY; + if (has_min && !unwrap_double(minval_op, &minval)) { + PyErr_SetString(PyExc_TypeError, "Expected a number"); + return NULL; + } + if (has_max && !unwrap_double(maxval_op, &maxval)) { + PyErr_SetString(PyExc_TypeError, "Expected a number"); + return NULL; + } + + if (has_min && has_max && maxval < minval) { + PyErr_SetString(PyExc_AssertionError, "maxval < minval"); + return NULL; + } + + PyObject *out_op = NULL; + matrix_impl *out = set_output(impl, op, &out_op, out_target, in_place); + if (out == NULL) { + return NULL; + } + + double *src = impl->data; + double *dst = out->data; + for (size_t i = 0; i < impl->size; ++i, ++src, ++dst) { + double value = *src; + if (value < minval) { + value = minval; + } + if (value > maxval) { + value = maxval; + } + + *dst = value; + } + + return out_op; +} + +/// @brief The "no valid index" sentinel written by the bin / digitize kernels +/// for a NaN input cell. Mirrors the argextreme masked-empty convention +/// (-1): a NaN value has no meaningful bin, so it maps to -1 rather than +/// a valid ``[0, n]`` index. Stored as a double in the result Matrix +/// (read back as ``-1.0``). +#define INDEX_SENTINEL (-1.0) + +/// @brief Data range ignoring NaN, mirroring the masked-aggregate skip rule. +/// @details ``v < lo`` / ``v > hi`` are both false for a NaN ``v``, so NaN +/// never updates either bound. An all-NaN (or empty) matrix leaves +/// ``lo > hi`` (DBL_MAX / -DBL_MAX), which callers treat as a +/// degenerate range. +static void impl_data_range(const matrix_impl *m, double *lo_out, + double *hi_out) { + double lo = DBL_MAX; + double hi = -DBL_MAX; + const double *p = m->data; + for (size_t i = 0; i < m->size; ++i) { + const double v = p[i]; + if (v < lo) { + lo = v; + } + if (v > hi) { + hi = v; } - dispatch_bin_scalar(impl, m, out, Divide); - return 0; } + *lo_out = lo; + *hi_out = hi; +} - if (axis.axis == 0) { - matrix_impl *divisor = impl_new(1, impl->columns); - if (divisor == NULL) { - return -1; +/// @brief Equal-width binning kernel: map each value to its bin index. +/// @details Bins partition ``[lo, hi]`` into ``n`` equal-width intervals. +/// ``idx = floor((v - lo) * n / (hi - lo))``; ``right`` uses +/// ``ceil(...) - 1`` so a value on an interior boundary drops into +/// the lower bin. The index is clamped to ``[0, n-1]`` (so ``v == +/// hi`` lands in the last bin). A NaN value maps to the -1 "no valid +/// index" sentinel (matching the argextreme masked-empty convention). +/// A degenerate range (``hi <= lo`` -- every value equal) sends each +/// non-NaN value to bin 0. Indices are published as doubles (Matrix +/// stores only doubles). Each output cell depends solely on the +/// co-located input cell plus the pre-computed range, so an in-place +/// run (``out`` aliasing the source) is safe. +static void impl_bin_ewise(const matrix_impl *m, matrix_impl *out, double lo, + double hi, long n, bool right) { + const double range = hi - lo; + const double scale = range > 0.0 ? (double)n / range : 0.0; + const double top = (double)(n - 1); + const double *sp = m->data; + double *dp = out->data; + for (size_t i = 0; i < m->size; ++i, ++sp, ++dp) { + const double v = *sp; + if (isnan(v)) { + *dp = INDEX_SENTINEL; + continue; } - dispatch_agg_columnwise(impl, Magnitude, divisor); - sanitize_divisor(divisor); - dispatch_bin_rowwise(impl, divisor, out, Divide); - impl_free(divisor); - return 0; + if (scale == 0.0) { + *dp = 0.0; + continue; + } + const double pos = (v - lo) * scale; + double idx = right ? ceil(pos) - 1.0 : floor(pos); + if (idx < 0.0) { + idx = 0.0; + } else if (idx > top) { + idx = top; + } + *dp = idx; } +} - // axis == 1 (row-wise): parse_validate_normalise_axis restricts axis - // to {0, 1}. - matrix_impl *divisor = impl_new(impl->rows, 1); - if (divisor == NULL) { - return -1; +/// @brief Digitize kernel: map each value to the index of its edge interval. +/// @details ``edges`` is a strictly increasing array of ``nedges`` interior +/// boundaries. With ``right`` false the index is the count of edges +/// ``<= v`` (bin ``i`` means ``edges[i-1] <= v < edges[i]``); with +/// ``right`` true it is the count of edges ``< v`` (``edges[i-1] < v +/// <= edges[i]``), matching ``numpy.digitize``. Indices span ``[0, +/// nedges]``. A NaN value maps to the -1 "no valid index" sentinel. +/// The scan is linear in ``nedges``: edge counts are expected to be +/// far smaller than the cell count, keeping it below the +/// binary-search crossover. A large-histogram fast path can be added +/// behind a flag later. +static void impl_digitize_ewise(const matrix_impl *m, matrix_impl *out, + const double *edges, size_t nedges, + bool right) { + const double *sp = m->data; + double *dp = out->data; + for (size_t i = 0; i < m->size; ++i, ++sp, ++dp) { + const double v = *sp; + if (isnan(v)) { + *dp = INDEX_SENTINEL; + continue; + } + size_t idx = 0; + if (right) { + while (idx < nedges && edges[idx] < v) { + ++idx; + } + } else { + while (idx < nedges && edges[idx] <= v) { + ++idx; + } + } + *dp = (double)idx; } - dispatch_agg_rowwise(impl, Magnitude, divisor); - sanitize_divisor(divisor); - dispatch_bin_columnwise(impl, divisor, out, Divide); - impl_free(divisor); - return 0; } -static PyObject *Matrix_normalize(PyObject *op, PyObject *args, - PyObject *kwds) { +/// @brief Map every element to its equal-width bin index over a range. +/// @details ``bins`` (positional, must be >= 1) is the number of equal-width +/// bins spanning the range. ``bounds`` is an optional ``(min, max)`` +/// tuple that fixes the range and skips the O(size) data scan; when +/// omitted (``None``) the range is the matrix's own min/max, both +/// computed ignoring NaN. A value outside a fixed range clamps into +/// the nearest edge bin (0 or ``bins-1``). Output routing matches the +/// unary elementwise family via ``set_output``: ``in_place=True`` +/// bins ``self`` in place and returns it, the keyword-only ``out`` +/// writes into a caller-supplied same-shape matrix, and the two are +/// mutually exclusive. ``right`` selects the boundary side. See +/// ``impl_bin_ewise`` for the index formula and the NaN / +/// degenerate-range handling. The range is resolved before +/// ``set_output`` so an in-place run bins against the original +/// values. +static PyObject *Matrix_bin(PyObject *op, PyObject *args, PyObject *kwds) { MatrixObject *self = (MatrixObject *)op; - PyObject *axis_obj = NULL; + matrix_impl *impl = self->impl; + + long bins = 0; + PyObject *bounds_op = Py_None; + int right = 0; int in_place = 0; - PyObject *out_op = NULL; + PyObject *out_target = NULL; + static char *kwlist[] = {"bins", "bounds", "right", "in_place", "out", NULL}; - static char *kwlist[] = {"axis", "in_place", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Op", kwlist, &axis_obj, - &in_place)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "l|$OppO", kwlist, &bins, + &bounds_op, &right, &in_place, + &out_target)) { return NULL; } - if (!impl_check_acquired(self->impl, true)) { + + if (!impl_check_acquired(impl, true)) { return NULL; } - AxisArg axis; - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + if (bins < 1) { + PyErr_SetString(PyExc_ValueError, "bins must be >= 1"); return NULL; } - matrix_impl *out = set_output(self->impl, op, &out_op, NULL, in_place); - if (out == NULL) { + if (in_place && out_target != NULL && out_target != Py_None) { + PyErr_SetString(PyExc_ValueError, + "out and in_place are mutually exclusive"); return NULL; } - if (do_normalize(self->impl, axis, out) < 0) { - Py_DECREF(out_op); + + // Resolve the binning range BEFORE set_output: an in-place run makes out + // alias impl, and each cell must bin against the original range, not a + // partially overwritten buffer. A `bounds=(lo, hi)` pair fixes the range and + // skips the O(size) scan; bounds=None falls back to the matrix's own min/max + // (ignoring NaN). + double lo = 0.0; + double hi = 0.0; + if (bounds_op == Py_None) { + impl_data_range(impl, &lo, &hi); + } else { + PyObject *seq = + PySequence_Fast(bounds_op, "bounds must be a (min, max) pair"); + if (seq == NULL) { + return NULL; + } + if (PySequence_Fast_GET_SIZE(seq) != 2) { + Py_DECREF(seq); + PyErr_SetString(PyExc_ValueError, + "bounds must be a (min, max) pair of length 2"); + return NULL; + } + if (!unwrap_double(PySequence_Fast_GET_ITEM(seq, 0), &lo) || + !unwrap_double(PySequence_Fast_GET_ITEM(seq, 1), &hi)) { + Py_DECREF(seq); + PyErr_SetString(PyExc_TypeError, "bounds values must be real numbers"); + return NULL; + } + Py_DECREF(seq); + if (!isfinite(lo) || !isfinite(hi)) { + PyErr_SetString(PyExc_ValueError, "bounds values must be finite"); + return NULL; + } + if (hi < lo) { + PyErr_SetString(PyExc_ValueError, "bounds max must be >= min"); + return NULL; + } + } + + PyObject *out_op = NULL; + matrix_impl *out = set_output(impl, op, &out_op, out_target, in_place); + if (out == NULL) { return NULL; } + + impl_bin_ewise(impl, out, lo, hi, bins, (bool)right); return out_op; } -enum Vec2Axis { - VEC2_SCALAR_1x2, - VEC2_SCALAR_2x1, - VEC2_ROWS_Nx2, - VEC2_COLS_2xN, - VEC2_INVALID -}; +/// @brief Map every element to the index of the edge interval it falls in. +/// @details ``edges`` (positional) is a strictly increasing vector of interior +/// boundaries -- a Matrix vector (1xN or Mx1) or any sequence that +/// ``unwrap_matrix`` coerces. Output routing matches the unary +/// elementwise family via ``set_output``: ``in_place`` and the +/// keyword-only ``out`` are mutually exclusive. ``right`` selects the +/// boundary side (``numpy.digitize`` convention). See +/// ``impl_digitize_ewise`` for the index rule and NaN handling. +static PyObject *Matrix_digitize(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; -/// @brief Classify a matrix as a 2D vector or batch of 2D vectors. -/// @details ``has_axis`` / ``explicit_axis`` carry an optional caller- -/// supplied axis (already normalised to 0 or 1). The ``2x2`` -/// shape is doubly-valid and ``explicit_axis`` picks the -/// orientation (``0`` -> columns, default -> rows). For all -/// other shapes only one orientation is valid; supplying an -/// ``explicit_axis`` that contradicts that orientation returns -/// ``VEC2_INVALID``. Returns ``VEC2_INVALID`` for any shape -/// that is not a 2D vector or Nx2 / 2xN batch. -static enum Vec2Axis classify_vec2_axis(const matrix_impl *impl, bool has_axis, - int explicit_axis) { - const size_t M = impl->rows; - const size_t N = impl->columns; - if (M == 1 && N == 2) { - if (has_axis && explicit_axis == 0) { - return VEC2_INVALID; - } - return VEC2_SCALAR_1x2; + PyObject *edges_op = NULL; + int right = 0; + int in_place = 0; + PyObject *out_target = NULL; + static char *kwlist[] = {"edges", "right", "in_place", "out", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|$ppO", kwlist, &edges_op, + &right, &in_place, &out_target)) { + return NULL; } - if (M == 2 && N == 1) { - if (has_axis && explicit_axis == 1) { - return VEC2_INVALID; - } - return VEC2_SCALAR_2x1; + + if (!impl_check_acquired(impl, true)) { + return NULL; } - if (M == 2 && N == 2) { - return (has_axis && explicit_axis == 0) ? VEC2_COLS_2xN : VEC2_ROWS_Nx2; + + if (in_place && out_target != NULL && out_target != Py_None) { + PyErr_SetString(PyExc_ValueError, + "out and in_place are mutually exclusive"); + return NULL; } - if (N == 2) { - if (has_axis && explicit_axis == 0) { - return VEC2_INVALID; - } - return VEC2_ROWS_Nx2; + + matrix_impl *edges = unwrap_matrix(edges_op, false); + if (edges == NULL) { + return NULL; } - if (M == 2) { - if (has_axis && explicit_axis == 1) { - return VEC2_INVALID; + + // A vector of any orientation reads as a flat contiguous boundary list; + // reject a genuine 2-D matrix rather than silently flattening it. + if (edges->rows != 1 && edges->columns != 1) { + PyErr_SetString(PyExc_ValueError, "edges must be a vector"); + IMPL_DECREF(edges); + return NULL; + } + + // Strictly increasing and NaN-free: both the count-of-edges index rule and + // the scan's early-out assume monotone boundaries. + for (size_t i = 0; i < edges->size; ++i) { + if (isnan(edges->data[i]) || + (i > 0 && edges->data[i] <= edges->data[i - 1])) { + PyErr_SetString(PyExc_ValueError, + "edges must be strictly increasing and non-NaN"); + IMPL_DECREF(edges); + return NULL; } - return VEC2_COLS_2xN; } - return VEC2_INVALID; + + // Snapshot edges into an owned buffer so the kernel is decoupled from the + // source object: without this, m.digitize(m, in_place=True) would read + // boundaries out of the very buffer it is overwriting. Edge counts are small + // by design, so the copy is negligible. + size_t nedges = edges->size; + double *edge_buf = PyMem_RawMalloc(nedges * sizeof(double)); + if (edge_buf == NULL) { + IMPL_DECREF(edges); + PyErr_NoMemory(); + return NULL; + } + memcpy(edge_buf, edges->data, nedges * sizeof(double)); + IMPL_DECREF(edges); + + PyObject *out_op = NULL; + matrix_impl *out = set_output(impl, op, &out_op, out_target, in_place); + if (out == NULL) { + PyMem_RawFree(edge_buf); + return NULL; + } + + impl_digitize_ewise(impl, out, edge_buf, nedges, (bool)right); + PyMem_RawFree(edge_buf); + return out_op; } -/// @brief Fill @p out with the 2D perpendicular of every vector in @p impl. -/// @details Row-batch and ``1x2`` scalar share one pointer walk; column- -/// batch and ``2x1`` scalar share another. Self-aliasing is NOT -/// supported here \u2014 callers needing in-place must use the -/// dedicated in-place helper. -static void impl_perpendicular_out_of_place(const matrix_impl *impl, - matrix_impl *out, - enum Vec2Axis flavor) { - const size_t M = impl->rows; - const size_t N = impl->columns; - if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_ROWS_Nx2) { - const double *src = impl->data; - double *dst = out->data; - for (size_t r = 0; r < M; ++r) { - const double sx = *src++; - const double sy = *src++; - *dst++ = -sy; - *dst++ = sx; - } - return; +static PyObject *Matrix_copy(PyObject *op, PyObject *Py_UNUSED(dummy)) { + MatrixObject *matrix = (MatrixObject *)op; + matrix_impl *impl = matrix->impl; + + if (!impl_check_acquired(impl, true)) { + return NULL; } - const double *src_x = impl->data; - const double *src_y = impl->data + N; - double *dst_x = out->data; - double *dst_y = out->data + N; - for (size_t c = 0; c < N; ++c, ++src_x, ++src_y, ++dst_x, ++dst_y) { - *dst_x = -*src_y; - *dst_y = *src_x; + + matrix_impl *copy = impl_new(impl->rows, impl->columns); + if (copy == NULL) { + return NULL; + } + + memcpy(copy->data, impl->data, impl->size * sizeof(double)); + + return (PyObject *)wrap_impl_or_free(copy); +} + +// Matrix-side factory materialising a MatrixView into an independent, owned +// Matrix -- the classmethod twin of MatrixView.copy(), sharing view_materialize +// and fitting the Matrix.zeros / ones / full / vector / concat factory family. +static PyObject *Matrix_from_view(PyObject *Py_UNUSED(cls), PyObject *view_op) { + if (Py_TYPE(view_op) != LOCAL_STATE->view_type) { + PyErr_Format(PyExc_TypeError, "from_view expects a MatrixView, not %.200s", + Py_TYPE(view_op)->tp_name); + return NULL; + } + matrix_impl *out = view_materialize((MatrixViewObject *)view_op); + if (out == NULL) { + return NULL; + } + return (PyObject *)wrap_impl_or_free(out); +} + +/// @brief Resolve one fancy-index item to an in-bounds row/column number. +/// @details Applies Python-style negative wrapping (``-1`` -> ``dim-1``) +/// then an explicit bounds check, so a list/tuple subscript can +/// never read outside the matrix. ``axis_name`` ("row"/"column") +/// and the *original* (pre-wrap) value are named in the +/// ``IndexError`` so a bad subscript is self-describing. +/// ``bool`` items are accepted as ``0``/``1`` (bool is an ``int`` +/// subclass), matching the scalar-int subscript path. +/// @return 0 with ``*out`` set on success; -1 with an exception set on +/// failure: ``TypeError`` for a non-int item, ``OverflowError`` +/// for a value outside ``Py_ssize_t`` (raised by +/// ``PyLong_AsSsize_t`` before the bounds check), or +/// ``IndexError`` for an out-of-range index. +static int resolve_gather_index(PyObject *item, size_t dim, + const char *axis_name, size_t *out) { + Py_ssize_t value = PyLong_AsSsize_t(item); + if (value == -1 && PyErr_Occurred()) { + return -1; + } + + Py_ssize_t resolved = value; + if (resolved < 0) { + resolved += (Py_ssize_t)dim; } + + if (resolved < 0 || (size_t)resolved >= dim) { + PyErr_Format(PyExc_IndexError, + "%s index %zd out of range for dimension of size %zu", + axis_name, value, dim); + return -1; + } + + *out = (size_t)resolved; + return 0; } -/// @brief In-place 2D perpendicular: swap each (x, y) pair to (-y, x). -static void impl_perpendicular_in_place(matrix_impl *impl, - enum Vec2Axis flavor) { - const size_t M = impl->rows; - const size_t N = impl->columns; - if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_ROWS_Nx2) { - double *p = impl->data; - for (size_t r = 0; r < M; ++r, p += 2) { - const double temp = p[0]; - p[0] = -p[1]; - p[1] = temp; +/// @brief Resolve a gather's output matrix: a caller ``out=`` or a fresh one. +/// @details Shared by @ref gather_axis and @ref gather_along_axis. When +/// @p out_target is a real object it is shape-validated via +/// @ref use_out_target and rejected if it aliases @p impl (a +/// reordering gather would read cells it has already overwritten); +/// otherwise a fresh matrix of the requested shape is allocated and +/// wrapped as @p type (so a subclass is preserved). On success the +/// impl is returned and @p *result is set to the owning ``PyObject``; +/// on failure NULL is returned with an exception set (the caller +/// still owns any index buffer it must free). +static matrix_impl *gather_output(matrix_impl *impl, PyObject *out_target, + PyTypeObject *type, size_t want_rows, + size_t want_cols, PyObject **result) { + if (out_target != NULL && out_target != Py_None) { + matrix_impl *out = use_out_target(out_target, result, want_rows, want_cols); + if (out == NULL) { + return NULL; } - return; + if (out == impl) { + Py_DECREF(*result); + PyErr_SetString(PyExc_ValueError, + "out must not alias the matrix being taken from"); + return NULL; + } + return out; } - double *p = impl->data; - for (size_t c = 0; c < N; ++c, ++p) { - const double temp = p[0]; - p[0] = -p[N]; - p[N] = temp; + + matrix_impl *out = impl_new(want_rows, want_cols); + if (out == NULL) { + return NULL; } + *result = wrap_matrix(type, out); + if (*result == NULL) { + impl_free(out); + return NULL; + } + return out; } -static PyObject *Matrix_perpendicular(PyObject *op, PyObject *args, - PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - PyObject *axis_obj = NULL; - int in_place = 0; - PyObject *out_op = NULL; +/// @brief Gather rows (``axis == 0``) or columns (``axis == 1``) named by a +/// sequence of indices into a matrix. +/// @details Shared by ``Matrix.take`` and the subscript gather path so +/// both surfaces resolve indices identically (negative-aware and +/// bounds-checked via ``resolve_gather_index``) rather than +/// keeping two copies in sync. Row gather copies each selected +/// row with ``memcpy`` (source and destination rows are both +/// contiguous); column gather copies element-wise (inherently +/// strided). An empty ``indices`` sequence raises ``IndexError`` +/// and is rejected *before* any allocation — ``impl_new`` only +/// ``assert``s non-zero dimensions, so a ``0``-length axis must +/// never reach it. +/// +/// Every index is resolved up front, before any element is written, +/// so a bad index can never leave a caller-supplied ``out_target`` +/// partially overwritten (the validate-then-write contract shared +/// with ``put``). When ``out_target`` is NULL / ``Py_None`` the +/// result is a fresh matrix wrapped as ``type`` (so a subclass is +/// preserved); otherwise it is written into that caller-owned +/// matrix, which must have the result's shape and must not alias the +/// source (a gather reorders elements, so an in-place permutation +/// would read rows it had already clobbered). +static PyObject *gather_axis(matrix_impl *impl, PyObject *indices, int axis, + PyTypeObject *type, PyObject *out_target) { + const char *err_msg = + "Indices must be specified as a list or a tuple of ints"; + PyObject *fast = PySequence_Fast(indices, err_msg); + if (fast == NULL) { + return NULL; + } - static char *kwlist[] = {"axis", "in_place", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|Op", kwlist, &axis_obj, - &in_place)) { + Py_ssize_t count = PySequence_Fast_GET_SIZE(fast); + if (count == 0) { + Py_DECREF(fast); + PyErr_SetString(PyExc_IndexError, "index sequence must not be empty"); return NULL; } - if (!impl_check_acquired(self->impl, true)) { + + const size_t dim = (axis == 0) ? impl->rows : impl->columns; + const char *axis_name = (axis == 0) ? "row" : "column"; + + // Validate phase: resolve every index before any write. + size_t *resolved = PyMem_Malloc((size_t)count * sizeof(size_t)); + if (resolved == NULL) { + Py_DECREF(fast); + return PyErr_NoMemory(); + } + for (Py_ssize_t i = 0; i < count; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + if (resolve_gather_index(item, dim, axis_name, &resolved[i]) < 0) { + PyMem_Free(resolved); + Py_DECREF(fast); + return NULL; + } + } + Py_DECREF(fast); + + const size_t want_rows = (axis == 0) ? (size_t)count : impl->rows; + const size_t want_cols = (axis == 0) ? impl->columns : (size_t)count; + + PyObject *result = NULL; + matrix_impl *out = + gather_output(impl, out_target, type, want_rows, want_cols, &result); + if (out == NULL) { + PyMem_Free(resolved); return NULL; } - AxisArg axis; - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + // Write phase. + if (axis == 0) { + for (Py_ssize_t i = 0; i < count; ++i) { + memcpy(out->row_ptrs[i], impl->row_ptrs[resolved[i]], + impl->columns * sizeof(double)); + } + } else { + for (Py_ssize_t i = 0; i < count; ++i) { + const size_t c = resolved[i]; + for (size_t r = 0; r < impl->rows; ++r) { + out->row_ptrs[r][i] = impl->row_ptrs[r][c]; + } + } + } + + PyMem_Free(resolved); + return result; +} + +/// @brief Gather one element per row (``axis == 1``) or per column +/// (``axis == 0``) into a fresh matrix, the indices running *along* +/// the named axis. +/// @details This is the ``np.take_along_axis`` counterpart to ``gather_axis`` +/// (which selects whole rows/columns). For ``axis == 1`` the +/// sequence holds one column index per row — its length must equal +/// ``rows`` — and the result is ``rows x 1`` with +/// ``out[r] = self[r][indices[r]]``. For ``axis == 0`` it holds one +/// row index per column — length ``columns`` — and the result is +/// ``1 x columns`` with ``out[c] = self[indices[c]][c]``. This pairs +/// with ``argmin``/``argmax`` along the same axis: the index list +/// they return feeds straight back in to gather the reduced values. +/// Each index is resolved negative-aware and bounds-checked against +/// the *gathered* axis via ``resolve_gather_index``. +/// +/// Every index is resolved up front, before any element is written, +/// so a bad index can never leave a caller-supplied ``out_target`` +/// partially overwritten (the validate-then-write contract shared +/// with ``take``). When ``out_target`` is NULL / ``Py_None`` the +/// result is a fresh matrix wrapped as ``type`` (so a subclass is +/// preserved); otherwise it is written into that caller-owned +/// matrix, which must have the result's shape and must not alias the +/// source. +static PyObject *gather_along_axis(matrix_impl *impl, PyObject *indices, + int axis, PyTypeObject *type, + PyObject *out_target) { + const char *err_msg = + "Indices must be specified as a list or a tuple of ints"; + PyObject *fast = PySequence_Fast(indices, err_msg); + if (fast == NULL) { return NULL; } - enum Vec2Axis flavor = - classify_vec2_axis(self->impl, axis.has_axis, axis.axis); - if (flavor == VEC2_INVALID) { - PyErr_SetString(PyExc_ValueError, - "perpendicular requires a 2D vector or Nx2 or 2xN matrix"); + Py_ssize_t count = PySequence_Fast_GET_SIZE(fast); + size_t expected = axis == 0 ? impl->columns : impl->rows; + size_t bound = axis == 0 ? impl->rows : impl->columns; + const char *bound_name = axis == 0 ? "row" : "column"; + if ((size_t)count != expected) { + PyErr_Format(PyExc_ValueError, + "take_along_axis on axis %d expects %zu indices (one per %s), " + "got %zd", + axis, expected, axis == 0 ? "column" : "row", count); + Py_DECREF(fast); return NULL; } - if (in_place) { - impl_perpendicular_in_place(self->impl, flavor); - return Py_NewRef(op); + // Validate phase: resolve every index before any write. + size_t *resolved = PyMem_Malloc((size_t)count * sizeof(size_t)); + if (resolved == NULL) { + Py_DECREF(fast); + return PyErr_NoMemory(); + } + for (Py_ssize_t i = 0; i < count; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + if (resolve_gather_index(item, bound, bound_name, &resolved[i]) < 0) { + PyMem_Free(resolved); + Py_DECREF(fast); + return NULL; + } } + Py_DECREF(fast); - matrix_impl *out = set_output(self->impl, op, &out_op, NULL, false); + const size_t want_rows = axis == 0 ? 1 : impl->rows; + const size_t want_cols = axis == 0 ? impl->columns : 1; + + PyObject *result = NULL; + matrix_impl *out = + gather_output(impl, out_target, type, want_rows, want_cols, &result); if (out == NULL) { + PyMem_Free(resolved); return NULL; } - impl_perpendicular_out_of_place(self->impl, out, flavor); - return out_op; + + // Write phase. + for (Py_ssize_t i = 0; i < count; ++i) { + const size_t k = resolved[i]; + if (axis == 0) { + out->row_ptrs[0][i] = impl->row_ptrs[k][i]; + } else { + out->row_ptrs[i][0] = impl->row_ptrs[i][k]; + } + } + + PyMem_Free(resolved); + return result; } -/// @brief Angle of every 2D vector in @p impl, computed via ``atan2``. -/// @details Returns a Python float for a single vector input, an ``M\xc3\x971`` -/// column matrix for an ``Nx2`` row batch, or a ``1\xc3\x97N`` row -/// matrix for a ``2xN`` column batch. The ``2x2`` ambiguous shape -/// defaults to per-row; pass ``axis=0`` to force per-column. -static PyObject *Matrix_angle(PyObject *op, PyObject *args, PyObject *kwds) { +/* Defined alongside Matrix_ass_subscript (the scatter machinery lives next + to the write path); forward-declared here so Matrix_put can share it. */ +static int scatter_axis(matrix_impl *impl, PyObject *indices, int axis, + PyObject *value_op, int accumulate); + +/* The along-axis scatter counterpart to scatter_axis, likewise defined next + to the write path and forward-declared so Matrix_put_along_axis can use it. + */ +static int scatter_along_axis(matrix_impl *impl, PyObject *indices, int axis, + PyObject *value_op, int accumulate); + +PyObject *Matrix_take(PyObject *op, PyObject *args, PyObject *kwds) { MatrixObject *self = (MatrixObject *)op; - PyObject *axis_obj = NULL; + matrix_impl *impl = self->impl; - static char *kwlist[] = {"axis", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &axis_obj)) { + if (!impl_check_acquired(impl, true)) { return NULL; } - if (!impl_check_acquired(self->impl, true)) { + + PyObject *indices = NULL; + int axis = 0; + PyObject *out = NULL; + static char *keywords[] = {"indices", "axis", "out", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|i$O", keywords, &indices, + &axis, &out)) { return NULL; } - AxisArg axis; - if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { - return NULL; + if (axis < 0) { + axis = 2 + axis; } - enum Vec2Axis flavor = - classify_vec2_axis(self->impl, axis.has_axis, axis.axis); - if (flavor == VEC2_INVALID) { - PyErr_SetString(PyExc_ValueError, - "angle requires a 2D vector or Nx2 or 2xN matrix"); + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); return NULL; } + return gather_axis(impl, indices, axis, Py_TYPE(self), out); +} + +PyObject *Matrix_put(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; - if (flavor == VEC2_SCALAR_1x2 || flavor == VEC2_SCALAR_2x1) { - return PyFloat_FromDouble(atan2(impl->data[1], impl->data[0])); - } - if (flavor == VEC2_ROWS_Nx2) { - const size_t M = impl->rows; - matrix_impl *out = impl_new(M, 1); - if (out == NULL) { - return NULL; - } - const double *src = impl->data; - double *dst = out->data; - for (size_t r = 0; r < M; ++r) { - *dst++ = atan2(src[1], src[0]); - src += 2; - } - return wrap_impl_or_free(out); + if (!impl_check_acquired(impl, true)) { + return NULL; } - // Fall-through is the 2xN column-batch case: classify_vec2_axis already - // rejected every shape other than VEC2_ROWS_Nx2 and this one. - const size_t N = impl->columns; - matrix_impl *out = impl_new(1, N); - if (out == NULL) { + PyObject *indices = NULL; + PyObject *value = NULL; + int axis = 0; + int accumulate = 0; + static char *keywords[] = {"indices", "value", "axis", "accumulate", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|ip", keywords, &indices, + &value, &axis, &accumulate)) { return NULL; } - const double *xp = impl->data; - const double *yp = impl->data + N; - double *dst = out->data; - for (size_t c = 0; c < N; ++c, ++xp, ++yp, ++dst) { - *dst = atan2(*yp, *xp); - } - return wrap_impl_or_free(out); -} -/* MATRIX_UNARY_METHOD stamps a Python METH_VARARGS|METH_KEYWORDS wrapper - that calls the per-op kernel impl__ewise directly — no runtime - dispatch. The ``in_place`` kwarg routes the output through - ``set_output``: when true, the kernel aliases its input and output - buffers and the method returns ``self`` (refcount-incremented); when - false, a fresh matrix is allocated and returned. The keyword-only - ``out`` target writes the result into a caller-supplied same-shape - matrix (allocation-free, numpy convention) and returns it; ``out`` and - ``in_place`` are mutually exclusive. See the BOC_UNARY_OPS top-of-family - block comment for the full template. */ -#define MATRIX_UNARY_METHOD(ENUM, STAMP) \ - static PyObject *Matrix_##ENUM##_method(PyObject *op, PyObject *args, \ - PyObject *kwds) { \ - MatrixObject *self = (MatrixObject *)op; \ - matrix_impl *impl = self->impl; \ - int in_place = 0; \ - PyObject *out_target = NULL; \ - static char *kwlist[] = {"in_place", "out", NULL}; \ - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|p$O", kwlist, &in_place, \ - &out_target)) { \ - return NULL; \ - } \ - if (!impl_check_acquired(impl, true)) { \ - return NULL; \ - } \ - if (in_place && out_target != NULL && out_target != Py_None) { \ - PyErr_SetString(PyExc_ValueError, \ - "out and in_place are mutually exclusive"); \ - return NULL; \ - } \ - PyObject *out_op = NULL; \ - matrix_impl *out = set_output(impl, op, &out_op, out_target, in_place); \ - if (out == NULL) { \ - return NULL; \ - } \ - impl_##STAMP##_ewise(impl, out); \ - return out_op; \ + if (axis < 0) { + axis = 2 + axis; } -#define X(E, S, EX) MATRIX_UNARY_METHOD(E, S) -BOC_UNARY_OPS(X) -#undef X + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); + return NULL; + } -/// @brief Clamp every element to ``[min, max]``; either bound may be omitted -/// @details The first positional argument is the lower bound, the second the -/// upper. An omitted bound (``None``) is realised as -/// ``-INFINITY`` / ``+INFINITY`` so the inner loop stays branch-free; -/// NaN elements pass through unchanged. Raises ``ValueError`` if both -/// bounds are omitted and ``AssertionError`` if ``max < min``. -static PyObject *Matrix_clip(PyObject *op, PyObject *args, PyObject *kwds) { + if (scatter_axis(impl, indices, axis, value, accumulate) < 0) { + return NULL; + } + + Py_INCREF(self); + return (PyObject *)self; +} + +PyObject *Matrix_take_along_axis(PyObject *op, PyObject *args, PyObject *kwds) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -2946,210 +5004,368 @@ static PyObject *Matrix_clip(PyObject *op, PyObject *args, PyObject *kwds) { return NULL; } - PyObject *minval_op = Py_None; - PyObject *maxval_op = Py_None; - static char *kwlist[] = {"min", "max", NULL}; + PyObject *indices = NULL; + int axis = 0; + PyObject *out = NULL; + static char *keywords[] = {"indices", "axis", "out", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO", kwlist, &minval_op, - &maxval_op)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|i$O", keywords, &indices, + &axis, &out)) { return NULL; } - bool has_min = minval_op != Py_None; - bool has_max = maxval_op != Py_None; - if (!has_min && !has_max) { - PyErr_SetString(PyExc_ValueError, "clip: must provide min and/or max"); - return NULL; + if (axis < 0) { + axis = 2 + axis; } - double minval = -INFINITY; - double maxval = INFINITY; - if (has_min && !unwrap_double(minval_op, &minval)) { - PyErr_SetString(PyExc_TypeError, "Expected a number"); + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); return NULL; } - if (has_max && !unwrap_double(maxval_op, &maxval)) { - PyErr_SetString(PyExc_TypeError, "Expected a number"); + + return gather_along_axis(impl, indices, axis, Py_TYPE(self), out); +} + +PyObject *Matrix_put_along_axis(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + + if (!impl_check_acquired(impl, true)) { return NULL; } - if (has_min && has_max && maxval < minval) { - PyErr_SetString(PyExc_AssertionError, "maxval < minval"); + PyObject *indices = NULL; + PyObject *value = NULL; + int axis = 0; + int accumulate = 0; + static char *keywords[] = {"indices", "value", "axis", "accumulate", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|ip", keywords, &indices, + &value, &axis, &accumulate)) { return NULL; } - PyTypeObject *type = Py_TYPE(self); - MatrixObject *out = (MatrixObject *)type->tp_alloc(type, 0); - if (out == NULL) { + if (axis < 0) { + axis = 2 + axis; + } + + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); return NULL; } - out->impl = impl_new(impl->rows, impl->columns); - if (out->impl == NULL) { - Py_DECREF(out); + if (scatter_along_axis(impl, indices, axis, value, accumulate) < 0) { return NULL; } - IMPL_INCREF(out->impl); + Py_INCREF(self); + return (PyObject *)self; +} - double *src = impl->data; - double *dst = out->impl->data; - for (size_t i = 0; i < impl->size; ++i, ++src, ++dst) { - double value = *src; - if (value < minval) { - value = minval; +/// @brief Repeat each element/row/column ``repeats`` times consecutively. +/// @details The ``np.repeat`` / ``torch.repeat_interleave`` shape: copies run +/// *interleaved*, not tiled — ``[a, b]`` with ``repeats == 2`` becomes +/// ``[a, a, b, b]`` (contrast tiling, which gives ``[a, b, a, b]``). +/// ``axis == 0`` repeats whole rows into a ``(rows*repeats) x +/// columns`` result (each source row ``memcpy``'d ``repeats`` times); +/// ``axis == 1`` repeats each column into ``rows x +/// (columns*repeats)``; +/// ``flatten`` (the no-axis path) walks the row-major buffer and emits +/// a ``1 x (size*repeats)`` row vector. The *total* output element +/// count (``size * repeats``) is validated for ``size_t`` overflow +/// before allocation — and ``impl_new`` re-checks the final +/// ``rows * columns`` product — so a huge ``repeats`` can never +/// under-allocate and write out of bounds. +static matrix_impl *impl_repeat_interleave(matrix_impl *m, size_t repeats, + int axis, bool flatten) { + // Bound the total element count, not just the repeated dimension: for + // axis 0/1 the result is rows*columns*repeats == size*repeats, and a + // per-dimension guard (columns*repeats) would let the other dimension + // wrap the product. impl_new re-checks rows*columns as a backstop. + if (repeats != 0 && m->size > SIZE_MAX / repeats) { + PyErr_SetString(PyExc_OverflowError, "repeat_interleave result too large"); + return NULL; + } + + if (flatten) { + matrix_impl *out = impl_new(1, m->size * repeats); + if (out == NULL) { + return NULL; } - if (value > maxval) { - value = maxval; + double *op = out->data; + for (size_t i = 0; i < m->size; ++i) { + const double v = m->data[i]; + for (size_t t = 0; t < repeats; ++t, ++op) { + *op = v; + } } + return out; + } - *dst = value; + if (axis == 0) { + matrix_impl *out = impl_new(m->rows * repeats, m->columns); + if (out == NULL) { + return NULL; + } + for (size_t r = 0; r < m->rows; ++r) { + const double *src = m->row_ptrs[r]; + for (size_t t = 0; t < repeats; ++t) { + memcpy(out->row_ptrs[r * repeats + t], src, + m->columns * sizeof(double)); + } + } + return out; } - return (PyObject *)out; + matrix_impl *out = impl_new(m->rows, m->columns * repeats); + if (out == NULL) { + return NULL; + } + for (size_t r = 0; r < m->rows; ++r) { + const double *src = m->row_ptrs[r]; + double *dst = out->row_ptrs[r]; + for (size_t c = 0; c < m->columns; ++c) { + const double v = src[c]; + for (size_t t = 0; t < repeats; ++t, ++dst) { + *dst = v; + } + } + } + return out; } -static PyObject *Matrix_copy(PyObject *op, PyObject *Py_UNUSED(dummy)) { - MatrixObject *matrix = (MatrixObject *)op; - matrix_impl *impl = matrix->impl; +PyObject *Matrix_repeat_interleave(PyObject *op, PyObject *args, + PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; if (!impl_check_acquired(impl, true)) { return NULL; } - matrix_impl *copy = impl_new(impl->rows, impl->columns); - if (copy == NULL) { + Py_ssize_t repeats = 0; + PyObject *axis_obj = NULL; + static char *keywords[] = {"repeats", "axis", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "n|O", keywords, &repeats, + &axis_obj)) { return NULL; } - memcpy(copy->data, impl->data, impl->size * sizeof(double)); + if (repeats < 1) { + PyErr_SetString(PyExc_ValueError, "repeats must be a positive integer"); + return NULL; + } - return (PyObject *)wrap_impl_or_free(copy); + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { + return NULL; + } + + matrix_impl *out = impl_repeat_interleave( + impl, (size_t)repeats, axis.has_axis ? axis.axis : 0, !axis.has_axis); + if (out == NULL) { + return NULL; + } + + return wrap_impl_or_free(out); } -/// @brief Resolve one fancy-index item to an in-bounds row/column number. -/// @details Applies Python-style negative wrapping (``-1`` -> ``dim-1``) -/// then an explicit bounds check, so a list/tuple subscript can -/// never read outside the matrix. ``axis_name`` ("row"/"column") -/// and the *original* (pre-wrap) value are named in the -/// ``IndexError`` so a bad subscript is self-describing. -/// ``bool`` items are accepted as ``0``/``1`` (bool is an ``int`` -/// subclass), matching the scalar-int subscript path. -/// @return 0 with ``*out`` set on success; -1 with an exception set on -/// failure: ``TypeError`` for a non-int item, ``OverflowError`` -/// for a value outside ``Py_ssize_t`` (raised by -/// ``PyLong_AsSsize_t`` before the bounds check), or -/// ``IndexError`` for an out-of-range index. -static int resolve_gather_index(PyObject *item, size_t dim, - const char *axis_name, size_t *out) { - Py_ssize_t value = PyLong_AsSsize_t(item); - if (value == -1 && PyErr_Occurred()) { +/* -------------------------------------------------------------------------- + topk: the k extreme elements per reduction group, in sorted order. + + A group is gathered into a scratch array of (value, original-index) pairs + and sorted by qsort with one of two comparators. Both put any NaN last + (a NaN is never a "top" value, matching np.sort) and break ties between + equal values by the smaller original index, so qsort's lack of stability + does not matter: the comparator is a strict total order and the first + occurrence of a tied value always wins (NumPy tie-break). + + With a where= mask only the included cells (mask cell != 0.0; NaN counts + as included) are gathered. A group with fewer than k included elements + fills its leading slots with the available extremes and pads the rest with + NaN values and -1 indices -- the same "no element" sentinels used by the + masked aggregates (NaN) and masked argmin/argmax (-1). + + The indices come back in one of two shapes. By default (as_matrix=False) + they are Python ints -- a flat list for axis=None, or a list of per-group + lists for axis=0/1 -- directly usable in fancy indexing. With + as_matrix=True they are instead a Matrix of the same shape as the values + matrix, each cell the (double-valued) source index aligned with the value + beside it; the -1 pad becomes -1.0. + -------------------------------------------------------------------------- */ +typedef struct { + double value; + Py_ssize_t index; +} topk_entry; + +static int topk_cmp_desc(const void *pa, const void *pb) { + const topk_entry *a = (const topk_entry *)pa; + const topk_entry *b = (const topk_entry *)pb; + const int na = isnan(a->value); + const int nb = isnan(b->value); + if (na || nb) { + if (na && nb) { + return (a->index < b->index) ? -1 : 1; + } + return na ? 1 : -1; /* NaN sorts last */ + } + if (a->value > b->value) { return -1; } - - Py_ssize_t resolved = value; - if (resolved < 0) { - resolved += (Py_ssize_t)dim; + if (a->value < b->value) { + return 1; } + return (a->index < b->index) ? -1 : 1; /* first occurrence wins ties */ +} - if (resolved < 0 || (size_t)resolved >= dim) { - PyErr_Format(PyExc_IndexError, - "%s index %zd out of range for dimension of size %zu", - axis_name, value, dim); +static int topk_cmp_asc(const void *pa, const void *pb) { + const topk_entry *a = (const topk_entry *)pa; + const topk_entry *b = (const topk_entry *)pb; + const int na = isnan(a->value); + const int nb = isnan(b->value); + if (na || nb) { + if (na && nb) { + return (a->index < b->index) ? -1 : 1; + } + return na ? 1 : -1; /* NaN sorts last */ + } + if (a->value < b->value) { return -1; } - - *out = (size_t)resolved; - return 0; + if (a->value > b->value) { + return 1; + } + return (a->index < b->index) ? -1 : 1; /* first occurrence wins ties */ } -/// @brief Gather rows (``axis == 0``) or columns (``axis == 1``) named by a -/// sequence of indices into a fresh matrix. -/// @details Shared by ``Matrix.take`` and the subscript gather path so -/// both surfaces resolve indices identically (negative-aware and -/// bounds-checked via ``resolve_gather_index``) rather than -/// keeping two copies in sync. Row gather copies each selected -/// row with ``memcpy`` (source and destination rows are both -/// contiguous); column gather copies element-wise (inherently -/// strided). The result is wrapped as ``type`` so a subclass is -/// preserved. An empty ``indices`` sequence raises ``IndexError`` -/// and is rejected *before* any allocation — ``impl_new`` only -/// ``assert``s non-zero dimensions, so a ``0``-length axis must -/// never reach it. -static PyObject *gather_axis(matrix_impl *impl, PyObject *indices, int axis, - PyTypeObject *type) { - const char *err_msg = - "Indices must be specified as a list or a tuple of ints"; - PyObject *fast = PySequence_Fast(indices, err_msg); - if (fast == NULL) { - return NULL; +// Sort the first `m` gathered entries and emit the top `k`: the j-th slot +// takes entry j when j < m, else the (NaN, -1) pad. out_vals/out_idx are +// contiguous length-k scratch buffers the caller scatters into its result. +static void topk_sort_pad(topk_entry *scratch, size_t m, size_t k, bool largest, + double *out_vals, Py_ssize_t *out_idx) { + qsort(scratch, m, sizeof(topk_entry), largest ? topk_cmp_desc : topk_cmp_asc); + for (size_t j = 0; j < k; ++j) { + if (j < m) { + out_vals[j] = scratch[j].value; + out_idx[j] = scratch[j].index; + } else { + out_vals[j] = NAN; + out_idx[j] = -1; + } } +} - Py_ssize_t count = PySequence_Fast_GET_SIZE(fast); - if (count == 0) { - Py_DECREF(fast); - PyErr_SetString(PyExc_IndexError, "index sequence must not be empty"); +static PyObject *topk_build_int_list(const Py_ssize_t *idx, size_t k) { + PyObject *list = PyList_New((Py_ssize_t)k); + if (list == NULL) { return NULL; } - - matrix_impl *out; - if (axis == 0) { - out = impl_new((size_t)count, impl->columns); - if (out == NULL) { - Py_DECREF(fast); + for (size_t j = 0; j < k; ++j) { + PyObject *o = PyLong_FromSsize_t(idx[j]); + if (o == NULL) { + Py_DECREF(list); return NULL; } + PyList_SET_ITEM(list, (Py_ssize_t)j, o); + } + return list; +} - for (Py_ssize_t i = 0; i < count; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - size_t r; - if (resolve_gather_index(item, impl->rows, "row", &r) < 0) { - Py_DECREF(fast); - impl_free(out); - return NULL; +/// @brief Gather the included candidate cells of one topk reduction group. +/// @details Group @p g is the whole buffer for axis=None, column @p g for +/// axis=0, or row @p g for axis=1. A masked-out cell (mask == 0.0; +/// NaN counts as included) is skipped. Each kept cell is written to +/// @p scratch as (value, original-index-along-the-reduced-axis). +/// @return the number of included cells (the group's effective length). +static size_t topk_gather_group(const matrix_impl *impl, + const matrix_impl *mask, AxisArg axis, size_t g, + topk_entry *scratch) { + size_t m = 0; + if (!axis.has_axis) { + for (size_t i = 0; i < impl->size; ++i) { + if (mask != NULL && mask->data[i] == 0.0) { + continue; } - - memcpy(out->row_ptrs[i], impl->row_ptrs[r], - impl->columns * sizeof(double)); + scratch[m].value = impl->data[i]; + scratch[m].index = (Py_ssize_t)i; + ++m; + } + } else if (axis.axis == 0) { + for (size_t r = 0; r < impl->rows; ++r) { + if (mask != NULL && mask->row_ptrs[r][g] == 0.0) { + continue; + } + scratch[m].value = impl->row_ptrs[r][g]; + scratch[m].index = (Py_ssize_t)r; + ++m; } } else { - out = impl_new(impl->rows, (size_t)count); - if (out == NULL) { - Py_DECREF(fast); - return NULL; + for (size_t c = 0; c < impl->columns; ++c) { + if (mask != NULL && mask->row_ptrs[g][c] == 0.0) { + continue; + } + scratch[m].value = impl->row_ptrs[g][c]; + scratch[m].index = (Py_ssize_t)c; + ++m; } + } + return m; +} - for (Py_ssize_t i = 0; i < count; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - size_t c; - if (resolve_gather_index(item, impl->columns, "column", &c) < 0) { - Py_DECREF(fast); - impl_free(out); - return NULL; +/// @brief Scatter one group's k sorted (value, index) pairs into the outputs. +/// @details Mirrors @ref topk_gather_group's group layout. The value matrix +/// always receives the k values; @p idx_impl (the as_matrix index +/// matrix) receives the k indices as doubles when non-NULL — it is +/// NULL in the default (list) path, where indices are emitted as +/// Python lists by the caller instead. +static void topk_write_group(matrix_impl *values, matrix_impl *idx_impl, + AxisArg axis, size_t g, size_t kk, + const double *tmp_vals, + const Py_ssize_t *tmp_idx) { + for (size_t j = 0; j < kk; ++j) { + const double iv = (double)tmp_idx[j]; + if (!axis.has_axis) { + values->data[j] = tmp_vals[j]; + if (idx_impl != NULL) { + idx_impl->data[j] = iv; } - - for (size_t r = 0; r < impl->rows; ++r) { - out->row_ptrs[r][i] = impl->row_ptrs[r][c]; + } else if (axis.axis == 0) { + values->row_ptrs[j][g] = tmp_vals[j]; + if (idx_impl != NULL) { + idx_impl->row_ptrs[j][g] = iv; + } + } else { + values->row_ptrs[g][j] = tmp_vals[j]; + if (idx_impl != NULL) { + idx_impl->row_ptrs[g][j] = iv; } } } +} - Py_DECREF(fast); - - PyObject *result = wrap_matrix(type, out); - if (result == NULL) { - impl_free(out); +/// @brief Pack the (values, indices) result tuple, consuming both operands. +/// @details Wraps @p values into a Matrix and packs it with the already-built +/// @p index_obj. Balances references on every path: on any failure +/// the half-built pieces are released and NULL is returned. @p values +/// is consumed (wrapped or freed); @p index_obj's reference is stolen +/// into the tuple (or released on failure). +static PyObject *topk_pack_result(matrix_impl *values, PyObject *index_obj) { + if (index_obj == NULL) { + impl_free(values); + return NULL; } - - return result; + PyObject *values_obj = wrap_impl_or_free(values); + if (values_obj == NULL) { + Py_DECREF(index_obj); + return NULL; + } + PyObject *tuple = PyTuple_Pack(2, values_obj, index_obj); + Py_DECREF(values_obj); + Py_DECREF(index_obj); + return tuple; } -/* Defined alongside Matrix_ass_subscript (the scatter machinery lives next - to the write path); forward-declared here so Matrix_put can share it. */ -static int scatter_axis(matrix_impl *impl, PyObject *indices, int axis, - PyObject *value_op, int accumulate); - -PyObject *Matrix_take(PyObject *op, PyObject *args, PyObject *kwds) { +PyObject *Matrix_topk(PyObject *op, PyObject *args, PyObject *kwds) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -3157,61 +5373,144 @@ PyObject *Matrix_take(PyObject *op, PyObject *args, PyObject *kwds) { return NULL; } - PyObject *indices = NULL; - int axis = 0; - static char *keywords[] = {"indices", "axis", NULL}; - - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|i", keywords, &indices, - &axis)) { + Py_ssize_t k = 0; + PyObject *axis_obj = NULL; + int largest = 1; + PyObject *where_obj = NULL; + int as_matrix = 0; + static char *keywords[] = {"k", "axis", "largest", + "where", "as_matrix", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "n|OpOp", keywords, &k, + &axis_obj, &largest, &where_obj, + &as_matrix)) { return NULL; } - if (axis < 0) { - axis = 2 + axis; + if (k < 1) { + PyErr_SetString(PyExc_ValueError, "k must be a positive integer"); + return NULL; } - if (axis < 0 || axis >= 2) { - PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); + AxisArg axis; + if (parse_validate_normalise_axis(axis_obj, &axis) < 0) { return NULL; } - return gather_axis(impl, indices, axis, Py_TYPE(self)); -} + matrix_impl *mask = NULL; + if (where_obj != NULL && where_obj != Py_None) { + mask = unwrap_mask(where_obj, impl); + if (mask == NULL) { + return NULL; + } + } -PyObject *Matrix_put(PyObject *op, PyObject *args, PyObject *kwds) { - MatrixObject *self = (MatrixObject *)op; - matrix_impl *impl = self->impl; + const size_t kk = (size_t)k; + const size_t rows = impl->rows; + const size_t cols = impl->columns; - if (!impl_check_acquired(impl, true)) { + // The reduced axis must be at least k long (counting every cell, masked or + // not): flattened size for axis=None, rows for axis=0, columns for axis=1. + const size_t axis_len = + !axis.has_axis ? impl->size : (axis.axis == 0 ? rows : cols); + if (kk > axis_len) { + PyErr_Format(PyExc_ValueError, + "k (%zu) cannot exceed the length of the reduced axis (%zu)", + kk, axis_len); return NULL; } - PyObject *indices = NULL; - PyObject *value = NULL; - int axis = 0; - int accumulate = 0; - static char *keywords[] = {"indices", "value", "axis", "accumulate", NULL}; + // One scratch (value, index) buffer sized to the largest group, plus two + // contiguous length-k staging buffers reused across groups. + const size_t group_max = + !axis.has_axis ? impl->size : (axis.axis == 0 ? rows : cols); + topk_entry *scratch = PyMem_Malloc(group_max * sizeof(topk_entry)); + double *tmp_vals = PyMem_Malloc(kk * sizeof(double)); + Py_ssize_t *tmp_idx = PyMem_Malloc(kk * sizeof(Py_ssize_t)); + if (scratch == NULL || tmp_vals == NULL || tmp_idx == NULL) { + PyMem_Free(scratch); + PyMem_Free(tmp_vals); + PyMem_Free(tmp_idx); + return PyErr_NoMemory(); + } + + PyObject *result = NULL; + + // Unified over all three axis layouts: `groups` reduction groups, each + // producing k (value, index) pairs scattered into a values/index matrix of + // shape (vrows x vcols). axis=None is the single flat group -> (1 x k); + // axis=0 reduces down the rows -> (k x cols); axis=1 across the columns -> + // (rows x k). + size_t groups, vrows, vcols; + if (!axis.has_axis) { + groups = 1; + vrows = 1; + vcols = kk; + } else if (axis.axis == 0) { + groups = cols; + vrows = kk; + vcols = cols; + } else { + groups = rows; + vrows = rows; + vcols = kk; + } - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|ip", keywords, &indices, - &value, &axis, &accumulate)) { - return NULL; + matrix_impl *values = impl_new(vrows, vcols); + if (values == NULL) { + goto done; } - if (axis < 0) { - axis = 2 + axis; + // Default index form is a list (a same-length list[int] for axis=None, or a + // list of `groups` per-group int lists for an axis). Pass as_matrix=True for + // a same-shape Matrix of float index positions (idx_impl) instead. + matrix_impl *idx_impl = NULL; + PyObject *index_lists = NULL; + if (as_matrix) { + idx_impl = impl_new(vrows, vcols); + if (idx_impl == NULL) { + impl_free(values); + goto done; + } + } else if (axis.has_axis) { + index_lists = PyList_New((Py_ssize_t)groups); + if (index_lists == NULL) { + impl_free(values); + goto done; + } } - if (axis < 0 || axis >= 2) { - PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); - return NULL; + for (size_t g = 0; g < groups; ++g) { + size_t m = topk_gather_group(impl, mask, axis, g, scratch); + topk_sort_pad(scratch, m, kk, largest, tmp_vals, tmp_idx); + topk_write_group(values, idx_impl, axis, g, kk, tmp_vals, tmp_idx); + if (!as_matrix && axis.has_axis) { + PyObject *sub = topk_build_int_list(tmp_idx, kk); + if (sub == NULL) { + Py_DECREF(index_lists); + impl_free(values); + goto done; + } + PyList_SET_ITEM(index_lists, (Py_ssize_t)g, sub); + } } - if (scatter_axis(impl, indices, axis, value, accumulate) < 0) { - return NULL; + PyObject *index_obj; + if (as_matrix) { + index_obj = wrap_impl_or_free(idx_impl); + } else if (axis.has_axis) { + index_obj = index_lists; + } else { + // Flat list form: groups == 1, so tmp_idx still holds the single group's + // sorted indices. + index_obj = topk_build_int_list(tmp_idx, kk); } + result = topk_pack_result(values, index_obj); - Py_INCREF(self); - return (PyObject *)self; +done: + PyMem_Free(scratch); + PyMem_Free(tmp_vals); + PyMem_Free(tmp_idx); + return result; } static PyObject *Matrix_allclose(PyObject *cls, PyObject *args, @@ -3232,6 +5531,11 @@ static PyObject *Matrix_allclose(PyObject *cls, PyObject *args, MatrixObject *lhs = (MatrixObject *)lhs_op; MatrixObject *rhs = (MatrixObject *)rhs_op; + if (!impl_check_acquired(lhs->impl, true) || + !impl_check_acquired(rhs->impl, true)) { + return NULL; + } + if (impl_allclose(lhs->impl, rhs->impl, rtol, atol, equal_nan)) { Py_RETURN_TRUE; } @@ -3284,12 +5588,15 @@ static int where_resolve_operand(PyObject *op, size_t rows, size_t columns, /// ``ValueError``. A 1x1 matrix is treated as a matrix (not a scalar) /// and so must match the mask shape. NaN mask elements are non-zero /// and select ``a``. -static PyObject *Matrix_where(PyObject *cls, PyObject *args) { +static PyObject *Matrix_where(PyObject *cls, PyObject *args, PyObject *kwds) { PyObject *mask_op = NULL; PyObject *a_op = NULL; PyObject *b_op = NULL; + PyObject *out_target = NULL; + static char *kwlist[] = {"", "", "", "out", NULL}; - if (!PyArg_ParseTuple(args, "O!OO", cls, &mask_op, &a_op, &b_op)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|$O", kwlist, cls, &mask_op, + &a_op, &b_op, &out_target)) { return NULL; } @@ -3314,7 +5621,12 @@ static PyObject *Matrix_where(PyObject *cls, PyObject *args) { goto done; } - out = impl_new(rows, columns); + const bool fresh = (out_target == NULL || out_target == Py_None); + if (fresh) { + out = impl_new(rows, columns); + } else { + out = use_out_target(out_target, &result, rows, columns); + } if (out == NULL) { goto done; } @@ -3327,7 +5639,9 @@ static PyObject *Matrix_where(PyObject *cls, PyObject *args) { *outp = (*mp != 0.0) ? av : bv; } - result = wrap_impl_or_free(out); + if (fresh) { + result = wrap_impl_or_free(out); + } done: IMPL_DECREF(a_mat); @@ -3407,19 +5721,65 @@ static PyObject *Matrix_ones(PyObject *cls, PyObject *args) { return wrap_impl_or_free(impl); } -const double RAND_MAX_D = (double)RAND_MAX; +static PyObject *Matrix_full(PyObject *cls, PyObject *args) { + PyObject *size = NULL; + double value; + + if (!PyArg_ParseTuple(args, "Od", &size, &value)) { + return NULL; + } + + size_t rows; + size_t columns; + + if (parse_dims(size, &rows, &columns) < 0) { + return NULL; + } + + matrix_impl *impl = impl_new(rows, columns); + if (impl == NULL) { + return NULL; + } + + double *ptr = impl->data; + for (size_t i = 0; i < impl->size; ++i, ++ptr) { + *ptr = value; + } + + return wrap_impl_or_free(impl); +} + +/// @brief splitmix64: draw the next 64-bit value and advance the state. +/// @details A single-word PRNG kept per interpreter in the module state, so +/// parallel worker sub-interpreters draw from independent streams +/// and never share mutable RNG state (the old ``rand()``/``srand()`` +/// were process-global -- non-reproducible across workers and a data +/// race off glibc / on free-threaded builds). splitmix64 is chosen +/// for its tiny footprint and good distribution, and mixes even +/// low-entropy seeds well so ``seed(1)`` and ``seed(2)`` diverge. +static inline uint64_t prng_next_u64(uint64_t *state) { + uint64_t z = (*state += 0x9E3779B97F4A7C15ULL); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} + +/// @brief A uniform double in [0, 1) using the top 53 bits (full mantissa). +static inline double prng_next_unit(uint64_t *state) { + return (double)(prng_next_u64(state) >> 11) * (1.0 / 9007199254740992.0); +} -static double sample_uniform(double min, double max) { - double val = (double)rand() / RAND_MAX_D; +static double sample_uniform(uint64_t *rng, double min, double max) { + double val = prng_next_unit(rng); return (val * (max - min)) + min; } -static void sample_normal(double *values, size_t n, double mean, +static void sample_normal(uint64_t *rng, double *values, size_t n, double mean, double stddev) { size_t i = 0; while (i < n) { - double u = sample_uniform(-1, 1); - double v = sample_uniform(-1, 1); + double u = sample_uniform(rng, -1, 1); + double v = sample_uniform(rng, -1, 1); double s = (u * u) + (v * v); if (s > 0 && s < 1) { double factor = sqrt(-2 * log(s) / s); @@ -3446,9 +5806,14 @@ static PyObject *Matrix_normal(PyObject *cls, PyObject *args, return NULL; } + _math_module_state *state = math_local_state(); + if (state == NULL) { + return NULL; + } + if (Py_IsNone(size)) { double value = 0; - sample_normal(&value, 1, mean, stddev); + sample_normal(&state->prng_state, &value, 1, mean, stddev); return PyFloat_FromDouble(value); } @@ -3462,7 +5827,7 @@ static PyObject *Matrix_normal(PyObject *cls, PyObject *args, return NULL; } - sample_normal(impl->data, impl->size, mean, stddev); + sample_normal(&state->prng_state, impl->data, impl->size, mean, stddev); return wrap_impl_or_free(impl); } @@ -3480,8 +5845,14 @@ static PyObject *Matrix_uniform(PyObject *cls, PyObject *args, return NULL; } + _math_module_state *state = math_local_state(); + if (state == NULL) { + return NULL; + } + if (Py_IsNone(size)) { - return PyFloat_FromDouble(sample_uniform(minval, maxval)); + return PyFloat_FromDouble( + sample_uniform(&state->prng_state, minval, maxval)); } size_t rows, columns; @@ -3496,7 +5867,7 @@ static PyObject *Matrix_uniform(PyObject *cls, PyObject *args, double *ptr = impl->data; for (size_t i = 0; i < impl->size; ++i, ptr++) { - *ptr = sample_uniform(minval, maxval); + *ptr = sample_uniform(&state->prng_state, minval, maxval); } return wrap_impl_or_free(impl); @@ -3509,7 +5880,14 @@ static PyObject *Matrix_seed(PyObject *cls, PyObject *args) { return NULL; } - srand((unsigned int)value); + _math_module_state *state = math_local_state(); + if (state == NULL) { + return NULL; + } + + // Seed this interpreter's stream only. splitmix64 mixes the raw seed + // internally, so distinct seeds diverge immediately. + state->prng_state = (uint64_t)value; Py_RETURN_NONE; } @@ -3540,493 +5918,211 @@ static int unwrap_and_get_shape(PyObject *object, shape *shape, if (Py_TYPE(object) == LOCAL_STATE->matrix_type) { MatrixObject *matrix = (MatrixObject *)object; matrix_impl *impl = matrix->impl; - if (!impl_check_acquired(impl, true)) { - return -1; - } - - shape->rows = impl->rows; - shape->columns = impl->columns; - return 0; - } - - PyObject *fast = - PySequence_Fast(object, "object must be a Matrix, List, or Tuple"); - if (fast == NULL) { - return -1; - } - - if (seq_to_column) { - shape->rows = PySequence_Fast_GET_SIZE(fast); - shape->columns = 1; - } else { - shape->rows = 1; - shape->columns = PySequence_Fast_GET_SIZE(fast); - } - Py_DECREF(fast); - return 0; -} - -static PyObject *Matrix_concat(PyObject *cls, PyObject *args, - PyObject *kwargs) { - PyObject *matrices = NULL; - int axis = 0; - - static char *kwlist[] = {"values", "axis", NULL}; - - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|i", kwlist, &matrices, - &axis)) { - return NULL; - } - - if (axis < 0) { - axis = 2 + axis; - } - - if (axis >= 2) { - PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); - return NULL; - } - - const char *err_msg = - "Matrices must be specified as a list or a tuple of Matrix"; - PyObject *fast = PySequence_Fast(matrices, err_msg); - if (fast == NULL) { - return NULL; - } - - Py_ssize_t size = PySequence_Fast_GET_SIZE(fast); - if (size <= 0) { - Py_DECREF(fast); - Py_RETURN_NONE; - } - - shape shape; - size_t rows = 0; - size_t columns = 0; - range rows_range; - range columns_range; - matrix_impl *out; - if (axis == 0) { - for (Py_ssize_t i = 0; i < size; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - if (unwrap_and_get_shape(item, &shape, false) < 0) { - Py_DECREF(fast); - return NULL; - } - - if (i == 0) { - columns = shape.columns; - } else if (shape.columns != columns) { - PyErr_SetString( - PyExc_AssertionError, - "all sub-matrices must have the same number of columns"); - Py_DECREF(fast); - return NULL; - } - - rows += shape.rows; - } - - out = impl_new(rows, columns); - if (out == NULL) { - Py_DECREF(fast); - return NULL; - } - - rows_range.start = 0; - rows_range.step = 1; - columns_range.start = 0; - columns_range.stop = (Py_ssize_t)columns; - columns_range.count = columns; - columns_range.step = 1; - for (Py_ssize_t i = 0; i < size; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - matrix_impl *impl = unwrap_matrix(item, false); - if (impl == NULL) { - Py_DECREF(fast); - impl_free(out); - return NULL; - } - - rows_range.count = impl->rows; - rows_range.stop = rows_range.start + (Py_ssize_t)impl->rows; - impl_set(out, &rows_range, &columns_range, impl); - IMPL_DECREF(impl); - rows_range.start = rows_range.stop; - } - - Py_DECREF(fast); - return wrap_impl_or_free(out); - } - - for (Py_ssize_t i = 0; i < size; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - if (unwrap_and_get_shape(item, &shape, true) < 0) { - Py_DECREF(fast); - return NULL; - } - - if (i == 0) { - rows = shape.rows; - } else if (shape.rows != rows) { - PyErr_SetString(PyExc_AssertionError, - "all sub-matrices must have the same number of rows"); - Py_DECREF(fast); - return NULL; - } - - columns += shape.columns; - } - - out = impl_new(rows, columns); - if (out == NULL) { - Py_DECREF(fast); - return NULL; - } - - columns_range.start = 0; - columns_range.step = 1; - rows_range.start = 0; - rows_range.stop = (Py_ssize_t)rows; - rows_range.count = rows; - rows_range.step = 1; - for (Py_ssize_t i = 0; i < size; ++i) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - matrix_impl *impl = unwrap_matrix(item, true); - if (impl == NULL) { - Py_DECREF(fast); - impl_free(out); - return NULL; - } - - columns_range.count = impl->columns; - columns_range.stop = columns_range.start + (Py_ssize_t)impl->columns; - impl_set(out, &rows_range, &columns_range, impl); - IMPL_DECREF(impl); - columns_range.start = columns_range.stop; - } - - Py_DECREF(fast); - return wrap_impl_or_free(out); -} - -/// @brief Pickle support: reduce a matrix to its raw double buffer. -/// @details Returns ``(_matrix_unpickle, (rows, columns, payload))`` where -/// ``payload`` is the native-endian byte image of the contiguous -/// row-major ``double`` data. Reconstruction is a single ``memcpy``, -/// so pickling cost is linear in the element count with no per-element -/// Python object churn. ``copy.copy`` and ``copy.deepcopy`` route -/// through the same path. The current interpreter must own the matrix. -static PyObject *Matrix_reduce(PyObject *op, PyObject *Py_UNUSED(dummy)) { - MatrixObject *self = (MatrixObject *)op; - matrix_impl *impl = self->impl; + if (!impl_check_acquired(impl, true)) { + return -1; + } - if (impl == NULL) { - PyErr_SetString(PyExc_ValueError, "Cannot pickle an uninitialized matrix"); - return NULL; + shape->rows = impl->rows; + shape->columns = impl->columns; + return 0; } - if (!impl_check_acquired(impl, true)) { - return NULL; + PyObject *fast = + PySequence_Fast(object, "object must be a Matrix, List, or Tuple"); + if (fast == NULL) { + return -1; } - PyObject *rebuild = LOCAL_STATE->matrix_unpickle; - - PyObject *payload = PyBytes_FromStringAndSize( - (const char *)impl->data, (Py_ssize_t)(impl->size * sizeof(double))); - if (payload == NULL) { - return NULL; + if (seq_to_column) { + shape->rows = PySequence_Fast_GET_SIZE(fast); + shape->columns = 1; + } else { + shape->rows = 1; + shape->columns = PySequence_Fast_GET_SIZE(fast); } - - return Py_BuildValue("(O(nnN))", rebuild, (Py_ssize_t)impl->rows, - (Py_ssize_t)impl->columns, payload); + Py_DECREF(fast); + return 0; } -static PyObject *Matrix_Less_compare(PyObject *self, PyObject *other); -static PyObject *Matrix_LessEqual_compare(PyObject *self, PyObject *other); -static PyObject *Matrix_Greater_compare(PyObject *self, PyObject *other); -static PyObject *Matrix_GreaterEqual_compare(PyObject *self, PyObject *other); -static PyObject *Matrix_Equal_compare(PyObject *self, PyObject *other); -static PyObject *Matrix_NotEqual_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_concat(PyObject *cls, PyObject *args, + PyObject *kwargs) { + PyObject *matrices = NULL; + int axis = 0; -static PyObject *Matrix_add_method(PyObject *self, PyObject *args, - PyObject *kwds); -static PyObject *Matrix_subtract_method(PyObject *self, PyObject *args, - PyObject *kwds); -static PyObject *Matrix_multiply_method(PyObject *self, PyObject *args, - PyObject *kwds); -static PyObject *Matrix_divide_method(PyObject *self, PyObject *args, - PyObject *kwds); + static char *kwlist[] = {"values", "axis", NULL}; -static PyMethodDef Matrix_methods[] = { - {"transpose", (PyCFunction)Matrix_transpose, METH_VARARGS | METH_KEYWORDS, - "transpose($self, /, in_place=False)\n--\n\n" - "Return a transposed copy, or transpose ``self`` in place when " - "``in_place=True`` (in which case ``self`` is returned)."}, - {"sum", (PyCFunction)Matrix_Sum_method, METH_VARARGS | METH_KEYWORDS, - "sum($self, /, axis=None)\n--\n\nSum of elements."}, - {"mean", (PyCFunction)Matrix_Mean_method, METH_VARARGS | METH_KEYWORDS, - "mean($self, /, axis=None)\n--\n\nMean of elements."}, - {"magnitude", (PyCFunction)Matrix_Magnitude_method, - METH_VARARGS | METH_KEYWORDS, - "magnitude($self, /, axis=None)\n--\n\nEuclidean magnitude."}, - {"magnitude_squared", (PyCFunction)Matrix_MagnitudeSquared_method, - METH_VARARGS | METH_KEYWORDS, - "magnitude_squared($self, /, axis=None)\n--\n\n" - "Sum of squared elements (Euclidean magnitude without the sqrt)."}, - {"vecdot", (PyCFunction)Matrix_vecdot, METH_VARARGS | METH_KEYWORDS, - "vecdot($self, other, /, axis=None)\n--\n\n" - "Axis-aware inner product: sum of element-wise products. " - "Equivalent to numpy.linalg.vecdot for 1-D inputs with axis=None; " - "**not** equivalent to numpy.dot."}, - {"fma", (PyCFunction)Matrix_fma, METH_VARARGS | METH_KEYWORDS, - "fma($self, b, c, /, in_place=False)\n--\n\n" - "Fused multiply-add: single-rounding ``self*b + c`` (libc fma()). " - "b and c may each be a same-shape matrix, a 1x1 matrix, a row or column " - "vector that broadcasts, or a scalar; other shapes raise ValueError. " - "The single rounding differs from ``self*b + c`` (which rounds twice) by " - "up to half a ULP, so compare results with allclose(), not ==."}, - {"scaled_add", (PyCFunction)Matrix_scaled_add, METH_VARARGS | METH_KEYWORDS, - "scaled_add($self, s, x, /, in_place=False)\n--\n\n" - "Scaled add ``self + s * x`` with two roundings; the two-rounding " - "sibling of fma().\n\n" - "s is the scale -- a scalar, a 1x1 matrix, a row or column vector that " - "broadcasts, or a same-shape matrix (the same operand rules as fma's " - "multiplier) -- and x is a same-shape matrix. The arithmetic rounds " - "twice (``round(round(s*x) + self)``), so the result is bit-for-bit " - "identical to ``self + s * x`` and -- unlike fma() -- never fuses to a " - "single rounding. With in_place=True the result is written into self's " - "buffer (allocating nothing) and self is returned; otherwise a new " - "matrix is returned. s and x may alias self. A shape mismatch raises " - "ValueError before any write."}, - {"add", (PyCFunction)Matrix_add_method, METH_VARARGS | METH_KEYWORDS, - "add($self, other, /, *, out=None)\n--\n\n" - "Element-wise ``self + other`` with the same broadcasting as ``+`` " - "(other may be a scalar). With ``out`` (a same-shape matrix) the result " - "is written there and returned instead of allocating a new matrix; out " - "may alias an input. Bit-for-bit identical to ``self + other``. A shape " - "mismatch between out and the result raises ValueError before any write."}, - {"subtract", (PyCFunction)Matrix_subtract_method, - METH_VARARGS | METH_KEYWORDS, - "subtract($self, other, /, *, out=None)\n--\n\n" - "Element-wise ``self - other`` with the same broadcasting as ``-`` " - "(other may be a scalar). With ``out`` (a same-shape matrix) the result " - "is written there and returned instead of allocating a new matrix; out " - "may alias an input. Bit-for-bit identical to ``self - other``. A shape " - "mismatch between out and the result raises ValueError before any write."}, - {"multiply", (PyCFunction)Matrix_multiply_method, - METH_VARARGS | METH_KEYWORDS, - "multiply($self, other, /, *, out=None)\n--\n\n" - "Element-wise ``self * other`` with the same broadcasting as ``*`` " - "(other may be a scalar). With ``out`` (a same-shape matrix) the result " - "is written there and returned instead of allocating a new matrix; out " - "may alias an input. Bit-for-bit identical to ``self * other``. A shape " - "mismatch between out and the result raises ValueError before any write."}, - {"divide", (PyCFunction)Matrix_divide_method, METH_VARARGS | METH_KEYWORDS, - "divide($self, other, /, *, out=None)\n--\n\n" - "Element-wise ``self / other`` with the same broadcasting as ``/`` " - "(other may be a scalar). With ``out`` (a same-shape matrix) the result " - "is written there and returned instead of allocating a new matrix; out " - "may alias an input. Bit-for-bit identical to ``self / other``. A shape " - "mismatch between out and the result raises ValueError before any write."}, - {"cross", (PyCFunction)Matrix_cross, METH_VARARGS | METH_KEYWORDS, - "cross($self, other, /, axis=None)\n--\n\n" - "2D (scalar z-component) or 3D cross product against another " - "vector or batch. 1x2 / 2x1 inputs return a float; 1x3 / 3x1 return " - "a Matrix preserving self's orientation. Nx2 / 2xN row/column " - "batches return per-vector scalars (Mx1 / 1xN); Nx3 / 3xN return " - "same-shape batches. Batch operands accept either a same-shape " - "other or a single 2D/3D vector (1xK / Kx1) broadcast against " - "every per-vector slot \u2014 ``self`` must be the batch (cross is " - "anticommutative). ``axis`` disambiguates the 2x2 / 3x3 squares " - "(default rows, ``axis=0`` for columns)."}, - {"normalize", (PyCFunction)Matrix_normalize, METH_VARARGS | METH_KEYWORDS, - "normalize($self, /, axis=None, in_place=False)\n--\n\n" - "Divide elements by their magnitude. ``axis=None`` divides by the " - "matrix's total magnitude; ``axis=0`` divides each column by its own " - "magnitude; ``axis=1`` divides each row by its own magnitude. Rows or " - "columns whose magnitude is zero are left as the all-zero vector. " - "Sub-normal magnitudes may overflow during division; threshold with " - "magnitude_squared() if safety matters. When ``in_place=True``, mutates " - "``self`` and returns it."}, - {"perpendicular", (PyCFunction)Matrix_perpendicular, - METH_VARARGS | METH_KEYWORDS, - "perpendicular($self, /, axis=None, in_place=False)\n--\n\n" - "Rotate every 2D vector 90 degrees counter-clockwise: ``(x, y) -> " - "(-y, x)``. Accepts a single 2D vector (``1x2`` or ``2x1``), a row " - "batch (``Nx2``), or a column batch (``2xN``). On the ambiguous " - "``2x2`` shape the default is per-row; pass ``axis=0`` to force " - "per-column. When ``in_place=True``, mutates ``self`` and returns it."}, - {"angle", (PyCFunction)Matrix_angle, METH_VARARGS | METH_KEYWORDS, - "angle($self, /, axis=None)\n--\n\n" - "Polar angle (``atan2(y, x)``) of every 2D vector. Returns a float " - "for a single 2D vector, an ``Mx1`` column matrix for an ``Nx2`` row " - "batch, or a ``1xN`` row matrix for a ``2xN`` column batch. On the " - "ambiguous ``2x2`` shape the default is per-row; pass ``axis=0`` to " - "force per-column."}, - {"min", (PyCFunction)Matrix_Minimum_method, METH_VARARGS | METH_KEYWORDS, - "min($self, /, axis=None)\n--\n\nMinimum of elements."}, - {"max", (PyCFunction)Matrix_Maximum_method, METH_VARARGS | METH_KEYWORDS, - "max($self, /, axis=None)\n--\n\nMaximum of elements."}, - {"argmin", (PyCFunction)Matrix_argmin_method, METH_VARARGS | METH_KEYWORDS, - "argmin($self, /, axis=None)\n--\n\n" - "Index of the minimum element (first occurrence on ties).\n\n" - "NaN elements are skipped unless the running extreme starts at NaN\n" - "(element 0 along the reduced axis), which pins the result to that\n" - "position. This differs from NumPy, which propagates NaN."}, - {"argmax", (PyCFunction)Matrix_argmax_method, METH_VARARGS | METH_KEYWORDS, - "argmax($self, /, axis=None)\n--\n\n" - "Index of the maximum element (first occurrence on ties).\n\n" - "NaN elements are skipped unless the running extreme starts at NaN\n" - "(element 0 along the reduced axis), which pins the result to that\n" - "position. This differs from NumPy, which propagates NaN."}, - {"ceil", (PyCFunction)Matrix_Ceil_method, METH_VARARGS | METH_KEYWORDS, - "ceil($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise ceiling. With ``out`` (a same-shape matrix) the result " - "is written there and returned; ``out`` and ``in_place`` are mutually " - "exclusive."}, - {"floor", (PyCFunction)Matrix_Floor_method, METH_VARARGS | METH_KEYWORDS, - "floor($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise floor. With ``out`` (a same-shape matrix) the result " - "is written there and returned; ``out`` and ``in_place`` are mutually " - "exclusive."}, - {"round", (PyCFunction)Matrix_Round_method, METH_VARARGS | METH_KEYWORDS, - "round($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise rounding (banker's; IEEE round-half-to-even). With ``out`` " - "(a same-shape matrix) the result is written there and returned; " - "``out`` and ``in_place`` are mutually exclusive."}, - {"negate", (PyCFunction)Matrix_Negate_method, METH_VARARGS | METH_KEYWORDS, - "negate($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise negation. With ``out`` (a same-shape matrix) the result " - "is written there and returned; ``out`` and ``in_place`` are mutually " - "exclusive."}, - {"abs", (PyCFunction)Matrix_Abs_method, METH_VARARGS | METH_KEYWORDS, - "abs($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise absolute value. With ``out`` (a same-shape matrix) the " - "result is written there and returned; ``out`` and ``in_place`` are " - "mutually exclusive."}, - {"sqrt", (PyCFunction)Matrix_Sqrt_method, METH_VARARGS | METH_KEYWORDS, - "sqrt($self, /, in_place=False, *, out=None)\n--\n\n" - "Element-wise square root. Negative elements yield NaN. With ``out`` " - "(a same-shape matrix) the result is written there and returned; " - "``out`` and ``in_place`` are mutually exclusive."}, - {"less", (PyCFunction)Matrix_Less_compare, METH_O, - "less($self, other, /)\n--\n\n" - "Element-wise ``self < other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 0. Distinct from the ``<`` operator, which returns a " - "single bool."}, - {"less_equal", (PyCFunction)Matrix_LessEqual_compare, METH_O, - "less_equal($self, other, /)\n--\n\n" - "Element-wise ``self <= other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 0. Distinct from the ``<=`` operator, which returns a " - "single bool."}, - {"greater", (PyCFunction)Matrix_Greater_compare, METH_O, - "greater($self, other, /)\n--\n\n" - "Element-wise ``self > other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 0. Distinct from the ``>`` operator, which returns a " - "single bool."}, - {"greater_equal", (PyCFunction)Matrix_GreaterEqual_compare, METH_O, - "greater_equal($self, other, /)\n--\n\n" - "Element-wise ``self >= other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 0. Distinct from the ``>=`` operator, which returns a " - "single bool."}, - {"equal", (PyCFunction)Matrix_Equal_compare, METH_O, - "equal($self, other, /)\n--\n\n" - "Element-wise ``self == other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 0. Distinct from the ``==`` operator, which returns a " - "single bool."}, - {"not_equal", (PyCFunction)Matrix_NotEqual_compare, METH_O, - "not_equal($self, other, /)\n--\n\n" - "Element-wise ``self != other`` as a 0/1 mask matrix. ``other`` may be a " - "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " - "row/column vector that broadcasts, or a list/tuple of numbers. NaN " - "comparisons yield 1. Distinct from the ``!=`` operator, which returns a " - "single bool."}, - {"clip", (PyCFunction)Matrix_clip, METH_VARARGS | METH_KEYWORDS, - "clip($self, min=None, max=None)\n--\n\n" - "Clamp elements to [min, max]; either bound may be omitted.\n\n" - "The first argument is the lower bound and the second the upper\n" - "bound. Pass None (or omit) to leave a side unbounded, so\n" - "clip(min=0) clamps only below and clip(max=255) only above. Raises\n" - "ValueError if both bounds are omitted."}, - {"copy", Matrix_copy, METH_NOARGS, - "copy($self, /)\n--\n\nReturn a deep copy."}, - {"__reduce__", Matrix_reduce, METH_NOARGS, - "__reduce__($self, /)\n--\n\n" - "Pickle helper: serialize the matrix to its raw double buffer."}, - {"take", (PyCFunction)Matrix_take, METH_VARARGS | METH_KEYWORDS, - "take($self, indices, axis=0)\n--\n\n" - "Take rows or columns by index into a new matrix.\n\n" - "indices is a 1-D list or tuple of ints selecting whole rows\n" - "(axis=0) or columns (axis=1). Negative indices count from the\n" - "end; duplicates repeat the row or column; an out-of-range index\n" - "raises IndexError, and an empty index sequence raises IndexError."}, - {"put", (PyCFunction)Matrix_put, METH_VARARGS | METH_KEYWORDS, - "put($self, indices, value, axis=0, accumulate=False)\n--\n\n" - "Assign value into the selected rows or columns in place.\n\n" - "The write-side counterpart of take(): value may be a scalar, a\n" - "1x1 matrix, or a matrix matching the selection shape. All indices\n" - "and the value shape are validated before any write, so a rejected\n" - "call leaves the matrix unchanged. Negative indices count from the\n" - "end; an out-of-range or empty index raises IndexError. With\n" - "accumulate=False (the default) duplicate indices follow\n" - "last-write-wins; with accumulate=True the values are added into the\n" - "selection, so duplicate indices fold additively (scatter-add).\n" - "Returns self."}, - {"allclose", (PyCFunction)Matrix_allclose, - METH_VARARGS | METH_KEYWORDS | METH_CLASS, - "allclose($type, lhs, rhs, /, rtol=1e-05, atol=1e-08, " - "equal_nan=False)\n--\n\n" - "Check element-wise equality within tolerance."}, - {"where", (PyCFunction)Matrix_where, METH_VARARGS | METH_CLASS, - "where($type, mask, a, b, /)\n--\n\n" - "Select element-wise from a or b on a truthy mask.\n\n" - "Returns a fresh matrix taking a where the corresponding mask\n" - "element is non-zero and b elsewhere. a and b may each be a scalar,\n" - "a matrix matching the mask's shape, or a list/tuple of numbers\n" - "(taken as a 1xN row vector that must then match the mask shape);\n" - "other shapes raise ValueError. A 1x1 matrix is treated as a matrix\n" - "(not a scalar) and so must match the mask shape. NaN mask elements\n" - "are non-zero and select a."}, - {"zeros", Matrix_zeros, METH_VARARGS | METH_CLASS, - "zeros($type, size, /)\n--\n\nCreate a zero-filled matrix."}, - {"ones", Matrix_ones, METH_VARARGS | METH_CLASS, - "ones($type, size, /)\n--\n\nCreate a matrix of ones."}, - {"normal", (PyCFunction)Matrix_normal, - METH_VARARGS | METH_KEYWORDS | METH_CLASS, - "normal($type, mean=0.0, stddev=1.0, /, size=None)\n--\n\n" - "Sample from a normal distribution."}, - {"uniform", (PyCFunction)Matrix_uniform, - METH_VARARGS | METH_KEYWORDS | METH_CLASS, - "uniform($type, minval=0.0, maxval=1.0, /, size=None)\n--\n\n" - "Sample from a uniform distribution."}, - {"seed", Matrix_seed, METH_VARARGS | METH_CLASS, - "seed($type, value, /)\n--\n\n" - "Seed the random generator used by normal() and uniform().\n\n" - "The generator is the process-global C library PRNG shared by every\n" - "sub-interpreter, so a seed only makes subsequent draws reproducible\n" - "when random generation stays on a single thread; concurrent draws\n" - "interleave on the shared state. The sequence is also not portable\n" - "across platforms."}, - {"vector", Matrix_vector, METH_VARARGS | METH_CLASS, - "vector($type, values, /, as_column=False)\n--\n\n" - "Create a vector from a sequence."}, - {"concat", (PyCFunction)Matrix_concat, - METH_VARARGS | METH_KEYWORDS | METH_CLASS, - "concat($type, values, /, axis=0)\n--\n\n" - "Concatenate matrices along an axis."}, - {NULL} /* Sentinel */ -}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|i", kwlist, &matrices, + &axis)) { + return NULL; + } + + if (axis < 0) { + axis = 2 + axis; + } + + if (axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); + return NULL; + } + + const char *err_msg = + "Matrices must be specified as a list or a tuple of Matrix"; + PyObject *fast = PySequence_Fast(matrices, err_msg); + if (fast == NULL) { + return NULL; + } + + Py_ssize_t size = PySequence_Fast_GET_SIZE(fast); + if (size <= 0) { + Py_DECREF(fast); + Py_RETURN_NONE; + } + + shape shape; + size_t rows = 0; + size_t columns = 0; + range rows_range; + range columns_range; + matrix_impl *out; + if (axis == 0) { + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + if (unwrap_and_get_shape(item, &shape, false) < 0) { + Py_DECREF(fast); + return NULL; + } + + if (i == 0) { + columns = shape.columns; + } else if (shape.columns != columns) { + PyErr_SetString( + PyExc_AssertionError, + "all sub-matrices must have the same number of columns"); + Py_DECREF(fast); + return NULL; + } + + rows += shape.rows; + } + + out = impl_new(rows, columns); + if (out == NULL) { + Py_DECREF(fast); + return NULL; + } + + rows_range.start = 0; + rows_range.step = 1; + columns_range.start = 0; + columns_range.stop = (Py_ssize_t)columns; + columns_range.count = columns; + columns_range.step = 1; + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + matrix_impl *impl = unwrap_matrix(item, false); + if (impl == NULL) { + Py_DECREF(fast); + impl_free(out); + return NULL; + } + + rows_range.count = impl->rows; + rows_range.stop = rows_range.start + (Py_ssize_t)impl->rows; + impl_set(out, &rows_range, &columns_range, impl); + IMPL_DECREF(impl); + rows_range.start = rows_range.stop; + } + + Py_DECREF(fast); + return wrap_impl_or_free(out); + } + + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + if (unwrap_and_get_shape(item, &shape, true) < 0) { + Py_DECREF(fast); + return NULL; + } + + if (i == 0) { + rows = shape.rows; + } else if (shape.rows != rows) { + PyErr_SetString(PyExc_AssertionError, + "all sub-matrices must have the same number of rows"); + Py_DECREF(fast); + return NULL; + } + + columns += shape.columns; + } + + out = impl_new(rows, columns); + if (out == NULL) { + Py_DECREF(fast); + return NULL; + } + + columns_range.start = 0; + columns_range.step = 1; + rows_range.start = 0; + rows_range.stop = (Py_ssize_t)rows; + rows_range.count = rows; + rows_range.step = 1; + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + matrix_impl *impl = unwrap_matrix(item, true); + if (impl == NULL) { + Py_DECREF(fast); + impl_free(out); + return NULL; + } + + columns_range.count = impl->columns; + columns_range.stop = columns_range.start + (Py_ssize_t)impl->columns; + impl_set(out, &rows_range, &columns_range, impl); + IMPL_DECREF(impl); + columns_range.start = columns_range.stop; + } + + Py_DECREF(fast); + return wrap_impl_or_free(out); +} + +/// @brief Pickle support: reduce a matrix to its raw double buffer. +/// @details Returns ``(_matrix_unpickle, (rows, columns, payload))`` where +/// ``payload`` is the native-endian byte image of the contiguous +/// row-major ``double`` data. Reconstruction is a single ``memcpy``, +/// so pickling cost is linear in the element count with no per-element +/// Python object churn. ``copy.copy`` and ``copy.deepcopy`` route +/// through the same path. The current interpreter must own the matrix. +static PyObject *Matrix_reduce(PyObject *op, PyObject *Py_UNUSED(dummy)) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + + if (impl == NULL) { + PyErr_SetString(PyExc_ValueError, "Cannot pickle an uninitialized matrix"); + return NULL; + } + + if (!impl_check_acquired(impl, true)) { + return NULL; + } + + _math_module_state *state = math_local_state(); + if (state == NULL) { + return NULL; + } + PyObject *rebuild = state->matrix_unpickle; + + PyObject *payload = PyBytes_FromStringAndSize( + (const char *)impl->data, (Py_ssize_t)(impl->size * sizeof(double))); + if (payload == NULL) { + return NULL; + } + + return Py_BuildValue("(O(nnN))", rebuild, (Py_ssize_t)impl->rows, + (Py_ssize_t)impl->columns, payload); +} static PyObject *Matrix_get_rows(PyObject *op, void *Py_UNUSED(dummy)) { MatrixObject *self = (MatrixObject *)op; @@ -4097,6 +6193,9 @@ static int Matrix_set_x(PyObject *op, PyObject *arg, void *Py_UNUSED(dummy)) { } if (!unwrap_double(arg, impl->data)) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "expected a real number"); + } return -1; } @@ -4131,6 +6230,9 @@ static int Matrix_set_y(PyObject *op, PyObject *arg, void *Py_UNUSED(dummy)) { } if (!unwrap_double(arg, impl->data + 1)) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "expected a real number"); + } return -1; } @@ -4165,6 +6267,9 @@ static int Matrix_set_z(PyObject *op, PyObject *arg, void *Py_UNUSED(dummy)) { } if (!unwrap_double(arg, impl->data + 2)) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "expected a real number"); + } return -1; } @@ -4199,6 +6304,9 @@ static int Matrix_set_w(PyObject *op, PyObject *arg, void *Py_UNUSED(dummy)) { } if (!unwrap_double(arg, impl->data + 3)) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "expected a real number"); + } return -1; } @@ -4234,6 +6342,7 @@ static PyGetSetDef Matrix_getset[] = { {"z", (getter)Matrix_get_z, (setter)Matrix_set_z, NULL, NULL}, {"w", (getter)Matrix_get_w, (setter)Matrix_set_w, NULL, NULL}, {"shape", (getter)Matrix_get_shape, NULL, NULL, NULL}, + {"view", (getter)Matrix_get_view, NULL, NULL, NULL}, {NULL} /* Sentinel */ }; @@ -4467,11 +6576,21 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, reuses Matrix_binary_op for the full broadcast routing (same-shape, scalar, 1x1, row-vector, column-vector). Always invoked as self.(other), so self is the left operand and other the right. Returns a fresh 0/1 mask - matrix; the comparison kernels compile branch-free to cmppd + andpd. */ + matrix; the comparison kernels compile branch-free to cmppd + andpd. A + keyword-only ``out=`` target writes the mask into a caller-supplied + same-shape matrix (allocation-free) and returns it; out may alias self. */ #define MATRIX_COMPARE_METHOD(ENUM) \ - static PyObject *Matrix_##ENUM##_compare(PyObject *self, PyObject *other) { \ + static PyObject *Matrix_##ENUM##_compare(PyObject *self, PyObject *args, \ + PyObject *kwds) { \ + PyObject *other = NULL; \ + PyObject *out_target = NULL; \ + static char *kwlist[] = {"", "out", NULL}; \ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|$O", kwlist, &other, \ + &out_target)) { \ + return NULL; \ + } \ PyObject *out = NULL; \ - if (Matrix_binary_op(self, other, &out, NULL, ENUM, false) < 0) { \ + if (Matrix_binary_op(self, other, &out, out_target, ENUM, false) < 0) { \ return NULL; \ } \ return out; \ @@ -4766,10 +6885,12 @@ static int read_key_ranges(PyObject *key, range *rows, range *columns, rows->stop = (Py_ssize_t)impl->rows; rows->step = 1; rows->count = impl->rows; + rows->scalar = true; columns->start = 0; columns->stop = (Py_ssize_t)impl->columns; columns->step = 1; columns->count = impl->columns; + columns->scalar = true; if (PyTuple_Check(key)) { if (PyTuple_GET_SIZE(key) != 2) { PyErr_SetString( @@ -4861,9 +6982,9 @@ PyObject *Matrix_subscript(PyObject *op, PyObject *key) { PyObject *index_list; switch (classify_fancy_key(key, &index_list)) { case FANCY_KEY_AXIS0: - return gather_axis(impl, index_list, 0, Py_TYPE(op)); + return gather_axis(impl, index_list, 0, Py_TYPE(op), NULL); case FANCY_KEY_AXIS1: - return gather_axis(impl, index_list, 1, Py_TYPE(op)); + return gather_axis(impl, index_list, 1, Py_TYPE(op), NULL); case FANCY_KEY_UNSUPPORTED: PyErr_SetString(PyExc_IndexError, "fancy indexing supports only m[[rows]], m[[rows], :], or " @@ -4879,7 +7000,10 @@ PyObject *Matrix_subscript(PyObject *op, PyObject *key) { return NULL; } - if (rows.count == 1 && columns.count == 1) { + // Collapse to a Python float only for an all-integer key that selects a + // single cell. A slice anywhere in the key (even a length-1 one) keeps the + // result a Matrix, so m[0:1, 0:1] stays 1x1 while m[i, j] stays a scalar. + if (rows.scalar && columns.scalar && rows.count == 1 && columns.count == 1) { return PyFloat_FromDouble(impl->row_ptrs[rows.start][columns.start]); } @@ -4890,6 +7014,10 @@ PyObject *Matrix_subscript(PyObject *op, PyObject *key) { } out->impl = impl_new(rows.count, columns.count); + if (out->impl == NULL) { + Py_DECREF(out); + return NULL; + } IMPL_INCREF(out->impl); impl_get(impl, &rows, &columns, out->impl); return (PyObject *)out; @@ -4922,6 +7050,10 @@ static PyObject *Matrix_item(PyObject *op, Py_ssize_t index) { } out->impl = impl_new(rows.count, columns.count); + if (out->impl == NULL) { + Py_DECREF(out); + return NULL; + } IMPL_INCREF(out->impl); impl_get(impl, &rows, &columns, out->impl); return (PyObject *)out; @@ -4929,6 +7061,72 @@ static PyObject *Matrix_item(PyObject *op, Py_ssize_t index) { static PyObject *Matrix_iter(PyObject *op) { return PySeqIter_New(op); } +// Lazy iterator returned by Matrix.values(): a strong ref to the source matrix +// plus a row-major cursor. One float is boxed per step; the backing store is +// already row-major so the cursor is a flat data[0..size) walk. +typedef struct { + PyObject_HEAD MatrixObject *matrix; + size_t cursor; +} MatrixValuesIterObject; + +static void MatrixValuesIter_dealloc(PyObject *op) { + MatrixValuesIterObject *it = (MatrixValuesIterObject *)op; + PyTypeObject *type = Py_TYPE(op); + Py_XDECREF(it->matrix); + type->tp_free(op); + Py_DECREF(type); +} + +static PyObject *MatrixValuesIter_next(PyObject *op) { + MatrixValuesIterObject *it = (MatrixValuesIterObject *)op; + matrix_impl *impl = it->matrix->impl; + + // Re-check ownership every step: the matrix may be released between yields + // (e.g. a values() iterator outliving the @when body that produced it). + if (!impl_check_acquired(impl, true)) { + return NULL; + } + if (it->cursor >= impl->size) { + return NULL; // StopIteration + } + return PyFloat_FromDouble(impl->data[it->cursor++]); +} + +static PyObject *MatrixValuesIter_self(PyObject *op) { + Py_INCREF(op); + return op; +} + +static PyType_Slot MatrixValuesIter_slots[] = { + {Py_tp_dealloc, MatrixValuesIter_dealloc}, + {Py_tp_iter, MatrixValuesIter_self}, + {Py_tp_iternext, MatrixValuesIter_next}, + {0, NULL}, +}; + +static PyType_Spec MatrixValuesIter_Spec = { + .name = "bocpy._math.MatrixValuesIter", + .basicsize = sizeof(MatrixValuesIterObject), + .itemsize = 0, + .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE, + .slots = MatrixValuesIter_slots}; + +static PyObject *Matrix_values(PyObject *op, PyObject *Py_UNUSED(ignored)) { + if (!impl_check_acquired(((MatrixObject *)op)->impl, true)) { + return NULL; + } + PyTypeObject *type = LOCAL_STATE->values_iter_type; + MatrixValuesIterObject *it = + (MatrixValuesIterObject *)type->tp_alloc(type, 0); + if (it == NULL) { + return NULL; + } + Py_INCREF(op); + it->matrix = (MatrixObject *)op; + it->cursor = 0; + return (PyObject *)it; +} + /// @brief Resolve an entire index list into a validated ``size_t`` array. /// @details The scatter (write) counterpart of ``gather_axis``'s per-element /// resolution. Every index is validated (negative-aware, @@ -5198,6 +7396,106 @@ static int scatter_axis(matrix_impl *impl, PyObject *indices, int axis, return 0; } +/// @brief Three-phase scatter-store for ``put_along_axis``. +/// @details The write counterpart to ``gather_along_axis``: ``indices`` runs +/// *along* ``axis``, one per row (``axis == 1``, length ``rows``) or +/// one per column (``axis == 0``, length ``columns``). Resolve all +/// indices, classify/shape-check the RHS, then write — so a rejected +/// scatter leaves the receiver unmodified. The RHS is a scalar, a +/// ``1x1`` matrix, or a vector matching the selection shape +/// (``rows x 1`` for ``axis == 1``, ``1 x columns`` for +/// ``axis == 0``); element ``i`` lands in ``self[i][indices[i]]`` +/// (``axis == 1``) or ``self[indices[i]][i]`` (``axis == 0``). When +/// ``accumulate`` the RHS is added in (else it overwrites). A +/// self-aliased matrix RHS is snapshotted before the write. +/// @pre Callers must ``impl_check_acquired(impl)`` before calling; the RHS is +/// gated internally by ``classify_scatter_rhs``. +/// @param impl The receiver matrix impl (written in place). +/// @param indices The along-axis index sequence. +/// @param axis 0 to index rows per column, 1 to index columns per row. +/// @param value_op The assigned scalar or matrix. +/// @param accumulate When non-zero, add into the selection instead of +/// overwriting. +/// @return 0 on success; -1 with an exception set on failure. +static int scatter_along_axis(matrix_impl *impl, PyObject *indices, int axis, + PyObject *value_op, int accumulate) { + size_t stack[64]; + size_t *idx; + size_t count; + size_t bound = axis == 0 ? impl->rows : impl->columns; + const char *bound_name = axis == 0 ? "row" : "column"; + if (resolve_index_list(indices, bound, bound_name, stack, + sizeof(stack) / sizeof(stack[0]), &idx, &count) < 0) { + return -1; + } + + size_t expected = axis == 0 ? impl->columns : impl->rows; + if (count != expected) { + PyErr_Format(PyExc_ValueError, + "put_along_axis on axis %d expects %zu indices (one per %s), " + "got %zu", + axis, expected, axis == 0 ? "column" : "row", count); + if (idx != stack) { + PyMem_Free(idx); + } + return -1; + } + + size_t want_rows = axis == 0 ? 1 : impl->rows; + size_t want_cols = axis == 0 ? impl->columns : 1; + double scalar; + matrix_impl *value; + if (classify_scatter_rhs(value_op, want_rows, want_cols, &scalar, &value) < + 0) { + if (idx != stack) { + PyMem_Free(idx); + } + return -1; + } + + // Snapshot a self-aliased RHS first: each element is read once and written + // to a (possibly) different cell, so an in-place write could clobber a cell + // still to be read. + matrix_impl *src = value; + matrix_impl *snapshot = NULL; + if (value != NULL && value == impl) { + snapshot = impl_new(value->rows, value->columns); + if (snapshot == NULL) { + IMPL_DECREF(value); + if (idx != stack) { + PyMem_Free(idx); + } + return -1; + } + memcpy(snapshot->data, value->data, value->size * sizeof(double)); + src = snapshot; + } + + for (size_t i = 0; i < count; ++i) { + double v = value != NULL + ? (axis == 0 ? src->row_ptrs[0][i] : src->row_ptrs[i][0]) + : scalar; + double *cell = + axis == 0 ? &impl->row_ptrs[idx[i]][i] : &impl->row_ptrs[i][idx[i]]; + if (accumulate) { + *cell += v; + } else { + *cell = v; + } + } + + if (snapshot != NULL) { + impl_free(snapshot); + } + if (value != NULL) { + IMPL_DECREF(value); + } + if (idx != stack) { + PyMem_Free(idx); + } + return 0; +} + int Matrix_ass_subscript(PyObject *op, PyObject *key, PyObject *value_op) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -5411,6 +7709,483 @@ static PyObject *Matrix_repr(PyObject *op) { return str; } +static PyMethodDef Matrix_methods[] = { + {"row_view", (PyCFunction)Matrix_row_view, METH_VARARGS | METH_KEYWORDS, + "row_view($self, /, row, columns=None, *, read_only=False)\n--\n\n" + "Return a MatrixView of row ``row`` (a 1 x N row), optionally narrowed\n" + "to the column ``slice`` ``columns``. The view shares storage with this\n" + "matrix; call ``.copy()`` for an independent Matrix. With\n" + "``read_only=True`` the view (and any view derived from it via ``.T`` or\n" + "slicing) rejects writes with ``TypeError``."}, + {"column_view", (PyCFunction)Matrix_column_view, + METH_VARARGS | METH_KEYWORDS, + "column_view($self, /, column, rows=None, *, read_only=False)\n--\n\n" + "Return a MatrixView of column ``column`` (an M x 1 column), optionally\n" + "narrowed to the row ``slice`` ``rows``. The view shares storage with\n" + "this matrix; call ``.copy()`` for an independent Matrix. With\n" + "``read_only=True`` the view (and any view derived from it) rejects\n" + "writes with ``TypeError``."}, + {"row_views", (PyCFunction)Matrix_row_views, METH_VARARGS | METH_KEYWORDS, + "row_views($self, /, *, read_only=False)\n--\n\n" + "Return an iterator yielding a MatrixView over each row in turn (a\n" + "1 x N row view per row), each aliasing this matrix's storage. Cleaner\n" + "than ``for i in range(m.rows): m.row_view(i)``; each yielded view is a\n" + "distinct object. Pass ``read_only=True`` for immutable row views."}, + {"column_views", (PyCFunction)Matrix_column_views, + METH_VARARGS | METH_KEYWORDS, + "column_views($self, /, *, read_only=False)\n--\n\n" + "Return an iterator yielding a MatrixView over each column in turn (an\n" + "M x 1 column view per column), each aliasing this matrix's storage.\n" + "Each yielded view is a distinct object. Pass ``read_only=True`` for\n" + "immutable column views."}, + {"transpose", (PyCFunction)Matrix_transpose, METH_VARARGS | METH_KEYWORDS, + "transpose($self, /, in_place=False)\n--\n\n" + "Return a transposed copy, or transpose ``self`` in place when " + "``in_place=True`` (in which case ``self`` is returned)."}, + {"sum", (PyCFunction)Matrix_Sum_method, METH_VARARGS | METH_KEYWORDS, + "sum($self, /, axis=None, where=None)\n--\n\nSum of elements.\n\n" + "``where`` is an optional same-shape Matrix mask: an element is included\n" + "only where its mask cell is non-zero (NaN counts as included). An\n" + "all-excluded group sums to 0."}, + {"mean", (PyCFunction)Matrix_Mean_method, METH_VARARGS | METH_KEYWORDS, + "mean($self, /, axis=None, where=None)\n--\n\nMean of elements.\n\n" + "``where`` is an optional same-shape Matrix mask: the mean is taken over\n" + "only the elements whose mask cell is non-zero (NaN counts as included).\n" + "An all-excluded group yields NaN (matching NumPy's empty-slice mean)."}, + {"magnitude", (PyCFunction)Matrix_Magnitude_method, + METH_VARARGS | METH_KEYWORDS, + "magnitude($self, /, axis=None, where=None)\n--\n\nEuclidean " + "magnitude.\n\n" + "``where`` is an optional same-shape Matrix mask: only elements whose\n" + "mask cell is non-zero contribute (NaN counts as included). An\n" + "all-excluded group yields 0."}, + {"magnitude_squared", (PyCFunction)Matrix_MagnitudeSquared_method, + METH_VARARGS | METH_KEYWORDS, + "magnitude_squared($self, /, axis=None, where=None)\n--\n\n" + "Sum of squared elements (Euclidean magnitude without the sqrt).\n\n" + "``where`` is an optional same-shape Matrix mask: only elements whose\n" + "mask cell is non-zero contribute (NaN counts as included). An\n" + "all-excluded group yields 0."}, + {"vecdot", (PyCFunction)Matrix_vecdot, METH_VARARGS | METH_KEYWORDS, + "vecdot($self, other, /, axis=None)\n--\n\n" + "Axis-aware inner product: sum of element-wise products. " + "Equivalent to numpy.linalg.vecdot for 1-D inputs with axis=None; " + "**not** equivalent to numpy.dot."}, + {"fma", (PyCFunction)Matrix_fma, METH_VARARGS | METH_KEYWORDS, + "fma($self, b, c, /, in_place=False)\n--\n\n" + "Fused multiply-add: single-rounding ``self*b + c`` (libc fma()). " + "b and c may each be a same-shape matrix, a 1x1 matrix, a row or column " + "vector that broadcasts, or a scalar; other shapes raise ValueError. " + "The single rounding differs from ``self*b + c`` (which rounds twice) by " + "up to half a ULP, so compare results with allclose(), not ==."}, + {"scaled_add", (PyCFunction)Matrix_scaled_add, METH_VARARGS | METH_KEYWORDS, + "scaled_add($self, s, x, /, in_place=False)\n--\n\n" + "Scaled add ``self + s * x`` with two roundings; the two-rounding " + "sibling of fma().\n\n" + "s is the scale -- a scalar, a 1x1 matrix, a row or column vector that " + "broadcasts, or a same-shape matrix (the same operand rules as fma's " + "multiplier) -- and x is a same-shape matrix. The arithmetic rounds " + "twice (``round(round(s*x) + self)``), so the result is bit-for-bit " + "identical to ``self + s * x`` and -- unlike fma() -- never fuses to a " + "single rounding. With in_place=True the result is written into self's " + "buffer (allocating nothing) and self is returned; otherwise a new " + "matrix is returned. s and x may alias self. A shape mismatch raises " + "ValueError before any write."}, + {"add", (PyCFunction)Matrix_add_method, METH_VARARGS | METH_KEYWORDS, + "add($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self + other`` with the same broadcasting as ``+`` " + "(other may be a scalar). With ``out`` (a same-shape matrix) the result " + "is written there and returned instead of allocating a new matrix; out " + "may alias an input. Bit-for-bit identical to ``self + other``. A shape " + "mismatch between out and the result raises ValueError before any write."}, + {"subtract", (PyCFunction)Matrix_subtract_method, + METH_VARARGS | METH_KEYWORDS, + "subtract($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self - other`` with the same broadcasting as ``-`` " + "(other may be a scalar). With ``out`` (a same-shape matrix) the result " + "is written there and returned instead of allocating a new matrix; out " + "may alias an input. Bit-for-bit identical to ``self - other``. A shape " + "mismatch between out and the result raises ValueError before any write."}, + {"multiply", (PyCFunction)Matrix_multiply_method, + METH_VARARGS | METH_KEYWORDS, + "multiply($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self * other`` with the same broadcasting as ``*`` " + "(other may be a scalar). With ``out`` (a same-shape matrix) the result " + "is written there and returned instead of allocating a new matrix; out " + "may alias an input. Bit-for-bit identical to ``self * other``. A shape " + "mismatch between out and the result raises ValueError before any write."}, + {"divide", (PyCFunction)Matrix_divide_method, METH_VARARGS | METH_KEYWORDS, + "divide($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self / other`` with the same broadcasting as ``/`` " + "(other may be a scalar). With ``out`` (a same-shape matrix) the result " + "is written there and returned instead of allocating a new matrix; out " + "may alias an input. Bit-for-bit identical to ``self / other``. A shape " + "mismatch between out and the result raises ValueError before any write."}, + {"cross", (PyCFunction)Matrix_cross, METH_VARARGS | METH_KEYWORDS, + "cross($self, other, /, axis=None)\n--\n\n" + "2D (scalar z-component) or 3D cross product against another " + "vector or batch. 1x2 / 2x1 inputs return a float; 1x3 / 3x1 return " + "a Matrix preserving self's orientation. Nx2 / 2xN row/column " + "batches return per-vector scalars (Mx1 / 1xN); Nx3 / 3xN return " + "same-shape batches. Batch operands accept either a same-shape " + "other or a single 2D/3D vector (1xK / Kx1) broadcast against " + "every per-vector slot \u2014 ``self`` must be the batch (cross is " + "anticommutative). ``axis`` disambiguates the 2x2 / 3x3 squares " + "(default rows, ``axis=0`` for columns)."}, + {"normalize", (PyCFunction)Matrix_normalize, METH_VARARGS | METH_KEYWORDS, + "normalize($self, /, axis=None, in_place=False)\n--\n\n" + "Divide elements by their magnitude. ``axis=None`` divides by the " + "matrix's total magnitude; ``axis=0`` divides each column by its own " + "magnitude; ``axis=1`` divides each row by its own magnitude. Rows or " + "columns whose magnitude is zero are left as the all-zero vector. " + "Sub-normal magnitudes may overflow during division; threshold with " + "magnitude_squared() if safety matters. When ``in_place=True``, mutates " + "``self`` and returns it."}, + {"perpendicular", (PyCFunction)Matrix_perpendicular, + METH_VARARGS | METH_KEYWORDS, + "perpendicular($self, /, axis=None, in_place=False)\n--\n\n" + "Rotate every 2D vector 90 degrees counter-clockwise: ``(x, y) -> " + "(-y, x)``. Accepts a single 2D vector (``1x2`` or ``2x1``), a row " + "batch (``Nx2``), or a column batch (``2xN``). On the ambiguous " + "``2x2`` shape the default is per-row; pass ``axis=0`` to force " + "per-column. When ``in_place=True``, mutates ``self`` and returns it."}, + {"angle", (PyCFunction)Matrix_angle, METH_VARARGS | METH_KEYWORDS, + "angle($self, /, axis=None)\n--\n\n" + "Polar angle (``atan2(y, x)``) of every 2D vector. Returns a float " + "for a single 2D vector, an ``Mx1`` column matrix for an ``Nx2`` row " + "batch, or a ``1xN`` row matrix for a ``2xN`` column batch. On the " + "ambiguous ``2x2`` shape the default is per-row; pass ``axis=0`` to " + "force per-column."}, + {"min", (PyCFunction)Matrix_Minimum_method, METH_VARARGS | METH_KEYWORDS, + "min($self, /, axis=None, where=None)\n--\n\nMinimum of elements.\n\n" + "``where`` is an optional same-shape Matrix mask: only elements whose\n" + "mask cell is non-zero are considered (NaN counts as included). An\n" + "all-excluded group yields NaN."}, + {"max", (PyCFunction)Matrix_Maximum_method, METH_VARARGS | METH_KEYWORDS, + "max($self, /, axis=None, where=None)\n--\n\nMaximum of elements.\n\n" + "``where`` is an optional same-shape Matrix mask: only elements whose\n" + "mask cell is non-zero are considered (NaN counts as included). An\n" + "all-excluded group yields NaN."}, + {"argmin", (PyCFunction)Matrix_argmin_method, METH_VARARGS | METH_KEYWORDS, + "argmin($self, /, axis=None, where=None, as_matrix=False)\n--\n\n" + "Index of the minimum element (first occurrence on ties).\n\n" + "With axis=None returns a single int. With axis=0/1 returns a list of\n" + "ints (directly usable in fancy indexing); pass as_matrix=True to get a\n" + "Matrix vector of index positions instead. NaN elements are skipped\n" + "unless the first included element along the reduced axis is NaN, which\n" + "pins the result to that position. This differs from NumPy, which\n" + "propagates NaN.\n\n" + "``where`` is an optional same-shape Matrix mask: a cell == 0.0\n" + "excludes that element (NaN counts as included). A group with no\n" + "included element yields -1, the integer analog of a masked min/max\n" + "returning NaN."}, + {"argmax", (PyCFunction)Matrix_argmax_method, METH_VARARGS | METH_KEYWORDS, + "argmax($self, /, axis=None, where=None, as_matrix=False)\n--\n\n" + "Index of the maximum element (first occurrence on ties).\n\n" + "With axis=None returns a single int. With axis=0/1 returns a list of\n" + "ints (directly usable in fancy indexing); pass as_matrix=True to get a\n" + "Matrix vector of index positions instead. NaN elements are skipped\n" + "unless the first included element along the reduced axis is NaN, which\n" + "pins the result to that position. This differs from NumPy, which\n" + "propagates NaN.\n\n" + "``where`` is an optional same-shape Matrix mask: a cell == 0.0\n" + "excludes that element (NaN counts as included). A group with no\n" + "included element yields -1, the integer analog of a masked min/max\n" + "returning NaN."}, + {"bin", (PyCFunction)Matrix_bin, METH_VARARGS | METH_KEYWORDS, + "bin($self, /, bins, *, bounds=None, right=False, " + "in_place=False, out=None)\n--\n\n" + "Map each element to its equal-width bin index over a range.\n\n" + "``bins`` (>= 1) equal-width intervals span the range. ``bounds`` is an\n" + "optional (min, max) tuple: when given it fixes the range and skips the\n" + "O(size) min/max scan (a value outside it clamps into the first or last\n" + "bin); when omitted the range is the matrix's own min and max, both\n" + "computed ignoring NaN. The index is clamped to [0, bins-1], so a value\n" + "equal to max lands in the last bin; ``right=True`` puts a value on an\n" + "interior boundary in the lower bin. A NaN element maps to -1 (the\n" + "invalid-index sentinel), and a degenerate range (min == max) sends\n" + "every element to bin 0. ``in_place`` bins self and returns it; the\n" + "keyword-only ``out`` writes into a same-shape Matrix; the two are\n" + "mutually exclusive. Raises ValueError if bounds is not a length-2\n" + "(min, max) pair of finite numbers with max >= min."}, + {"digitize", (PyCFunction)Matrix_digitize, METH_VARARGS | METH_KEYWORDS, + "digitize($self, /, edges, *, right=False, in_place=False, out=None)\n" + "--\n\n" + "Map each element to the index of the edge interval it falls in.\n\n" + "``edges`` is a strictly increasing vector of interior boundaries (a\n" + "Matrix vector or any coercible sequence). With ``right=False`` the\n" + "index is the count of edges <= value; with ``right=True`` the count of\n" + "edges < value, matching numpy.digitize. Indices span [0, len(edges)].\n" + "A NaN element maps to -1 (the invalid-index sentinel). ``in_place`` /\n" + "keyword-only ``out`` are mutually exclusive."}, + {"ceil", (PyCFunction)Matrix_Ceil_method, METH_VARARGS | METH_KEYWORDS, + "ceil($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise ceiling. With ``out`` (a same-shape matrix) the result " + "is written there and returned; ``out`` and ``in_place`` are mutually " + "exclusive."}, + {"floor", (PyCFunction)Matrix_Floor_method, METH_VARARGS | METH_KEYWORDS, + "floor($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise floor. With ``out`` (a same-shape matrix) the result " + "is written there and returned; ``out`` and ``in_place`` are mutually " + "exclusive."}, + {"round", (PyCFunction)Matrix_Round_method, METH_VARARGS | METH_KEYWORDS, + "round($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise rounding (banker's; IEEE round-half-to-even). With ``out`` " + "(a same-shape matrix) the result is written there and returned; " + "``out`` and ``in_place`` are mutually exclusive."}, + {"negate", (PyCFunction)Matrix_Negate_method, METH_VARARGS | METH_KEYWORDS, + "negate($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise negation. With ``out`` (a same-shape matrix) the result " + "is written there and returned; ``out`` and ``in_place`` are mutually " + "exclusive."}, + {"abs", (PyCFunction)Matrix_Abs_method, METH_VARARGS | METH_KEYWORDS, + "abs($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise absolute value. With ``out`` (a same-shape matrix) the " + "result is written there and returned; ``out`` and ``in_place`` are " + "mutually exclusive."}, + {"sqrt", (PyCFunction)Matrix_Sqrt_method, METH_VARARGS | METH_KEYWORDS, + "sqrt($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise square root. Negative elements yield NaN. With ``out`` " + "(a same-shape matrix) the result is written there and returned; " + "``out`` and ``in_place`` are mutually exclusive."}, + {"reciprocal", (PyCFunction)Matrix_Reciprocal_method, + METH_VARARGS | METH_KEYWORDS, + "reciprocal($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise reciprocal (``1 / x``). Zero elements yield +/-inf. With " + "``out`` (a same-shape matrix) the result is written there and " + "returned; ``out`` and ``in_place`` are mutually exclusive."}, + {"sign", (PyCFunction)Matrix_Sign_method, METH_VARARGS | METH_KEYWORDS, + "sign($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise sign: -1, 0, or 1 (NaN maps to 0). With ``out`` " + "(a same-shape matrix) the result is written there and returned; " + "``out`` and ``in_place`` are mutually exclusive."}, + {"cos", (PyCFunction)Matrix_Cos_method, METH_VARARGS | METH_KEYWORDS, + "cos($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise cosine (radians). With ``out`` (a same-shape matrix) the " + "result is written there and returned; ``out`` and ``in_place`` are " + "mutually exclusive."}, + {"sin", (PyCFunction)Matrix_Sin_method, METH_VARARGS | METH_KEYWORDS, + "sin($self, /, *, in_place=False, out=None)\n--\n\n" + "Element-wise sine (radians). With ``out`` (a same-shape matrix) the " + "result is written there and returned; ``out`` and ``in_place`` are " + "mutually exclusive."}, + {"less", (PyCFunction)Matrix_Less_compare, METH_VARARGS | METH_KEYWORDS, + "less($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self < other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``<`` operator, which " + "returns a single bool."}, + {"less_equal", (PyCFunction)Matrix_LessEqual_compare, + METH_VARARGS | METH_KEYWORDS, + "less_equal($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self <= other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``<=`` operator, which " + "returns a single bool."}, + {"greater", (PyCFunction)Matrix_Greater_compare, + METH_VARARGS | METH_KEYWORDS, + "greater($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self > other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``>`` operator, which " + "returns a single bool."}, + {"greater_equal", (PyCFunction)Matrix_GreaterEqual_compare, + METH_VARARGS | METH_KEYWORDS, + "greater_equal($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self >= other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``>=`` operator, which " + "returns a single bool."}, + {"equal", (PyCFunction)Matrix_Equal_compare, METH_VARARGS | METH_KEYWORDS, + "equal($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self == other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``==`` operator, which " + "returns a single bool."}, + {"not_equal", (PyCFunction)Matrix_NotEqual_compare, + METH_VARARGS | METH_KEYWORDS, + "not_equal($self, other, /, *, out=None)\n--\n\n" + "Element-wise ``self != other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 1. With ``out`` (a same-shape matrix) the mask is " + "written there and returned. Distinct from the ``!=`` operator, which " + "returns a single bool."}, + {"clip", (PyCFunction)Matrix_clip, METH_VARARGS | METH_KEYWORDS, + "clip($self, min=None, max=None, *, in_place=False, out=None)\n--\n\n" + "Clamp elements to [min, max]; either bound may be omitted.\n\n" + "The first argument is the lower bound and the second the upper\n" + "bound. Pass None (or omit) to leave a side unbounded, so\n" + "clip(min=0) clamps only below and clip(max=255) only above. Raises\n" + "ValueError if both bounds are omitted.\n\n" + "With in_place=True the matrix is clamped in place and returned; with\n" + "out (a pre-allocated same-shape matrix) the result is written there\n" + "and returned instead of allocating a new matrix. in_place and out\n" + "are mutually exclusive."}, + {"copy", Matrix_copy, METH_NOARGS, + "copy($self, /)\n--\n\nReturn a deep copy."}, + {"from_view", Matrix_from_view, METH_O | METH_CLASS, + "from_view($type, view, /)\n--\n\n" + "Materialise a MatrixView into a new, independent Matrix (the " + "classmethod twin of MatrixView.copy())."}, + {"values", Matrix_values, METH_NOARGS, + "values($self, /)\n--\n\n" + "Yield every element as a float in row-major (row/column) order. " + "Lazy: one float is boxed per step, so streaming a large matrix never " + "materialises a list."}, + {"__reduce__", Matrix_reduce, METH_NOARGS, + "__reduce__($self, /)\n--\n\n" + "Pickle helper: serialize the matrix to its raw double buffer."}, + {"take", (PyCFunction)Matrix_take, METH_VARARGS | METH_KEYWORDS, + "take($self, indices, axis=0, *, out=None)\n--\n\n" + "Take rows or columns by index into a new matrix.\n\n" + "indices is a 1-D list or tuple of ints selecting whole rows\n" + "(axis=0) or columns (axis=1). Negative indices count from the\n" + "end; duplicates repeat the row or column; an out-of-range index\n" + "raises IndexError, and an empty index sequence raises IndexError.\n" + "With out (a pre-allocated matrix of the selection shape -- " + "len(indices) x columns for axis=0, rows x len(indices) for axis=1) " + "the result is written there and returned instead of allocating a new " + "matrix. All indices are validated before any write, so a rejected " + "call leaves out untouched; out must not alias self."}, + {"put", (PyCFunction)Matrix_put, METH_VARARGS | METH_KEYWORDS, + "put($self, indices, value, axis=0, accumulate=False)\n--\n\n" + "Assign value into the selected rows or columns in place.\n\n" + "The write-side counterpart of take(): value may be a scalar, a\n" + "1x1 matrix, or a matrix matching the selection shape. All indices\n" + "and the value shape are validated before any write, so a rejected\n" + "call leaves the matrix unchanged. Negative indices count from the\n" + "end; an out-of-range or empty index raises IndexError. With\n" + "accumulate=False (the default) duplicate indices follow\n" + "last-write-wins; with accumulate=True the values are added into the\n" + "selection, so duplicate indices fold additively (scatter-add).\n" + "Returns self."}, + {"take_along_axis", (PyCFunction)Matrix_take_along_axis, + METH_VARARGS | METH_KEYWORDS, + "take_along_axis($self, indices, axis=0, *, out=None)\n--\n\n" + "Gather one element per row or column along an axis.\n\n" + "The np.take_along_axis counterpart of take() (which selects whole\n" + "rows or columns). With axis=1 the indices give one column index per\n" + "row (so len(indices) must equal the row count) and the result is an\n" + "rows x 1 column vector with out[r] = self[r][indices[r]]. With\n" + "axis=0 they give one row index per column (len(indices) must equal\n" + "the column count) and the result is a 1 x columns row vector with\n" + "out[c] = self[indices[c]][c]. This pairs directly with argmin/argmax\n" + "along the same axis: feed their returned index list straight back in\n" + "to gather the reduced values. Negative indices count from the end;\n" + "an out-of-range index raises IndexError, and a wrong index count\n" + "raises ValueError. With out (a pre-allocated matrix of the result\n" + "shape -- 1 x columns for axis=0, rows x 1 for axis=1) the result is\n" + "written there and returned instead of allocating a new matrix. All\n" + "indices are validated before any write, so a rejected call leaves\n" + "out untouched; out must not alias self."}, + {"put_along_axis", (PyCFunction)Matrix_put_along_axis, + METH_VARARGS | METH_KEYWORDS, + "put_along_axis($self, indices, value, axis=0, accumulate=False)\n--\n\n" + "Assign one element per row or column along an axis in place.\n\n" + "The write-side counterpart of take_along_axis(). With axis=1 element\n" + "i lands in self[i][indices[i]] (len(indices) equals the row count);\n" + "with axis=0 it lands in self[indices[i]][i] (len(indices) equals the\n" + "column count). value may be a scalar, a 1x1 matrix, or a vector\n" + "matching the selection shape (rows x 1 for axis=1, 1 x columns for\n" + "axis=0). All indices and the value shape are validated before any\n" + "write, so a rejected call leaves the matrix unchanged. Negative\n" + "indices count from the end; an out-of-range index raises IndexError,\n" + "and a wrong index count raises ValueError. With accumulate=False (the\n" + "default) duplicate indices follow last-write-wins; with\n" + "accumulate=True the values are added in (scatter-add). Returns self."}, + {"repeat_interleave", (PyCFunction)Matrix_repeat_interleave, + METH_VARARGS | METH_KEYWORDS, + "repeat_interleave($self, repeats, axis=None)\n--\n\n" + "Repeat each element, row, or column consecutively into a new matrix.\n\n" + "Interleaved (np.repeat / torch.repeat_interleave), not tiled: [a, b]\n" + "with repeats=2 gives [a, a, b, b]. With axis=0 each row is repeated\n" + "into a (rows*repeats) x columns result; with axis=1 each column is\n" + "repeated into rows x (columns*repeats); with axis=None (the default)\n" + "the row-major buffer is flattened to a 1 x (size*repeats) row vector.\n" + "repeats must be a positive integer. Returns a new matrix."}, + {"topk", (PyCFunction)Matrix_topk, METH_VARARGS | METH_KEYWORDS, + "topk($self, k, axis=None, largest=True, where=None, " + "as_matrix=False)\n--\n\n" + "The k extreme elements per reduction group, in sorted order.\n\n" + "Returns a (values, indices) tuple. ``largest=True`` (the default)\n" + "selects the k greatest in descending order; ``largest=False`` selects\n" + "the k smallest in ascending order. Ties keep the first occurrence\n" + "(NumPy tie-break) and any NaN sorts last.\n\n" + "With axis=None the whole row-major buffer is one group and values is\n" + "a 1 x k row vector. With axis=0 each column is reduced down the rows\n" + "and values is k x cols. With axis=1 each row is reduced across the\n" + "columns and values is rows x k.\n\n" + "By default (as_matrix=False) indices is Python ints: a flat list[int]\n" + "for axis=None, or a list of `cols`/`rows` lists (each k indices) for\n" + "axis=0/axis=1. With as_matrix=True indices is instead a Matrix of the\n" + "same shape as values, each cell the source index (as a float) of the\n" + "value beside it; the -1 pad becomes -1.0.\n\n" + "k must be a positive integer no larger than the reduced axis length\n" + "(every cell counts, masked or not), else ValueError. ``where`` is an\n" + "optional same-shape Matrix mask: a cell == 0.0 excludes that element\n" + "(NaN counts as included). A group with fewer than k included elements\n" + "fills its leading slots and pads the rest with NaN values and -1\n" + "indices."}, + {"allclose", (PyCFunction)Matrix_allclose, + METH_VARARGS | METH_KEYWORDS | METH_CLASS, + "allclose($type, lhs, rhs, /, rtol=1e-05, atol=1e-08, " + "equal_nan=False)\n--\n\n" + "Check element-wise equality within tolerance."}, + {"where", (PyCFunction)Matrix_where, + METH_VARARGS | METH_KEYWORDS | METH_CLASS, + "where($type, mask, a, b, /, *, out=None)\n--\n\n" + "Select element-wise from a or b on a truthy mask.\n\n" + "Returns a fresh matrix taking a where the corresponding mask\n" + "element is non-zero and b elsewhere. a and b may each be a scalar,\n" + "a matrix matching the mask's shape, or a list/tuple of numbers\n" + "(taken as a 1xN row vector that must then match the mask shape);\n" + "other shapes raise ValueError. A 1x1 matrix is treated as a matrix\n" + "(not a scalar) and so must match the mask shape. NaN mask elements\n" + "are non-zero and select a. With out (a same-shape matrix) the result\n" + "is written there and returned; out may alias mask, a, or b."}, + {"zeros", Matrix_zeros, METH_VARARGS | METH_CLASS, + "zeros($type, size, /)\n--\n\nCreate a zero-filled matrix."}, + {"ones", Matrix_ones, METH_VARARGS | METH_CLASS, + "ones($type, size, /)\n--\n\nCreate a matrix of ones."}, + {"full", Matrix_full, METH_VARARGS | METH_CLASS, + "full($type, size, value, /)\n--\n\n" + "Create a matrix filled with a constant value."}, + {"normal", (PyCFunction)Matrix_normal, + METH_VARARGS | METH_KEYWORDS | METH_CLASS, + "normal($type, mean=0.0, stddev=1.0, /, size=None)\n--\n\n" + "Sample from a normal distribution."}, + {"uniform", (PyCFunction)Matrix_uniform, + METH_VARARGS | METH_KEYWORDS | METH_CLASS, + "uniform($type, minval=0.0, maxval=1.0, /, size=None)\n--\n\n" + "Sample from a uniform distribution."}, + {"seed", Matrix_seed, METH_VARARGS | METH_CLASS, + "seed($type, value, /)\n--\n\n" + "Seed the random generator used by normal() and uniform().\n\n" + "Each interpreter owns an independent splitmix64 stream, so a seed\n" + "makes that interpreter's subsequent draws reproducible; parallel\n" + "workers seed their own streams independently. The sequence is not\n" + "portable across platforms."}, + {"vector", Matrix_vector, METH_VARARGS | METH_CLASS, + "vector($type, values, /, as_column=False)\n--\n\n" + "Create a vector from a sequence."}, + {"concat", (PyCFunction)Matrix_concat, + METH_VARARGS | METH_KEYWORDS | METH_CLASS, + "concat($type, values, /, axis=0)\n--\n\n" + "Concatenate matrices along an axis."}, + {NULL} /* Sentinel */ +}; + static PyType_Slot Matrix_slots[] = { {Py_tp_doc, "Matrix(rows, columns, values=None)\n--\n\n" "A dense 2-D matrix of double-precision floats."}, @@ -5458,6 +8233,11 @@ static PyType_Spec Matrix_Spec = {.name = "bocpy._math.Matrix", static PyObject *_new_matrix_object(XIDATA_T *xidata) { matrix_impl *impl = (matrix_impl *)xidata->data; + _math_module_state *state = math_local_state(); + if (state == NULL) { + return NULL; + } + int_least64_t expected = BOCPY_NO_OWNER; int_least64_t desired = bocpy_interpid(); if (!atomic_compare_exchange_strong(&impl->owner, &expected, desired)) { @@ -5468,7 +8248,7 @@ static PyObject *_new_matrix_object(XIDATA_T *xidata) { return NULL; } - PyTypeObject *type = LOCAL_STATE->matrix_type; + PyTypeObject *type = state->matrix_type; MatrixObject *matrix = (MatrixObject *)type->tp_alloc(type, 0); if (matrix == NULL) { int_least64_t rollback_expected = desired; @@ -5578,6 +8358,45 @@ static int _math_module_exec(PyObject *module) { return -1; } + state->values_iter_type = (PyTypeObject *)PyType_FromModuleAndSpec( + module, &MatrixValuesIter_Spec, NULL); + if (state->values_iter_type == NULL) { + return -1; + } + + state->view_type = + (PyTypeObject *)PyType_FromModuleAndSpec(module, &MatrixView_Spec, NULL); + if (state->view_type == NULL) { + return -1; + } + + if (PyModule_AddType(module, state->view_type) < 0) { + return -1; + } + // MatrixView is intentionally NOT passed to XIDATA_REGISTERCLASS: a view is a + // local borrow of its base matrix and cannot cross interpreters (a view is + // not shippable). Callers .copy() to obtain an owned Matrix to ship. + + // The two companion types are private (NOT PyModule_AddType'd), like + // values_iter_type: the values() iterator and the m.view[...] accessor. + state->view_values_iter_type = (PyTypeObject *)PyType_FromModuleAndSpec( + module, &MatrixViewValuesIter_Spec, NULL); + if (state->view_values_iter_type == NULL) { + return -1; + } + + state->view_indexer_type = (PyTypeObject *)PyType_FromModuleAndSpec( + module, &MatrixViewIndexer_Spec, NULL); + if (state->view_indexer_type == NULL) { + return -1; + } + + state->views_iter_type = (PyTypeObject *)PyType_FromModuleAndSpec( + module, &MatrixViewsIter_Spec, NULL); + if (state->views_iter_type == NULL) { + return -1; + } + if (XIDATA_REGISTERCLASS(state->matrix_type, _matrix_shared)) { Py_FatalError( "could not register MatrixObject for cross-interpreter sharing"); @@ -5589,6 +8408,14 @@ static int _math_module_exec(PyObject *module) { return -1; } + // Seed this interpreter's PRNG with a value distinct per interpreter (and + // per run): mixing the interpreter id and the state pointer ensures parallel + // workers created within the same clock tick still draw independent streams. + // Matrix.seed() overrides this for reproducibility within an interpreter. + state->prng_state = ((uint64_t)time(NULL) << 20) ^ + ((uint64_t)bocpy_interpid() * 0x9E3779B97F4A7C15ULL) ^ + (uint64_t)(uintptr_t)state; + assert(LOCAL_STATE == NULL); LOCAL_STATE = state; @@ -5598,6 +8425,11 @@ static int _math_module_exec(PyObject *module) { static int _math_module_clear(PyObject *module) { _math_module_state *state = (_math_module_state *)PyModule_GetState(module); Py_CLEAR(state->matrix_type); + Py_CLEAR(state->values_iter_type); + Py_CLEAR(state->view_type); + Py_CLEAR(state->view_values_iter_type); + Py_CLEAR(state->view_indexer_type); + Py_CLEAR(state->views_iter_type); Py_CLEAR(state->matrix_unpickle); return 0; } @@ -5609,6 +8441,11 @@ static void _math_module_free(void *module) { static int _math_module_traverse(PyObject *module, visitproc visit, void *arg) { _math_module_state *state = (_math_module_state *)PyModule_GetState(module); Py_VISIT(state->matrix_type); + Py_VISIT(state->values_iter_type); + Py_VISIT(state->view_type); + Py_VISIT(state->view_values_iter_type); + Py_VISIT(state->view_indexer_type); + Py_VISIT(state->views_iter_type); Py_VISIT(state->matrix_unpickle); return 0; } diff --git a/src/bocpy/behaviors.py b/src/bocpy/behaviors.py index 0f70bd9..5991e26 100644 --- a/src/bocpy/behaviors.py +++ b/src/bocpy/behaviors.py @@ -1870,7 +1870,9 @@ def whencall(func, args: list[Union[Cown, list[Cown]]], captures: list[Any]) -> group_id += 1 - behavior = _core.BehaviorCapsule(key, result.impl, cowns, captures) + # Empty groups leave no cown entry, so pass the top-level arg count for tuple sizing and [] fill. + behavior = _core.BehaviorCapsule(key, result.impl, cowns, captures, + len(args)) # Take the terminator hold before scheduling so a concurrent stop()/terminator_close() refuses the schedule # rather than racing teardown; the matching dec runs on the worker thread once the body completes. if _core.terminator_inc() < 0: diff --git a/templates/c_abi_consumer/pyproject.toml b/templates/c_abi_consumer/pyproject.toml index 6c7657d..4c4098c 100644 --- a/templates/c_abi_consumer/pyproject.toml +++ b/templates/c_abi_consumer/pyproject.toml @@ -8,7 +8,7 @@ # the public C ABI it was authored against; bump it in lock-step with # ``[project].version`` in the root ``pyproject.toml`` (see the # ``finalize-pr`` skill). -requires = ["setuptools", "wheel", "bocpy~=0.13"] +requires = ["setuptools", "wheel", "bocpy~=0.14"] build-backend = "setuptools.build_meta" [project] @@ -16,4 +16,4 @@ name = "bocpy-c-abi-consumer" version = "0.0.0" description = "Smoke test and canonical downstream template for the bocpy public C ABI." requires-python = ">=3.10" -dependencies = ["bocpy~=0.13"] +dependencies = ["bocpy~=0.14"] diff --git a/test/test_boc.py b/test/test_boc.py index 1c091f2..effe73a 100644 --- a/test/test_boc.py +++ b/test/test_boc.py @@ -10,7 +10,7 @@ import pytest -from bocpy import (Cown, drain, notice_write, +from bocpy import (Cown, drain, notice_write, PinnedCown, quiesce, receive, send, start, TIMEOUT, wait, when) from bocpy._core import CownCapsule from bocpy_test import collide_a, collide_b, dunders @@ -172,6 +172,137 @@ def single_group_single(single0: Cown[int], group: list[Cown[int]], single1: Cow return expected, [group, group_single, single_group, group_single_group, single_group_single] +class TestEmptyGroups: + """An empty cown-list @when argument passes an empty list to its parameter.""" + + @classmethod + def teardown_class(cls): + """Ensure runtime is drained after suite.""" + wait() + + def test_trailing_empty_group(self): + """An empty list as the last @when arg yields [] for that parameter.""" + a = Cown(1) + + @when(a, []) + def result(a, group): + return (a.value, list(group)) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (1, []) + + def test_leading_empty_group(self): + """An empty list as the first @when arg keeps later slots correct.""" + b = Cown(2) + + @when([], b) + def result(group, b): + return (list(group), b.value) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == ([], 2) + + def test_middle_empty_group(self): + """An empty list between two singletons yields [] in the middle slot.""" + a = Cown(1) + b = Cown(2) + + @when(a, [], b) + def result(a, group, b): + return (a.value, list(group), b.value) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (1, [], 2) + + def test_all_empty_groups_zero_cowns(self): + """A behavior whose only args are empty groups still fires.""" + + @when([], []) + def result(g0, g1): + return (list(g0), list(g1)) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == ([], []) + + def test_mixed_empty_and_nonempty_groups(self): + """Empty and non-empty groups reconstruct at their correct slots.""" + a = Cown(1) + grp = [Cown(2), Cown(3)] + c = Cown(4) + + @when(a, [], grp, c) + def result(a, empty, group, c): + return (a.value, list(empty), + sorted(x.value for x in group), c.value) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (1, [], [2, 3], 4) + + def test_empty_group_with_capture(self): + """An empty group coexists with a trailing capture parameter.""" + a = Cown(1) + factor = 10 + + @when(a, []) + def result(a, group, factor=factor): + return (a.value * factor, list(group)) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (10, []) + + def test_multi_member_group_with_capture(self): + """A multi-member group plus an empty group and a trailing capture.""" + a = Cown(1) + grp = [Cown(2), Cown(3)] + factor = 10 + + @when(a, [], grp) + def result(a, empty, group, factor=factor): + return (a.value * factor, list(empty), + sorted(x.value for x in group)) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (10, [], [2, 3]) + + def test_tuple_form_groups(self): + """Groups passed as tuples reconstruct identically to lists.""" + a = Cown(1) + b = Cown(2) + + @when(a, (), (b,)) + def result(a, empty, group): + return (a.value, list(empty), [x.value for x in group]) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (1, [], [2]) + + def test_empty_group_result_feeds_chained_behavior(self): + """A behavior with an empty group feeds its result downstream.""" + a = Cown(1) + + @when(a, []) + def stage1(a, group): + return a.value + len(group) + + @when(stage1) + def stage2(s): + return s.value * 100 + + quiesce(QUIESCE_TIMEOUT) + assert stage2.unwrap() == 100 + + def test_pinned_empty_group_pump_path(self): + """An empty group on a pinned-cown behavior (pump path) yields [].""" + pc = PinnedCown({"n": 5}) + + @when(pc, []) + def result(pc, group): + return (pc.value["n"], list(group)) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (5, []) + + class TestBOC: """Integration-style tests for bocpy behaviors.""" @@ -1124,6 +1255,7 @@ def test_zero_args_behavior_capsule(self): result.impl, [], [], + 0, ) assert capsule is not None @@ -1135,16 +1267,41 @@ def test_large_args_behavior_capsule(self): result = Cown(None) cowns = [Cown(i) for i in range(32)] - args = [(i, c.impl) for i, c in enumerate(cowns)] + # group_ids are 1-based (as whencall emits them), one slot per cown. + args = [(i + 1, c.impl) for i, c in enumerate(cowns)] capsule = BehaviorCapsule( "__behavior_large_args__", result.impl, args, [], + len(args), ) assert capsule is not None + def test_group_id_out_of_range_rejected(self): + """A group_id outside [1, num_arg_slots] is rejected at construction. + + Guards the raw (public) BehaviorCapsule constructor against an + out-of-bounds argument-tuple write at execute time: the reconstruction + places each cown at slot abs(group_id) - 1 in a tuple of num_arg_slots. + """ + from bocpy import start as _start_runtime + from bocpy._core import BehaviorCapsule + _start_runtime() + + result = Cown(None) + c = Cown(1) + # group_id 5 but only one slot -> would index a 1-tuple at slot 4. + with pytest.raises(ValueError): + BehaviorCapsule("__oob__", result.impl, [(5, c.impl)], [], 1) + # group_id 0 -> slot -1. + with pytest.raises(ValueError): + BehaviorCapsule("__zero__", result.impl, [(0, c.impl)], [], 1) + # A negative num_arg_slots is rejected outright. + with pytest.raises(ValueError): + BehaviorCapsule("__neg__", result.impl, [(1, c.impl)], [], -1) + class TestExceptionFlag: """Tests for the Cown.exception flag distinguishing thrown vs returned.""" diff --git a/test/test_matrix.py b/test/test_matrix.py index 9cc048d..a25bb72 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -11,7 +11,7 @@ import pytest -from bocpy import Cown, Matrix, quiesce, wait, when +from bocpy import Cown, Matrix, quiesce, send, wait, when QUIESCE_TIMEOUT = 5 @@ -31,9 +31,9 @@ def ref_fma(a, b, c): return float(Fraction(a) * Fraction(b) + Fraction(c)) -def _flatten(m): +def flatten(m): """Row-major list of every element in a Matrix.""" - return [m[i, j] for i in range(m.rows) for j in range(m.columns)] + return list(m.values()) MATRIX_SIZES = [ @@ -152,6 +152,21 @@ def test_ones(self, shape): assert m.columns == cols assert m.sum() == pytest.approx(rows * cols) + def test_full(self, shape): + """Verify Matrix.full() creates a matrix filled with a constant.""" + rows, cols = shape + m = Matrix.full(shape, 3.5) + assert m.rows == rows + assert m.columns == cols + assert m.sum() == pytest.approx(rows * cols * 3.5) + assert all(m[r, c] == 3.5 for r in range(rows) for c in range(cols)) + + def test_full_negative(self, shape): + """Verify Matrix.full() handles negative fill values.""" + rows, cols = shape + m = Matrix.full(shape, -2.0) + assert m.sum() == pytest.approx(rows * cols * -2.0) + def test_normal_shape(self, shape): """Verify Matrix.normal() produces a matrix of the given shape.""" rows, cols = shape @@ -201,6 +216,47 @@ def test_seed_requires_argument(self): Matrix.seed() +class TestFactoriesOnWorker: + """RNG factories driven from inside @when behaviors (worker interpreters). + + The per-interpreter PRNG state must be reachable from worker threads, + which never ran the _math module init on their own thread. + """ + + @classmethod + def teardown_class(cls): + """Tear the runtime down after all behavior tests.""" + wait() + + def test_rng_runs_on_worker_interpreter(self): + """normal()/uniform() draw from the worker PRNG in a behavior.""" + c = Cown(0) + + @when(c) + def result(c): # noqa: D401 — short behavior body + u = Matrix.uniform(0.0, 1.0, size=(3, 2)) + n = Matrix.normal(0.0, 1.0, size=(2, 4)) + return (u.rows, u.columns, n.rows, n.columns) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() == (3, 2, 2, 4) + + def test_seed_reproducible_on_worker(self): + """Seeding in a behavior reproduces the same draw on that interpreter.""" + c = Cown(0) + + @when(c) + def result(c): # noqa: D401 — short behavior body + Matrix.seed(4242) + a = Matrix.uniform(0.0, 1.0, size=(4, 4)) + Matrix.seed(4242) + b = Matrix.uniform(0.0, 1.0, size=(4, 4)) + return Matrix.allclose(a, b) + + quiesce(QUIESCE_TIMEOUT) + assert result.unwrap() is True + + class TestIndexing: """Tests for element and row indexing.""" @@ -1401,7 +1457,7 @@ def test_parity_with_operator(self, name, lhs_shape, rhs_shape, rng): expected = op(lhs, rhs) actual = getattr(lhs, name)(rhs) assert actual is not lhs - assert _flatten(actual) == _flatten(expected) + assert flatten(actual) == flatten(expected) def test_parity_with_scalar_operand(self, name, rng): """A Python scalar operand matches the operator form.""" @@ -1409,7 +1465,7 @@ def test_parity_with_scalar_operand(self, name, rng): lhs = Matrix(2, 3, [rng.uniform(1.0, 10.0) for _ in range(6)]) expected = op(lhs, 2.5) actual = getattr(lhs, name)(2.5) - assert _flatten(actual) == _flatten(expected) + assert flatten(actual) == flatten(expected) @pytest.mark.parametrize("lhs_shape,rhs_shape", _BROADCAST_PAIRS) def test_out_writes_in_place_and_returns_it(self, name, lhs_shape, @@ -1422,28 +1478,28 @@ def test_out_writes_in_place_and_returns_it(self, name, lhs_shape, [0.0] * (expected.rows * expected.columns)) result = getattr(lhs, name)(rhs, out=out) assert result is out - assert _flatten(out) == _flatten(expected) + assert flatten(out) == flatten(expected) def test_out_may_alias_an_input(self, name, rng): """out= aliasing an operand still produces the correct result.""" op = _NAMED_OPS[name] a = Matrix(2, 2, [rng.uniform(1.0, 10.0) for _ in range(4)]) b = Matrix(2, 2, [rng.uniform(1.0, 10.0) for _ in range(4)]) - expected = _flatten(op(a, b)) + expected = flatten(op(a, b)) result = getattr(a, name)(b, out=a) assert result is a - assert _flatten(a) == expected + assert flatten(a) == expected def test_out_shape_mismatch_raises_and_leaves_target(self, name, rng): """A wrong-shape out= raises ValueError before any write.""" a = Matrix(2, 2, [rng.uniform(1.0, 10.0) for _ in range(4)]) b = Matrix(2, 2, [rng.uniform(1.0, 10.0) for _ in range(4)]) bad = Matrix(3, 3, [7.0] * 9) - before = _flatten(bad) + before = flatten(bad) with pytest.raises(ValueError, match=r"out shape 3x3 does not match result 2x2"): getattr(a, name)(b, out=bad) - assert _flatten(bad) == before + assert flatten(bad) == before def test_out_wrong_type_raises(self, name): """A non-Matrix out= raises TypeError.""" @@ -1483,26 +1539,26 @@ def test_named_add_1x1_with_list(self): """1x1 receiver + list operand via the named method (new surface).""" r = Matrix(1, 1, [5.0]).add([1.0, 2.0, 3.0]) assert isinstance(r, Matrix) - assert _flatten(r) == [6.0, 7.0, 8.0] + assert flatten(r) == [6.0, 7.0, 8.0] def test_named_add_1x1_with_list_out(self): """1x1 receiver + list operand with an out= target (new surface).""" out = Matrix(1, 3, [0.0, 0.0, 0.0]) r = Matrix(1, 1, [5.0]).add([1.0, 2.0, 3.0], out=out) assert r is out - assert _flatten(r) == [6.0, 7.0, 8.0] + assert flatten(r) == [6.0, 7.0, 8.0] def test_1x1_plus_list_operator(self): """1x1 receiver + list via the + operator (scalar-broadcast branch).""" r = Matrix(1, 1, [5.0]) + [1.0, 2.0, 3.0] assert isinstance(r, Matrix) - assert _flatten(r) == [6.0, 7.0, 8.0] + assert flatten(r) == [6.0, 7.0, 8.0] def test_list_plus_matrix_same_shape(self): """list + 1xN matrix via __radd__ (same-shape ewise branch).""" r = [1.0, 2.0, 3.0] + Matrix(1, 3, [4.0, 5.0, 6.0]) assert isinstance(r, Matrix) - assert _flatten(r) == [5.0, 7.0, 9.0] + assert flatten(r) == [5.0, 7.0, 9.0] def test_list_times_column_outer(self): """list * column vector via __rmul__ (reflected outer-product branch).""" @@ -1512,7 +1568,7 @@ def test_list_times_column_outer(self): expected = [col * row for col in (10.0, 20.0, 30.0) for row in (1.0, 2.0, 3.0)] - assert _flatten(r) == expected + assert flatten(r) == expected def test_list_matmul_column(self): """list @ column vector via __rmatmul__ (matmul wrap path).""" @@ -1535,7 +1591,7 @@ def test_list_minus_matrix_operand_order(self): """ r = [10.0, 20.0, 30.0] - Matrix(1, 3, [1.0, 2.0, 3.0]) assert isinstance(r, Matrix) - assert _flatten(r) == [9.0, 18.0, 27.0] + assert flatten(r) == [9.0, 18.0, 27.0] def test_list_div_matrix_operand_order(self): """[seq] / matrix computes seq / matrix (not matrix / seq). @@ -1545,7 +1601,7 @@ def test_list_div_matrix_operand_order(self): """ r = [10.0, 20.0, 30.0] / Matrix(1, 3, [2.0, 4.0, 5.0]) assert isinstance(r, Matrix) - assert _flatten(r) == [5.0, 5.0, 6.0] + assert flatten(r) == [5.0, 5.0, 6.0] class TestUnaryOutTarget: @@ -1558,8 +1614,14 @@ class TestUnaryOutTarget: "floor": math.floor, "round": round, "sqrt": math.sqrt, + "reciprocal": lambda v: 1.0 / v, } + # Every unary elementwise op with keyword-only in_place/out (the UNARY + # dict above omits sign/cos/sin, which lack simple value-reference fns). + KEYWORD_ONLY_UNARY = ["negate", "abs", "ceil", "floor", "round", "sqrt", + "reciprocal", "sign", "cos", "sin"] + @pytest.mark.parametrize("name", list(UNARY)) def test_out_writes_and_returns_target(self, name): """out= writes the result and returns the target object.""" @@ -1569,9 +1631,9 @@ def test_out_writes_and_returns_target(self, name): out = Matrix(2, 2, [0.0, 0.0, 0.0, 0.0]) result = getattr(m, name)(out=out) assert result is out - assert _flatten(out) == pytest.approx([ref(v) for v in vals]) + assert flatten(out) == pytest.approx([ref(v) for v in vals]) # self is untouched (out is a distinct buffer). - assert _flatten(m) == vals + assert flatten(m) == vals @pytest.mark.parametrize("name", list(UNARY)) def test_out_and_in_place_mutually_exclusive(self, name): @@ -1580,25 +1642,32 @@ def test_out_and_in_place_mutually_exclusive(self, name): out = Matrix(2, 2, [0.0, 0.0, 0.0, 0.0]) with pytest.raises(ValueError, match=r"out and in_place are mutually exclusive"): - getattr(m, name)(True, out=out) + getattr(m, name)(in_place=True, out=out) - @pytest.mark.parametrize("name", list(UNARY)) + @pytest.mark.parametrize("name", KEYWORD_ONLY_UNARY) + def test_in_place_is_keyword_only(self, name): + """in_place cannot be passed positionally on the unary methods.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + getattr(m, name)(True) + + @pytest.mark.parametrize("name", KEYWORD_ONLY_UNARY) def test_out_is_keyword_only(self, name): """out cannot be passed positionally on the unary methods.""" m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) out = Matrix(2, 2, [0.0, 0.0, 0.0, 0.0]) with pytest.raises(TypeError): - getattr(m, name)(False, out) + getattr(m, name)(out) def test_out_shape_mismatch_raises(self): """A wrong-shape out= raises ValueError before any write.""" m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) bad = Matrix(3, 3, [9.0] * 9) - before = _flatten(bad) + before = flatten(bad) with pytest.raises(ValueError, match=r"out shape 3x3 does not match result 2x2"): m.negate(out=bad) - assert _flatten(bad) == before + assert flatten(bad) == before def test_out_wrong_type_raises(self): """A non-Matrix out= raises TypeError.""" @@ -3008,6 +3077,22 @@ def test_within_tolerance(self, shape, rng): b = Matrix(rows, cols, perturbed) assert Matrix.allclose(a, b) + def test_lhs_not_acquired_raises(self): + """allclose on a cown-resident left operand raises RuntimeError.""" + a = Matrix(2, 2, 1.0) + b = Matrix(2, 2, 1.0) + Cown(a) + with pytest.raises(RuntimeError): + Matrix.allclose(a, b) + + def test_rhs_not_acquired_raises(self): + """allclose on a cown-resident right operand raises RuntimeError.""" + a = Matrix(2, 2, 1.0) + b = Matrix(2, 2, 1.0) + Cown(b) + with pytest.raises(RuntimeError): + Matrix.allclose(a, b) + class TestRepr: """Smoke tests for string representations.""" @@ -3275,6 +3360,107 @@ def test_take_axis_below_neg2_raises(self): m.take([0], -3) +class TestTakeOut: + """Tests for take(out=) — writing the selection into a pre-allocated matrix.""" + + def test_rows_into_out_returns_out(self, mat, shape): + """take(out=) writes the selected rows and returns the out matrix.""" + rows, cols = shape + indices = [rows - 1, 0] + out = Matrix(len(indices), cols, 0.0) + result = mat.take(indices, axis=0, out=out) + assert result is out + for out_r, src_r in enumerate(indices): + for c in range(cols): + assert out[out_r, c] == pytest.approx(mat[src_r, c]) + + def test_columns_into_out_returns_out(self, mat, shape): + """take(out=) writes the selected columns and returns the out matrix.""" + rows, cols = shape + indices = [cols - 1, 0] + out = Matrix(rows, len(indices), 0.0) + result = mat.take(indices, axis=1, out=out) + assert result is out + for r in range(rows): + for out_c, src_c in enumerate(indices): + assert out[r, out_c] == pytest.approx(mat[r, src_c]) + + def test_out_matches_fresh_result(self, mat, shape): + """The out= result equals the freshly-allocated result.""" + rows, cols = shape + indices = [0, 0, rows - 1] + out = Matrix(len(indices), cols, 0.0) + mat.take(indices, axis=0, out=out) + assert Matrix.allclose(out, mat.take(indices, axis=0)) + + def test_out_is_keyword_only(self, mat, shape): + """out cannot be passed positionally (it is keyword-only).""" + rows, cols = shape + with pytest.raises(TypeError): + mat.take([0], 0, Matrix(1, cols, 0.0)) + + def test_out_wrong_shape_raises(self): + """An out of the wrong shape raises ValueError.""" + m = Matrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + with pytest.raises(ValueError): + m.take([0, 1], axis=0, out=Matrix(3, 2, 0.0)) + with pytest.raises(ValueError): + m.take([0], axis=1, out=Matrix(3, 2, 0.0)) + + def test_out_non_matrix_raises(self): + """A non-Matrix out raises TypeError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(TypeError): + m.take([0], axis=0, out=[1, 2]) + + def test_out_alias_self_raises(self): + """out aliasing the source matrix raises ValueError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(ValueError): + m.take([0, 1, 2], axis=0, out=m) + + def test_out_unacquired_raises(self): + """A cown-resident out raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + out = Matrix(2, 2, 0.0) + Cown(out) + with pytest.raises(RuntimeError): + m.take([0, 1], axis=0, out=out) + + def test_bad_index_leaves_out_untouched(self): + """A bad index is rejected before any write (validate-then-write).""" + m = Matrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + out = Matrix(2, 2, 9.0) + with pytest.raises(IndexError): + m.take([0, 99], axis=0, out=out) + for r in range(2): + for c in range(2): + assert out[r, c] == pytest.approx(9.0) + + def test_out_none_allocates_fresh(self, mat, shape): + """out=None behaves like the default (fresh allocation).""" + rows, cols = shape + indices = [0] + result = mat.take(indices, axis=0, out=None) + assert result is not mat + assert Matrix.allclose(result, mat.take(indices, axis=0)) + + def test_out_boc_roundtrip(self): + """take(out=) runs inside a @when behavior over cowns.""" + a = Cown(Matrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) + b = Cown(Matrix(2, 2, 0.0)) + + @when(a, b) + def result(a, b): # noqa: D401 — short behavior + """Take rows into a pre-allocated cown matrix.""" + a.value.take([2, 0], axis=0, out=b.value) + return list(b.value.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([5.0, 6.0, 1.0, 2.0]) + + class TestFancyIndexing: """List-key gather through __getitem__ (rows and columns).""" @@ -3422,6 +3608,87 @@ def test_getitem_fuzz_permutation(self, mat, shape, rng): assert result[out_r, c] == pytest.approx(mat[src_r, c]) +class TestSliceForcesMatrix: + """A slice anywhere in the key keeps a Matrix, even at 1x1.""" + + @pytest.fixture + def m(self): + """A 3x3 matrix filled with 0..8 row-major.""" + return Matrix(3, 3, [float(i) for i in range(9)]) + + def test_int_int_collapses_to_float(self, m): + """m[i, j] with two integer selectors returns a Python float.""" + value = m[1, 2] + assert isinstance(value, float) + assert value == pytest.approx(5.0) + + def test_slice_slice_stays_matrix(self, m): + """m[0:1, 0:1] returns a 1x1 Matrix, not a float.""" + sub = m[0:1, 0:1] + assert isinstance(sub, Matrix) + assert sub.rows == 1 and sub.columns == 1 + assert sub[0, 0] == pytest.approx(0.0) + + def test_int_slice_stays_matrix(self, m): + """m[i, c:c+1] keeps the column axis as a length-1 Matrix.""" + sub = m[1, 0:1] + assert isinstance(sub, Matrix) + assert sub.rows == 1 and sub.columns == 1 + assert sub[0, 0] == pytest.approx(3.0) + + def test_slice_int_stays_matrix(self, m): + """m[r:r+1, j] keeps the row axis as a length-1 Matrix.""" + sub = m[0:1, 2] + assert isinstance(sub, Matrix) + assert sub.rows == 1 and sub.columns == 1 + assert sub[0, 0] == pytest.approx(2.0) + + def test_1x1_matrix_int_int_is_float(self): + """A whole-matrix (int, int) read on a 1x1 still collapses.""" + assert isinstance(Matrix(1, 1, [5.0])[0, 0], float) + + def test_1x1_matrix_slice_slice_is_matrix(self): + """m[0:1, 0:1] on a 1x1 keeps it a Matrix.""" + sub = Matrix(1, 1, [5.0])[0:1, 0:1] + assert isinstance(sub, Matrix) + assert sub.rows == 1 and sub.columns == 1 + assert sub[0, 0] == pytest.approx(5.0) + + def test_column_vector_integer_still_float(self): + """A bare integer on a column vector keeps the 1-D float ergonomic.""" + v = Matrix(4, 1, [10.0, 20.0, 30.0, 40.0]) + assert isinstance(v[2], float) + assert v[2] == pytest.approx(30.0) + + def test_column_vector_slice_stays_matrix(self): + """A slice on a column vector keeps a Matrix, even at 1x1.""" + v = Matrix(4, 1, [10.0, 20.0, 30.0, 40.0]) + sub = v[2:3, 0:1] + assert isinstance(sub, Matrix) + assert sub.rows == 1 and sub.columns == 1 + assert sub[0, 0] == pytest.approx(30.0) + + def test_multi_cell_slice_unchanged(self, m): + """A slice selecting more than one cell is unaffected.""" + sub = m[0:2, 1:3] + assert isinstance(sub, Matrix) + assert sub.rows == 2 and sub.columns == 2 + + def test_boc_roundtrip(self): + """The length-1 slice rule holds inside a @when behavior.""" + a = Cown(Matrix(3, 3, [float(i) for i in range(9)])) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Read a length-1 slice and report its type and value.""" + sub = a.value[1, 0:1] + return (isinstance(sub, Matrix), sub.rows, sub.columns, sub[0, 0]) + + wait() + assert result.exception is False + assert result.value == (True, 1, 1, pytest.approx(3.0)) + + def _snapshot(m, rows, cols): """Return a flat list snapshot of a matrix's elements.""" return [m[r, c] for r in range(rows) for c in range(cols)] @@ -3965,157 +4232,1120 @@ def result(a): # noqa: D401 — short behavior assert Matrix.allclose(result.value, expected) -VECTOR_LENGTHS = [1, 3, 5, 10, 32] +class TestTakeAlongAxis: + """take_along_axis() — per-row/per-column gather (np.take_along_axis).""" + + def test_gather_axis1_rows(self): + """axis=1 gathers one column index per row into a column vector.""" + m = Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0]) + out = m.take_along_axis([1, 0, 2], axis=1) + assert out.rows == 3 + assert out.columns == 1 + assert Matrix.allclose(out, Matrix(3, 1, [1.0, 0.0, 1.0])) + + def test_gather_axis0_columns(self): + """axis=0 gathers one row index per column into a row vector.""" + m = Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0]) + out = m.take_along_axis([1, 0, 2], axis=0) + assert out.rows == 1 + assert out.columns == 3 + # out[c] = m[idx[c]][c]: m[1][0]=0, m[0][1]=1, m[2][2]=1. + assert Matrix.allclose(out, Matrix(1, 3, [0.0, 1.0, 1.0])) + + def test_gather_pairs_with_argmin(self, mat, shape): + """argmin(axis=k) feeds take_along_axis(axis=k) to read the minima.""" + rows, cols = shape + idx = mat.argmin(axis=1) + out = mat.take_along_axis(idx, axis=1) + assert out.rows == rows + assert out.columns == 1 + for r in range(rows): + assert out[r, 0] == pytest.approx(min(mat[r, c] + for c in range(cols))) + def test_gather_pairs_with_argmin_axis0(self, mat, shape): + """argmin(axis=0) feeds take_along_axis(axis=0) for column minima.""" + rows, cols = shape + idx = mat.argmin(axis=0) + out = mat.take_along_axis(idx, axis=0) + assert out.rows == 1 + assert out.columns == cols + for c in range(cols): + assert out[0, c] == pytest.approx(min(mat[r, c] + for r in range(rows))) -class TestVectorLen: - """len() returns the number of rows for all matrices, including vectors.""" + def test_gather_accepts_tuple(self): + """A tuple of indices is accepted like a list.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.take_along_axis((1, 0), axis=1) + assert Matrix.allclose(out, Matrix(2, 1, [2.0, 3.0])) - @pytest.mark.parametrize("n", VECTOR_LENGTHS) - def test_row_vector_len(self, n): - """len() of a 1xN row vector returns 1 (the number of rows).""" - v = Matrix(1, n, [float(i) for i in range(n)]) - assert len(v) == 1 + def test_gather_negative_axis(self): + """A negative axis maps -1 to columns (axis=1).""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.take_along_axis([1, 0], axis=-1) + assert Matrix.allclose(out, Matrix(2, 1, [2.0, 3.0])) - @pytest.mark.parametrize("n", VECTOR_LENGTHS) - def test_column_vector_len(self, n): - """len() of an Nx1 column vector returns N (the number of rows).""" - v = Matrix(n, 1, [float(i) for i in range(n)]) - assert len(v) == n + def test_gather_negative_index(self): + """A negative index counts from the end of the gathered axis.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + out = m.take_along_axis([-1, -2], axis=1) + assert Matrix.allclose(out, Matrix(2, 1, [3.0, 5.0])) - def test_matrix_len_returns_rows(self): - """len() of a non-vector matrix returns rows.""" - m = Matrix(3, 4) - assert len(m) == 3 + def test_gather_bool_index(self): + """A bool index element is taken as the integer 0/1.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.take_along_axis([True, False], axis=1) + assert Matrix.allclose(out, Matrix(2, 1, [2.0, 3.0])) + def test_gather_wrong_count_raises(self): + """An index count that does not match the axis raises ValueError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(ValueError): + m.take_along_axis([0, 1], axis=1) + with pytest.raises(ValueError): + m.take_along_axis([0, 1], axis=0) -class TestVectorItemAccess: - """Integer indexing and iteration on vectors.""" + def test_gather_out_of_range_raises(self): + """An out-of-range index raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m.take_along_axis([0, 1, 9], axis=1) - @pytest.mark.parametrize("n", VECTOR_LENGTHS) - def test_row_vector_item(self, n): - """Indexing a 1xN row vector with [0] returns the row as a Matrix.""" - vals = [float(i) * 1.5 for i in range(n)] - v = Matrix(1, n, vals) - row = v[0] - if n == 1: - assert isinstance(row, float) - assert row == pytest.approx(vals[0]) - else: - assert isinstance(row, Matrix) - assert row.rows == 1 - assert row.columns == n - for i in range(n): - assert v[0, i] == pytest.approx(vals[i]) + def test_gather_invalid_axis_raises(self): + """axis >= 2 raises KeyError, matching take().""" + m = Matrix(3, 3, 1.0) + with pytest.raises(KeyError): + m.take_along_axis([0, 1, 2], 2) - @pytest.mark.parametrize("n", VECTOR_LENGTHS) - def test_column_vector_item(self, n): - """Indexing an Nx1 column vector with an integer returns the element as a float.""" - vals = [float(i) * 2.0 for i in range(n)] - v = Matrix(n, 1, vals) - for i in range(n): - item = v[i] - assert isinstance(item, float) - assert item == pytest.approx(vals[i]) + def test_gather_bad_type_raises(self): + """A non-int index raises TypeError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(TypeError): + m.take_along_axis([0.5, 1.0], axis=1) - @pytest.mark.parametrize("n", [3, 5, 10]) - def test_row_vector_iteration(self, n): - """Iterating a 1xN row vector yields one Matrix (the single row).""" - vals = [float(i) for i in range(n)] - v = Matrix(1, n, vals) - collected = list(v) - assert len(collected) == 1 - assert isinstance(collected[0], Matrix) - assert collected[0].rows == 1 - assert collected[0].columns == n + def test_gather_boc_roundtrip(self): + """take_along_axis runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0])) - @pytest.mark.parametrize("n", [3, 5, 10]) - def test_column_vector_iteration(self, n): - """Iterating an Nx1 column vector yields individual float elements.""" - vals = [float(i) for i in range(n)] - v = Matrix(n, 1, vals) - collected = list(v) - assert len(collected) == n - for got, expected in zip(collected, vals): - assert isinstance(got, float) - assert got == pytest.approx(expected) + @when(a) + def result(a): # noqa: D401 — short behavior + """Gather per-row minima inside a behavior.""" + return a.value.take_along_axis( + a.value.argmin(axis=1), axis=1) + wait() + assert result.exception is False + assert Matrix.allclose(result.value, Matrix(3, 1, [1.0, 0.0, 1.0])) -class TestRowVectorIndexing: - """Subscript on a 1-row matrix behaves like any other matrix.""" - def test_row_vector_two_index_reads_element(self): - """v[0, i] on a row vector reads column i.""" - v = Matrix(1, 5, [10.0, 20.0, 30.0, 40.0, 50.0]) - assert v[0, 2] == pytest.approx(30.0) - assert v[0, 0] == pytest.approx(10.0) - assert v[0, 4] == pytest.approx(50.0) +class TestTakeAlongAxisOut: + """take_along_axis(out=) — gather into a pre-allocated matrix.""" - def test_row_vector_row_slice(self): - """v[0:1] on a 1-row matrix returns the whole row.""" - v = Matrix(1, 6, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - sub = v[0:1] - assert sub.rows == 1 - assert sub.columns == 6 - assert sub[0, 0] == pytest.approx(1.0) - assert sub[0, 5] == pytest.approx(6.0) + def test_axis1_into_out_returns_out(self): + """axis=1 writes the gathered column vector into out and returns it.""" + m = Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0]) + out = Matrix(3, 1, 0.0) + result = m.take_along_axis([1, 0, 2], axis=1, out=out) + assert result is out + assert Matrix.allclose(out, Matrix(3, 1, [1.0, 0.0, 1.0])) + def test_axis0_into_out_returns_out(self): + """axis=0 writes the gathered row vector into out and returns it.""" + m = Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0]) + out = Matrix(1, 3, 0.0) + result = m.take_along_axis([1, 0, 2], axis=0, out=out) + assert result is out + assert Matrix.allclose(out, Matrix(1, 3, [0.0, 1.0, 1.0])) -class TestVectorBroadcastArithmetic: - """Arithmetic broadcasting with row and column vectors.""" + def test_out_matches_fresh_result(self, mat, shape): + """The out= result equals the freshly-allocated result.""" + rows, cols = shape + idx = mat.argmin(axis=1) + out = Matrix(rows, 1, 0.0) + mat.take_along_axis(idx, axis=1, out=out) + assert Matrix.allclose(out, mat.take_along_axis(idx, axis=1)) - @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) - def test_row_vector_add_broadcast(self, rows, cols): - """A (1xcols) row vector added to (rowsxcols) broadcasts across rows.""" - m_vals = [float(i) for i in range(rows * cols)] - m = Matrix(rows, cols, m_vals) - v_vals = [float(j) * 100 for j in range(cols)] - v = Matrix(1, cols, v_vals) - result = m + v - assert result.rows == rows - assert result.columns == cols - for i in range(rows): - for j in range(cols): - expected = m_vals[i * cols + j] + v_vals[j] - assert result[i, j] == pytest.approx(expected) + def test_out_is_keyword_only(self): + """out cannot be passed positionally (it is keyword-only).""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + m.take_along_axis([1, 0], 1, Matrix(2, 1, 0.0)) - @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) - def test_row_vector_subtract_broadcast(self, rows, cols): - """Subtraction broadcasts a row vector across all rows.""" - m = Matrix.uniform(0.0, 10.0, size=(rows, cols)) - v = Matrix.uniform(0.0, 5.0, size=(1, cols)) - result = m - v - for i in range(rows): - for j in range(cols): - assert result[i, j] == pytest.approx(m[i, j] - v[0, j]) + def test_out_wrong_shape_raises(self): + """An out of the wrong shape raises ValueError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(ValueError): + m.take_along_axis([0, 1, 2], axis=1, out=Matrix(3, 3, 0.0)) + with pytest.raises(ValueError): + m.take_along_axis([0, 1, 2], axis=0, out=Matrix(3, 1, 0.0)) - @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) - def test_row_vector_multiply_broadcast(self, rows, cols): - """Multiplication broadcasts a row vector across all rows.""" - m = Matrix.uniform(1.0, 10.0, size=(rows, cols)) - v = Matrix.uniform(1.0, 5.0, size=(1, cols)) - result = m * v - for i in range(rows): - for j in range(cols): - assert result[i, j] == pytest.approx(m[i, j] * v[0, j]) + def test_out_non_matrix_raises(self): + """A non-Matrix out raises TypeError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(TypeError): + m.take_along_axis([0, 1], axis=1, out=[1, 2]) - @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) - def test_row_vector_divide_broadcast(self, rows, cols): - """Division broadcasts a row vector across all rows.""" - m = Matrix.uniform(1.0, 10.0, size=(rows, cols)) - v = Matrix.uniform(1.0, 5.0, size=(1, cols)) - result = m / v - for i in range(rows): - for j in range(cols): - assert result[i, j] == pytest.approx(m[i, j] / v[0, j]) + def test_out_alias_self_raises(self): + """out aliasing the source matrix raises ValueError.""" + m = Matrix(1, 3, [1.0, 2.0, 3.0]) + with pytest.raises(ValueError): + m.take_along_axis([0, 0, 0], axis=0, out=m) - @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) - def test_column_vector_add_broadcast(self, rows, cols): - """An (rowsx1) column vector added to (rowsxcols) broadcasts across columns.""" - m_vals = [float(i) for i in range(rows * cols)] - m = Matrix(rows, cols, m_vals) - v_vals = [float(i) * 100 for i in range(rows)] - v = Matrix(rows, 1, v_vals) + def test_out_unacquired_raises(self): + """A cown-resident out raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + out = Matrix(2, 1, 0.0) + Cown(out) + with pytest.raises(RuntimeError): + m.take_along_axis([0, 1], axis=1, out=out) + + def test_bad_index_leaves_out_untouched(self): + """A bad index is rejected before any write (validate-then-write).""" + m = Matrix(3, 3, 1.0) + out = Matrix(3, 1, 9.0) + with pytest.raises(IndexError): + m.take_along_axis([0, 1, 9], axis=1, out=out) + for r in range(3): + assert out[r, 0] == pytest.approx(9.0) + + def test_out_none_allocates_fresh(self): + """out=None behaves like the default (fresh allocation).""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + result = m.take_along_axis([1, 0], axis=1, out=None) + assert result is not m + assert Matrix.allclose(result, Matrix(2, 1, [2.0, 3.0])) + + def test_out_boc_roundtrip(self): + """take_along_axis(out=) runs inside a @when behavior over cowns.""" + a = Cown(Matrix(3, 3, [3.0, 1.0, 2.0, 0.0, 5.0, 4.0, 9.0, 8.0, 1.0])) + b = Cown(Matrix(3, 1, 0.0)) + + @when(a, b) + def result(a, b): # noqa: D401 — short behavior + """Gather per-row minima into a pre-allocated cown matrix.""" + a.value.take_along_axis(a.value.argmin(axis=1), + axis=1, out=b.value) + return list(b.value.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([1.0, 0.0, 1.0]) + + +class TestPutAlongAxis: + """put_along_axis() — the write-side counterpart of take_along_axis().""" + + def test_put_axis1_scalar(self): + """axis=1 writes a scalar to one column per row.""" + m = Matrix(3, 3, 1.0) + m.put_along_axis([0, 1, 2], 9.0, axis=1) + assert Matrix.allclose( + m, Matrix(3, 3, [9.0, 1.0, 1.0, 1.0, 9.0, 1.0, 1.0, 1.0, 9.0])) + + def test_put_axis1_vector(self): + """axis=1 writes a rows x 1 vector, one element per row.""" + m = Matrix(3, 3, 1.0) + m.put_along_axis([0, 1, 2], Matrix(3, 1, [7.0, 8.0, 9.0]), axis=1) + assert Matrix.allclose( + m, Matrix(3, 3, [7.0, 1.0, 1.0, 1.0, 8.0, 1.0, 1.0, 1.0, 9.0])) + + def test_put_axis0_vector(self): + """axis=0 writes a 1 x columns vector, one element per column.""" + m = Matrix(3, 3, 1.0) + m.put_along_axis([0, 1, 2], Matrix(1, 3, [7.0, 8.0, 9.0]), axis=0) + assert Matrix.allclose( + m, Matrix(3, 3, [7.0, 1.0, 1.0, 1.0, 8.0, 1.0, 1.0, 1.0, 9.0])) + + def test_put_round_trips_take(self, mat, shape): + """put_along_axis writes back exactly what take_along_axis gathered.""" + rows, cols = shape + idx = mat.argmin(axis=1) + gathered = mat.take_along_axis(idx, axis=1) + target = Matrix(rows, cols, 0.0) + target.put_along_axis(idx, gathered, axis=1) + for r in range(rows): + assert target[r, idx[r]] == pytest.approx(gathered[r, 0]) + + def test_put_1x1_as_scalar(self): + """A 1x1 matrix value broadcasts like a scalar.""" + m = Matrix(2, 2, 0.0) + m.put_along_axis([0, 1], Matrix(1, 1, 8.0), axis=1) + assert Matrix.allclose(m, Matrix(2, 2, [8.0, 0.0, 0.0, 8.0])) + + def test_put_returns_self(self): + """put_along_axis returns self to allow chaining.""" + m = Matrix(2, 2, 0.0) + assert m.put_along_axis([0, 1], 1.0, axis=1) is m + + def test_put_accumulate(self): + """accumulate=True folds duplicate target cells additively.""" + m = Matrix(3, 3, 1.0) + m.put_along_axis([0, 0, 0], 5.0, axis=1, accumulate=True) + assert Matrix.allclose( + m, Matrix(3, 3, [6.0, 1.0, 1.0, 6.0, 1.0, 1.0, 6.0, 1.0, 1.0])) + + def test_put_negative_index(self): + """A negative index counts from the end of the indexed axis.""" + m = Matrix(2, 3, 0.0) + m.put_along_axis([-1, -1], 4.0, axis=1) + assert Matrix.allclose(m, Matrix(2, 3, [0.0, 0.0, 4.0, 0.0, 0.0, 4.0])) + + def test_put_wrong_count_raises_nowrite(self): + """A bad index count raises ValueError and writes nothing.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(ValueError): + m.put_along_axis([0, 1], 5.0, axis=1) + assert _snapshot(m, 3, 3) == before + + def test_put_out_of_range_raises_nowrite(self): + """An out-of-range index raises IndexError and writes nothing.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(IndexError): + m.put_along_axis([0, 1, 9], 5.0, axis=1) + assert _snapshot(m, 3, 3) == before + + def test_put_shape_mismatch_raises_nowrite(self): + """A wrong-shaped value raises ValueError and writes nothing.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(ValueError): + m.put_along_axis([0, 1, 2], Matrix(1, 3, [1.0, 2.0, 3.0]), axis=1) + assert _snapshot(m, 3, 3) == before + + def test_put_invalid_axis_raises(self): + """axis >= 2 raises KeyError, matching put().""" + m = Matrix(3, 3, 1.0) + with pytest.raises(KeyError): + m.put_along_axis([0, 1, 2], 1.0, 2) + + def test_put_bad_value_type_raises(self): + """A non-numeric, non-matrix value raises TypeError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(TypeError): + m.put_along_axis([0, 1], "x", axis=1) + + def test_put_self_alias(self): + """A self-aliased RHS reads pre-write values via the snapshot.""" + # m is 2x1, so a put_along_axis(axis=1) selection (2x1) can alias self. + m = Matrix(2, 1, [3.0, 5.0]) + m.put_along_axis([0, 0], m, axis=1, accumulate=True) + # Each row reads its snapshot value then adds it: doubles. + assert Matrix.allclose(m, Matrix(2, 1, [6.0, 10.0])) + + def test_put_boc_roundtrip(self): + """put_along_axis runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(3, 3, 1.0)) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Scatter one element per row inside a behavior.""" + a.value.put_along_axis([0, 1, 2], Matrix(3, 1, [7.0, 8.0, 9.0]), + axis=1) + return a.value.copy() + + wait() + assert result.exception is False + assert Matrix.allclose( + result.value, + Matrix(3, 3, [7.0, 1.0, 1.0, 1.0, 8.0, 1.0, 1.0, 1.0, 9.0])) + + +class TestRepeatInterleave: + """repeat_interleave() — interleaved row/column/element repetition.""" + + def test_flatten_default(self): + """axis=None flattens row-major and repeats each element.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.repeat_interleave(2) + assert out.rows == 1 + assert out.columns == 8 + assert Matrix.allclose( + out, Matrix(1, 8, [1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0])) + + def test_axis0_repeats_rows(self): + """axis=0 repeats whole rows consecutively.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.repeat_interleave(2, axis=0) + assert out.rows == 4 + assert out.columns == 2 + assert Matrix.allclose( + out, Matrix(4, 2, [1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.0])) + + def test_axis1_repeats_columns(self): + """axis=1 repeats each column consecutively.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = m.repeat_interleave(3, axis=1) + assert out.rows == 2 + assert out.columns == 6 + assert Matrix.allclose( + out, Matrix(2, 6, [1.0, 1.0, 1.0, 2.0, 2.0, 2.0, + 3.0, 3.0, 3.0, 4.0, 4.0, 4.0])) + + def test_repeat_one_is_copy(self, mat, shape): + """repeats=1 returns an equal but independent matrix.""" + rows, cols = shape + out = mat.repeat_interleave(1, axis=0) + assert out.rows == rows + assert out.columns == cols + assert Matrix.allclose(out, mat) + assert out is not mat + + def test_negative_axis(self): + """axis=-1 maps to columns, axis=-2 to rows.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert Matrix.allclose(m.repeat_interleave(2, axis=-1), + m.repeat_interleave(2, axis=1)) + assert Matrix.allclose(m.repeat_interleave(2, axis=-2), + m.repeat_interleave(2, axis=0)) + + def test_shapes_fuzzed(self, mat, shape): + """Result dimensions and element multiplicity hold across shapes.""" + rows, cols = shape + r0 = mat.repeat_interleave(3, axis=0) + assert (r0.rows, r0.columns) == (rows * 3, cols) + for r in range(rows): + for t in range(3): + for c in range(cols): + assert r0[r * 3 + t, c] == pytest.approx(mat[r, c]) + r1 = mat.repeat_interleave(3, axis=1) + assert (r1.rows, r1.columns) == (rows, cols * 3) + for r in range(rows): + for c in range(cols): + for t in range(3): + assert r1[r, c * 3 + t] == pytest.approx(mat[r, c]) + + def test_zero_repeats_raises(self): + """repeats < 1 raises ValueError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(ValueError): + m.repeat_interleave(0) + with pytest.raises(ValueError): + m.repeat_interleave(-1, axis=0) + + def test_invalid_axis_raises(self): + """An out-of-range axis raises ValueError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(ValueError): + m.repeat_interleave(2, axis=2) + + @pytest.mark.parametrize("rows, cols, axis", [ + (3, 1, 1), # axis=1: columns*repeats fits, but total wraps size_t + (1, 3, 0), # axis=0: rows*repeats fits, but total wraps size_t + (3, 1, None), # flatten: size*repeats wraps size_t + ]) + def test_repeats_overflow_raises(self, rows, cols, axis): + """Overflowing the total output element count raises OverflowError. + + The bound is on the whole rows*columns*repeats product, not just the + repeated dimension: a large repeats with a small repeated axis but the + *other* dimension > 1 would otherwise wrap the product to a tiny size + and write past the heap allocation. sys.maxsize is the largest value + the ``n`` (Py_ssize_t) argument accepts; for size >= 3 the total + product exceeds SIZE_MAX while the single-dimension product stays in + range. + """ + m = Matrix(rows, cols, 1.0) + with pytest.raises(OverflowError): + if axis is None: + m.repeat_interleave(sys.maxsize) + else: + m.repeat_interleave(sys.maxsize, axis=axis) + + def test_not_acquired_raises(self): + """repeat_interleave on a cown-resident matrix raises.""" + m = Matrix(2, 2, 1.0) + Cown(m) + with pytest.raises(RuntimeError): + m.repeat_interleave(2) + + def test_boc_roundtrip(self): + """repeat_interleave runs inside a @when behavior over a Cown.""" + a = Cown(Matrix(2, 2, [1.0, 2.0, 3.0, 4.0])) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Repeat rows inside a behavior.""" + return a.value.repeat_interleave(2, axis=0) + + wait() + assert result.exception is False + assert Matrix.allclose( + result.value, + Matrix(4, 2, [1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.0])) + + +def _topk_group(values, mask, k, largest): + """Reference top-k of one group: sorted extremes, NaN last, -1/NaN pad. + + Mirrors the C kernel: only included cells (``mask[i] != 0.0``; NaN counts + as included) are considered, the result is sorted (descending for + ``largest``, ascending otherwise), ties keep the first occurrence, NaN + sorts last, and a group with fewer than *k* included elements pads the + trailing slots with ``NaN`` values and ``-1`` indices. + """ + entries = [(v, i) for i, v in enumerate(values) + if mask is None or mask[i] != 0.0] + + def key(entry): + v, i = entry + is_nan = math.isnan(v) + sort_value = 0.0 if is_nan else (-v if largest else v) + return (is_nan, sort_value, i) + + entries.sort(key=key) + out_vals, out_idx = [], [] + for j in range(k): + if j < len(entries): + out_vals.append(entries[j][0]) + out_idx.append(entries[j][1]) + else: + out_vals.append(float("nan")) + out_idx.append(-1) + return out_vals, out_idx + + +def _assert_values_equal(got, expected): + """Compare value lists allowing NaN == NaN.""" + assert len(got) == len(expected) + for g, e in zip(got, expected): + if math.isnan(e): + assert math.isnan(g) + else: + assert g == pytest.approx(e) + + +class TestTopK: + """topk() — k extreme elements per group, sorted, with optional mask.""" + + def test_flat_largest_golden(self): + """axis=None returns the k greatest, descending, with flat indices.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2) + assert values.rows == 1 + assert values.columns == 2 + assert list(values.values()) == pytest.approx([9.0, 4.0]) + assert indices == [4, 5] + + def test_flat_smallest_golden(self): + """largest=False returns the k smallest, ascending.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2, largest=False) + assert list(values.values()) == pytest.approx([0.0, 1.0]) + assert indices == [3, 1] + + def test_axis1_rowwise_golden(self): + """axis=1 reduces across columns into rows x k values.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2, axis=1) + assert values.rows == 2 + assert values.columns == 2 + assert values[0, 0] == pytest.approx(3.0) + assert values[0, 1] == pytest.approx(2.0) + assert values[1, 0] == pytest.approx(9.0) + assert values[1, 1] == pytest.approx(4.0) + assert indices == [[0, 2], [1, 2]] + + def test_axis0_columnwise_golden(self): + """axis=0 reduces down rows into k x columns values.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2, axis=0) + assert values.rows == 2 + assert values.columns == 3 + # Column 0: [3, 0] -> [3, 0]; col 1: [1, 9] -> [9, 1]; col 2: [2, 4] + # -> [4, 2]. + assert values[0, 0] == pytest.approx(3.0) + assert values[1, 0] == pytest.approx(0.0) + assert values[0, 1] == pytest.approx(9.0) + assert values[1, 1] == pytest.approx(1.0) + assert indices == [[0, 1], [1, 0], [1, 0]] + + def test_ties_keep_first_occurrence(self): + """Equal values resolve to ascending original index.""" + m = Matrix(1, 4, [5.0, 5.0, 5.0, 5.0]) + _, indices = m.topk(3) + assert indices == [0, 1, 2] + + def test_nan_sorts_last(self): + """A NaN is never a top value; it lands after the real elements.""" + nan = float("nan") + m = Matrix(1, 4, [1.0, nan, 3.0, 2.0]) + values, indices = m.topk(4) + _assert_values_equal(list(values.values()), [3.0, 2.0, 1.0, nan]) + assert indices == [2, 3, 0, 1] + + def test_k_equals_axis_length_full_sort(self): + """k == axis length performs a complete sort of the group.""" + m = Matrix(1, 5, [4.0, 2.0, 5.0, 1.0, 3.0]) + values, indices = m.topk(5) + assert list(values.values()) == pytest.approx([5.0, 4.0, 3.0, 2.0, 1.0]) + assert indices == [2, 0, 4, 1, 3] + + def test_masked_axis1_pads_short_group(self): + """A masked row with fewer than k included cells pads with NaN/-1.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0]) + values, indices = m.topk(3, axis=1, where=mask) + _assert_values_equal([values[0, j] for j in range(3)], + [3.0, 2.0, float("nan")]) + _assert_values_equal([values[1, j] for j in range(3)], + [9.0, 4.0, float("nan")]) + assert indices == [[0, 2, -1], [1, 2, -1]] + + def test_masked_flatten(self): + """A flat masked top-k only considers included cells.""" + m = Matrix(1, 5, [3.0, 0.0, 9.0, 1.0, 4.0]) + mask = Matrix(1, 5, [1.0, 0.0, 0.0, 1.0, 1.0]) + values, indices = m.topk(2, where=mask) + assert list(values.values()) == pytest.approx([4.0, 3.0]) + assert indices == [4, 0] + + def test_all_excluded_group_all_pad(self): + """A fully-masked group is all NaN values and -1 indices.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + zero = Matrix(2, 3, 0.0) + values, indices = m.topk(2, axis=1, where=zero) + for r in range(2): + _assert_values_equal([values[r, j] for j in range(2)], + [float("nan"), float("nan")]) + assert indices == [[-1, -1], [-1, -1]] + + def test_nan_mask_cell_is_included(self): + """A NaN mask cell counts as included (truthy).""" + m = Matrix(1, 3, [5.0, 2.0, 8.0]) + mask = Matrix(1, 3, [0.0, float("nan"), 1.0]) + values, indices = m.topk(2, where=mask) + assert list(values.values()) == pytest.approx([8.0, 2.0]) + assert indices == [2, 1] + + def test_negative_axis(self): + """axis=-1 behaves like axis=1 and axis=-2 like axis=0.""" + m = Matrix(3, 2, [4.0, 1.0, 2.0, 8.0, 7.0, 3.0]) + v1, i1 = m.topk(2, axis=-1) + v2, i2 = m.topk(2, axis=1) + assert Matrix.allclose(v1, v2) + assert i1 == i2 + + def test_k_too_large_raises(self): + """k larger than the reduced axis length raises ValueError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(ValueError): + m.topk(7) # flat size is 6 + with pytest.raises(ValueError): + m.topk(3, axis=0) # only 2 rows + with pytest.raises(ValueError): + m.topk(4, axis=1) # only 3 columns + + def test_k_not_positive_raises(self): + """k < 1 raises ValueError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(ValueError): + m.topk(0) + with pytest.raises(ValueError): + m.topk(-2, axis=1) + + def test_invalid_axis_raises(self): + """An out-of-range axis raises ValueError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(ValueError): + m.topk(1, axis=2) + + def test_non_matrix_mask_raises(self): + """A non-Matrix where= argument raises TypeError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(TypeError): + m.topk(1, where=[1, 0, 1, 0, 1, 0]) + + def test_shape_mismatch_mask_raises(self): + """A mask of the wrong shape raises ValueError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(ValueError): + m.topk(1, where=Matrix(3, 2, 1.0)) + + def test_not_acquired_raises(self): + """topk on a cown-resident matrix raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + Cown(m) + with pytest.raises(RuntimeError): + m.topk(1) + + def test_unacquired_mask_raises(self): + """A cown-resident mask raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + mask = Matrix(2, 2, 1.0) + Cown(mask) + with pytest.raises(RuntimeError): + m.topk(1, where=mask) + + @pytest.mark.parametrize("largest", [True, False]) + def test_fuzzed_flatten(self, mat, shape, random_values, largest): + """Flat top-k matches a pure-Python reference across shapes.""" + rows, cols = shape + k = max(1, (rows * cols) // 2) + values, indices = mat.topk(k, largest=largest) + exp_vals, exp_idx = _topk_group(random_values, None, k, largest) + _assert_values_equal(list(values.values()), exp_vals) + assert indices == exp_idx + + @pytest.mark.parametrize("largest", [True, False]) + def test_fuzzed_axis1(self, mat, shape, random_values, largest): + """Per-row top-k matches a pure-Python reference across shapes.""" + rows, cols = shape + k = max(1, cols // 2) + values, indices = mat.topk(k, axis=1, largest=largest) + assert values.rows == rows + assert values.columns == k + for r in range(rows): + row = [random_values[r * cols + c] for c in range(cols)] + exp_vals, exp_idx = _topk_group(row, None, k, largest) + _assert_values_equal([values[r, j] for j in range(k)], exp_vals) + assert indices[r] == exp_idx + + @pytest.mark.parametrize("largest", [True, False]) + def test_fuzzed_axis0(self, mat, shape, random_values, largest): + """Per-column top-k matches a pure-Python reference across shapes.""" + rows, cols = shape + k = max(1, rows // 2) + values, indices = mat.topk(k, axis=0, largest=largest) + assert values.rows == k + assert values.columns == cols + for c in range(cols): + column = [random_values[r * cols + c] for r in range(rows)] + exp_vals, exp_idx = _topk_group(column, None, k, largest) + _assert_values_equal([values[j, c] for j in range(k)], exp_vals) + assert indices[c] == exp_idx + + @pytest.mark.parametrize("largest", [True, False]) + def test_fuzzed_masked_axis1(self, mat, shape, random_values, rng, + largest): + """Masked per-row top-k matches a pure-Python reference.""" + rows, cols = shape + k = cols # full width so masking can leave short groups + mask_flat = [1.0 if rng.random() < 0.6 else 0.0 + for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_flat) + values, indices = mat.topk(k, axis=1, largest=largest, where=mask) + for r in range(rows): + row = [random_values[r * cols + c] for c in range(cols)] + row_mask = [mask_flat[r * cols + c] for c in range(cols)] + exp_vals, exp_idx = _topk_group(row, row_mask, k, largest) + _assert_values_equal([values[r, j] for j in range(k)], exp_vals) + assert indices[r] == exp_idx + + def test_boc_roundtrip(self): + """topk runs inside a @when behavior over cowns.""" + a = Cown(Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0])) + b = Cown(Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0])) + + @when(a, b) + def result(a, b): # noqa: D401 — short behavior + """Masked top-k inside a behavior.""" + values, indices = a.value.topk(2, axis=1, where=b.value) + return list(values.values()), indices + + wait() + assert result.exception is False + vals, idx = result.value + assert vals == pytest.approx([3.0, 2.0, 9.0, 4.0]) + assert idx == [[0, 2], [1, 2]] + + +def _matrix_idx_to_lists(indices, axis, k): + """Reshape a matrix-form topk index result back to the list form. + + Mirrors the index Matrix layout: axis=None -> a flat list of k ints from + the 1 x k row; axis=0 -> per column, k row indices read down the k x cols + matrix; axis=1 -> per row, k column indices read across the rows x k + matrix. Float cells (including the -1.0 pad) are cast back to ints. + """ + if axis is None: + return [int(indices[0, j]) for j in range(k)] + if axis == 0: + cols = indices.columns + return [[int(indices[j, c]) for j in range(k)] for c in range(cols)] + rows = indices.rows + return [[int(indices[r, j]) for j in range(k)] for r in range(rows)] + + +class TestTopKIndexMatrix: + """topk(as_matrix=True) — indices returned as a same-shape index Matrix.""" + + def test_default_is_matrix(self): + """With as_matrix=True the indices come back as a Matrix, not a list.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + _, indices = m.topk(2, as_matrix=True) + assert isinstance(indices, Matrix) + + def test_flat_index_matrix_shape_and_values(self): + """axis=None indices is a 1 x k matrix of flat indices.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + _, indices = m.topk(2, as_matrix=True) + assert indices.rows == 1 + assert indices.columns == 2 + assert list(indices.values()) == pytest.approx([4.0, 5.0]) + + def test_axis1_index_matrix_matches_values_shape(self): + """axis=1 indices is a rows x k matrix aligned with the values.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2, axis=1, as_matrix=True) + assert (indices.rows, indices.columns) == (values.rows, values.columns) + assert list(indices.values()) == pytest.approx([0.0, 2.0, 1.0, 2.0]) + + def test_axis0_index_matrix_matches_values_shape(self): + """axis=0 indices is a k x cols matrix aligned with the values.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + values, indices = m.topk(2, axis=0, as_matrix=True) + assert (indices.rows, indices.columns) == (values.rows, values.columns) + assert list(indices.values()) == pytest.approx([0.0, 1.0, 1.0, + 1.0, 0.0, 0.0]) + + def test_pad_is_minus_one_float(self): + """A short masked group pads the index matrix with -1.0.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0]) + _, indices = m.topk(3, axis=1, where=mask, as_matrix=True) + assert [indices[0, j] for j in range(3)] == pytest.approx([0.0, 2.0, + -1.0]) + assert [indices[1, j] for j in range(3)] == pytest.approx([1.0, 2.0, + -1.0]) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("largest", [True, False]) + def test_matrix_form_matches_list_form(self, mat, shape, axis, largest): + """The index matrix carries the same indices as the list form.""" + rows, cols = shape + if axis is None: + k = max(1, (rows * cols) // 2) + elif axis == 0: + k = max(1, rows // 2) + else: + k = max(1, cols // 2) + v_mat, i_mat = mat.topk(k, axis=axis, largest=largest, as_matrix=True) + v_list, i_list = mat.topk(k, axis=axis, largest=largest) + assert Matrix.allclose(v_mat, v_list) + assert _matrix_idx_to_lists(i_mat, axis, k) == i_list + + def test_index_matrix_boc_roundtrip(self): + """topk(as_matrix=True) returns a cown-able index Matrix in a behavior.""" + a = Cown(Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0])) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Read the index matrix inside a behavior.""" + _, indices = a.value.topk(2, axis=1, as_matrix=True) + return list(indices.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([0.0, 2.0, 1.0, 2.0]) + + +class TestAggregateWhere: + """where= masking on sum/mean/magnitude/magnitude_squared/min/max.""" + + def test_sum_masks_elements(self): + """Only mask-selected elements contribute to the sum.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 0.0]) + assert m.sum(where=mask) == pytest.approx(1.0 + 3.0 + 5.0) + + def test_mean_uses_included_count(self): + """The mean divides by the number of included elements only.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 0.0]) + assert m.mean(where=mask) == pytest.approx((1.0 + 3.0 + 5.0) / 3.0) + + def test_min_max_ignore_excluded(self): + """min/max only consider included elements.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + # Exclude the global extremes (1 and 6). + mask = Matrix(2, 3, [0.0, 1.0, 1.0, 1.0, 1.0, 0.0]) + assert m.min(where=mask) == pytest.approx(2.0) + assert m.max(where=mask) == pytest.approx(5.0) + + def test_magnitude_masks_elements(self): + """magnitude and magnitude_squared respect the mask.""" + m = Matrix(2, 2, [3.0, 0.0, 0.0, 4.0]) + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + assert m.magnitude_squared(where=mask) == pytest.approx(25.0) + assert m.magnitude(where=mask) == pytest.approx(5.0) + + def test_nan_mask_cell_is_included(self): + """A NaN mask cell counts as included (truthy).""" + m = Matrix(1, 3, [1.0, 2.0, 3.0]) + mask = Matrix(1, 3, [float("nan"), 0.0, 1.0]) + assert m.sum(where=mask) == pytest.approx(1.0 + 3.0) + + def test_axis1_rowwise(self): + """where= masks per-row reductions along axis=1.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0]) + out = m.sum(axis=1, where=mask) + assert out.rows == 2 + assert out.columns == 1 + assert list(out.values()) == pytest.approx([1.0 + 3.0, 5.0 + 6.0]) + + def test_axis0_columnwise(self): + """where= masks per-column reductions along axis=0.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0]) + out = m.sum(axis=0, where=mask) + assert out.rows == 1 + assert out.columns == 3 + assert list(out.values()) == pytest.approx([1.0, 5.0, 3.0 + 6.0]) + + def test_all_excluded_sum_is_zero(self): + """An all-masked additive group collapses to 0; mean yields NaN.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + zero = Matrix(2, 3, 0.0) + assert m.sum(where=zero) == pytest.approx(0.0) + assert m.magnitude(where=zero) == pytest.approx(0.0) + assert m.magnitude_squared(where=zero) == pytest.approx(0.0) + # mean of an empty group is NaN (NumPy semantics), distinguishing it + # from a genuine zero mean. + assert math.isnan(m.mean(where=zero)) + + def test_all_excluded_mean_axis_is_nan(self): + """A fully-masked group yields NaN for mean along an axis.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 1.0, 1.0, 0.0, 0.0, 0.0]) + out = m.mean(axis=1, where=mask) + values = list(out.values()) + assert values[0] == pytest.approx((1.0 + 2.0 + 3.0) / 3.0) + assert math.isnan(values[1]) + + def test_all_excluded_min_max_is_nan(self): + """An all-masked min/max group yields NaN.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + zero = Matrix(2, 3, 0.0) + assert math.isnan(m.min(where=zero)) + assert math.isnan(m.max(where=zero)) + + def test_all_excluded_row_is_nan(self): + """A fully-masked row yields NaN for min/max along axis=1.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 1.0, 1.0, 0.0, 0.0, 0.0]) + out = m.min(axis=1, where=mask) + values = list(out.values()) + assert values[0] == pytest.approx(1.0) + assert math.isnan(values[1]) + + def test_no_mask_matches_unmasked(self, mat, shape): + """where=None (and an all-ones mask) reproduce the unmasked result.""" + rows, cols = shape + ones = Matrix(rows, cols, 1.0) + assert mat.sum(where=ones) == pytest.approx(mat.sum()) + assert mat.mean(where=ones) == pytest.approx(mat.mean()) + assert mat.min(where=ones) == pytest.approx(mat.min()) + assert mat.max(where=ones) == pytest.approx(mat.max()) + + def test_shape_mismatch_raises(self): + """A mask whose shape differs from the matrix raises ValueError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(ValueError): + m.sum(where=Matrix(3, 2, 1.0)) + + def test_non_matrix_mask_raises(self): + """A non-Matrix where= argument raises TypeError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(TypeError): + m.sum(where=[1, 0, 1, 0, 1, 0]) + + def test_unacquired_mask_raises(self): + """A cown-resident mask raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + mask = Matrix(2, 2, 1.0) + Cown(mask) + with pytest.raises(RuntimeError): + m.sum(where=mask) + + def test_fuzzed_against_reference(self, mat, shape, rng): + """Masked sum/mean/min/max match a pure-Python reference.""" + rows, cols = shape + flat = list(mat.values()) + mask_flat = [1.0 if rng.random() < 0.6 else 0.0 + for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_flat) + included = [v for v, k in zip(flat, mask_flat) if k != 0.0] + if included: + assert mat.sum(where=mask) == pytest.approx(sum(included)) + assert mat.mean(where=mask) == pytest.approx( + sum(included) / len(included)) + assert mat.min(where=mask) == pytest.approx(min(included)) + assert mat.max(where=mask) == pytest.approx(max(included)) + else: + assert mat.sum(where=mask) == pytest.approx(0.0) + assert math.isnan(mat.min(where=mask)) + assert math.isnan(mat.max(where=mask)) + + def test_boc_roundtrip(self): + """A masked aggregate runs inside a @when behavior over cowns.""" + a = Cown(Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) + b = Cown(Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 0.0])) + + @when(a, b) + def result(a, b): # noqa: D401 — short behavior + """Masked sum inside a behavior.""" + return a.value.sum(where=b.value) + + wait() + assert result.exception is False + assert result.value == pytest.approx(1.0 + 3.0 + 5.0) + + +VECTOR_LENGTHS = [1, 3, 5, 10, 32] + + +class TestVectorLen: + """len() returns the number of rows for all matrices, including vectors.""" + + @pytest.mark.parametrize("n", VECTOR_LENGTHS) + def test_row_vector_len(self, n): + """len() of a 1xN row vector returns 1 (the number of rows).""" + v = Matrix(1, n, [float(i) for i in range(n)]) + assert len(v) == 1 + + @pytest.mark.parametrize("n", VECTOR_LENGTHS) + def test_column_vector_len(self, n): + """len() of an Nx1 column vector returns N (the number of rows).""" + v = Matrix(n, 1, [float(i) for i in range(n)]) + assert len(v) == n + + def test_matrix_len_returns_rows(self): + """len() of a non-vector matrix returns rows.""" + m = Matrix(3, 4) + assert len(m) == 3 + + +class TestVectorItemAccess: + """Integer indexing and iteration on vectors.""" + + @pytest.mark.parametrize("n", VECTOR_LENGTHS) + def test_row_vector_item(self, n): + """Indexing a 1xN row vector with [0] returns the row as a Matrix.""" + vals = [float(i) * 1.5 for i in range(n)] + v = Matrix(1, n, vals) + row = v[0] + if n == 1: + assert isinstance(row, float) + assert row == pytest.approx(vals[0]) + else: + assert isinstance(row, Matrix) + assert row.rows == 1 + assert row.columns == n + for i in range(n): + assert v[0, i] == pytest.approx(vals[i]) + + @pytest.mark.parametrize("n", VECTOR_LENGTHS) + def test_column_vector_item(self, n): + """Indexing an Nx1 column vector with an integer returns the element as a float.""" + vals = [float(i) * 2.0 for i in range(n)] + v = Matrix(n, 1, vals) + for i in range(n): + item = v[i] + assert isinstance(item, float) + assert item == pytest.approx(vals[i]) + + @pytest.mark.parametrize("n", [3, 5, 10]) + def test_row_vector_iteration(self, n): + """Iterating a 1xN row vector yields one Matrix (the single row).""" + vals = [float(i) for i in range(n)] + v = Matrix(1, n, vals) + collected = list(v) + assert len(collected) == 1 + assert isinstance(collected[0], Matrix) + assert collected[0].rows == 1 + assert collected[0].columns == n + + @pytest.mark.parametrize("n", [3, 5, 10]) + def test_column_vector_iteration(self, n): + """Iterating an Nx1 column vector yields individual float elements.""" + vals = [float(i) for i in range(n)] + v = Matrix(n, 1, vals) + collected = list(v) + assert len(collected) == n + for got, expected in zip(collected, vals): + assert isinstance(got, float) + assert got == pytest.approx(expected) + + +class TestRowVectorIndexing: + """Subscript on a 1-row matrix behaves like any other matrix.""" + + def test_row_vector_two_index_reads_element(self): + """v[0, i] on a row vector reads column i.""" + v = Matrix(1, 5, [10.0, 20.0, 30.0, 40.0, 50.0]) + assert v[0, 2] == pytest.approx(30.0) + assert v[0, 0] == pytest.approx(10.0) + assert v[0, 4] == pytest.approx(50.0) + + def test_row_vector_row_slice(self): + """v[0:1] on a 1-row matrix returns the whole row.""" + v = Matrix(1, 6, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + sub = v[0:1] + assert sub.rows == 1 + assert sub.columns == 6 + assert sub[0, 0] == pytest.approx(1.0) + assert sub[0, 5] == pytest.approx(6.0) + + +class TestVectorBroadcastArithmetic: + """Arithmetic broadcasting with row and column vectors.""" + + @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) + def test_row_vector_add_broadcast(self, rows, cols): + """A (1xcols) row vector added to (rowsxcols) broadcasts across rows.""" + m_vals = [float(i) for i in range(rows * cols)] + m = Matrix(rows, cols, m_vals) + v_vals = [float(j) * 100 for j in range(cols)] + v = Matrix(1, cols, v_vals) + result = m + v + assert result.rows == rows + assert result.columns == cols + for i in range(rows): + for j in range(cols): + expected = m_vals[i * cols + j] + v_vals[j] + assert result[i, j] == pytest.approx(expected) + + @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) + def test_row_vector_subtract_broadcast(self, rows, cols): + """Subtraction broadcasts a row vector across all rows.""" + m = Matrix.uniform(0.0, 10.0, size=(rows, cols)) + v = Matrix.uniform(0.0, 5.0, size=(1, cols)) + result = m - v + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(m[i, j] - v[0, j]) + + @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) + def test_row_vector_multiply_broadcast(self, rows, cols): + """Multiplication broadcasts a row vector across all rows.""" + m = Matrix.uniform(1.0, 10.0, size=(rows, cols)) + v = Matrix.uniform(1.0, 5.0, size=(1, cols)) + result = m * v + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(m[i, j] * v[0, j]) + + @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) + def test_row_vector_divide_broadcast(self, rows, cols): + """Division broadcasts a row vector across all rows.""" + m = Matrix.uniform(1.0, 10.0, size=(rows, cols)) + v = Matrix.uniform(1.0, 5.0, size=(1, cols)) + result = m / v + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(m[i, j] / v[0, j]) + + @pytest.mark.parametrize("rows,cols", [(3, 4), (5, 3), (10, 10)]) + def test_column_vector_add_broadcast(self, rows, cols): + """An (rowsx1) column vector added to (rowsxcols) broadcasts across columns.""" + m_vals = [float(i) for i in range(rows * cols)] + m = Matrix(rows, cols, m_vals) + v_vals = [float(i) * 100 for i in range(rows)] + v = Matrix(rows, 1, v_vals) result = m + v assert result.rows == rows assert result.columns == cols @@ -4545,27 +5775,417 @@ def test_clip_no_change(self, shape): c = m.clip(0.0, 10.0) assert Matrix.allclose(c, m) - def test_clip_all_below(self, shape): - """All values below min are clamped to min.""" - rows, cols = shape - m = Matrix(rows, cols, -5.0) - c = m.clip(0.0, 10.0) - expected = Matrix(rows, cols, 0.0) - assert Matrix.allclose(c, expected) + def test_clip_all_below(self, shape): + """All values below min are clamped to min.""" + rows, cols = shape + m = Matrix(rows, cols, -5.0) + c = m.clip(0.0, 10.0) + expected = Matrix(rows, cols, 0.0) + assert Matrix.allclose(c, expected) + + def test_clip_all_above(self, shape): + """All values above max are clamped to max.""" + rows, cols = shape + m = Matrix(rows, cols, 20.0) + c = m.clip(0.0, 10.0) + expected = Matrix(rows, cols, 10.0) + assert Matrix.allclose(c, expected) + + def test_clip_invalid_range(self): + """clip() raises AssertionError when max < min.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(AssertionError): + m.clip(10.0, 0.0) + + def test_clip_in_place_returns_self(self): + """clip(in_place=True) clamps self and returns it.""" + m = Matrix(2, 2, [-5.0, 0.0, 5.0, 20.0]) + result = m.clip(0.0, 10.0, in_place=True) + assert result is m + assert Matrix.allclose(m, Matrix(2, 2, [0.0, 0.0, 5.0, 10.0])) + + def test_clip_in_place_min_only(self): + """clip(min=..., in_place=True) clamps only below, in place.""" + m = Matrix(1, 3, [-2.0, 1.0, 3.0]) + result = m.clip(min=0.0, in_place=True) + assert result is m + assert Matrix.allclose(m, Matrix(1, 3, [0.0, 1.0, 3.0])) + + def test_clip_out_writes_and_returns_target(self): + """clip(out=...) writes the result and returns the target object.""" + m = Matrix(2, 2, [-5.0, 0.0, 5.0, 20.0]) + out = Matrix(2, 2, 0.0) + result = m.clip(0.0, 10.0, out=out) + assert result is out + assert Matrix.allclose(out, Matrix(2, 2, [0.0, 0.0, 5.0, 10.0])) + # self is untouched (out is a distinct buffer). + assert Matrix.allclose(m, Matrix(2, 2, [-5.0, 0.0, 5.0, 20.0])) + + def test_clip_out_and_in_place_mutually_exclusive(self): + """Passing both out= and in_place raises ValueError.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = Matrix(2, 2, 0.0) + with pytest.raises(ValueError, + match=r"out and in_place are mutually exclusive"): + m.clip(0.0, 10.0, in_place=True, out=out) + + def test_clip_in_place_and_out_are_keyword_only(self): + """in_place and out cannot be passed positionally after min/max.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + m.clip(0.0, 10.0, True) + + def test_clip_out_shape_mismatch_raises(self): + """A wrong-shape out= raises ValueError before any write.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + bad = Matrix(3, 3, 9.0) + with pytest.raises(ValueError, + match=r"out shape 3x3 does not match result 2x2"): + m.clip(0.0, 10.0, out=bad) + assert Matrix.allclose(bad, Matrix(3, 3, 9.0)) + + def test_clip_out_wrong_type_raises(self): + """A non-Matrix out= raises TypeError.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError, match=r"out must be a Matrix"): + m.clip(0.0, 10.0, out=[0.0, 0.0, 0.0, 0.0]) + + def test_clip_out_on_unacquired_cown_raises(self): + """An out= target resident in a cown raises RuntimeError.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = Matrix(2, 2, 0.0) + Cown(out) + with pytest.raises(RuntimeError): + m.clip(0.0, 10.0, out=out) + + def test_clip_in_place_boc_roundtrip(self): + """clip(in_place=True) mutates the matrix held by a cown.""" + v = Cown(Matrix(2, 2, [-5.0, 0.0, 5.0, 20.0])) + + @when(v) + def result(v): # noqa: D401 — short behavior + """Clamp the cown-resident matrix in place.""" + v.value.clip(0.0, 10.0, in_place=True) + return list(v.value.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([0.0, 0.0, 5.0, 10.0]) + + +def ref_digitize(v, edges, right): + """Reference: index of the edge interval ``v`` falls in (NaN -> -1).""" + if math.isnan(v): + return -1.0 + if right: + return float(sum(1 for e in edges if e < v)) + return float(sum(1 for e in edges if e <= v)) + + +def ref_bin(v, lo, hi, n, right): + """Reference: equal-width bin index of ``v`` over ``[lo, hi]`` (NaN -> -1).""" + if math.isnan(v): + return -1.0 + rng = hi - lo + if rng <= 0.0: + return 0.0 + pos = (v - lo) * (n / rng) + idx = math.ceil(pos) - 1 if right else math.floor(pos) + return float(max(0, min(n - 1, idx))) + + +class TestDigitize: + """Tests for the digitize() method.""" + + @pytest.mark.parametrize("right", [False, True], ids=["left", "right"]) + def test_matches_reference(self, mat, random_values, right): + """digitize agrees cell-for-cell with the pure-Python oracle.""" + edges = [-50.0, 0.0, 50.0] + got = flatten(mat.digitize(edges, right=right)) + expected = [ref_digitize(v, edges, right) for v in random_values] + assert got == pytest.approx(expected) + + def test_left_vs_right_boundary(self): + """A value on an edge goes to the upper bin (left) or lower (right).""" + m = Matrix.vector([1.0, 2.0, 3.0]) + edges = [1.0, 2.0, 3.0] + assert flatten(m.digitize(edges)) == pytest.approx([1.0, 2.0, 3.0]) + assert flatten(m.digitize(edges, right=True)) == \ + pytest.approx([0.0, 1.0, 2.0]) + + def test_below_and_above_range(self): + """Values below the first / above the last edge get 0 / len(edges).""" + m = Matrix.vector([-10.0, 100.0]) + assert flatten(m.digitize([0.0, 1.0, 2.0])) == pytest.approx([0.0, 3.0]) + + def test_accepts_matrix_vector(self): + """edges may be a Matrix vector as well as a plain sequence.""" + m = Matrix.vector([0.5, 1.5, 2.5]) + edges = Matrix.vector([1.0, 2.0]) + assert flatten(m.digitize(edges)) == pytest.approx([0.0, 1.0, 2.0]) + + def test_column_vector_edges(self): + """An Mx1 edge vector reads as the same flat boundary list.""" + m = Matrix.vector([0.5, 1.5, 2.5]) + edges = Matrix(2, 1, [1.0, 2.0]) + assert flatten(m.digitize(edges)) == pytest.approx([0.0, 1.0, 2.0]) + + def test_nan_maps_to_minus_one(self): + """A NaN element maps to -1; neighbours are unaffected.""" + m = Matrix.vector([1.0, float("nan"), 3.0]) + _assert_values_equal(flatten(m.digitize([2.0])), [0.0, -1.0, 1.0]) + + def test_not_a_vector_raises(self): + """A 2-D edges matrix raises rather than being flattened.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"edges must be a vector"): + m.digitize(Matrix(2, 2, [1.0, 2.0, 3.0, 4.0])) + + @pytest.mark.parametrize("bad", [[2.0, 1.0], [1.0, 1.0], + [1.0, float("nan")]]) + def test_non_monotonic_or_nan_edges_raise(self, bad): + """Edges must be strictly increasing and NaN-free.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"strictly increasing"): + m.digitize(bad) + + def test_in_place_returns_self(self): + """digitize(in_place=True) rewrites self and returns it.""" + m = Matrix.vector([0.5, 1.5, 2.5]) + result = m.digitize([1.0, 2.0], in_place=True) + assert result is m + assert flatten(m) == pytest.approx([0.0, 1.0, 2.0]) + + def test_in_place_self_as_edges(self): + """m.digitize(m, in_place=True) reads a snapshot of its own edges.""" + m = Matrix.vector([1.0, 2.0, 3.0]) + result = m.digitize(m, in_place=True) + assert result is m + # edges == [1,2,3]; original values [1,2,3] -> counts of edges <= v. + assert flatten(m) == pytest.approx([1.0, 2.0, 3.0]) + + def test_out_writes_and_returns_target(self): + """digitize(out=...) writes into and returns the target; self intact.""" + m = Matrix.vector([0.5, 1.5, 2.5]) + out = Matrix(1, 3, 0.0) + result = m.digitize([1.0, 2.0], out=out) + assert result is out + assert flatten(out) == pytest.approx([0.0, 1.0, 2.0]) + assert flatten(m) == pytest.approx([0.5, 1.5, 2.5]) + + def test_out_and_in_place_mutually_exclusive(self): + """Passing both out= and in_place raises ValueError.""" + m = Matrix.vector([0.5, 1.5]) + out = Matrix(1, 2, 0.0) + with pytest.raises(ValueError, + match=r"out and in_place are mutually exclusive"): + m.digitize([1.0], in_place=True, out=out) + + def test_right_in_place_out_are_keyword_only(self): + """right/in_place/out cannot be passed positionally after edges.""" + m = Matrix.vector([0.5, 1.5]) + with pytest.raises(TypeError): + m.digitize([1.0], True) + + def test_out_on_unacquired_cown_raises(self): + """An out= target resident in a cown raises RuntimeError.""" + m = Matrix.vector([0.5, 1.5]) + out = Matrix(1, 2, 0.0) + Cown(out) + with pytest.raises(RuntimeError): + m.digitize([1.0], out=out) + + def test_boc_roundtrip(self): + """digitize(in_place=True) mutates a cown-resident matrix.""" + v = Cown(Matrix.vector([0.5, 1.5, 2.5])) + + @when(v) + def result(v): # noqa: D401 — short behavior + """Digitize the cown-resident matrix in place.""" + v.value.digitize([1.0, 2.0], in_place=True) + return list(v.value.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([0.0, 1.0, 2.0]) + + +class TestBin: + """Tests for the bin() method.""" + + @pytest.mark.parametrize("bins", [1, 2, 5, 7]) + @pytest.mark.parametrize("right", [False, True], ids=["left", "right"]) + def test_matches_reference(self, random_values, shape, bins, right): + """bin agrees cell-for-cell with the pure-Python oracle.""" + rows, cols = shape + m = Matrix(rows, cols, random_values) + lo = min(random_values) + hi = max(random_values) + got = flatten(m.bin(bins, right=right)) + expected = [ref_bin(v, lo, hi, bins, right) for v in random_values] + assert got == pytest.approx(expected) + + def test_indices_within_range(self, mat, shape, random_values): + """Every bin index lies in [0, bins-1].""" + bins = 4 + got = flatten(mat.bin(bins)) + assert all(0.0 <= x <= bins - 1 for x in got) + + def test_equal_width_split(self): + """bin over an even range splits values into equal-width bins.""" + m = Matrix.vector([0.0, 1.0, 2.0, 3.0, 4.0]) + assert flatten(m.bin(2)) == pytest.approx([0.0, 0.0, 1.0, 1.0, 1.0]) + assert flatten(m.bin(4)) == pytest.approx([0.0, 1.0, 2.0, 3.0, 3.0]) + + def test_max_lands_in_last_bin(self): + """A value equal to the max is clamped into the final bin.""" + m = Matrix.vector([0.0, 10.0]) + assert flatten(m.bin(5)) == pytest.approx([0.0, 4.0]) + + def test_right_boundary(self): + """right=True puts interior-boundary values in the lower bin.""" + m = Matrix.vector([0.0, 1.0, 2.0, 3.0, 4.0]) + assert flatten(m.bin(4, right=True)) == \ + pytest.approx([0.0, 0.0, 1.0, 2.0, 3.0]) + + def test_degenerate_range_all_zero(self): + """When every value is equal, all elements map to bin 0.""" + m = Matrix.vector([5.0, 5.0, 5.0]) + assert flatten(m.bin(3)) == pytest.approx([0.0, 0.0, 0.0]) + + @pytest.mark.parametrize("right", [False, True], ids=["left", "right"]) + def test_fixed_bounds_matches_reference(self, random_values, shape, right): + """bin(bounds=) bins against the fixed range, not the data range.""" + rows, cols = shape + m = Matrix(rows, cols, random_values) + lo, hi, bins = -100.0, 100.0, 5 + got = flatten(m.bin(bins, bounds=(lo, hi), right=right)) + expected = [ref_bin(v, lo, hi, bins, right) for v in random_values] + assert got == pytest.approx(expected) + + def test_fixed_bounds_differ_from_data_range(self): + """A wider fixed range shifts values away from the extreme bins.""" + m = Matrix.vector([0.0, 5.0, 10.0]) + # Data range [0, 10] with 2 bins would give [0, 1, 1]; the fixed range + # [0, 100] puts all three values in the lowest bin. + assert flatten(m.bin(2, bounds=(0.0, 100.0))) == \ + pytest.approx([0.0, 0.0, 0.0]) + + def test_values_below_bounds_clamp_to_first_bin(self): + """Values below the fixed min clamp into bin 0.""" + m = Matrix.vector([-5.0, 0.0, 5.0]) + assert flatten(m.bin(2, bounds=(0.0, 10.0))) == \ + pytest.approx([0.0, 0.0, 1.0]) + + def test_values_above_bounds_clamp_to_last_bin(self): + """Values above the fixed max clamp into the final bin.""" + m = Matrix.vector([2.0, 10.0, 50.0]) + assert flatten(m.bin(2, bounds=(0.0, 10.0))) == \ + pytest.approx([0.0, 1.0, 1.0]) + + def test_fixed_bounds_nan_maps_to_minus_one(self): + """NaN still maps to -1 with a fixed range.""" + m = Matrix.vector([0.0, float("nan"), 10.0]) + _assert_values_equal(flatten(m.bin(2, bounds=(0.0, 10.0))), + [0.0, -1.0, 1.0]) + + def test_bounds_max_less_than_min_raises(self): + """A bounds pair with max below min raises ValueError.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"bounds max must be >= min"): + m.bin(2, bounds=(10.0, 0.0)) + + def test_bounds_is_keyword_only(self): + """bounds cannot be passed positionally after bins.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(TypeError): + m.bin(2, (0.0, 10.0)) + + def test_bounds_wrong_length_raises(self): + """A bounds pair that is not length 2 raises ValueError.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"length 2"): + m.bin(2, bounds=(0.0, 1.0, 2.0)) + + def test_bounds_wrong_type_raises(self): + """Non-numeric bounds values raise TypeError.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(TypeError, match=r"bounds values must be real"): + m.bin(2, bounds=("lo", "hi")) + + def test_nan_maps_to_minus_one(self): + """A NaN element maps to -1 and is excluded from the range.""" + m = Matrix.vector([0.0, float("nan"), 4.0]) + _assert_values_equal(flatten(m.bin(2)), [0.0, -1.0, 1.0]) + + @pytest.mark.parametrize("bad", [float("nan"), float("inf"), + float("-inf")]) + def test_bounds_non_finite_raises(self, bad): + """A non-finite bounds value raises ValueError.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"bounds values must be finite"): + m.bin(2, bounds=(0.0, bad)) + with pytest.raises(ValueError, match=r"bounds values must be finite"): + m.bin(2, bounds=(bad, 10.0)) + + @pytest.mark.parametrize("bad", [0, -1]) + def test_bins_below_one_raises(self, bad): + """bins must be >= 1.""" + m = Matrix.vector([1.0, 2.0]) + with pytest.raises(ValueError, match=r"bins must be >= 1"): + m.bin(bad) + + def test_in_place_returns_self(self): + """bin(in_place=True) rewrites self and returns it.""" + m = Matrix.vector([0.0, 5.0, 10.0]) + result = m.bin(2, in_place=True) + assert result is m + assert flatten(m) == pytest.approx([0.0, 1.0, 1.0]) + + def test_out_writes_and_returns_target(self): + """bin(out=...) writes into and returns the target; self intact.""" + m = Matrix.vector([0.0, 5.0, 10.0]) + out = Matrix(1, 3, 0.0) + result = m.bin(2, out=out) + assert result is out + assert flatten(out) == pytest.approx([0.0, 1.0, 1.0]) + assert flatten(m) == pytest.approx([0.0, 5.0, 10.0]) + + def test_out_and_in_place_mutually_exclusive(self): + """Passing both out= and in_place raises ValueError.""" + m = Matrix.vector([0.0, 5.0]) + out = Matrix(1, 2, 0.0) + with pytest.raises(ValueError, + match=r"out and in_place are mutually exclusive"): + m.bin(2, in_place=True, out=out) + + def test_right_in_place_out_are_keyword_only(self): + """right/in_place/out cannot be passed positionally after bins.""" + m = Matrix.vector([0.0, 5.0]) + with pytest.raises(TypeError): + m.bin(2, True) - def test_clip_all_above(self, shape): - """All values above max are clamped to max.""" - rows, cols = shape - m = Matrix(rows, cols, 20.0) - c = m.clip(0.0, 10.0) - expected = Matrix(rows, cols, 10.0) - assert Matrix.allclose(c, expected) + def test_out_on_unacquired_cown_raises(self): + """An out= target resident in a cown raises RuntimeError.""" + m = Matrix.vector([0.0, 5.0]) + out = Matrix(1, 2, 0.0) + Cown(out) + with pytest.raises(RuntimeError): + m.bin(2, out=out) - def test_clip_invalid_range(self): - """clip() raises AssertionError when max < min.""" - m = Matrix(2, 2, 1.0) - with pytest.raises(AssertionError): - m.clip(10.0, 0.0) + def test_boc_roundtrip(self): + """bin(in_place=True) mutates a cown-resident matrix.""" + v = Cown(Matrix.vector([0.0, 5.0, 10.0])) + + @when(v) + def result(v): # noqa: D401 — short behavior + """Bin the cown-resident matrix in place.""" + v.value.bin(2, in_place=True) + return list(v.value.values()) + + wait() + assert result.exception is False + assert result.value == pytest.approx([0.0, 1.0, 1.0]) class TestCopy: @@ -4979,26 +6599,48 @@ def test_argmax_ties_first_occurrence(self): assert m.argmax() == 0 def test_argmin_axis0_golden(self): - """argmin(axis=0) returns per-column row indices as a 1xcols matrix.""" + """argmin(axis=0) returns per-column row indices as a list of ints.""" m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) result = m.argmin(axis=0) - assert result.rows == 1 - assert result.columns == 3 - assert [result[0, c] for c in range(3)] == [1.0, 0.0, 0.0] + assert result == [1, 0, 0] + assert all(isinstance(i, int) for i in result) def test_argmax_axis1_golden(self): - """argmax(axis=1) returns per-row column indices as a rowsx1 matrix.""" + """argmax(axis=1) returns per-row column indices as a list of ints.""" m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) result = m.argmax(axis=1) - assert result.rows == 2 - assert result.columns == 1 - assert [result[r, 0] for r in range(2)] == [0.0, 1.0] + assert result == [0, 1] + assert all(isinstance(i, int) for i in result) def test_argmin_negative_axis(self): """axis=-1 behaves like axis=1 and axis=-2 like axis=0.""" m = Matrix(3, 2, [4.0, 1.0, 2.0, 8.0, 7.0, 3.0]) - assert Matrix.allclose(m.argmin(axis=-1), m.argmin(axis=1)) - assert Matrix.allclose(m.argmin(axis=-2), m.argmin(axis=0)) + assert m.argmin(axis=-1) == m.argmin(axis=1) + assert m.argmin(axis=-2) == m.argmin(axis=0) + + def test_argmin_axis0_as_matrix(self): + """argmin(axis=0, as_matrix=True) returns a Matrix vector of indices.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + result = m.argmin(axis=0, as_matrix=True) + assert isinstance(result, Matrix) + assert result.rows == 1 + assert result.columns == 3 + assert list(result.values()) == pytest.approx([1.0, 0.0, 0.0]) + + def test_argmax_axis1_as_matrix(self): + """argmax(axis=1, as_matrix=True) returns a Matrix vector of indices.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + result = m.argmax(axis=1, as_matrix=True) + assert isinstance(result, Matrix) + assert result.rows == 2 + assert result.columns == 1 + assert list(result.values()) == pytest.approx([0.0, 1.0]) + + def test_argmin_axis_none_ignores_as_matrix(self): + """axis=None always returns a single int regardless of as_matrix.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + assert m.argmin(as_matrix=True) == 3 + assert isinstance(m.argmin(as_matrix=True), int) @pytest.mark.parametrize("want_max", [False, True]) def test_argextreme_no_axis_fuzz(self, mat, shape, random_values, want_max): @@ -5011,22 +6653,20 @@ def test_argextreme_axis0_fuzz(self, mat, shape, random_values, want_max): """Per-column arg-extreme matches a Python reference.""" rows, cols = shape result = mat.argmax(axis=0) if want_max else mat.argmin(axis=0) - assert result.rows == 1 - assert result.columns == cols + assert len(result) == cols for c in range(cols): column = [random_values[r * cols + c] for r in range(rows)] - assert result[0, c] == _flat_argextreme(column, want_max) + assert result[c] == _flat_argextreme(column, want_max) @pytest.mark.parametrize("want_max", [False, True]) def test_argextreme_axis1_fuzz(self, mat, shape, random_values, want_max): """Per-row arg-extreme matches a Python reference.""" rows, cols = shape result = mat.argmax(axis=1) if want_max else mat.argmin(axis=1) - assert result.rows == rows - assert result.columns == 1 + assert len(result) == rows for r in range(rows): row = [random_values[r * cols + c] for c in range(cols)] - assert result[r, 0] == _flat_argextreme(row, want_max) + assert result[r] == _flat_argextreme(row, want_max) def test_argmin_invalid_axis_raises(self): """An out-of-range axis raises ValueError.""" @@ -5049,6 +6689,175 @@ def test_argextreme_leading_nan_pins_result(self): assert m.argmax() == 0 +def _masked_argextreme(values, mask, want_max): + """Reference masked arg-extreme: first strict extreme among included cells. + + Mirrors the C kernel: a mask cell == 0.0 excludes the element (NaN counts + as included), the first included element seeds the running extreme, strict + comparisons keep the first occurrence on a tie, and an all-excluded group + yields the -1 "no argument" sentinel. + """ + best_i = -1 + best = None + for i, (v, k) in enumerate(zip(values, mask)): + if k == 0.0: + continue + if best_i < 0 or ((v > best) if want_max else (v < best)): + best = v + best_i = i + return best_i + + +class TestArgExtremeWhere: + """where= masking on Matrix.argmin and Matrix.argmax.""" + + def test_argmin_no_axis_masks_elements(self): + """Flat argmin only considers mask-included elements.""" + m = Matrix(1, 5, [3.0, 0.0, 9.0, 1.0, 4.0]) + mask = Matrix(1, 5, [1.0, 0.0, 1.0, 0.0, 1.0]) + # The global min (0.0) and the second-smallest (1.0) are excluded. + assert m.argmin(where=mask) == 0 + + def test_argmax_no_axis_masks_elements(self): + """Flat argmax only considers mask-included elements.""" + m = Matrix(1, 5, [3.0, 0.0, 9.0, 1.0, 4.0]) + mask = Matrix(1, 5, [1.0, 1.0, 0.0, 1.0, 1.0]) + # The global max (9.0) is excluded, leaving 4.0 at index 4. + assert m.argmax(where=mask) == 4 + + def test_ties_first_included_occurrence(self): + """A tied extreme resolves to the first included occurrence.""" + m = Matrix(1, 4, [1.0, 1.0, 1.0, 1.0]) + mask = Matrix(1, 4, [0.0, 0.0, 1.0, 1.0]) + assert m.argmin(where=mask) == 2 + assert m.argmax(where=mask) == 2 + + def test_nan_mask_cell_is_included(self): + """A NaN mask cell counts as included (truthy).""" + m = Matrix(1, 3, [5.0, 2.0, 8.0]) + mask = Matrix(1, 3, [0.0, float("nan"), 1.0]) + assert m.argmin(where=mask) == 1 + + def test_axis0_columnwise(self): + """where= masks per-column arg-reductions along axis=0.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + # Exclude row 1 of column 0 (the 0.0) so the column min becomes row 0. + mask = Matrix(2, 3, [1.0, 1.0, 1.0, 0.0, 1.0, 1.0]) + result = m.argmin(axis=0, where=mask) + assert result == [0, 0, 0] + assert all(isinstance(i, int) for i in result) + + def test_axis1_rowwise(self): + """where= masks per-row arg-reductions along axis=1.""" + m = Matrix(2, 3, [3.0, 1.0, 2.0, 0.0, 9.0, 4.0]) + # Exclude the row minima (col 1 of row 0, col 0 of row 1). + mask = Matrix(2, 3, [1.0, 0.0, 1.0, 0.0, 1.0, 1.0]) + result = m.argmax(axis=1, where=mask) + assert result == [0, 1] + + def test_all_excluded_no_axis_is_minus_one(self): + """A fully-masked flat arg-reduction yields the -1 sentinel.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + zero = Matrix(2, 3, 0.0) + assert m.argmin(where=zero) == -1 + assert m.argmax(where=zero) == -1 + + def test_all_excluded_group_is_minus_one(self): + """A fully-masked row yields -1 for that group along axis=1.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + mask = Matrix(2, 3, [1.0, 1.0, 1.0, 0.0, 0.0, 0.0]) + result = m.argmin(axis=1, where=mask) + assert result[0] == 0 + assert result[1] == -1 + + def test_no_mask_matches_unmasked(self, mat, shape): + """where=None (and an all-ones mask) reproduce the unmasked result.""" + rows, cols = shape + ones = Matrix(rows, cols, 1.0) + assert mat.argmin(where=ones) == mat.argmin() + assert mat.argmax(where=ones) == mat.argmax() + assert (mat.argmin(axis=0, where=ones) + == mat.argmin(axis=0)) + assert (mat.argmax(axis=1, where=ones) + == mat.argmax(axis=1)) + + def test_shape_mismatch_raises(self): + """A mask whose shape differs from the matrix raises ValueError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(ValueError): + m.argmin(where=Matrix(3, 2, 1.0)) + + def test_non_matrix_mask_raises(self): + """A non-Matrix where= argument raises TypeError.""" + m = Matrix(2, 3, 1.0) + with pytest.raises(TypeError): + m.argmax(where=[1, 0, 1, 0, 1, 0]) + + def test_unacquired_mask_raises(self): + """A cown-resident mask raises RuntimeError.""" + m = Matrix(2, 2, 1.0) + mask = Matrix(2, 2, 1.0) + Cown(mask) + with pytest.raises(RuntimeError): + m.argmin(where=mask) + + @pytest.mark.parametrize("want_max", [False, True]) + def test_fuzzed_no_axis(self, mat, shape, random_values, rng, want_max): + """Masked flat arg-extreme matches a pure-Python reference.""" + rows, cols = shape + mask_flat = [1.0 if rng.random() < 0.6 else 0.0 + for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_flat) + expected = _masked_argextreme(random_values, mask_flat, want_max) + result = (mat.argmax(where=mask) if want_max + else mat.argmin(where=mask)) + assert result == expected + + @pytest.mark.parametrize("want_max", [False, True]) + def test_fuzzed_axis1(self, mat, shape, random_values, rng, want_max): + """Masked per-row arg-extreme matches a pure-Python reference.""" + rows, cols = shape + mask_flat = [1.0 if rng.random() < 0.6 else 0.0 + for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_flat) + result = (mat.argmax(axis=1, where=mask) if want_max + else mat.argmin(axis=1, where=mask)) + assert len(result) == rows + for r in range(rows): + row = [random_values[r * cols + c] for c in range(cols)] + row_mask = [mask_flat[r * cols + c] for c in range(cols)] + assert result[r] == _masked_argextreme(row, row_mask, want_max) + + @pytest.mark.parametrize("want_max", [False, True]) + def test_fuzzed_axis0(self, mat, shape, random_values, rng, want_max): + """Masked per-column arg-extreme matches a pure-Python reference.""" + rows, cols = shape + mask_flat = [1.0 if rng.random() < 0.6 else 0.0 + for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_flat) + result = (mat.argmax(axis=0, where=mask) if want_max + else mat.argmin(axis=0, where=mask)) + assert len(result) == cols + for c in range(cols): + column = [random_values[r * cols + c] for r in range(rows)] + col_mask = [mask_flat[r * cols + c] for r in range(rows)] + assert result[c] == _masked_argextreme(column, col_mask, want_max) + + def test_boc_roundtrip(self): + """A masked arg-reduction runs inside a @when behavior over cowns.""" + a = Cown(Matrix(1, 5, [3.0, 0.0, 9.0, 1.0, 4.0])) + b = Cown(Matrix(1, 5, [1.0, 0.0, 1.0, 0.0, 1.0])) + + @when(a, b) + def result(a, b): # noqa: D401 — short behavior + """Masked argmin inside a behavior.""" + return a.value.argmin(where=b.value) + + wait() + assert result.exception is False + assert result.value == 0 + + def _outer_op(op, row_vals, col_vals): """Reference RxC outer broadcast: out[r,c] = op(col[r], row[c]).""" return [[op(col_vals[r], row_vals[c]) for c in range(len(row_vals))] @@ -5783,6 +7592,13 @@ def test_xyzw_getter_parametrized(self, size): assert m.z == pytest.approx(vals[2]) assert m.w == pytest.approx(vals[3]) + @pytest.mark.parametrize("comp", ["x", "y", "z", "w"]) + def test_setter_non_number_raises_type_error(self, comp): + """A non-number RHS raises TypeError (not SystemError).""" + m = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + setattr(m, comp, "nope") + class TestNegativeIndexing: """Tests for negative integer indices in __getitem__ and __setitem__. @@ -6322,12 +8138,6 @@ def check(v): assert r2 == pytest.approx(-3.0) -def _flat(m): - """Row-major flat list of a matrix's elements.""" - rows, cols = m.rows, m.columns - return [m[i, j] for i in range(rows) for j in range(cols)] - - class TestSqrt: """Element-wise square root.""" @@ -6376,6 +8186,145 @@ def test_sqrt_in_place(self, shape, rng): assert m[i, j] == pytest.approx(math.sqrt(vals[i * cols + j])) +class TestReciprocal: + """Element-wise reciprocal (1 / x).""" + + def test_reciprocal_matches(self, shape, rng): + """reciprocal() matches 1 / x element-wise for non-zero inputs.""" + rows, cols = shape + vals = [rng.uniform(-100, 100) or 1.0 for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.reciprocal() + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(1.0 / vals[i * cols + j]) + + def test_reciprocal_zero_is_inf(self): + """Zero elements yield +/-inf rather than raising.""" + m = Matrix(1, 3, [0.0, -0.0, 4.0]) + result = m.reciprocal() + assert result[0, 0] == math.inf + assert result[0, 1] == -math.inf + assert result[0, 2] == pytest.approx(0.25) + + def test_reciprocal_out_of_place_preserves_source(self, shape, rng): + """Default reciprocal() returns a new matrix and leaves source intact.""" + rows, cols = shape + vals = [rng.uniform(-100, 100) or 1.0 for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + _ = m.reciprocal() + for i in range(rows): + for j in range(cols): + assert m[i, j] == pytest.approx(vals[i * cols + j]) + + def test_reciprocal_in_place(self, shape, rng): + """reciprocal(in_place=True) mutates self and returns it.""" + rows, cols = shape + vals = [rng.uniform(-100, 100) or 1.0 for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.reciprocal(in_place=True) + assert result is m + for i in range(rows): + for j in range(cols): + assert m[i, j] == pytest.approx(1.0 / vals[i * cols + j]) + + +class TestSignCosSin: + """Element-wise sign, cosine, and sine.""" + + def test_sign_matches(self, shape, rng): + """sign() returns -1, 0, or 1 by element sign.""" + rows, cols = shape + vals = [rng.uniform(-5, 5) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.sign() + for i in range(rows): + for j in range(cols): + v = vals[i * cols + j] + assert result[i, j] == ((v > 0) - (v < 0)) + + def test_sign_zero_and_nan(self): + """sign(0) is 0 and sign(NaN) is 0.""" + m = Matrix(1, 3, [0.0, float("nan"), -0.0]) + result = m.sign() + assert result[0, 0] == 0.0 + assert result[0, 1] == 0.0 + assert result[0, 2] == 0.0 + + def test_cos_matches_math(self, shape, rng): + """cos() matches math.cos element-wise.""" + rows, cols = shape + vals = [rng.uniform(-math.pi, math.pi) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.cos() + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(math.cos(vals[i * cols + j])) + + def test_sin_matches_math(self, shape, rng): + """sin() matches math.sin element-wise.""" + rows, cols = shape + vals = [rng.uniform(-math.pi, math.pi) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.sin() + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(math.sin(vals[i * cols + j])) + + def test_sign_in_place(self, shape, rng): + """sign(in_place=True) mutates self and returns it.""" + rows, cols = shape + vals = [rng.uniform(-5, 5) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.sign(in_place=True) + assert result is m + + +class TestValues: + """Lazy row-major float stream via Matrix.values().""" + + def test_golden_row_major(self): + """values() yields elements in row-major order as floats.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + assert list(m.values()) == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + assert all(isinstance(v, float) for v in m.values()) + + def test_iterator_protocol(self): + """values() returns a self-iterable iterator.""" + m = Matrix(1, 3, [7.0, 8.0, 9.0]) + it = m.values() + assert iter(it) is it + assert next(it) == 7.0 + assert next(it) == 8.0 + assert next(it) == 9.0 + with pytest.raises(StopIteration): + next(it) + + def test_sum(self): + """values() composes with sum().""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert sum(m.values()) == 10.0 + + def test_lazy_single_advance(self): + """Each next() advances by exactly one element.""" + m = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + it = m.values() + assert next(it) == 1.0 + assert list(it) == [2.0, 3.0, 4.0] + + def test_matches_matrix_reads(self, mat, shape, random_values): + """values() agrees with m[i, j] reads in row-major order.""" + rows, cols = shape + assert list(mat.values()) == pytest.approx(random_values) + + def test_matches_indexed(self, mat, shape): + """Element-for-element agreement with subscript access.""" + rows, cols = shape + streamed = list(mat.values()) + indexed = [mat[i, j] for i in range(rows) for j in range(cols)] + assert streamed == indexed + + COMPARE_OPS = [ ("less", lambda a, b: a < b), ("less_equal", lambda a, b: a <= b), @@ -6471,9 +8420,27 @@ def test_masks_are_zero_or_one(self, shape, rng): a = Matrix(rows, cols, [rng.uniform(-9, 9) for _ in range(rows * cols)]) b = Matrix(rows, cols, [rng.uniform(-9, 9) for _ in range(rows * cols)]) for name, _ in COMPARE_OPS: - for v in _flat(getattr(a, name)(b)): + for v in flatten(getattr(a, name)(b)): assert v in (0.0, 1.0) + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_out_writes_into_target(self, name, ref, shape, rng): + """out= writes the mask into the target and returns it.""" + rows, cols = shape + a = Matrix(rows, cols, [float(rng.randint(0, 5)) for _ in range(rows * cols)]) + b = Matrix(rows, cols, [float(rng.randint(0, 5)) for _ in range(rows * cols)]) + out = Matrix.zeros(shape) + result = getattr(a, name)(b, out=out) + assert result is out + assert flatten(out) == flatten(getattr(a, name)(b)) + + def test_out_wrong_shape_raises(self): + """out= with a mismatched shape raises before any write.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = Matrix.zeros((1, 3)) + with pytest.raises(ValueError): + a.greater(2.0, out=out) + RICHCOMPARE_OPS = [ ("lt", lambda a, b: a < b), @@ -6620,19 +8587,46 @@ def test_scalar_both(self): """where accepts scalars for both a and b.""" mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) result = Matrix.where(mask, 1.0, 0.0) - assert _flat(result) == [1.0, 0.0, 0.0, 1.0] + assert flatten(result) == [1.0, 0.0, 0.0, 1.0] + + def test_out_writes_into_target(self, shape, rng): + """out= writes the result into the target and returns it.""" + rows, cols = shape + mask = Matrix(rows, cols, [float(rng.randint(0, 1)) for _ in range(rows * cols)]) + a = Matrix.full(shape, 7.0) + b = Matrix.full(shape, -7.0) + out = Matrix.zeros(shape) + result = Matrix.where(mask, a, b, out=out) + assert result is out + assert flatten(out) == flatten(Matrix.where(mask, a, b)) + + def test_out_can_alias_a(self): + """out= may alias operand a.""" + mask = Matrix(1, 4, [1.0, 0.0, 1.0, 0.0]) + a = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(1, 4, [9.0, 9.0, 9.0, 9.0]) + result = Matrix.where(mask, a, b, out=a) + assert result is a + assert flatten(a) == [1.0, 9.0, 3.0, 9.0] + + def test_out_wrong_shape_raises(self): + """out= with a mismatched shape raises before any write.""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + out = Matrix.zeros((1, 3)) + with pytest.raises(ValueError): + Matrix.where(mask, 1.0, 0.0, out=out) def test_mask_from_comparison(self): """A mask produced by a comparison method drives the selection.""" a = Matrix(1, 4, [1.0, 5.0, 2.0, 8.0]) result = Matrix.where(a.greater(3.0), a, 0.0) - assert _flat(result) == [0.0, 5.0, 0.0, 8.0] + assert flatten(result) == [0.0, 5.0, 0.0, 8.0] def test_nan_mask_selects_a(self): """A NaN mask element is non-zero and selects a.""" mask = Matrix(1, 2, [float("nan"), 0.0]) result = Matrix.where(mask, 1.0, 2.0) - assert _flat(result) == [1.0, 2.0] + assert flatten(result) == [1.0, 2.0] def test_shape_mismatch_raises(self): """A matrix operand whose shape differs from the mask raises.""" @@ -6725,14 +8719,14 @@ def test_tuple_operand_same_shape(self): """A tuple operand matching the flat shape compares element-wise.""" a = Matrix(1, 3, [1.0, 3.0, 5.0]) result = a.equal((1.0, 0.0, 5.0)) - assert _flat(result) == [1.0, 0.0, 1.0] + assert flatten(result) == [1.0, 0.0, 1.0] def test_bool_operand_is_scalar(self): """A bool operand is the scalar 1.0 (True) or 0.0 (False).""" a = Matrix(1, 3, [0.0, 1.0, 2.0]) - assert _flat(a.less(True)) == [1.0, 0.0, 0.0] - assert _flat(a.greater(False)) == [0.0, 1.0, 1.0] - assert _flat(a.equal(True)) == [0.0, 1.0, 0.0] + assert flatten(a.less(True)) == [1.0, 0.0, 0.0] + assert flatten(a.greater(False)) == [0.0, 1.0, 1.0] + assert flatten(a.equal(True)) == [0.0, 1.0, 0.0] def test_empty_list_operand_raises(self): """An empty list/tuple cannot be coerced and raises ValueError.""" @@ -6786,7 +8780,7 @@ def test_equal_mask_vs_eq_operator_nan_divergence(self): nan = float("nan") a = Matrix(1, 2, [nan, nan]) b = Matrix(1, 2, [nan, nan]) - assert _flat(a.equal(b)) == [0.0, 0.0] + assert flatten(a.equal(b)) == [0.0, 0.0] assert (a == b) is True @@ -6900,13 +8894,13 @@ def test_list_operand(self): """A list value operand is taken as a row vector matching the mask.""" mask = Matrix(1, 3, [1.0, 0.0, 1.0]) result = Matrix.where(mask, [10.0, 20.0, 30.0], 0.0) - assert _flat(result) == [10.0, 0.0, 30.0] + assert flatten(result) == [10.0, 0.0, 30.0] def test_tuple_operand(self): """A tuple value operand is coerced the same as a list.""" mask = Matrix(1, 3, [0.0, 1.0, 0.0]) result = Matrix.where(mask, 9.0, (10.0, 20.0, 30.0)) - assert _flat(result) == [10.0, 9.0, 30.0] + assert flatten(result) == [10.0, 9.0, 30.0] def test_list_shape_mismatch_raises(self): """A list whose shape differs from the mask raises ValueError.""" @@ -6918,7 +8912,7 @@ def test_bool_operand(self): """Bool value operands are scalars (True -> 1.0, False -> 0.0).""" mask = Matrix(1, 2, [1.0, 0.0]) result = Matrix.where(mask, True, False) - assert _flat(result) == [1.0, 0.0] + assert flatten(result) == [1.0, 0.0] def test_empty_list_operand_raises(self): """An empty list value operand cannot be coerced and raises ValueError.""" @@ -6997,21 +8991,21 @@ def test_row_b_and_column_c(self): def test_broadcast_in_place(self): """A broadcast operand still supports in_place=True.""" a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - av = _flat(a) + av = flatten(a) row = Matrix(1, 3, [10.0, 20.0, 30.0]) expected = [ref_fma(av[i * 3 + j], row[0, j], 0.0) for i in range(2) for j in range(3)] out = a.fma(row, 0.0, in_place=True) assert out is a - assert _flat(a) == expected + assert flatten(a) == expected def test_broadcast_leaves_operand_unmodified(self): """Materialising a broadcast operand does not mutate the source vector.""" a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) row = Matrix(1, 3, [10.0, 20.0, 30.0]) - before = _flat(row) + before = flatten(row) a.fma(row, 0.0) - assert _flat(row) == before + assert flatten(row) == before def test_wrong_length_row_vector_raises(self): """A 1xN b whose width mismatches self's columns raises ValueError.""" @@ -7026,3 +9020,987 @@ def test_wrong_length_column_vector_raises(self): bad = Matrix(3, 1, [1.0, 2.0, 3.0]) with pytest.raises(ValueError): a.fma(bad, 0.0) + + +def _grid(rows, cols): + """A rows x cols Matrix filled with 0..rows*cols-1 in row-major order.""" + return Matrix(rows, cols, [float(v) for v in range(rows * cols)]) + + +class TestMatrixViewConstruction: + """Matrix.row_view / column_view and the m.view[...] accessor.""" + + def test_row_view_full(self): + """row_view(r) is the whole row r as a 1 x N row view.""" + m = _grid(3, 4) + rv = m.row_view(1) + assert type(rv).__name__ == "MatrixView" + assert rv.is_row is True + assert rv.shape == (1, 4) + assert (rv.rows, rv.columns, rv.size) == (1, 4, 4) + assert list(rv.values()) == [4.0, 5.0, 6.0, 7.0] + + def test_row_view_partial_slice(self): + """row_view(r, cols) narrows to a column slice of the row.""" + m = _grid(3, 4) + rv = m.row_view(2, slice(1, 3)) + assert rv.shape == (1, 2) + assert list(rv.values()) == [9.0, 10.0] + + def test_row_view_negative_index(self): + """A negative row index counts from the end.""" + m = _grid(3, 4) + assert list(m.row_view(-1).values()) == [8.0, 9.0, 10.0, 11.0] + + def test_row_view_keyword_args(self): + """row_view accepts the documented ``row`` / ``columns`` keywords.""" + m = _grid(3, 4) + rv = m.row_view(row=2, columns=slice(1, 3)) + assert rv.shape == (1, 2) + assert list(rv.values()) == [9.0, 10.0] + + def test_column_view_full(self): + """column_view(c) is the whole column c as an M x 1 column view.""" + m = _grid(3, 4) + cv = m.column_view(2) + assert cv.is_row is False + assert cv.shape == (3, 1) + assert list(cv.values()) == [2.0, 6.0, 10.0] + + def test_column_view_partial_slice(self): + """column_view(c, rows) narrows to a row slice of the column.""" + m = _grid(3, 4) + cv = m.column_view(3, slice(0, 2)) + assert cv.shape == (2, 1) + assert list(cv.values()) == [3.0, 7.0] + + def test_column_view_negative_index(self): + """A negative column index counts from the end.""" + m = _grid(3, 4) + assert list(m.column_view(-1).values()) == [3.0, 7.0, 11.0] + + def test_column_view_keyword_args(self): + """column_view accepts the documented ``column`` / ``rows`` keywords.""" + m = _grid(3, 4) + cv = m.column_view(column=3, rows=slice(0, 2)) + assert cv.shape == (2, 1) + assert list(cv.values()) == [3.0, 7.0] + + def test_row_view_out_of_range(self): + """An out-of-range row index raises IndexError.""" + m = _grid(3, 4) + with pytest.raises(IndexError): + m.row_view(3) + + def test_column_view_out_of_range(self): + """An out-of-range column index raises IndexError.""" + m = _grid(3, 4) + with pytest.raises(IndexError): + m.column_view(4) + + def test_empty_slice_raises(self): + """A zero-length column slice raises IndexError (via range_read).""" + m = _grid(3, 4) + with pytest.raises(IndexError): + m.row_view(0, slice(2, 2)) + + def test_giant_step_collapses_to_single_element(self): + """A giant slice step yields a 1-element view, matching m[...]. + + ``PySlice_Unpack`` clamps a step to ``PY_SSIZE_T_MAX`` while collapsing + the count to 1, so the view must produce the single selected element + (never raise) -- consistent with the plain ``m[...]`` subscript. + """ + m = _grid(3, 4) + big = 10**30 + cv = m.column_view(1, slice(0, 3, big)) + assert cv.shape == (1, 1) + assert list(cv.values()) == [m[0, 1]] + rv = m.row_view(2, slice(0, 4, big)) + assert rv.shape == (1, 1) + assert list(rv.values()) == [m[2, 0]] + + def test_direct_construction_rejected(self): + """MatrixView / its companion types cannot be constructed directly.""" + import bocpy._math as _math + + m = _grid(2, 2) + with pytest.raises(TypeError): + _math.MatrixView() + with pytest.raises(TypeError): + type(m.view)() + with pytest.raises(TypeError): + type(m.row_view(0).values())() + + +class TestMatrixViewReads: + """Element reads, sub-views, .T, len, repr/str on a MatrixView.""" + + def test_scalar_getitem(self): + """An int key returns the element as a float.""" + rv = _grid(3, 4).row_view(1) + assert rv[0] == 4.0 + assert rv[3] == 7.0 + assert isinstance(rv[0], float) + + def test_scalar_getitem_negative(self): + """A negative element index counts from the end.""" + rv = _grid(3, 4).row_view(1) + assert rv[-1] == 7.0 + + def test_scalar_getitem_out_of_range(self): + """An out-of-range element index raises IndexError.""" + rv = _grid(3, 4).row_view(1) + with pytest.raises(IndexError): + rv[4] + + def test_tuple_key_rejected(self): + """A view is 1-D: a 2-D tuple key raises TypeError.""" + rv = _grid(3, 4).row_view(1) + with pytest.raises(TypeError): + rv[0, 1] + + def test_subview_slice(self): + """A slice key returns a narrowed sub-view over the same storage.""" + rv = _grid(3, 4).row_view(1) + sub = rv[1:3] + assert type(sub).__name__ == "MatrixView" + assert list(sub.values()) == [5.0, 6.0] + + def test_subview_reversed(self): + """A negative-step slice reverses the view.""" + rv = _grid(3, 4).row_view(1) + assert list(rv[::-1].values()) == [7.0, 6.0, 5.0, 4.0] + + def test_subview_strided(self): + """A strided slice composes stride against the root.""" + cv = _grid(4, 2).column_view(0) # [0, 2, 4, 6] + assert list(cv[::2].values()) == [0.0, 4.0] + + def test_subview_of_subview(self): + """Sub-views compose (re-based to the root, single indirection).""" + rv = _grid(1, 8).row_view(0) # 0..7 + assert list(rv[2:7][1:3].values()) == [3.0, 4.0] + + def test_subview_giant_step_single_element(self): + """A giant-step sub-slice collapses to the single first element.""" + cv = _grid(4, 2).column_view(0) # [0, 2, 4, 6], stride 2 + sub = cv[0:4:10**30] + assert list(sub.values()) == [0.0] + + def test_iter_yields_elements(self): + """``for x in view`` yields the view's elements as floats.""" + rv = _grid(3, 4).row_view(1) + assert [x for x in rv] == [4.0, 5.0, 6.0, 7.0] + assert all(isinstance(x, float) for x in rv) + + def test_iter_column_view_strided(self): + """Iterating a column view walks the strided elements in order.""" + cv = _grid(3, 4).column_view(2) + assert list(cv) == [2.0, 6.0, 10.0] + + def test_iter_matches_values(self): + """Iteration equals values() for row, column, and reversed views.""" + m = _grid(4, 3) + for v in (m.row_view(2), m.column_view(1), m.row_view(0)[::-1]): + assert list(v) == list(v.values()) + + def test_transpose_is_free_flip(self): + """.T flips orientation over the same storage (no copy).""" + rv = _grid(3, 4).row_view(1) + t = rv.T + assert t.is_row is False + assert t.shape == (4, 1) + assert list(t.values()) == list(rv.values()) + # double transpose restores orientation + assert rv.T.T.is_row is True + + def test_len(self): + """len(view) is the element count.""" + assert len(_grid(3, 4).row_view(1)) == 4 + assert len(_grid(3, 4).column_view(2)) == 3 + + def test_repr_and_str_carry_marker(self): + """repr/str render with a MatrixView( marker.""" + rv = _grid(3, 4).row_view(1, slice(1, 3)) + assert repr(rv).startswith("MatrixView(") + assert str(rv).startswith("MatrixView(") + + def test_values_iterator_protocol(self): + """values() returns a self-iterable lazy iterator.""" + rv = _grid(3, 4).row_view(1) + it = rv.values() + assert iter(it) is it + assert next(it) == 4.0 + assert list(it) == [5.0, 6.0, 7.0] + + +class TestMatrixViewAccessor: + """The m.view[...] slice-notation accessor.""" + + def test_int_key_whole_row(self): + """m.view[i] is the whole row i (== m.view[i, :]).""" + m = _grid(3, 4) + rv = m.view[2] + assert rv.is_row is True + assert list(rv.values()) == [8.0, 9.0, 10.0, 11.0] + + def test_row_int_col_slice(self): + """m.view[i, a:b] is a partial row view.""" + m = _grid(3, 4) + assert list(m.view[2, 1:3].values()) == [9.0, 10.0] + + def test_row_slice_col_int(self): + """m.view[:, j] is a whole-column view.""" + m = _grid(3, 4) + cv = m.view[:, 1] + assert cv.is_row is False + assert list(cv.values()) == [1.0, 5.0, 9.0] + + def test_strided_column(self): + """m.view[a:b:s, j] is a strided column view.""" + m = _grid(4, 3) + assert list(m.view[0:4:2, 2].values()) == [2.0, 8.0] + + def test_reversed_row(self): + """m.view[i, ::-1] reverses a row view.""" + m = _grid(3, 4) + assert list(m.view[1, ::-1].values()) == [7.0, 6.0, 5.0, 4.0] + + def test_int_int_is_one_element_row_view(self): + """m.view[i, j] is a 1-element row view (not a float).""" + m = _grid(3, 4) + e = m.view[1, 2] + assert type(e).__name__ == "MatrixView" + assert e.is_row is True + assert e.shape == (1, 1) + assert list(e.values()) == [6.0] + + def test_two_slices_rejected(self): + """m.view[a:b, c:d] is a 2-D block, not a 1-D view -> TypeError.""" + m = _grid(3, 4) + with pytest.raises(TypeError): + m.view[0:2, 0:2] + + def test_bare_slice_rejected(self): + """A bare 1-D slice is an ambiguous 2-D row range -> TypeError.""" + m = _grid(3, 4) + with pytest.raises(TypeError): + m.view[0:2] + + def test_giant_step_collapses_to_single_element(self): + """m.view[...] with a giant step yields a 1-element view, not an error.""" + m = _grid(3, 4) + big = 10**30 + assert list(m.view[1, 0:4:big].values()) == [m[1, 0]] + assert list(m.view[0:3:big, 2].values()) == [m[0, 2]] + + def test_accessor_matches_methods(self): + """The accessor agrees with the explicit row_view/column_view forms.""" + m = _grid(5, 6) + assert list(m.view[3].values()) == list(m.row_view(3).values()) + assert list(m.view[3, 1:4].values()) == list( + m.row_view(3, slice(1, 4)).values() + ) + assert list(m.view[:, 2].values()) == list(m.column_view(2).values()) + assert list(m.view[1:4, 2].values()) == list( + m.column_view(2, slice(1, 4)).values() + ) + + +class TestMatrixViewArithmetic: + """Views participate in arithmetic by materialising; results are Matrix.""" + + def test_view_plus_matrix(self): + """view + Matrix and Matrix + view both yield a fresh Matrix.""" + m = _grid(3, 4) + rv = m.row_view(1) # [4, 5, 6, 7] + ones = Matrix(1, 4, [1.0, 1.0, 1.0, 1.0]) + assert type(rv + ones).__name__ == "Matrix" + assert rv + ones == [5, 6, 7, 8] + assert ones + rv == [5, 6, 7, 8] + + def test_view_plus_view(self): + """view + view (row and column) materialises both operands.""" + m = _grid(3, 4) + assert m.row_view(1) + m.row_view(2) == [12, 14, 16, 18] + col = _grid(3, 1) + cv = _grid(3, 4).column_view(2) # [2, 6, 10] + assert cv + col == Matrix(3, 1, [2, 7, 12]) + + def test_view_scalar_ops(self): + """view scalar and scalar view work in both orders.""" + rv = _grid(3, 4).row_view(1) # [4, 5, 6, 7] + assert rv * 2 == [8, 10, 12, 14] + assert 3 * rv == [12, 15, 18, 21] + assert 10 - rv == [6, 5, 4, 3] + assert rv / 2 == [2.0, 2.5, 3.0, 3.5] + + def test_view_unary(self): + """-view and abs(view) materialise then negate/abs, yielding Matrix.""" + rv = _grid(3, 4).row_view(1) + assert type(-rv).__name__ == "Matrix" + assert -rv == [-4, -5, -6, -7] + assert abs(rv * -1) == [4, 5, 6, 7] + + def test_view_matmul(self): + """view @ Matrix and view @ view.T contract correctly.""" + rv = _grid(3, 4).row_view(1) # 1x4 [4,5,6,7] + col = Matrix(4, 1, [1.0, 2.0, 3.0, 4.0]) + assert rv @ col == [4 + 10 + 18 + 28] + assert rv @ rv.T == [16 + 25 + 36 + 49] + + def test_augmented_assign_rebinds(self): + """view += m rebinds to a fresh Matrix; the view/base is untouched.""" + m = _grid(3, 4) + x = m.row_view(1) + x += Matrix(1, 4, [1.0, 1.0, 1.0, 1.0]) + assert type(x).__name__ == "Matrix" + assert x == [5, 6, 7, 8] + # the base row is unchanged (the view never mutated storage) + assert list(m.row_view(1).values()) == [4, 5, 6, 7] + + def test_matches_equivalent_matrix(self): + """Arithmetic on a view equals the same op on an equivalent Matrix.""" + rv = _grid(3, 4).row_view(2) # [8, 9, 10, 11] + eq = Matrix(1, 4, [8.0, 9.0, 10.0, 11.0]) + assert rv * 3 + 1 == eq * 3 + 1 + cv = _grid(3, 4).column_view(1) # [1, 5, 9] + eqc = Matrix(3, 1, [1.0, 5.0, 9.0]) + assert cv - 2 == eqc - 2 + + +class TestMatrixViewCopy: + """MatrixView.copy() and Matrix.from_view() materialise an owned Matrix.""" + + def test_copy_returns_owned_matrix(self): + """copy() returns a canonical Matrix with the view's values/shape.""" + c = _grid(3, 4).row_view(1).copy() + assert type(c).__name__ == "Matrix" + assert c.shape == (1, 4) + assert c == [4, 5, 6, 7] + + def test_copy_is_independent(self): + """Mutating the copy or the base leaves the other unchanged.""" + m = _grid(3, 4) + c = m.row_view(1).copy() + c[0, 0] = 99.0 + assert list(m.row_view(1).values()) == [4, 5, 6, 7] + m[1, 1] = 55.0 + assert c == [99, 5, 6, 7] + + def test_copy_column_view(self): + """copy() of a column view is M x 1.""" + c = _grid(3, 4).column_view(2).copy() + assert c.shape == (3, 1) + assert c == Matrix(3, 1, [2, 6, 10]) + + def test_copy_strided_reversed_view(self): + """copy() materialises a strided/reversed view in order.""" + c = _grid(1, 8).row_view(0)[::-2].copy() # 7,5,3,1 + assert c == [7, 5, 3, 1] + + def test_from_view_matches_copy(self): + """Matrix.from_view(v) equals v.copy().""" + v = _grid(4, 3).column_view(1) + assert Matrix.from_view(v) == v.copy() + assert type(Matrix.from_view(v)).__name__ == "Matrix" + + def test_from_view_rejects_non_view(self): + """from_view raises TypeError for a non-MatrixView argument.""" + m = _grid(2, 2) + for bad in (m, [1.0, 2.0], 5.0): + with pytest.raises(TypeError): + Matrix.from_view(bad) + + +class TestMatrixViewFuzz: + """Fuzzed views oracled against the untouched m[i, j] copy path.""" + + def test_row_view_matches_source(self, mat, shape): + """Every row_view (full and sliced, incl. steps) matches m[r, c].""" + rng = random.Random(1234) + rows, cols = shape + for _ in range(20): + r = rng.randrange(rows) + if rng.random() < 0.5: + view = mat.row_view(r) + expected = [mat[r, c] for c in range(cols)] + else: + start = rng.randrange(cols) + stop = rng.randrange(start + 1, cols + 1) + step = rng.choice([1, 2, -1]) + sl = slice(start, stop, step) if step > 0 else slice( + stop - 1, start - 1 if start > 0 else None, step + ) + idx = list(range(*sl.indices(cols))) + if not idx: + continue + view = mat.row_view(r, sl) + expected = [mat[r, c] for c in idx] + assert list(view.values()) == pytest.approx(expected) + + def test_column_view_matches_source(self, mat, shape): + """Every column_view (full and sliced, incl. steps) matches m[r, c].""" + rng = random.Random(4321) + rows, cols = shape + for _ in range(20): + c = rng.randrange(cols) + if rng.random() < 0.5: + view = mat.column_view(c) + expected = [mat[r, c] for r in range(rows)] + else: + start = rng.randrange(rows) + stop = rng.randrange(start + 1, rows + 1) + step = rng.choice([1, 2, -1]) + sl = slice(start, stop, step) if step > 0 else slice( + stop - 1, start - 1 if start > 0 else None, step + ) + idx = list(range(*sl.indices(rows))) + if not idx: + continue + view = mat.column_view(c, sl) + expected = [mat[r, c] for r in idx] + assert list(view.values()) == pytest.approx(expected) + + +class TestMatrixViewNotShippable: + """A view aliases its base, so it cannot be pickled or sent; .copy() can.""" + + def test_pickle_dumps_raises_with_copy_hint(self): + """pickle.dumps(view) fails eagerly and names the .copy() remedy.""" + view = _grid(3, 4).row_view(1) + with pytest.raises(TypeError) as excinfo: + pickle.dumps(view) + assert ".copy()" in str(excinfo.value) + + def test_cown_rejects_view_with_copy_hint(self): + """Putting a view in a Cown fails eagerly and names .copy().""" + view = _grid(3, 4).column_view(2) + with pytest.raises(TypeError) as excinfo: + Cown(view) + assert ".copy()" in str(excinfo.value) + + def test_send_rejects_view_with_copy_hint(self): + """send()ing a view fails eagerly and names .copy().""" + view = _grid(3, 4).row_view(0) + with pytest.raises(TypeError) as excinfo: + send("view-not-shippable", view) + assert ".copy()" in str(excinfo.value) + + def test_copy_is_shippable(self): + """The .copy() escape hatch produces a picklable, owned Matrix.""" + view = _grid(3, 4).row_view(1) + restored = pickle.loads(pickle.dumps(view.copy())) + assert restored == [4, 5, 6, 7] + + +class TestMatrixViewWrite: + """Item and slice assignment write THROUGH the view's shared storage.""" + + def test_item_write_through(self): + """view[i] = x mutates the underlying base matrix.""" + m = _grid(3, 4) + v = m.row_view(1) + v[0] = 99.0 + assert m[1, 0] == 99.0 + assert list(v.values()) == [99, 5, 6, 7] + + def test_item_negative_index(self): + """A negative int key addresses from the end.""" + m = _grid(1, 4) + v = m.row_view(0) + v[-1] = 42.0 + assert m[0, 3] == 42.0 + + def test_item_out_of_range_raises(self): + """An out-of-range int key raises IndexError.""" + v = _grid(1, 4).row_view(0) + with pytest.raises(IndexError): + v[4] = 1.0 + + def test_item_non_number_raises(self): + """An int key requires a scalar RHS; a sequence raises TypeError.""" + v = _grid(1, 4).row_view(0) + with pytest.raises(TypeError): + v[0] = [1.0] + + def test_slice_scalar_broadcast(self): + """view[:] = scalar broadcasts to every selected element.""" + m = _grid(3, 4) + v = m.column_view(2) + v[:] = 7.0 + assert [m[r, 2] for r in range(3)] == [7, 7, 7] + + def test_slice_sequence_assign(self): + """view[a:b] = sequence copies element-wise.""" + m = _grid(1, 6) + v = m.row_view(0) + v[1:4] = [10, 20, 30] + assert list(v.values()) == [0, 10, 20, 30, 4, 5] + + def test_slice_from_matrix(self): + """The RHS may be a Matrix; matching is length-only.""" + m = _grid(1, 4) + v = m.row_view(0) + v[:] = Matrix(1, 4, [9, 8, 7, 6]) + assert list(v.values()) == [9, 8, 7, 6] + + def test_slice_from_view_orientation_agnostic(self): + """A column-view RHS assigns into a row view by length alone.""" + src = _grid(4, 1) # column 0,1,2,3 + m = _grid(1, 4) + v = m.row_view(0) + v[:] = src.column_view(0) + assert list(v.values()) == [0, 1, 2, 3] + + def test_length_mismatch_raises(self): + """A length mismatch on slice assignment raises ValueError.""" + v = _grid(1, 4).row_view(0) + with pytest.raises(ValueError): + v[0:2] = [1, 2, 3] + + def test_reversed_slice_assign(self): + """A negative-step slice assigns in reverse order.""" + m = _grid(1, 5) + v = m.row_view(0) + v[::-1] = [1, 2, 3, 4, 5] + assert list(v.values()) == [5, 4, 3, 2, 1] + + def test_transpose_write_through(self): + """Writing through .T mutates the same base storage.""" + m = _grid(2, 3) + t = m.row_view(0).T + t[1] = 88.0 + assert m[0, 1] == 88.0 + + def test_subview_write_through(self): + """A sub-view aliases the base; writing it mutates the base.""" + m = _grid(1, 6) + sub = m.row_view(0)[1:4] + sub[0] = 55.0 + assert m[0, 1] == 55.0 + + def test_column_view_write_through(self): + """Column-view assignment writes down the column.""" + m = _grid(3, 4) + v = m.column_view(1) + v[1] = 33.0 + assert m[1, 1] == 33.0 + + def test_del_rejected(self): + """A view cannot resize, so deletion raises TypeError.""" + v = _grid(1, 4).row_view(0) + with pytest.raises(TypeError): + del v[0] + + def test_tuple_key_rejected(self): + """A view is 1-D; a tuple key raises TypeError.""" + v = _grid(1, 4).row_view(0) + with pytest.raises(TypeError): + v[0, 1] = 1.0 + + def test_inplace_add_rebinds_not_writes(self): + """view += m rebinds the name to a fresh Matrix; the base is untouched.""" + m = _grid(1, 4) + v = m.row_view(0) + v += _grid(1, 4) + assert type(v).__name__ == "Matrix" + assert v == [0, 2, 4, 6] + assert list(m.row_view(0).values()) == [0, 1, 2, 3] + + def test_write_fuzz(self): + """Fuzzed item/slice writes match a Python list model of the base.""" + rng = random.Random(20260704) + for _ in range(200): + cols = rng.randrange(1, 9) + m = Matrix(1, cols, [float(v) for v in range(cols)]) + model = list(range(cols)) + v = m.row_view(0) + if rng.random() < 0.5: + i = rng.randrange(-cols, cols) + x = float(rng.randrange(-100, 100)) + v[i] = x + model[i] = x + else: + start = rng.randrange(cols) + stop = rng.randrange(start + 1, cols + 1) + step = rng.choice([1, 2, -1]) + sl = slice(start, stop, step) if step > 0 else slice( + stop - 1, start - 1 if start > 0 else None, step + ) + idx = list(range(*sl.indices(cols))) + if not idx: + continue + if rng.random() < 0.5: + x = float(rng.randrange(-100, 100)) + v[sl] = x + for j in idx: + model[j] = x + else: + payload = [float(rng.randrange(-100, 100)) for _ in idx] + v[sl] = payload + for j, x in zip(idx, payload): + model[j] = x + assert list(m.row_view(0).values()) == pytest.approx(model) + + +class TestMatrixViewOwnership: + """A view over a cown-resident (unacquired) base raises RuntimeError.""" + + @staticmethod + def _unowned_row_view(): + """A row view whose base was moved into a Cown (owner cleared).""" + m = _grid(3, 4) + v = m.row_view(1) + Cown(m) # clears m's owner; v still references the same impl + return v + + def test_read_values_raises(self): + """Reading values() on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + list(v.values()) + + def test_read_item_raises(self): + """An int read on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + v[0] + + def test_write_item_raises(self): + """Writing through an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + v[0] = 9.0 + + def test_write_slice_raises(self): + """Slice assignment on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + v[:] = 0.0 + + def test_arithmetic_raises(self): + """Arithmetic on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + v + _grid(1, 4) + + def test_negate_raises(self): + """Unary negation on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + -v + + def test_copy_raises(self): + """copy() on an unowned view raises RuntimeError.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + v.copy() + + def test_shape_read_raises(self): + """The (conservatively guarded) shape read raises when unowned.""" + v = self._unowned_row_view() + with pytest.raises(RuntimeError): + _ = v.shape + + def test_is_row_is_exempt(self): + """is_row is pure orientation metadata: readable even when unowned.""" + v = self._unowned_row_view() + assert v.is_row is True + + def test_values_iterator_reguards_midstream(self): + """values() re-guards every step, so losing ownership mid-scan raises.""" + m = _grid(1, 4) + v = m.row_view(0) + it = v.values() + assert next(it) == 0.0 # first element read while still owned + Cown(m) # clears the owner mid-iteration + with pytest.raises(RuntimeError): + next(it) + + +class TestMatrixViewComponents: + """MatrixView .x/.y/.z/.w mirror Matrix's convenience accessors.""" + + def test_getters_row_view(self): + """x/y/z/w read the view's first four elements.""" + v = _grid(1, 6).row_view(0) # 0,1,2,3,4,5 + assert v.x == pytest.approx(0.0) + assert v.y == pytest.approx(1.0) + assert v.z == pytest.approx(2.0) + assert v.w == pytest.approx(3.0) + + def test_getters_column_view_strided(self): + """x/y/z/w follow the view's stride down a column.""" + v = _grid(4, 3).column_view(1) # 1,4,7,10 + assert v.x == pytest.approx(1.0) + assert v.y == pytest.approx(4.0) + assert v.z == pytest.approx(7.0) + assert v.w == pytest.approx(10.0) + + def test_getters_reversed_subview(self): + """Components honour a negative-step sub-view's order.""" + v = _grid(1, 5).row_view(0)[::-1] # 4,3,2,1,0 + assert v.x == pytest.approx(4.0) + assert v.w == pytest.approx(1.0) + + def test_setters_write_through(self): + """Setting x/y/z/w writes through to the base matrix.""" + m = _grid(1, 4) + v = m.row_view(0) + v.x = 10.0 + v.y = 20.0 + v.z = 30.0 + v.w = 40.0 + assert list(m.row_view(0).values()) == [10, 20, 30, 40] + + def test_setter_column_view_write_through(self): + """A column-view setter mutates the correct base cell.""" + m = _grid(4, 3) + v = m.column_view(2) # rows: 2,5,8,11 + v.y = 99.0 + assert m[1, 2] == pytest.approx(99.0) + + def test_getter_absent_component_raises(self): + """A component past the view's length raises IndexError.""" + v = _grid(1, 2).row_view(0) + with pytest.raises(IndexError): + _ = v.z + + def test_setter_absent_component_raises(self): + """Setting a component past the view's length raises IndexError.""" + v = _grid(1, 2).row_view(0) + with pytest.raises(IndexError): + v.z = 1.0 + + def test_setter_non_number_raises_type_error(self): + """A non-number RHS raises TypeError (not SystemError).""" + v = _grid(1, 4).row_view(0) + with pytest.raises(TypeError): + v.x = "nope" + + def test_getter_unowned_raises(self): + """A component read on an unowned view raises RuntimeError.""" + m = _grid(1, 4) + v = m.row_view(0) + Cown(m) + with pytest.raises(RuntimeError): + _ = v.x + + def test_setter_unowned_raises(self): + """A component write on an unowned view raises RuntimeError.""" + m = _grid(1, 4) + v = m.row_view(0) + Cown(m) + with pytest.raises(RuntimeError): + v.x = 1.0 + + +class TestMatrixViewsIterators: + """Matrix.row_views() / column_views() yield fresh per-row/column views.""" + + def test_row_views_count_and_shape(self): + """row_views() yields one 1 x N row view per row.""" + m = _grid(3, 4) + views = list(m.row_views()) + assert len(views) == 3 + assert all(v.shape == (1, 4) for v in views) + assert all(v.is_row for v in views) + + def test_row_views_values(self): + """Each row view exposes that row's elements in order.""" + m = _grid(3, 4) + assert [list(v.values()) for v in m.row_views()] == [ + [0, 1, 2, 3], + [4, 5, 6, 7], + [8, 9, 10, 11], + ] + + def test_column_views_count_and_shape(self): + """column_views() yields one M x 1 column view per column.""" + m = _grid(3, 4) + views = list(m.column_views()) + assert len(views) == 4 + assert all(v.shape == (3, 1) for v in views) + assert all(v.is_row is False for v in views) + + def test_column_views_values(self): + """Each column view walks its column top to bottom.""" + m = _grid(3, 4) + assert [list(v.values()) for v in m.column_views()] == [ + [0, 4, 8], + [1, 5, 9], + [2, 6, 10], + [3, 7, 11], + ] + + def test_views_are_distinct_objects(self): + """Yielded views are distinct objects, not one reused cursor.""" + m = _grid(3, 4) + views = list(m.row_views()) + assert len({id(v) for v in views}) == 3 + # The first view still shows its own row after the loop completes. + assert list(views[0].values()) == [0, 1, 2, 3] + + def test_row_views_write_through(self): + """Writing through a yielded row view mutates the base matrix.""" + m = _grid(3, 4) + for v in m.row_views(): + v[0] = 99.0 + assert [m[r, 0] for r in range(3)] == [99, 99, 99] + + def test_column_views_write_through(self): + """Writing through a yielded column view mutates the base matrix.""" + m = _grid(3, 4) + for v in m.column_views(): + v[0] = 7.0 + assert [m[0, c] for c in range(4)] == [7, 7, 7, 7] + + def test_column_vector_row_views(self): + """row_views() of an M x 1 column yields M single-element rows.""" + m = Matrix(4, 1, [1.0, 2.0, 3.0, 4.0]) + assert [list(v.values()) for v in m.row_views()] == [[1], [2], [3], [4]] + + def test_row_vector_column_views(self): + """column_views() of a 1 x N row yields N single-element columns.""" + m = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + got = [list(v.values()) for v in m.column_views()] + assert got == [[1], [2], [3], [4]] + + def test_iterator_is_self_iterable(self): + """The returned iterator is its own iterator.""" + it = _grid(2, 2).row_views() + assert iter(it) is it + + def test_direct_construction_rejected(self): + """The iterator type cannot be constructed directly.""" + it = _grid(2, 2).row_views() + with pytest.raises(TypeError): + type(it)() + + def test_unowned_matrix_raises(self): + """row_views() on an unowned matrix raises RuntimeError.""" + m = _grid(3, 4) + Cown(m) + with pytest.raises(RuntimeError): + m.row_views() + + def test_reguards_midstream(self): + """Losing ownership mid-iteration raises on the next step.""" + m = _grid(3, 4) + it = m.row_views() + first = next(it) + assert list(first.values()) == [0, 1, 2, 3] + Cown(m) + with pytest.raises(RuntimeError): + next(it) + + +class TestMatrixViewReadOnly: + """The read_only flag on row_view / column_view / row_views / column_views.""" + + def test_default_is_writable(self): + """Views are read/write by default (read_only is False).""" + m = _grid(3, 4) + assert m.row_view(0).read_only is False + assert m.column_view(0).read_only is False + + def test_row_view_read_only_flag(self): + """row_view(read_only=True) sets the read_only property.""" + assert _grid(3, 4).row_view(0, read_only=True).read_only is True + + def test_column_view_read_only_flag(self): + """column_view(read_only=True) sets the read_only property.""" + assert _grid(3, 4).column_view(0, read_only=True).read_only is True + + def test_read_only_reads_still_work(self): + """A read-only view still reads elements, values(), and .x.""" + v = _grid(3, 4).row_view(1, read_only=True) + assert v[0] == pytest.approx(4.0) + assert list(v.values()) == [4, 5, 6, 7] + assert v.x == pytest.approx(4.0) + + def test_read_only_item_write_raises(self): + """Item assignment through a read-only view raises TypeError.""" + v = _grid(3, 4).row_view(0, read_only=True) + with pytest.raises(TypeError, match=r"read-only MatrixView"): + v[0] = 1.0 + + def test_read_only_slice_write_raises(self): + """Slice assignment through a read-only view raises TypeError.""" + v = _grid(3, 4).row_view(0, read_only=True) + with pytest.raises(TypeError, match=r"read-only MatrixView"): + v[:] = 0.0 + + def test_read_only_component_setter_raises(self): + """A .x/.y/.z/.w setter on a read-only view raises TypeError.""" + v = _grid(3, 4).row_view(0, read_only=True) + with pytest.raises(TypeError, match=r"read-only MatrixView"): + v.x = 1.0 + + def test_read_only_leaves_base_unchanged(self): + """A rejected write does not mutate the base matrix.""" + m = _grid(3, 4) + v = m.row_view(0, read_only=True) + with pytest.raises(TypeError): + v[0] = 999.0 + assert m[0, 0] == pytest.approx(0.0) + + def test_transpose_inherits_read_only(self): + """.T of a read-only view stays read-only and blocks writes.""" + v = _grid(3, 4).row_view(0, read_only=True) + assert v.T.read_only is True + with pytest.raises(TypeError, match=r"read-only MatrixView"): + v.T[0] = 1.0 + + def test_slice_inherits_read_only(self): + """A sub-slice of a read-only view stays read-only and blocks writes.""" + v = _grid(3, 4).row_view(0, read_only=True) + sub = v[1:3] + assert sub.read_only is True + with pytest.raises(TypeError, match=r"read-only MatrixView"): + sub[0] = 1.0 + + def test_copy_is_writable(self): + """copy() detaches a read-only view into a writable owned Matrix.""" + v = _grid(3, 4).row_view(0, read_only=True) + c = v.copy() + c[0, 0] = 7.0 # not a MatrixView; an owned Matrix + assert c[0, 0] == pytest.approx(7.0) + + def test_read_only_is_keyword_only(self): + """read_only cannot be passed positionally after columns.""" + m = _grid(3, 4) + with pytest.raises(TypeError): + m.row_view(0, None, True) + + def test_row_views_read_only(self): + """row_views(read_only=True) yields immutable row views.""" + m = _grid(3, 4) + views = list(m.row_views(read_only=True)) + assert all(v.read_only for v in views) + with pytest.raises(TypeError, match=r"read-only MatrixView"): + views[0][0] = 1.0 + + def test_column_views_read_only(self): + """column_views(read_only=True) yields immutable column views.""" + m = _grid(3, 4) + views = list(m.column_views(read_only=True)) + assert all(v.read_only for v in views) + with pytest.raises(TypeError, match=r"read-only MatrixView"): + views[0][0] = 1.0 + + def test_iterators_default_writable(self): + """The view iterators default to read/write views.""" + m = _grid(3, 4) + assert next(iter(m.row_views())).read_only is False + assert next(iter(m.column_views())).read_only is False + + def test_read_only_iteration_is_readable(self): + """A read-only row iteration still reads through to the base.""" + m = _grid(2, 3) + assert [list(v.values()) for v in m.row_views(read_only=True)] == [ + [0, 1, 2], + [3, 4, 5], + ] diff --git a/test/test_noticeboard.py b/test/test_noticeboard.py index 982b9dd..8438a7e 100644 --- a/test/test_noticeboard.py +++ b/test/test_noticeboard.py @@ -4,7 +4,7 @@ import pytest -from bocpy import (Cown, notice_delete, notice_read, +from bocpy import (Cown, Matrix, notice_delete, notice_read, notice_seed, notice_update, notice_write, noticeboard, quiesce, REMOVED, start, wait, when) import bocpy._core as _core @@ -445,6 +445,15 @@ def _conditionally_remove(x): return x + 1 +def _make_matrix_payload(current): + """Return a container holding a Matrix. + + Used to force the noticeboard mutator thread to *pickle* a Matrix + (``Matrix.__reduce__``) when it commits the update result. + """ + return {"m": Matrix.zeros((4, 3))} + + class TestNoticeUpdate: """Tests for notice_update atomic read-modify-write.""" @@ -562,6 +571,62 @@ def step1(x): assert snap.get("best") == 42 +class TestNoticeboardMatrixEntry: + """Matrix-valued noticeboard entries reconstructed on the mutator thread. + + The mutator thread is a second OS thread in the primary interpreter that + never ran the ``_math`` module init, so its cached module-state pointer + starts NULL. A ``notice_update`` read-modify-write snapshots every entry + (including Matrix-valued ones) on that thread, so Matrix reconstruction + must resolve the module state via the main-interpreter fallback rather + than dereferencing that NULL pointer. + """ + + @classmethod + def teardown_class(cls): + """Ensure runtime is drained after suite.""" + wait() + + def setup_method(self): + """Clear the noticeboard before each test.""" + _core.noticeboard_clear() + + def test_update_survives_matrix_entry_on_board(self): + """Concurrent RMW updates coexist with a pickled Matrix board entry.""" + # A pickled container of Matrices: each RMW snapshot on the + # noticeboard thread reconstructs it via Matrix.__reduce__'s + # rebuild (_matrix_unpickle) -- the path that used to segfault. + notice_seed("bulky", {i: Matrix.zeros((4, 3)) for i in range(8)}) + + n = 8 + cowns = [Cown(i) for i in range(n)] + for i in range(n): + + @when(cowns[i]) + def bump(c): + notice_update("acc", _increment, default=0) + + snap = quiesce(QUIESCE_TIMEOUT, noticeboard=True) + assert snap.get("acc") == n + bulky = snap.get("bulky") + assert len(bulky) == 8 + assert bulky[0].rows == 4 + assert bulky[0].columns == 3 + + def test_update_returning_matrix_value(self): + """An update fn returning a Matrix-bearing value pickles it safely.""" + x = Cown(0) + + @when(x) + def step(x): + notice_update("payload", _make_matrix_payload, default=None) + + snap = quiesce(QUIESCE_TIMEOUT, noticeboard=True) + payload = snap.get("payload") + assert payload["m"].rows == 4 + assert payload["m"].columns == 3 + + class TestNoticeboardReadOnly: """Tests that the snapshot is read-only (MappingProxyType).""" diff --git a/test/test_scheduling_stress.py b/test/test_scheduling_stress.py index 545c4a6..59a2c17 100644 --- a/test/test_scheduling_stress.py +++ b/test/test_scheduling_stress.py @@ -619,6 +619,7 @@ def test_set_drop_exception_marks_result_cown(self): result.impl, [(1, arg.impl)], [], + 1, ) drop = RuntimeError("orphaned during stop()") @@ -731,6 +732,7 @@ def test_schedule_after_runtime_stop_raises(self): result.impl, [(1, c.impl)], [], + 1, ) prior_count, prior_seeded = _core.terminator_reset()