From 4873af2d4c513b20e526a2bd6e1167384c74e937 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Sun, 15 Dec 2024 17:58:35 +0100
Subject: [PATCH] added major and minor radius attributes

---
 src/paramak/assemblies/assembly.py            | 21 ++++++++++++++-----
 src/paramak/assemblies/spherical_tokamak.py   |  3 +++
 src/paramak/assemblies/tokamak.py             |  5 +++++
 .../test_assemblies/test_spherical_tokamak.py |  6 ++++--
 4 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/paramak/assemblies/assembly.py b/src/paramak/assemblies/assembly.py
index 5e65b920..ff556eb2 100644
--- a/src/paramak/assemblies/assembly.py
+++ b/src/paramak/assemblies/assembly.py
@@ -8,9 +8,21 @@
 class Assembly(cq.Assembly):
     """Nested assembly of Workplane and Shape objects defining their relative positions."""
 
-    def remove(self, name:str):
-        new_assembly = Assembly()
-        part_found=False
+    def __init__(self, elongation=None, triangularity=None, major_radius=None, minor_radius=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.elongation = elongation
+        self.triangularity = triangularity
+        self.major_radius = major_radius
+        self.minor_radius = minor_radius
+
+    def remove(self, name: str):
+        new_assembly = Assembly(
+            elongation=self.elongation,
+            triangularity=self.triangularity,
+            major_radius=self.major_radius,
+            minor_radius=self.minor_radius
+        )
+        part_found = False
         for part in self:
             if part[1].endswith(f'/{name}'):
                 part_found = True
@@ -18,11 +30,10 @@ def remove(self, name:str):
             else:
                 # print('adding' , part)
                 new_assembly.add(part[0], name=part[1], color=part[3], loc=part[2])
-        if part_found == False:
+        if not part_found:
             warnings.warn(f'Part with name {name} not found')
         return new_assembly
 
-
     def names(self):
         names = []
         for part in self:
diff --git a/src/paramak/assemblies/spherical_tokamak.py b/src/paramak/assemblies/spherical_tokamak.py
index 6ef7d279..a9ae5bb2 100644
--- a/src/paramak/assemblies/spherical_tokamak.py
+++ b/src/paramak/assemblies/spherical_tokamak.py
@@ -289,4 +289,7 @@ def spherical_tokamak(
 
     my_assembly.elongation = elongation
     my_assembly.triangularity = triangularity
+    my_assembly.major_radius = major_radius
+    my_assembly.minor_radius = minor_radius
+
     return my_assembly
diff --git a/src/paramak/assemblies/tokamak.py b/src/paramak/assemblies/tokamak.py
index 68c38a9e..687b3984 100644
--- a/src/paramak/assemblies/tokamak.py
+++ b/src/paramak/assemblies/tokamak.py
@@ -323,4 +323,9 @@ def tokamak(
 
     my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
+    my_assembly.elongation = elongation
+    my_assembly.triangularity = triangularity
+    my_assembly.major_radius = major_radius
+    my_assembly.minor_radius = minor_radius
+
     return my_assembly
diff --git a/tests/test_assemblies/test_spherical_tokamak.py b/tests/test_assemblies/test_spherical_tokamak.py
index 134f897d..447e9ebf 100644
--- a/tests/test_assemblies/test_spherical_tokamak.py
+++ b/tests/test_assemblies/test_spherical_tokamak.py
@@ -19,7 +19,7 @@ def test_transport_with_magnets(rotation_angle):
         [10, 15],
         [20, 50],
         [20, 50],
-        [(500, 300), (610, 100)],
+        [(500, 300), (590, 100)],
     ):
         poloidal_field_coils.append(
             paramak.poloidal_field_coil(
@@ -155,4 +155,6 @@ def test_attributes():
     )
 
     assert my_reactor.elongation == 2
-    assert my_reactor.triangularity == 0.55
\ No newline at end of file
+    assert my_reactor.triangularity == 0.55
+    assert my_reactor.major_radius == 275
+    assert my_reactor.minor_radius == 150
\ No newline at end of file