#########################
# gs2_correlation #
# Ferdinand van Wyk #
#########################
###############################################################################
# This file is part of gs2_correlation.
#
# gs2_correlation_analysis is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# gs2_correlation is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gs2_correlation.
# If not, see <http://www.gnu.org/licenses/>.
###############################################################################
"""
.. class:: Simulation
:platform: Unix
:synopsis: Class describing the GS2 simulation.
.. moduleauthor:: Ferdinand van Wyk <ferdinandvwyk@gmail.com>
"""
# Standard
import os
import sys
import gc
import configparser
import logging
import operator
import warnings
import json
# Third Party
import numpy as np
from netCDF4 import Dataset
import scipy.interpolate as interp
import scipy.integrate as integrate
import scipy.optimize as opt
import scipy.signal as sig
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from mpl_toolkits.axes_grid1 import make_axes_locatable
import f90nml as nml
import lmfit as lm
import pyfftw
from progressbar import ProgressBar, Percentage, Bar
plt.rcParams.update({'figure.autolayout': True})
mpl.rcParams['axes.unicode_minus'] = False
pal = sns.color_palette('deep')
# Local
import gs2_correlation.fitting_functions as fit
import gs2_correlation.plot_style as plot_style
[docs]class Simulation(object):
"""
Class containing all simulation information.
The class mainly reads from the simulation NetCDF file and operates on the
field specified in the configuration file, such as performing correlations,
FFTs, plotting, etc.
An exhaustive list of instance variables is provided separately in the
documentation since it is very long.
"""
[docs] def __init__(self, config_file):
"""
Initialized by object using information from configuration file.
Initialization does the following, depending on the parameters in the
configuration file:
* Calculates sizes of kx, ky, x, and y.
* Reads data from the NetCDF file.
* Interpolates onto a regular time grid.
* Zeros out BES scales.
* Zeros out ZF scales.
* Transforms to lab frame.
* Calculates the real space field.
* Reduce domain size to according to *box_size* if domain is specified
as 'middle', otherwise do nothing.
* Ensures real space field has odd points.
Parameters
----------
config_file : str
Filename of configuration file and path if not in the same
directory.
Notes
-----
The configuration file should contain the following namelists:
* 'general': information such as the analysis to be performed and
where to find the NetCDF file.
* 'perp' : parameters relating to the perpendicular correlation analysis.
* 'time' : parameters relating to the time analysis.
* 'par' : parameters relating to the parallel analysis.
* 'normalization': normalization parameters for the simulation/experiment.
* 'output': parameters which relate to the output produced by the code.
The exact parameters read in are documented in the documentation.
"""
self.config_file = config_file
self.read_config()
# Set plot options
sns.set_context(self.seaborn_context)
self.read_input_file()
self.read_geometry_file()
self.read_netcdf()
if self.out_dir not in os.listdir():
os.system("mkdir -p " + self.out_dir)
self.nt = len(self.t)
self.nkx = len(self.kx)
self.nky = len(self.ky)
self.nx = self.nkx
self.ny = 2*(self.nky - 1)
self.ntheta = self.field.shape[3]
self.config_checks()
self.R = self.geometry[:,1]*self.amin
self.Z = self.geometry[:,2]*self.amin
self.alpha = self.geometry[:,3]
self.tor_phi = -self.geometry[:,3]
self.dR_drho = self.geometry[:,4]*self.amin
self.dZ_drho = self.geometry[:,5]*self.amin
self.dalpha_drho = self.geometry[:,6]*self.amin
self.geo_ntheta = len(self.R)
self.rho_star = self.rho_ref/self.amin
self.rmaj = self.R[int(self.geo_ntheta/2)]
self.pitch_angle = np.arctan(self.Z[int(self.geo_ntheta/2)+1]/ \
(self.R[int(self.geo_ntheta/2)] * \
self.tor_phi[int(self.geo_ntheta/2)+1]))
self.t = self.t*self.amin/self.vth
# Determine the box size in terms of flux surface label rho
self.n0 = int(np.around(self.ky[1]*(self.amin/self.rho_ref)))
delta_rho = (self.rho_tor/self.qinp) * (self.jtwist/(self.n0*self.shat))
self.x_box_size = (self.r_prime[int(self.geo_ntheta/2)] * delta_rho *
self.amin)
self.x = np.linspace(0, self.x_box_size, self.nx, endpoint=False)
self.y_tor_box_size = self.rmaj * 2 * np.pi / self.n0
self.y_pol_box_size = self.y_tor_box_size * np.tan(self.pitch_angle)
self.y_perp_box_size = self.y_tor_box_size * np.sin(self.pitch_angle)
self.y = np.linspace(0, self.y_perp_box_size, self.ny, endpoint=False)
self.btor = self.bref*self.r_geo/self.R[int(self.ntheta/2)]
self.bmag = np.sqrt(self.btor**2 + self.bpol**2)
self.field_to_complex()
self.fourier_correction()
if self.time_interpolate_bool or self.lab_frame:
self.time_interpolate()
self.nt_slices = int(self.nt/self.time_slice)
if self.zero_bes_scales_bool:
self.zero_bes_scales()
if self.zero_zf_scales_bool:
self.zero_zf_scales()
if self.lab_frame:
self.to_lab_frame()
self.field_to_real_space()
if self.domain == 'middle':
self.domain_reduce()
self.field_odd_pts()
if self.analysis != 'par' and self.analysis != 'write_field_full':
self.field_real_space = self.field_real_space[:,:,:,0]
[docs] def find_file_with_ext(self, ext):
"""
Find a file in the run_folder with the extension ext
Parameters
----------
ext : str
Extension of the file to be searched for. Of the form '.ext'.
"""
dir_files = os.listdir(self.run_folder)
found = False
for s in dir_files:
if s.endswith(ext):
found_file = self.run_folder + s
found = True
break
if found:
return(found_file)
else:
raise NameError('No file found ending in ' + ext)
[docs] def read_config(self):
"""
Reads analysis and normalization parameters from self.config_file.
The full list of possible configuration parameters is provided
separately in the documentation since it is quite long and are
technically not parameters of this function.
"""
logging.info('Started read_config...')
config_parse = configparser.ConfigParser()
config_parse.read(self.config_file)
##########################
# Normalization Namelist #
##########################
self.amin = float(config_parse['normalization']['a_minor'])
self.bref = float(config_parse['normalization']['bref'])
self.dpsi_da = float(config_parse.get('normalization', 'dpsi_da',
fallback=0))
self.nref = float(config_parse.get('normalization', 'nref',
fallback=1))
self.omega = float(config_parse.get('normalization', 'omega',
fallback=0))
self.rho_ref = float(config_parse['normalization']['rho_ref'])
self.rho_tor = float(config_parse['normalization']['rho_tor'])
self.tref = float(config_parse.get('normalization', 'tref',
fallback=1))
self.vth = float(config_parse['normalization']['vth_ref'])
####################
# General Namelist #
####################
self.domain = config_parse.get('general', 'domain',
fallback='full')
self.file_ext = '.out.nc'
self.run_folder = str(config_parse['general']['run_folder'])
self.cdf_file = config_parse.get('general', 'cdf_file', fallback='None')
if self.cdf_file == "None":
self.cdf_file = self.find_file_with_ext(self.file_ext)
if self.domain == 'full':
self.out_dir = self. run_folder + 'correlation_analysis/full'
elif self.domain == 'middle':
self.out_dir = self.run_folder + 'correlation_analysis/middle'
self.out_dir = config_parse.get('general', 'out_dir',
fallback=self.out_dir)
self.g_file = config_parse.get('general', 'g_file', fallback='None')
if self.g_file == 'None':
self.g_file = self.find_file_with_ext('.g')
self.extract_input_file()
if 'input_file.in' in os.listdir(self.run_folder):
self.in_file = self.run_folder + 'input_file.in'
else:
NameError('Could not extract input file from NetCDF file. '
'Make sure GS2 is using new diagnostic output.')
self.in_field = str(config_parse['general']['field'])
self.analysis = config_parse.get('general', 'analysis',
fallback='all')
if self.analysis not in ['all', 'perp', 'par', 'time', 'write_field',
'write_field_full']:
raise ValueError('Analysis must be one of (perp, time, par, '
'write_field, write_field_full)')
self.time_interpolate_bool = config_parse.getboolean('general',
'time_interpolate',
fallback=True)
self.time_interp_fac = int(config_parse.get('general',
'time_interp_fac',
fallback=1))
self.zero_bes_scales_bool = config_parse.getboolean('general',
'zero_bes_scales', fallback=False)
self.zero_zf_scales_bool = config_parse.getboolean('general',
'zero_zf_scales', fallback=False)
self.lab_frame = config_parse.getboolean('general',
'lab_frame', fallback=False)
self.spec_idx = str(config_parse['general']['species_index'])
if self.spec_idx == "None":
self.spec_idx = None
else:
self.spec_idx = int(self.spec_idx)
self.theta_idx = str(config_parse['general']['theta_index'])
if self.theta_idx == "None":
self.theta_idx = None
elif self.theta_idx == "-1":
self.theta_idx = [0, None]
elif type(eval(self.theta_idx)) == int:
self.theta_idx = [int(self.theta_idx), int(self.theta_idx)+1]
else:
raise ValueError("theta_idx can only be one of: None, -1(all), int")
if ((self.analysis == 'par' or self.analysis == 'write_field_full') and
self.theta_idx != [0, None]):
warnings.warn('Analysis = par/write_field_full, change to reading '
'full theta grid.')
self.theta_idx = [0, None]
self.time_slice = int(config_parse.get('general', 'time_slice',
fallback=49))
self.box_size = str(config_parse.get('general',
'box_size', fallback='[0.2,0.2]'))
self.box_size = self.box_size[1:-1].split(',')
self.box_size = [float(s) for s in self.box_size]
self.time_range = str(config_parse.get('general',
'time_range', fallback='[0,-1]'))
self.time_range = self.time_range[1:-1].split(',')
self.time_range = [int(s) for s in self.time_range]
if self.time_range[1] == -1:
self.time_range[1] = None
#################
# Perp Namelist #
#################
self.ky_free = config_parse.getboolean('perp','ky_free', fallback=False)
perp_guess = str(config_parse.get('perp',
'perp_guess', fallback='[0.05,0.1,1]'))
perp_guess = perp_guess[1:-1].split(',')
perp_guess = [float(s) for s in perp_guess]
self.perp_guess_x = perp_guess[0]
self.perp_guess_y = perp_guess[1]
if self.ky_free:
try:
self.perp_guess_ky = perp_guess[2]
except IndexError:
warnings.warn('ky guess not specified in perp_guess, '
'setting perp_guess_ky = 1')
self.perp_guess_ky = 1.0
#################
# Time Namelist #
#################
self.npeaks_fit = int(config_parse.get('time',
'npeaks_fit', fallback=5))
self.time_guess = str(config_parse.get('time',
'time_guess', fallback='[1e-5,100]'))
self.time_guess = self.time_guess[1:-1].split(',')
self.time_guess = [float(s) for s in self.time_guess]
self.time_guess_dec = self.time_guess[0]
self.time_guess_grow = self.time_guess[0]
self.time_guess_osc = np.array([self.time_guess[0], self.time_guess[1],
0.0])
self.time_max = float(config_parse.get('time',
'time_max', fallback=1))
################
# Par Namelist #
################
self.par_guess = str(config_parse.get('par',
'par_guess', fallback='[1,0.1]'))
self.par_guess = self.par_guess[1:-1].split(',')
self.par_guess = [float(s) for s in self.par_guess]
###################
# Output Namelist #
###################
self.seaborn_context = str(config_parse.get('output', 'seaborn_context',
fallback='talk'))
self.write_field_interp_x = config_parse.getboolean('output',
'write_field_interp_x',
fallback=True)
# Log the variables
logging.info('The following values were read from ' + self.config_file)
logging.info(vars(self))
logging.info('Finished read_config.')
[docs] def config_checks(self):
"""
This function contains consistency checks of configurations parameters.
"""
if self.time_slice%2 != 1:
warnings.warn('time_slice should be odd, reducing by one...')
self.time_slice -= 1
if self.lab_frame and self.omega == 0:
warnings.warn('Changing to lab frame but omega = 0 (default).')
if self.lab_frame and self.dpsi_da == 0:
warnings.warn('Changing to lab frame but dpsi_da = 0 (default).')
if self.lab_frame and self.time_interp_fac == 1:
warnings.warn('Transforming to lab frame, but time_interp_fac = 1. '
'This is probably not high enough. Recommend 4.')
if not self.lab_frame and self.time_interp_fac > 1:
warnings.warn('Not transforming to lab frame, but time_interp_fac > 1. '
'This is probably not needed.')
if self.theta_idx == None and self.in_field[-2:] == '_t':
raise ValueError('You have specified a field with theta info but '
'left theta_idx=None. Specify theta_idx as -1 '
'for full theta info or pick a specific theta.')
if ((self.analysis != 'par' and self.analysis != 'write_field_full') and
self.ntheta > 1):
raise ValueError('You are not doing parallel analysis and have more '
'than one theta point. Can only handle one theta '
'value at the moment!')
if self.analysis == 'perp' and self.zero_zf_scales_bool == False:
warnings.warn('Doing perp analysis but not zeroing ZF scales. This '
'is required for radial correlation. Changing '
'zero_zf_scales_bool to True')
self.zero_zf_scales_bool = True
[docs] def read_netcdf(self):
"""
Read array from NetCDF file.
Read array specified in configuration file as 'in_field'. Function
uses information from the configuration object passed to it.
Notes
-----
* An additional axis is added in the case of no theta_idx. This allows
usual loops in theta to be kept but have no effect. If only one
element in dimension initialization will be performed regardless,
however dimension will be removed for perp and time analysis.
"""
logging.info('Start reading from NetCDf file...')
with Dataset(self.cdf_file, 'r') as ncfile:
# NetCDF order is [t, species, ky, kx, theta, r]
if self.theta_idx == None:
self.field = np.array(ncfile.variables[self.in_field]
[self.time_range[0]:self.time_range[1],
self.spec_idx,:,:,:])
else:
self.field = np.array(ncfile.variables[self.in_field]
[self.time_range[0]:self.time_range[1],
self.spec_idx,:,:,
self.theta_idx[0]:self.theta_idx[1],
:])
self.t = np.array(ncfile.variables['t'][self.time_range[0]:
self.time_range[1]])
self.field = np.squeeze(self.field)
self.field = np.swapaxes(self.field, 1, 2)
if len(self.field.shape) < 5:
self.field = self.field[:,:,:,np.newaxis,:]
self.drho_dpsi = float(ncfile.variables['drhodpsi'][:])
self.kx = np.array(ncfile.variables['kx'][:])/self.drho_dpsi
self.ky = np.array(ncfile.variables['ky'][:])/self.drho_dpsi
self.theta = np.array(ncfile.variables['theta'][:])
self.gradpar = np.array(ncfile.variables['gradpar'][:])/self.amin
self.r_prime = np.array(ncfile.variables['Rprime'][:])
try:
self.bpol = np.array(ncfile.variables['bpol'][:])*self.bref
except KeyError:
self.bpol = self.geometry[:,7]*self.bref
logging.info('Finished reading from NetCDf file.')
[docs] def read_geometry_file(self):
"""
Read the geometry file for the GS2 run.
"""
logging.info('Reading geometry file...')
self.geometry = np.loadtxt(self.g_file)
logging.info('Finished reading geometry file.')
[docs] def field_to_complex(self):
"""
Converts field to a complex array.
Field is in the following format: field[t, kx, ky, ri] where ri
represents a dimension of length 2.
* ri = 0 - Real part of the field.
* ri = 1 - Imaginary part of the field.
"""
self.field = self.field[:,:,:,:,0] + 1j*self.field[:,:,:,:,1]
[docs] def fourier_correction(self):
"""
Correction to GS2s Fourier components to regular Fourier components.
Notes
-----
GS2 fourier components are NOT simple fourier components obtained from
regular FFT functions. Due to GS2's legacy as a linear code, instead
of just getting fourier components g_k the output is:
G_k = {g_k for ky = 0, 2g_k for ky > 0}
Therfore converting to regular fourier components simply means dividing
all non-zonal components by 2.
"""
self.field[:,:,1:,:] = self.field[:,:,1:,:]/2
[docs] def time_interpolate(self):
"""
Interpolates in time onto a regular grid
Depending on whether the user specified time_interpolate, the time grid
is interpolated into a regular grid. This is required in order to do
FFTs in time. Interpolation is done by default if not specified.
time_interp_fac sets the multiple of interpolation.
"""
logging.info('Started interpolating onto a regular time grid...')
t_reg = np.linspace(min(self.t), max(self.t), self.time_interp_fac*self.nt)
tmp_field = np.empty([self.time_interp_fac*self.nt, self.nkx, self.nky,
self.ntheta], dtype=complex)
for ikx in range(self.nkx):
for iky in range(self.nky):
for ith in range(self.ntheta):
f = interp.interp1d(self.t, self.field[:, ikx, iky, ith])
tmp_field[:, ikx, iky, ith] = f(t_reg)
self.t = t_reg
self.nt = len(self.t)
self.field = tmp_field
tmp_field = None
f = None
gc.collect()
logging.info('Finished interpolating onto a regular time grid.')
[docs] def zero_bes_scales(self):
"""
Sets modes larger than the BES to zero.
The BES is approximately 160x80mm(rad x pol), so we would set kx < 0.25
and ky < 0.5 to zero, since k = 2 pi / L.
"""
for ikx in range(self.nkx):
for iky in range(self.nky):
for ith in range(self.ntheta):
# Roughly the size of BES (160x80mm)
if abs(self.kx[ikx]) < 0.25 and self.ky[iky] < 0.5:
self.field[:,ikx,iky,ith] = 0.0
[docs] def zero_zf_scales(self):
"""
Sets zonal flow (ky = 0) modes to zero.
"""
self.field[:,:,0,:] = 0.0
[docs] def to_lab_frame(self):
"""
Transforms from the rotating frame to the lab frame.
Notes
-----
The important thing here is that ky is NOT the toroidal wavenumber *n0*.
It is related to n0 by [1]_:
ky_gs2 = n0*(rho_ref/a_min)*dpsi_da
.. [1] C. M. Roach, "Equilibrium flow shear implementation in GS2",
http://gyrokinetics.sourceforge.net/wiki/index.php/Documents,
http://svn.code.sf.net/p/gyrokinetics/code/wikifiles/CMR/ExB_GS2.pdf
"""
for ix in range(self.nkx):
for iy in range(self.nky):
for ith in range(self.ntheta):
self.field[:,ix,iy,ith] = self.field[:,ix,iy,ith] * \
np.exp(1j * self.n0 * iy *
self.omega * self.t)
[docs] def field_to_real_space(self):
"""
Converts field from (kx, ky) to (x, y) and saves as new array attribute.
Notes
-----
* Since python defines x = IFFT[FFT(x)] need to undo the implicit
normalization by multiplying by the size of the arrays.
* GS2 fluctuations are O(rho_star) and must be multiplied by rho_star
to get their true values.
* In order to avoid memory overloads, the fourier space field is
cleared after the real space field is calculated.
"""
logging.info('Calculating real space field...')
self.field_real_space = np.empty([self.nt,self.nx,self.ny,self.ntheta],
dtype=float)
pyfftw.n_byte_align(self.field, 16)
self.field_real_space = pyfftw.interfaces.numpy_fft.irfft2(self.field,
axes=[1,2])
if self.analysis == 'par' or self.analysis == 'write_field_full':
self.field = None
gc.collect()
self.field_real_space = np.roll(self.field_real_space,
int(self.nx/2), axis=1)
self.field_real_space = self.field_real_space*self.nx*self.ny
self.field_real_space = self.field_real_space*self.rho_star
logging.info('Finished calculating real space field.')
[docs] def domain_reduce(self):
"""
Initialization consists of:
* Calculating radial and poloidal coordinates r, z.
* Using input parameter *box_size* to determine the index range to
perform the correlation analysis on.
* Reduce extent of real space field using this index.
* Recalculate some real space arrays such as x, y, dx, dy, etc.
"""
logging.info('Reducing domain size to %f x %f m'%(self.box_size[0],
self.box_size[1]))
# Switch box size to length either side of 0
self.box_size = np.array(self.box_size)/2
# Calculate coords r, z
self.r = self.x[:] - self.x[-1]/2 + self.rmaj
self.z = self.y[:] - self.y[-1]/2
# Find index range
r_min_idx, r_min = min(enumerate(abs(self.r - (self.box_size[0] +
self.rmaj))),
key=operator.itemgetter(1))
r_box_idx = r_min_idx-int(self.nx/2) + 1
z_min_idx, z_min = min(enumerate(abs(self.z - self.box_size[1])),
key=operator.itemgetter(1))
z_box_idx = z_min_idx-int(self.ny/2) + 1
# Reduce extent
self.r = self.r[int(self.nx/2)-r_box_idx+1:int(self.nx/2)+r_box_idx]
self.z = self.z[int(self.ny/2)-z_box_idx+1:int(self.ny/2)+z_box_idx]
self.field_real_space = self.field_real_space[
:,int(self.nx/2)-r_box_idx+1:int(self.nx/2)+r_box_idx,
int(self.ny/2)-z_box_idx+1:int(self.ny/2)+z_box_idx,:]
# Recalculate real space arrays
self.nx = len(self.r)
self.ny = len(self.z)
self.x = np.linspace(0, self.r[-1] - self.r[0], self.nx)
self.y = np.linspace(0, 2*self.z[-1], self.ny)
logging.info('Finished reducing domain size.')
[docs] def field_odd_pts(self):
"""
Ensures real space field has odd number of points in x and y.
This is done so that when ``sig.fftconvolve`` is called on the field
with the 'same' option, the resulting correlation function also has an
odd number of points. If this wasn't done, one axis *might* have an
even number of points meaning the correlation function will not have a
convenient middle point to use as the zero separation point, which
would then require more general fitting functions. This way we can
force dx=0, dy=0 at the self.nx/2, self.ny/2 location and avoid extra
parameters in the fitting functions.
"""
logging.info('Ensuring field has odd points in space')
if self.nx%2 != 1:
self.field_real_space = self.field_real_space[:,:-1,:,:]
self.x = self.x[:-1]
if self.ny%2 != 1:
self.field_real_space = self.field_real_space[:,:,:-1,:]
self.y = self.y[:-1]
self.nx = len(self.x)
self.ny = len(self.y)
self.dx = np.linspace(-self.x[-1]/2, self.x[-1]/2, self.nx)
self.dy = np.linspace(-self.y[-1]/2, self.y[-1]/2, self.ny)
[docs] def perp_analysis(self):
"""
Performs a perpendicular correlation analysis on the field.
Notes
-----
* Remove theta dimension before doing analysis to avoid pointless code
which accounts for a theta information. Will produce error if more
than one element during config_checks.
* Uses sig.fftconvolve to calculate the 2D perp correlation function.
* Splits correlation function into time slices and fits each time
slice with a tilted Gaussian using the perp_fit function.
* The fit parameters for the previous time slice is used as the initial
guess for the next time slice.
* Also writes information on the mean fluctuation levels
"""
logging.info('Start perpendicular correlation analysis...')
if not self.ky_free:
self.perp_dir = 'perp/ky_fixed'
else:
self.perp_dir = 'perp/ky_free'
if self.perp_dir not in os.listdir(self.out_dir):
os.system("mkdir -p " + self.out_dir + '/' + self.perp_dir)
if 'corr_fns_x' not in os.listdir(self.out_dir + '/' + self.perp_dir):
os.system("mkdir -p " + self.out_dir+'/'+self.perp_dir+'/corr_fns_x')
if 'corr_fns_y' not in os.listdir(self.out_dir + '/' + self.perp_dir):
os.system("mkdir -p " + self.out_dir+'/'+self.perp_dir+'/corr_fns_y')
os.system('rm -f ' + self.out_dir + '/' + self.perp_dir + '/corr_fns_x/*')
os.system('rm -f ' + self.out_dir + '/' + self.perp_dir + '/corr_fns_y/*')
self.field_normalize_perp()
self.calculate_perp_corr()
self.perp_norm_mask()
self.perp_fit_x = np.empty([self.nt_slices], dtype=float)
self.perp_fit_x_err = np.empty([self.nt_slices], dtype=float)
self.perp_fit_y = np.empty([self.nt_slices], dtype=float)
self.perp_fit_y_err = np.empty([self.nt_slices], dtype=float)
if self.ky_free:
self.perp_fit_ky = np.empty([self.nt_slices], dtype=float)
self.perp_fit_ky_err = np.empty([self.nt_slices], dtype=float)
pbar = ProgressBar(widgets=['Progress: ', Percentage(), Bar()])
for it in pbar(range(self.nt_slices)):
self.perp_corr_fit(it)
self.perp_analysis_summary()
logging.info('Finished perpendicular correlation analysis.')
[docs] def field_normalize_perp(self):
"""
Defines normalized field for the perpandicular correlation by
subtracting the mean and dividing by the RMS value.
"""
logging.info('Normalizing the real space field...')
self.field_real_space_norm_x = \
np.empty([self.nt,self.nx,self.ny],dtype=float)
self.field_real_space_norm_y = \
np.empty([self.nt,self.nx,self.ny],dtype=float)
for it in range(self.nt):
for iy in range(self.ny):
self.field_real_space_norm_x[it,:,iy] = \
self.field_real_space[it,:,iy] - \
np.mean(self.field_real_space[it,:,iy])
self.field_real_space_norm_x[it,:,iy] /= \
np.std(self.field_real_space_norm_x[it,:,iy])
for ix in range(self.nx):
self.field_real_space_norm_y[it,ix,:] = \
self.field_real_space[it,ix,:] - \
np.mean(self.field_real_space[it,ix,:])
self.field_real_space_norm_y[it,ix,:] /= \
np.std(self.field_real_space_norm_y[it,ix,:])
logging.info('Finished normalizing the real space field.')
[docs] def calculate_perp_corr(self):
"""
Calculates the perpendicular correlation function from the real space
field.
"""
logging.info("Calculating perpendicular correlation function...")
self.perp_corr_x = np.empty([self.nt, self.nx, self.ny], dtype=float)
self.perp_corr_y = np.empty([self.nt, self.nx, self.ny], dtype=float)
for it in range(self.nt):
for iy in range(self.ny):
self.perp_corr_x[it,:,iy] = \
sig.correlate(self.field_real_space_norm_x[it,:,iy],
self.field_real_space_norm_x[it,:,iy],
mode='same')
for ix in range(self.nx):
self.perp_corr_y[it,ix,:] = \
sig.correlate(self.field_real_space_norm_y[it,ix,:],
self.field_real_space_norm_y[it,ix,:],
mode='same')
logging.info("Finished calculating perpendicular correlation "
"function...")
[docs] def perp_norm_mask(self):
"""
Applies the appropriate normalization to the perpendicular correlation
function.
Notes
-----
After calling ``sig.correlate`` to calculate ``perp_corr``, we are
left with an unnormalized correlation function as a function of dx
or dy. This function applies a nomalization mask to ``perp_corr_x`` and
``perp_corr_y`` which is dependent on the number of points that
``field_real_space_norm`` has in common with itself for a given dx or dy.
``field_real_space_norm`` is already normalized to the standard
deviation of the time signal, so the only difference between correlate
and np.corrcoef is the number of points in common in the convolution
(that aren't the zero padded values and after averaging over many time
steps).
"""
logging.info('Applying perp normalization mask...')
x = np.ones([self.nx])
y = np.ones([self.ny])
mask_x = sig.correlate(x,x,'same')
mask_y = sig.correlate(y,y,'same')
self.perp_corr_x /= mask_x[np.newaxis, :, np.newaxis]
self.perp_corr_y /= mask_y
logging.info('Finised applying perp normalization mask...')
[docs] def perp_corr_fit(self, it):
"""
Fits the appropriate Gaussian to the radial and poloidal correlation
functions.
Parameters
----------
it : int
This is the index of the time slice currently being fitted.
Notes
-----
* The radial correlation function is fitted with a Gaussian.
* The poloidal correlation function is fitted with an oscillating
Gaussian.
"""
corr_fn_x = \
np.array(self.perp_corr_x[it*self.time_slice:(it+1)*self.time_slice,:,:])
corr_fn_y = \
np.array(self.perp_corr_y[it*self.time_slice:(it+1)*self.time_slice,:,:])
# Average corr_fn over time
corr_std_x = np.empty([self.nx])
corr_std_y = np.empty([self.ny])
for ix in range(self.nx):
corr_std_x[ix] = np.std(corr_fn_x[:,ix,:])
for iy in range(self.ny):
corr_std_y[iy] = np.std(corr_fn_y[:,:,iy])
avg_corr_x = np.mean(np.mean(corr_fn_x, axis=0), axis=1)
avg_corr_y = np.mean(np.mean(corr_fn_y, axis=0), axis=0)
gmod_gauss = lm.Model(fit.gauss)
gmod_osc_gauss = lm.Model(fit.osc_gauss)
params_x = lm.Parameters()
params_x.add('l', value=self.perp_guess_x)
params_x.add('p', value=0.0, vary=False)
fit_x = gmod_gauss.fit(avg_corr_x, params_x, x=self.dx)
params_y = lm.Parameters()
params_y.add('l', value=self.perp_guess_y)
if not self.ky_free:
params_y.add('k', value=1, expr='2*3.141592653589793/l')
else:
params_y.add('k', value=self.perp_guess_ky)
params_y.add('p', value=0.0, vary=False)
fit_y = gmod_osc_gauss.fit(avg_corr_y, params_y, x=self.dy)
self.perp_fit_x[it] = fit_x.best_values['l']
self.perp_fit_y[it] = fit_y.best_values['l']
if self.ky_free:
self.perp_fit_ky[it] = fit_y.best_values['k']
if fit_x.errorbars:
self.perp_fit_x_err[it] = np.sqrt(fit_x.covar[0,0])
else:
self.perp_fit_x_err[it] = 0
if fit_y.errorbars:
self.perp_fit_y_err[it] = np.sqrt(fit_y.covar[0,0])
if self.ky_free:
self.perp_fit_ky_err[it] = np.sqrt(fit_y.covar[1,1])
else:
self.perp_fit_y_err[it] = 0
if self.ky_free:
self.perp_fit_ky_err[it] = 0
self.perp_guess_x = fit_x.best_values['l']
self.perp_guess_y = fit_y.best_values['l']
if self.ky_free:
self.perp_guess_ky = fit_y.best_values['k']
self.perp_plots_x(it, avg_corr_x, corr_std_x, fit_x)
self.perp_plots_y(it, avg_corr_y, corr_std_y, fit_y)
[docs] def perp_plots_x(self, it, corr_fn, corr_std, corr_fit):
"""
Plot radial correlation function and fitted Gaussian.
"""
plot_style.white()
plt.clf()
fig, ax = plt.subplots(1, 1)
plt.scatter(self.dx, corr_fn, color=pal[0],
label=r'$C(\Delta x)$')
plt.plot(self.dx, corr_fit.best_fit, color=pal[2],
label=r'$\exp(-(\Delta x / \ell_x)^2)$')
plt.fill_between(self.dx, corr_fn-corr_std, corr_fn+corr_std,
alpha=0.3)
plt.legend()
plt.xlabel(r'$\Delta x$ (m)')
plt.xlim([self.dx[0], self.dx[-1]])
plt.ylabel(r'$C(\Delta x)$')
plt.ylim(ymax=1)
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/' +self.perp_dir +
'/corr_fns_x/corr_x_fit_it_' + str(it) + '.pdf')
plt.close(fig)
[docs] def perp_plots_y(self, it, corr_fn, corr_std, corr_fit):
"""
Plot radial correlation function and fitted Gaussian.
"""
plot_style.white()
plt.clf()
fig, ax = plt.subplots(1, 1)
plt.scatter(self.dy, corr_fn, color=pal[0], label=r'$C(\Delta y)$')
plt.plot(self.dy, np.exp(-(self.dy/corr_fit.best_values['l'])**2),
'k--', label=r'$\exp(-(\Delta y / \ell_y)^2)$')
if not self.ky_free:
fit_label=r'$\exp(-(\Delta y / \ell_y)^2) \cos(2 \pi \Delta y/ \ell_y)$'
else:
fit_label=r'$\exp(-(\Delta y / \ell_y)^2) \cos(k_y \Delta y)$'
plt.plot(self.dy, corr_fit.best_fit, color=pal[2], label=fit_label)
plt.fill_between(self.dy, corr_fn-corr_std, corr_fn+corr_std,
alpha=0.3)
plt.legend()
plt.xlabel(r'$\Delta y$ (m)')
plt.xlim([self.dy[0], self.dy[-1]])
plt.ylabel(r'$C(\Delta y)$')
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/' +self.perp_dir +
'/corr_fns_y/corr_y_fit_it_' + str(it) + '.pdf')
plt.close(fig)
[docs] def perp_analysis_summary(self):
"""
Prints out a summary of the perpendicular analysis.
* Plots fitting parameters as a function of time window.
* Averages them in time and calculates a standard deviation.
* Writes summary to a text file.
"""
logging.info("Writing perp_analysis summary...")
plt.clf()
plot_style.white()
fig, ax = plt.subplots(1, 1)
plt.errorbar(range(self.nt_slices), np.abs(self.perp_fit_x),
yerr=self.perp_fit_x_err, capthick=1, capsize=5)
plt.xlabel('Time Window')
plt.ylabel(r'$l_x$ (m)')
plt.ylim(ymin=0, ymax=2*np.mean(np.abs(self.perp_fit_x[0])))
plt.xticks(range(self.nt_slices))
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/'+self.perp_dir+'/perp_fit_x_vs_time_slice.pdf')
plt.close(fig)
plt.clf()
plot_style.white()
fig, ax = plt.subplots(1, 1)
plt.errorbar(range(self.nt_slices), np.abs(self.perp_fit_y),
yerr=self.perp_fit_y_err, capthick=1, capsize=5)
plt.xlabel('Time Window')
plt.ylabel(r'$l_y$ (m)')
plt.ylim(ymin=0, ymax=2*np.mean(np.abs(self.perp_fit_y)))
plt.xticks(range(self.nt_slices))
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/'+self.perp_dir+
'/perp_fit_y_vs_time_slice.pdf')
plt.close(fig)
perp_results = {}
current_analysis = 'perp'
perp_results['lx_t'] = np.abs(self.perp_fit_x).tolist()
perp_results['lx'] = abs(np.nanmean(self.perp_fit_x))
perp_results['lx_t_err'] = np.abs(self.perp_fit_x_err).tolist()
perp_results['lx_err'] = abs(np.nanmean(self.perp_fit_x_err))
perp_results['ly_t'] = np.abs(self.perp_fit_y).tolist()
perp_results['ly'] = abs(np.nanmean(self.perp_fit_y))
perp_results['ly_t_err'] = np.abs(self.perp_fit_y_err).tolist()
perp_results['ly_err'] = abs(np.nanmean(self.perp_fit_y_err))
if self.ky_free:
plt.clf()
plot_style.white()
fig, ax = plt.subplots(1, 1)
plt.errorbar(range(self.nt_slices), np.abs(self.perp_fit_ky),
yerr=self.perp_fit_ky_err, capthick=1, capsize=5)
plt.xlabel('Time Window')
plt.ylabel(r'$k_y (m^{-1})$')
plt.ylim(ymin=0, ymax=2*np.mean(np.abs(self.perp_fit_ky)))
plt.xticks(range(self.nt_slices))
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/'+self.perp_dir+
'/perp_fit_ky_vs_time_slice.pdf')
plt.close(fig)
perp_results['ky_t'] = self.perp_fit_ky.tolist()
perp_results['ky'] = np.nanmean(self.perp_fit_ky)
perp_results['ky_t_err'] = self.perp_fit_ky_err.tolist()
perp_results['ky_err'] = np.nanmean(self.perp_fit_ky_err)
current_analysis = 'perp_ky_free'
self.write_results(current_analysis, perp_results)
logging.info("Finished writing perp_analysis summary...")
[docs] def write_results(self, analysis, result_dict):
"""
Write results to the main results file.
Parameters
----------
result_dict : dict
Dictionary containing results from a given analysis. Will overwrite
any existing results for that analysis.
"""
try:
with open(self.out_dir + '/' + 'results.json', 'r') as fp:
results = json.load(fp)
except FileNotFoundError:
results = {}
with open(self.out_dir + '/' + 'results.json', 'w') as fp:
results[analysis] = result_dict
json.dump(results, fp)
[docs] def time_analysis(self):
"""
Performs a time correlation analysis on the field.
Notes
-----
* Split into time windows and perform correlation analysis on each
window separately.
"""
logging.info("Starting time_analysis...")
if self.lab_frame:
self.time_dir = 'time_lab_frame'
elif not self.lab_frame:
self.time_dir = 'time'
if self.time_dir not in os.listdir(self.out_dir):
os.system("mkdir -p " + self.out_dir + '/' + self.time_dir)
if 'corr_fns' not in os.listdir(self.out_dir+'/'+self.time_dir):
os.system("mkdir -p " + self.out_dir + '/'+self.time_dir+'/corr_fns')
os.system('rm -f ' + self.out_dir + '/'+self.time_dir+'/corr_fns/*')
self.time_corr = np.empty([self.nt_slices, self.time_slice, self.nx,
self.ny], dtype=float)
self.corr_time = np.empty([self.nt_slices, self.nx], dtype=float)
self.corr_time_err = np.empty([self.nt_slices, self.nx], dtype=float)
self.field_normalize_time()
pbar = ProgressBar(widgets=['Progress: ', Percentage(), Bar()])
for it in pbar(range(self.nt_slices)):
self.calculate_time_corr(it)
self.time_norm_mask(it)
self.time_corr_fit(it)
self.time_analysis_summary()
logging.info("Finished time_analysis...")
[docs] def field_normalize_time(self):
"""
Defines normalized field for the time correlation by subtracting the
mean and dividing by the RMS value.
"""
logging.info('Normalizing the real space field...')
self.field_real_space_norm = \
np.empty([self.nt,self.nx,self.ny],dtype=float)
for it in range(self.nt_slices):
t_min = it*self.time_slice
t_max = (it+1)*self.time_slice
field_window = self.field_real_space[t_min:t_max,:,:]
for ix in range(self.nx):
self.field_real_space_norm[t_min:t_max,ix,:] = \
field_window[:,ix,:] - \
np.mean(field_window[:,ix,:])
self.field_real_space_norm[t_min:t_max,ix,:] /= \
np.std(self.field_real_space_norm[t_min:t_max,ix,:])
logging.info('Finished normalizing the real space field.')
[docs] def calculate_time_corr(self, it):
"""
Calculate the time correlation for a given time window at each x.
Parameters
----------
it : int
This is the index of the time slice currently being calculated.
"""
field_window = self.field_real_space_norm[it*self.time_slice:(it+1)*
self.time_slice,:,:]
for ix in range(self.nx):
self.time_corr[it,:,ix,:] = sig.fftconvolve(field_window[:,ix,:],
field_window[::-1,ix,::-1],
'same')
[docs] def time_norm_mask(self, it):
"""
Applies the appropriate normalization to the time correlation function.
Notes
-----
After calling ``sig.fftconvolve`` to calculate ``time_corr``, we are
left with an unnormalized correlation function as a function of dt
and dy. This function applies a 2D nomalization mask to ``time_corr``
which is dependent on the number of points that ``field_real_space_norm``
has in common with itself for a given dt, dy, and time window.
``field_real_space_norm`` is already normalized to the standard
deviation of the time signal, so the only difference between fftconvolve
and np.corrcoef is the number of points in common in the convolution
(that aren't the zero padded values and after averaging over many time
steps).
Parameters
----------
it : int
This is the index of the time slice currently being calculated.
"""
logging.info('Applying time normalization mask...')
x = np.ones([self.time_slice, self.ny])
mask = sig.fftconvolve(x,x,'same')
for ix in range(self.nx):
self.time_corr[it,:,ix,:] /= mask
logging.info('Finised applying time normalization mask.')
[docs] def time_corr_fit(self, it):
"""
Fit the time correlation function with exponentials to find correlation
time.
Notes
-----
The fitting procedure consists of the following steps:
* Loop over radial points
* Identify the time correlation function peaks
* Determine what type of function to fit to the peaks
* Fitting the appropriate function
The peaks can be fitted with a growing/decaying exponential depending
on the direction of decrease of the peaks, or an oscillating exponential
if the peaks don't monotically increase or decrease.
Parameters
----------
it : int
This is the index of the time slice currently being fitted.
"""
t = self.t[it*self.time_slice:(it+1)*self.time_slice]
self.dt = np.linspace((-max(t)+t[0])/2, (max(t)-t[0])/2, self.time_slice)
peaks = np.zeros([self.nx, self.npeaks_fit], dtype=float)
max_index = np.empty([self.nx, self.npeaks_fit], dtype=int);
mid_idx = int(self.ny/2)
for ix in range(self.nx):
for iy in range(mid_idx,mid_idx+self.npeaks_fit):
max_index[ix, iy-mid_idx], peaks[ix, iy-mid_idx] = \
max(enumerate(self.time_corr[it,:,ix,iy]),
key=operator.itemgetter(1))
if (fit.strictly_increasing(max_index[ix,:]) == True or
fit.strictly_increasing(max_index[ix,::-1]) == True):
if max_index[ix, self.npeaks_fit-1] > max_index[ix, 0]:
try:
gmod_decay = lm.Model(fit.decaying_exp)
params_t = lm.Parameters()
params_t.add('tau_c', value=self.time_guess_dec)
fit_t = gmod_decay.fit(peaks[ix,:], params_t,
t=self.dt[max_index[ix,:]])
if np.abs(fit_t.best_values['tau_c']) > self.time_max:
self.corr_time[it,ix] = np.nan
self.corr_time_err[it,ix] = np.nan
else:
self.corr_time[it,ix] = fit_t.best_values['tau_c']
if fit_t.errorbars:
self.corr_time_err[it,ix] = np.sqrt(fit_t.covar[0,0])
else:
self.corr_time_err[it,ix] = np.nan
self.time_guess_dec = self.corr_time[it,ix]
self.time_plot(it, ix, max_index, peaks, 'decaying')
logging.info("(" + str(it) + "," + str(ix) + ") was fitted "
"with decaying exponential. tau = "
+ str(self.corr_time[it,ix]) + " s\n")
except RuntimeError:
logging.info("(" + str(it) + "," + str(ix) + ") "
"RuntimeError - max fitting iterations reached, "
"skipping this case with tau = NaN\n")
self.corr_time[it, ix] = np.nan
else:
try:
gmod_grow = lm.Model(fit.growing_exp)
params_t = lm.Parameters()
params_t.add('tau_c', value=self.time_guess_grow)
fit_t = gmod_grow.fit(peaks[ix,:], params_t,
t=self.dt[max_index[ix,:]])
if np.abs(fit_t.best_values['tau_c']) > self.time_max:
self.corr_time[it,ix] = np.nan
self.corr_time_err[it,ix] = np.nan
else:
self.corr_time[it,ix] = fit_t.best_values['tau_c']
if fit_t.errorbars:
self.corr_time_err[it,ix] = np.sqrt(fit_t.covar[0,0])
else:
self.corr_time_err[it,ix] = np.nan
self.time_guess_grow = self.corr_time[it,ix]
self.time_plot(it, ix, max_index, peaks, 'growing')
logging.info("(" + str(it) + "," + str(ix) + ") was fitted "
"with growing exponential. tau = "
+ str(self.corr_time[it,ix]) + " s\n")
except RuntimeError:
logging.info("(" + str(it) + "," + str(ix) + ") "
"RuntimeError - max fitting iterations reached, "
"skipping this case with tau = NaN\n")
self.corr_time[it, ix] = np.nan
else:
try:
# If abs(max_index) is not monotonically increasing, this
# usually means that there is no flow and that the above method
# cannot be used to calculate the correlation time. Try fitting
# a decaying oscillating exponential to the central peak.
gmod_osc = lm.Model(fit.osc_gauss)
params_t = lm.Parameters()
params_t.add('l', value=self.time_guess_osc[0])
params_t.add('k', value=self.time_guess_osc[1])
params_t.add('p', value=self.time_guess_osc[2], vary=False)
fit_t = gmod_osc.fit(self.time_corr[it,:,ix,mid_idx],
params_t, x=self.dt)
# Note l = tau_c sinc fitting function specification is for
# general l, k, p.
if np.abs(fit_t.best_values['l']) > self.time_max:
self.corr_time[it,ix] = np.nan
self.corr_time_err[it,ix] = np.nan
logging.info("(" + str(it) + "," + str(ix) + ") was "
"fitted with an oscillating Gaussian to "
"the central peak. (tau, omega) = "
+ str([np.nan, np.nan]) + "\n")
else:
self.corr_time[it,ix] = fit_t.best_values['l']
if fit_t.errorbars:
self.corr_time_err[it,ix] = np.sqrt(fit_t.covar[0,0])
else:
self.corr_time_err[it,ix] = np.nan
self.time_guess_osc = np.array([fit_t.best_values['l'],
fit_t.best_values['k'],
fit_t.best_values['p']])
self.time_plot(it, ix, max_index, peaks, 'oscillating',
omega=fit_t.best_values['k'])
logging.info("(" + str(it) + "," + str(ix) + ") was "
"fitted with an oscillating Gaussian to "
"the central peak. (tau, omega) = "
+ str([fit_t.best_values['l'],
fit_t.best_values['k']]) + "\n")
except RuntimeError:
logging.info("(" + str(it) + "," + str(ix) + ") "
"RuntimeError - max fitting iterations reached, "
"skipping this case with (tau, omega) = NaN\n")
self.corr_time[it, ix] = np.nan
[docs] def time_plot(self, it, ix, max_index, peaks, plot_type, **kwargs):
"""
Plots the time correlation peaks as well as the apprpriate fitting
function.
Parameters
----------
it : int
Time slice currently being fitted.
ix : int
Radial index currently being fitted.
max_index : array_like
Array containing radial indices of the peaks. Size: (nx, npeaks_fit)
peaks : array_like
Array of peak values at the max_index values. Size: (nx, npeaks_fit)
plot_type : str
Type of fitting function to plot. One of: 'growing'/'decaying'/
'oscillating'
"""
plot_style.white()
mid_idx = int(self.ny/2)
plt.clf()
fig, ax = plt.subplots(1, 1)
if plot_type == 'decaying':
plt.plot(self.dt*1e6, self.time_corr[it,:,ix,mid_idx:mid_idx+self.npeaks_fit])
plt.hold(True)
plt.plot(self.dt[max_index[ix,:]]*1e6, peaks[ix,:], 'o', color='#7A1919')
plt.hold(True)
plt.plot(self.dt[int(self.time_slice/2):]*1e6,
fit.decaying_exp(self.dt[int(self.time_slice/2):],self.corr_time[it,ix]),
'k--', lw=2,
label=r'$\exp[-|\Delta t_{peak} / \tau_c|]$')
plt.legend()
if plot_type == 'growing':
plt.plot(self.dt*1e6, self.time_corr[it,:,ix,mid_idx:mid_idx+self.npeaks_fit])
plt.hold(True)
plt.plot(self.dt[max_index[ix,:]]*1e6, peaks[ix,:], 'o', color='#7A1919')
plt.hold(True)
plt.plot(self.dt[:int(self.time_slice/2)]*1e6,
fit.growing_exp(self.dt[:int(self.time_slice/2)],self.corr_time[it,ix]),
'k--', lw=2,
label=r'$\exp[|\Delta t_{peak} / \tau_c|]$')
plt.legend()
if plot_type == 'oscillating':
plt.plot(self.dt*1e6, self.time_corr[it,:,ix,mid_idx])
plt.hold(True)
plt.plot(self.dt*1e6, fit.osc_gauss(self.dt,self.corr_time[it,ix],
kwargs['omega'], 0), 'k--', lw=2,
label=r'$\exp[- (\Delta t_{peak} / \tau_c)^2] '
'\cos(\omega \Delta t) $')
plt.legend()
plt.xlabel(r'$\Delta t (\mu s)})$')
plt.ylabel(r'$C_{\Delta y}(\Delta t)$')
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/'+self.time_dir+'/corr_fns/time_fit_it_' +
str(it) + '_ix_' + str(ix) + '.pdf')
plt.close(fig)
[docs] def time_analysis_summary(self):
"""
Prints out a summary of the time analysis.
* Plots average correlation time as a function of radius.
* Calculates standard deviation.
* Writes summary to a text file.
"""
logging.info("Writing time_analysis summary...")
self.corr_time = np.abs(self.corr_time)
time_results = {}
if self.lab_frame:
current_analysis = 'time_lab_frame'
else:
current_analysis = 'time'
time_results['corr_time'] = self.corr_time.tolist()
time_results['corr_time_err'] = self.corr_time_err.tolist()
time_results['tau_c'] = np.nanmean(self.corr_time)*1e6
time_results['tau_c_err'] = np.nanstd(self.corr_time)*1e6
self.write_results(current_analysis, time_results)
# Plot corr_time as a function of radius, average over time window
plt.clf()
fig, ax = plt.subplots(1, 1)
t_error = np.nanstd(self.corr_time*1e6, axis=0)
plt.errorbar(self.x, np.nanmean(self.corr_time*1e6, axis=0),
yerr=t_error, capthick=1, capsize=5)
plt.ylim(ymin=0)
plt.xlabel("Radius (m)")
plt.ylabel(r'Correlation Time $\tau_c$ ($\mu$ s)')
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/'+self.time_dir+'/corr_time.pdf')
logging.info("Finished writing time_analysis summary...")
[docs] def par_analysis(self):
"""
Calculates the parallel correlation function and fits with a Gaussian
to find the parallel correlation length.
"""
logging.info("Starting par_analysis...")
if 'parallel' not in os.listdir(self.out_dir):
os.system("mkdir -p " + self.out_dir + '/parallel')
if 'corr_fns' not in os.listdir(self.out_dir + '/parallel'):
os.system("mkdir -p " + self.out_dir + '/parallel/corr_fns')
os.system('rm -f ' + self.out_dir + '/parallel/corr_fns/*')
self.calculate_l_par()
self.calculate_par_corr()
self.par_fit_params = np.empty([self.nt_slices, 2],
dtype=float)
self.par_fit_params_err = np.empty([self.nt_slices, 2],
dtype=float)
for it in range(self.nt_slices):
self.par_corr_fit(it)
self.par_analysis_summary()
logging.info("Finished par_analysis...")
[docs] def calculate_l_par(self):
"""
Calculates the real space parallel grid.
"""
logging.info('Start calculating parallel length...')
dR_dtheta = np.gradient(self.R)/np.gradient(self.theta)
dZ_dtheta = np.gradient(self.Z)/np.gradient(self.theta)
dphi_dtheta = self.r_geo*self.bref / (self.R**2 * \
self.bmag[int(self.ntheta/2)] * self.gradpar)
dl_dtheta = np.sqrt(dR_dtheta**2 + dZ_dtheta**2 + (self.R*dphi_dtheta)**2)
self.l_par = np.append(0, integrate.cumtrapz(dl_dtheta, x=self.theta))
logging.info('Finished calculating parallel length.')
[docs] def calculate_par_corr(self):
"""
Calculate the parallel correlation function and apply normalization mask.
Interpolation onto a regular parallel grid, correlation calculation,
and normalization are done in one function to avoid unnecessary looping
over x, y, and theta.
"""
logging.info('Start calculating parallel correlation function...')
x = np.ones([self.ntheta])
mask = sig.correlate(x, x, 'same')
self.par_corr = np.empty([self.nt, self.nx, self.ny, self.ntheta],
dtype=float)
l_par_reg = np.linspace(0, self.l_par[-1], self.ntheta)
pbar = ProgressBar(widgets=['Progress: ', Percentage(), Bar()])
for it in pbar(range(self.nt)):
logging.info('Parallel correlation timestep: %d of %d'%(it,self.nt))
for ix in range(self.nx):
for iy in range(self.ny):
f = interp.interp1d(self.l_par,
self.field_real_space[it,ix,iy,:])
self.field_real_space[it,ix,iy,:] = f(l_par_reg)
self.field_real_space[it,ix,iy,:] = \
self.field_real_space[it,ix,iy,:] - \
np.mean(self.field_real_space[it,ix,iy,:])
self.field_real_space[it,ix,iy,:] = \
self.field_real_space[it,ix,iy,:] / \
np.std(self.field_real_space[it,ix,iy,:])
self.par_corr[it,ix,iy,:] = \
sig.correlate(self.field_real_space[it,ix,iy,:],
self.field_real_space[it,ix,iy,:],
'same')/mask
self.l_par = l_par_reg
self.dl_par = np.linspace(-self.l_par[-1]/2, self.l_par[-1]/2, self.ntheta)
logging.info('Finished calculating parallel correlation function.')
[docs] def par_corr_fit(self, it):
"""
Fit the parallel correlation function with an oscillatory Gaussian
function for time slice it.
Parameters
----------
it : int
Time slice to average over and fit.
Before fitting average over time, x, and y.
"""
corr_fn = self.par_corr[it*self.time_slice:(it+1)*self.time_slice,:,:,:]
corr_std = np.empty([self.ntheta])
for i in range(self.ntheta):
corr_std[i] = np.std(corr_fn[:,:,:,i])
corr_fn = np.mean(np.mean(np.mean(corr_fn, axis=0), axis=0), axis=0)
try:
gmod_osc = lm.Model(fit.osc_gauss)
params = lm.Parameters()
params.add('l', value=self.par_guess[0], min=self.l_par[1],
max=100)
params.add('k', value=self.par_guess[1],
max=(2*np.pi/self.dl_par[-1]*len(self.dl_par)/2))
params.add('p', value=0, vary=False)
par_fit = gmod_osc.fit(corr_fn, params, x=self.dl_par)
self.par_fit_params[it, :] = np.abs([par_fit.best_values['l'],
par_fit.best_values['k']])
self.par_fit_params_err[it, :] = np.sqrt(np.diag(par_fit.covar))
self.par_plot(it, corr_fn, corr_std)
except RuntimeError:
logging.info("(" + str(it) + ") RuntimeError - max fitting iterations reached, "
"skipping this case with (l_par, k_par) = NaN\n")
self.par_fit_params[it, :] = np.nan
self.par_fit_params_err[it, :] = np.nan
self.par_guess = self.par_fit_params[it, :]
[docs] def par_plot(self, it, corr, corr_std):
"""
Plots and saves the parallel correlation function and its fit for each
time window.
Parameters
----------
it : int
Time window being plotted.
corr : array_like
Parallel correlation averaged in t, x, and y.
corr_std : array_like
Standard deviation in the correlation function as a function of
dl_par.
"""
plot_style.white()
plt.clf()
fig, ax = plt.subplots(1, 1)
plt.scatter(self.dl_par, corr, color=pal[0],
label=r'$C(\Delta t = 0, \Delta x = 0, \Delta y = 0, \Delta z)$')
plt.fill_between(self.dl_par, corr-corr_std, corr+corr_std,
alpha=0.3)
plt.plot(self.dl_par, fit.osc_gauss(self.dl_par, self.par_fit_params[it,0],
self.par_fit_params[it,1], 0), color=pal[2] ,
label=r'$p_\parallel + (1-p_\parallel)\exp[- (\Delta z / l_{\parallel})^2] '
'\cos(k_{\parallel} \Delta z) $')
plt.plot(self.dl_par, np.exp(-(self.dl_par/self.par_fit_params[it,0])**2),
'k--', label='Gaussian Envelope')
plt.legend()
plt.xlabel(r'$\Delta z$ (m)')
plt.ylabel(r'$C(\Delta z)$')
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/parallel/corr_fns/par_fit_it_' +
str(it) + '.pdf')
plt.close(fig)
[docs] def par_analysis_summary(self):
"""
Summarize parallel correlation analysis by plotting parallel correlation
fitting parameters along with associated errors and writing to results
file.
"""
plot_style.white()
par_results = {}
current_analysis = 'par'
par_results['par_fit_params'] = self.par_fit_params.tolist()
par_results['l_par'] = np.nanmean(self.par_fit_params[:,0])
par_results['l_par_err'] = np.nanmean(self.par_fit_params_err[:,0])
par_results['k_par'] = np.nanmean(self.par_fit_params[:,1])
par_results['k_par_err'] = np.nanmean(self.par_fit_params_err[:,1])
self.write_results(current_analysis, par_results)
plt.clf()
fig, ax = plt.subplots(1, 1)
plt.errorbar(range(self.nt_slices), np.abs(self.par_fit_params[:,0]),
yerr=self.par_fit_params_err[:,0], capthick=1,
capsize=5)
plt.xlabel('Time Window')
plt.ylabel(r'Parallel Correlation Length $l_{\parallel} (m)$')
plt.ylim(ymin=0, ymax=2*np.mean(np.abs(self.par_fit_params[:,0])))
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/parallel/par_fit_length_vs_time_slice.pdf')
plt.close(fig)
plt.clf()
fig, ax = plt.subplots(1, 1)
plt.errorbar(range(self.nt_slices), np.abs(self.par_fit_params[:,1]),
yerr=self.par_fit_params_err[:,1], capthick=1, capsize=5)
plt.xlabel('Time Window')
plt.ylabel(r'Parallel Correlation Wavenumber $k_{\parallel} (m^{-1})$')
plt.ylim(ymin=0, ymax=2*np.mean(np.abs(self.par_fit_params[:,1])))
plot_style.minor_grid(ax)
plot_style.ticks_bottom_left(ax)
plt.savefig(self.out_dir + '/parallel/par_fit_wavenumber_vs_time_slice.pdf')
plt.close(fig)
[docs] def write_field(self):
"""
Outputs the field to NetCDF in real space.
Notes
-----
* The radial and poloidal coordinates are centered at 0.
* The radial coordinate is interpolated if neccessary to ensure a
0.5cm resolution consistent with the BES.
"""
logging.info("Starting write_field...")
if 'write_field' not in os.listdir(self.out_dir):
os.system("mkdir -p " + self.out_dir + '/write_field')
if self.write_field_interp_x:
#interpolate radial coordinate to be approx 0.5cm
interp_fac = int(np.ceil(self.x[1]/0.005))
x_nc = np.linspace(min(self.x), max(self.x), interp_fac*self.nx)
field_real_space_nc = np.empty([self.nt, len(x_nc), self.ny],
dtype=float)
for it in range(self.nt):
for iy in range(self.ny):
f = interp.interp1d(self.x,
self.field_real_space[it,:,iy])
field_real_space_nc[it,:,iy] = f(x_nc)
else:
x_nc = self.x
field_real_space_nc = self.field_real_space
if self.lab_frame:
nc_file = Dataset(self.out_dir + '/write_field/' +
self.in_field +'_lab_frame.cdf', 'w')
elif not self.lab_frame:
nc_file = Dataset(self.out_dir + '/write_field/' +
self.in_field +'.cdf', 'w')
nc_file.createDimension('x', len(x_nc))
nc_file.createDimension('y', self.ny)
nc_file.createDimension('t', self.nt)
nc_file.createDimension('none', 1)
nc_nref = nc_file.createVariable('nref','d', ('none',))
nc_tref = nc_file.createVariable('tref','d', ('none',))
nc_x = nc_file.createVariable('x','d',('x',))
nc_y = nc_file.createVariable('y','d',('y',))
nc_t = nc_file.createVariable('t','d',('t',))
nc_field = nc_file.createVariable(self.in_field[:self.in_field.find('_')],
'd',('t', 'x', 'y'))
nc_field[:,:,:] = field_real_space_nc[:,:,:]
nc_nref[:] = self.nref
nc_tref[:] = self.tref
nc_x[:] = x_nc[:] - x_nc[-1]/2
nc_y[:] = self.y[:] - self.y[-1]/2
nc_t[:] = self.t[:] - self.t[0]
nc_file.close()
logging.info("Finished write_field...")
[docs] def write_field_full(self):
"""
Outputs the full 3D field to NetCDF in real space.
Notes
-----
* The radial, poloidal, and parallel coordinates are centered at 0.
* The radial coordinate is interpolated if neccessary to ensure a
0.5cm resolution consistent with the BES.
"""
logging.info("Starting write_field_full...")
if 'write_field_full' not in os.listdir(self.out_dir):
os.system("mkdir -p " + self.out_dir + '/write_field_full')
self.calculate_l_par()
if self.write_field_interp_x:
#interpolate radial coordinate to be approx 0.5cm
interp_fac = int(np.ceil(self.x[1]/0.005))
x_nc = np.linspace(min(self.x), max(self.x), interp_fac*self.nx)
field_real_space_nc = np.empty([self.nt, len(x_nc), self.ny,
self.ntheta], dtype=float)
for it in range(self.nt):
for iy in range(self.ny):
for iz in range(self.ntheta):
f = interp.interp1d(self.x,
self.field_real_space[it,:,iy,iz])
field_real_space_nc[it,:,iy,iz] = f(x_nc)
else:
x_nc = self.x
field_real_space_nc = self.field_real_space
if self.lab_frame:
nc_file = Dataset(self.out_dir + '/write_field_full/' +
self.in_field +'_lab_frame.cdf', 'w')
elif not self.lab_frame:
nc_file = Dataset(self.out_dir + '/write_field_full/' +
self.in_field +'.cdf', 'w')
nc_file.createDimension('x', len(x_nc))
nc_file.createDimension('y', self.ny)
nc_file.createDimension('z', self.ntheta)
nc_file.createDimension('t', self.nt)
nc_file.createDimension('none', 1)
nc_nref = nc_file.createVariable('nref','d', ('none',))
nc_tref = nc_file.createVariable('tref','d', ('none',))
nc_x = nc_file.createVariable('x','d',('x',))
nc_y = nc_file.createVariable('y','d',('y',))
nc_z = nc_file.createVariable('z','d',('z',))
nc_t = nc_file.createVariable('t','d',('t',))
nc_field = nc_file.createVariable(self.in_field[:self.in_field.find('_')],
'd',('t', 'x', 'y', 'z'))
nc_field[:,:,:,:] = field_real_space_nc[:,:,:,:]
nc_nref[:] = self.nref
nc_tref[:] = self.tref
nc_x[:] = x_nc[:] - x_nc[-1]/2
nc_y[:] = self.y[:] - self.y[-1]/2
nc_z[:] = self.l_par[:] - self.l_par[-1]/2
nc_t[:] = self.t[:] - self.t[0]
nc_file.close()
logging.info("Finished write_field_full...")