try:
# framework is running
from .startup_choice import *
except ImportError as _excp:
# class is imported by itself
if (
'attempted relative import with no known parent package' in str(_excp)
or 'No module named \'omfit_classes\'' in str(_excp)
or "No module named '__main__.startup_choice'" in str(_excp)
):
from startup_choice import *
else:
raise
from matplotlib import pyplot
import numpy as np
from omfit_classes.omfit_ascii import OMFITascii
from omfit_classes.omfit_base import OMFITtree
__all__ = [
'OMFITeliteGamma',
'OMFITeliteEigenfunction',
'OMFITelite2DEigenfunction',
'OMFITelitefun2d',
'OMFITelitextopsi',
'OMFITeliteAggregated',
'OMFITeqdat',
'OMFITbalstab',
'OMFIToutbal',
]
[docs]class OMFITeliteGamma(SortedDict, OMFITascii):
"""ELITE growth rates data file (Gamma file)"""
def __init__(self, filename, **kw):
OMFITascii.__init__(self, filename, **kw)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
self.clear()
with open(self.filename, 'r') as f:
lines = f.readlines()
for line in lines[1:]:
tmp = line.split()
nn = int(tmp[0])
items = tmp[1:]
self[nn] = SortedDict()
for k, item in enumerate(['cent_gam/w_A', 'gam/(omegspi_max/4)', 'gamsq', 'del', 'gamr', 'gamr2', 'tmatasym']):
self[nn][item] = float(items[k])
[docs]class OMFITeliteEigenfunction(OMFITascii, SortedDict):
def __init__(self, filename, symmetric=False):
OMFITobject.__init__(self, filename)
SortedDict.__init__(self)
self.symmetric = symmetric
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
lines = self.read().splitlines()
self['date'] = lines[0]
self['version'] = lines[1]
self['eq_type'] = lines[2]
self['n'] = int(lines[3])
self['eq_data'] = np.array(list(map(float, (lines[4] + ' ' + lines[5]).split())))
self['eq_par'] = np.array(list(map(int, (' '.join(lines[6:8])).split())))
self['dummy'] = np.array(list(map(float, (' '.join(lines[8:12])).split())))
tmp = np.array(list(map(float, (' '.join(lines[12:])).split())))
tmp = np.reshape(tmp, (-1, self['eq_par'][9])).T
self['xx'] = tmp[:, 0]
self['psi'] = tmp[:, -1]
if self.symmetric:
self['eigh'] = tmp[:, 1:-1:1] + 1j * np.zeros(tmp[:, 1:-1:1].shape)
else:
self['eigh'] = tmp[:, 1:-2:2] + 1j * tmp[:, 2:-1:2]
[docs] def peak(self, decimate=1):
f = np.zeros(self['eigh'][::decimate].shape)
for k1 in np.linspace(0, 2 * np.pi, 64)[:-1]:
m1 = np.exp(1j * k1 * 2)
f = np.max([f, np.real(self['eigh'][::decimate] * m1)], 0)
return f
[docs] def envelope(self, decimate=1):
return np.max(self.peak(decimate), 1)
[docs] def plot(self, fancy=False, decimate=None):
from omfit_classes.utils_plot import plotc
if decimate is None:
decimate = int(np.ceil(1 + len(self['eigh']) / 10000.0))
if fancy:
tmp = self.peak(decimate)
plotc(self['psi'][::decimate], tmp)
pyplot.plot(self['psi'][::decimate], np.max(tmp, 1), color='k')
else:
y = self.envelope(decimate)
x = self['psi'][::decimate]
tmp = abs(np.trapz(y[np.where(x < 0.98)[0]], x[np.where(x < 0.98)[0]]))
pyplot.plot(x, y / tmp)
pyplot.ylim([1e-8, pyplot.ylim()[-1]])
pyplot.xlabel('$\\psi$')
pyplot.title('Eigenfunction n=%d' % self['eq_par'][5])
[docs]class OMFITelite2DEigenfunction(OMFITtree):
def __init__(self, filename, psicontour=np.asarray([0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1])):
self.dynaLoad = True
self.filename = filename
self.psicontour = psicontour
[docs] @dynaLoad
def load(self):
self['xtopsi'] = OMFITelitextopsi(self.filename + '.xtopsi')
self['fun2d'] = OMFITelitefun2d(self.filename + '.fun2d')
self['xtopsi'].load()
self['fun2d'].load()
if len(self.psicontour) > 0:
self['psiind'] = []
for psi in self.psicontour:
self['psiind'].append(np.argmin(abs(self['xtopsi']['psiinterp'] - psi)))
self['rind'] = []
for i in self['psiind']:
self['rind'].append(
np.argmin(
abs(
np.asarray(self['fun2d']['xx'][1 : self['fun2d']['ns'] * self['fun2d']['nxplt'] : self['fun2d']['ns']])
- self['xtopsi']['rrmax'][i]
)
)
)
[docs] @dynaSave
def save(self):
self['xtopsi'].save()
self['fun2d'].save()
[docs] def plot(self):
pyplot.tricontourf(self['fun2d']['xx'], self['fun2d']['zz'], self['fun2d']['f2d'], 30, cmap='rainbow')
if len(self.psicontour) > 0:
for i in self['rind']:
pyplot.plot(
self['fun2d']['xx'][i * self['fun2d']['ns'] + 1 : (i + 1) * self['fun2d']['ns']],
self['fun2d']['zz'][i * self['fun2d']['ns'] + 1 : (i + 1) * self['fun2d']['ns']],
'k-',
linewidth=0.5,
)
pyplot.axis('equal')
[docs]class OMFITelitefun2d(OMFITascii, SortedDict):
def __init__(self, filename):
OMFITobject.__init__(self, filename)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
lines = self.read().splitlines()
self['cdatetime'] = lines[0]
self['codver'] = lines[1]
self['shape'] = lines[2]
self['npts'] = int(lines[3])
set3 = np.array(list(map(float, (' '.join(lines[4:6])).split())))
set2 = np.array(list(map(int, (' '.join(lines[6:8])).split())))
self['nx'] = set2[9]
self['nm'] = set2[4]
set1 = np.array(list(map(float, (' '.join(lines[8:12])).split())))
set0 = np.array(list(map(int, (lines[12]).split())))
self['nxplt'] = set0[0]
self['ns'] = set0[1]
tmp = np.array(list(map(float, (' '.join(lines[13:])).split())))
tmp = np.reshape(tmp, (-1, self['nxplt'] * self['ns'])).T
self['xx'] = tmp[:, 0] / 100.0
self['zz'] = tmp[:, 1] / 100.0
self['bp'] = tmp[:, 2]
self['f2d'] = tmp[:, 3]
[docs] def plot(self):
from matplotlib import pyplot
pyplot.tricontourf(self['xx'], self['zz'], self['f2d'], 30, cmap='rainbow')
pyplot.axis('equal')
[docs]class OMFITelitextopsi(OMFITascii, SortedDict):
def __init__(self, filename):
OMFITobject.__init__(self, filename)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
lines = self.read().splitlines()
self['nxinterp'] = int(lines[0])
tmp = np.array(list(map(float, (' '.join(lines[1:-1])).split())))
tmp = np.reshape(tmp, (-1, self['nxinterp'])).T
self['xinterp'] = tmp[:, 0]
self['psiinterp'] = tmp[:, 1]
self['rrmax'] = tmp[:, 2] / 100.0
self['rmax'] = tmp[:, 3] / 100.0
[docs]class OMFITeliteAggregated(OMFITascii, SortedDict):
def __init__(self, filename):
OMFITascii.__init__(self, filename)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
with open(self.filename, 'r') as f:
tmp = f.readlines()
dum = ' '.join([line.strip() for line in tmp])
dum = list(map(float, dum.split()))
try:
dum = np.reshape(dum, (-1, 12))
except Exception:
dum = np.reshape(dum, (-1, 11))
# in case the ELITE run did not finish
self['alpha'] = dum[:, 0]
self['pprime'] = dum[:, 1]
self['jmax'] = dum[:, 2]
self['jsep'] = dum[:, 3]
self['gamma'] = dum[:, 4]
self['ngam'] = dum[:, 5].astype('int')
self['gamws'] = dum[:, 6]
self['ngamws'] = dum[:, 7].astype('int')
self['javg'] = dum[:, 8]
self['gamws16'] = dum[:, 9]
self['ngamws16'] = dum[:, 10].astype('int')
if dum.shape[1] == 12:
self['icount'] = dum[:, 11].astype('int')
[docs]class OMFITeqdat(SortedDict, OMFITascii):
r"""
OMFIT class used to interface to equilibria files generated by ELITE and BALOO (.dskbal files)
NOTE: this object is "READ ONLY", meaning that the changes to the entries of this object will not be saved to a file. Method .save() could be written if becomes necessary
:param filename: filename passed to OMFITobject class
:param \**kw: keyword dictionary passed to OMFITobject class
"""
def __init__(self, filename='', **kw):
OMFITascii.__init__(self, filename, **kw)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
with open(self.filename, 'r') as f:
lines = f.readlines()
for k, line in enumerate(lines):
if 'npsi,nthet' in line:
break
if k == len(lines) - 1:
raise ValueError('eqdat file should have a line with `npsi,nthet` in it')
k += 1
self['npsi'], self['npts'] = list(map(int, lines[k].split()))
lsegment = int(np.ceil(self['npsi'] / 5.0))
while k < len(lines) - 1:
k += 1
name = lines[k].strip()
self[name] = []
if name not in ['X', 'Z', 'poloidal field--bpol']:
for seg in range(lsegment):
k += 1
self[name].extend(re.sub('D', 'E', lines[k]).split())
self[name] = np.flipud(np.array(list(map(eval, self[name]))))
else:
for x in range(self['npts']):
self[name].append([])
for seg in range(lsegment):
k += 1
self[name][-1].extend(re.sub('D', 'E', lines[k]).split())
self[name][-1] = np.array(list(map(eval, self[name][-1])))
self[name] = np.array(self[name])
for name in self:
if isinstance(self[name], np.ndarray) and len(self[name].shape) == 1:
self[name] = self[name][::-1]
[docs] @dynaLoad
def plot(self, **kw):
from omfit_classes.utils_plot import title_inside, autofmt_sharex
items = [item for item in self if item not in ['npsi', 'npts', 'X', 'Z', 'poloidal field--bpol']]
ax = None
r = np.sqrt(len(items) + 1)
c = np.ceil((len(items)) / np.sqrt(len(items))) + 1
for k, item in enumerate(items):
x = self.get('poloidal flux--psiv', np.arange(len(self[item])))
x = (x - min(x)) / (max(x) - min(x))
ax = pyplot.subplot(r, c, k + 1 + k // (c - 1) + 1)
ax.plot(x, self[item], **kw)
title_inside(item, y=0.8)
ax.set_xlim([0, 1])
nskip = self['npsi'] // 10
ax0 = ax
ax = pyplot.subplot(1, c, 1)
ax.plot(self['X'][:, ::-nskip] / 100.0, self['Z'][:, ::-nskip] / 100.0, color=ax0.lines[-1].get_color())
ax.set_aspect('equal')
ax.set_frame_on(False)
autofmt_sharex()
[docs]class OMFITbalstab(SortedDict, OMFITascii):
r"""
OMFIT class used to interface to `balstab` file from BALOO
NOTE: this object is "READ ONLY", meaning that the changes to the entries of this object will not be saved to a file. Method .save() could be written if becomes necessary
:param filename: filename passed to OMFITobject class
:param \**kw: keyword dictionary passed to OMFITobject class
"""
def __init__(self, filename='', **kw):
OMFITascii.__init__(self, filename, **kw)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
with open(self.filename, 'r') as f:
tmp = f.read().strip()
tmp = re.sub('% ', '%_', tmp)
tmp = re.sub(' s ', ' 0 ', tmp)
tmp = re.sub(' u ', ' 1 ', tmp)
tmp = tmp.split('\n')
numbers = list(map(int, tmp[0].split()))
headers = tmp[1].split()
self.clear()
for k in range(2, numbers[0] + 2):
res = list(map(float, tmp[k].split()))
res[0] = int(res[0])
res[4] = int(res[4])
self[res[0]] = {}
for k1, h in enumerate(headers):
if k1 == 0:
continue
self[res[0]][h] = res[k1]
[docs]class OMFIToutbal(SortedDict, OMFITascii):
r"""
OMFIT class used to interface to `outbal` file from BALOO
NOTE: this object is "READ ONLY", meaning that the changes to the entries of this object will not be saved to a file. Method .save() could be written if becomes necessary
:param filename: filename passed to OMFITobject class
:param \**kw: keyword dictionary passed to OMFITobject class
"""
def __init__(self, filename='', **kw):
OMFITascii.__init__(self, filename, **kw)
SortedDict.__init__(self)
self.dynaLoad = True
[docs] @dynaLoad
def load(self):
with open(self.filename, 'r') as f:
lines = f.read().strip().split('\n')
if 'PARZERO' in lines[0]:
lines = lines[1:] # skip loading this line if there is a parzero error
tmp = []
for k in range(int(len(lines) // 2)):
tmp.append(lines[k * 2] + lines[k * 2 + 1])
names = tmp[0].split()
tmp = tmp[1:]
for k, line in enumerate(tmp):
tmp[k] = list(map(float, line.split()))
tmp = np.array(tmp)
for k, name in enumerate(names):
self[name] = tmp[:, k]
[docs] @dynaLoad
def plot(self, **kw):
from omfit_classes.utils_plot import title_inside
items = [item for item in self if item not in ['surf', 'psi_norm']]
r = np.sqrt(len(items) + 1)
c = np.ceil((len(items)) / np.sqrt(len(items)))
for k, item in enumerate(items):
x = self.get('psi_norm', np.arange(len(self[item])))
x = (x - min(x)) / (max(x) - min(x))
ax = pyplot.subplot(r, c, k + 1)
ax.plot(x, self[item], **kw)
title_inside(item, y=0.8)
ax.set_xlim([0, 1])
