init
This commit is contained in:
@@ -0,0 +1,363 @@
|
||||
"""
|
||||
inertia.py
|
||||
-------------
|
||||
|
||||
Functions for dealing with inertia tensors.
|
||||
|
||||
Results validated against known geometries and checked for
|
||||
internal consistency.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .typed import ArrayLike, NDArray, Number, Optional, Union, float64
|
||||
from .util import multi_dot
|
||||
|
||||
|
||||
def cylinder_inertia(
|
||||
mass: Number, radius: Number, height: Number, transform: Optional[ArrayLike] = None
|
||||
) -> NDArray[float64]:
|
||||
"""
|
||||
Return the inertia tensor of a cylinder.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
mass : float
|
||||
Mass of cylinder
|
||||
radius : float
|
||||
Radius of cylinder
|
||||
height : float
|
||||
Height of cylinder
|
||||
transform : (4, 4) float
|
||||
Transformation of cylinder
|
||||
|
||||
Returns
|
||||
------------
|
||||
inertia : (3, 3) float
|
||||
Inertia tensor
|
||||
"""
|
||||
h2, r2 = height**2, radius**2
|
||||
diagonal = np.array(
|
||||
[
|
||||
((mass * h2) / 12) + ((mass * r2) / 4),
|
||||
((mass * h2) / 12) + ((mass * r2) / 4),
|
||||
(mass * r2) / 2,
|
||||
]
|
||||
)
|
||||
inertia = diagonal * np.eye(3)
|
||||
|
||||
if transform is not None:
|
||||
inertia = transform_inertia(transform, inertia)
|
||||
|
||||
return inertia
|
||||
|
||||
|
||||
def sphere_inertia(mass: Number, radius: Number) -> NDArray[float64]:
|
||||
"""
|
||||
Return the inertia tensor of a sphere.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
mass : float
|
||||
Mass of sphere
|
||||
radius : float
|
||||
Radius of sphere
|
||||
|
||||
Returns
|
||||
------------
|
||||
inertia : (3, 3) float
|
||||
Inertia tensor
|
||||
"""
|
||||
return (2.0 / 5.0) * (radius**2) * mass * np.eye(3)
|
||||
|
||||
|
||||
def points_inertia(
|
||||
points: ArrayLike,
|
||||
weights: Union[None, ArrayLike, Number] = None,
|
||||
at_center_mass: bool = True,
|
||||
) -> NDArray[float64]:
|
||||
"""
|
||||
Calculate an inertia tensor for an array of point masses
|
||||
at the center of mass.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points : (n, 3)
|
||||
Points in space.
|
||||
weights : (n,) or number
|
||||
Per-point weight to use.
|
||||
at_center_mass
|
||||
Calculate at the center of mass of the points, or if False
|
||||
at the original origin.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
tensor : (3, 3)
|
||||
Inertia tensor for point masses.
|
||||
"""
|
||||
if weights is None:
|
||||
# by default make the total weight 1.0 to match
|
||||
# the default mass in other functions, and so that
|
||||
# if a user didn't specify anything it doesn't blow
|
||||
# up the scale depending on the number of points
|
||||
weights = np.full(len(points), 1.0 / float(len(points)), dtype=np.float64)
|
||||
elif isinstance(weights, (float, np.integer, int)):
|
||||
# "is it a number" check
|
||||
weights = np.full(len(points), float(weights), dtype=np.float64)
|
||||
else:
|
||||
weights = np.array(weights)
|
||||
if len(weights) != len(points):
|
||||
raise ValueError(
|
||||
f"Weights must correspond to points! {len(weights)} != {len(points)}"
|
||||
)
|
||||
|
||||
# make sure the points are an array of correct shape
|
||||
points = np.asanyarray(points, dtype=np.float64)
|
||||
if len(points.shape) != 2 or points.shape[1] != 3:
|
||||
raise ValueError(f"Points must be `(n, 3)` not {points.shape}")
|
||||
|
||||
if at_center_mass:
|
||||
# get the center of mass of the points
|
||||
center_mass = np.average(points, weights=weights, axis=0)
|
||||
# get the points with the origin at their center of mass
|
||||
points_com = points - center_mass
|
||||
else:
|
||||
# calculate at original origin
|
||||
points_com = points
|
||||
|
||||
# expand into shorthand for the expressions
|
||||
x, y, z = points_com.T
|
||||
x2, y2, z2 = (points_com**2).T
|
||||
|
||||
# calculate tensors per-point in a flattened (9, n) array
|
||||
# from physics.stackexchange.com/questions/614094
|
||||
tensors = np.array(
|
||||
[y2 + z2, -x * y, -x * z, -x * y, x2 + z2, -y * z, -x * z, -y * z, x2 + y2],
|
||||
dtype=np.float64,
|
||||
)
|
||||
|
||||
# combine the weighted tensors and reshape
|
||||
tensor = (tensors * weights).sum(axis=1).reshape((3, 3))
|
||||
|
||||
return tensor
|
||||
|
||||
|
||||
def principal_axis(inertia: ArrayLike):
|
||||
"""
|
||||
Find the principal components and principal axis
|
||||
of inertia from the inertia tensor.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
inertia : (3, 3) float
|
||||
Inertia tensor
|
||||
|
||||
Returns
|
||||
------------
|
||||
components : (3,) float
|
||||
Principal components of inertia
|
||||
vectors : (3, 3) float
|
||||
Row vectors pointing along the
|
||||
principal axes of inertia
|
||||
"""
|
||||
inertia = np.asanyarray(inertia, dtype=np.float64)
|
||||
if inertia.shape != (3, 3):
|
||||
raise ValueError("inertia tensor must be (3, 3)!")
|
||||
|
||||
# you could any of the following to calculate this:
|
||||
# np.linalg.svd, np.linalg.eig, np.linalg.eigh
|
||||
# moment of inertia is square symmetric matrix
|
||||
# eigh has the best precision in tests
|
||||
components, vectors = np.linalg.eigh(inertia)
|
||||
|
||||
# eigh returns them as column vectors, change them to row vectors
|
||||
vectors = vectors.T
|
||||
|
||||
return components, vectors
|
||||
|
||||
|
||||
def transform_inertia(
|
||||
transform: ArrayLike,
|
||||
inertia_tensor: ArrayLike,
|
||||
parallel_axis: bool = False,
|
||||
mass: Optional[Number] = None,
|
||||
):
|
||||
"""
|
||||
Transform an inertia tensor to a new frame.
|
||||
|
||||
Note that in trimesh `mesh.moment_inertia` is *axis aligned*
|
||||
and at `mesh.center_mass`.
|
||||
|
||||
So to transform to a new frame and get the moment of inertia at
|
||||
the center of mass the translation should be ignored and only
|
||||
rotation applied.
|
||||
|
||||
If parallel axis is enabled it will compute the inertia
|
||||
about a new location.
|
||||
|
||||
More details in the MIT OpenCourseWare PDF:
|
||||
` MIT16_07F09_Lec26.pdf`
|
||||
|
||||
|
||||
Parameters
|
||||
------------
|
||||
transform : (3, 3) or (4, 4) float
|
||||
Transformation matrix
|
||||
inertia_tensor : (3, 3) float
|
||||
Inertia tensor.
|
||||
parallel_axis : bool
|
||||
Apply the parallel axis theorum or not.
|
||||
If the passed inertia tensor is at the center of mass
|
||||
and you want the new post-transform tensor also at the
|
||||
center of mass you DON'T want this enabled as you *only*
|
||||
want to apply the rotation. Use this to get moment of
|
||||
inertia at an arbitrary frame that isn't the center of mass.
|
||||
|
||||
Returns
|
||||
------------
|
||||
transformed : (3, 3) float
|
||||
Inertia tensor in new frame.
|
||||
"""
|
||||
# check inputs and extract rotation
|
||||
transform = np.asanyarray(transform, dtype=np.float64)
|
||||
if transform.shape == (4, 4):
|
||||
rotation = transform[:3, :3]
|
||||
elif transform.shape == (3, 3):
|
||||
rotation = transform
|
||||
else:
|
||||
raise ValueError("transform must be (3, 3) or (4, 4)!")
|
||||
|
||||
inertia_tensor = np.asanyarray(inertia_tensor, dtype=np.float64)
|
||||
if inertia_tensor.shape != (3, 3):
|
||||
raise ValueError("inertia_tensor must be (3, 3)!")
|
||||
|
||||
if parallel_axis:
|
||||
if transform.shape == (3, 3):
|
||||
# shorthand for "translation"
|
||||
a = np.zeros(3, dtype=np.float64)
|
||||
else:
|
||||
# get the translation
|
||||
a = transform[:3, 3]
|
||||
# First the changed origin of the new transform is taken into
|
||||
# account. To calculate the inertia tensor
|
||||
# the parallel axis theorem is used
|
||||
M = np.array(
|
||||
[
|
||||
[a[1] ** 2 + a[2] ** 2, -a[0] * a[1], -a[0] * a[2]],
|
||||
[-a[0] * a[1], a[0] ** 2 + a[2] ** 2, -a[1] * a[2]],
|
||||
[-a[0] * a[2], -a[1] * a[2], a[0] ** 2 + a[1] ** 2],
|
||||
]
|
||||
)
|
||||
aligned_inertia = inertia_tensor + mass * M
|
||||
|
||||
return multi_dot([rotation.T, aligned_inertia, rotation])
|
||||
|
||||
return multi_dot([rotation, inertia_tensor, rotation.T])
|
||||
|
||||
|
||||
def radial_symmetry(mesh):
|
||||
"""
|
||||
Check whether a mesh has radial symmetry.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
symmetry : None or str
|
||||
None No rotational symmetry
|
||||
'radial' Symmetric around an axis
|
||||
'spherical' Symmetric around a point
|
||||
axis : None or (3,) float
|
||||
Rotation axis or point
|
||||
section : None or (3, 2) float
|
||||
If radial symmetry provide vectors
|
||||
to get cross section
|
||||
"""
|
||||
|
||||
# shortcuts to avoid typing and hitting cache
|
||||
scalar = mesh.principal_inertia_components.copy()
|
||||
|
||||
# exit early if inertia components are all zero
|
||||
if (scalar < 1e-30).any():
|
||||
return None, None, None
|
||||
|
||||
# normalize the PCI so we can compare them
|
||||
scalar = scalar / np.linalg.norm(scalar)
|
||||
vector = mesh.principal_inertia_vectors
|
||||
# the sorted order of the principal components
|
||||
order = scalar.argsort()
|
||||
|
||||
# we are checking if a geometry has radial symmetry
|
||||
# if 2 of the PCI are equal, it is a revolved 2D profile
|
||||
# if 3 of the PCI (all of them) are equal it is a sphere
|
||||
diff = np.abs(np.diff(scalar[order]))
|
||||
# diffs that are within tol of zero
|
||||
diff_zero = diff < 1e-4
|
||||
|
||||
if diff_zero.all():
|
||||
# this is the case where all 3 PCI are identical
|
||||
# this means that the geometry is symmetric about a point
|
||||
# examples of this are a sphere, icosahedron, etc
|
||||
axis = vector[0]
|
||||
section = vector[1:]
|
||||
|
||||
return "spherical", axis, section
|
||||
|
||||
elif diff_zero.any():
|
||||
# this is the case for 2/3 PCI are identical
|
||||
# this means the geometry is symmetric about an axis
|
||||
# probably a revolved 2D profile
|
||||
|
||||
# we know that only 1/2 of the diff values are True
|
||||
# if the first diff is 0, it means if we take the first element
|
||||
# in the ordered PCI we will have one of the non- revolve axis
|
||||
# if the second diff is 0, we take the last element of
|
||||
# the ordered PCI for the section axis
|
||||
# if we wanted the revolve axis we would just switch [0,-1] to
|
||||
# [-1,0]
|
||||
|
||||
# since two vectors are the same, we know the middle
|
||||
# one is one of those two
|
||||
section_index = order[np.array([[0, 1], [1, -1]])[diff_zero]].flatten()
|
||||
section = vector[section_index]
|
||||
|
||||
# we know the rotation axis is the sole unique value
|
||||
# and is either first or last of the sorted values
|
||||
axis_index = order[np.array([-1, 0])[diff_zero]][0]
|
||||
axis = vector[axis_index]
|
||||
return "radial", axis, section
|
||||
|
||||
return None, None, None
|
||||
|
||||
|
||||
def scene_inertia(scene, transform: Optional[ArrayLike] = None) -> NDArray[float64]:
|
||||
"""
|
||||
Calculate the inertia of a scene about a specific frame.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
scene : trimesh.Scene
|
||||
Scene with geometry.
|
||||
transform : None or (4, 4) float
|
||||
Homogeneous transform to compute inertia at.
|
||||
|
||||
Returns
|
||||
----------
|
||||
moment : (3, 3)
|
||||
Inertia tensor about requested frame
|
||||
"""
|
||||
# shortcuts for tight loop
|
||||
graph = scene.graph
|
||||
geoms = scene.geometry
|
||||
|
||||
# get the matrix ang geometry name for
|
||||
nodes = [graph[n] for n in graph.nodes_geometry]
|
||||
# get the moment of inertia with the mesh moved to a location
|
||||
moments = np.array(
|
||||
[
|
||||
geoms[g].moment_inertia_frame(np.dot(np.linalg.inv(mat), transform))
|
||||
for mat, g in nodes
|
||||
if hasattr(geoms[g], "moment_inertia_frame")
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
|
||||
return moments.sum(axis=0)
|
||||
Reference in New Issue
Block a user