Implement numpy's triu and tril methods #1323
Description
Hi, I want to propose an implementation of numpy.triu and numpy.tril methods (link here) for arrays that are atleast 2 dimensions.
import numpy as np
np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=0)returns:
array([[ 1, 2, 3],
[ 0, 5, 6],
[ 0, 0, 9],
[ 0, 0, 0]])and
np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=1)returns:
array([[ 0, 2, 3],
[ 0, 0, 6],
[ 0, 0, 0],
[ 0, 0, 0]])and
np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=-1)returns:
array([[ 1, 2, 3],
[ 4, 5, 6],
[ 0, 8, 9],
[ 0, 0, 12]])I implemented a basic version of np.triu and np.tril for 2D owned float square ndarrays for personal use, however I did not write it to support negative values of k.
pub trait Tri {
fn triu(&self, k: usize) -> Self;
fn tril(&self, k: usize) -> Self;
}
impl Tri for ArrayBase<OwnedRepr<f64>, Ix2> {
fn triu(&self, k: usize) -> Self {
let cols = self.len_of(Axis(1));
let copied_arr = self.clone();
let mut result_arr = Array2::<f64>::zeros((cols, cols).f());
for (i, row) in copied_arr.axis_iter(Axis(0)).enumerate() {
let slice_main = (i as usize + k)..self.len_of(Axis(0));
row.to_owned()
.slice(s![slice_main.clone()])
.assign_to(result_arr.slice_mut(s![i, slice_main.clone()]));
}
return result_arr;
}
fn tril(&self, k: usize) -> Self {
return self.clone().t().to_owned().triu(k).t().to_owned();
}
}Reactions are currently unavailable