Skip to content

liblaf.apple ¤

Modules:

Classes:

Attributes:

__version__ module-attribute ¤

__version__: str = '0.7.0'

__version_tuple__ module-attribute ¤

__version_tuple__: tuple[int | str, ...] = (0, 7, 0)

Arap ¤

Bases: Hyperelastic


              flowchart TD
              liblaf.apple.Arap[Arap]
              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic[Hyperelastic]
              liblaf.apple.warp.model._energy.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic --> liblaf.apple.Arap
                                liblaf.apple.warp.model._energy.WarpEnergy --> liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.warp.model._energy.WarpEnergy
                




              click liblaf.apple.Arap href "" "liblaf.apple.Arap"
              click liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic href "" "liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic"
              click liblaf.apple.warp.model._energy.WarpEnergy href "" "liblaf.apple.warp.model._energy.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (Sequence[str], default: () ) –

  • cells ¤

    (Integer[array, 'c a']) –

  • dhdX ¤

    (Integer[array, 'c q a J']) –

  • dV ¤

    (Integer[array, 'c q']) –

  • params ¤

    (Struct) –

  • clamp_hess_diag ¤

    (bool, default: True ) –

  • clamp_hess_quad ¤

    (bool, default: True ) –

  • clamp_lambda ¤

    (bool, default: True ) –
Subclassed by:

Classes:

Methods:

Attributes:

cells instance-attribute ¤

cells: Integer[array, 'c a']

clamp_hess_diag class-attribute instance-attribute ¤

clamp_hess_diag: bool = True

clamp_hess_quad class-attribute instance-attribute ¤

clamp_hess_quad: bool = True

clamp_lambda class-attribute instance-attribute ¤

clamp_lambda: bool = True

dV instance-attribute ¤

dV: Integer[array, 'c q']

dhdX instance-attribute ¤

dhdX: Integer[array, 'c q a J']

fun_kernel cached property ¤

fun_kernel: Kernel

grad_and_hess_diag_kernel cached property ¤

grad_and_hess_diag_kernel: Kernel

grad_kernel cached property ¤

grad_kernel: Kernel

hess_diag_kernel cached property ¤

hess_diag_kernel: Kernel

hess_prod_kernel cached property ¤

hess_prod_kernel: Kernel

hess_quad_kernel cached property ¤

hess_quad_kernel: Kernel

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

n_cells property ¤

n_cells: int

n_quadrature_points property ¤

n_quadrature_points: int

params instance-attribute ¤

params: Params

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

value_and_grad_kernel cached property ¤

value_and_grad_kernel: Kernel

Params ¤

Used by:

Attributes:

mu instance-attribute ¤

mu: array(dtype=float_)

ParamsElem ¤

Returned by: Used by:

Attributes:

mu instance-attribute ¤

mu: float_

energy_density_func staticmethod ¤

