liblaf.apple.region.tetra

Classes:

• RegionTetra –

RegionTetra

Methods:

• fill –
• fun –
• fun_jac –
• hess –
• hess_diag –
• jac –
• prepare –

Attributes:

• aux (PyTree) –
• cells (Integer[Array, 'c a']) –
• dirichlet_mask (Bool[ndarray, 'V 3']) –
• dirichlet_values (Float[Array, 'V 3']) –
• free_mask (Bool[ndarray, 'V 3']) –
• material (MaterialTetra) –
• mesh (UnstructuredGrid) –
• n_cells (int) –
• n_dof (int) –
• n_points (int) –
• params (PyTree) –
• points (Float[Array, 'V 3']) –

aux class-attribute instance-attribute

aux: PyTree = field(factory=dict)

cells property

cells: Integer[Array, 'c a']

dirichlet_mask class-attribute instance-attribute

dirichlet_mask: Bool[ndarray, "V 3"] = field(
    metadata={"static": True}, converter=asarray
)

dirichlet_values class-attribute instance-attribute

dirichlet_values: Float[Array, "V 3"] = field(
    converter=asarray
)

free_mask property

free_mask: Bool[ndarray, 'V 3']

material class-attribute instance-attribute

material: MaterialTetra = field(metadata={'static': True})

mesh class-attribute instance-attribute

mesh: UnstructuredGrid = field(metadata={'static': True})

n_cells property

n_cells: int

n_dof property

n_dof: int

n_points property

n_points: int

params class-attribute instance-attribute

params: PyTree = field(factory=dict)

points property

points: Float[Array, 'V 3']

fill

fill(u_free: Float[Array, ' F']) -> Float[Array, 'V 3']

Source code in src/liblaf/apple/region/tetra.py

def fill(self, u_free: Float[jax.Array, " F"]) -> Float[jax.Array, "V 3"]:
    u: Float[jax.Array, "V 3"] = self.dirichlet_values.copy()
    u = u.at[self.free_mask].set(u_free)
    return u

fun

fun(u: Float[Array, ' F'], q: PyTree) -> Float[Array, '']

Source code in src/liblaf/apple/region/tetra.py

@apple.jit()
def fun(self, u: Float[jax.Array, " F"], q: PyTree) -> Float[jax.Array, ""]:
    u: Float[jax.Array, "V 3"] = self.fill(u)
    return self.material.fun(u[self.cells], q, self.aux)

fun_jac

fun_jac(
    u: Float[Array, " F"], q: PyTree
) -> tuple[Float[Array, ""], Float[Array, " F"]]

Source code in src/liblaf/apple/region/tetra.py

@apple.jit()
def fun_jac(
    self, u: Float[jax.Array, " F"], q: PyTree
) -> tuple[Float[jax.Array, ""], Float[jax.Array, " F"]]:
    u: Float[jax.Array, "V 3"] = self.fill(u)
    fun: Float[jax.Array, ""]
    jac: Float[jax.Array, "C 4 3"]
    fun, jac = self.material.fun_jac(u[self.cells], q, self.aux)
    jac = apple.elem.tetra.segment_sum(jac, self.cells, self.n_points)
    return fun, jac[self.free_mask]

hess

hess(
    u: Float[Array, " F"], q: PyTree
) -> Float[LinearOperator, "F F"]

Source code in src/liblaf/apple/region/tetra.py

def hess(
    self, u: Float[jax.Array, " F"], q: PyTree
) -> Float[pylops.LinearOperator, "F F"]:
    u: Float[jax.Array, "V 3"] = self.fill(u)
    hess: Float[jax.Array, "C 4 3 4 3"] = self.material.hess(
        u[self.cells], q, self.aux
    )

    def matvec(v: Float[jax.Array, " F"]) -> Float[jax.Array, " F"]:
        v: Float[jax.Array, " F"] = jnp.asarray(v, dtype=hess.dtype)
        return self._hvp(hess, v)

    return pylops.FunctionOperator(
        matvec, matvec, self.n_dof, self.n_dof, dtype=hess.dtype, name="hessian"
    )

hess_diag

hess_diag(
    u: Float[Array, " F"], q: PyTree
) -> Float[Array, " F"]

Source code in src/liblaf/apple/region/tetra.py

@apple.jit()
def hess_diag(self, u: Float[jax.Array, " F"], q: PyTree) -> Float[jax.Array, " F"]:
    u: Float[jax.Array, "V 3"] = self.fill(u)
    hess_diag: Float[jax.Array, "C 4 3"] = self.material.hess_diag(
        u[self.cells], q, self.aux
    )
    hess_diag: Float[jax.Array, "V 3"] = apple.elem.tetra.segment_sum(
        hess_diag, self.cells, self.n_points
    )
    return hess_diag[self.free_mask]

jac

jac(u: Float[Array, ' F'], q: PyTree) -> Float[Array, ' F']

Source code in src/liblaf/apple/region/tetra.py

@apple.jit()
def jac(self, u: Float[jax.Array, " F"], q: PyTree) -> Float[jax.Array, " F"]:
    u: Float[jax.Array, "V 3"] = self.fill(u)
    jac: Float[jax.Array, "C 3"] = self.material.jac(u[self.cells], q, self.aux)
    jac: Float[jax.Array, "V 3"] = apple.elem.tetra.segment_sum(
        jac, self.cells, self.n_points
    )
    return jac[self.free_mask]

prepare

prepare() -> None

Source code in src/liblaf/apple/region/tetra.py

def prepare(self) -> None:
    self.aux = {
        "dh_dX": apple.elem.tetra.dh_dX(self.points[self.cells]),
        "dV": apple.elem.tetra.dV(self.points[self.cells]),
    }