From 69c33bf74bcdc1d9781bd5db27f942f6d676c032 Mon Sep 17 00:00:00 2001 From: Frederic Date: Mon, 11 Jun 2012 16:23:17 -0400 Subject: fix the wrapping problem of fill_diagonal with tall matrix. --- numpy/lib/index_tricks.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'numpy/lib/index_tricks.py') diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py index c29f3a6d3..e248bfaea 100644 --- a/numpy/lib/index_tricks.py +++ b/numpy/lib/index_tricks.py @@ -719,10 +719,13 @@ def fill_diagonal(a, val): """ if a.ndim < 2: raise ValueError("array must be at least 2-d") + end = None if a.ndim == 2: # Explicit, fast formula for the common case. For 2-d arrays, we # accept rectangular ones. step = a.shape[1] + 1 + #This is needed to don't have tall matrix have the diagonal wrap. + end = a.shape[1] * a.shape[1] else: # For more than d=2, the strided formula is only valid for arrays with # all dimensions equal, so we check first. @@ -731,7 +734,7 @@ def fill_diagonal(a, val): step = 1 + (cumprod(a.shape[:-1])).sum() # Write the value out into the diagonal. - a.flat[::step] = val + a.flat[:end:step] = val def diag_indices(n, ndim=2): -- cgit v1.2.1