energy_density_func(
    F: mat33, params: ParamsElem
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
39
40
41
42
43
44
45
46
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_func(F: mat33, params: ParamsElem) -> scalar:
    R, _ = math.polar_rv(F)
    Psi = F.dtype(0.5) * params.mu * math.fro_norm_square(F - R)
    return Psi

energy_density_hess_diag_func staticmethod ¤

energy_density_hess_diag_func(
    F: mat33,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
61
62
63
64
65
66
67
68
69
70
71
72
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_diag_func(
    F: mat33, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    h4_diag = func.h4_diag(dhdX, U, s, V, clamp=clamp)  # mat43
    h5_diag = func.h5_diag(dhdX)  # mat43
    h_diag = -F.dtype(2.0) * h4_diag + h5_diag  # mat43
    return F.dtype(0.5) * params.mu * h_diag  # mat43

energy_density_hess_prod_func staticmethod ¤

energy_density_hess_prod_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
74
75
76
77
78
79
80
81
82
83
84
85
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_prod_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    h4_prod = func.h4_prod(p, dhdX, U, s, V, clamp=clamp)  # mat43
    h5_prod = func.h5_prod(p, dhdX)  # mat43
    h_prod = -F.dtype(2.0) * h4_prod + h5_prod  # mat43
    return F.dtype(0.5) * params.mu * h_prod  # mat43

energy_density_hess_quad_func staticmethod ¤

energy_density_hess_quad_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
87
88
89
90
91
92
93
94
95
96
97
98
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_quad_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> scalar:
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    h4_quad = func.h4_quad(p, dhdX, U, s, V, clamp=clamp)  # float
    h5_quad = func.h5_quad(p, dhdX)  # float
    h_quad = -F.dtype(2.0) * h4_quad + h5_quad
    return F.dtype(0.5) * params.mu * h_quad

first_piola_kirchhoff_stress_func staticmethod ¤

first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
48
49
50
51
52
53
54
55
56
57
58
59
@override
@staticmethod
@no_type_check
@wp.func
def first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33:
    R, _ = math.polar_rv(F)
    g1 = func.g1(R)  # mat33
    g2 = func.g2(F)  # mat33
    PK1 = F.dtype(0.5) * params.mu * (g2 - F.dtype(2.0) * g1)  # mat33
    return PK1

from_pyvista classmethod ¤

from_pyvista(
    obj: DataObject,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
59
60
61
62
@classmethod
def from_pyvista(cls, obj: pv.DataObject, **kwargs) -> Self:
    region = Region.from_pyvista(obj, grad=True)
    return cls.from_region(region, **kwargs)

from_region classmethod ¤

from_region(
    region: Region,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def from_region(
    cls, region: Region, *, requires_grad: Sequence[str] = (), **kwargs
) -> Self:
    self: Self = cls(
        cells=wpu.to_warp(region.cells, wpt.vec4i),
        dhdX=wpu.to_warp(region.dhdX, wpt.mat43),
        dV=wpu.to_warp(region.dV, wpt.float_),
        params=cls.make_params(region, requires_grad),
        requires_grad=requires_grad,
        **kwargs,
    )
    return self

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
110
111
112
113
114
115
116
117
@override
def fun(self, u: Vector, output: Scalar) -> None:
    wp.launch(
        self.fun_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

get_cell_params staticmethod ¤

get_cell_params(params: Params, cid: int) -> ParamsElem
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py 
32
33
34
35
36
37
@override
@staticmethod
@no_type_check
@wp.func
def get_cell_params(params: Params, cid: int) -> ParamsElem:
    return Arap.ParamsElem(mu=params.mu[cid])

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
119
120
121
122
123
124
125
126
@override
def grad(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
179
180
181
182
183
184
185
186
@override
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    wp.launch(
        self.grad_and_hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[grad, hess_diag],
    )

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
128
129
130
131
132
133
134
135
@override
def hess_diag(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
137
138
139
140
141
142
143
144
@override
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_prod_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
146
147
148
149
150
151
152
153
@override
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    wp.launch(
        self.hess_quad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

make_params classmethod ¤

make_params(
    region: Region, requires_grad: Sequence[str] = ()
) -> Params
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
90
91
92
93
94
@classmethod
@no_type_check
def make_params(cls, region: Region, requires_grad: Sequence[str] = ()) -> Params:
    fields: Mapping[str, wp.array] = cls._params_fields_from_region(region)
    return cls._params_from_fields(fields, requires_grad)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: array, p: array
) -> dict[str, array]
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@override
def mixed_derivative_prod(self, u: wp.array, p: wp.array) -> dict[str, wp.array]:
    if not self.requires_grad:
        return {}
    for name in self.requires_grad:
        getattr(self.params, name).grad.zero_()
    output: wp.array = wp.zeros_like(u)
    with wp.Tape() as tape:
        self.grad(u, output)
    tape.backward(grads={output: p})
    outputs: dict[str, wp.array] = {
        name: getattr(self.params, name).grad for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
104
105
106
107
108
@override
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self.params, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
170
171
172
173
174
175
176
177
@override
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    wp.launch(
        self.value_and_grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[value, grad],
    )

ArapActive ¤

Bases: Hyperelastic


              flowchart TD
              liblaf.apple.ArapActive[ArapActive]
              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic[Hyperelastic]
              liblaf.apple.warp.model._energy.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic --> liblaf.apple.ArapActive
                                liblaf.apple.warp.model._energy.WarpEnergy --> liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.warp.model._energy.WarpEnergy
                




              click liblaf.apple.ArapActive href "" "liblaf.apple.ArapActive"
              click liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic href "" "liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic"
              click liblaf.apple.warp.model._energy.WarpEnergy href "" "liblaf.apple.warp.model._energy.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (Sequence[str], default: () ) –

  • cells ¤

    (Integer[array, 'c a']) –

  • dhdX ¤

    (Integer[array, 'c q a J']) –

  • dV ¤

    (Integer[array, 'c q']) –

  • params ¤

    (Struct) –

  • clamp_hess_diag ¤

    (bool, default: True ) –

  • clamp_hess_quad ¤

    (bool, default: True ) –

  • clamp_lambda ¤

    (bool, default: True ) –

Classes:

Methods:

Attributes:

cells instance-attribute ¤

cells: Integer[array, 'c a']

clamp_hess_diag class-attribute instance-attribute ¤

clamp_hess_diag: bool = True

clamp_hess_quad class-attribute instance-attribute ¤

clamp_hess_quad: bool = True

clamp_lambda class-attribute instance-attribute ¤

clamp_lambda: bool = True

dV instance-attribute ¤

dV: Integer[array, 'c q']

dhdX instance-attribute ¤

dhdX: Integer[array, 'c q a J']

fun_kernel cached property ¤

fun_kernel: Kernel

grad_and_hess_diag_kernel cached property ¤

grad_and_hess_diag_kernel: Kernel

grad_kernel cached property ¤

grad_kernel: Kernel

hess_diag_kernel cached property ¤

hess_diag_kernel: Kernel

hess_prod_kernel cached property ¤

hess_prod_kernel: Kernel

hess_quad_kernel cached property ¤

hess_quad_kernel: Kernel

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

n_cells property ¤

n_cells: int

n_quadrature_points property ¤

n_quadrature_points: int

params instance-attribute ¤

params: Params

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

value_and_grad_kernel cached property ¤

value_and_grad_kernel: Kernel

Params ¤

Used by:

Attributes:

activation instance-attribute ¤

activation: array(dtype=vec6)

mu instance-attribute ¤

mu: array(dtype=float_)

muscle_fraction instance-attribute ¤

muscle_fraction: array(dtype=float_)

ParamsElem ¤

Returned by: Used by:

Attributes:

activation instance-attribute ¤

activation: vec6

mu instance-attribute ¤

mu: float_

muscle_fraction instance-attribute ¤

muscle_fraction: float_

energy_density_func staticmethod ¤

energy_density_func(
    F: mat33, params: ParamsElem
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_func(F: mat33, params: ParamsElem) -> scalar:
    Psi_active = ArapMuscle.energy_density_func(
        F, ArapActive._arap_active_params(params)
    )  # float
    Psi_passive = Arap.energy_density_func(
        F, ArapActive._arap_params(params)
    )  # float
    Psi = (
        params.muscle_fraction * Psi_active
        + (F.dtype(1.0) - params.muscle_fraction) * Psi_passive
    )  # float
    return Psi

energy_density_hess_diag_func staticmethod ¤

energy_density_hess_diag_func(
    F: mat33,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_diag_func(
    F: mat33, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    diag_active = ArapMuscle.energy_density_hess_diag_func(
        F, dhdX, ArapActive._arap_active_params(params), clamp=clamp
    )  # mat43
    diag_passive = Arap.energy_density_hess_diag_func(
        F, dhdX, ArapActive._arap_params(params), clamp=clamp
    )  # mat43
    diag = (
        params.muscle_fraction * diag_active
        + (F.dtype(1.0) - params.muscle_fraction) * diag_passive
    )  # mat43
    return diag

energy_density_hess_prod_func staticmethod ¤

energy_density_hess_prod_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_prod_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    prod_active = ArapMuscle.energy_density_hess_prod_func(
        F, p, dhdX, ArapActive._arap_active_params(params), clamp=clamp
    )  # mat43
    prod_passive = Arap.energy_density_hess_prod_func(
        F, p, dhdX, ArapActive._arap_params(params), clamp=clamp
    )  # mat43
    prod = (
        params.muscle_fraction * prod_active
        + (F.dtype(1.0) - params.muscle_fraction) * prod_passive
    )  # mat43
    return prod

energy_density_hess_quad_func staticmethod ¤

energy_density_hess_quad_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_quad_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> scalar:
    quad_active = ArapMuscle.energy_density_hess_quad_func(
        F, p, dhdX, ArapActive._arap_active_params(params), clamp=clamp
    )
    quad_passive = Arap.energy_density_hess_quad_func(
        F, p, dhdX, ArapActive._arap_params(params), clamp=clamp
    )
    quad = (
        params.muscle_fraction * quad_active
        + (F.dtype(1.0) - params.muscle_fraction) * quad_passive
    )
    return quad

first_piola_kirchhoff_stress_func staticmethod ¤

first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@override
@staticmethod
@no_type_check
@wp.func
def first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33:
    PK1_active = ArapMuscle.first_piola_kirchhoff_stress_func(
        F, ArapActive._arap_active_params(params), clamp=clamp
    )  # mat33
    PK1_passive = Arap.first_piola_kirchhoff_stress_func(
        F, ArapActive._arap_params(params), clamp=clamp
    )  # mat33
    PK1 = (
        params.muscle_fraction * PK1_active
        + (F.dtype(1.0) - params.muscle_fraction) * PK1_passive
    )  # mat33
    return PK1

from_pyvista classmethod ¤

from_pyvista(
    obj: DataObject,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
59
60
61
62
@classmethod
def from_pyvista(cls, obj: pv.DataObject, **kwargs) -> Self:
    region = Region.from_pyvista(obj, grad=True)
    return cls.from_region(region, **kwargs)

from_region classmethod ¤

from_region(
    region: Region,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def from_region(
    cls, region: Region, *, requires_grad: Sequence[str] = (), **kwargs
) -> Self:
    self: Self = cls(
        cells=wpu.to_warp(region.cells, wpt.vec4i),
        dhdX=wpu.to_warp(region.dhdX, wpt.mat43),
        dV=wpu.to_warp(region.dV, wpt.float_),
        params=cls.make_params(region, requires_grad),
        requires_grad=requires_grad,
        **kwargs,
    )
    return self

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
110
111
112
113
114
115
116
117
@override
def fun(self, u: Vector, output: Scalar) -> None:
    wp.launch(
        self.fun_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

get_cell_params staticmethod ¤

get_cell_params(params: Params, cid: int) -> ParamsElem
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py 
37
38
39
40
41
42
43
44
45
46
@override
@staticmethod
@no_type_check
@wp.func
def get_cell_params(params: Params, cid: int) -> ParamsElem:
    return ArapActive.ParamsElem(
        activation=params.activation[cid],
        mu=params.mu[cid],
        muscle_fraction=params.muscle_fraction[cid],
    )

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
119
120
121
122
123
124
125
126
@override
def grad(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
179
180
181
182
183
184
185
186
@override
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    wp.launch(
        self.grad_and_hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[grad, hess_diag],
    )

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
128
129
130
131
132
133
134
135
@override
def hess_diag(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
137
138
139
140
141
142
143
144
@override
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_prod_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
146
147
148
149
150
151
152
153
@override
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    wp.launch(
        self.hess_quad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

make_params classmethod ¤

make_params(
    region: Region, requires_grad: Sequence[str] = ()
) -> Params
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
90
91
92
93
94
@classmethod
@no_type_check
def make_params(cls, region: Region, requires_grad: Sequence[str] = ()) -> Params:
    fields: Mapping[str, wp.array] = cls._params_fields_from_region(region)
    return cls._params_from_fields(fields, requires_grad)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: array, p: array
) -> dict[str, array]
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@override
def mixed_derivative_prod(self, u: wp.array, p: wp.array) -> dict[str, wp.array]:
    if not self.requires_grad:
        return {}
    for name in self.requires_grad:
        getattr(self.params, name).grad.zero_()
    output: wp.array = wp.zeros_like(u)
    with wp.Tape() as tape:
        self.grad(u, output)
    tape.backward(grads={output: p})
    outputs: dict[str, wp.array] = {
        name: getattr(self.params, name).grad for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
104
105
106
107
108
@override
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self.params, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
170
171
172
173
174
175
176
177
@override
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    wp.launch(
        self.value_and_grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[value, grad],
    )

ArapMuscle ¤

Bases: Arap


              flowchart TD
              liblaf.apple.ArapMuscle[ArapMuscle]
              liblaf.apple.warp.energies.elastic.hyperelastic._arap.Arap[Arap]
              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic[Hyperelastic]
              liblaf.apple.warp.model._energy.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.warp.energies.elastic.hyperelastic._arap.Arap --> liblaf.apple.ArapMuscle
                                liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic --> liblaf.apple.warp.energies.elastic.hyperelastic._arap.Arap
                                liblaf.apple.warp.model._energy.WarpEnergy --> liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.warp.model._energy.WarpEnergy
                





              click liblaf.apple.ArapMuscle href "" "liblaf.apple.ArapMuscle"
              click liblaf.apple.warp.energies.elastic.hyperelastic._arap.Arap href "" "liblaf.apple.warp.energies.elastic.hyperelastic._arap.Arap"
              click liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic href "" "liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic"
              click liblaf.apple.warp.model._energy.WarpEnergy href "" "liblaf.apple.warp.model._energy.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (Sequence[str], default: () ) –

  • cells ¤

    (Integer[array, 'c a']) –

  • dhdX ¤

    (Integer[array, 'c q a J']) –

  • dV ¤

    (Integer[array, 'c q']) –

  • params ¤

    (Struct) –

  • clamp_hess_diag ¤

    (bool, default: True ) –

  • clamp_hess_quad ¤

    (bool, default: True ) –

  • clamp_lambda ¤

    (bool, default: True ) –

Classes:

Methods:

Attributes:

cells instance-attribute ¤

cells: Integer[array, 'c a']

clamp_hess_diag class-attribute instance-attribute ¤

clamp_hess_diag: bool = True

clamp_hess_quad class-attribute instance-attribute ¤

clamp_hess_quad: bool = True

clamp_lambda class-attribute instance-attribute ¤

clamp_lambda: bool = True

dV instance-attribute ¤

dV: Integer[array, 'c q']

dhdX instance-attribute ¤

dhdX: Integer[array, 'c q a J']

fun_kernel cached property ¤

fun_kernel: Kernel

grad_and_hess_diag_kernel cached property ¤

grad_and_hess_diag_kernel: Kernel

grad_kernel cached property ¤

grad_kernel: Kernel

hess_diag_kernel cached property ¤

hess_diag_kernel: Kernel

hess_prod_kernel cached property ¤

hess_prod_kernel: Kernel

hess_quad_kernel cached property ¤

hess_quad_kernel: Kernel

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

n_cells property ¤

n_cells: int

n_quadrature_points property ¤

n_quadrature_points: int

params instance-attribute ¤

params: Params

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

value_and_grad_kernel cached property ¤

value_and_grad_kernel: Kernel

Params ¤

Used by:

Attributes:

activation instance-attribute ¤

activation: array(dtype=vec6)

mu instance-attribute ¤

mu: array(dtype=float_)

ParamsElem ¤

Returned by: Used by:

Attributes:

activation instance-attribute ¤

activation: vec6

mu instance-attribute ¤

mu: float_

energy_density_func staticmethod ¤

energy_density_func(
    F: mat33, params: ParamsElem
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
43
44
45
46
47
48
49
50
51
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_func(F: mat33, params: ParamsElem) -> scalar:
    A = func.make_activation_mat33(params.activation)  # mat33
    R, _ = math.polar_rv(F)
    Psi = F.dtype(0.5) * params.mu * math.fro_norm_square(F - R @ A)
    return Psi

energy_density_hess_diag_func staticmethod ¤

energy_density_hess_diag_func(
    F: mat33,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_diag_func(
    F: mat33, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    A = func.make_activation_mat33(params.activation)  # mat33
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    lambdas = func.lambdas(s, clamp=clamp)  # vec3
    Q0, Q1, Q2 = func.Qs(U, V)  # mat33, mat33, mat33
    h4_diag = (
        lambdas[0]
        * wp.cw_mul(
            func.deformation_gradient_vjp(dhdX, Q0 @ A),
            func.deformation_gradient_vjp(dhdX, Q0),
        )
        + lambdas[1]
        * wp.cw_mul(
            func.deformation_gradient_vjp(dhdX, Q1 @ A),
            func.deformation_gradient_vjp(dhdX, Q1),
        )
        + lambdas[2]
        * wp.cw_mul(
            func.deformation_gradient_vjp(dhdX, Q2 @ A),
            func.deformation_gradient_vjp(dhdX, Q2),
        )
    )  # mat43
    h5_diag = func.h5_diag(dhdX)  # mat43
    h_diag = -F.dtype(2.0) * h4_diag + h5_diag  # mat43
    return F.dtype(0.5) * params.mu * h_diag  # mat43

energy_density_hess_prod_func staticmethod ¤

energy_density_hess_prod_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_prod_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    A = func.make_activation_mat33(params.activation)  # mat33
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    lambdas = func.lambdas(s, clamp=clamp)  # vec3
    Q0, Q1, Q2 = func.Qs(U, V)  # mat33, mat33, mat33
    dFdx_p = func.deformation_gradient_jvp(dhdX, p)  # mat33
    h4_prod = (
        lambdas[0]
        * func.deformation_gradient_vjp(dhdX, Q0 @ A)
        * wp.ddot(Q0, dFdx_p)
        + lambdas[1]
        * func.deformation_gradient_vjp(dhdX, Q1 @ A)
        * wp.ddot(Q1, dFdx_p)
        + lambdas[2]
        * func.deformation_gradient_vjp(dhdX, Q2 @ A)
        * wp.ddot(Q2, dFdx_p)
    )  # mat43
    # h4_prod = func.h4_prod(p, dhdX, U, s, V, clamp=clamp)  # mat43
    h5_prod = func.h5_prod(p, dhdX)  # mat43
    h_prod = -F.dtype(2.0) * h4_prod + h5_prod  # mat43
    return F.dtype(0.5) * params.mu * h_prod  # mat43

energy_density_hess_quad_func staticmethod ¤

energy_density_hess_quad_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_quad_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> scalar:
    A = func.make_activation_mat33(params.activation)  # mat33
    U, s, V = math.svd_rv(F)  # mat33, vec3, mat33
    lambdas = func.lambdas(s, clamp=clamp)  # vec3
    Q0, Q1, Q2 = func.Qs(U, V)  # mat33, mat33, mat33
    h4_quad = (
        lambdas[0]
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q0 @ A), p)
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q0), p)
        + lambdas[1]
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q1 @ A), p)
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q1), p)
        + lambdas[2]
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q2 @ A), p)
        * wp.ddot(func.deformation_gradient_vjp(dhdX, Q2), p)
    )  # float
    h5_quad = func.h5_quad(p, dhdX)  # float
    h_quad = -F.dtype(2.0) * h4_quad + h5_quad
    return F.dtype(0.5) * params.mu * h_quad

first_piola_kirchhoff_stress_func staticmethod ¤

first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
53
54
55
56
57
58
59
60
61
62
63
@override
@staticmethod
@no_type_check
@wp.func
def first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33:
    A = func.make_activation_mat33(params.activation)  # mat33
    R, _ = math.polar_rv(F)  # mat33
    PK1 = params.mu * (F - R @ A)
    return PK1

from_pyvista classmethod ¤

from_pyvista(
    obj: DataObject,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
59
60
61
62
@classmethod
def from_pyvista(cls, obj: pv.DataObject, **kwargs) -> Self:
    region = Region.from_pyvista(obj, grad=True)
    return cls.from_region(region, **kwargs)

from_region classmethod ¤

from_region(
    region: Region,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def from_region(
    cls, region: Region, *, requires_grad: Sequence[str] = (), **kwargs
) -> Self:
    self: Self = cls(
        cells=wpu.to_warp(region.cells, wpt.vec4i),
        dhdX=wpu.to_warp(region.dhdX, wpt.mat43),
        dV=wpu.to_warp(region.dV, wpt.float_),
        params=cls.make_params(region, requires_grad),
        requires_grad=requires_grad,
        **kwargs,
    )
    return self

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
110
111
112
113
114
115
116
117
@override
def fun(self, u: Vector, output: Scalar) -> None:
    wp.launch(
        self.fun_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

get_cell_params staticmethod ¤

get_cell_params(params: Params, cid: int) -> ParamsElem
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py 
34
35
36
37
38
39
40
41
@override
@staticmethod
@no_type_check
@wp.func
def get_cell_params(params: Params, cid: int) -> ParamsElem:
    return ArapMuscle.ParamsElem(
        activation=params.activation[cid], mu=params.mu[cid]
    )

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
119
120
121
122
123
124
125
126
@override
def grad(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
179
180
181
182
183
184
185
186
@override
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    wp.launch(
        self.grad_and_hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[grad, hess_diag],
    )

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
128
129
130
131
132
133
134
135
@override
def hess_diag(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
137
138
139
140
141
142
143
144
@override
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_prod_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
146
147
148
149
150
151
152
153
@override
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    wp.launch(
        self.hess_quad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

make_params classmethod ¤

make_params(
    region: Region, requires_grad: Sequence[str] = ()
) -> Params
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
90
91
92
93
94
@classmethod
@no_type_check
def make_params(cls, region: Region, requires_grad: Sequence[str] = ()) -> Params:
    fields: Mapping[str, wp.array] = cls._params_fields_from_region(region)
    return cls._params_from_fields(fields, requires_grad)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: array, p: array
) -> dict[str, array]
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@override
def mixed_derivative_prod(self, u: wp.array, p: wp.array) -> dict[str, wp.array]:
    if not self.requires_grad:
        return {}
    for name in self.requires_grad:
        getattr(self.params, name).grad.zero_()
    output: wp.array = wp.zeros_like(u)
    with wp.Tape() as tape:
        self.grad(u, output)
    tape.backward(grads={output: p})
    outputs: dict[str, wp.array] = {
        name: getattr(self.params, name).grad for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
104
105
106
107
108
@override
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self.params, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
170
171
172
173
174
175
176
177
@override
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    wp.launch(
        self.value_and_grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[value, grad],
    )

Dirichlet ¤

Parameters:

  • dim ¤

    (int) –

  • dirichlet_index ¤

    (Integer[Array, dirichlet]) –

  • dirichlet_value ¤

    (Float[Array, dirichlet]) –

  • free_index ¤

    (Integer[Array, free]) –

  • n_points ¤

    (int) –
Returned by: Used by:

Methods:

Attributes:

dim instance-attribute ¤

dim: int

dirichlet_index instance-attribute ¤

dirichlet_index: Integer[Array, ' dirichlet']

dirichlet_value instance-attribute ¤

dirichlet_value: Float[Array, ' dirichlet']

free_index instance-attribute ¤

free_index: Integer[Array, ' free']

n_dirichlet property ¤

n_dirichlet: int

n_free property ¤

n_free: int

n_full property ¤

n_full: int

n_points instance-attribute ¤

n_points: int

get_dirichlet ¤

get_dirichlet(
    full: Float[Array, "points dim"],
) -> Float[Array, " dirichlet"]
Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py 
27
28
29
30
31
@eqx.filter_jit
def get_dirichlet(
    self, full: Float[Array, "points dim"]
) -> Float[Array, " dirichlet"]:
    return full.flatten()[self.dirichlet_index]

get_free ¤

get_free(
    full: Float[Array, "points dim"],
) -> Float[Array, " free"]
Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py 
33
34
35
@eqx.filter_jit
def get_free(self, full: Float[Array, "points dim"]) -> Float[Array, " free"]:
    return full.flatten()[self.free_index]

set_dirichlet ¤

set_dirichlet(
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]
Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py 
37
38
39
40
41
42
43
44
45
@eqx.filter_jit
def set_dirichlet(
    self,
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]:
    if values is None:
        values = self.dirichlet_value
    return full.flatten().at[self.dirichlet_index].set(values).reshape(full.shape)

set_free ¤

set_free(
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " free"],
) -> Float[Array, "points dim"]
Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py 
47
48
49
50
51
@eqx.filter_jit
def set_free(
    self, full: Float[Array, "points dim"], values: Float[ArrayLike, " free"]
) -> Float[Array, "points dim"]:
    return full.flatten().at[self.free_index].set(values).reshape(full.shape)

to_full ¤

to_full(
    free: Float[Array, " free"],
    dirichlet: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]
Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py 
53
54
55
56
57
58
59
60
61
62
63
64
@eqx.filter_jit
def to_full(
    self,
    free: Float[Array, " free"],
    dirichlet: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]:
    full: Float[Array, "points dim"] = jnp.empty(
        (self.n_points, self.dim), free.dtype
    )
    full = self.set_free(full, free)
    full = self.set_dirichlet(full, dirichlet)
    return full

DirichletBuilder ¤

DirichletBuilder(dim: int = 3)

Parameters:

  • mask ¤

    (Bool[ndarray, 'points dim']) –

  • value ¤

    (Float[ndarray, 'points dim']) –
Returned by: Used by:

Methods:

Attributes:

Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py 
19
20
21
22
def __init__(self, dim: int = 3) -> None:
    mask: Bool[np.ndarray, "points dim"] = np.empty((0, dim), bool)
    value: Float[np.ndarray, "points dim"] = np.empty((0, dim))
    self.__attrs_init__(mask=mask, value=value)  # pyright: ignore[reportAttributeAccessIssue]

dim property ¤

dim: int

mask instance-attribute ¤

mask: Bool[ndarray, 'points dim']

n_points property ¤

n_points: int

value instance-attribute ¤

value: Float[ndarray, 'points dim']

add_pyvista ¤

add_pyvista(obj: DataSet) -> None
Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py 
32
33
34
35
36
37
38
39
40
41
42
def add_pyvista(self, obj: pv.DataSet) -> None:
    point_id = obj.point_data[POINT_ID]
    self.resize(point_id.max() + 1)
    dirichlet_mask: Bool[Array, "points dim"] = self._left_broadcast_to(
        obj.point_data[DIRICHLET_MASK], obj.n_points
    )
    dirichlet_value: Float[Array, "points dim"] = self._left_broadcast_to(
        obj.point_data[DIRICHLET_VALUE], obj.n_points
    )
    self.mask[point_id] = dirichlet_mask
    self.value[point_id] = dirichlet_value

finalize ¤

finalize() -> Dirichlet
Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py 
44
45
46
47
48
49
50
51
52
53
def finalize(self) -> Dirichlet:
    mask: Bool[Array, "points dim"] = jnp.asarray(self.mask)
    dirichlet_index: Integer[Array, " dirichlet"] = jnp.flatnonzero(mask)
    return Dirichlet(
        dim=self.dim,
        dirichlet_index=dirichlet_index,
        dirichlet_value=jnp.asarray(self.value.flat[dirichlet_index]),
        free_index=jnp.flatnonzero(~mask),
        n_points=self.n_points,
    )

resize ¤

resize(n_points: int) -> None
Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py 
55
56
57
58
59
60
def resize(self, n_points: int) -> None:
    pad_after: int = n_points - self.n_points
    if pad_after <= 0:
        return
    self.mask = np.pad(self.mask, ((0, pad_after), (0, 0)), constant_values=False)
    self.value = np.pad(self.value, ((0, pad_after), (0, 0)), constant_values=0.0)

Forward ¤

Parameters:

  • model ¤

    (Model) –

  • optimizer ¤

    (Optimizer, default: PNCG(jit=False, timer=False, max_steps=1000, norm=None, atol=Array(1.e-28, dtype=float32, weak_type=True), rtol=Array(1.e-05, dtype=float32, weak_type=True), d_hat=Array(inf, dtype=float32, weak_type=True)) ) –
Returned by: Used by:

Methods:

Attributes:

model instance-attribute ¤

model: Model

optimizer class-attribute instance-attribute ¤

optimizer: Optimizer = field(
    factory=lambda: PNCG(max_steps=1000)
)

u_full property ¤

u_full: Float[Array, 'points dim']

step ¤

step(callback: Callback | None = None) -> Solution
Source code in src/liblaf/apple/model/_forward.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def step(self, callback: Callback | None = None) -> Optimizer.Solution:
    objective = Objective(
        fun=self.model.fun,
        grad=self.model.grad,
        hess_diag=self.model.hess_diag,
        hess_prod=self.model.hess_prod,
        hess_quad=self.model.hess_quad,
        value_and_grad=self.model.value_and_grad,
        grad_and_hess_diag=self.model.grad_and_hess_diag,
    )
    solution: Optimizer.Solution = self.optimizer.minimize(
        objective, self.model.u_free, callback=callback
    )
    if not solution.success:
        logger.warning("Forward fail: %r", solution)
    self.model.u_free = solution.params
    return solution

update_params ¤

update_params(params: ModelParams) -> None
Source code in src/liblaf/apple/model/_forward.py 
26
27
def update_params(self, params: ModelParams) -> None:
    self.model.update_params(params)

Gravity ¤

Bases: JaxEnergy


              flowchart TD
              liblaf.apple.Gravity[Gravity]
              liblaf.apple.jax.model._energy.JaxEnergy[JaxEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.jax.model._energy.JaxEnergy --> liblaf.apple.Gravity
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.jax.model._energy.JaxEnergy
                



              click liblaf.apple.Gravity href "" "liblaf.apple.Gravity"
              click liblaf.apple.jax.model._energy.JaxEnergy href "" "liblaf.apple.jax.model._energy.JaxEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (frozenset[str], default: frozenset() ) –

  • gravity ¤

    (Float[Array, dim]) –

  • indices ¤

    (Integer[Array, points]) –

  • mass ¤

    (Float[Array, points]) –

Methods:

Attributes:

gravity instance-attribute ¤

gravity: Float[Array, ' dim']

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

indices instance-attribute ¤

indices: Integer[Array, ' points']

mass instance-attribute ¤

mass: Float[Array, ' points']

requires_grad class-attribute instance-attribute ¤

requires_grad: frozenset[str] = field(
    default=frozenset(), kw_only=True
)

from_pyvista classmethod ¤

from_pyvista(
    obj: DataSet,
    gravity: Float[ArrayLike, " dim"] | None = None,
) -> Self
Source code in src/liblaf/apple/jax/energies/_gravity.py 
23
24
25
26
27
28
29
30
31
32
33
@classmethod
def from_pyvista(
    cls, obj: pv.DataSet, gravity: Float[ArrayLike, " dim"] | None = None
) -> Self:
    if gravity is None:
        gravity = jnp.asarray([0.0, -9.81, 0.0])
    return cls(
        gravity=jnp.asarray(gravity),
        indices=jnp.asarray(obj.point_data[POINT_ID]),
        mass=jnp.asarray(obj.point_data[MASS]),
    )

fun ¤

fun(u: Vector) -> Scalar
Source code in src/liblaf/apple/jax/energies/_gravity.py 
35
36
37
38
@override
def fun(self, u: Vector) -> Scalar:
    u = u[self.indices]
    return -jnp.vdot(self.mass, jnp.vecdot(u, self.gravity, axis=-1))

grad ¤

grad(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
33
34
35
36
@eqx.filter_jit
def grad(self, u: Vector) -> Updates:
    values: Vector = eqx.filter_grad(self.fun)(u)
    return values, jnp.arange(u.shape[0])

grad_and_hess_diag ¤

grad_and_hess_diag(u: Vector) -> tuple[Updates, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
61
62
63
@eqx.filter_jit
def grad_and_hess_diag(self, u: Vector) -> tuple[Updates, Updates]:
    return self.grad(u), self.hess_diag(u)

hess_diag ¤

hess_diag(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/energies/_gravity.py 
40
41
42
@override
def hess_diag(self, u: Vector) -> Updates:
    return jnp.zeros_like(u[self.indices]), self.indices

hess_prod ¤

hess_prod(u: Vector, p: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
41
42
43
44
45
@eqx.filter_jit
def hess_prod(self, u: Vector, p: Vector) -> Updates:
    values: Vector
    _, values = jax.jvp(jax.grad(self.fun), (u,), (p,))
    return values, jnp.arange(u.shape[0])

hess_quad ¤

hess_quad(u: Vector, p: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_energy.py 
47
48
49
50
51
52
@eqx.filter_jit
def hess_quad(self, u: Vector, p: Vector) -> Scalar:
    values: Vector
    index: Index
    values, index = self.hess_prod(u, p)
    return jnp.vdot(p[index], values)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> EnergyParams
Source code in src/liblaf/apple/jax/model/_energy.py 
65
66
67
68
69
70
71
@eqx.filter_jit
def mixed_derivative_prod(self, u: Vector, p: Vector) -> EnergyParams:
    outputs: EnergyParams = {
        name: getattr(self, f"mixed_derivative_prod_{name}")(u, p)
        for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
22
23
24
@eqx.filter_jit
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: Mapping[str, Array]) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
26
27
28
def update_params(self, params: Mapping[str, Array]) -> None:
    for name, value in params.items():
        setattr(self, name, value)

value_and_grad ¤

value_and_grad(u: Vector) -> tuple[Scalar, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
54
55
56
57
58
59
@eqx.filter_jit
def value_and_grad(self, u: Vector) -> tuple[Scalar, Updates]:
    value: Scalar
    grad: Vector
    value, grad = jax.value_and_grad(self.fun)(u)
    return value, (grad, jnp.arange(u.shape[0]))

Hyperelastic ¤

Bases: WarpEnergy


              flowchart TD
              liblaf.apple.Hyperelastic[Hyperelastic]
              liblaf.apple.warp.model._energy.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.warp.model._energy.WarpEnergy --> liblaf.apple.Hyperelastic
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.warp.model._energy.WarpEnergy
                



              click liblaf.apple.Hyperelastic href "" "liblaf.apple.Hyperelastic"
              click liblaf.apple.warp.model._energy.WarpEnergy href "" "liblaf.apple.warp.model._energy.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (Sequence[str], default: () ) –

  • cells ¤

    (Integer[array, 'c a']) –

  • dhdX ¤

    (Integer[array, 'c q a J']) –

  • dV ¤

    (Integer[array, 'c q']) –

  • params ¤

    (Struct) –

  • clamp_hess_diag ¤

    (bool, default: True ) –

  • clamp_hess_quad ¤

    (bool, default: True ) –

  • clamp_lambda ¤

    (bool, default: True ) –

Classes:

Methods:

Attributes:

cells instance-attribute ¤

cells: Integer[array, 'c a']

clamp_hess_diag class-attribute instance-attribute ¤

clamp_hess_diag: bool = True

clamp_hess_quad class-attribute instance-attribute ¤

clamp_hess_quad: bool = True

clamp_lambda class-attribute instance-attribute ¤

clamp_lambda: bool = True

dV instance-attribute ¤

dV: Integer[array, 'c q']

dhdX instance-attribute ¤

dhdX: Integer[array, 'c q a J']

fun_kernel cached property ¤

fun_kernel: Kernel

grad_and_hess_diag_kernel cached property ¤

grad_and_hess_diag_kernel: Kernel

grad_kernel cached property ¤

grad_kernel: Kernel

hess_diag_kernel cached property ¤

hess_diag_kernel: Kernel

hess_prod_kernel cached property ¤

hess_prod_kernel: Kernel

hess_quad_kernel cached property ¤

hess_quad_kernel: Kernel

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

n_cells property ¤

n_cells: int

n_quadrature_points property ¤

n_quadrature_points: int

params instance-attribute ¤

params: Params

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

value_and_grad_kernel cached property ¤

value_and_grad_kernel: Kernel

Params ¤

ParamsElem ¤

energy_density_func staticmethod ¤

energy_density_func(
    F: mat33, params: ParamsElem
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
450
451
452
453
454
@staticmethod
@no_type_check
@wp.func
def energy_density_func(F: mat33, params: ParamsElem) -> scalar:
    raise NotImplementedError

energy_density_hess_diag_func staticmethod ¤

energy_density_hess_diag_func(
    F: mat33,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
464
465
466
467
468
469
470
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_diag_func(
    F: mat33, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    raise NotImplementedError

energy_density_hess_prod_func staticmethod ¤

energy_density_hess_prod_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
472
473
474
475
476
477
478
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_prod_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    raise NotImplementedError

energy_density_hess_quad_func staticmethod ¤

energy_density_hess_quad_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
480
481
482
483
484
485
486
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_quad_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> scalar:
    raise NotImplementedError

first_piola_kirchhoff_stress_func staticmethod ¤

first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
456
457
458
459
460
461
462
@staticmethod
@no_type_check
@wp.func
def first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33:
    raise NotImplementedError

from_pyvista classmethod ¤

from_pyvista(
    obj: DataObject,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
59
60
61
62
@classmethod
def from_pyvista(cls, obj: pv.DataObject, **kwargs) -> Self:
    region = Region.from_pyvista(obj, grad=True)
    return cls.from_region(region, **kwargs)

from_region classmethod ¤

from_region(
    region: Region,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def from_region(
    cls, region: Region, *, requires_grad: Sequence[str] = (), **kwargs
) -> Self:
    self: Self = cls(
        cells=wpu.to_warp(region.cells, wpt.vec4i),
        dhdX=wpu.to_warp(region.dhdX, wpt.mat43),
        dV=wpu.to_warp(region.dV, wpt.float_),
        params=cls.make_params(region, requires_grad),
        requires_grad=requires_grad,
        **kwargs,
    )
    return self

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
110
111
112
113
114
115
116
117
@override
def fun(self, u: Vector, output: Scalar) -> None:
    wp.launch(
        self.fun_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

get_cell_params staticmethod ¤

get_cell_params(params: Params, cid: int) -> ParamsElem
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
444
445
446
447
448
@staticmethod
@no_type_check
@wp.func
def get_cell_params(params: Params, cid: int) -> ParamsElem:
    raise NotImplementedError

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
119
120
121
122
123
124
125
126
@override
def grad(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
179
180
181
182
183
184
185
186
@override
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    wp.launch(
        self.grad_and_hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[grad, hess_diag],
    )

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
128
129
130
131
132
133
134
135
@override
def hess_diag(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
137
138
139
140
141
142
143
144
@override
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_prod_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
146
147
148
149
150
151
152
153
@override
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    wp.launch(
        self.hess_quad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

make_params classmethod ¤

make_params(
    region: Region, requires_grad: Sequence[str] = ()
) -> Params
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
90
91
92
93
94
@classmethod
@no_type_check
def make_params(cls, region: Region, requires_grad: Sequence[str] = ()) -> Params:
    fields: Mapping[str, wp.array] = cls._params_fields_from_region(region)
    return cls._params_from_fields(fields, requires_grad)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: array, p: array
) -> dict[str, array]
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@override
def mixed_derivative_prod(self, u: wp.array, p: wp.array) -> dict[str, wp.array]:
    if not self.requires_grad:
        return {}
    for name in self.requires_grad:
        getattr(self.params, name).grad.zero_()
    output: wp.array = wp.zeros_like(u)
    with wp.Tape() as tape:
        self.grad(u, output)
    tape.backward(grads={output: p})
    outputs: dict[str, wp.array] = {
        name: getattr(self.params, name).grad for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
104
105
106
107
108
@override
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self.params, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
170
171
172
173
174
175
176
177
@override
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    wp.launch(
        self.value_and_grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[value, grad],
    )

Inverse ¤

Bases: ABC


              flowchart TD
              liblaf.apple.Inverse[Inverse]

              

              click liblaf.apple.Inverse href "" "liblaf.apple.Inverse"

Parameters:

  • forward ¤

    (Forward) –

  • adjoint_solver ¤

    (LinearSolver, default: <dynamic> ) –

  • optimizer ¤

    (Optimizer, default: ScipyOptimizer(max_steps=256, jit=False, timer=False, method='L-BFGS-B', tol=1e-05, options=None) ) –

Methods:

Attributes:

adjoint_solver class-attribute instance-attribute ¤

adjoint_solver: LinearSolver = field(
    factory=CompositeSolver, kw_only=True
)

forward instance-attribute ¤

forward: Forward

model property ¤

model: Model

optimizer class-attribute instance-attribute ¤

optimizer: Optimizer = field(
    factory=lambda: ScipyOptimizer(
        method="L-BFGS-B", tol=1e-05
    ),
    kw_only=True,
)

adjoint ¤

adjoint(u: Full, dLdu: Full) -> Full
Source code in src/liblaf/apple/inverse/_inverse.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def adjoint(self, u: Full, dLdu: Full) -> Full:
    u_free: Free = self.model.to_free(u)
    preconditioner: Free = jnp.reciprocal(self.model.hess_diag(u_free))

    def matvec(p_free: Free) -> Free:
        return self.model.hess_prod(u_free, p_free)

    def preconditioner_fn(p_free: Free) -> Free:
        return preconditioner * p_free

    system = LinearSystem(
        matvec=matvec,
        b=-self.model.to_free(dLdu),
        rmatvec=matvec,
        preconditioner=preconditioner_fn,
        rpreconditioner=preconditioner_fn,
    )
    solution: LinearSolver.Solution = self.adjoint_solver.solve(
        system, jnp.zeros_like(u_free)
    )
    if not solution.success:
        logger.warning("Adjoint fail: %r", solution)
    logger.info("Adjoint time: %g sec", solution.stats.time)
    return self.model.to_full(solution.params, 0.0)

fun ¤

fun(params: ParamsT) -> tuple[Scalar, AuxT]
Source code in src/liblaf/apple/inverse/_inverse.py 
67
68
69
70
71
72
def fun(self, params: ParamsT) -> tuple[Scalar, AuxT]:
    model_params: ModelParams = self.make_params(params)
    self.model.update_params(model_params)
    solution: Optimizer.Solution = self.forward.step()
    logger.info("Forward time: %g sec", solution.stats.time)
    return self.loss(self.model.u_full, model_params)

loss abstractmethod ¤

loss(u: Full, params: ModelParams) -> tuple[Scalar, AuxT]
Source code in src/liblaf/apple/inverse/_inverse.py 
74
75
76
@abc.abstractmethod
def loss(self, u: Full, params: ModelParams) -> tuple[Scalar, AuxT]:
    raise NotImplementedError

loss_and_grad ¤

loss_and_grad(
    u: Full, params: ModelParams
) -> tuple[Scalar, Full, ModelParams, AuxT]
Source code in src/liblaf/apple/inverse/_inverse.py 
78
79
80
81
82
83
84
85
86
87
88
89
@eqx.filter_jit
def loss_and_grad(
    self, u: Full, params: ModelParams
) -> tuple[Scalar, Full, ModelParams, AuxT]:
    loss: Scalar
    aux: AuxT
    dLdu: Full
    dLdq: ModelParams
    (loss, aux), (dLdu, dLdq) = jax.value_and_grad(
        self.loss, argnums=(0, 1), has_aux=True
    )(u, params)
    return loss, dLdu, dLdq, aux

make_params abstractmethod ¤

make_params(params: ParamsT) -> ModelParams
Source code in src/liblaf/apple/inverse/_inverse.py 
91
92
93
@abc.abstractmethod
def make_params(self, params: ParamsT) -> ModelParams:
    raise NotImplementedError

solve ¤

solve(
    params: ParamsT, callback: Callback | None = None
) -> Solution
Source code in src/liblaf/apple/inverse/_inverse.py 
 95
 96
 97
 98
 99
100
101
102
103
104
def solve(
    self, params: ParamsT, callback: Callback | None = None
) -> Optimizer.Solution:
    objective = Objective(value_and_grad=self.value_and_grad)
    optimizer_solution: Optimizer.Solution = self.optimizer.minimize(
        objective, params, callback=callback
    )
    if not optimizer_solution.success:
        logger.warning("Inverse fail: %r", optimizer_solution)
    return optimizer_solution

value_and_grad ¤

value_and_grad(
    params: ParamsT,
) -> tuple[Scalar, ParamsT, AuxT]
Source code in src/liblaf/apple/inverse/_inverse.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def value_and_grad(self, params: ParamsT) -> tuple[Scalar, ParamsT, AuxT]:
    model_params: ModelParams
    model_params_vjp: Callable[[ModelParams], ParamsT]
    model_params, model_params_vjp = jax.vjp(self.make_params, params)
    self.model.update_params(model_params)
    solution: Optimizer.Solution = self.forward.step()
    logger.info("Forward time: %g sec", solution.stats.time)
    u_full: Full = self.model.u_full
    loss: Scalar
    dLdu: Full
    dLdq: ModelParams
    aux: AuxT
    loss, dLdu, dLdq, aux = self.loss_and_grad(u_full, model_params)
    p: Full = self.adjoint(u_full, dLdu)
    prod: ModelParams = self.model.mixed_derivative_prod(u_full, p)
    model_params_grad: ModelParams = jax.tree.map(operator.add, dLdq, prod)
    grad: ParamsT = model_params_vjp(model_params_grad)
    return loss, grad, aux

JaxEnergy ¤

Bases: IdMixin


              flowchart TD
              liblaf.apple.JaxEnergy[JaxEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.JaxEnergy
                


              click liblaf.apple.JaxEnergy href "" "liblaf.apple.JaxEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (frozenset[str], default: frozenset() ) –
Subclassed by: Used by:

Methods:

Attributes:

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

requires_grad class-attribute instance-attribute ¤

requires_grad: frozenset[str] = field(
    default=frozenset(), kw_only=True
)

fun ¤

fun(u: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_energy.py 
30
31
def fun(self, u: Vector) -> Scalar:
    raise NotImplementedError

grad ¤

grad(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
33
34
35
36
@eqx.filter_jit
def grad(self, u: Vector) -> Updates:
    values: Vector = eqx.filter_grad(self.fun)(u)
    return values, jnp.arange(u.shape[0])

grad_and_hess_diag ¤

grad_and_hess_diag(u: Vector) -> tuple[Updates, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
61
62
63
@eqx.filter_jit
def grad_and_hess_diag(self, u: Vector) -> tuple[Updates, Updates]:
    return self.grad(u), self.hess_diag(u)

hess_diag ¤

hess_diag(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
38
39
def hess_diag(self, u: Vector) -> Updates:
    raise NotImplementedError

hess_prod ¤

hess_prod(u: Vector, p: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
41
42
43
44
45
@eqx.filter_jit
def hess_prod(self, u: Vector, p: Vector) -> Updates:
    values: Vector
    _, values = jax.jvp(jax.grad(self.fun), (u,), (p,))
    return values, jnp.arange(u.shape[0])

hess_quad ¤

hess_quad(u: Vector, p: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_energy.py 
47
48
49
50
51
52
@eqx.filter_jit
def hess_quad(self, u: Vector, p: Vector) -> Scalar:
    values: Vector
    index: Index
    values, index = self.hess_prod(u, p)
    return jnp.vdot(p[index], values)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> EnergyParams
Source code in src/liblaf/apple/jax/model/_energy.py 
65
66
67
68
69
70
71
@eqx.filter_jit
def mixed_derivative_prod(self, u: Vector, p: Vector) -> EnergyParams:
    outputs: EnergyParams = {
        name: getattr(self, f"mixed_derivative_prod_{name}")(u, p)
        for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
22
23
24
@eqx.filter_jit
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: Mapping[str, Array]) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
26
27
28
def update_params(self, params: Mapping[str, Array]) -> None:
    for name, value in params.items():
        setattr(self, name, value)

value_and_grad ¤

value_and_grad(u: Vector) -> tuple[Scalar, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
54
55
56
57
58
59
@eqx.filter_jit
def value_and_grad(self, u: Vector) -> tuple[Scalar, Updates]:
    value: Scalar
    grad: Vector
    value, grad = jax.value_and_grad(self.fun)(u)
    return value, (grad, jnp.arange(u.shape[0]))

JaxModel ¤

Parameters:

  • energies ¤

    (dict[str, JaxEnergy], default: <class 'dict'> ) –

    dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object’s (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

Returned by: Used by:

Methods:

Attributes:

energies class-attribute instance-attribute ¤

energies: dict[str, JaxEnergy] = field(
    factory=dict, kw_only=True
)

fun ¤

fun(x: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_model.py 
29
30
31
32
33
34
@eqx.filter_jit
def fun(self, x: Vector) -> Scalar:
    output: Scalar = jnp.zeros(())
    for energy in self.energies.values():
        output += energy.fun(x)
    return output

grad ¤

grad(x: Vector) -> Vector
Source code in src/liblaf/apple/jax/model/_model.py 
36
37
38
39
40
41
42
43
44
@eqx.filter_jit
def grad(self, x: Vector) -> Vector:
    output: Vector = jnp.zeros_like(x)
    for energy in self.energies.values():
        grad: Vector
        index: Index
        grad, index = energy.grad(x)
        output = output.at[index].add(grad)
    return output

grad_and_hess_diag ¤

grad_and_hess_diag(x: Vector) -> tuple[Vector, Vector]
Source code in src/liblaf/apple/jax/model/_model.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@eqx.filter_jit
def grad_and_hess_diag(self, x: Vector) -> tuple[Vector, Vector]:
    grad: Vector = jnp.zeros_like(x)
    hess_diag: Vector = jnp.zeros_like(x)
    for energy in self.energies.values():
        grad_i: Vector
        index_g: Index
        hess_diag_i: Vector
        index_h: Index
        (grad_i, index_g), (hess_diag_i, index_h) = energy.grad_and_hess_diag(x)
        grad = grad.at[index_g].add(grad_i)
        hess_diag = hess_diag.at[index_h].add(hess_diag_i)
    return grad, hess_diag

hess_diag ¤

hess_diag(x: Vector) -> Vector
Source code in src/liblaf/apple/jax/model/_model.py 
46
47
48
49
50
51
52
53
54
@eqx.filter_jit
def hess_diag(self, x: Vector) -> Vector:
    output: Vector = jnp.zeros_like(x)
    for energy in self.energies.values():
        diag: Vector
        index: Index
        diag, index = energy.hess_diag(x)
        output = output.at[index].add(diag)
    return output

hess_prod ¤

hess_prod(x: Vector, p: Vector) -> Vector
Source code in src/liblaf/apple/jax/model/_model.py 
56
57
58
59
60
61
62
63
64
@eqx.filter_jit
def hess_prod(self, x: Vector, p: Vector) -> Vector:
    output: Vector = jnp.zeros_like(x)
    for energy in self.energies.values():
        prod: Vector
        index: Index
        prod, index = energy.hess_prod(x, p)
        output = output.at[index].add(prod)
    return output

hess_quad ¤

hess_quad(x: Vector, p: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_model.py 
66
67
68
69
70
71
@eqx.filter_jit
def hess_quad(self, x: Vector, p: Vector) -> Scalar:
    output: Scalar = jnp.zeros(())
    for energy in self.energies.values():
        output += energy.hess_quad(x, p)
    return output

mixed_derivative_prod ¤

mixed_derivative_prod(
    x: Vector, p: Vector
) -> ModelParams
Source code in src/liblaf/apple/jax/model/_model.py 
73
74
75
76
77
78
@eqx.filter_jit
def mixed_derivative_prod(self, x: Vector, p: Vector) -> ModelParams:
    return {
        name: energy.mixed_derivative_prod(x, p)
        for name, energy in self.energies.items()
    }

update ¤

update(x: Vector) -> None
Source code in src/liblaf/apple/jax/model/_model.py 
21
22
23
def update(self, x: Vector) -> None:
    for energy in self.energies.values():
        energy.update(x)

update_params ¤

update_params(params: ModelParams) -> None
Source code in src/liblaf/apple/jax/model/_model.py 
25
26
27
def update_params(self, params: ModelParams) -> None:
    for name, energy_params in params.items():
        self.energies[name].update_params(energy_params)

value_and_grad ¤

value_and_grad(x: Vector) -> tuple[Scalar, Vector]
Source code in src/liblaf/apple/jax/model/_model.py 
80
81
82
83
84
85
86
87
88
89
90
@eqx.filter_jit
def value_and_grad(self, x: Vector) -> tuple[Scalar, Vector]:
    value: Scalar = jnp.zeros(())
    grad: Vector = jnp.zeros_like(x)
    for energy in self.energies.values():
        value_i: Scalar
        grad_i: Vector
        value_i, (grad_i, index) = energy.value_and_grad(x)
        value += value_i
        grad = grad.at[index].add(grad_i)
    return value, grad

JaxModelBuilder ¤

Parameters:

  • energies ¤

    (dict[str, JaxEnergy], default: <class 'dict'> ) –

    dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object’s (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

Returned by: Used by:

Methods:

Attributes:

energies class-attribute instance-attribute ¤

energies: dict[str, JaxEnergy] = field(
    factory=dict, kw_only=True
)

add_energy ¤

add_energy(energy: JaxEnergy) -> None
Source code in src/liblaf/apple/jax/model/_builder.py 
11
12
def add_energy(self, energy: JaxEnergy) -> None:
    self.energies[energy.id] = energy

finalize ¤

finalize() -> JaxModel
Source code in src/liblaf/apple/jax/model/_builder.py 
14
15
def finalize(self) -> JaxModel:
    return JaxModel(energies=self.energies)

MassSpring ¤

Bases: JaxEnergy


              flowchart TD
              liblaf.apple.MassSpring[MassSpring]
              liblaf.apple.jax.model._energy.JaxEnergy[JaxEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.jax.model._energy.JaxEnergy --> liblaf.apple.MassSpring
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.jax.model._energy.JaxEnergy
                



              click liblaf.apple.MassSpring href "" "liblaf.apple.MassSpring"
              click liblaf.apple.jax.model._energy.JaxEnergy href "" "liblaf.apple.jax.model._energy.JaxEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (frozenset[str], default: frozenset() ) –

  • edges ¤

    (Integer[Array, 'edges 2']) –

  • length ¤

    (Float[Array, edges]) –

  • points ¤

    (Float[Array, 'edges 2 3']) –

  • stiffness ¤

    (Float[Array, edges]) –

Methods:

Attributes:

edges instance-attribute ¤

edges: Integer[Array, ' edges 2']

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

length instance-attribute ¤

length: Float[Array, ' edges']

n_edges property ¤

n_edges: int

points instance-attribute ¤

points: Float[Array, 'edges 2 3']

requires_grad class-attribute instance-attribute ¤

requires_grad: frozenset[str] = field(
    default=frozenset(), kw_only=True
)

stiffness instance-attribute ¤

stiffness: Float[Array, ' edges']

from_pyvista classmethod ¤

from_pyvista(obj: PolyData) -> Self
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@classmethod
def from_pyvista(cls, obj: pv.PolyData) -> Self:
    if LENGTH not in obj.cell_data:
        obj = obj.compute_cell_sizes(length=True, area=False, volume=False)  # pyright: ignore[reportAssignmentType]
    point_id: Integer[np.ndarray, " points"] = obj.point_data[POINT_ID]
    edges: Integer[np.ndarray, "edges 2"] = obj.lines.reshape((-1, 3))[:, 1:]
    length: Float[Array, " edges"] = jnp.asarray(obj.cell_data[LENGTH])
    if jnp.any(length < 0.0):
        logger.warning("Length < 0")
    return cls(
        edges=jnp.asarray(point_id[edges]),
        length=length,
        points=jnp.asarray(obj.points[edges]),
        stiffness=jnp.asarray(obj.cell_data[STIFFNESS]),
    )

fun ¤

fun(u: Vector) -> Scalar
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
49
50
51
52
53
54
55
56
57
58
@override
def fun(self, u: Vector) -> Scalar:
    x: Float[Array, "edges 2 3"] = self.points + u[self.edges]
    delta: Float[Array, "edges 3"] = x[:, 1, :] - x[:, 0, :]
    energy: Float[Array, " edges"] = (
        0.5
        * self.stiffness
        * jnp.square(jnp.linalg.norm(delta, axis=-1) - self.length)
    )
    return jnp.sum(energy)

grad ¤

grad(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@override
def grad(self, u: Vector) -> Updates:
    x: Float[Array, "edges 2 3"] = self.points + u[self.edges]
    delta: Float[Array, "edges 3"] = x[:, 1, :] - x[:, 0, :]
    length: Float[Array, " edges"] = jnp.linalg.norm(delta, axis=-1)
    direction: Float[Array, "edges 3"] = (
        delta / jnp.where(length > 0, length, 1.0)[:, jnp.newaxis]
    )
    force: Float[Array, "edges 3"] = (
        self.stiffness[:, jnp.newaxis]
        * (length - self.length)[:, jnp.newaxis]
        * direction
    )
    grad: Float[Array, "edges 2 3"] = jnp.stack([-force, force], axis=1)
    return grad, self.edges.flatten()

grad_and_hess_diag ¤

grad_and_hess_diag(u: Vector) -> tuple[Updates, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
61
62
63
@eqx.filter_jit
def grad_and_hess_diag(self, u: Vector) -> tuple[Updates, Updates]:
    return self.grad(u), self.hess_diag(u)

hess_diag ¤

hess_diag(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
76
77
78
79
80
81
@override
def hess_diag(self, u: Vector) -> Updates:
    values: Float[Array, "edges*2 3"] = einops.repeat(
        self.stiffness, "edges -> (edges i) j", i=2, j=3
    )
    return values, self.edges.flatten()

hess_prod ¤

hess_prod(u: Vector, p: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
41
42
43
44
45
@eqx.filter_jit
def hess_prod(self, u: Vector, p: Vector) -> Updates:
    values: Vector
    _, values = jax.jvp(jax.grad(self.fun), (u,), (p,))
    return values, jnp.arange(u.shape[0])

hess_quad ¤

hess_quad(u: Vector, p: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_energy.py 
47
48
49
50
51
52
@eqx.filter_jit
def hess_quad(self, u: Vector, p: Vector) -> Scalar:
    values: Vector
    index: Index
    values, index = self.hess_prod(u, p)
    return jnp.vdot(p[index], values)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> EnergyParams
Source code in src/liblaf/apple/jax/model/_energy.py 
65
66
67
68
69
70
71
@eqx.filter_jit
def mixed_derivative_prod(self, u: Vector, p: Vector) -> EnergyParams:
    outputs: EnergyParams = {
        name: getattr(self, f"mixed_derivative_prod_{name}")(u, p)
        for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
22
23
24
@eqx.filter_jit
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: Mapping[str, Array]) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
26
27
28
def update_params(self, params: Mapping[str, Array]) -> None:
    for name, value in params.items():
        setattr(self, name, value)

value_and_grad ¤

value_and_grad(u: Vector) -> tuple[Scalar, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
54
55
56
57
58
59
@eqx.filter_jit
def value_and_grad(self, u: Vector) -> tuple[Scalar, Updates]:
    value: Scalar
    grad: Vector
    value, grad = jax.value_and_grad(self.fun)(u)
    return value, (grad, jnp.arange(u.shape[0]))

MassSpringPrestrain ¤

Bases: MassSpring


              flowchart TD
              liblaf.apple.MassSpringPrestrain[MassSpringPrestrain]
              liblaf.apple.jax.energies._mass_spring.MassSpring[MassSpring]
              liblaf.apple.jax.model._energy.JaxEnergy[JaxEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.jax.energies._mass_spring.MassSpring --> liblaf.apple.MassSpringPrestrain
                                liblaf.apple.jax.model._energy.JaxEnergy --> liblaf.apple.jax.energies._mass_spring.MassSpring
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.jax.model._energy.JaxEnergy
                




              click liblaf.apple.MassSpringPrestrain href "" "liblaf.apple.MassSpringPrestrain"
              click liblaf.apple.jax.energies._mass_spring.MassSpring href "" "liblaf.apple.jax.energies._mass_spring.MassSpring"
              click liblaf.apple.jax.model._energy.JaxEnergy href "" "liblaf.apple.jax.model._energy.JaxEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (frozenset[str], default: frozenset() ) –

  • edges ¤

    (Integer[Array, 'edges 2']) –

  • length ¤

    (Float[Array, edges]) –

  • points ¤

    (Float[Array, 'edges 2 3']) –

  • stiffness ¤

    (Float[Array, edges]) –

Methods:

Attributes:

edges instance-attribute ¤

edges: Integer[Array, ' edges 2']

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

length instance-attribute ¤

length: Float[Array, ' edges']

n_edges property ¤

n_edges: int

points instance-attribute ¤

points: Float[Array, 'edges 2 3']

requires_grad class-attribute instance-attribute ¤

requires_grad: frozenset[str] = field(
    default=frozenset(), kw_only=True
)

stiffness instance-attribute ¤

stiffness: Float[Array, ' edges']

from_pyvista classmethod ¤

from_pyvista(obj: PolyData) -> Self
Source code in src/liblaf/apple/jax/energies/_mass_spring_prestrain.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
@classmethod
def from_pyvista(cls, obj: pv.PolyData) -> Self:
    if LENGTH not in obj.cell_data:
        obj = obj.compute_cell_sizes(length=True, area=False, volume=False)  # pyright: ignore[reportAssignmentType]
    point_id: Integer[np.ndarray, " points"] = obj.point_data[POINT_ID]
    edges: Integer[np.ndarray, "edges 2"] = obj.lines.reshape((-1, 3))[:, 1:]
    length: Float[Array, " edges"] = jnp.asarray(obj.cell_data[LENGTH])
    if jnp.any(length < 0.0):
        logger.warning("Length < 0")
    prestrain: Float[Array, " edges"] = jnp.asarray(obj.cell_data[PRESTRAIN])
    return cls(
        edges=jnp.asarray(point_id[edges]),
        length=length * (1.0 + prestrain),
        points=jnp.asarray(obj.points[edges]),
        stiffness=jnp.asarray(obj.cell_data[STIFFNESS]),
    )

fun ¤

fun(u: Vector) -> Scalar
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
49
50
51
52
53
54
55
56
57
58
@override
def fun(self, u: Vector) -> Scalar:
    x: Float[Array, "edges 2 3"] = self.points + u[self.edges]
    delta: Float[Array, "edges 3"] = x[:, 1, :] - x[:, 0, :]
    energy: Float[Array, " edges"] = (
        0.5
        * self.stiffness
        * jnp.square(jnp.linalg.norm(delta, axis=-1) - self.length)
    )
    return jnp.sum(energy)

grad ¤

grad(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@override
def grad(self, u: Vector) -> Updates:
    x: Float[Array, "edges 2 3"] = self.points + u[self.edges]
    delta: Float[Array, "edges 3"] = x[:, 1, :] - x[:, 0, :]
    length: Float[Array, " edges"] = jnp.linalg.norm(delta, axis=-1)
    direction: Float[Array, "edges 3"] = (
        delta / jnp.where(length > 0, length, 1.0)[:, jnp.newaxis]
    )
    force: Float[Array, "edges 3"] = (
        self.stiffness[:, jnp.newaxis]
        * (length - self.length)[:, jnp.newaxis]
        * direction
    )
    grad: Float[Array, "edges 2 3"] = jnp.stack([-force, force], axis=1)
    return grad, self.edges.flatten()

grad_and_hess_diag ¤

grad_and_hess_diag(u: Vector) -> tuple[Updates, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
61
62
63
@eqx.filter_jit
def grad_and_hess_diag(self, u: Vector) -> tuple[Updates, Updates]:
    return self.grad(u), self.hess_diag(u)

hess_diag ¤

hess_diag(u: Vector) -> Updates
Source code in src/liblaf/apple/jax/energies/_mass_spring.py 
76
77
78
79
80
81
@override
def hess_diag(self, u: Vector) -> Updates:
    values: Float[Array, "edges*2 3"] = einops.repeat(
        self.stiffness, "edges -> (edges i) j", i=2, j=3
    )
    return values, self.edges.flatten()

hess_prod ¤

hess_prod(u: Vector, p: Vector) -> Updates
Source code in src/liblaf/apple/jax/model/_energy.py 
41
42
43
44
45
@eqx.filter_jit
def hess_prod(self, u: Vector, p: Vector) -> Updates:
    values: Vector
    _, values = jax.jvp(jax.grad(self.fun), (u,), (p,))
    return values, jnp.arange(u.shape[0])

hess_quad ¤

hess_quad(u: Vector, p: Vector) -> Scalar
Source code in src/liblaf/apple/jax/model/_energy.py 
47
48
49
50
51
52
@eqx.filter_jit
def hess_quad(self, u: Vector, p: Vector) -> Scalar:
    values: Vector
    index: Index
    values, index = self.hess_prod(u, p)
    return jnp.vdot(p[index], values)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> EnergyParams
Source code in src/liblaf/apple/jax/model/_energy.py 
65
66
67
68
69
70
71
@eqx.filter_jit
def mixed_derivative_prod(self, u: Vector, p: Vector) -> EnergyParams:
    outputs: EnergyParams = {
        name: getattr(self, f"mixed_derivative_prod_{name}")(u, p)
        for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
22
23
24
@eqx.filter_jit
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: Mapping[str, Array]) -> None
Source code in src/liblaf/apple/jax/model/_energy.py 
26
27
28
def update_params(self, params: Mapping[str, Array]) -> None:
    for name, value in params.items():
        setattr(self, name, value)

value_and_grad ¤

value_and_grad(u: Vector) -> tuple[Scalar, Updates]
Source code in src/liblaf/apple/jax/model/_energy.py 
54
55
56
57
58
59
@eqx.filter_jit
def value_and_grad(self, u: Vector) -> tuple[Scalar, Updates]:
    value: Scalar
    grad: Vector
    value, grad = jax.value_and_grad(self.fun)(u)
    return value, (grad, jnp.arange(u.shape[0]))

Model ¤

Parameters:

Returned by:

Methods:

Attributes:

dim property ¤

dim: int

dirichlet instance-attribute ¤

dirichlet: Dirichlet

jax instance-attribute ¤

jax: JaxModel

n_free property ¤

n_free: int

n_full property ¤

n_full: int

n_points property ¤

n_points: int

u_free property writable ¤

u_free: Free

u_full instance-attribute ¤

u_full: Full

warp instance-attribute ¤

warp: WarpModelAdapter

fun ¤

fun(u: FreeOrFull) -> Scalar
Source code in src/liblaf/apple/model/_model.py 
84
85
86
87
88
89
90
def fun(self, u: FreeOrFull) -> Scalar:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    output_jax: Scalar = self.jax.fun(u_full)
    output_wp: Scalar = self.warp.fun(u_full)
    output: Scalar = output_jax + output_wp
    return output

grad ¤

grad(u: FreeOrFull) -> FreeOrFull
Source code in src/liblaf/apple/model/_model.py 
92
93
94
95
96
97
98
def grad(self, u: FreeOrFull) -> FreeOrFull:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    output_jax: Full = self.jax.grad(u_full)
    output_wp: Full = self.warp.grad(u_full)
    output: Full = output_jax + output_wp
    return self.to_shape_like(output, u)

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: FreeOrFull,
) -> tuple[FreeOrFull, FreeOrFull]
Source code in src/liblaf/apple/model/_model.py 
148
149
150
151
152
153
154
155
156
157
158
159
def grad_and_hess_diag(self, u: FreeOrFull) -> tuple[FreeOrFull, FreeOrFull]:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    grad_jax: Full
    hess_diag_jax: Full
    grad_jax, hess_diag_jax = self.jax.grad_and_hess_diag(u_full)
    grad_wp: Full
    hess_diag_wp: Full
    grad_wp, hess_diag_wp = self.warp.grad_and_hess_diag(u_full)
    grad: Full = grad_jax + grad_wp
    hess_diag: Full = hess_diag_jax + hess_diag_wp
    return self.to_shape_like(grad, u), self.to_shape_like(hess_diag, u)

hess_diag ¤

hess_diag(u: FreeOrFull) -> FreeOrFull
Source code in src/liblaf/apple/model/_model.py 
100
101
102
103
104
105
106
def hess_diag(self, u: FreeOrFull) -> FreeOrFull:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    output_jax: Full = self.jax.hess_diag(u_full)
    output_wp: Full = self.warp.hess_diag(u_full)
    output: Full = output_jax + output_wp
    return self.to_shape_like(output, u)

hess_prod ¤

hess_prod(u: FreeOrFull, p: FreeOrFull) -> FreeOrFull
Source code in src/liblaf/apple/model/_model.py 
108
109
110
111
112
113
114
115
def hess_prod(self, u: FreeOrFull, p: FreeOrFull) -> FreeOrFull:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    p_full: Full = self.to_full(p, 0.0)
    output_jax: Full = self.jax.hess_prod(u_full, p_full)
    output_wp: Full = self.warp.hess_prod(u_full, p_full)
    output: Full = output_jax + output_wp
    return self.to_shape_like(output, u)

hess_quad ¤

hess_quad(u: FreeOrFull, p: FreeOrFull) -> Scalar
Source code in src/liblaf/apple/model/_model.py 
117
118
119
120
121
122
123
124
def hess_quad(self, u: FreeOrFull, p: FreeOrFull) -> Scalar:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    p_full: Full = self.to_full(p, 0.0)
    output_jax: Scalar = self.jax.hess_quad(u_full, p_full)
    output_wp: Scalar = self.warp.hess_quad(u_full, p_full)
    output: Scalar = output_jax + output_wp
    return output

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: FreeOrFull, p: FreeOrFull
) -> ModelParams
Source code in src/liblaf/apple/model/_model.py 
126
127
128
129
130
131
132
133
def mixed_derivative_prod(self, u: FreeOrFull, p: FreeOrFull) -> ModelParams:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    p_full: Full = self.to_full(p, 0.0)
    outputs_jax: ModelParams = self.jax.mixed_derivative_prod(u_full, p_full)
    outputs_wp: ModelParams = self.warp.mixed_derivative_prod(u_full, p_full)
    outputs: ModelParams = tlz.merge(outputs_jax, outputs_wp)
    return outputs

to_free ¤

to_free(u: FreeOrFull) -> Free
Source code in src/liblaf/apple/model/_model.py 
50
51
52
53
def to_free(self, u: FreeOrFull) -> Free:
    if u.size == self.n_free:
        return u.reshape((self.n_free,))
    return self.dirichlet.get_free(u)

to_full ¤

to_full(
    u: FreeOrFull,
    dirichlet: Float[ArrayLike, " dirichlet"] | None = None,
) -> Full
Source code in src/liblaf/apple/model/_model.py 
55
56
57
58
59
60
def to_full(
    self, u: FreeOrFull, dirichlet: Float[ArrayLike, " dirichlet"] | None = None
) -> Full:
    if u.size == self.n_full:
        return u.reshape((self.n_points, self.dim))
    return self.dirichlet.to_full(u, dirichlet)

to_shape_like ¤

to_shape_like(u_full: Full, like: FreeOrFull) -> FreeOrFull
Source code in src/liblaf/apple/model/_model.py 
62
63
64
65
def to_shape_like(self, u_full: Full, like: FreeOrFull) -> FreeOrFull:
    if u_full.size == like.size:
        return u_full.reshape(like.shape)
    return self.dirichlet.get_free(u_full)

update ¤

update(u: FreeOrFull) -> None
Source code in src/liblaf/apple/model/_model.py 
67
68
69
70
71
72
73
def update(self, u: FreeOrFull) -> None:
    u_full: Full = self.to_full(u)
    if jnp.array_equiv(u_full, self.u_full):
        return
    self.u_full = u_full
    self.jax.update(u_full)
    self.warp.update(u_full)

update_params ¤

update_params(params: ModelParams) -> None
Source code in src/liblaf/apple/model/_model.py 
75
76
77
78
79
80
81
82
def update_params(self, params: ModelParams) -> None:
    def pick(allowlist: Container[str], d: ModelParams) -> ModelParams:
        return tlz.keyfilter(lambda name: name in allowlist, d)

    params_jax: ModelParams = pick(self.jax.energies, params)
    params_warp: ModelParams = pick(self.warp.energies, params)
    self.jax.update_params(params_jax)
    self.warp.update_params(params_warp)

value_and_grad ¤

value_and_grad(u: FreeOrFull) -> tuple[Scalar, FreeOrFull]
Source code in src/liblaf/apple/model/_model.py 
135
136
137
138
139
140
141
142
143
144
145
146
def value_and_grad(self, u: FreeOrFull) -> tuple[Scalar, FreeOrFull]:
    u_full: Full = self.to_full(u)
    self.update(u_full)
    value_jax: Scalar
    grad_jax: Full
    value_jax, grad_jax = self.jax.value_and_grad(u_full)
    value_wp: Scalar
    grad_wp: Full
    value_wp, grad_wp = self.warp.value_and_grad(u_full)
    value: Scalar = value_jax + value_wp
    grad: Full = grad_jax + grad_wp
    return value, self.to_shape_like(grad, u)

ModelBuilder ¤

ModelBuilder(dim: int = 3)

Parameters:

Methods:

Attributes:

Source code in src/liblaf/apple/model/_builder.py 
27
28
29
def __init__(self, dim: int = 3) -> None:
    dirichlet: DirichletBuilder = DirichletBuilder(dim=dim)
    self.__attrs_init__(dirichlet=dirichlet)  # pyright: ignore[reportAttributeAccessIssue]

dirichlet class-attribute instance-attribute ¤

dirichlet: DirichletBuilder = field(
    factory=DirichletBuilder
)

jax class-attribute instance-attribute ¤

jax: JaxModelBuilder = field(factory=JaxModelBuilder)

n_points property ¤

n_points: int

warp class-attribute instance-attribute ¤

warp: WarpModelBuilder = field(factory=WarpModelBuilder)

add_dirichlet ¤

add_dirichlet(obj: DataSet) -> None
Source code in src/liblaf/apple/model/_builder.py 
35
36
def add_dirichlet(self, obj: pv.DataSet) -> None:
    self.dirichlet.add_pyvista(obj)

add_energy ¤

add_energy(energy: JaxEnergy | WarpEnergy) -> None
Source code in src/liblaf/apple/model/_builder.py 
38
39
40
41
42
43
44
def add_energy(self, energy: JaxEnergy | WarpEnergy) -> None:
    if isinstance(energy, JaxEnergy):
        self.jax.add_energy(energy)
    elif isinstance(energy, WarpEnergy):
        self.warp.add_energy(energy)
    else:
        raise TypeError(energy)

assign_global_ids ¤

assign_global_ids[T: DataSet](obj: T) -> T
Source code in src/liblaf/apple/model/_builder.py 
46
47
48
49
50
51
def assign_global_ids[T: pv.DataSet](self, obj: T) -> T:
    start: int = self.n_points
    stop: int = start + obj.n_points
    self.dirichlet.resize(stop)
    obj.point_data[POINT_ID] = np.arange(start, stop)
    return obj

finalize ¤

finalize() -> Model
Source code in src/liblaf/apple/model/_builder.py 
53
54
55
56
57
58
59
60
61
62
def finalize(self) -> Model:
    dirichlet: Dirichlet = self.dirichlet.finalize()
    u_full: Full = jnp.zeros((self.dirichlet.n_points, self.dirichlet.dim))
    u_full = dirichlet.set_dirichlet(u_full)
    return Model(
        dirichlet=dirichlet,
        u_full=u_full,
        jax=self.jax.finalize(),
        warp=WarpModelAdapter(self.warp.finalize()),
    )

Phace ¤

Bases: Hyperelastic


              flowchart TD
              liblaf.apple.Phace[Phace]
              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic[Hyperelastic]
              liblaf.apple.warp.model._energy.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic --> liblaf.apple.Phace
                                liblaf.apple.warp.model._energy.WarpEnergy --> liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic
                                liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.warp.model._energy.WarpEnergy
                




              click liblaf.apple.Phace href "" "liblaf.apple.Phace"
              click liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic href "" "liblaf.apple.warp.energies.elastic.hyperelastic._base.Hyperelastic"
              click liblaf.apple.warp.model._energy.WarpEnergy href "" "liblaf.apple.warp.model._energy.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

  • id ¤

    (str, default: <dynamic> ) –

  • requires_grad ¤

    (Sequence[str], default: () ) –

  • cells ¤

    (Integer[array, 'c a']) –

  • dhdX ¤

    (Integer[array, 'c q a J']) –

  • dV ¤

    (Integer[array, 'c q']) –

  • params ¤

    (Struct) –

  • clamp_hess_diag ¤

    (bool, default: True ) –

  • clamp_hess_quad ¤

    (bool, default: True ) –

  • clamp_lambda ¤

    (bool, default: True ) –

Classes:

Methods:

Attributes:

cells instance-attribute ¤

cells: Integer[array, 'c a']

clamp_hess_diag class-attribute instance-attribute ¤

clamp_hess_diag: bool = True

clamp_hess_quad class-attribute instance-attribute ¤

clamp_hess_quad: bool = True

clamp_lambda class-attribute instance-attribute ¤

clamp_lambda: bool = True

dV instance-attribute ¤

dV: Integer[array, 'c q']

dhdX instance-attribute ¤

dhdX: Integer[array, 'c q a J']

fun_kernel cached property ¤

fun_kernel: Kernel

grad_and_hess_diag_kernel cached property ¤

grad_and_hess_diag_kernel: Kernel

grad_kernel cached property ¤

grad_kernel: Kernel

hess_diag_kernel cached property ¤

hess_diag_kernel: Kernel

hess_prod_kernel cached property ¤

hess_prod_kernel: Kernel

hess_quad_kernel cached property ¤

hess_quad_kernel: Kernel

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

n_cells property ¤

n_cells: int

n_quadrature_points property ¤

n_quadrature_points: int

params instance-attribute ¤

params: Params

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

value_and_grad_kernel cached property ¤

value_and_grad_kernel: Kernel

Params ¤

Used by:

Attributes:

activation instance-attribute ¤

activation: array(dtype=vec6)

lambda_ instance-attribute ¤

lambda_: array(dtype=float_)

mu instance-attribute ¤

mu: array(dtype=float_)

muscle_fraction instance-attribute ¤

muscle_fraction: array(dtype=float_)

ParamsElem ¤

Returned by: Used by:

Attributes:

activation instance-attribute ¤

activation: vec6

lambda_ instance-attribute ¤

lambda_: float_

mu instance-attribute ¤

mu: float_

muscle_fraction instance-attribute ¤

muscle_fraction: float_

energy_density_func staticmethod ¤

energy_density_func(
    F: mat33, params: ParamsElem
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_func(F: mat33, params: ParamsElem) -> scalar:
    _1 = F.dtype(1.0)
    _2 = F.dtype(2.0)
    J = func.I3(F)  # float
    Psi_ARAP_active = ArapMuscle.energy_density_func(
        F, Phace._arap_active_params(params)
    )  # float
    Psi_ARAP_passive = Arap.energy_density_func(
        F, Phace._arap_params(params)
    )  # float
    Psi_ARAP = (
        params.muscle_fraction * Psi_ARAP_active
        + (_1 - params.muscle_fraction) * Psi_ARAP_passive
    )  # float
    Psi_VP = params.lambda_ * math.square(J - _1)  # float
    Psi = _2 * Psi_ARAP + Psi_VP  # float
    return Psi

energy_density_hess_diag_func staticmethod ¤

energy_density_hess_diag_func(
    F: mat33,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_diag_func(
    F: mat33, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    _1 = F.dtype(1.0)
    _2 = F.dtype(2.0)
    J = func.I3(F)  # float
    g3 = func.g3(F)  # mat33
    diag_arap_active = ArapMuscle.energy_density_hess_diag_func(
        F, dhdX, Phace._arap_active_params(params), clamp=clamp
    )  # mat43
    diag_arap_passive = Arap.energy_density_hess_diag_func(
        F, dhdX, Phace._arap_params(params), clamp=clamp
    )  # mat43
    diag_arap = (
        params.muscle_fraction * diag_arap_active
        + (_1 - params.muscle_fraction) * diag_arap_passive
    )  # mat43
    d2Psi_dI32 = _2 * params.lambda_  # float
    dPsi_dI3 = _2 * params.lambda_ * (J - _1)  # float
    h3_diag = func.h3_diag(dhdX, g3)  # mat43
    h6_diag = func.h6_diag(dhdX, F)  # mat43
    diag_vp = d2Psi_dI32 * h3_diag + dPsi_dI3 * h6_diag  # mat43
    diag = _2 * diag_arap + diag_vp  # mat43
    return diag

energy_density_hess_prod_func staticmethod ¤

energy_density_hess_prod_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_prod_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> mat33:
    _1 = F.dtype(1.0)
    _2 = F.dtype(2.0)
    J = func.I3(F)  # float
    g3 = func.g3(F)  # mat33
    prod_arap_active = ArapMuscle.energy_density_hess_prod_func(
        F, p, dhdX, Phace._arap_active_params(params), clamp=clamp
    )  # mat43
    prod_arap_passive = Arap.energy_density_hess_prod_func(
        F, p, dhdX, Phace._arap_params(params), clamp=clamp
    )  # mat43
    prod_arap = (
        params.muscle_fraction * prod_arap_active
        + (_1 - params.muscle_fraction) * prod_arap_passive
    )  # mat43
    d2Psi_dI32 = _2 * params.lambda_  # float
    dPsi_dI3 = _2 * params.lambda_ * (J - _1)  # float
    h3_prod = func.h3_prod(p, dhdX, g3)  # mat43
    h6_prod = func.h6_prod(p, dhdX, F)  # mat43
    prod_vp = d2Psi_dI32 * h3_prod + dPsi_dI3 * h6_prod  # mat43
    prod = _2 * prod_arap + prod_vp  # mat43
    return prod

energy_density_hess_quad_func staticmethod ¤

energy_density_hess_quad_func(
    F: mat33,
    p: mat43,
    dhdX: mat43,
    params: ParamsElem,
    *,
    clamp: bool = True,
) -> scalar
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
@override
@staticmethod
@no_type_check
@wp.func
def energy_density_hess_quad_func(
    F: mat33, p: mat43, dhdX: mat43, params: ParamsElem, *, clamp: bool = True
) -> scalar:
    _1 = F.dtype(1.0)
    _2 = F.dtype(2.0)
    J = func.I3(F)  # float
    g3 = func.g3(F)  # mat33
    quad_arap_active = ArapMuscle.energy_density_hess_quad_func(
        F, p, dhdX, Phace._arap_active_params(params), clamp=clamp
    )
    quad_arap_passive = Arap.energy_density_hess_quad_func(
        F, p, dhdX, Phace._arap_params(params), clamp=clamp
    )
    quad_arap = (
        params.muscle_fraction * quad_arap_active
        + (_1 - params.muscle_fraction) * quad_arap_passive
    )
    d2Psi_dI32 = _2 * params.lambda_  # float
    dPsi_dI3 = _2 * params.lambda_ * (J - _1)  # float
    h3_quad = func.h3_quad(p, dhdX, g3)  # float
    h6_quad = func.h6_quad(p, dhdX, F)  # float
    quad_vp = d2Psi_dI32 * h3_quad + dPsi_dI3 * h6_quad  # float
    quad = _2 * quad_arap + quad_vp  # float
    return quad

first_piola_kirchhoff_stress_func staticmethod ¤

first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
@override
@staticmethod
@no_type_check
@wp.func
def first_piola_kirchhoff_stress_func(
    F: mat33, params: ParamsElem, *, clamp: bool = False
) -> mat33:
    _1 = F.dtype(1.0)
    _2 = F.dtype(2.0)
    J = func.I3(F)  # float
    g3 = func.g3(F)  # mat33
    PK1_ARAP_active = ArapMuscle.first_piola_kirchhoff_stress_func(
        F, Phace._arap_active_params(params), clamp=clamp
    )  # mat33
    PK1_ARAP_passive = Arap.first_piola_kirchhoff_stress_func(
        F, Phace._arap_params(params), clamp=clamp
    )  # mat33
    PK1_ARAP = (
        params.muscle_fraction * PK1_ARAP_active
        + (_1 - params.muscle_fraction) * PK1_ARAP_passive
    )  # mat33
    PK1_VP = _2 * params.lambda_ * (J - _1) * g3  # mat33
    PK1 = _2 * PK1_ARAP + PK1_VP  # mat33
    return PK1

from_pyvista classmethod ¤

from_pyvista(
    obj: DataObject,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
59
60
61
62
@classmethod
def from_pyvista(cls, obj: pv.DataObject, **kwargs) -> Self:
    region = Region.from_pyvista(obj, grad=True)
    return cls.from_region(region, **kwargs)

from_region classmethod ¤

from_region(
    region: Region,
    *,
    clamp_hess_diag: bool = True,
    clamp_hess_quad: bool = True,
    clamp_lambda: bool = True,
    requires_grad: Sequence[str] = (),
    **kwargs,
) -> Self
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def from_region(
    cls, region: Region, *, requires_grad: Sequence[str] = (), **kwargs
) -> Self:
    self: Self = cls(
        cells=wpu.to_warp(region.cells, wpt.vec4i),
        dhdX=wpu.to_warp(region.dhdX, wpt.mat43),
        dV=wpu.to_warp(region.dV, wpt.float_),
        params=cls.make_params(region, requires_grad),
        requires_grad=requires_grad,
        **kwargs,
    )
    return self

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
110
111
112
113
114
115
116
117
@override
def fun(self, u: Vector, output: Scalar) -> None:
    wp.launch(
        self.fun_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

get_cell_params staticmethod ¤

get_cell_params(params: Params, cid: int) -> ParamsElem
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py 
40
41
42
43
44
45
46
47
48
49
50
@override
@staticmethod
@no_type_check
@wp.func
def get_cell_params(params: Params, cid: int) -> ParamsElem:
    return Phace.ParamsElem(
        activation=params.activation[cid],
        lambda_=params.lambda_[cid],
        mu=params.mu[cid],
        muscle_fraction=params.muscle_fraction[cid],
    )

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
119
120
121
122
123
124
125
126
@override
def grad(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
179
180
181
182
183
184
185
186
@override
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    wp.launch(
        self.grad_and_hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[grad, hess_diag],
    )

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
128
129
130
131
132
133
134
135
@override
def hess_diag(self, u: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_diag_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
137
138
139
140
141
142
143
144
@override
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    wp.launch(
        self.hess_prod_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
146
147
148
149
150
151
152
153
@override
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    wp.launch(
        self.hess_quad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, p, self.cells, self.dhdX, self.dV, self.params],
        outputs=[output],
    )

make_params classmethod ¤

make_params(
    region: Region, requires_grad: Sequence[str] = ()
) -> Params
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
90
91
92
93
94
@classmethod
@no_type_check
def make_params(cls, region: Region, requires_grad: Sequence[str] = ()) -> Params:
    fields: Mapping[str, wp.array] = cls._params_fields_from_region(region)
    return cls._params_from_fields(fields, requires_grad)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: array, p: array
) -> dict[str, array]
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@override
def mixed_derivative_prod(self, u: wp.array, p: wp.array) -> dict[str, wp.array]:
    if not self.requires_grad:
        return {}
    for name in self.requires_grad:
        getattr(self.params, name).grad.zero_()
    output: wp.array = wp.zeros_like(u)
    with wp.Tape() as tape:
        self.grad(u, output)
    tape.backward(grads={output: p})
    outputs: dict[str, wp.array] = {
        name: getattr(self.params, name).grad for name in self.requires_grad
    }
    return outputs

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
104
105
106
107
108
@override
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self.params, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py 
170
171
172
173
174
175
176
177
@override
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    wp.launch(
        self.value_and_grad_kernel,
        dim=(self.n_cells, self.n_quadrature_points),
        inputs=[u, self.cells, self.dhdX, self.dV, self.params],
        outputs=[value, grad],
    )

WarpEnergy ¤

Bases: IdMixin


              flowchart TD
              liblaf.apple.WarpEnergy[WarpEnergy]
              liblaf.apple.utils._id_mixin.IdMixin[IdMixin]

                              liblaf.apple.utils._id_mixin.IdMixin --> liblaf.apple.WarpEnergy
                


              click liblaf.apple.WarpEnergy href "" "liblaf.apple.WarpEnergy"
              click liblaf.apple.utils._id_mixin.IdMixin href "" "liblaf.apple.utils._id_mixin.IdMixin"

Parameters:

Subclassed by: Used by:

Methods:

Attributes:

id class-attribute instance-attribute ¤

id: str = field(
    default=Factory(_default_id, takes_self=True),
    kw_only=True,
)

requires_grad class-attribute instance-attribute ¤

requires_grad: Sequence[str] = field(
    default=(), kw_only=True
)

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
27
28
def fun(self, u: Vector, output: Scalar) -> None:
    raise NotImplementedError

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
30
31
def grad(self, u: Vector, output: Vector) -> None:
    raise NotImplementedError

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
49
50
51
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    self.grad(u, grad)
    self.hess_diag(u, hess_diag)

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
33
34
def hess_diag(self, u: Vector, output: Vector) -> None:
    raise NotImplementedError

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
36
37
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    raise NotImplementedError

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
39
40
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    raise NotImplementedError

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> dict[str, array]
Source code in src/liblaf/apple/warp/model/_energy.py 
42
43
def mixed_derivative_prod(self, u: Vector, p: Vector) -> dict[str, wp.array]:
    raise NotImplementedError

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
19
20
def update(self, u: Vector) -> None:
    pass

update_params ¤

update_params(params: EnergyParams) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
22
23
24
25
def update_params(self, params: EnergyParams) -> None:
    for name, value in params.items():
        param: wp.array = getattr(self, name)
        wp.copy(param, wpu.to_warp(value, param.dtype))

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/model/_energy.py 
45
46
47
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    self.fun(u, value)
    self.grad(u, grad)

WarpModel ¤

Parameters:

  • energies ¤

    (dict[str, WarpEnergy], default: <class 'dict'> ) –

    dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object’s (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

Methods:

Attributes:

energies class-attribute instance-attribute ¤

energies: dict[str, WarpEnergy] = field(factory=dict)

fun ¤

fun(u: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
27
28
29
def fun(self, u: Vector, output: Scalar) -> None:
    for energy in self.energies.values():
        energy.fun(u, output)

grad ¤

grad(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
31
32
33
def grad(self, u: Vector, output: Vector) -> None:
    for energy in self.energies.values():
        energy.grad(u, output)

grad_and_hess_diag ¤

grad_and_hess_diag(
    u: Vector, grad: Vector, hess_diag: Vector
) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
60
61
62
def grad_and_hess_diag(self, u: Vector, grad: Vector, hess_diag: Vector) -> None:
    for energy in self.energies.values():
        energy.grad_and_hess_diag(u, grad, hess_diag)

hess_diag ¤

hess_diag(u: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
35
36
37
def hess_diag(self, u: Vector, output: Vector) -> None:
    for energy in self.energies.values():
        energy.hess_diag(u, output)

hess_prod ¤

hess_prod(u: Vector, p: Vector, output: Vector) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
39
40
41
def hess_prod(self, u: Vector, p: Vector, output: Vector) -> None:
    for energy in self.energies.values():
        energy.hess_prod(u, p, output)

hess_quad ¤

hess_quad(u: Vector, p: Vector, output: Scalar) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
43
44
45
def hess_quad(self, u: Vector, p: Vector, output: Scalar) -> None:
    for energy in self.energies.values():
        energy.hess_quad(u, p, output)

mixed_derivative_prod ¤

mixed_derivative_prod(
    u: Vector, p: Vector
) -> dict[str, dict[str, array]]
Source code in src/liblaf/apple/warp/model/_model.py 
47
48
49
50
51
52
53
54
def mixed_derivative_prod(
    self, u: Vector, p: Vector
) -> dict[str, dict[str, wp.array]]:
    output: dict[str, dict[str, wp.array]] = {
        name: energy.mixed_derivative_prod(u, p)
        for name, energy in self.energies.items()
    }
    return output

update ¤

update(u: Vector) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
19
20
21
def update(self, u: Vector) -> None:
    for energy in self.energies.values():
        energy.update(u)

update_params ¤

update_params(params: ModelParams) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
23
24
25
def update_params(self, params: ModelParams) -> None:
    for name, energy_params in params.items():
        self.energies[name].update_params(energy_params)

value_and_grad ¤

value_and_grad(
    u: Vector, value: Scalar, grad: Vector
) -> None
Source code in src/liblaf/apple/warp/model/_model.py 
56
57
58
def value_and_grad(self, u: Vector, value: Scalar, grad: Vector) -> None:
    for energy in self.energies.values():
        energy.value_and_grad(u, value, grad)

WarpModelBuilder ¤

Parameters:

  • energies ¤

    (dict[str, WarpEnergy], default: <class 'dict'> ) –

    dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object’s (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

Returned by: Used by:

Methods:

Attributes:

energies class-attribute instance-attribute ¤

energies: dict[str, WarpEnergy] = field(factory=dict)

add_energy ¤

add_energy(energy: WarpEnergy) -> None
Source code in src/liblaf/apple/warp/model/_builder.py 
11
12
def add_energy(self, energy: WarpEnergy) -> None:
    self.energies[energy.id] = energy

finalize ¤

finalize() -> WarpModel
Source code in src/liblaf/apple/warp/model/_builder.py 
14
15
def finalize(self) -> WarpModel:
    return WarpModel(energies=self.energies)