283 lines
7.7 KiB
Python
283 lines
7.7 KiB
Python
"""
|
|
Utility functions for cadquery objects
|
|
"""
|
|
import functools
|
|
import math
|
|
from typing import Optional
|
|
import cadquery as Cq
|
|
from cadquery.occ_impl.solver import ConstraintSpec
|
|
from nhf import Role
|
|
from typing import Union, Tuple, cast
|
|
from nhf.materials import KEY_ITEM, KEY_MATERIAL
|
|
|
|
# Bug fixes
|
|
def _subloc(self, name: str) -> Tuple[Cq.Location, str]:
|
|
"""
|
|
Calculate relative location of an object in a subassembly.
|
|
|
|
Returns the relative positions as well as the name of the top assembly.
|
|
"""
|
|
|
|
rv = Cq.Location()
|
|
obj = self.objects[name]
|
|
name_out = name
|
|
|
|
if obj not in self.children and obj is not self:
|
|
locs = []
|
|
while not obj.parent is self:
|
|
locs.append(obj.loc)
|
|
obj = cast(Cq.Assembly, obj.parent)
|
|
name_out = obj.name
|
|
|
|
rv = functools.reduce(lambda l1, l2: l2 * l1, locs)
|
|
|
|
return (rv, name_out)
|
|
Cq.Assembly._subloc = _subloc
|
|
|
|
### Vector arithmetic
|
|
|
|
def location_sub(self: Cq.Location, rhs: Cq.Location) -> Cq.Vector:
|
|
(x1, y1, z1), _ = self.toTuple()
|
|
(x2, y2, z2), _ = rhs.toTuple()
|
|
return Cq.Vector(x1 - x2, y1 - y2, z1 - z2)
|
|
Cq.Location.__sub__ = location_sub
|
|
|
|
def from2d(x: float, y: float, rotate: float=0.0) -> Cq.Location:
|
|
return Cq.Location((x, y, 0), (0, 0, 1), rotate)
|
|
Cq.Location.from2d = from2d
|
|
|
|
def rot2d(angle: float) -> Cq.Location:
|
|
return Cq.Location((0, 0, 0), (0, 0, 1), angle)
|
|
Cq.Location.rot2d = rot2d
|
|
|
|
def is2d(self: Cq.Location) -> bool:
|
|
(_, _, z), (rx, ry, _) = self.toTuple()
|
|
return z == 0 and rx == 0 and ry == 0
|
|
Cq.Location.is2d = is2d
|
|
|
|
def to2d(self: Cq.Location) -> Tuple[Tuple[float, float], float]:
|
|
"""
|
|
Returns position and angle
|
|
"""
|
|
(x, y, z), (rx, ry, rz) = self.toTuple()
|
|
assert z == 0
|
|
assert rx == 0
|
|
assert ry == 0
|
|
return (x, y), rz
|
|
Cq.Location.to2d = to2d
|
|
|
|
def to2d_pos(self: Cq.Location) -> Tuple[float, float]:
|
|
"""
|
|
Returns position and angle
|
|
"""
|
|
(x, y), _ = self.to2d()
|
|
return x, y
|
|
Cq.Location.to2d_pos = to2d_pos
|
|
|
|
def to2d_rot(self: Cq.Location) -> float:
|
|
"""
|
|
Returns position and angle
|
|
"""
|
|
_, r = self.to2d()
|
|
return r
|
|
Cq.Location.to2d_rot = to2d_rot
|
|
|
|
|
|
def with_angle_2d(self: Cq.Location, angle: float) -> Tuple[float, float]:
|
|
"""
|
|
Returns position and angle
|
|
"""
|
|
x, y = self.to2d_pos()
|
|
return Cq.Location.from2d(x, y, angle)
|
|
Cq.Location.with_angle_2d = with_angle_2d
|
|
|
|
def flip_x(self: Cq.Location) -> Cq.Location:
|
|
(x, y), a = self.to2d()
|
|
return Cq.Location.from2d(-x, y, 90 - a)
|
|
Cq.Location.flip_x = flip_x
|
|
def flip_y(self: Cq.Location) -> Cq.Location:
|
|
(x, y), a = self.to2d()
|
|
return Cq.Location.from2d(x, -y, -a)
|
|
Cq.Location.flip_y = flip_y
|
|
|
|
def boolean(self: Cq.Sketch, obj, **kwargs) -> Cq.Sketch:
|
|
return (
|
|
self
|
|
.reset()
|
|
.push([(0, 0)])
|
|
.each(lambda _: obj, **kwargs)
|
|
)
|
|
Cq.Sketch.boolean = boolean
|
|
|
|
### Tags
|
|
|
|
def tagPoint(self, tag: str):
|
|
"""
|
|
Adds a vertex that can be used in `Point` constraints.
|
|
"""
|
|
vertex = Cq.Vertex.makeVertex(0, 0, 0)
|
|
self.eachpoint(vertex.moved, useLocalCoordinates=True).tag(tag)
|
|
|
|
Cq.Workplane.tagPoint = tagPoint
|
|
|
|
def tagPlane(self, tag: str,
|
|
direction: Union[str, Cq.Vector, Tuple[float, float, float]] = '+Z'):
|
|
"""
|
|
Adds a phantom `Cq.Edge` in the given location which can be referenced in a
|
|
`Axis`, `Point`, or `Plane` constraint.
|
|
"""
|
|
if isinstance(direction, str):
|
|
x, y, z = 0, 0, 0
|
|
assert len(direction) == 2
|
|
sign, axis = direction
|
|
if axis in ('z', 'Z'):
|
|
z = 1
|
|
elif axis in ('y', 'Y'):
|
|
y = 1
|
|
elif axis in ('x', 'X'):
|
|
x = 1
|
|
else:
|
|
assert False, "Axis must be one of x,y,z"
|
|
if sign == '+':
|
|
sign = 1
|
|
elif sign == '-':
|
|
sign = -1
|
|
else:
|
|
assert False, "Sign must be one of +/-"
|
|
v = Cq.Vector(x, y, z) * sign
|
|
else:
|
|
v = Cq.Vector(direction)
|
|
edge = Cq.Edge.makeLine(v * (-1), v)
|
|
return self.eachpoint(edge.located, useLocalCoordinates=True).tag(tag)
|
|
|
|
Cq.Workplane.tagPlane = tagPlane
|
|
|
|
def make_sphere(r: float = 2) -> Cq.Solid:
|
|
"""
|
|
Makes a full sphere. The default function makes a hemisphere
|
|
"""
|
|
return Cq.Solid.makeSphere(r, angleDegrees1=-90)
|
|
def make_arrow(size: float = 2) -> Cq.Workplane:
|
|
cone = Cq.Solid.makeCone(
|
|
radius1 = size,
|
|
radius2 = 0,
|
|
height=size)
|
|
result = (
|
|
Cq.Workplane("XY")
|
|
.cylinder(radius=size / 2, height=size, centered=(True, True, False))
|
|
.union(cone.located(Cq.Location((0, 0, size))))
|
|
)
|
|
result.faces("<Z").tag("dir_rev")
|
|
return result
|
|
|
|
def to_marker_name(tag: str) -> str:
|
|
return tag.replace("?", "__T").replace("/", "__Z") + "_marker"
|
|
|
|
COLOR_MARKER = Cq.Color(0, 1, 1, 1)
|
|
|
|
def mark_point(self: Cq.Assembly,
|
|
tag: str,
|
|
size: float = 2,
|
|
color: Cq.Color = COLOR_MARKER) -> Cq.Assembly:
|
|
"""
|
|
Adds a marker to make a point visible
|
|
"""
|
|
name = to_marker_name(tag)
|
|
return (
|
|
self
|
|
.add(make_sphere(size), name=name, color=color)
|
|
.constrain(tag, name, "Point")
|
|
)
|
|
|
|
Cq.Assembly.markPoint = mark_point
|
|
|
|
def mark_plane(self: Cq.Assembly,
|
|
tag: str,
|
|
size: float = 2,
|
|
color: Cq.Color = COLOR_MARKER) -> Cq.Assembly:
|
|
"""
|
|
Adds a marker to make a plane visible
|
|
"""
|
|
name = to_marker_name(tag)
|
|
return (
|
|
self
|
|
.add(make_arrow(size), name=name, color=color)
|
|
.constrain(tag, f"{name}?dir_rev", "Plane", param=180)
|
|
)
|
|
|
|
Cq.Assembly.markPlane = mark_plane
|
|
|
|
def get_abs_location(self: Cq.Assembly,
|
|
tag: str) -> Cq.Location:
|
|
"""
|
|
Gets the location of a tag
|
|
|
|
BUG: Currently bugged. See `nhf/test.py` for example
|
|
"""
|
|
name, shape = self._query(tag)
|
|
loc_self = Cq.Location(shape.Center())
|
|
loc_parent, _ = self._subloc(name)
|
|
loc = loc_parent * loc_self
|
|
return loc
|
|
|
|
Cq.Assembly.get_abs_location = get_abs_location
|
|
|
|
def get_abs_direction(self: Cq.Assembly,
|
|
tag: str) -> Cq.Location:
|
|
"""
|
|
Gets the location of a tag
|
|
"""
|
|
name, shape = self._query(tag)
|
|
# Must match `cadquery.occ_impl.solver.ConstraintSpec._getAxis`
|
|
if isinstance(shape, Cq.Face):
|
|
vec_dir = shape.normalAt()
|
|
elif isinstance(shape, Cq.Edge) and shape.geomType() != "CIRCLE":
|
|
vec_dir = shape.tangentAt()
|
|
elif isinstance(shape, Cq.Edge) and shape.geomType() == "CIRCLE":
|
|
vec_dir = shape.normal()
|
|
else:
|
|
raise ValueError(f"Cannot construct Axis for {shape}")
|
|
loc_self = Cq.Location(vec_dir)
|
|
loc_parent, _ = self._subloc(name)
|
|
loc = loc_parent * loc_self
|
|
return loc
|
|
Cq.Assembly.get_abs_direction = get_abs_direction
|
|
|
|
|
|
# Tallying functions
|
|
|
|
def assembly_this_mass(self: Cq.Assembly) -> Optional[float]:
|
|
"""
|
|
Gets the mass of an assembly, without considering its components.
|
|
"""
|
|
if item := self.metadata.get(KEY_ITEM):
|
|
return item.mass
|
|
elif material := self.metadata.get(KEY_MATERIAL):
|
|
vol = self.toCompound().Volume()
|
|
return (vol / 1000) * material.density
|
|
else:
|
|
return None
|
|
|
|
def total_mass(self: Cq.Assembly) -> float:
|
|
"""
|
|
Calculates the total mass in units of g
|
|
"""
|
|
total = 0.0
|
|
for _, a in self.traverse():
|
|
if m := assembly_this_mass(a):
|
|
total += m
|
|
return total
|
|
Cq.Assembly.total_mass = total_mass
|
|
|
|
def centre_of_mass(self: Cq.Assembly) -> Optional[float]:
|
|
moment = Cq.Vector()
|
|
total = 0.0
|
|
for n, a in self.traverse():
|
|
if m := assembly_this_mass(a):
|
|
moment += m * a.toCompound().Center()
|
|
total += m
|
|
if total == 0.0:
|
|
return None
|
|
return moment / total
|
|
Cq.Assembly.centre_of_mass = centre_of_mass
|