Speed up linearization construction

examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb

@@ -216,6 +216,7 @@
     "- `plasma_resistivity`: resistivity of the plasma (which here is assumed to be constant during the time evolution but can be made time-dependent if known).\n",
     "- `min_dIy_dI`: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<`min_dIy_dI` are dropped, which filters out modes that do not actually couple with the plasma.\n",
     "- `max_mode_frequency`: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).\n",
+    "- **`n_linearization_workers`**: number of independent worker processes used to build the current-response (`dIydI`) and profile-parameter-response (`dIydtheta`) Jacobian columns. The default, `1`, retains serial construction. Values larger than one parallelise both the initial linearisation and later relinearisations without limiting the inner threading of ordinary time-step GS solves.\n",
     "\n",
     "Other customisable inputs are available, do see the documentation or the later notebook on \"Growth Rates\" for more details. For example, one may explicitly set your own resistance and inductance matrices for the tokamak, rather than the geometrical value calculated internally in FreeGSNKE.\n",
     "\n",
@@ -239,6 +240,7 @@
     "    full_timestep=5e-4, \n",
     "    plasma_resistivity=1e-6,\n",
     "    max_mode_frequency=10**2.5,\n",
+    "    n_linearization_workers=1,  # set >1 to parallelise Jacobian construction\n",
     ")"
    ]
   },
@@ -651,7 +653,8 @@
     "    profiles=profiles, \n",
     "    full_timestep=.5e-3, \n",
     "    plasma_resistivity=1e-6, \n",
-    "    GSStaticSolver=GSStaticSolver   \n",
+    "    GSStaticSolver=GSStaticSolver,\n",
+    "    n_linearization_workers=1,  # set >1 to parallelise Jacobian construction\n",
     ")"
    ]
   },

examples/example05b - linear_evolution_without_GS.ipynb

@@ -171,7 +171,9 @@
    "source": [
     "### Time evolution\n",
     "\n",
-    "Having defined the plasma descriptors, we can now instantiate the evolutive solver object. By including the `plasma_descriptors` function as an argument, the relevant Jacobians will be calculated to enable this evolution. "
+    "Having defined the plasma descriptors, we can now instantiate the evolutive solver object. By including the `plasma_descriptors` function as an argument, the relevant Jacobians will be calculated to enable this evolution.\n",
+    "\n",
+    "The new **`n_linearization_workers`** argument controls the number of independent worker processes used to construct both the current-response and profile-parameter-response Jacobian columns. `1` retains serial construction; a larger value parallelises the initial linearisation and any later relinearisations. It does not limit the inner threading of ordinary time-step GS solves."
    ]
   },
   {
@@ -189,6 +191,7 @@
     "    full_timestep=5e-4, \n",
     "    plasma_resistivity=1e-6,\n",
     "    max_mode_frequency=10**2.5,\n",
+    "    n_linearization_workers=1,  # set >1 to parallelise Jacobian construction\n",
     "    plasma_descriptor_function=plasma_descriptors\n",
     ")"
    ]

examples/example05c - linear_evolution_with_relinearisation.ipynb

@@ -168,10 +168,18 @@
     "    full_timestep=5e-4, \n",
     "    plasma_resistivity=1e-6,\n",
     "    max_mode_frequency=10**2.5,\n",
+    "    n_linearization_workers=1,  # set >1 to parallelise Jacobian construction\n",
     "    plasma_descriptor_function=plasma_descriptors\n",
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The new **`n_linearization_workers`** argument controls how many independent worker processes build both the current-response and profile-parameter-response columns. It applies to the first linearisation constructed by `nl_solver` and to later relinearisations. The default value of `1` retains the serial calculation. Values larger than one use that many CPU threads in total for Jacobian construction, with one native numerical-library thread per worker to avoid oversubscription. This limit is scoped to Jacobian construction; it does not alter the inner threading used by the ordinary GS solves that advance the evolution."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

examples/example11 - pulse_design_tool.ipynb

@@ -972,7 +972,11 @@
     "\n",
     "The simulation timestep was chosen based on the vertical instability timescale in the output below (you need to always choose a value that is 5-10x smaller than this timescale) - this keeps the simulation numerically stable.\n",
     "\n",
-    "This example explicitly uses coupling-based mode selection. The coupling ranking is calculated about the initial equilibrium and can evolve as the plasma moves or changes shape. The 30-mode reduction used here is therefore specific to this trajectory; for larger shape changes, validate the retained modes across representative equilibria or use `mode_selection=\"timescale\"`."
+    "This example explicitly uses coupling-based mode selection. The coupling ranking is calculated about the initial equilibrium and can evolve as the plasma moves or changes shape. The 30-mode reduction used here is therefore specific to this trajectory; for larger shape changes, validate the retained modes across representative equilibria or use `mode_selection=\"timescale\"`.\n",
+    "\n",
+    "The simulation timestep was chosen based on the vertical instability timescale in the output below (you need to always choose a value that is 5-10x smaller than this timescale) - this keeps the simulation numerically stable.\n",
+    "\n",
+    "The new **`n_linearization_workers`** argument controls the number of independent worker processes used to build both the current-response (`dIydI`) and profile-parameter-response (`dIydtheta`) Jacobian columns. `1` retains serial construction; set it higher to parallelise both the initial linearisation and later relinearisations. This setting is scoped to Jacobian construction and does not limit the inner threading of the GS solves used during ordinary evolution steps."
    ]
   },
   {
@@ -992,6 +996,7 @@
     "    # Coupling-ranked reduction, validated for this example trajectory.\n",
     "    mode_selection=\"coupling\",\n",
     "    fix_n_vessel_modes=30,\n",
+    "    n_linearization_workers=4,  # set >1 to parallelise Jacobian construction\n",
     "    plasma_descriptor_function=plasma_descriptors,\n",
     "    )"
    ]