Skip to content

Commit 6e19f04

Browse files
jturner314bluss
authored andcommitted
Add from_diag_elem method
1 parent 562104a commit 6e19f04

File tree

1 file changed

+23
-0
lines changed

1 file changed

+23
-0
lines changed

src/impl_constructors.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,29 @@ where
230230
arr.diag_mut().assign(diag);
231231
arr
232232
}
233+
234+
/// Create a square 2D matrix of the specified size, with the specified
235+
/// element along the diagonal and zeros elsewhere.
236+
///
237+
/// **Panics** if `n * n` would overflow `isize`.
238+
///
239+
/// ```rust
240+
/// use ndarray::{array, Array2};
241+
///
242+
/// let array = Array2::from_diag_elem(2, 5.);
243+
/// assert_eq!(array, array![[5., 0.], [0., 5.]]);
244+
/// ```
245+
pub fn from_diag_elem(n: usize, elem: A) -> Self
246+
where
247+
S: DataMut,
248+
A: Clone + Zero,
249+
{
250+
let mut eye = Self::zeros((n, n));
251+
for a_ii in eye.diag_mut() {
252+
*a_ii = elem.clone();
253+
}
254+
eye
255+
}
233256
}
234257

235258
#[cfg(not(debug_assertions))]

0 commit comments

Comments
 (0)