"""Frictionless Incremental Potential Contact for small CPU simulations.
The collision geometry, barrier derivatives, and continuous collision detection
come from the official IPC Toolkit. PyTorch supplies the tetrahedral elasticity
and the implicit adjoint of the converged time step.
"""
from __future__ import annotations
import numpy as np
import torch
from .diff_physics import ImplicitDifferentiation, SpatiallyVaryingMaterial
from .material import StableNeoHookean
def _load_ipctk():
try:
import ipctk
except ImportError as error:
raise ImportError(
"IPC contact requires the optional dependency. "
"Install it with `pip install torch-diffsim[ipc]`."
) from error
version = tuple(int(part) for part in ipctk.__version__.split(".")[:2])
if version != (1, 6):
raise ImportError(
"torch-diffsim requires ipctk>=1.6,<1.7; "
f"found ipctk {ipctk.__version__}."
)
return ipctk
def _boundary_faces(tetrahedra):
"""Return all boundary triangles without sampling or topology simplification."""
faces_by_key = {}
for tet in tetrahedra.detach().cpu().tolist():
for local_face in ((0, 2, 1), (0, 1, 3), (0, 3, 2), (1, 2, 3)):
face = tuple(tet[index] for index in local_face)
faces_by_key.setdefault(tuple(sorted(face)), []).append(face)
if any(len(faces) > 2 for faces in faces_by_key.values()):
raise ValueError("IPC requires a manifold tetrahedral boundary")
boundary = [faces[0] for faces in faces_by_key.values() if len(faces) == 1]
if not boundary:
raise ValueError("The tetrahedral mesh has no boundary faces")
return np.asarray(boundary, dtype=np.int64)
class _IPCContact:
"""Thin, CPU-double wrapper around the official IPC Toolkit."""
def __init__(self, mesh, dhat, stiffness, ground_height):
self.ipctk = _load_ipctk()
self.dhat = float(dhat)
self.ground_height = ground_height
rest_positions = np.ascontiguousarray(
mesh.vertices.detach().cpu().numpy(), dtype=np.float64
)
faces = _boundary_faces(mesh.tetrahedra)
edges = self.ipctk.edges(faces)
self.mesh = self.ipctk.CollisionMesh.build_from_full_mesh(
rest_positions, edges, faces
)
if ground_height is not None:
self.mesh.planes = [
self.ipctk.Hyperplane(
np.asarray([0.0, 1.0, 0.0]),
np.asarray([0.0, float(ground_height), 0.0]),
)
]
self.potential = self.ipctk.BarrierPotential(self.dhat, float(stiffness))
self._check_derivative_layout()
def _check_derivative_layout(self):
"""Fail clearly for an ipctk build with a nonstandard DOF ordering."""
probe = np.arange(1, self.mesh.ndof + 1, dtype=np.float64)
mapped = np.asarray(self.mesh.to_full_dof(probe)).reshape(-1)
expected = np.zeros(self.mesh.full_ndof, dtype=np.float64)
for surface_id in range(self.mesh.num_vertices):
full_id = self.mesh.to_full_vertex_id(surface_id)
expected[3 * full_id : 3 * full_id + 3] = probe[
3 * surface_id : 3 * surface_id + 3
]
if not np.array_equal(mapped, expected):
raise RuntimeError(
"This ipctk build does not use the supported vertex-major "
"derivative layout"
)
def _surface_positions(self, full_positions):
full = np.ascontiguousarray(
full_positions.detach().cpu().numpy(), dtype=np.float64
)
return np.ascontiguousarray(self.mesh.vertices(full), dtype=np.float64)
def _collisions(self, surface_positions):
candidates = self.ipctk.Candidates()
candidates.build(self.mesh, surface_positions, self.dhat)
collisions = self.ipctk.NormalCollisions()
collisions.build(
candidates,
self.mesh,
surface_positions,
self.dhat,
)
return collisions
def value(self, full_positions):
surface = self._surface_positions(full_positions)
collisions = self._collisions(surface)
return float(self.potential(collisions, self.mesh, surface))
def gradient_hessian(self, full_positions):
surface = self._surface_positions(full_positions)
collisions = self._collisions(surface)
gradient = self.potential.gradient(collisions, self.mesh, surface)
hessian = self.potential.hessian(
collisions,
self.mesh,
surface,
project_hessian_to_psd=self.ipctk.PSDProjectionMethod.NONE,
)
full_gradient = np.asarray(self.mesh.to_full_dof(gradient)).reshape(-1)
full_hessian = self.mesh.to_full_dof(hessian).toarray()
return full_gradient, full_hessian
def collision_free_stepsize(self, full_start, full_end):
start = self._surface_positions(full_start)
end = self._surface_positions(full_end)
return float(self.ipctk.compute_collision_free_stepsize(self.mesh, start, end))
def is_step_collision_free(self, full_start, full_end):
start = self._surface_positions(full_start)
end = self._surface_positions(full_end)
return bool(self.ipctk.is_step_collision_free(self.mesh, start, end))
def validate_start(self, full_positions):
surface = self._surface_positions(full_positions)
if self.ipctk.has_intersections(self.mesh, surface):
raise ValueError("IPC requires an initially intersection-free surface")
if self.ground_height is not None:
signed_distance = surface[:, 1] - float(self.ground_height)
if np.any(signed_distance <= 0.0):
raise ValueError(
"IPC requires every surface vertex to start above the ground plane"
)
if not np.isfinite(self.value(full_positions)):
raise ValueError("IPC requires a finite, strictly feasible initial state")
class _IPCBarrierGradient(torch.autograd.Function):
"""Expose the toolkit gradient with its exact unprojected Hessian as VJP."""
@staticmethod
def forward(ctx, positions, contact):
gradient, hessian = contact.gradient_hessian(positions)
hessian_tensor = torch.as_tensor(
hessian, dtype=positions.dtype, device=positions.device
)
ctx.save_for_backward(hessian_tensor)
return torch.as_tensor(
gradient, dtype=positions.dtype, device=positions.device
).reshape_as(positions)
@staticmethod
def backward(ctx, grad_output):
(hessian,) = ctx.saved_tensors
gradient = hessian.transpose(0, 1).matmul(grad_output.reshape(-1))
return gradient.reshape_as(grad_output), None
def _determinant_and_cofactor(F):
a, b, c = F.unbind(dim=2)
cofactor = torch.stack(
(
torch.cross(b, c, dim=1),
torch.cross(c, a, dim=1),
torch.cross(a, b, dim=1),
),
dim=2,
)
return torch.sum(a * cofactor[:, :, 0], dim=1), cofactor
def _tetrahedron_jacobian_coefficients(mesh, start, end):
"""Return cubic coefficients proportional to det(F(alpha)) for each tet."""
tetrahedra = mesh.tetrahedra.detach().cpu().numpy()
start_np = np.asarray(start.detach().cpu(), dtype=np.float64)
end_np = np.asarray(end.detach().cpu(), dtype=np.float64)
start_tets = start_np[tetrahedra]
step_tets = (end_np - start_np)[tetrahedra]
a = start_tets[:, 1] - start_tets[:, 0]
b = start_tets[:, 2] - start_tets[:, 0]
c = start_tets[:, 3] - start_tets[:, 0]
da = step_tets[:, 1] - step_tets[:, 0]
db = step_tets[:, 2] - step_tets[:, 0]
dc = step_tets[:, 3] - step_tets[:, 0]
def determinant(column_0, column_1, column_2):
return np.einsum("ij,ij->i", column_0, np.cross(column_1, column_2))
coefficients = np.stack(
(
determinant(a, b, c),
determinant(da, b, c) + determinant(a, db, c) + determinant(a, b, dc),
determinant(da, db, c) + determinant(da, b, dc) + determinant(a, db, dc),
determinant(da, db, dc),
),
axis=1,
)
rest_orientation = np.sign(
np.linalg.det(np.asarray(mesh.Dm.detach().cpu(), dtype=np.float64))
)
return coefficients * rest_orientation[:, None]
def _is_step_inversion_free(mesh, start, end):
"""Check positivity over the entire linear segment, not only its endpoint."""
coefficients = _tetrahedron_jacobian_coefficients(mesh, start, end)
for c0, c1, c2, c3 in coefficients:
values = [c0, c0 + c1 + c2 + c3]
derivative = np.trim_zeros(np.asarray([3.0 * c3, 2.0 * c2, c1]), "f")
if derivative.size > 1:
for root in np.roots(derivative):
real_tolerance = 1e-7 * max(1.0, abs(root.real))
if abs(root.imag) <= real_tolerance and 0.0 < root.real < 1.0:
alpha = root.real
values.append(c0 + alpha * (c1 + alpha * (c2 + alpha * c3)))
if min(values) <= 0.0:
return False
return True
def _inversion_free_stepsize(mesh, start, end, safety=0.8):
"""Find a conservative fraction before the first tetrahedron degeneracy."""
coefficients = _tetrahedron_jacobian_coefficients(mesh, start, end)
if np.any(coefficients[:, 0] <= 0.0):
raise ValueError("IPC requires initially positive tetrahedron Jacobians")
earliest_root = np.inf
for coefficients_for_tet in coefficients:
scale = np.max(np.abs(coefficients_for_tet))
if not np.isfinite(scale) or scale == 0.0:
raise RuntimeError("Could not compute a finite tetrahedron step bound")
polynomial = np.trim_zeros((coefficients_for_tet / scale)[::-1], "f")
if polynomial.size <= 1:
continue
for root in np.roots(polynomial):
real_tolerance = 1e-7 * max(1.0, abs(root.real))
if abs(root.imag) <= real_tolerance and 0.0 < root.real <= 1.0:
earliest_root = min(earliest_root, float(root.real))
if np.isfinite(earliest_root):
return safety * earliest_root
return 1.0
def _validate_positive_tetrahedra(mesh, positions):
jacobians = torch.linalg.det(mesh.compute_deformation_gradient(positions))
if torch.any(jacobians <= 0.0):
raise ValueError("IPC requires initially positive tetrahedron Jacobians")
def _stable_energy_density(F, youngs_modulus, poissons_ratio):
mu = youngs_modulus / (2.0 * (1.0 + poissons_ratio))
lam = (
youngs_modulus
* poissons_ratio
/ ((1.0 + poissons_ratio) * (1.0 - 2.0 * poissons_ratio))
)
stable_mu = 4.0 * mu / 3.0
stable_lam = lam + 5.0 * mu / 6.0
Ic = torch.sum(F * F, dim=(1, 2))
J, _ = _determinant_and_cofactor(F)
Ic_offset = Ic - 3.0
return (
stable_mu / 2.0 * (Ic_offset - torch.log1p(Ic_offset / 4.0))
+ stable_lam / 2.0 * (J - 1.0) ** 2
- 3.0 * stable_mu / 4.0 * (J - 1.0)
)
def _elastic_forces(mesh, positions, youngs_modulus, poissons_ratio):
F = mesh.compute_deformation_gradient(positions)
mu = youngs_modulus / (2.0 * (1.0 + poissons_ratio))
lam = (
youngs_modulus
* poissons_ratio
/ ((1.0 + poissons_ratio) * (1.0 - 2.0 * poissons_ratio))
)
stable_mu = 4.0 * mu / 3.0
stable_lam = lam + 5.0 * mu / 6.0
Ic = torch.sum(F * F, dim=(1, 2))
J, cofactor = _determinant_and_cofactor(F)
stress = (stable_mu * (1.0 - 1.0 / (Ic + 1.0)))[..., None, None] * F + (
stable_lam * (J - 1.0) - 3.0 * stable_mu / 4.0
)[..., None, None] * cofactor
force_matrix = -mesh.rest_volume[:, None, None] * torch.bmm(
stress, mesh.Dm_inv.transpose(1, 2)
)
element_forces = torch.stack(
(
-force_matrix.sum(dim=2),
force_matrix[:, :, 0],
force_matrix[:, :, 1],
force_matrix[:, :, 2],
),
dim=1,
)
forces = torch.zeros_like(positions)
for local_vertex in range(4):
forces = forces.index_add(
0,
mesh.tetrahedra[:, local_vertex],
element_forces[:, local_vertex],
)
return forces
[docs]
class IPCImplicitEulerSolver:
"""Dense, differentiable, frictionless IPC for small tetrahedral simulations.
Each substep minimizes the backward-Euler incremental potential containing
inertia, Stable Neo-Hookean elasticity, and the official IPC barrier in one
nonlinear solve. Every accepted Newton segment is limited by conservative CCD.
This solver is intentionally CPU-only and requires ``torch.float64``. It protects
accepted trajectories from surface intersections and tetrahedral inversion (and
optionally includes an analytic ground plane), but it does not implement friction.
"""
def __init__(
self,
dt=0.01,
gravity=-9.8,
damping=1.0,
substeps=1,
dhat=0.01,
barrier_stiffness=1e5,
ground_height=None,
newton_tolerance=1e-8,
max_newton_iterations=50,
max_line_search_iterations=25,
armijo=1e-4,
hessian_regularization=1e-9,
):
scalar_settings = {
"dt": dt,
"gravity": gravity,
"damping": damping,
"dhat": dhat,
"barrier_stiffness": barrier_stiffness,
"newton_tolerance": newton_tolerance,
"armijo": armijo,
"hessian_regularization": hessian_regularization,
}
if ground_height is not None:
scalar_settings["ground_height"] = ground_height
if not all(np.isfinite(value) for value in scalar_settings.values()):
raise ValueError("IPC solver settings must be finite")
if dt <= 0:
raise ValueError("dt must be positive")
if substeps < 1:
raise ValueError("substeps must be at least one")
if dhat <= 0 or barrier_stiffness <= 0:
raise ValueError("dhat and barrier_stiffness must be positive")
if newton_tolerance <= 0:
raise ValueError("newton_tolerance must be positive")
if max_newton_iterations < 0 or max_line_search_iterations < 1:
raise ValueError("Newton and line-search iteration limits are invalid")
if not 0.0 < armijo < 1.0:
raise ValueError("armijo must be between zero and one")
if hessian_regularization <= 0:
raise ValueError("hessian_regularization must be positive")
self.dt = float(dt)
self.gravity_value = float(gravity)
self.damping = float(damping)
self.substeps = int(substeps)
self.dhat = float(dhat)
self.barrier_stiffness = float(barrier_stiffness)
self.ground_height = ground_height
self.newton_tolerance = float(newton_tolerance)
self.max_newton_iterations = int(max_newton_iterations)
self.max_line_search_iterations = int(max_line_search_iterations)
self.armijo = float(armijo)
self.hessian_regularization = float(hessian_regularization)
self._contact_key = None
self._contact_backend = None
self.last_diagnostics = None
def _validate_inputs(self, mesh, material, positions, velocities, masses):
if positions.device.type != "cpu" or positions.dtype != torch.float64:
raise TypeError("IPC requires CPU torch.float64 positions")
if velocities.device.type != "cpu" or velocities.dtype != torch.float64:
raise TypeError("IPC requires CPU torch.float64 velocities")
if velocities.shape != positions.shape or positions.ndim != 2:
raise ValueError("positions and velocities must both have shape (N, 3)")
if positions.shape[1] != 3 or masses.shape != positions.shape[:1]:
raise ValueError("positions must have shape (N, 3) and masses shape (N,)")
if positions.shape[0] != mesh.num_vertices:
raise ValueError("IPC state size must match mesh.num_vertices")
if any(
tensor.device.type != "cpu"
for tensor in (
mesh.vertices,
mesh.tetrahedra,
mesh.Dm,
mesh.Dm_inv,
mesh.rest_volume,
masses,
)
):
raise TypeError("IPC requires the mesh and masses on the CPU")
if mesh.tetrahedra.dtype != torch.long:
raise TypeError("IPC requires torch.long tetrahedral indices")
if any(
tensor.dtype != torch.float64
for tensor in (
mesh.vertices,
mesh.Dm,
mesh.Dm_inv,
mesh.rest_volume,
)
):
raise TypeError("IPC requires a torch.float64 mesh")
if masses.dtype != torch.float64:
raise TypeError("IPC requires torch.float64 masses")
if torch.any(masses <= 0):
raise ValueError("All vertex masses must be positive")
if not all(
torch.isfinite(tensor).all() for tensor in (positions, velocities, masses)
):
raise ValueError("IPC state and masses must be finite")
if mesh.Dm_inv.requires_grad or mesh.rest_volume.requires_grad:
raise ValueError("IPC currently treats the rest mesh as fixed")
if not isinstance(material, (StableNeoHookean, SpatiallyVaryingMaterial)):
from .diff_physics import DifferentiableMaterial
if not isinstance(material, DifferentiableMaterial):
raise TypeError(
"IPC supports StableNeoHookean, DifferentiableMaterial, "
"and SpatiallyVaryingMaterial"
)
def _contact(self, mesh):
key = (
id(mesh),
self.dhat,
self.barrier_stiffness,
self.ground_height,
)
if key != self._contact_key:
self._contact_backend = _IPCContact(
mesh,
self.dhat,
self.barrier_stiffness,
self.ground_height,
)
self._contact_key = key
return self._contact_backend
@staticmethod
def _material_parameters(material, reference):
if isinstance(material, SpatiallyVaryingMaterial):
return "log", material.log_E, material.nu
youngs = material.E
poisson = material.nu
if not isinstance(youngs, torch.Tensor):
youngs = reference.new_tensor(youngs)
if not isinstance(poisson, torch.Tensor):
poisson = reference.new_tensor(poisson)
return "direct", youngs, poisson
@staticmethod
def _material_values(parameterization, youngs_parameter, poisson):
youngs = (
torch.exp(youngs_parameter)
if parameterization == "log"
else youngs_parameter
)
return youngs, poisson
def _stationarity(
self,
mesh,
contact,
positions,
predictor,
masses,
youngs,
poisson,
free_mask,
h,
):
elastic_gradient = -_elastic_forces(mesh, positions, youngs, poisson)
contact_gradient = _IPCBarrierGradient.apply(positions, contact)
physical = masses[:, None] * (positions - predictor) + h * h * (
elastic_gradient + contact_gradient
)
return torch.where(free_mask, physical, positions - predictor)
def _objective(
self,
mesh,
contact,
positions,
predictor,
masses,
youngs,
poisson,
h,
):
inertia = 0.5 * torch.sum(masses[:, None] * (positions - predictor) ** 2)
deformation = mesh.compute_deformation_gradient(positions)
elastic = torch.sum(
_stable_energy_density(deformation, youngs, poisson) * mesh.rest_volume
)
barrier = positions.new_tensor(contact.value(positions))
return inertia + h * h * (elastic + barrier)
def _search_direction(self, hessian, gradient):
symmetric = 0.5 * (hessian + hessian.transpose(0, 1))
eigenvalues, eigenvectors = torch.linalg.eigh(symmetric)
scale = torch.clamp(eigenvalues.abs().max(), min=1.0)
floor = self.hessian_regularization * scale
search_hessian = (eigenvectors * torch.clamp(eigenvalues, min=floor)) @ (
eigenvectors.transpose(0, 1)
)
return torch.linalg.solve(search_hessian, -gradient)
def _solve_substep(
self,
mesh,
contact,
start,
predictor,
masses,
youngs,
poisson,
free_mask,
h,
):
contact.validate_start(start)
_validate_positive_tetrahedra(mesh, start)
positions = start.detach().clone()
free_dofs = free_mask.expand_as(positions).reshape(-1)
accepted_steps = 0
ccd_limited_steps = 0
inversion_limited_steps = 0
for iteration in range(self.max_newton_iterations + 1):
positions_var = positions.detach().requires_grad_(True)
def residual_flat(candidate):
return self._stationarity(
mesh,
contact,
candidate,
predictor.detach(),
masses.detach(),
youngs.detach(),
poisson.detach(),
free_mask,
h,
).reshape(-1)
residual = residual_flat(positions_var)
free_residual = residual[free_dofs]
residual_norm = (
float(free_residual.detach().abs().max())
if free_residual.numel()
else 0.0
)
if residual_norm <= self.newton_tolerance:
self.last_diagnostics = {
"iterations": iteration,
"stationarity_norm": residual_norm,
"accepted_steps": accepted_steps,
"ccd_limited_steps": ccd_limited_steps,
"inversion_limited_steps": inversion_limited_steps,
"objective": float(
self._objective(
mesh,
contact,
positions,
predictor,
masses,
youngs,
poisson,
h,
)
),
}
return positions
if iteration == self.max_newton_iterations:
break
jacobian = torch.autograd.functional.jacobian(
residual_flat, positions_var
).reshape(positions.numel(), positions.numel())
reduced_hessian = jacobian[free_dofs][:, free_dofs]
reduced_direction = self._search_direction(
reduced_hessian, free_residual.detach()
)
direction = torch.zeros_like(positions).reshape(-1)
direction[free_dofs] = reduced_direction
direction = direction.reshape_as(positions)
directional_derivative = float(
torch.dot(free_residual.detach(), reduced_direction)
)
if directional_derivative >= 0:
direction = torch.zeros_like(positions).reshape(-1)
direction[free_dofs] = -free_residual.detach()
direction = direction.reshape_as(positions)
directional_derivative = -float(
torch.dot(free_residual.detach(), free_residual.detach())
)
ccd_stepsize = contact.collision_free_stepsize(
positions, positions + direction
)
if ccd_stepsize < 1.0:
ccd_limited_steps += 1
if ccd_stepsize <= 0.0:
raise RuntimeError("IPC CCD found no positive feasible Newton step")
inversion_stepsize = _inversion_free_stepsize(
mesh, positions, positions + direction
)
if inversion_stepsize < 1.0:
inversion_limited_steps += 1
if inversion_stepsize <= 0.0:
raise RuntimeError("IPC found no positive inversion-free Newton step")
stepsize = min(1.0, ccd_stepsize, inversion_stepsize)
current_objective = float(
self._objective(
mesh,
contact,
positions,
predictor,
masses,
youngs,
poisson,
h,
)
)
accepted = False
for _ in range(self.max_line_search_iterations):
trial = positions + stepsize * direction
if contact.is_step_collision_free(
positions, trial
) and _is_step_inversion_free(mesh, positions, trial):
trial_objective = float(
self._objective(
mesh,
contact,
trial,
predictor,
masses,
youngs,
poisson,
h,
)
)
armijo_bound = (
current_objective
+ self.armijo * stepsize * directional_derivative
)
if np.isfinite(trial_objective) and trial_objective <= armijo_bound:
positions = trial.detach()
accepted = True
accepted_steps += 1
break
stepsize *= 0.5
if not accepted:
raise RuntimeError("IPC Newton line search failed")
raise RuntimeError(
"IPC Newton solve did not converge: "
f"stationarity norm {residual_norm:.3e} exceeds "
f"{self.newton_tolerance:.3e}"
)
[docs]
def step(
self,
mesh,
material,
positions,
velocities,
masses,
fixed_vertices=None,
external_forces=None,
):
"""Advance one differentiable IPC time step.
``external_forces`` may come from a neural network; gradients propagate
through the converged implicit solution by the implicit function theorem.
"""
self._validate_inputs(mesh, material, positions, velocities, masses)
if external_forces is None:
external_forces = torch.zeros_like(positions)
if external_forces.shape != positions.shape:
raise ValueError("external_forces must have shape (N, 3)")
if (
external_forces.device.type != "cpu"
or not torch.isfinite(external_forces).all()
):
raise ValueError("external_forces must be finite and on the CPU")
free_mask = torch.ones(
(positions.shape[0], 1), dtype=torch.bool, device=positions.device
)
if fixed_vertices is not None:
fixed_vertices = torch.as_tensor(
fixed_vertices, dtype=torch.long, device=positions.device
)
if fixed_vertices.ndim != 1 or torch.any(
(fixed_vertices < 0) | (fixed_vertices >= positions.shape[0])
):
raise ValueError("fixed_vertices must contain valid vertex indices")
free_mask[fixed_vertices] = False
parameterization, youngs_parameter, poisson = self._material_parameters(
material, positions
)
youngs, poisson_value = self._material_values(
parameterization, youngs_parameter, poisson
)
if youngs.device.type != "cpu" or poisson_value.device.type != "cpu":
raise TypeError("IPC requires material parameters on the CPU")
if (
not torch.isfinite(youngs.detach()).all()
or not torch.isfinite(poisson_value.detach()).all()
):
raise ValueError("Material parameters must be finite")
if torch.any(youngs.detach() < 0) or torch.any(
(poisson_value.detach() <= -1.0) | (poisson_value.detach() >= 0.5)
):
raise ValueError("Material parameters require E >= 0 and -1 < nu < 0.5")
if youngs.numel() not in (
1,
mesh.num_elements,
) or poisson_value.numel() not in (
1,
mesh.num_elements,
):
raise ValueError("Material parameters must be scalar or per-element")
contact = self._contact(mesh)
h = self.dt / self.substeps
gravity = positions.new_tensor([0.0, self.gravity_value, 0.0])
current_positions = positions
current_velocities = velocities
for _ in range(self.substeps):
predictor = (
current_positions
+ h * self.damping * current_velocities
+ h * h * (gravity[None, :] + external_forces / masses[:, None])
)
predictor = torch.where(free_mask, predictor, current_positions)
youngs, poisson_value = self._material_values(
parameterization, youngs_parameter, poisson
)
solution = self._solve_substep(
mesh,
contact,
current_positions.detach(),
predictor.detach(),
masses.detach(),
youngs.detach(),
poisson_value.detach(),
free_mask,
h,
)
def residual_fn(
candidate,
differentiable_predictor,
differentiable_masses,
differentiable_youngs_parameter,
differentiable_poisson,
):
differentiable_youngs, differentiable_poisson_value = (
self._material_values(
parameterization,
differentiable_youngs_parameter,
differentiable_poisson,
)
)
return self._stationarity(
mesh,
contact,
candidate,
differentiable_predictor,
differentiable_masses,
differentiable_youngs,
differentiable_poisson_value,
free_mask,
h,
)
current_positions_next = ImplicitDifferentiation.implicit_backward(
residual_fn,
solution,
(predictor, masses, youngs_parameter, poisson),
atol=self.newton_tolerance,
)
current_velocities = (current_positions_next - current_positions) / h
current_positions = current_positions_next
return current_positions, current_velocities