Merge pull request #47 from timothy-nunn/copy-objects

Add copy methods to Machine, Circuit, Coil, and Coil subclasses

Author: kpentland

freegs4e/coil.py

@@ -114,6 +114,20 @@ def __init__(
         self.control = control
         self.area = area
 
+    def copy(self):
+        """Creates a new object that has identical attributes to self (a copy).
+
+        The copy method will need to be implemented for any subclasses of Coil
+        to ensure the correct __init__ signature is called and to ensure
+        lists/other object attributes are copied where appropriate (not passed
+        by reference).
+        """
+        new_obj = type(self)(
+            self.R, self.Z, self.current, self.turns, self.area
+        )
+        new_obj._area = self.area
+        return new_obj
+
     def psi(self, R, Z):
         """
         Calculate poloidal flux at (R,Z)
@@ -283,11 +297,11 @@ def area(self):
         The cross-section area of the coil in m^2
         """
         if isinstance(self._area, numbers.Number):
-            assert self._area > 0
+            assert self._area >= 0
             return self._area
         # Calculate using functor
         area = self._area(self)
-        assert area > 0
+        assert area >= 0
         return area
 
     @area.setter

freegs4e/machine.py

@@ -87,6 +87,15 @@ def __init__(self, coils, current=0.0, control=True):
         self.current = current
         self.control = control
 
+    def copy(self):
+        """Creates a copy of the circuit by initialising a new Circuit object.
+
+        The coils forming the circuit are copied by calling their individual
+        `copy` methods.
+        """
+        coils = [(label, c.copy(), m) for label, c, m in self.coils]
+        return type(self)(coils, self.current, self.control)
+
     def psi(self, R, Z):
         """
         Poloidal flux due to coils in the circuit (at chosen R and Z).
@@ -985,6 +994,14 @@ def __init__(self, coils, wall=None):
         for i, coil in enumerate(self.coil_names):
             self.coil_order[coil] = i
 
+    def copy(self):
+        """Creates a copy of the machine by initialising a new object.
+
+        The coils are copied by calling the `copy` method for the coil/circuit.
+        """
+        coils = [(label, c.copy()) for label, c in self.coils]
+        return type(self)(coils, self.wall)
+
     def __repr__(self):
         """
         Return a string representation of the Machine object.

freegs4e/multi_coil.py

@@ -81,7 +81,7 @@ def __init__(
         turns=1,
         control=True,
         mirror=False,
-        polarity=[1.0, 1.0],
+        polarity=None,
         area=AreaCurrentLimit(),
     ):
         """
@@ -123,13 +123,25 @@ def __init__(
         self.current = current
         self.control = control
         self.mirror = mirror
-        self.polarity = polarity
+        self.polarity = [1.0, 1.0] if polarity is None else polarity
         self.area = area
 
         # Internal (R,Z) value, should not be modified directly
         self._R_centre = np.mean(self.Rfil)
         self._Z_centre = np.mean(self.Zfil)
 
+    def copy(self):
+        return type(self)(
+            np.copy(self.Rfil),
+            np.copy(self.Zfil),
+            self.current,
+            self.turns,
+            self.control,
+            self.mirror,
+            self.polarity,
+            self.area,
+        )
+
     def controlPsi(self, R, Z):
         """
         Calculate poloidal flux at (R,Z) due to a unit current

freegs4e/optimise.py

@@ -24,9 +24,9 @@
 from math import sqrt
 
 import matplotlib.pyplot as plt
-from freegs.plotting import plotEquilibrium
 
 from . import optimiser, picard
+from .plotting import plotEquilibrium
 
 # Measures which operate on Equilibrium objects
 

freegs4e/shaped_coil.py

@@ -88,6 +88,25 @@ def __init__(self, shape, current=0.0, turns=1, control=True, npoints=6):
         self.npoints_per_triangle = npoints
         self._points = quadrature.polygon_quad(shape, n=npoints)
 
+    def copy(self):
+        R_centre = np.copy(self._R_centre)
+        Z_centre = np.copy(self._Z_centre)
+        points = np.copy(self._points)
+
+        new_obj = type(self)(
+            np.copy(self.shape),
+            self.current,
+            self.turns,
+            self.control,
+            self.npoints_per_triangle,
+        )
+
+        new_obj._R_centre = R_centre
+        new_obj._Z_centre = Z_centre
+        new_obj._points = points
+
+        return new_obj
+
     def controlPsi(self, R, Z):
         """
         Calculate poloidal flux at (R,Z) due to a unit current

freegs4e/test_machine.py

@@ -23,10 +23,19 @@ def test_coil_axis():
     def analytic_Bz(dZ):
         return (mu0 / 2) * Rcoil**2 * current / (dZ**2 + Rcoil**2) ** 1.5
 
-    # Note: Can't evaluate at R=0,
-    assert math.isclose(coil.Br(0.0001, 2.0), 0.0, abs_tol=1e-8)
-    assert math.isclose(coil.Bz(0.001, 2.0), analytic_Bz(1.0), abs_tol=1e-8)
-    assert math.isclose(coil.Bz(0.001, -1.0), analytic_Bz(-2.0), abs_tol=1e-8)
+    # run the test twice: once on the original coil and once on a copy
+    # to check both produce the same results
+    for _ in range(2):
+        # Note: Can't evaluate at R=0,
+        assert math.isclose(coil.Br(0.0001, 2.0), 0.0, abs_tol=1e-8)
+        assert math.isclose(
+            coil.Bz(0.001, 2.0), analytic_Bz(1.0), abs_tol=1e-8
+        )
+        assert math.isclose(
+            coil.Bz(0.001, -1.0), analytic_Bz(-2.0), abs_tol=1e-8
+        )
+
+        coil = coil.copy()
 
 
 def test_coil_forces():

freegs4e/test_multi_coil.py

@@ -13,9 +13,17 @@ def test_single():
     mcoil = MultiCoil(1.1, 0.2, current=100.0, mirror=False)
     coil = MultiCoil(1.1, 0.2, current=100.0)
 
-    assert np.isclose(coil.controlPsi(0.3, 0.1), mcoil.controlPsi(0.3, 0.1))
+    # run the test twice: once on the original coil and once on a copy
+    # to check both produce the same results
+    for _ in range(2):
+        assert np.isclose(
+            coil.controlPsi(0.3, 0.1), mcoil.controlPsi(0.3, 0.1)
+        )
 
-    assert np.isclose(coil.controlBr(0.3, 0.1), mcoil.controlBr(0.3, 0.1))
+        assert np.isclose(coil.controlBr(0.3, 0.1), mcoil.controlBr(0.3, 0.1))
+
+        mcoil = mcoil.copy()
+        coil = coil.copy()
 
 
 def test_two_turns():
@@ -52,11 +60,23 @@ def test_mirrored():
         ]
     )
 
-    assert np.isclose(circuit.controlPsi(0.3, 0.1), mcoil.controlPsi(0.3, 0.1))
+    # run the test twice: once on the original coil and once on a copy
+    # to check both produce the same results
+    for _ in range(2):
+        assert np.isclose(
+            circuit.controlPsi(0.3, 0.1), mcoil.controlPsi(0.3, 0.1)
+        )
 
-    assert np.isclose(circuit.controlBr(0.3, 0.1), mcoil.controlBr(0.3, 0.1))
+        assert np.isclose(
+            circuit.controlBr(0.3, 0.1), mcoil.controlBr(0.3, 0.1)
+        )
 
-    assert np.isclose(circuit.controlBz(0.3, 0.1), mcoil.controlBz(0.3, 0.1))
+        assert np.isclose(
+            circuit.controlBz(0.3, 0.1), mcoil.controlBz(0.3, 0.1)
+        )
+
+        mcoil = mcoil.copy()
+        circuit = circuit.copy()
 
 
 def test_move_R():

freegs4e/test_shaped_coil.py

@@ -29,14 +29,31 @@ def test_move_R():
         current=100.0,
     )
 
+    # change the copy by a different amount to ensure changing the copy
+    # does not impact the original coil
+    coil1_copy = coil1.copy()
+    coil1_copy.R += dR + 0.4
+
     # Shift coil1 to same location as coil2
     coil1.R += dR
 
-    assert np.isclose(coil1.controlPsi(0.4, 0.5), coil2.controlPsi(0.4, 0.5))
+    # run the test twice: once on the original coil and once on a copy
+    # to check both produce the same results
+    for _ in range(2):
+        assert np.isclose(
+            coil1.controlPsi(0.4, 0.5), coil2.controlPsi(0.4, 0.5)
+        )
 
-    assert np.isclose(coil1.controlBr(0.3, -0.2), coil2.controlBr(0.3, -0.2))
+        assert np.isclose(
+            coil1.controlBr(0.3, -0.2), coil2.controlBr(0.3, -0.2)
+        )
 
-    assert np.isclose(coil1.controlBz(1.75, 1.2), coil2.controlBz(1.75, 1.2))
+        assert np.isclose(
+            coil1.controlBz(1.75, 1.2), coil2.controlBz(1.75, 1.2)
+        )
+
+        coil1 = coil1.copy()
+        coil2 = coil2.copy()
 
 
 def test_move_Z():