ENH: add k offsets to diagonal helper functions#31493
Conversation
GoncaloTorcato
force-pushed
the
enh/add-diagonal-offset-args
branch
from
e6e3410 to
251cd1f
Compare
Add support for a k argument to diagonal helper functions so they can
operate on diagonals above or below the main diagonal. This extends
fill_diagonal, diag_indices, diag_indices_from, identity, and
ma.identity.
For example, np.fill_diagonal(a, value, k=1) now fills the first diagonal
above the main diagonal, while k=-1 targets the first diagonal below it.
Also update documentation, typing stubs, release notes, and tests for the new
offset behavior.
Closes numpy#28310.
Co-authored-by: Ricardo Curveira Santos <ricardo.curveira.santos@tecnico.ulisboa.pt>
ricas28
force-pushed
the
enh/add-diagonal-offset-args
branch
from
251cd1f to
8572463
Compare
|
Tested on macOS (Apple M-series), Python 3.14.5, numpy 2.6.0.dev0. import numpy as np # type: ignore Result: 0.0 |
seberg
left a comment
seberg
left a comment
There was a problem hiding this comment.
This may still need a bit of API decisions. fill_diagonal makes a lot of sense to me to match diagonal. identity to me doesn't make sense.
The others seem fine, although I would tend to go go with *, offset= generally (yes, existing functions use k=, but I think they usually pass it by position).
| @finalize_array_function_like | ||
| @set_module('numpy') | ||
| def identity(n, dtype=None, *, like=None): | ||
| def identity(n, dtype=None, *, k=0, like=None): |
There was a problem hiding this comment.
Personally, I am very much against adding it to identity, eye is the same and identity+offset is a misnomer.
| @array_function_dispatch(_fill_diagonal_dispatcher) | ||
| def fill_diagonal(a, val, wrap=False): | ||
| """Fill the main diagonal of the given array of any dimensionality. | ||
| def fill_diagonal(a, val, wrap=False, *, k=0): |
There was a problem hiding this comment.
Diagonal uses offset=. If we add it now and maybe also for the other functions (not currently in the PR) add it as a kwarg only, should we use that name? This needs a final API decision.
| def test_offset_out_of_bounds(self): | ||
| a = np.zeros((3, 3), int) | ||
| fill_diagonal(a, 5, k=3) | ||
| assert_array_equal(a, np.zeros((3, 3), int)) |
There was a problem hiding this comment.
I think you are missing a test for wrap=True?
5 participants
PR summary
Add support for a
kargument to diagonal helper functions so they can operate on diagonals above or below the main diagonal. This extendsfill_diagonal,diag_indices,diag_indices_from,identity, andma.identity.For example,
np.fill_diagonal(a, value, k=1)now fills the first diagonal above the main diagonal, whilek=-1targets the first diagonal below it.Also update documentation, typing stubs, release notes, and tests for the new offset behavior.
Closes #28310.
Co-authored-by: Ricardo Curveira Santos ricardo.curveira.santos@tecnico.ulisboa.pt
AI Disclosure
Yes. I used OpenAI Codex/ChatGPT to help inspect the existing NumPy code paths, update tests/docs/typing stubs.
I reviewed the generated changes, adjusted the implementation where
needed, and verified the final result locally with
spin test.