use super::*;
use ndarray::{Array1, Ix1, Ix2, IntoDimension};
use std::ops::{Add, AddAssign, Mul, MulAssign};
type GradMap<C> = ::std::collections::HashMap<usize, C>;
#[derive(Clone, Debug, Default)]
#[cfg_attr(
feature = "serde",
derive(Serialize, Deserialize),
serde(crate = "serde_crate")
)]
pub struct Columnar<C: Buffer<Dim = Ix1> = Array1<f64>> {
dim: Ix2,
grads: GradMap<C>,
}
impl<C: Buffer<Dim = Ix1>> Columnar<C> {
pub fn new<D: IntoDimension<Dim = Ix2>>(dim: D, grads: GradMap<C>) -> Self {
Columnar {
dim: dim.into_dimension(),
grads,
}
}
pub fn from_column(n_cols: usize, index: usize, grad: C) -> Self {
let dim = (grad.dim(), n_cols);
let mut grads = GradMap::new();
grads.insert(index, grad);
Self::new(dim, grads)
}
}
impl<C: Buffer<Dim = Ix1>> Buffer for Columnar<C> {
type Dim = Ix2;
fn dim(&self) -> (usize, usize) { self.dim.into_pattern() }
fn raw_dim(&self) -> Ix2 { self.dim }
fn addto<D: DataMut<Elem = f64>>(&self, weights: &mut ArrayBase<D, Ix2>) {
for (&c, pds) in self.grads.iter() {
pds.addto(&mut weights.column_mut(c));
}
}
fn scaled_addto<D: DataMut<Elem = f64>>(&self, alpha: f64, weights: &mut ArrayBase<D, Ix2>) {
for (&c, buf) in self.grads.iter() {
buf.scaled_addto(alpha, &mut weights.column_mut(c));
}
}
fn to_dense(&self) -> Array2<f64> {
let mut g_matrix = Array2::zeros(self.dim);
self.addto(&mut g_matrix);
g_matrix
}
}
impl<C: BufferMut<Dim = Ix1>> BufferMut for Columnar<C> {
fn zeros<D: IntoDimension<Dim = Ix2>>(dim: D) -> Self { Self::new(dim, GradMap::new()) }
fn map(&self, f: impl Fn(f64) -> f64) -> Self {
Columnar {
dim: self.dim,
grads: self.grads.iter().map(|(c, j)| (*c, j.map(&f))).collect(),
}
}
fn map_into(mut self, f: impl Fn(f64) -> f64) -> Self {
self.map_inplace(f);
self
}
fn map_inplace(&mut self, f: impl Fn(f64) -> f64) {
self.grads.values_mut().for_each(|buf| buf.map_inplace(&f));
}
fn merge(&self, other: &Self, f: impl Fn(f64, f64) -> f64) -> Self {
Columnar {
dim: self.dim,
grads: self
.grads
.iter()
.map(|(c, j1)| match other.grads.get(c) {
Some(j2) => (*c, j1.merge(j2, &f)),
None => (*c, j1.map(|x| f(x, 0.0))),
})
.collect(),
}
}
fn merge_inplace(&mut self, other: &Self, f: impl Fn(f64, f64) -> f64) {
for (c, xs) in self.grads.iter_mut() {
if let Some(ys) = other.grads.get(c) {
xs.merge_inplace(ys, &f)
} else {
xs.map_inplace(|x| f(x, 0.0))
}
}
for (&c, ys) in other.grads.iter() {
self.grads.entry(c).or_insert_with(|| ys.map(|y| f(0.0, y)));
}
}
}
impl<C: Buffer<Dim = Ix1>> Into<Array2<f64>> for Columnar<C> {
fn into(self) -> Array2<f64> { self.to_dense() }
}
impl<C: BufferMut<Dim = Ix1>> Add<Columnar<C>> for Columnar<C> {
type Output = Columnar<C>;
fn add(mut self, other: Columnar<C>) -> Columnar<C> {
self.add_assign(&other);
self
}
}
impl<C: BufferMut<Dim = Ix1>> AddAssign<Columnar<C>> for Columnar<C> {
fn add_assign(&mut self, other: Columnar<C>) {
for (c, buf_other) in other.grads.into_iter() {
if let Some(buf) = self.grads.get_mut(&c) {
buf.merge_inplace(&buf_other, |x, y| x + y);
} else {
self.grads.insert(c, buf_other);
}
}
}
}
impl<C: BufferMut<Dim = Ix1>> AddAssign<&Columnar<C>> for Columnar<C> {
fn add_assign(&mut self, other: &Columnar<C>) {
for (&c, buf_other) in other.grads.iter() {
if let Some(buf) = self.grads.get_mut(&c) {
buf.merge_inplace(buf_other, |x, y| x + y);
} else {
self.grads.insert(c, buf_other.to_owned());
}
}
}
}
impl<C: BufferMut<Dim = Ix1>> Mul<f64> for Columnar<C> {
type Output = Columnar<C>;
fn mul(mut self, alpha: f64) -> Columnar<C> {
self.mul_assign(alpha);
self
}
}
impl<C: BufferMut<Dim = Ix1>> MulAssign<f64> for Columnar<C> {
fn mul_assign(&mut self, alpha: f64) {
for buf in self.grads.values_mut() {
buf.map_inplace(|x| alpha * x);
}
}
}