src/dolfinx_access/function.rs

/*!
Python and C++ wrapper for Dolfinx `Function<f64>`.
*/

use super::ffi;
use super::ffi::Function_f64;
use alg_tools::bounds::{Bounds, GlobalAnalysis, MinMaxMapping};
use alg_tools::euclidean::Euclidean;
use alg_tools::fe_model::base::LocalModel;
use alg_tools::instance::Instance;
use alg_tools::loc::Loc;
use alg_tools::mapping::{BasicDecomposition, DifferentiableImpl, Mapping, Space};
use alg_tools::norms::{Dist, HasDual, Norm, Normed, L2};
use itertools::izip;
use pyo3::ffi::c_str;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyFunction};
use std::ffi::CStr;

/// A helper structure of dealing with dolfinx functions.
/// `N` is the domain dimension, `O` the order, and `D` is the codomain dimension.
#[allow(non_camel_case_types)]
pub struct DolfinxPyFunction_f64<'py, const N: u32, const O: u32 = 1, const D: u32 = 1> {
    #[allow(dead_code)]
    /// Python object.
    ///
    /// This is crucial for maintaining [`cxx`] live, when we obtained the object from Python.
    ///
    /// If we needed to copy-on-write, or initialise without obtaining and object from Python,
    /// this will become `None`.
    pub u: Option<Bound<'py, PyAny>>,
    /// We need to maintain a reference to the FunctionSpace Python wrapper, as it's difficult
    /// to recreate if we loose `py`. Fenics is a one-way hell full of pointless wrappers upon
    /// wrappers upon wrappers,very difficult to create having with just the unwrapped object.
    function_space: Bound<'py, PyAny>,
    /// Do we have the only reference?
    ///
    /// This is check on construction only, when we have the GIL.
    /// If we have the only reference at that time, this will not change during the lifetime.
    /// We use this to implement copy-on-write optimisations.
    owned: bool,
    /// C++ object pointer
    cxx: *mut Function_f64,
}

/// References to Python side helper routines
struct Helpers {
    /// Dot product
    dot: Py<PyFunction>,
    /// Squared 2-norm
    norm2_squared: Py<PyFunction>,
    /// Squared distance
    dist2_squared: Py<PyFunction>,
}

use std::sync::OnceLock;
/// A populated-once structure for referencing Python-side helper routines.
static HELPERS: OnceLock<Helpers> = OnceLock::new();

/// Gets [`Helpers`], populates [`HELPERS`] if needed.
fn get_helpers(py: Python<'_>) -> PyResult<&Helpers> {
    HELPERS.get_or_try_init(|| {
        let extras = PyModule::import(py, "pointsource_pde.dolfinx_extras")?;
        let g = |s| -> PyResult<Py<PyFunction>> { Ok(extras.getattr(s)?.cast_into()?.unbind()) };
        Ok(Helpers {
            dot: g("dot")?,
            norm2_squared: g("norm2_squared")?,
            dist2_squared: g("dist2_squared")?,
        })
    })
}

