From e23cb5cc47a7da1fa522ad7b8401c9306661e7c6 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Sun, 21 Jul 2024 00:04:59 -0700
Subject: [PATCH] feat: Centre of mass

---
 nhf/parts/item.py |  4 ++--
 nhf/test.py       | 27 +++++++++++++++++++++++++++
 nhf/utils.py      | 32 +++++++++++++++++++++++++++-----
 3 files changed, 56 insertions(+), 7 deletions(-)

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