liblaf.apple.energy ¤

Modules:

collision –
elastic –
zero –

Classes:

ARAP –

As-Rigid-As-Possible.
CollisionCandidatesVertFace –
CollisionVertFace –
EnergyZero –
PhaceStatic –

As-Rigid-As-Possible.

ARAP ¤

Bases: Elastic

As-Rigid-As-Possible.

\[ \Psi = \frac{\mu}{2} \|F - R\|_F^2 = \frac{\mu}{2} (I_2 - 2 I_1 + 3) \]

Parameters:

id (str, default: <dynamic> ) –
actor (Actor) –
hess_diag_filter (bool, default: True ) –
hess_quad_filter (bool, default: True ) –

Methods:

__pdoc__ –

...
__repr__ –
energy_density –
energy_density_hess_diag –
energy_density_hess_quad –
energy_density_jac –
evolve –
first_piola_kirchhoff_stress –
from_actor –
fun –
fun_and_jac –
hess_diag –
hess_quad –
hessp –
jac –
jac_and_hess_diag –
make_field –
pre_optim_iter –
pre_time_step –
tree_at –
with_actors –

Attributes:

actor (Actor) –
actors (NodeContainer[Actor]) –
hess_diag_filter (bool) –
hess_quad_filter (bool) –
id (str) –
mu (Float[Array, ' c']) –

actor `class-attribute` `instance-attribute` ¤

actor: Actor = data()

actors `abstractmethod` `property` ¤

actors: NodeContainer[Actor]

hess_diag_filter `class-attribute` `instance-attribute` ¤

hess_diag_filter: bool = static(default=True, kw_only=True)

hess_quad_filter `class-attribute` `instance-attribute` ¤

hess_quad_filter: bool = static(default=True, kw_only=True)

id `class-attribute` `instance-attribute` ¤

id: str = static(
    default=Factory(uniq_id, takes_self=True), kw_only=True
)

mu `property` ¤

mu: Float[Array, ' c']

pdoc ¤

__pdoc__(**kwargs) -> AbstractDoc

...

References

wadler_lindig._definitions._pformat_dataclass()

Source code in src/liblaf/apple/energy/elastic/arap.py