impl<'py, const N: u32, const O: u32, const D: u32> Drop for DolfinxPyFunction_f64<'py, N, O, D> {
    fn drop(&mut self) {
        // If Python is not managing our C++ object, we need to free it
        if let None = self.u {
            if self.cxx != std::ptr::null_mut() {
                unsafe { ffi::drop_Function_f64(self.cxx) };
            }
        }
        self.cxx = std::ptr::null_mut();
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> Clone for DolfinxPyFunction_f64<'py, N, O, D> {
    fn clone(&self) -> Self {
        DolfinxPyFunction_f64 {
            u: None, // It's a different new function, so no Python instance
            function_space: self.function_space.clone(),
            owned: true,
            cxx: unsafe { ffi::clone_Function_f64(self.cxx) },
        }
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> Space for DolfinxPyFunction_f64<'py, N, O, D> {
    type Principal = Self;
    type Decomp = BasicDecomposition;
}

impl<'py, const N: u32, const O: u32, const D: u32> DolfinxPyFunction_f64<'py, N, O, D> {
    /// Create a new instance, assuming `cxx` has been correctly extracted from `u`
    /// and checked to match `N`, `O`, `D`, and mesh restrictions.
    pub(crate) fn new_prechecked(u: Bound<'py, PyAny>, cxx: *mut Function_f64) -> PyResult<Self> {
        let function_space = u.getattr("_V")?;
        let owned = u.get_refcnt() <= 1;
        // TODO: check dimensionality
        Ok(DolfinxPyFunction_f64 { u: Some(u), function_space, owned, cxx })
    }

    /// Ensures that `self.u` exists.
    /// Panics if there's a failure to create it if needed.
    pub(crate) fn _make_py(&mut self, py: Python<'py>) {
        if let None = self.u {
            // We need to do a lot of work here, due to Fenics one-wayness.

            // This takes ownership of our pointer, so if we fail to create the Python
            // object, we have a lot of restoring to do.
            let cpp_object = unsafe {
                Py::<PyAny>::from_owned_ptr(
                    py,
                    ffi::wrap_Function_f64(self.cxx) as *mut pyo3::ffi::PyObject,
                )
            };
            let locals = PyDict::new(py);
            match locals.set_item("cpp_object", cpp_object).and_then(|()| {
                locals.set_item("V", self.function_space.as_borrowed())?;
                // TODO: precompile
                py.run(CREATE_FUNCTION_HACK, None, Some(&locals))?;
                locals.get_item("wrapper")
            }) {
                Err(e) => {
                    // We should restore to recover
                    panic!("Failure to pass object to Python: {:?}", e);
                }
                Ok(maybe_p) => {
                    self.u = Some(maybe_p.unwrap());
                }
            }
        }
    }

    /// Returns a (possibly new) Python presentation, without consuming current object.
    ///
    /// Please use this functionality through [`pyo3::IntoPyObject`].
    pub(crate) fn _into_py_ref<'a>(
        &'a mut self,
        py: Python<'py>,
    ) -> PyResult<pyo3::Borrowed<'a, 'py, PyAny>> {
        self._make_py(py);
        self.owned = false;
        match self.u {
            Some(ref p) => Ok(p.as_borrowed()),
            None => unreachable!("This cannot happen"),
        }
    }

    /// Returns a (possibly new) Python presentation, consuming current object.
    ///
    /// Please use this functionality through [`pyo3::IntoPyObject`].
    pub(crate) fn _into_py(mut self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, PyAny>> {
        self._make_py(py);
        self.cxx = std::ptr::null_mut();
        self.owned = false;
        match std::mem::replace(&mut self.u, None) {
            Some(p) => Ok(p),
            None => unreachable!("This cannot happen"),
        }
    }

    /// Create a new FEM function on the same function space, with unintialised data.
    pub fn similar(&self) -> Self {
        DolfinxPyFunction_f64 {
            u: None,
            function_space: self.function_space.clone(),
            owned: true,
            cxx: unsafe { ffi::similar_Function_f64(self.cxx) },
        }
    }

    /// Create a new FEM function on the same function space, with zero data.
    pub fn similar_zeros(&self) -> Self {
        let new = self.similar();
        for x in unsafe { ffi::data_mut_Function_f64(new.cxx) } {
            *x = 0.0;
        }
        new
    }
}

/// An extremely durty hack to create a Python-side Dolfinx `Function` referencing
/// a C++-side Dolfinx `Function` that we already have. For some unfathomable reason,
/// Dolfinx provies no way to do this cleanly.
pub(crate) static CREATE_FUNCTION_HACK: &CStr = c_str!(
    "\
from dolfinx.fem import Function
from dolfinx.la import Vector
x = Vector(cpp_object.x)
# Fenics is very one-way. It does not provide a proper way to create the wrapper,
# directly from the unwrapped C++
wrapper = Function.__new__(Function, V, x, \"?\", x.array.dtype)
# This all just repeats Function.__init__ without  creating a new ._cpp_object.
super(Function, wrapper).__init__(V.ufl_function_space())
wrapper._cpp_object = cpp_object;
wrapper._V = V
wrapper._x = x
wrapper.name = \"?\"
"
);

/// Trait for points that can be expanded into Dolfinx' always-3D presentation
trait ExpandCoordDolphinx {
    /// Convert to always-3D presentation
    fn to_dolphinx(self) -> [f64; 3];
    #[allow(dead_code)]
    /// Convert from always-3D presentation
    fn from_dolphinx(val: [f64; 3]) -> Self;
}

impl ExpandCoordDolphinx for Loc<1, f64> {
    #[inline]
    fn to_dolphinx(self) -> [f64; 3] {
        let Loc([x1]) = self;
        [x1, 0.0, 0.0]
    }

    #[inline]
    fn from_dolphinx([x1, _, _]: [f64; 3]) -> Self {
        [x1].into()
    }
}

impl ExpandCoordDolphinx for Loc<2, f64> {
    #[inline]
    fn to_dolphinx(self) -> [f64; 3] {
        let Loc([x1, x2]) = self;
        [x1, x2, 0.0]
    }

    #[inline]
    fn from_dolphinx([x1, x2, _]: [f64; 3]) -> Self {
        [x1, x2].into()
    }
}

impl ExpandCoordDolphinx for Loc<3, f64> {
    #[inline]
    fn to_dolphinx(self) -> [f64; 3] {
        let Loc([x1, x2, x3]) = self;
        [x1, x2, x3]
    }

    #[inline]
    fn from_dolphinx([x1, x2, x3]: [f64; 3]) -> Self {
        [x1, x2, x3].into()
    }
}

/// Implements [`DifferentiableImpl`]
macro_rules! impl_diff {
    ($n:literal ; $($o:literal)+) => { $(
        impl<'py> DifferentiableImpl<Loc<$n, f64>> for DolfinxPyFunction_f64<'py, $n, $o, 1> {
            type Derivative = Loc<$n, f64>;

            fn differential_impl<I: Instance<Loc<$n, f64>>>(&self, x: I) -> Self::Derivative {
                let x_own = x.own();
                let f = self.cxx;
                let expand_x = x_own.to_dolphinx();
                // Find cell containing `x`
                let cell = unsafe { ffi::cell_Function_f64(f, &expand_x) };
                if cell < 0 {
                    Loc::ORIGIN
                } else {
                    let (model, _) = unsafe {
                        // Find the coordinates of the corresponding triangle. We guarantee a
                        // triangular mesh in the constructor of `DolfinxPyFunction_f64`.
                        let simplex_coords = ffi::cell_coords_Function_f64_triangle(f, cell);
                        // Construct a new P2 model on that cell, as this is easier than dealing
                        // with the awful documentation of Dolphinx and C++, and probably having
                        // to reimplemen our P2 stuff anyway, to do any optimisation with the basix
                        // basis functions.
                        super::model_dolfinx_p2_cell(f, &simplex_coords, cell, false)
                    };
                    model.differential(&x_own)
                }
            }
        }
    )+ };
}

/// Implements [`Mapping`]
macro_rules! impl_mapping {
    ($($n:literal)+) => { $(
        impl<'py, const O: u32> Mapping<Loc<$n, f64>> for DolfinxPyFunction_f64<'py, $n, O, 1> {
            type Codomain = f64;

            fn apply<I: Instance<Loc<$n, f64>>>(&self, x: I) -> f64 {
                let expand_x = x.own().to_dolphinx();
                let f = self.cxx;
                unsafe {
                    // Find cell containing `x`
                    return ffi::eval_Function_f64(f, &expand_x);
                }
            }
        }
        //impl_diff!($n; 2 3 4 5 6 7 8);
    )+ };
}

impl_mapping!(1 2 3);
// Extension to 1D should be easy.
//impl_diff!(1; 2 3 4 5 6 7 8);
impl_diff!(2; 2 3 4 5 6 7 8);

impl<'py> GlobalAnalysis<f64, Bounds<f64>> for DolfinxPyFunction_f64<'py, 2, 2, 1> {
    fn global_analysis(&self) -> Bounds<f64> {
        let min =
            unsafe { ffi::min_Function_f64_p2(self.cxx) }.map_or(f64::NEG_INFINITY, |res| res.v);
        let max = unsafe { ffi::max_Function_f64_p2(self.cxx) }.map_or(f64::INFINITY, |res| res.v);
        Bounds(min, max)
    }
}

/// Only second-order polynomials on a triangular mesh are implemented.
///
/// May panic if assumptions of the underlying functionn do not hold despite
/// verification in the constructors of [`DolfinxPyFunction_f64`].
impl<'py> MinMaxMapping<Loc<2>, f64> for DolfinxPyFunction_f64<'py, 2, 2, 1> {
    fn minimise(&mut self, _tolerance: f64, _max_steps: usize) -> (Loc<2>, f64) {
        let res = unsafe { ffi::min_Function_f64_p2(self.cxx) }.unwrap();
        (res.x.into(), res.v)
    }

    fn maximise(&mut self, _tolerance: f64, _max_steps: usize) -> (Loc<2>, f64) {
        let res = unsafe { ffi::max_Function_f64_p2(self.cxx) }.unwrap();
        (res.x.into(), res.v)
    }
}

/*
enum MaybeValid<A> {
    Valid(A),
    Invalid(A),
}

use MaybeValid::*;

impl<A> MaybeValid<A> {
    #[inline]
    fn uninitialised_ok(self) -> A {
        match self {
            Valid(a) => a,
            Invalid(a) => a,
        }
    }
}
*/

impl<'py, const N: u32, const O: u32, const D: u32> DolfinxPyFunction_f64<'py, N, O, D> {
    fn new_cxx(&mut self, new: *mut Function_f64) {
        if let None = self.u {
            if self.cxx != std::ptr::null_mut() {
                unsafe { ffi::drop_Function_f64(self.cxx) };
            }
        }
        self.cxx = new;
        self.u = None;
        self.owned = true;
    }

    /// If owned, call `f_valid` with mutable data.
    /// If not owned, create a new C++ fem.Function, uninitialised, and call `f_invalid` with its
    /// mutable data, as well as the original immutable data.
    #[inline]
    fn with_mut_data_and_orig<F, G, Z>(&mut self, f_valid: F, f_invalid: G) -> Z
    where
        F: Fn(&mut [f64]) -> Z,
        G: Fn(&mut [f64], &[f64]) -> Z,
    {
        if self.owned {
            let res = f_valid(unsafe { ffi::data_mut_Function_f64(self.cxx) });
            unsafe { ffi::scatter_fwd_Function_f64(self.cxx) };
            return res;
        } else {
            // Create a completely new Function
            let old_data = unsafe { ffi::data_mut_Function_f64(self.cxx) };
            let new = unsafe { ffi::similar_Function_f64(self.cxx) };
            let new_data = unsafe { ffi::data_mut_Function_f64(new) };
            let res = f_invalid(new_data, old_data);
            unsafe { ffi::scatter_fwd_Function_f64(new) };
            // Invalidate Python reference; we're managing the object now
            self.new_cxx(new);
            return res;
        }
    }

    /// If owned, call `f` with mutable data.
    /// If not owned, create a new C++ fem.Function, uninitialised, and call `f` with its
    /// mutable data.
    #[inline]
    fn with_mut_data<F, Z>(&mut self, f: F) -> Z
    where
        F: Fn(&mut [f64]) -> Z,
    {
        if self.owned {
            let res = f(unsafe { ffi::data_mut_Function_f64(self.cxx) });
            unsafe { ffi::scatter_fwd_Function_f64(self.cxx) };
            return res;
        } else {
            // Create a completely new Function
            let new = unsafe { ffi::similar_Function_f64(self.cxx) };
            let new_data = unsafe { ffi::data_mut_Function_f64(new) };
            let res = f(new_data);
            unsafe { ffi::scatter_fwd_Function_f64(new) };
            // Invalidate Python reference; we're managing the object now
            self.new_cxx(new);
            return res;
        }
    }

    /// Create a similar new function, and call `f` with both the mutable (unininitialised)
    /// new data, and original data.
    ///
    /// Returns object for the new data.
    #[inline]
    fn with_new_data_and_orig<F>(&self, f: F) -> Self
    where
        F: Fn(&mut [f64], &[f64]),
    {
        // Create a completely new Function
        let new = self.similar();
        let old_data = unsafe { ffi::data_Function_f64(self.cxx) };
        let new_data = unsafe { ffi::data_mut_Function_f64(new.cxx) };
        f(new_data, old_data);
        unsafe { ffi::scatter_fwd_Function_f64(new.cxx) };
        new
    }

    /// Similar to [`Self::with_mut_data`], but check compatibility with `other` and also pass its
    /// immutable data slice as the final argument.
    #[inline]
    fn with_mut_data_and_orig_binop<F, G, Z>(
        &mut self,
        other: &DolfinxPyFunction_f64<'py, N, O, D>,
        f_valid: F,
        f_invalid: G,
    ) -> Z
    where
        F: Fn(&mut [f64], &[f64]) -> Z,
        G: Fn(&mut [f64], &[f64], &[f64]) -> Z,
    {
        let data_other = unsafe {
            ffi::check_compat_Function_f64(self.cxx, other.cxx).unwrap();
            ffi::data_Function_f64(other.cxx)
        };
        let res = self.with_mut_data_and_orig(
            |data_self_valid| {
                assert_eq!(data_self_valid.len(), data_other.len());
                f_valid(data_self_valid, data_other)
            },
            |data_self_invalid, data_orig| {
                assert_eq!(data_self_invalid.len(), data_other.len());
                assert_eq!(data_orig.len(), data_other.len());
                f_invalid(data_self_invalid, data_orig, data_other)
            },
        );
        unsafe { ffi::scatter_fwd_Function_f64(self.cxx) };
        return res;
    }

    /// Similar to [`Self::with_new_data_and_orig`], but check compatibility with `other` and also
    /// pass its immutable data slice as the final argument.
    #[inline]
    fn with_new_data_and_orig_binop<F>(
        &self,
        other: &DolfinxPyFunction_f64<'py, N, O, D>,
        f: F,
    ) -> Self
    where
        F: Fn(&mut [f64], &[f64], &[f64]),
    {
        let data_other = unsafe {
            ffi::check_compat_Function_f64(self.cxx, other.cxx).unwrap();
            ffi::data_Function_f64(other.cxx)
        };
        self.with_new_data_and_orig(|data, data_orig| {
            assert_eq!(data.len(), data_other.len());
            assert_eq!(data_orig.len(), data_other.len());
            f(data, data_orig, data_other)
        })
    }
}

/// Implements unary operations
macro_rules! impl_unop {
    ($trait:ident, $fn:ident) => {
        impl<'py, const N: u32, const O: u32, const D: u32> $trait
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = Self;

            fn $fn(mut self) -> Self {
                self.with_mut_data_and_orig(
                    |a| a.iter_mut().for_each(|x| *x = x.$fn()),
                    |a, a_orig| {
                        izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| {
                            *x = x_orig.$fn();
                        })
                    },
                );
                return self;
            }
        }

        impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait
            for &'a DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = DolfinxPyFunction_f64<'py, N, O, D>;

            fn $fn(self) -> Self::Output {
                self.with_new_data_and_orig(|a, a_orig| {
                    izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| {
                        *x = x_orig.$fn();
                    })
                })
            }
        }
    };
}

