Skip to content

Commit 251cd1f

Browse files
GoncaloTorcatoGoncaloTorcato
committed
ENH: add k offsets to diagonal helper functions
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 #28310. Co-authored-by: Ricardo Curveira Santos ricardo.curveira.santos@tecnico.ulisboa.pt
1 parent bf8007a commit 251cd1f

14 files changed

Lines changed: 222 additions & 36 deletions

File tree

‎doc/release/upcoming_changes/28310.new_feature.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
``k`` offsets for diagonal helper functions
2+
-------------------------------------------
3+
`fill_diagonal`, `diag_indices`, `diag_indices_from`, `identity`, and
4+
`numpy.ma.identity` now accept a ``k`` argument to operate on diagonals above
5+
or below the main diagonal.

‎numpy/_core/numeric.py‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2176,27 +2176,30 @@ def _maketup(descr, val):
21762176

21772177
@finalize_array_function_like
21782178
@set_module('numpy')
2179-
def identity(n, dtype=None, *, like=None):
2179+
def identity(n, dtype=None, *, k=0, like=None):
21802180
"""
2181-
Return the identity array.
2181+
Return a square array with ones on a diagonal.
21822182
2183-
The identity array is a square array with ones on
2184-
the main diagonal.
2183+
When ``k=0``, this is the identity array.
21852184
21862185
Parameters
21872186
----------
21882187
n : int
21892188
Number of rows (and columns) in `n` x `n` output.
21902189
dtype : data-type, optional
21912190
Data-type of the output. Defaults to ``float``.
2191+
k : int, optional
2192+
Index of the diagonal: 0 (the default) refers to the main diagonal,
2193+
a positive value refers to an upper diagonal, and a negative value
2194+
to a lower diagonal.
21922195
${ARRAY_FUNCTION_LIKE}
21932196
21942197
.. versionadded:: 1.20.0
21952198
21962199
Returns
21972200
-------
21982201
out : ndarray
2199-
`n` x `n` array with its main diagonal set to one,
2202+
`n` x `n` array with its `k`-th diagonal set to one,
22002203
and all other elements 0.
22012204
22022205
Examples
@@ -2206,13 +2209,17 @@ def identity(n, dtype=None, *, like=None):
22062209
array([[1., 0., 0.],
22072210
[0., 1., 0.],
22082211
[0., 0., 1.]])
2212+
>>> np.identity(3, k=1)
2213+
array([[0., 1., 0.],
2214+
[0., 0., 1.],
2215+
[0., 0., 0.]])
22092216
22102217
"""
22112218
if like is not None:
2212-
return _identity_with_like(like, n, dtype=dtype)
2219+
return _identity_with_like(like, n, dtype=dtype, k=k)
22132220

22142221
from numpy import eye
2215-
return eye(n, dtype=dtype, like=like)
2222+
return eye(n, k=k, dtype=dtype, like=like)
22162223

22172224

22182225
_identity_with_like = array_function_dispatch()(identity)

