seqnmf-linalg 0.1.4

Linalg functions used in seqNMF.
Documentation
use ndarray::*;
use num_traits::Float;

use crate::linalg::*;

pub fn update_wh<T>(
    k: usize,
    l: usize,
    lambda: T,
    w: &mut Array3<T>,
    h: &mut Array2<T>,
    x_bar: &mut Array2<T>,
    wtx: &mut Array2<T>,
    x: &Array2<T>,
    x_pad: &Padded2D<T>,
    cost: &mut Vec<T>,
    epsilon: T,
    update_w: bool,
    update_h: bool,
) where
    T: Float + std::ops::AddAssign + ScalarOperand + std::ops::MulAssign + 'static,
{
    if !update_w && !update_h {
        return;
    }
    let one = T::one();
    let ten = T::from(10).unwrap();
    let nine_ten = (ten - one) / ten;

    let mut w_new = w.to_owned();
    let mut h_new = h.to_owned();

    let h_pad = Padded2D::from_arr(h, l);
    let wtx_bar = w.transpose_convolution_headtail(&Padded2D::from_arr(x_bar, l));

    let one_eye = Array2::ones((k, k)) - Array2::eye(k);
    let l_sh_one_eye = Array2::smooth_x_headtail(&h.t().dot(&one_eye), l) * lambda;
    if update_w {
        for i in 0..l {
            let hlt = h_pad
                .shift_view(i, ShiftDirect::Right)
                .unwrap()
                .t()
                .to_owned();
            let n0 = x.dot(&hlt);
            let d0 = x_bar.dot(&hlt);
            let d1 = x_pad
                .shift_view(i, ShiftDirect::Right)
                .unwrap()
                .dot(&l_sh_one_eye);
            let dw1 = n0 / (d0 + d1 + epsilon);
            Zip::from(&mut w_new.slice_mut(s![.., .., i]))
                .and(&dw1)
                .apply(|a, &b| {
                    *a *= b;
                });
        }
    }
    if update_h {
        let dh = wtx.to_owned() / (wtx_bar + one_eye.dot(&wtx.smooth_headtail(l)) * lambda + epsilon);
        Zip::from(&mut h_new).and(&dh).apply(|a, &b| {
            *a *= b;
        });
    }

    let cost_last = cost[cost.len() - 1];
    let w_old_nine_ten = w.to_owned() * nine_ten;
    let h_old_nine_ten = h.to_owned() * nine_ten;
    for i in 0..11 {
        *x_bar = w_new.convolution(&Padded2D::from_arr(&h_new, l));
        let c = reconstruction_cost(x_bar, x);
        if c < cost_last || i == 10 {
            cost.push(c);
            break;
        } else {
            if update_w {
                w_new = &w_new / ten + &w_old_nine_ten;
            }
            if update_h {
                h_new = &h_new / ten + &h_old_nine_ten;
            }
        }
    }
    if update_w {
        *w = w_new;
    }
    if update_h {
        *h = h_new;
    }
}