diff --git a/nhf/parts/item.py b/nhf/parts/item.py index e7a153e..3fd5450 100644 --- a/nhf/parts/item.py +++ b/nhf/parts/item.py @@ -31,7 +31,7 @@ class Item: def role(self) -> Optional[Role]: return None - def generate(self, **kwargs) -> Union[Cq.Assembly, Cq.Workplane]: + def generate(self, **kwargs) -> Union[Cq.Solid, Cq.Assembly, Cq.Workplane]: """ Creates an assembly for this item. Subclass should implement this """ @@ -42,7 +42,7 @@ class Item: Interface for creating assembly with the necessary metadata """ a = self.generate(**kwargs) - if isinstance(a, Cq.Workplane): + if isinstance(a, Cq.Workplane) or isinstance(a, Cq.Solid): a = Cq.Assembly(a) if role := self.role: a.metadata[KEY_ROLE] = role diff --git a/nhf/test.py b/nhf/test.py index 7066a9a..b2dbb45 100644 --- a/nhf/test.py +++ b/nhf/test.py @@ -1,9 +1,11 @@ """ Unit tests for tooling """ +from dataclasses import dataclass import unittest import cadquery as Cq from nhf.build import Model, target +from nhf.parts.item import Item import nhf.checks import nhf.utils @@ -17,6 +19,20 @@ def makeSphere(r: float) -> Cq.Solid: """ return Cq.Solid.makeSphere(r, angleDegrees1=-90) +@dataclass(frozen=True) +class MassBall(Item): + """ + A ball with fixed mass + """ + radius: float = 0.2 + + @property + def name(self) -> str: + return f"MassBall {self.mass}" + def generate(self) -> Cq.Solid: + return makeSphere(self.radius) + + class BuildScaffold(Model): def __init__(self): @@ -180,5 +196,16 @@ class TestUtils(unittest.TestCase): self.assertAlmostEqual(ry, 1) self.assertAlmostEqual(rz, 0) + def test_centre_of_mass(self): + assembly = ( + Cq.Assembly() + .add(MassBall(mass=3).assembly(), name="s1", loc=Cq.Location((0, 0, 0))) + .add(MassBall(mass=7).assembly(), name="s2", loc=Cq.Location((0, 0, 10))) + ) + com = assembly.centre_of_mass() + self.assertAlmostEqual(com.x, 0) + self.assertAlmostEqual(com.y, 0) + self.assertAlmostEqual(com.z, 7) + if __name__ == '__main__': unittest.main() diff --git a/nhf/utils.py b/nhf/utils.py index 876f891..1d2c796 100644 --- a/nhf/utils.py +++ b/nhf/utils.py @@ -3,6 +3,7 @@ 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 @@ -236,16 +237,37 @@ 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 item := a.metadata.get(KEY_ITEM): - total += item.mass - elif material := a.metadata.get(KEY_MATERIAL): - vol = a.toCompound().Volume() - total += (vol / 1000) * material.density + 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