<commit_before>//! Define semi-implicit schemes

use ndarray::*;
use ndarray_linalg::*;

use super::traits::*;

/// Linear ODE with diagonalized matrix (exactly solvable)
#[derive(Debug, Clone)]
pub struct Diagonal<A, D>
where
    A: Scalar,
    D: Dimension,
{
    exp_diag: Array<A, D>,
    diag: Array<A, D>,
    dt: A::Real,
}

impl<A, D> TimeStep for Diagonal<A, D>
where
    A: Scalar,
    D: Dimension,
{
    type Time = A::Real;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }
    fn set_dt(&mut self, dt: Self::Time) {
        Zip::from(&mut self.exp_diag)
            .and(&self.diag)
            .apply(|a, &b| {
                *a = b.mul_real(dt).exp();
            });
    }
}

impl<A, D> ModelSpec for Diagonal<A, D>
where
    A: Scalar,
    D: Dimension,
{
    type Scalar = A;
    type Dim = D;

    fn model_size(&self) -> D::Pattern {
        self.exp_diag.dim()
    }
}

impl<A, D> TimeEvolution for Diagonal<A, D>
where
    A: Scalar,
    D: Dimension,
{
    fn iterate<'a, S>(&mut self, x: &'a mut ArrayBase<S, D>) -> &'a mut ArrayBase<S, D>
    where
        S: DataMut<Elem = A>,
    {
        for (val, d) in x.iter_mut().zip(self.exp_diag.iter()) {
            *val = *val * *d;
        }
        x
    }
}