$$
\Psi = \frac{\mu}{2} \|F - R\|_F^2 = \frac{\mu}{2} (I_2 - 2 I_1 + 3)
$$
"""

@property
def mu(self) -> Float[jax.Array, " c"]:
    return self.actor.cell_data["mu"]

@override
@utils.jit_method(inline=True)
def energy_density(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    Psi: Float[jax.Array, " cq"]
    (Psi,) = arap_energy_density_warp(F, self.mu)
    Psi: Float[jax.Array, "c q"] = region.unsqueeze_cq(Psi)
    return Psi

@override
@utils.jit_method(inline=True)
def first_piola_kirchhoff_stress(

repr ¤

__repr__() -> str

Source code in src/liblaf/apple/energy/elastic/arap.py

) -> Float[jax.Array, "c q J J"]:
    region: sim.Region = field.region

energy_density ¤

energy_density(
    field: Field, /, params: GlobalParams
) -> Float[Array, "c q"]

Source code in src/liblaf/apple/energy/elastic/arap.py

@override
@utils.jit_method(inline=True)
def energy_density(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    Psi: Float[jax.Array, " cq"]
    (Psi,) = arap_energy_density_warp(F, self.mu)
    Psi: Float[jax.Array, "c q"] = region.unsqueeze_cq(Psi)
    return Psi

energy_density_hess_diag ¤

energy_density_hess_diag(
    field: Field, /, params: GlobalParams
) -> Float[Array, "c q a J"]

Source code in src/liblaf/apple/energy/elastic/arap.py

@override
@utils.jit_method(inline=True)
def energy_density_hess_diag(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q a J"]:
    hess_diag: Float[jax.Array, "cells 4 3"]
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    dhdX: Float[jax.Array, "cq a J"] = region.squeeze_cq(region.dhdX)
    hess_diag: Float[jax.Array, "cq a J"]
    (hess_diag,) = arap_energy_density_hess_diag_warp(F, self.mu, dhdX)
    hess_diag: Float[jax.Array, "c q a J"] = region.unsqueeze_cq(hess_diag)
    return hess_diag

energy_density_hess_quad ¤

energy_density_hess_quad(
    field: Field, p: Field, /, params: GlobalParams
) -> Float[Array, "c q"]

Source code in src/liblaf/apple/energy/elastic/arap.py

@override
@utils.jit_method(inline=True)
def energy_density_hess_quad(
    self, field: sim.Field, p: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    dhdX: Float[jax.Array, "cq a J"] = region.squeeze_cq(region.dhdX)
    hess_quad: Float[jax.Array, " cq"]
    (hess_quad,) = arap_energy_density_hess_quad_warp(
        F, region.scatter(p.values), self.mu, dhdX
    )
    hess_quad: Float[jax.Array, "c q"] = region.unsqueeze_cq(hess_quad)
    return hess_quad

energy_density_jac ¤

energy_density_jac(
    field: Field, /, params: GlobalParams
) -> Float[Array, "c q a J"]

Source code in src/liblaf/apple/energy/elastic/arap.py

@no_type_check
@utils.jax_kernel
def arap_energy_density_hess_quad_warp(
    F: wp.array(dtype=mat33),
    p: wp.array(dtype=mat43),
    mu: wp.array(dtype=float),
    dh_dX: wp.array(dtype=mat43),
    hess_quad: wp.array(dtype=float),
) -> None:

evolve ¤

evolve(**changes) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
PK1: Float[jax.Array, "cq J J"]

first_piola_kirchhoff_stress ¤

first_piola_kirchhoff_stress(
    field: Field, /, params: GlobalParams
) -> Float[Array, "c q J J"]

Source code in src/liblaf/apple/energy/elastic/arap.py

@override
@utils.jit_method(inline=True)
def first_piola_kirchhoff_stress(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q J J"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    PK1: Float[jax.Array, "cq J J"]
    (PK1,) = arap_first_piola_kirchhoff_stress_warp(F, self.mu)
    PK1: Float[jax.Array, "c q J J"] = region.unsqueeze_cq(PK1)
    return PK1

from_actor `classmethod` ¤

from_actor(
    actor: Actor,
    *,
    hess_diag_filter: bool = True,
    hess_quad_filter: bool = True,
) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

$$
\Psi = \frac{\mu}{2} \|F - R\|_F^2 = \frac{\mu}{2} (I_2 - 2 I_1 + 3)
$$
"""

@property
def mu(self) -> Float[jax.Array, " c"]:
    return self.actor.cell_data["mu"]

