`ConstitutiveMaterial`: Threaded evaluation of the gradient/hessian of the strain energy function about felupe HOT 2 CLOSED

import felupe as fem
import numpy as np
from threading import Thread

mesh = fem.Cube(n=51)
region = fem.RegionHexahedron(mesh)
field = fem.FieldContainer([fem.Field(region, dim=3)])

boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)


class Threaded(fem.constitution.ConstitutiveMaterial):
    def __init__(self, material, axis, nthreads, **kwargs):
        self.material = material
        self.x = self.material.x
        self.axis = axis
        self._axes = [slice(None)] * abs(self.axis + 1)
        self.nthreads = nthreads
        self._P = None
        self._A = None
        for key, value in kwargs.items():
            setattr(self, key, value)

    def _kernel_grad(self, F, chunk):
        self.material.gradient([F, None], out=self._P[..., slice(*chunk), *self._axes])

    def _kernel_hess(self, F, chunk):
        self.material.hessian([F, None], out=self._A[..., slice(*chunk), *self._axes])

    def chunks(self, F):
        defgrads = np.array_split(F, self.nthreads, axis=self.axis)
        chunksizes = np.cumsum([chunk.shape[self.axis] for chunk in defgrads])
        chunks = np.vstack([np.append(0, chunksizes[:-1]), chunksizes]).T
        return defgrads, chunksizes, chunks

    def evaluate(self, target, defgrads, chunks):
        threads = [
            Thread(target=target, args=(F, chunk)) for F, chunk in zip(defgrads, chunks)
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

    def gradient(self, x):
        F, statevars = x
        defgrads, chunks, chunksizes = self.chunks(F)

        if self._P is None:
            self._P = np.zeros_like(F)

        self.evaluate(self._kernel_grad, defgrads, chunks)
        return [self._P, None]

    def hessian(self, x):
        F, statevars = x
        defgrads, chunks, chunksizes = self.chunks(F)

        if self._A is None:
            self._A = np.zeros((*F.shape[:2], *F.shape))

        self.evaluate(self._kernel_hess, defgrads, chunks)
        return [self._A]


umat_threaded = Threaded(fem.NeoHooke(mu=1, bulk=2), axis=-1, nthreads=16)
solid_threaded = fem.SolidBody(umat_threaded, field)

umat = fem.NeoHooke(mu=1, bulk=2)
solid = fem.SolidBody(umat, field)