Add TrilinearForm

[kinnala]

In #732 the following snippet was devised to assemble trilinear forms:

from typing import Optional, Tuple, Any

import numpy as np
from numpy import ndarray

from skfem import *
from skfem.assembly.form.form import Form
from skfem.assembly.form.form import FormExtraParams

class TrilinearForm(Form):

    def _assemble(self,
                  ubasis: Basis,
                  vbasis: Optional[Basis] = None,
                  wbasis: Optional[Basis] = None,
                  **kwargs) -> Tuple[ndarray,
                                     ndarray,
                                     ndarray,
                                     ndarray,
                                     Tuple[int, int, int]]:

        if vbasis is None:
            vbasis = ubasis
        if wbasis is None:
            wbasis = ubasis

        nt = ubasis.nelems
        dx = ubasis.dx
        wdict = FormExtraParams({
            **ubasis.default_parameters(),
            **self.dictify(kwargs, ubasis),
        })

        # initialize COO data structures
        sz = (ubasis.Nbfun, vbasis.Nbfun, wbasis.Nbfun, nt)
        data = np.zeros(sz, dtype=self.dtype)
        rows = np.zeros(sz, dtype=np.int64)
        cols = np.zeros(sz, dtype=np.int64)
        mats = np.zeros(sz, dtype=np.int64)

        # loop over the indices of local stiffness matrix
        for k in range(ubasis.Nbfun):
            for j in range(vbasis.Nbfun):
                for i in range(wbasis.Nbfun):
                    mats[k, j, i] = wbasis.element_dofs[i]
                    rows[k, j, i] = vbasis.element_dofs[j]
                    cols[k, j, i] = ubasis.element_dofs[k]
                    data[k, j, i] = self._kernel(
                        ubasis.basis[k],
                        vbasis.basis[j],
                        wbasis.basis[i],
                        wdict,
                        dx,
                    )

        return data, mats, rows, cols, (wbasis.N, vbasis.N, ubasis.N)

    def assemble(self, *args, **kwargs) -> Any:
        return self._assemble(*args, **kwargs)

    def _kernel(self, u, v, w, params, dx):
        return np.sum(self.form(*u, *v, *w, params) * dx, axis=1)

def asm():
    m = MeshTri().refined(7)
    basis = Basis(m, ElementTriP1())
    from skfem.helpers import dot, grad
    out = TrilinearForm(lambda u, v, w, p: dot(w * grad(u), grad(v))).assemble(basis)
    import sparse
    arr = sparse.COO(np.array([out[1].flatten(), out[2].flatten(), out[3].flatten()]), out[0].flatten(), out[4], has_duplicates=True)
    k_p = np.random.rand(100, basis.N)
    res = sparse.tensordot(arr, k_p.T, axes=1, return_type=sparse.COO)

    A = []
    for mat in res.T:
        A.append(mat.to_scipy_sparse().tocsr())
    return A

Adding this to scikit-fem is feasible but I think the output type above (variable out) is not optimal. It should be directly compatible with sparse.COO constructor and other sparse tensor libraries (tensorly, pytorch.sparse, tf.SparseTensor and so on). This means that the indices should be returned as a single ndims x ndata array.

Some reasonable tests must be figured out. Above the idea is to assemble 100 forms with different field. So therefore we assemble only 1 trilinear form and perform tensor dot products using sparse package.

[kinnala]

Let me add that this concept and approach quite easily generalizes to NLinearForm but I don't know if there are any use cases?

[AnyWayOne]

This will be great useful for parameter analysis and optimization design.