‎numpy/_core/numeric.pyi‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1252,19 +1252,19 @@ def base_repr(number: SupportsAbs[float], base: float = 2, padding: SupportsInde
12521252

12531253
#
12541254
@overload # dtype: None (default)
1255-
def identity(n: int, dtype: None = None, *, like: _SupportsArrayFunc | None = None) -> _Array2D[np.float64]: ...
1255+
def identity(n: int, dtype: None = None, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[np.float64]: ...
12561256
@overload # dtype: known scalar type
1257-
def identity[ScalarT: np.generic](n: int, dtype: _DTypeLike[ScalarT], *, like: _SupportsArrayFunc | None = None) -> _Array2D[ScalarT]: ...
1257+
def identity[ScalarT: np.generic](n: int, dtype: _DTypeLike[ScalarT], *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[ScalarT]: ...
12581258
@overload # dtype: like bool
1259-
def identity(n: int, dtype: _DTypeLikeBool, *, like: _SupportsArrayFunc | None = None) -> _Array2D[np.bool]: ...
1259+
def identity(n: int, dtype: _DTypeLikeBool, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[np.bool]: ...
12601260
@overload # dtype: like int_
1261-
def identity(n: int, dtype: _DTypeLikeInt, *, like: _SupportsArrayFunc | None = None) -> _Array2D[np.int_ | Any]: ...
1261+
def identity(n: int, dtype: _DTypeLikeInt, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[np.int_ | Any]: ...
12621262
@overload # dtype: like float64
1263-
def identity(n: int, dtype: _DTypeLikeFloat64, *, like: _SupportsArrayFunc | None = None) -> _Array2D[np.float64 | Any]: ...
1263+
def identity(n: int, dtype: _DTypeLikeFloat64, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[np.float64 | Any]: ...
12641264
@overload # dtype: like complex128
1265-
def identity(n: int, dtype: _DTypeLikeComplex128, *, like: _SupportsArrayFunc | None = None) -> _Array2D[np.complex128 | Any]: ...
1265+
def identity(n: int, dtype: _DTypeLikeComplex128, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[np.complex128 | Any]: ...
12661266
@overload # dtype: unknown
1267-
def identity(n: int, dtype: DTypeLike, *, like: _SupportsArrayFunc | None = None) -> _Array2D[Incomplete]: ...
1267+
def identity(n: int, dtype: DTypeLike, *, k: int = 0, like: _SupportsArrayFunc | None = None) -> _Array2D[Incomplete]: ...
12681268

12691269
#
12701270
def allclose(

‎numpy/_core/tests/test_numeric.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3381,6 +3381,18 @@ def test_full(self):
33813381
self.check_function(np.full, 0)
33823382
self.check_function(np.full, 1)
33833383

3384+
def test_identity_offset(self):
3385+
assert_equal(np.identity(4, k=1),
3386+
np.array([[0, 1, 0, 0],
3387+
[0, 0, 1, 0],
3388+
[0, 0, 0, 1],
3389+
[0, 0, 0, 0]], dtype=float))
3390+
assert_equal(np.identity(4, dtype=int, k=-2),
3391+
np.array([[0, 0, 0, 0],
3392+
[0, 0, 0, 0],
3393+
[1, 0, 0, 0],
3394+
[0, 1, 0, 0]]))
3395+
33843396
@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
33853397
def test_for_reference_leak(self):
33863398
# Make sure we have an object for reference

‎numpy/lib/_index_tricks_impl.py‎

Lines changed: 90 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import functools
22
import math
3+
import operator
34
import sys
45
from itertools import product
56

@@ -787,17 +788,18 @@ def __getitem__(self, item):
787788
# applicable to N-dimensions.
788789

789790

790-
def _fill_diagonal_dispatcher(a, val, wrap=None):
791+
def _fill_diagonal_dispatcher(a, val, wrap=None, *, k=None):
791792
return (a,)
792793

793794

794795
@array_function_dispatch(_fill_diagonal_dispatcher)
795-
def fill_diagonal(a, val, wrap=False):
796-
"""Fill the main diagonal of the given array of any dimensionality.
796+
def fill_diagonal(a, val, wrap=False, *, k=0):
797+
"""Fill a diagonal of the given array of any dimensionality.
797798
798799
For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
799-
values ``a[i, ..., i]`` with indices ``i`` all identical. This function
800-
modifies the input array in-place without returning a value.
800+
values ``a[i, ..., i]`` with indices ``i`` all identical when ``k=0``.
801+
This function modifies the input array in-place without returning a
802+
value.
801803
802804
Parameters
803805
----------
@@ -813,6 +815,11 @@ def fill_diagonal(a, val, wrap=False):
813815
For tall matrices in NumPy version up to 1.6.2, the
814816
diagonal "wrapped" after N columns. You can have this behavior
815817
with this option. This affects only tall matrices.
818+
k : int, optional
819+
Index of the diagonal to fill. 0 (the default) refers to the main
820+
diagonal, a positive value refers to an upper diagonal, and a negative
821+
value to a lower diagonal. For arrays with more than two dimensions,
822+
the offset is applied between the first two axes.
816823
817824
See also
818825
--------
@@ -834,6 +841,15 @@ def fill_diagonal(a, val, wrap=False):
834841
[0, 5, 0],
835842
[0, 0, 5]])
836843
844+
Fill an offset diagonal:
845+
846+
>>> a = np.zeros((3, 3), int)
847+
>>> np.fill_diagonal(a, 5, k=1)
848+
>>> a
849+
array([[0, 5, 0],
850+
[0, 0, 5],
851+
[0, 0, 0]])
852+
837853
The same function can operate on a 4-D array:
838854
839855
>>> a = np.zeros((3, 3, 3, 3), int)
@@ -904,29 +920,48 @@ def fill_diagonal(a, val, wrap=False):
904920
"""
905921
if a.ndim < 2:
906922
raise ValueError("array must be at least 2-d")
923+
k = operator.index(k)
907924
end = None
908925
if a.ndim == 2:
909926
# Explicit, fast formula for the common case. For 2-d arrays, we
910927
# accept rectangular ones.
911-
step = a.shape[1] + 1
912-
# This is needed to don't have tall matrix have the diagonal wrap.
928+
n, m = a.shape
929+
if k >= m or k <= -n:
930+
return
931+
step = m + 1
932+
if k >= 0:
933+
start = k
934+
count = min(n, m - k)
935+
else:
936+
start = (-k) * m
937+
count = min(n + k, m)
913938
if not wrap:
914-
end = a.shape[1] * a.shape[1]
939+
end = start + step * count
915940
else:
916941
# For more than d=2, the strided formula is only valid for arrays with
917942
# all dimensions equal, so we check first.
918943
if not np.all(diff(a.shape) == 0):
919944
raise ValueError("All dimensions of input must be of equal length")
945+
n = a.shape[0]
946+
if k >= n or k <= -n:
947+
return
948+
if k >= 0:
949+
start = k * math.prod(a.shape[2:])
950+
count = n - k
951+
else:
952+
start = (-k) * math.prod(a.shape[1:])
953+
count = n + k
920954
step = 1 + (np.cumprod(a.shape[:-1])).sum()
955+
end = start + step * count
921956

922957
# Write the value out into the diagonal.
923-
a.flat[:end:step] = val
958+
a.flat[start:end:step] = val
924959

925960

926961
@set_module('numpy')
927-
def diag_indices(n, ndim=2):
962+
def diag_indices(n, ndim=2, k=0):
928963
"""
929-
Return the indices to access the main diagonal of an array.
964+
Return the indices to access a diagonal of an array.
930965
931966
This returns a tuple of indices that can be used to access the main
932967
diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
@@ -943,6 +978,12 @@ def diag_indices(n, ndim=2):
943978
ndim : int, optional
944979
The number of dimensions.
945980
981+
k : int, optional
982+
Index of the diagonal. 0 (the default) refers to the main diagonal,
983+
a positive value refers to an upper diagonal, and a negative value to a
984+
lower diagonal. For arrays with more than two dimensions, the offset is
985+
applied between the first two axes.
986+
946987
See Also
947988
--------
948989
diag_indices_from
@@ -969,6 +1010,11 @@ def diag_indices(n, ndim=2):
9691010
[ 8, 9, 100, 11],
9701011
[ 12, 13, 14, 100]])
9711012
1013+
Create indices for an offset diagonal:
1014+
1015+
>>> np.diag_indices(4, k=1)
1016+
(array([0, 1, 2]), array([1, 2, 3]))
1017+
9721018
Now, we create indices to manipulate a 3-D array:
9731019
9741020
>>> d3 = np.diag_indices(2, 3)
@@ -987,23 +1033,46 @@ def diag_indices(n, ndim=2):
9871033
9881034
"""
9891035
idx = np.arange(n)
990-
return (idx,) * ndim
1036+
ndim = operator.index(ndim)
1037+
k = operator.index(k)
1038+
if k == 0:
1039+
return (idx,) * ndim
1040+
1041+
n = len(idx)
1042+
diag_len = max(n - abs(k), 0)
1043+
base = idx[:diag_len]
1044+
if ndim < 2:
1045+
return (base,) * ndim
1046+
1047+
if k > 0:
1048+
first = base
1049+
second = idx[k:k + diag_len]
1050+
else:
1051+
first = idx[-k:-k + diag_len]
1052+
second = base
9911053

1054+
return (first, second) + (base,) * (ndim - 2)
9921055

993-
def _diag_indices_from(arr):
1056+
1057+
def _diag_indices_from(arr, k=None):
9941058
return (arr,)
9951059

9961060

9971061
@array_function_dispatch(_diag_indices_from)
998-
def diag_indices_from(arr):
1062+
def diag_indices_from(arr, k=0):
9991063
"""
1000-
Return the indices to access the main diagonal of an n-dimensional array.
1064+
Return the indices to access a diagonal of an n-dimensional array.
10011065
10021066
See `diag_indices` for full details.
10031067
10041068
Parameters
10051069
----------
10061070
arr : array, at least 2-D
1071+
k : int, optional
1072+
Index of the diagonal. 0 (the default) refers to the main diagonal,
1073+
a positive value refers to an upper diagonal, and a negative value to a
1074+
lower diagonal. For arrays with more than two dimensions, the offset is
1075+
applied between the first two axes.
10071076
10081077
See Also
10091078
--------
@@ -1031,6 +1100,11 @@ def diag_indices_from(arr):
10311100
>>> a[di]
10321101
array([ 0, 5, 10, 15])
10331102
1103+
Get the indices for an offset diagonal:
1104+
1105+
>>> np.diag_indices_from(a, k=1)
1106+
(array([0, 1, 2]), array([1, 2, 3]))
1107+
10341108
This is simply syntactic sugar for diag_indices.
10351109
10361110
>>> np.diag_indices(a.shape[0])
@@ -1045,4 +1119,4 @@ def diag_indices_from(arr):
10451119
if not np.all(diff(arr.shape) == 0):
10461120
raise ValueError("All dimensions of input must be of equal length")
10471121

1048-
return diag_indices(arr.shape[0], arr.ndim)
1122+
return diag_indices(arr.shape[0], arr.ndim, k)

‎numpy/lib/_index_tricks_impl.pyi‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -310,9 +310,6 @@ def ix_(
310310
*args: Sequence[_ScalarLike_co] | _Array1D[Any],
311311
) -> tuple[NDArray[Any], ...]: ...
312312

313-
#
314-
def fill_diagonal(a: NDArray[Any], val: object, wrap: bool = False) -> None: ...
315-
316313
#
317314
@overload
318315
def diag_indices(n: _IntLike_co, ndim: L[0]) -> tuple[()]: ...
@@ -322,7 +319,14 @@ def diag_indices(n: _IntLike_co, ndim: L[1]) -> tuple[_Int1D]: ...
322319
def diag_indices(n: _IntLike_co, ndim: L[2] = 2) -> tuple[_Int1D, _Int1D]: ...
323320
@overload
324321
def diag_indices(n: _IntLike_co, ndim: L[3]) -> tuple[_Int1D, _Int1D, _Int1D]: ...
325-
@overload
322+
323+
324+
#
325+
def fill_diagonal(a: NDArray[Any], val: object, wrap: bool = False, *, k: int = 0) -> None: ...
326+
327+
#
328+
def diag_indices(n: int, ndim: int = 2, k: int = 0) -> tuple[NDArray[np.intp], ...]: ...
329+
def diag_indices_from(arr: ArrayLike, k: int = 0) -> tuple[NDArray[np.intp], ...]: ...
326330
def diag_indices(n: _IntLike_co, ndim: int) -> tuple[_Int1D, ...]: ...
327331

328332
#

0 commit comments

Comments
 (0)