use std::ops::Neg;

impl_unop!(Neg, neg);

/// Implements scalar operations
macro_rules! impl_scalarop {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<'py, const N: u32, const O: u32, const D: u32> $trait<f64>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = Self;

            fn $fn(mut self, α: f64) -> Self {
                self.with_mut_data_and_orig(
                    |a| {
                        a.iter_mut().for_each(|x| {
                            x.$fn_assign(α);
                        })
                    },
                    |a, a_orig| {
                        izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| {
                            *x = x_orig.$fn(α);
                        })
                    },
                );
                return self;
            }
        }

        impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait<f64>
            for &'a DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = DolfinxPyFunction_f64<'py, N, O, D>;

            fn $fn(self, α: f64) -> Self::Output {
                self.with_new_data_and_orig(|a, a_orig| {
                    izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| {
                        *x = x_orig.$fn(α);
                    })
                })
            }
        }

        impl<'py, const N: u32, const O: u32, const D: u32> $trait_assign<f64>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            fn $fn_assign(&mut self, α: f64) {
                self.with_mut_data_and_orig(
                    |a| {
                        a.iter_mut().for_each(|x| {
                            x.$fn_assign(α);
                        })
                    },
                    |a, a_orig| {
                        izip!(a.iter_mut(), a_orig.iter()).for_each(|(x, x_orig)| {
                            *x = x_orig.$fn(α);
                        })
                    },
                );
            }
        }
    };
}

