Merge pull request #80 from PrincetonUniversity/cleanupJypterAndPython

Python tools

Utilities/pythontools/plot_modB

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# coding: utf-8
+"""
+Plot of |B| in (R,Z) cross-section
+
+@author Ksenia Aleynikova [email protected]
+"""
+
+import numpy as np
+import py_spec as p
+import sys
+import matplotlib.pyplot as plt
+
+s = p.SPEC(sys.argv[1])
+
+sarr    = np.linspace(-0.999,1,4)
+tarr    = np.linspace(0,2*np.pi,5)
+zarr    = np.linspace(0,0,6)
+lvol    = 0
+
+R, Z, jacobian, g = p.get_grid_and_jacobian_and_metric(s,lvol=lvol,sarr=sarr,tarr=tarr,zarr=zarr)
+Bcontrav          = p.get_B(s,lvol=lvol,jacobian=jacobian,sarr=sarr,tarr=tarr,zarr=zarr)
+modB              = p.get_modB(Bcontrav,g)
+
+fig = plt.figure(figsize=(3*2,5*2))
+ax = fig.add_subplot(111)
+plot = ax.pcolormesh(R[:,:,0],Z[:,:,0],modB[:,:,0], vmin=2.8, vmax=3.5)#,linewidths=0.01/3,edgecolors='y')
+ax.set_aspect('equal')
+ax.set_xlabel(r"$R$")
+ax.set_ylabel(r"$Z$")
+cbar = fig.colorbar(plot, ax=ax)
+cbar.set_label(r"$|\vec B|$", rotation=0)
+plt.show()

Utilities/pythontools/plot_spec_poincare

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Plot the Poincare data from a SPEC output file
+
+@author: Jonathan Schilling ([email protected])
+"""
+
+# fancy command line argment parsing and help texts
+import argparse
+parser = argparse.ArgumentParser(description='Plot the Poincare data from a SPEC output file.')
+
+parser.add_argument('name', \
+                    default='', type=str, help="name of the SPEC output file, e.g. G3V01L1Fi.h5",\
+                    metavar='filename.h5')
+
+parser.add_argument('-t', \
+                    default=0, type=int, help="toroidal cutplane at which to plot the Poincare data",\
+                    metavar='0')
+
+args = parser.parse_args()
+
+# import all these only after successful parsing of command line arguments, since loading these takes a significant amount of time....
+import py_spec as p
+import matplotlib.pyplot as plt
+import os # for os.path.basename
+
+# read the data
+s = p.SPEC(args.name)
+
+# create a basic plot
+plt.figure()
+
+p.plot_poincare(s, args.t, plt)
+plt.title(os.path.basename(args.name))
+
+plt.tight_layout()
+
+plt.show()

Utilities/pythontools/py_spec/__init__.py

@@ -0,0 +1,8 @@
+# import of all SPEC-related python scripts. 
+
+from .read_spec import SPEC
+from .proc_spec import *
+del(read_spec)
+del(proc_spec)
+
+