Skip to content

Commit 5ad1afa

Browse files
rthjturner314
authored andcommitted
Add Array2::from_diag (#673)
* ENH Add Array2::from_diag * Lint * Address review comments * Add example and panic conditions * Remove cfg(feature = "approx") * Add missing closing code fence
1 parent 07702a1 commit 5ad1afa

File tree

2 files changed

+37
-0
lines changed

2 files changed

+37
-0
lines changed

src/impl_constructors.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,29 @@ where
171171
}
172172
eye
173173
}
174+
175+
/// Create a 2D matrix from its diagonal
176+
///
177+
/// **Panics** if `diag.len() * diag.len()` would overflow `isize`.
178+
///
179+
/// ```rust
180+
/// use ndarray::{Array2, arr1, arr2};
181+
///
182+
/// let diag = arr1(&[1, 2]);
183+
/// let array = Array2::from_diag(&diag);
184+
/// assert_eq!(array, arr2(&[[1, 0], [0, 2]]));
185+
/// ```
186+
pub fn from_diag<S2>(diag: &ArrayBase<S2, Ix1>) -> Self
187+
where
188+
A: Clone + Zero,
189+
S: DataMut,
190+
S2: Data<Elem = A>,
191+
{
192+
let n = diag.len();
193+
let mut arr = Self::zeros((n, n));
194+
arr.diag_mut().assign(&diag);
195+
arr
196+
}
174197
}
175198

176199
#[cfg(not(debug_assertions))]

tests/array.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1950,6 +1950,20 @@ fn test_array_clone_same_view() {
19501950
assert_eq!(a, b);
19511951
}
19521952

1953+
#[test]
1954+
fn test_array2_from_diag() {
1955+
let diag = arr1(&[0, 1, 2]);
1956+
let x = Array2::from_diag(&diag);
1957+
let x_exp = arr2(&[[0, 0, 0], [0, 1, 0], [0, 0, 2]]);
1958+
assert_eq!(x, x_exp);
1959+
1960+
// check 0 length array
1961+
let diag = Array1::<f64>::zeros(0);
1962+
let x = Array2::from_diag(&diag);
1963+
assert_eq!(x.ndim(), 2);
1964+
assert_eq!(x.shape(), [0, 0]);
1965+
}
1966+
19531967
#[test]
19541968
fn array_macros() {
19551969
// array

0 commit comments

Comments
 (0)