|
1 |
| -//! This module defines newtypes for ndarray's Array1. |
| 1 | +//! This module defines newtypes for ndarray's Array. |
2 | 2 | //!
|
3 | 3 | //! This is necessary to be able to abstract over floats (f32 and f64) so that
|
4 | 4 | //! the logistic regression code can be abstract in the float type it works
|
|
8 | 8 |
|
9 | 9 | use crate::float::Float;
|
10 | 10 | use argmin::prelude::*;
|
11 |
| -use ndarray::Array1; |
| 11 | +use ndarray::{Array, ArrayBase, Data, Dimension, Zip}; |
12 | 12 | use serde::{Deserialize, Serialize};
|
13 | 13 |
|
| 14 | +pub fn elem_dot<F: linfa::Float, A1: Data<Elem = F>, A2: Data<Elem = F>, D: Dimension>( |
| 15 | + a: &ArrayBase<A1, D>, |
| 16 | + b: &ArrayBase<A2, D>, |
| 17 | +) -> F { |
| 18 | + Zip::from(a) |
| 19 | + .and(b) |
| 20 | + .fold(F::zero(), |acc, &a, &b| acc + a * b) |
| 21 | +} |
| 22 | + |
14 | 23 | #[derive(Serialize, Clone, Deserialize, Debug, Default)]
|
15 |
| -pub struct ArgminParam<F>(pub Array1<F>); |
| 24 | +pub struct ArgminParam<F, D: Dimension>(pub Array<F, D>); |
16 | 25 |
|
17 |
| -impl<F> ArgminParam<F> { |
| 26 | +impl<F, D: Dimension> ArgminParam<F, D> { |
18 | 27 | #[inline]
|
19 |
| - pub fn as_array(&self) -> &Array1<F> { |
| 28 | + pub fn as_array(&self) -> &Array<F, D> { |
20 | 29 | &self.0
|
21 | 30 | }
|
22 | 31 | }
|
23 | 32 |
|
24 |
| -impl<F: Float> ArgminSub<ArgminParam<F>, ArgminParam<F>> for ArgminParam<F> { |
25 |
| - fn sub(&self, other: &ArgminParam<F>) -> ArgminParam<F> { |
| 33 | +impl<F: Float, D: Dimension> ArgminSub<ArgminParam<F, D>, ArgminParam<F, D>> for ArgminParam<F, D> { |
| 34 | + fn sub(&self, other: &ArgminParam<F, D>) -> ArgminParam<F, D> { |
26 | 35 | ArgminParam(&self.0 - &other.0)
|
27 | 36 | }
|
28 | 37 | }
|
29 | 38 |
|
30 |
| -impl<F: Float> ArgminAdd<ArgminParam<F>, ArgminParam<F>> for ArgminParam<F> { |
31 |
| - fn add(&self, other: &ArgminParam<F>) -> ArgminParam<F> { |
| 39 | +impl<F: Float, D: Dimension> ArgminAdd<ArgminParam<F, D>, ArgminParam<F, D>> for ArgminParam<F, D> { |
| 40 | + fn add(&self, other: &ArgminParam<F, D>) -> ArgminParam<F, D> { |
32 | 41 | ArgminParam(&self.0 + &other.0)
|
33 | 42 | }
|
34 | 43 | }
|
35 | 44 |
|
36 |
| -impl<F: Float> ArgminDot<ArgminParam<F>, F> for ArgminParam<F> { |
37 |
| - fn dot(&self, other: &ArgminParam<F>) -> F { |
38 |
| - self.0.dot(&other.0) |
| 45 | +impl<F: Float, D: Dimension> ArgminDot<ArgminParam<F, D>, F> for ArgminParam<F, D> { |
| 46 | + fn dot(&self, other: &ArgminParam<F, D>) -> F { |
| 47 | + elem_dot(&self.0, &other.0) |
39 | 48 | }
|
40 | 49 | }
|
41 | 50 |
|
42 |
| -impl<F: Float> ArgminNorm<F> for ArgminParam<F> { |
| 51 | +impl<F: Float, D: Dimension> ArgminNorm<F> for ArgminParam<F, D> { |
43 | 52 | fn norm(&self) -> F {
|
44 |
| - self.0.dot(&self.0) |
| 53 | + num_traits::Float::sqrt(elem_dot(&self.0, &self.0)) |
45 | 54 | }
|
46 | 55 | }
|
47 | 56 |
|
48 |
| -impl<F: Float> ArgminMul<F, ArgminParam<F>> for ArgminParam<F> { |
49 |
| - fn mul(&self, other: &F) -> ArgminParam<F> { |
| 57 | +impl<F: Float, D: Dimension> ArgminMul<F, ArgminParam<F, D>> for ArgminParam<F, D> { |
| 58 | + fn mul(&self, other: &F) -> ArgminParam<F, D> { |
50 | 59 | ArgminParam(&self.0 * *other)
|
51 | 60 | }
|
52 | 61 | }
|
53 | 62 |
|
54 |
| -impl<F: Float> ArgminMul<ArgminParam<F>, ArgminParam<F>> for ArgminParam<F> { |
55 |
| - fn mul(&self, other: &ArgminParam<F>) -> ArgminParam<F> { |
| 63 | +impl<F: Float, D: Dimension> ArgminMul<ArgminParam<F, D>, ArgminParam<F, D>> for ArgminParam<F, D> { |
| 64 | + fn mul(&self, other: &ArgminParam<F, D>) -> ArgminParam<F, D> { |
56 | 65 | ArgminParam(&self.0 * &other.0)
|
57 | 66 | }
|
58 | 67 | }
|
0 commit comments