impl<F: SemiImplicit> Scheme<F> for Diagonal<F::Scalar, F::Dim> {
    fn new(f: F, dt: Self::Time) -> Self {
        let diag = f.diag();
        let mut exp_diag = diag.to_owned();
        for v in exp_diag.iter_mut() {
            *v = v.mul_real(dt).exp();
        }
        Diagonal {
            exp_diag: exp_diag,
            diag: diag,
            dt: dt,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DiagRK4<A, D, NLin, Lin>
where
    A: Scalar,
    D: Dimension,
    NLin: ModelSpec<Scalar = A, Dim = D>,
    Lin: ModelSpec<Scalar = A, Dim = D> + TimeStep<Time = A::Real>,
{
    nlin: NLin,
    lin: Lin,
    dt: Lin::Time,
    x: Array<A, D>,
    lx: Array<A, D>,
    k1: Array<A, D>,
    k2: Array<A, D>,
    k3: Array<A, D>,
}

pub fn diag_rk4<A, D, NLin>(nlin: NLin, dt: A::Real) -> DiagRK4<A, D, NLin, Diagonal<A, D>>
where
    A: Scalar,
    D: Dimension,
    NLin: SemiImplicit<Scalar = A, Dim = D>,
{
    let lin = Diagonal::new(nlin.clone(), dt / into_scalar(2.0));
    let x = Array::zeros(lin.model_size());
    let lx = Array::zeros(lin.model_size());
    let k1 = Array::zeros(lin.model_size());
    let k2 = Array::zeros(lin.model_size());
    let k3 = Array::zeros(lin.model_size());
    DiagRK4 {
        nlin,
        lin,
        dt,
        x,
        lx,
        k1,
        k2,
        k3,
    }
}

impl<A, D, NLin, Lin> TimeStep for DiagRK4<A, D, NLin, Lin>
where
    A: Scalar,
    D: Dimension,
    NLin: ModelSpec<Scalar = A, Dim = D>,
    Lin: ModelSpec<Scalar = A, Dim = D> + TimeStep<Time = A::Real>,
{
    type Time = Lin::Time;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }

    fn set_dt(&mut self, dt: Self::Time) {
        self.lin.set_dt(dt / into_scalar(2.0));
    }
}

impl<A, D, NLin, Lin> ModelSpec for DiagRK4<A, D, NLin, Lin>
where
    A: Scalar,
    D: Dimension,
    NLin: ModelSpec<Scalar = A, Dim = D>,
    Lin: ModelSpec<Scalar = A, Dim = D> + TimeStep<Time = A::Real>,
{
    type Scalar = A;
    type Dim = D;

    fn model_size(&self) -> <Self::Dim as Dimension>::Pattern {
        self.nlin.model_size() // TODO check
    }
}

impl<A, D, NLin, Lin> TimeEvolution for DiagRK4<A, D, NLin, Lin>
where
    A: Scalar,
    D: Dimension,
    NLin: SemiImplicit<Scalar = A, Dim = D>,
    Lin: TimeEvolution<Scalar = A, Dim = D> + TimeStep<Time = A::Real>,
{
    fn iterate<'a, S>(
        &mut self,
        x: &'a mut ArrayBase<S, Self::Dim>,
    ) -> &'a mut ArrayBase<S, Self::Dim>
    where
        S: DataMut<Elem = A>,
    {
        // constants
        let dt = self.dt;
        let dt_2 = self.dt / into_scalar(2.0);
        let dt_3 = self.dt / into_scalar(3.0);
        let dt_6 = self.dt / into_scalar(6.0);
        // operators
        let l = &mut self.lin;
        let f = &mut self.nlin;
        // calc
        self.x.zip_mut_with(x, |buf, x| *buf = *x);
        self.lx.zip_mut_with(x, |buf, lx| *buf = *lx);
        l.iterate(&mut self.lx);
        let k1 = f.nlin(x);
        self.k1.zip_mut_with(k1, |buf, k1| *buf = *k1);
        Zip::from(&mut *k1).and(&self.x).apply(|k1, &x_| {
            *k1 = x_ + k1.mul_real(dt_2);
        });
        let k2 = f.nlin(l.iterate(k1));
        self.k2.zip_mut_with(k2, |buf, k| *buf = *k);
        Zip::from(&mut *k2).and(&self.lx).apply(|k2, &lx| {
            *k2 = lx + k2.mul_real(dt_2);
        });
        let k3 = f.nlin(k2);<commit_msg>Simplify type restrictions<commit_after>//! Define semi-implicit schemes

use ndarray::*;
use ndarray_linalg::*;

use super::traits::*;

/// Linear ODE with diagonalized matrix (exactly solvable)
#[derive(Debug, Clone)]
pub struct Diagonal<F: SemiImplicit> {
    exp_diag: Array<F::Scalar, F::Dim>,
    diag: Array<F::Scalar, F::Dim>,
    dt: <F::Scalar as AssociatedReal>::Real,
}

impl<F: SemiImplicit> TimeStep for Diagonal<F> {
    type Time = <F::Scalar as AssociatedReal>::Real;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }
    fn set_dt(&mut self, dt: Self::Time) {
        Zip::from(&mut self.exp_diag)
            .and(&self.diag)
            .apply(|a, &b| {
                *a = b.mul_real(dt).exp();
            });
    }
}

impl<F: SemiImplicit> ModelSpec for Diagonal<F> {
    type Scalar = F::Scalar;
    type Dim = F::Dim;

    fn model_size(&self) -> <Self::Dim as Dimension>::Pattern {
        self.exp_diag.dim()
    }
}

impl<F: SemiImplicit> TimeEvolution for Diagonal<F> {
    fn iterate<'a, S>(
        &mut self,
        x: &'a mut ArrayBase<S, Self::Dim>,
    ) -> &'a mut ArrayBase<S, Self::Dim>
    where
        S: DataMut<Elem = Self::Scalar>,
    {
        for (val, d) in x.iter_mut().zip(self.exp_diag.iter()) {
            *val = *val * *d;
        }
        x
    }
}

impl<F: SemiImplicit> Scheme<F> for Diagonal<F> {
    fn new(f: F, dt: Self::Time) -> Self {
        let diag = f.diag();
        let mut exp_diag = diag.to_owned();
        for v in exp_diag.iter_mut() {
            *v = v.mul_real(dt).exp();
        }
        Diagonal {
            exp_diag: exp_diag,
            diag: diag,
            dt: dt,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DiagRK4<F: SemiImplicit> {
    nlin: F,
    lin: Diagonal<F>,
    dt: <Diagonal<F> as TimeStep>::Time,
    x: Array<F::Scalar, F::Dim>,
    lx: Array<F::Scalar, F::Dim>,
    k1: Array<F::Scalar, F::Dim>,
    k2: Array<F::Scalar, F::Dim>,
    k3: Array<F::Scalar, F::Dim>,
}

impl<F: SemiImplicit> Scheme<F> for DiagRK4<F> {
    fn new(nlin: F, dt: Self::Time) -> Self {
        let lin = Diagonal::new(nlin.clone(), dt / into_scalar(2.0));
        let x = Array::zeros(lin.model_size());
        let lx = Array::zeros(lin.model_size());
        let k1 = Array::zeros(lin.model_size());
        let k2 = Array::zeros(lin.model_size());
        let k3 = Array::zeros(lin.model_size());
        DiagRK4 {
            nlin,
            lin,
            dt,
            x,
            lx,
            k1,
            k2,
            k3,
        }
    }
}

impl<F: SemiImplicit> TimeStep for DiagRK4<F> {
    type Time = <Diagonal<F> as TimeStep>::Time;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }

    fn set_dt(&mut self, dt: Self::Time) {
        self.lin.set_dt(dt / into_scalar(2.0));
    }
}

impl<F: SemiImplicit> ModelSpec for DiagRK4<F> {
    type Scalar = F::Scalar;
    type Dim = F::Dim;

    fn model_size(&self) -> <Self::Dim as Dimension>::Pattern {
        self.nlin.model_size() // TODO check
    }
}

impl<F: SemiImplicit> TimeEvolution for DiagRK4<F> {
    fn iterate<'a, S>(
        &mut self,
        x: &'a mut ArrayBase<S, Self::Dim>,
    ) -> &'a mut ArrayBase<S, Self::Dim>
    where
        S: DataMut<Elem = Self::Scalar>,
    {
        // constants
        let dt = self.dt;
        let dt_2 = self.dt / into_scalar(2.0);
        let dt_3 = self.dt / into_scalar(3.0);
        let dt_6 = self.dt / into_scalar(6.0);
        // operators
        let l = &mut self.lin;
        let f = &mut self.nlin;
        // calc
        self.x.zip_mut_with(x, |buf, x| *buf = *x);
        self.lx.zip_mut_with(x, |buf, lx| *buf = *lx);
        l.iterate(&mut self.lx);
        let k1 = f.nlin(x);
        self.k1.zip_mut_with(k1, |buf, k1| *buf = *k1);
        Zip::from(&mut *k1).and(&self.x).apply(|k1, &x_| {
            *k1 = x_ + k1.mul_real(dt_2);
        });
        let k2 = f.nlin(l.iterate(k1));
        self.k2.zip_mut_with(k2, |buf, k| *buf = *k);
        Zip::from(&mut *k2).and(&self.lx).apply(|k2, &lx| {
            *k2 = lx + k2.mul_real(dt_2);
        });
        let k3 = f.nlin(k2);
{'fraction_non_alphanumeric': 0.17282261741906066, 'fraction_numerical': 0.011285909712722298, 'mean_word_length': 2.0674825174825173, 'pattern_counts': {'":': 0, '<': 96, '<?xml version=': 0, '>': 112, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2601668', 'n_tokens_mistral': 3503, 'n_tokens_neox': 3422, 'n_words': 791}