diff_op_col as function factory
This commit is contained in:
		@@ -171,60 +171,61 @@ pub(crate) fn diff_op_1d(
 | 
			
		||||
 | 
			
		||||
#[inline(always)]
 | 
			
		||||
pub(crate) fn diff_op_col(
 | 
			
		||||
    block: &[&[Float]],
 | 
			
		||||
    diag: &[Float],
 | 
			
		||||
    block: &'static [&'static [Float]],
 | 
			
		||||
    diag: &'static [Float],
 | 
			
		||||
    symmetric: bool,
 | 
			
		||||
    is_h2: bool,
 | 
			
		||||
    prev: ArrayView2<Float>,
 | 
			
		||||
    mut fut: ArrayViewMut2<Float>,
 | 
			
		||||
) {
 | 
			
		||||
    assert_eq!(prev.shape(), fut.shape());
 | 
			
		||||
    let nx = prev.shape()[1];
 | 
			
		||||
    assert!(nx >= 2 * block.len());
 | 
			
		||||
) -> impl Fn(ArrayView2<Float>, ArrayViewMut2<Float>) {
 | 
			
		||||
    #[inline(always)]
 | 
			
		||||
    move |prev: ArrayView2<Float>, mut fut: ArrayViewMut2<Float>| {
 | 
			
		||||
        assert_eq!(prev.shape(), fut.shape());
 | 
			
		||||
        let nx = prev.shape()[1];
 | 
			
		||||
        assert!(nx >= 2 * block.len());
 | 
			
		||||
 | 
			
		||||
    assert_eq!(prev.strides()[0], 1);
 | 
			
		||||
    assert_eq!(fut.strides()[0], 1);
 | 
			
		||||
        assert_eq!(prev.strides()[0], 1);
 | 
			
		||||
        assert_eq!(fut.strides()[0], 1);
 | 
			
		||||
 | 
			
		||||
    let dx = if is_h2 {
 | 
			
		||||
        1.0 / (nx - 2) as Float
 | 
			
		||||
    } else {
 | 
			
		||||
        1.0 / (nx - 1) as Float
 | 
			
		||||
    };
 | 
			
		||||
    let idx = 1.0 / dx;
 | 
			
		||||
        let dx = if is_h2 {
 | 
			
		||||
            1.0 / (nx - 2) as Float
 | 
			
		||||
        } else {
 | 
			
		||||
            1.0 / (nx - 1) as Float
 | 
			
		||||
        };
 | 
			
		||||
        let idx = 1.0 / dx;
 | 
			
		||||
 | 
			
		||||
    fut.fill(0.0);
 | 
			
		||||
        fut.fill(0.0);
 | 
			
		||||
 | 
			
		||||
    // First block
 | 
			
		||||
    for (bl, mut fut) in block.iter().zip(fut.axis_iter_mut(ndarray::Axis(1))) {
 | 
			
		||||
        debug_assert_eq!(fut.len(), prev.shape()[0]);
 | 
			
		||||
        for (&bl, prev) in bl.iter().zip(prev.axis_iter(ndarray::Axis(1))) {
 | 
			
		||||
            debug_assert_eq!(prev.len(), fut.len());
 | 
			
		||||
            fut.scaled_add(idx * bl, &prev);
 | 
			
		||||
        }
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    let half_diag_width = (diag.len() - 1) / 2;
 | 
			
		||||
 | 
			
		||||
    // Diagonal entries
 | 
			
		||||
    for (ifut, mut fut) in fut
 | 
			
		||||
        .axis_iter_mut(ndarray::Axis(1))
 | 
			
		||||
        .enumerate()
 | 
			
		||||
        .skip(block.len())
 | 
			
		||||
        .take(nx - 2 * block.len())
 | 
			
		||||
    {
 | 
			
		||||
        for (id, d) in diag.iter().enumerate() {
 | 
			
		||||
            let offset = ifut - half_diag_width + id;
 | 
			
		||||
            fut.scaled_add(idx * d, &prev.slice(ndarray::s![.., offset]))
 | 
			
		||||
        }
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    // End block
 | 
			
		||||
    for (bl, mut fut) in block.iter().zip(fut.axis_iter_mut(ndarray::Axis(1)).rev()) {
 | 
			
		||||
        for (&bl, prev) in bl.iter().zip(prev.axis_iter(ndarray::Axis(1)).rev()) {
 | 
			
		||||
            if symmetric {
 | 
			
		||||
        // First block
 | 
			
		||||
        for (bl, mut fut) in block.iter().zip(fut.axis_iter_mut(ndarray::Axis(1))) {
 | 
			
		||||
            debug_assert_eq!(fut.len(), prev.shape()[0]);
 | 
			
		||||
            for (&bl, prev) in bl.iter().zip(prev.axis_iter(ndarray::Axis(1))) {
 | 
			
		||||
                debug_assert_eq!(prev.len(), fut.len());
 | 
			
		||||
                fut.scaled_add(idx * bl, &prev);
 | 
			
		||||
            } else {
 | 
			
		||||
                fut.scaled_add(-idx * bl, &prev);
 | 
			
		||||
            }
 | 
			
		||||
        }
 | 
			
		||||
 | 
			
		||||
        let half_diag_width = (diag.len() - 1) / 2;
 | 
			
		||||
 | 
			
		||||
        // Diagonal entries
 | 
			
		||||
        for (ifut, mut fut) in fut
 | 
			
		||||
            .axis_iter_mut(ndarray::Axis(1))
 | 
			
		||||
            .enumerate()
 | 
			
		||||
            .skip(block.len())
 | 
			
		||||
            .take(nx - 2 * block.len())
 | 
			
		||||
        {
 | 
			
		||||
            for (id, d) in diag.iter().enumerate() {
 | 
			
		||||
                let offset = ifut - half_diag_width + id;
 | 
			
		||||
                fut.scaled_add(idx * d, &prev.slice(ndarray::s![.., offset]))
 | 
			
		||||
            }
 | 
			
		||||
        }
 | 
			
		||||
 | 
			
		||||
        // End block
 | 
			
		||||
        for (bl, mut fut) in block.iter().zip(fut.axis_iter_mut(ndarray::Axis(1)).rev()) {
 | 
			
		||||
            for (&bl, prev) in bl.iter().zip(prev.axis_iter(ndarray::Axis(1)).rev()) {
 | 
			
		||||
                if symmetric {
 | 
			
		||||
                    fut.scaled_add(idx * bl, &prev);
 | 
			
		||||
                } else {
 | 
			
		||||
                    fut.scaled_add(-idx * bl, &prev);
 | 
			
		||||
                }
 | 
			
		||||
            }
 | 
			
		||||
        }
 | 
			
		||||
    }
 | 
			
		||||
 
 | 
			
		||||
@@ -43,7 +43,7 @@ impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &SBP4) {
 | 
			
		||||
                diff_op_row(SBP4::BLOCK, SBP4::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(SBP4::BLOCK, SBP4::DIAG, false, false, prev, fut);
 | 
			
		||||
                diff_op_col(SBP4::BLOCK, SBP4::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
 
 | 
			
		||||
@@ -47,7 +47,7 @@ impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &SBP8) {
 | 
			
		||||
                diff_op_row(SBP8::BLOCK, SBP8::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(SBP8::BLOCK, SBP8::DIAG, false, false, prev, fut);
 | 
			
		||||
                diff_op_col(SBP8::BLOCK, SBP8::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
 
 | 
			
		||||
@@ -61,7 +61,7 @@ impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &Upwind4h2) {
 | 
			
		||||
                diff_op_row(Upwind4h2::BLOCK, Upwind4h2::DIAG, false, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(Upwind4h2::BLOCK, Upwind4h2::DIAG, false, true, prev, fut);
 | 
			
		||||
                diff_op_col(Upwind4h2::BLOCK, Upwind4h2::DIAG, false, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
@@ -84,14 +84,7 @@ impl<UO: UpwindOperator1d> UpwindOperator2d for (&UO, &Upwind4h2) {
 | 
			
		||||
                diff_op_row(Upwind4h2::DISS_BLOCK, Upwind4h2::DISS_DIAG, true, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(
 | 
			
		||||
                    Upwind4h2::DISS_BLOCK,
 | 
			
		||||
                    Upwind4h2::DISS_DIAG,
 | 
			
		||||
                    true,
 | 
			
		||||
                    true,
 | 
			
		||||
                    prev,
 | 
			
		||||
                    fut,
 | 
			
		||||
                );
 | 
			
		||||
                diff_op_col(Upwind4h2::DISS_BLOCK, Upwind4h2::DISS_DIAG, true, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
 
 | 
			
		||||
@@ -66,7 +66,7 @@ impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &Upwind9) {
 | 
			
		||||
                diff_op_row(Upwind9::BLOCK, Upwind9::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(Upwind9::BLOCK, Upwind9::DIAG, false, false, prev, fut);
 | 
			
		||||
                diff_op_col(Upwind9::BLOCK, Upwind9::DIAG, false, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
@@ -99,14 +99,7 @@ impl<UO: UpwindOperator1d> UpwindOperator2d for (&UO, &Upwind9) {
 | 
			
		||||
                diff_op_row(Upwind9::DISS_BLOCK, Upwind9::DISS_DIAG, true, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(
 | 
			
		||||
                    Upwind9::DISS_BLOCK,
 | 
			
		||||
                    Upwind9::DISS_DIAG,
 | 
			
		||||
                    true,
 | 
			
		||||
                    false,
 | 
			
		||||
                    prev,
 | 
			
		||||
                    fut,
 | 
			
		||||
                );
 | 
			
		||||
                diff_op_col(Upwind9::DISS_BLOCK, Upwind9::DISS_DIAG, true, false)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
 
 | 
			
		||||
@@ -69,7 +69,7 @@ impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &Upwind9h2) {
 | 
			
		||||
                diff_op_row(Upwind9h2::BLOCK, Upwind9h2::DIAG, false, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(Upwind9h2::BLOCK, Upwind9h2::DIAG, false, true, prev, fut);
 | 
			
		||||
                diff_op_col(Upwind9h2::BLOCK, Upwind9h2::DIAG, false, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
@@ -126,14 +126,7 @@ impl<UO: UpwindOperator1d> UpwindOperator2d for (&UO, &Upwind9h2) {
 | 
			
		||||
                diff_op_row(Upwind9h2::DISS_BLOCK, Upwind9h2::DISS_DIAG, true, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([1, _], [1, _]) => {
 | 
			
		||||
                diff_op_col(
 | 
			
		||||
                    Upwind9h2::DISS_BLOCK,
 | 
			
		||||
                    Upwind9h2::DISS_DIAG,
 | 
			
		||||
                    true,
 | 
			
		||||
                    true,
 | 
			
		||||
                    prev,
 | 
			
		||||
                    fut,
 | 
			
		||||
                );
 | 
			
		||||
                diff_op_col(Upwind9h2::DISS_BLOCK, Upwind9h2::DISS_DIAG, true, true)(prev, fut);
 | 
			
		||||
            }
 | 
			
		||||
            ([_, _], [_, _]) => {
 | 
			
		||||
                // Fallback, work row by row
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user