From a045f0d35452afeb03740fd43cac8afb29249ed2 Mon Sep 17 00:00:00 2001
From: lazersos <lazersos@gmail.com>
Date: Mon, 9 Dec 2024 14:35:49 +0100
Subject: [PATCH] pySTEL: Magwell added to STELLOPT and VMECplot routines.

---
 pySTEL/STELLOPT.py | 16 ++++++++++++++++
 pySTEL/VMECplot.py | 11 ++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/pySTEL/STELLOPT.py b/pySTEL/STELLOPT.py
index 79fb3709..617ab5cc 100755
--- a/pySTEL/STELLOPT.py
+++ b/pySTEL/STELLOPT.py
@@ -923,6 +923,7 @@ def LoadSTELLOPT(self):
 			self.ui.ComboBoxOPTplot_type.addItem('q-prof')
 			self.ui.ComboBoxOPTplot_type.addItem('<j*B>')
 			self.ui.ComboBoxOPTplot_type.addItem('Mercier')
+			self.ui.ComboBoxOPTplot_type.addItem('Magwell')
 			wout_files = sorted([k for k in files if 'wout' in k])
 			self.wout_files = sorted([k for k in wout_files if '_opt' not in k])
 		# Handle Boozer Transformation
@@ -1905,6 +1906,21 @@ def UpdateOptplot(self):
 			self.ax2.set_ylabel('[Arb]')
 			self.ax2.set_title('Mercier Stability (>0 Stable)')
 			self.ax2.set_xlim((0,1))
+		elif (plot_name == 'Magwell'):
+			vmec_data = vmec.VMEC()
+			l=0
+			dl = len(self.wout_files)-1
+			for string in self.wout_files:
+				if 'wout' in string:
+					vmec_data.read_wout(self.workdir+string)
+					magwell = vmec_data.calc_magwell()
+					nflux = np.linspace(0.0,1.0,vmec_data.ns)
+					self.ax2.plot(nflux,magwell,color=_plt.cm.brg(l/dl))
+					l=l+1
+			self.ax2.set_xlabel('Norm Tor. Flux (s)')
+			self.ax2.set_ylabel('W')
+			self.ax2.set_title('Magnetic Well/Hill Stability (>0 Well)')
+			self.ax2.set_xlim((0,1))
 		elif (plot_name == 'Flux0'):
 			vmec_data = vmec.VMEC()
 			l=0
diff --git a/pySTEL/VMECplot.py b/pySTEL/VMECplot.py
index 2f095328..386c09b6 100755
--- a/pySTEL/VMECplot.py
+++ b/pySTEL/VMECplot.py
@@ -39,7 +39,8 @@ def __init__(self):
 		self.statusBar().showMessage('Ready')
 		self.ui.plot_list = ['Summary','-----1D-----','Iota','q','Pressure',\
 		'<Buco>','<Bvco>','<jcuru>','<jcurv>','<B.B>','<j.B>',  '-----3D------','|B|','sqrt(g)',\
-		'B^u','B^v','B_s','B_u','B_v','j^u','j^v', 'jll', 'j.B','---Special---','LPK','Mercier']
+		'B^u','B^v','B_s','B_u','B_v','j^u','j^v', 'jll', 'j.B','---Special---','LPK','Mercier',\
+		'Magwell']
 		files = sorted(os.listdir('.'))
 		for name in files:
 			if(name[0:4]=='wout'):
@@ -292,6 +293,14 @@ def update_plot(self,i):
 			self.ax.set_title('Mercier Stability (>0 Stable)')
 			self.ax.set_ylim(-0.1,0.1)
 			self.ax.legend()
+		elif (plot_name == 'Magwell'):
+			magwell = self.vmec_data.calc_magwell()
+			self.ax.plot(self.nflux,magwell)
+			self.ax.set_xlabel('Normalized Flux')
+			self.ax.set_ylabel('W')
+			self.ax.set_title('Magnetic Well/Hill (>0 Well)')
+			#self.ax.set_ylim(-0.1,0.1)
+			self.ax.legend()
 		elif (plot_name == 'LPK'):
 			self.ax.plot(self.r[self.ns-1,:,0],self.z[self.ns-1,:,0],color='red')
 			self.ax.plot(self.r[0,0,0],self.z[0,0,0],'+',color='red')