@override
@utils.jit_method(inline=True)
def energy_density(

fun ¤

fun(
    x: ArrayDict, /, params: GlobalParams
) -> Float[Array, ""]

Source code in src/liblaf/apple/energy/elastic/arap.py

    (Psi,) = arap_energy_density_warp(F, self.mu)
    Psi: Float[jax.Array, "c q"] = region.unsqueeze_cq(Psi)
    return Psi

@override
@utils.jit_method(inline=True)
def first_piola_kirchhoff_stress(

fun_and_jac ¤

fun_and_jac(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[Float[Array, ""], ArrayDict]

Source code in src/liblaf/apple/energy/elastic/arap.py

    F, region.scatter(p.values), self.mu, dhdX
)
hess_quad: Float[jax.Array, "c q"] = region.unsqueeze_cq(hess_quad)
return hess_quad

hess_diag ¤

hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/elastic/arap.py

    dhdX: Float[jax.Array, "cq a J"] = region.squeeze_cq(region.dhdX)
    hess_diag: Float[jax.Array, "cq a J"]
    (hess_diag,) = arap_energy_density_hess_diag_warp(F, self.mu, dhdX)
    hess_diag: Float[jax.Array, "c q a J"] = region.unsqueeze_cq(hess_diag)
    return hess_diag

@override
@utils.jit_method(inline=True)
def energy_density_hess_quad(

hess_quad ¤

hess_quad(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> Float[Array, ""]

Source code in src/liblaf/apple/energy/elastic/arap.py

) -> Float[jax.Array, "c q"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    dhdX: Float[jax.Array, "cq a J"] = region.squeeze_cq(region.dhdX)
    hess_quad: Float[jax.Array, " cq"]

hessp ¤

hessp(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/elastic/arap.py

def energy_density_hess_diag(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q a J"]:
    hess_diag: Float[jax.Array, "cells 4 3"]
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)

jac ¤

jac(x: ArrayDict, /, params: GlobalParams) -> ArrayDict

Source code in src/liblaf/apple/energy/elastic/arap.py

) -> Float[jax.Array, "c q J J"]:
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    PK1: Float[jax.Array, "cq J J"]
    (PK1,) = arap_first_piola_kirchhoff_stress_warp(F, self.mu)
    PK1: Float[jax.Array, "c q J J"] = region.unsqueeze_cq(PK1)
    return PK1

@override

jac_and_hess_diag ¤

jac_and_hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[ArrayDict, ArrayDict]

Source code in src/liblaf/apple/energy/elastic/arap.py

@utils.jax_kernel
def arap_energy_density_warp(
    F: wp.array(dtype=mat33), mu: wp.array(dtype=float), Psi: wp.array(dtype=float)
) -> None:
    tid = wp.tid()
    Psi[tid] = func.elastic.arap_energy_density(F=F[tid], mu=mu[tid])

make_field ¤

make_field(x: ArrayDict) -> Field

Source code in src/liblaf/apple/energy/elastic/arap.py

@override
@utils.jit_method(inline=True)
def first_piola_kirchhoff_stress(

pre_optim_iter ¤

pre_optim_iter(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

return self.actor.cell_data["mu"]

pre_time_step ¤

pre_time_step(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

21	`@property`

tree_at ¤

tree_at(
    where: Callable[[Self], Node | Sequence[Node]],
    replace: Any | Sequence[Any] = ...,
    replace_fn: Callable[[Node], Any] = ...,
    is_leaf: Callable[[Any], bool] | None = None,
) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

    PK1: Float[jax.Array, "c q J J"] = region.unsqueeze_cq(PK1)
    return PK1

@override
@utils.jit_method(inline=True)
def energy_density_hess_diag(
    self, field: sim.Field, /, params: sim.GlobalParams
) -> Float[jax.Array, "c q a J"]:
    hess_diag: Float[jax.Array, "cells 4 3"]
    region: sim.Region = field.region
    F: Float[jax.Array, "c q J J"] = region.deformation_gradient(field.values)
    F: Float[jax.Array, "cq J J"] = region.squeeze_cq(F)
    dhdX: Float[jax.Array, "cq a J"] = region.squeeze_cq(region.dhdX)
    hess_diag: Float[jax.Array, "cq a J"]
    (hess_diag,) = arap_energy_density_hess_diag_warp(F, self.mu, dhdX)

with_actors `abstractmethod` ¤

with_actors(actors: NodeContainer[Actor]) -> Self

Source code in src/liblaf/apple/energy/elastic/arap.py

@utils.jit_method(inline=True)
def energy_density(
    self, field: sim.Field, /, params: sim.GlobalParams

CollisionCandidatesVertFace ¤

Bases: PyTreeMixin

Parameters:

closest (Float[Array, 'points 3'], default: None ) –
collide (Bool[Array, points], default: None ) –
distance (Float[Array, points], default: None ) –
face_id (Integer[Array, points], default: None ) –
face_normal (Float[Array, 'points 3'], default: None ) –
uv (Float[Array, 'points 2'], default: None ) –

Methods:

__pdoc__ –

...
__repr__ –
evolve –
tree_at –

Attributes:

closest (Float[Array, 'points 3']) –
collide (Bool[Array, ' points']) –
distance (Float[Array, ' points']) –
face_id (Integer[Array, ' points']) –
face_normal (Float[Array, 'points 3']) –
uv (Float[Array, 'points 2']) –

closest `class-attribute` `instance-attribute` ¤

closest: Float[Array, 'points 3'] = array(default=None)

collide `class-attribute` `instance-attribute` ¤

collide: Bool[Array, ' points'] = array(default=None)

distance `class-attribute` `instance-attribute` ¤

distance: Float[Array, ' points'] = array(default=None)

face_id `class-attribute` `instance-attribute` ¤

face_id: Integer[Array, ' points'] = array(default=None)

face_normal `class-attribute` `instance-attribute` ¤

face_normal: Float[Array, 'points 3'] = array(default=None)

uv `class-attribute` `instance-attribute` ¤

uv: Float[Array, 'points 2'] = array(default=None)

pdoc ¤

__pdoc__(**kwargs) -> AbstractDoc

...

References

wadler_lindig._definitions._pformat_dataclass()

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@struct.pytree
class CollisionCandidatesVertFace(struct.PyTreeMixin):
    closest: Float[Array, "points 3"] = struct.array(default=None)
    collide: Bool[Array, " points"] = struct.array(default=None)
    distance: Float[Array, " points"] = struct.array(default=None)
    face_id: Integer[Array, " points"] = struct.array(default=None)
    face_normal: Float[Array, "points 3"] = struct.array(default=None)
    uv: Float[Array, "points 2"] = struct.array(default=None)


@struct.pytree
class CollisionVertFace(sim.Energy):
    rigid: sim.Actor = struct.data(default=None)
    soft: sim.Actor = struct.data(default=None)

    stiffness: Float[Array, ""] = struct.array(default=1e5)
    rest_length: Float[Array, ""] = struct.array(default=1e-3)
    max_dist: Float[Array, ""] = struct.array(default=1e-2)
    epsilon: Float[Array, ""] = struct.array(default=1e-3)
    filter_hess_diag: bool = struct.static(default=True, kw_only=True)
    filter_hess_quad: bool = struct.static(default=True, kw_only=True)

    candidates: CollisionCandidatesVertFace = struct.data(
        factory=CollisionCandidatesVertFace

repr ¤

__repr__() -> str

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

43	`@classmethod`

evolve ¤

evolve(**changes) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

cls,
rigid: sim.Actor,

tree_at ¤

tree_at(
    where: Callable[[Self], Node | Sequence[Node]],
    replace: Any | Sequence[Any] = ...,
    replace_fn: Callable[[Node], Any] = ...,
    is_leaf: Callable[[Any], bool] | None = None,
) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

    *,
    stiffness: float = 1e3,
    rest_length: float = 1e-3,
    max_dist: float | None = None,
    epsilon: float = 1e-3,
) -> Self:
    if max_dist is None:
        max_dist = 2.0 * rest_length
    return cls(
        rigid=rigid,
        soft=soft,
        stiffness=jnp.asarray(stiffness),
        rest_length=jnp.asarray(rest_length),
        max_dist=jnp.asarray(max_dist),
        epsilon=jnp.asarray(epsilon),

CollisionVertFace ¤

Bases: Energy

Parameters:

id (str, default: <dynamic> ) –
rigid (Actor, default: None ) –
soft (Actor, default: None ) –
stiffness (Float[Array, ''], default: 100000.0 ) –
rest_length (Float[Array, ''], default: 0.001 ) –
max_dist (Float[Array, ''], default: 0.01 ) –
epsilon (Float[Array, ''], default: 0.001 ) –
filter_hess_diag (bool, default: True ) –
filter_hess_quad (bool, default: True ) –
candidates (CollisionCandidatesVertFace, default: <dynamic> ) –

Methods:

__pdoc__ –

...
__repr__ –
collide –
evolve –
from_actors –
fun –
fun_and_jac –
hess_diag –
hess_quad –
hessp –
jac –
jac_and_hess_diag –
pre_optim_iter –
pre_time_step –
tree_at –
with_actors –

Attributes:

actors (NodeContainer[Actor]) –
candidates (CollisionCandidatesVertFace) –
epsilon (Float[Array, '']) –
filter_hess_diag (bool) –
filter_hess_quad (bool) –
id (str) –
max_dist (Float[Array, '']) –
rest_length (Float[Array, '']) –
rigid (Actor) –
soft (Actor) –
stiffness (Float[Array, '']) –

actors `property` ¤

actors: NodeContainer[Actor]

candidates `class-attribute` `instance-attribute` ¤

candidates: CollisionCandidatesVertFace = data(
    factory=CollisionCandidatesVertFace
)

epsilon `class-attribute` `instance-attribute` ¤

epsilon: Float[Array, ''] = array(default=0.001)

filter_hess_diag `class-attribute` `instance-attribute` ¤

filter_hess_diag: bool = static(default=True, kw_only=True)

filter_hess_quad `class-attribute` `instance-attribute` ¤

filter_hess_quad: bool = static(default=True, kw_only=True)

id `class-attribute` `instance-attribute` ¤

id: str = static(
    default=Factory(uniq_id, takes_self=True), kw_only=True
)

max_dist `class-attribute` `instance-attribute` ¤

max_dist: Float[Array, ''] = array(default=0.01)

rest_length `class-attribute` `instance-attribute` ¤

rest_length: Float[Array, ''] = array(default=0.001)

rigid `class-attribute` `instance-attribute` ¤

rigid: Actor = data(default=None)

soft `class-attribute` `instance-attribute` ¤

soft: Actor = data(default=None)

stiffness `class-attribute` `instance-attribute` ¤

stiffness: Float[Array, ''] = array(default=100000.0)

pdoc ¤

__pdoc__(**kwargs) -> AbstractDoc

...

References

wadler_lindig._definitions._pformat_dataclass()

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@struct.pytree
class CollisionCandidatesVertFace(struct.PyTreeMixin):
    closest: Float[Array, "points 3"] = struct.array(default=None)
    collide: Bool[Array, " points"] = struct.array(default=None)
    distance: Float[Array, " points"] = struct.array(default=None)
    face_id: Integer[Array, " points"] = struct.array(default=None)
    face_normal: Float[Array, "points 3"] = struct.array(default=None)
    uv: Float[Array, "points 2"] = struct.array(default=None)


@struct.pytree
class CollisionVertFace(sim.Energy):
    rigid: sim.Actor = struct.data(default=None)
    soft: sim.Actor = struct.data(default=None)

    stiffness: Float[Array, ""] = struct.array(default=1e5)
    rest_length: Float[Array, ""] = struct.array(default=1e-3)
    max_dist: Float[Array, ""] = struct.array(default=1e-2)
    epsilon: Float[Array, ""] = struct.array(default=1e-3)
    filter_hess_diag: bool = struct.static(default=True, kw_only=True)
    filter_hess_quad: bool = struct.static(default=True, kw_only=True)

    candidates: CollisionCandidatesVertFace = struct.data(
        factory=CollisionCandidatesVertFace

repr ¤

__repr__() -> str

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

43	`@classmethod`

collide ¤

collide() -> CollisionCandidatesVertFace

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@utils.jit_method()
def collide(self) -> CollisionCandidatesVertFace:
    (
        closest,
        collide,
        distance,
        face_id,
        face_normal,
        uv,
    ) = collision_detect_vert_face_kernel(
        self.soft.positions,
        np.uint64(self.rigid.collision_mesh.id),
        self.rest_length.reshape((1,)),
        self.max_dist.reshape((1,)),
        self.epsilon.reshape((1,)),
        output_dims={
            "closest": (self.soft.n_points,),
            "collide": (self.soft.n_points,),
            "distance": (self.soft.n_points,),
            "face_id": (self.soft.n_points,),
            "face_normal": (self.soft.n_points,),
            "uv": (self.soft.n_points,),
        },
        launch_dims=(self.soft.n_points,),
    )
    return CollisionCandidatesVertFace(
        closest=closest,
        collide=collide,
        distance=distance,
        face_id=face_id,
        face_normal=face_normal,
        uv=uv,
    )

evolve ¤

evolve(**changes) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

cls,
rigid: sim.Actor,

from_actors `classmethod` ¤

from_actors(
    rigid: Actor,
    soft: Actor,
    *,
    stiffness: float = 1000.0,
    rest_length: float = 0.001,
    max_dist: float | None = None,
    epsilon: float = 0.001,
) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@classmethod
def from_actors(
    cls,
    rigid: sim.Actor,
    soft: sim.Actor,
    *,
    stiffness: float = 1e3,
    rest_length: float = 1e-3,
    max_dist: float | None = None,
    epsilon: float = 1e-3,
) -> Self:
    if max_dist is None:
        max_dist = 2.0 * rest_length
    return cls(
        rigid=rigid,
        soft=soft,
        stiffness=jnp.asarray(stiffness),
        rest_length=jnp.asarray(rest_length),
        max_dist=jnp.asarray(max_dist),
        epsilon=jnp.asarray(epsilon),
    )

fun ¤

fun(
    x: ArrayDict, /, params: GlobalParams
) -> Float[Array, ""]

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
@utils.jit_method(inline=True)
def fun(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> Float[Array, ""]:
    points: Float[Array, "points dim"] = self.soft.points + x[self.soft.id]
    energy: Float[Array, " points"]
    (energy,) = collision_energy_vert_face_fun_kernel(
        points,
        self.candidates.closest,
        self.candidates.collide,
        self.candidates.distance,
        self.rest_length.reshape((1,)),
        self.stiffness.reshape((1,)),
        output_dims={"energy": (1,)},
        launch_dims=(self.soft.n_points,),
    )
    return energy.sum()

fun_and_jac ¤

fun_and_jac(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[Float[Array, ""], ArrayDict]

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

    # wl.pprint(candidates, short_arrays=False)
    return self.evolve(candidates=candidates)

@override
@utils.jit_method(inline=True)
def fun(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> Float[Array, ""]:

hess_diag ¤

hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
@utils.jit_method(inline=True)
def hess_diag(
    self, x: struct.ArrayDict, /, params: sim.GlobalParams
) -> struct.ArrayDict:
    points: Float[Array, "points dim"] = self.soft.points + x[self.soft.id]
    hess_diag: Float[Array, "points dim"]
    (hess_diag,) = collision_energy_vert_face_hess_diag_kernel(
        points,
        self.candidates.closest,
        self.candidates.collide,
        self.candidates.distance,
        self.rest_length.reshape((1,)),
        self.stiffness.reshape((1,)),
        output_dims={"hess_diag": (self.soft.n_points,)},
        launch_dims=(self.soft.n_points,),
    )
    if self.filter_hess_diag:
        hess_diag = jnp.clip(hess_diag, min=0.0)
    return struct.ArrayDict(
        {self.soft.id: hess_diag, self.rigid.id: jnp.zeros_like(x[self.rigid.id])}
    )

hess_quad ¤

hess_quad(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> Float[Array, ""]

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
@utils.jit_method(inline=True)
def hess_quad(
    self, x: struct.ArrayDict, p: struct.ArrayDict, /, params: sim.GlobalParams
) -> Float[Array, ""]:
    points: Float[Array, "points dim"] = self.soft.points + x[self.soft.id]
    hess_quad: Float[Array, " points"]
    (hess_quad,) = collision_energy_vert_face_hess_quad_kernel(
        points,
        p[self.soft.id],
        self.candidates.closest,
        self.candidates.collide,
        self.candidates.distance,
        self.rest_length.reshape((1,)),
        self.stiffness.reshape((1,)),
        output_dims={"hess_quad": (self.soft.n_points,)},
        launch_dims=(self.soft.n_points,),
    )
    if self.filter_hess_quad:
        hess_quad = jnp.clip(hess_quad, min=0.0)
    return hess_quad.sum()

hessp ¤

hessp(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

) -> Self:
    if max_dist is None:
        max_dist = 2.0 * rest_length
    return cls(
        rigid=rigid,
        soft=soft,

jac ¤

jac(x: ArrayDict, /, params: GlobalParams) -> ArrayDict

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
@utils.jit_method(inline=True)
def jac(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> struct.ArrayDict:
    points: Float[Array, "points dim"] = self.soft.points + x[self.soft.id]
    jac_soft: Float[Array, " points dim"]
    (jac_soft,) = collision_energy_vert_face_jac_kernel(
        points,
        self.candidates.closest,
        self.candidates.collide,
        self.candidates.distance,
        self.rest_length.reshape((1,)),
        self.stiffness.reshape((1,)),
        output_dims={"jac": (self.soft.n_points,)},
        launch_dims=(self.soft.n_points,),
    )
    # jax.debug.print("CollisionVertFace.jac = {}", jac_soft)
    return struct.ArrayDict(
        {self.soft.id: jac_soft, self.rigid.id: jnp.zeros_like(x[self.rigid.id])}
    )

jac_and_hess_diag ¤

jac_and_hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[ArrayDict, ArrayDict]

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

energy: Float[Array, " points"]
(energy,) = collision_energy_vert_face_fun_kernel(
    points,
    self.candidates.closest,
    self.candidates.collide,
    self.candidates.distance,

pre_optim_iter ¤

pre_optim_iter(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
def pre_optim_iter(self, params: sim.GlobalParams) -> Self:
    candidates: CollisionCandidatesVertFace = self.collide()
    # wl.pprint(candidates, short_arrays=False)
    return self.evolve(candidates=candidates)

pre_time_step ¤

pre_time_step(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

collide: Bool[Array, " points"] = struct.array(default=None)
distance: Float[Array, " points"] = struct.array(default=None)

tree_at ¤

tree_at(
    where: Callable[[Self], Node | Sequence[Node]],
    replace: Any | Sequence[Any] = ...,
    replace_fn: Callable[[Node], Any] = ...,
    is_leaf: Callable[[Any], bool] | None = None,
) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

    *,
    stiffness: float = 1e3,
    rest_length: float = 1e-3,
    max_dist: float | None = None,
    epsilon: float = 1e-3,
) -> Self:
    if max_dist is None:
        max_dist = 2.0 * rest_length
    return cls(
        rigid=rigid,
        soft=soft,
        stiffness=jnp.asarray(stiffness),
        rest_length=jnp.asarray(rest_length),
        max_dist=jnp.asarray(max_dist),
        epsilon=jnp.asarray(epsilon),

with_actors ¤

with_actors(actors: NodeContainer[Actor]) -> Self

Source code in src/liblaf/apple/energy/collision/vert_face/vert_face.py

@override
def with_actors(self, actors: struct.NodeContainer[sim.Actor]) -> Self:
    return self.evolve(rigid=actors[self.rigid.id], soft=actors[self.soft.id])

EnergyZero ¤

Bases: Energy

Parameters:

id (str, default: <dynamic> ) –
_actors (NodeContainer[Actor], default: <dynamic> ) –

Methods:

__pdoc__ –

...
__repr__ –
evolve –
from_actors –
fun –
fun_and_jac –
hess_diag –
hess_quad –
hessp –
jac –
jac_and_hess_diag –
pre_optim_iter –
pre_time_step –
tree_at –
with_actors –

Attributes:

actors (NodeContainer[Actor]) –
id (str) –

actors `property` ¤

actors: NodeContainer[Actor]

id `class-attribute` `instance-attribute` ¤

id: str = static(
    default=Factory(uniq_id, takes_self=True), kw_only=True
)

pdoc ¤

__pdoc__(**kwargs) -> AbstractDoc

...

References

wadler_lindig._definitions._pformat_dataclass()

Source code in src/liblaf/apple/energy/zero.py

@classmethod
def from_actors(cls, *actors: sim.Actor) -> Self:
    return cls(_actors=struct.NodeContainer(actors))

@property
@override
def actors(self) -> struct.NodeContainer[sim.Actor]:
    return self._actors

@override
def with_actors(self, actors: struct.NodeContainer[sim.Actor]) -> Self:
    return self.evolve(_actors=actors)

@override
@utils.jit_method(inline=True)
def fun(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> FloatScalar:
    return jnp.zeros(())

@override
@utils.jit_method(inline=True)
def jac(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> struct.ArrayDict:
    return struct.ArrayDict(
        {actor.id: jnp.zeros_like(x[actor.id]) for actor in self.actors.values()}
    )

repr ¤

__repr__() -> str

Source code in src/liblaf/apple/energy/zero.py

@override
@utils.jit_method(inline=True)

evolve ¤

evolve(**changes) -> Self

Source code in src/liblaf/apple/energy/zero.py

    self, x: struct.ArrayDict, /, params: sim.GlobalParams
) -> struct.ArrayDict:

from_actors `classmethod` ¤

from_actors(*actors: Actor) -> Self

Source code in src/liblaf/apple/energy/zero.py

@classmethod
def from_actors(cls, *actors: sim.Actor) -> Self:
    return cls(_actors=struct.NodeContainer(actors))

fun ¤

fun(x: ArrayDict, /, params: GlobalParams) -> FloatScalar

Source code in src/liblaf/apple/energy/zero.py

@override
@utils.jit_method(inline=True)
def fun(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> FloatScalar:
    return jnp.zeros(())

fun_and_jac ¤

fun_and_jac(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[Float[Array, ""], ArrayDict]

hess_diag ¤

hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/zero.py

@override
@utils.jit_method(inline=True)
def hess_diag(
    self, x: struct.ArrayDict, /, params: sim.GlobalParams
) -> struct.ArrayDict:
    return struct.ArrayDict(
        {actor.id: jnp.zeros_like(x[actor.id]) for actor in self.actors.values()}
    )

hess_quad ¤

hess_quad(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> FloatScalar

Source code in src/liblaf/apple/energy/zero.py

@override
@utils.jit_method(inline=True)
def hess_quad(
    self, x: struct.ArrayDict, p: struct.ArrayDict, /, params: sim.GlobalParams
) -> FloatScalar:
    return jnp.zeros(())

hessp ¤

hessp(
    x: ArrayDict, p: ArrayDict, /, params: GlobalParams
) -> ArrayDict

Source code in src/liblaf/apple/energy/zero.py

def hess_quad(
    self, x: struct.ArrayDict, p: struct.ArrayDict, /, params: sim.GlobalParams
) -> FloatScalar:
    return jnp.zeros(())

jac ¤

jac(x: ArrayDict, /, params: GlobalParams) -> ArrayDict

Source code in src/liblaf/apple/energy/zero.py

@override
@utils.jit_method(inline=True)
def jac(self, x: struct.ArrayDict, /, params: sim.GlobalParams) -> struct.ArrayDict:
    return struct.ArrayDict(
        {actor.id: jnp.zeros_like(x[actor.id]) for actor in self.actors.values()}
    )

jac_and_hess_diag ¤

jac_and_hess_diag(
    x: ArrayDict, /, params: GlobalParams
) -> tuple[ArrayDict, ArrayDict]

pre_optim_iter ¤

pre_optim_iter(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/zero.py

def actors(self) -> struct.NodeContainer[sim.Actor]:
    return self._actors

pre_time_step ¤

pre_time_step(params: GlobalParams) -> Self

Source code in src/liblaf/apple/energy/zero.py

21	`@property`

tree_at ¤

tree_at(
    where: Callable[[Self], Node | Sequence[Node]],
    replace: Any | Sequence[Any] = ...,
    replace_fn: Callable[[Node], Any] = ...,
    is_leaf: Callable[[Any], bool] | None = None,
) -> Self

Source code in src/liblaf/apple/energy/zero.py

        {actor.id: jnp.zeros_like(x[actor.id]) for actor in self.actors.values()}
    )

@override
@utils.jit_method(inline=True)
def hess_quad(
    self, x: struct.ArrayDict, p: struct.ArrayDict, /, params: sim.GlobalParams
) -> FloatScalar:
    return jnp.zeros(())

with_actors ¤

with_actors(actors: NodeContainer[Actor]) -> Self

Source code in src/liblaf/apple/energy/zero.py

@override
def with_actors(self, actors: struct.NodeContainer[sim.Actor]) -> Self:
    return self.evolve(_actors=actors)

PhaceStatic ¤

Bases: Elastic

As-Rigid-As-Possible.

\[ \Psi = \frac{\mu}{2} \|F - R\|_F^2 = \frac{\mu}{2} (I_2 - 2 I_1 + 3) \]

Parameters:

id (str, default: <dynamic> ) –
actor (Actor) –
hess_diag_filter (bool, default: True ) –
hess_quad_filter (bool, default: True ) –

Methods:

__pdoc__ –

...
__repr__ –
energy_density –
energy_density_hess_diag –
energy_density_hess_quad –
energy_density_jac –
evolve –
first_piola_kirchhoff_stress –
from_actor –
fun –
fun_and_jac –
hess_diag –
hess_quad –
hessp –
jac –
jac_and_hess_diag –
make_field –
pre_optim_iter –
pre_time_step –
tree_at –
with_actors –

Attributes:

actor (Actor) –
actors (NodeContainer[Actor]) –
hess_diag_filter (bool) –
hess_quad_filter (bool) –
id (str) –
lambda_ (Float[Array, ' cells']) –
mu (Float[Array, ' cells']) –
params (Float[Array, 'cells 2']) –

actor `class-attribute` `instance-attribute` ¤

actor: Actor = data()

actors `abstractmethod` `property` ¤

actors: NodeContainer[Actor]

hess_diag_filter `class-attribute` `instance-attribute` ¤

hess_diag_filter: bool = static(default=True, kw_only=True)

hess_quad_filter `class-attribute` `instance-attribute` ¤

hess_quad_filter: bool = static(default=True, kw_only=True)

id `class-attribute` `instance-attribute` ¤

id: str = static(
    default=Factory(uniq_id, takes_self=True), kw_only=True
)

lambda_ `property` ¤

lambda_: Float[Array, ' cells']

mu `property` ¤

mu: Float[Array, ' cells']

params `property` ¤

params: Float[Array, 'cells 2']

pdoc ¤

__pdoc__(**kwargs) -> AbstractDoc

...

References

wadler_lindig._definitions._pformat_dataclass()