use std::ops::{Div, DivAssign, Mul, MulAssign};

impl_scalarop!(Mul, mul, MulAssign, mul_assign);
impl_scalarop!(Div, div, DivAssign, div_assign);

/// Implements consuming binary operations
macro_rules! impl_binop_consume {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<'py, const N: u32, const O: u32, const D: u32> $trait<Self>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = Self;

            #[inline]
            fn $fn(self, other: Self) -> Self {
                self.$fn(&other)
            }
        }

        impl<'py, const N: u32, const O: u32, const D: u32> $trait_assign<Self>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            #[inline]
            fn $fn_assign(&mut self, other: Self) {
                self.$fn_assign(&other)
            }
        }

        impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait<&'a Self>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = Self;

            fn $fn(mut self, other: &'a Self) -> Self {
                self.with_mut_data_and_orig_binop(
                    other,
                    |data, data_other| {
                        izip!(data.iter_mut(), data_other.iter()).for_each(|(x, y)| {
                            x.$fn_assign(y);
                        })
                    },
                    |data, data_orig, data_other| {
                        izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each(
                            |(x, x_orig, y)| {
                                *x = x_orig.$fn(y);
                            },
                        )
                    },
                );
                return self;
            }
        }

        impl<'py, 'a, const N: u32, const O: u32, const D: u32> $trait_assign<&'a Self>
            for DolfinxPyFunction_f64<'py, N, O, D>
        {
            fn $fn_assign(&mut self, other: &'a Self) {
                self.with_mut_data_and_orig_binop(
                    other,
                    |data, data_other| {
                        izip!(data.iter_mut(), data_other.iter()).for_each(|(x, y)| {
                            x.$fn_assign(y);
                        })
                    },
                    |data, data_orig, data_other| {
                        izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each(
                            |(x, x_orig, y)| {
                                *x = x_orig.$fn(y);
                            },
                        )
                    },
                );
            }
        }
    };
}

