Remove deepcopying of profiles

examples/example02 - static_forward_solve_MASTU.ipynb

@@ -227,11 +227,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from copy import deepcopy\n",
-    "\n",
     "# copy the original eq object (for the new forward solves with modified currents)\n",
-    "eq_forward_1 = deepcopy(eq)\n",
-    "eq_forward_2 = deepcopy(eq)\n",
+    "eq_forward_1 = eq.create_auxiliary_equilibrium()\n",
+    "eq_forward_2 = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# modify the P4 current and solve\n",
     "eq_forward_1.tokamak.set_coil_current('P4', 1.4*eq.tokamak['P4'].current)\n",
@@ -317,7 +315,7 @@
    "outputs": [],
    "source": [
     "# instatiate new equilibrium object\n",
-    "eq_beta = deepcopy(eq)\n",
+    "eq_beta = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# call solver with new profile object\n",
     "GSStaticSolver.solve(eq=eq_beta, \n",
@@ -368,7 +366,7 @@
    "outputs": [],
    "source": [
     "# instatiate new equilibrium object\n",
-    "eq_topeol = deepcopy(eq)\n",
+    "eq_topeol = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# call solver with new profile object\n",
     "GSStaticSolver.solve(eq=eq_topeol, \n",
@@ -476,7 +474,7 @@
    "outputs": [],
    "source": [
     "# instatiate new equilibrium object\n",
-    "eq_lao = deepcopy(eq)\n",
+    "eq_lao = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# call solver with new profile object\n",
     "GSStaticSolver.solve(eq=eq_lao, \n",
@@ -691,7 +689,7 @@
    "outputs": [],
    "source": [
     "# instatiate new equilibrium object\n",
-    "eq_general = deepcopy(eq)\n",
+    "eq_general = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# call solver with new profile object\n",
     "GSStaticSolver.solve(eq=eq_general, \n",
@@ -726,8 +724,7 @@
    "outputs": [],
    "source": [
     "# initialise the equilibrium\n",
-    "eq_limiter = deepcopy(eq)\n",
-    "\n",
+    "eq_limiter = eq.create_auxiliary_equilibrium()\n",
     "\n",
     "# initialise the profiles\n",
     "profiles = ConstrainPaxisIp(\n",
@@ -767,6 +764,13 @@
     "ax1.set_ylim(-2.25, 2.25)\n",
     "plt.tight_layout()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

examples/example03 - extracting_equilibrium_quantites.ipynb

@@ -328,17 +328,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import shapely as sh\n",
+    "\n",
     "# visualising these quantities\n",
     "psi_bndry = eq.psi_bndry\n",
     "xpts = eq.xpt\n",
     "\n",
-    "# compute absolute difference between each point's ψ and psi_bndry and then\n",
-    "# get indices of the two smallest differences\n",
-    "closest_indices = np.argsort(np.abs(xpts[:, 2] - psi_bndry))[:2]\n",
+    "# define your polygon (example)\n",
+    "polygon = sh.Polygon(np.array([eq.tokamak.wall.R, eq.tokamak.wall.Z]).T)\n",
+    "\n",
+    "# vectorized-ish approach: boolean mask for points inside\n",
+    "mask = np.array([polygon.contains(sh.Point(x, y)) for x, y in xpts[:,0:2]])\n",
+    "\n",
+    "# select points inside as a NumPy array\n",
+    "xpts_inside_wall = xpts[mask,:]\n",
     "\n",
-    "# extract the corresponding rows (two X-points closest to psi_bndry, then sort by lowest z coord z-point)\n",
-    "closest_xpts = xpts[closest_indices]\n",
-    "closest_xpts_sorted = closest_xpts[np.argsort(closest_xpts[:, 1])]\n",
+    "# get indices of the two cloest to magnetic axis \n",
+    "closest_xpts_idxs = np.argsort(np.linalg.norm(xpts_inside_wall[:,0:2] - eq.opt[0,0:2], axis=1))\n",
+    "closest_xpts_sorted = xpts_inside_wall[closest_xpts_idxs[0:2],:]\n",
     "\n",
     "# plot the flux contours for the values of psi_boundary at each X-point\n",
     "fig1, axes = plt.subplots(1, 3, figsize=(15, 8), dpi=80)\n",

examples/example05 - evolutive_forward_solve.ipynb

@@ -33,10 +33,8 @@
    "source": [
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
-    "from copy import deepcopy\n",
     "from IPython.display import display, clear_output\n",
-    "import pickle\n",
-    "import copy"
+    "import pickle"
    ]
   },
   {
@@ -305,7 +303,7 @@
     "\n",
     "history_times = [t]\n",
     "history_currents = [stepping.currents_vec]\n",
-    "history_equilibria = [deepcopy(stepping.eq1)]\n",
+    "history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]\n",
     "history_o_points = [stepping.eq1.opt[0]]\n",
     "history_elongation = [stepping.eq1.geometricElongation()]\n",
     "history_triangularity = [stepping.eq1.triangularity()]\n",
@@ -373,7 +371,7 @@
     "\n",
     "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
     "    history_currents.append(stepping.currents_vec)\n",
-    "    history_equilibria.append(deepcopy(stepping.eq1))\n",
+    "    history_equilibria.append(stepping.eq1.create_auxiliary_equilibrium())\n",
     "    history_o_points.append(stepping.eq1.opt[0])\n",
     "    history_elongation.append(stepping.eq1.geometricElongation())\n",
     "    history_triangularity.append(stepping.eq1.triangularity())\n",
@@ -422,7 +420,7 @@
     "\n",
     "history_times_nl = [t]\n",
     "history_currents_nl = [stepping.currents_vec]\n",
-    "history_equilibria_nl = [deepcopy(stepping.eq1)]\n",
+    "history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
     "history_o_points_nl = [stepping.eq1.opt[0]]\n",
     "history_elongation_nl = [stepping.eq1.geometricElongation()]\n",
     "history_triangularity_nl = [stepping.eq1.triangularity()]\n",
@@ -451,7 +449,7 @@
     "    \n",
     "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
     "    history_currents_nl.append(stepping.currents_vec)\n",
-    "    history_equilibria_nl.append(deepcopy(stepping.eq1))\n",
+    "    history_equilibria_nl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
     "    history_o_points_nl.append(stepping.eq1.opt[0])\n",
     "    history_elongation_nl.append(stepping.eq1.geometricElongation())\n",
     "    history_triangularity_nl.append(stepping.eq1.triangularity())\n",
@@ -750,7 +748,7 @@
     "\n",
     "history_times = [t]\n",
     "history_currents = [stepping.currents_vec]\n",
-    "history_equilibria = [deepcopy(stepping.eq1)]\n",
+    "history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]\n",
     "history_o_points = [stepping.eq1.opt[0]]\n",
     "history_elongation = [stepping.eq1.geometricElongation()]\n",
     "history_triangularity = [stepping.eq1.triangularity()]\n",
@@ -783,7 +781,7 @@
     "    \n",
     "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
     "    history_currents.append(stepping.currents_vec)\n",
-    "    history_equilibria.append(deepcopy(stepping.eq1))\n",
+    "    history_equilibria.append(stepping.eq1.create_auxiliary_equilibrium())\n",
     "    history_o_points.append(stepping.eq1.opt[0])\n",
     "    history_elongation.append(stepping.eq1.geometricElongation())\n",
     "    history_triangularity.append(stepping.eq1.triangularity())\n",
@@ -826,7 +824,7 @@
     "\n",
     "history_times_nl = [t]\n",
     "history_currents_nl = [stepping.currents_vec]\n",
-    "history_equilibria_nl = [deepcopy(stepping.eq1)]\n",
+    "history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
     "history_o_points_nl = [stepping.eq1.opt[0]]\n",
     "history_elongation_nl = [stepping.eq1.geometricElongation()]\n",
     "history_triangularity_nl = [stepping.eq1.triangularity()]\n",
@@ -859,7 +857,7 @@
     "    \n",
     "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
     "    history_currents_nl.append(stepping.currents_vec)\n",
-    "    history_equilibria_nl.append(deepcopy(stepping.eq1))\n",
+    "    history_equilibria_nl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
     "    history_o_points_nl.append(stepping.eq1.opt[0])\n",
     "    history_elongation_nl.append(stepping.eq1.geometricElongation())\n",
     "    history_triangularity_nl.append(stepping.eq1.triangularity())\n",
@@ -978,6 +976,13 @@
    "outputs": [],
    "source": []
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,

examples/example06a - static_forward_MASTU_solves_from_magnetics_only_data.ipynb

@@ -115,7 +115,7 @@
    "outputs": [],
    "source": [
     "# set paths (choose symmetric or non-symmetric active coils)\n",
-    "symmetric_machine = False\n",
+    "symmetric_machine = True\n",
     "if symmetric_machine:\n",
     "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_active_coils.pickle\"\n",
     "else:\n",
@@ -268,7 +268,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "time_indices = time_indices[1:-1]\n",
+    "time_indices = time_indices[0:-1]\n",
     "times = efit_times[time_indices]\n",
     "n = len(times)"
    ]

examples/example09 - virtual_circuits_MASTU.ipynb

@@ -297,9 +297,9 @@
     "ax1.scatter(all_target_values[1], 0.0, s=100, color='blue', marker='x', zorder=20, label=all_target_names[1])\n",
     "ax1.scatter(all_target_values[2], all_target_values[3], s=100, color='m', marker='*', zorder=20, label=f\"({all_target_names[2]},{all_target_names[3]})\")\n",
     "ax1.scatter(all_target_values[4], all_target_values[5], s=100, color='k', marker='*', zorder=20, label=f\"({all_target_names[4]},{all_target_names[5]})\")\n",
-    "ax1.scatter(all_target_values[6], -1.9 , s=100, color='k', marker='x', zorder=20, label=f\"{all_target_names[6]}\")\n",
+    "ax1.scatter(all_target_values[6], -1.95 , s=100, color='k', marker='x', zorder=20, label=f\"{all_target_names[6]}\")\n",
     "ax1.scatter(all_target_values[7], 1.95 , s=100, color='r', marker='x', zorder=20, label=f\"{all_target_names[7]}\")\n",
-    "ax1.scatter(all_target_values[8], -1.95 , s=100, color='orange', marker='o', zorder=5, label=f\"{all_target_names[8]}\")\n",
+    "ax1.scatter(all_target_values[8], -1.5 , s=100, color='orange', marker='o', zorder=5, label=f\"{all_target_names[8]}\")\n",
     "\n",
     "ax1.set_xlim(0.1, 2.15)\n",
     "ax1.set_ylim(-2.25, 2.25)\n",
@@ -719,6 +719,13 @@
     "plt.tight_layout()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,

freegsnke/GSstaticsolver.py

@@ -233,7 +233,7 @@ def port_critical(self, eq, profiles):
         eq.flag_limiter = profiles.flag_limiter
 
         eq._current = np.sum(profiles.jtor) * self.dRdZ
-        eq._profiles = deepcopy(profiles)
+        eq._profiles = profiles.copy()
 
         try:
             eq.tokamak_psi = self.tokamak_psi.reshape(self.nx, self.ny)

freegsnke/copying.py

@@ -0,0 +1,64 @@
+import copy
+import logging
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def copy_into(
+    obj, new_obj, attr: str, *, mutable=False, strict=True, allow_deepcopy=False
+):
+    """Copies an attribute from one object to another.
+
+    Parameters
+    ==========
+    obj : object
+        The object to copy the attribute from
+
+    new_obj : object
+        The object to copy the attribute to
+
+    attr : str
+        The attribute name to copy (e.g. copy `obj.attr` into `new_obj`)
+
+    mutable : bool
+        If an attribute is mutable it needs to be copied explicitly between objects,
+        a reference is not sufficient.
+
+    strict : bool
+        Raise an error if `True` and `attr` is not an attribute of `obj`
+
+    allow_deepcopy : bool
+        Raise an error if `True` and a deepcopy is required to copy the attribute across.
+    """
+
+    if not hasattr(obj, attr) and not strict:
+        logger.info(f"{obj.__class__} has no attribute {attr} but not in strict mode")
+        # return without an error because we are not strict
+        return
+
+    # will error if strict and attribute doesnt exist
+    attribute_value = getattr(obj, attr)
+
+    if mutable:
+        if (
+            isinstance(attribute_value, np.ndarray)
+            and not attribute_value.dtype.hasobject
+        ):
+            attribute_value = np.copy(attribute_value)
+
+        else:
+            if not allow_deepcopy:
+                raise TypeError(
+                    f"Cannot copy {attribute_value.__class__} without deepcopying"
+                )
+
+            logger.info(
+                f"Deepcopying {attribute_value.__class__} because it is mutable but not a numpy array"
+                "of non-objects"
+            )
+
+            attribute_value = copy.deepcopy(attribute_value)
+
+    setattr(new_obj, attr, attribute_value)

freegsnke/equilibrium_update.py

@@ -27,8 +27,9 @@
 from freegs4e import critical
 from scipy import interpolate
 
-from . import limiter_func, virtual_circuits
+from . import limiter_func
 from .build_machine import copy_tokamak
+from .copying import copy_into
 
 
 class Equilibrium(freegs4e.equilibrium.Equilibrium):
@@ -88,8 +89,6 @@ def create_auxiliary_equilibrium(self):
         equilibrium.dR = self.dR
         equilibrium.dZ = self.dZ
         equilibrium._applyBoundary = self._applyBoundary
-        equilibrium._pgreen = self._pgreen
-        equilibrium._vgreen = self._vgreen
         equilibrium._current = self._current
         equilibrium.order = self.order
         equilibrium._solver = self._solver
@@ -112,6 +111,20 @@ def create_auxiliary_equilibrium(self):
         equilibrium.plasma_psi = np.copy(self.plasma_psi)
         equilibrium.mask_inside_limiter = np.copy(self.mask_inside_limiter)
         equilibrium.mask_outside_limiter = np.copy(self.mask_outside_limiter)
+        equilibrium._pgreen = self._pgreen.copy()
+        equilibrium._vgreen = self._vgreen.copy()
+        copy_into(self, equilibrium, "current_vec", mutable=True, strict=False)
+
+        copy_into(
+            self, equilibrium, "opt", mutable=True, strict=False, allow_deepcopy=True
+        )
+        copy_into(
+            self, equilibrium, "xpt", mutable=True, strict=False, allow_deepcopy=True
+        )
+        copy_into(self, equilibrium, "psi_bndry", strict=False)
+
+        if hasattr(self, "_profiles"):
+            equilibrium._profiles = self._profiles.copy()
 
         return equilibrium