/// Implements binary operations on references
macro_rules! impl_binop_ref {
    ($trait:ident, $fn:ident) => {
        impl<'py, 'b, const N: u32, const O: u32, const D: u32> $trait<Self>
            for &'b DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = DolfinxPyFunction_f64<'py, N, O, D>;

            fn $fn(self, other: Self) -> Self::Output {
                self.$fn(&other)
            }
        }

        impl<'py, 'a, 'b, const N: u32, const O: u32, const D: u32> $trait<&'a Self>
            for &'b DolfinxPyFunction_f64<'py, N, O, D>
        {
            type Output = DolfinxPyFunction_f64<'py, N, O, D>;

            fn $fn(self, other: &'a Self) -> Self::Output {
                self.with_new_data_and_orig_binop(other, |data_dest, data, data_other| {
                    izip!(data_dest.iter_mut(), data.iter(), data_other.iter()).for_each(
                        |(x, y, z)| {
                            *x = y.$fn(z);
                        },
                    )
                })
            }
        }
    };
}

/// Implements binary operations, both consuming and on references
macro_rules! impl_binop {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl_binop_consume!($trait, $fn, $trait_assign, $fn_assign);
        impl_binop_ref!($trait, $fn);
    };
}

use std::ops::{Add, AddAssign, Sub, SubAssign};

impl_binop!(Add, add, AddAssign, add_assign);
impl_binop!(Sub, sub, SubAssign, sub_assign);

use alg_tools::linops::{VectorSpace, AXPY};
use alg_tools::self_ownable;

self_ownable!(DolfinxPyFunction_f64<'py, N, O, D> where 'py, const N: u32, const O: u32, const D: u32);

impl<'py, const N: u32, const O: u32, const D: u32> Dist<L2, f64>
    for DolfinxPyFunction_f64<'py, N, O, D>
{
    fn dist<I: Instance<Self>>(&self, other: I, p: L2) -> f64 {
        other.eval_ref(|v| (self - v).norm(p))
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> Norm<L2, f64>
    for DolfinxPyFunction_f64<'py, N, O, D>
{
    fn norm(&self, _p: L2) -> f64 {
        self.norm2_squared().sqrt()
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> VectorSpace
    for DolfinxPyFunction_f64<'py, N, O, D>
{
    type Field = f64;
    type PrincipalV = Self;

    /// Return a similar zero as `self`.
    fn similar_origin(&self) -> Self {
        let mut new = self.similar();
        new.set_zero();
        new
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> AXPY for DolfinxPyFunction_f64<'py, N, O, D> {
    /// Computes  `y = βy + αx`, where `y` is `Self`.
    fn axpy<I: Instance<Self>>(&mut self, α: Self::Field, x: I, β: Self::Field) {
        x.eval(|r| {
            self.with_mut_data_and_orig_binop(
                r,
                |data, data_other| {
                    izip!(data.iter_mut(), data_other.iter()).for_each(|(y, x)| {
                        *y = β * (*y) + α * (*x);
                    })
                },
                |data, data_orig, data_other| {
                    izip!(data.iter_mut(), data_orig.iter(), data_other.iter()).for_each(
                        |(y, y_orig, x)| {
                            *y = β * (*y_orig) + α * (*x);
                        },
                    )
                },
            );
        })
    }

    /// Set self to zero.
    fn set_zero(&mut self) {
        self.with_mut_data(|data| {
            data.iter_mut().for_each(|y| *y = 0.0);
        })
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> HasDual
    for DolfinxPyFunction_f64<'py, N, O, D>
{
    type DualSpace = Self;

    fn dual_origin(&self) -> <Self::DualSpace as VectorSpace>::PrincipalV {
        self.similar_origin()
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> Normed for DolfinxPyFunction_f64<'py, N, O, D> {
    type NormExp = L2;

    fn norm_exponent(&self) -> L2 {
        // TODO: maybe H1 more approriate for O=2
        L2
    }
}

impl<'py, const N: u32, const O: u32, const D: u32> Euclidean
    for DolfinxPyFunction_f64<'py, N, O, D>
{
    type PrincipalE = Self;

    fn dot<I: Instance<Self>>(&self, other: I) -> f64 {
        let py = self.function_space.py();
        get_helpers(py)
            .unwrap()
            .dot
            .call1(py, (self.own(), other.own()))
            .unwrap()
            .extract(py)
            .unwrap()
        /*
        let data = unsafe { ffi::data_mut_Function_f64(self.cxx) };
        other.eval_decompose(|x| {
            let data_other = unsafe { ffi::data_mut_Function_f64(x.cxx) };
            izip!(data.iter(), data_other.iter())
                .map(|(y, x)| y * x)
                .sum()
        })*/
    }

    fn norm2_squared(&self) -> f64 {
        let py = self.function_space.py();
        get_helpers(py)
            .unwrap()
            .norm2_squared
            .call1(py, (self.own(),))
            .unwrap()
            .extract(py)
            .unwrap()
    }

    fn dist2_squared<I: Instance<Self>>(&self, other: I) -> f64 {
        //other.eval_ref(|v| self.norm2_squared() + 2.0 * self.dot(v) + v.norm2_squared_div2())
        let py = self.function_space.py();
        get_helpers(py)
            .unwrap()
            .dist2_squared
            .call1(py, (self.own(), other.own()))
            .unwrap()
            .extract(py)
            .unwrap()
    }
}