text:
"""
Module css calculates or plots shear strain and displacement vorticity as well
as their correlations.

Files are saved according to active_particles.naming.Css (shear strain) and
active_particles.naming.Ccc (displacement vorticity) standards.

A brief description of the algorithm can be found at:
https://yketa.github.io/UBC_2018_Wiki/#Shear%20strain%20and%20vorticity%20fields

Environment modes
-----------------
MODE : string
	Computation mode.
     __________________________________________________________________
	| Mode      | Computation                                          |
	|___________|______________________________________________________|
	| 'real'    | Real space (from coarse-grained displacements)       |
	|___________|______________________________________________________|
	| 'fourier' | Fourier space (from displacements Fourier transform) |
	|___________|______________________________________________________|
	| 'ovito'   | OVITO                                                |
	|___________|______________________________________________________|
	DEFAULT: fourier
COMPUTE : bool
	Compute shear strain (all modes) and displacement vorticity ('real' and
	'fourier' modes) correlations.
	DEFAULT: False
PLOT : bool
	Plot saved shear strain (all modes) and displacement vorticity ('real'
	mode) as well as their correlations.
	DEFAULT: False
SHOW [COMPUTE or PLOT mode] : bool
	Show graphs.
	DEFAULT: False
SAVE [COMPUTE or PLOT mode] : bool
	Save graphs.
	DEFAULT: False
LINEAR_INTERPOLATION [not('real') and SHOW mode] : bool
	Get value on grid by linear interpolation of neighbouring grid boxes.
	DEFAULT: False
SUPERIMPOSE_C44 [not('real') and SHOW mode] : bool
	Superimpose C44 curve to Css(r, theta) curve.
	DEFAULT: False
FITTING_LINE [not('real') and SHOW mode] : bool
	Display adjustable fitting line on graph of projection of shear strain on
	cos(4 \\theta).
	DEFAULT: False

Environment parameters
----------------------
DATA_DIRECTORY : string
	Data directory.
	DEFAULT: current working directory
PARAMETERS_FILE : string
	Simulation parameters file.
	DEFAULT: DATA_DIRECTORY/active_particles.naming.parameters_file
WRAPPED_FILE : string
	Wrapped trajectory file. (.gsd)
	DEFAULT: DATA_DIRECTORY/active_particles.naming.wrapped_trajectory_file
UNWRAPPED_FILE : string
	Unwrapped trajectory file. (.dat)
	NOTE: .dat files defined with active_particles.dat
	DEFAULT: DATA_DIRECTORY/active_particles.naming.unwrapped_trajectory_file
INITIAL_FRAME : int
	Frame to consider as initial.
	NOTE: INITIAL_FRAME < 0 will be interpreted as the initial frame being
	      the middle frame of the simulation.
	DEFAULT: -1
TIME : int
	Lag time for displacement.
	NOTE: TIME < 0 will be interpreted as a lag time corresponding to the total
	      number of simulation frames - INITIAL_FRAME + TIME.
	DEFAULT: -1
INTERVAL_MAXIMUM : int
	Maximum number of intervals of length dt considered in correlations
	calculations.
	DEFAULT: 1
R_CUT ['real' mode] : float
	Cut-off radius for coarse graining function.
	NOTE: Value of R_CUT is then multiplied by average particle diameter from
	      simulation parameters file.
	DEFAULT: active_particles.analysis.css._r_cut
SIGMA ['real' mode] : float
	Length scale of the spatial extent of the coarse graining function.
	DEFAULT: R_CUT
N_CASES : int
	Number of boxes in each direction to compute the shear strain and
	displacement vorticity grid.
	DEFAULT: smallest integer value greater than or equal to the square root of
	         the number of particles from the simulation parameters file.
BOX_SIZE : float
	Size of the square box to consider.
	DEFAULT: simulation box size
X_ZERO : float
	1st coordinate of the centre of the square box to consider.
	DEFAULT: 0
Y_ZERO : float
	2nd coordinate of the centre of the square box to consider.
	DEFAULT: 0
R_MAX [PLOT or SHOW mode] : float
	Half size of the box showed for 2D correlation.
	NOTE: R_MAX < 0 will be interpreted as the box shown being the actual
	      simulation box.
	DEFAULT: active_particles.analysis.css._r_max
DISPLAY_GRID [COMPUTE and 'real' mode] : int
	Index of map in list of variable maps to display.
	DEFAULT : 0
FONT_SIZE : int
    Plot font size.
    DEFAULT: active_particles.analysis.css._font_size
POINTS_X_C44 [not('real') and PLOT mode] : int
	Number of radii at which to evaluate integrated strain correlation.
	DEFAULT: active_particles.analysis.css._points_x_c44
POINTS_THETA_C44 [not('real') and PLOT mode] : int
	Number of angles to evaluate integrated strain correlation.
	DEFAULT: active_particles.analysis.css._points_theta_c44
Y_MIN_C44 [not('real') and PLOT mode] : float
	Minimum plot value for C44.
	DEFAULT: active_particles.analysis.css._y_min_c44
Y_MAX_C44 [not('real') and PLOT mode] : float
	Maximum plot value for C44.
	DEFAULT: active_particles.analysis.css._y_max_c44
R_MIN_C44 [not('real') and PLOT mode] : float
	Minimum radius in average particle separation for C44 calculation.
	DEFAULT: active_particles.analysis.css._r_min_c44
R_MAX_C44 [not('real') and PLOT mode] : float
	Maximum radius in average particle separation for C44 calculation.
	DEFAULT: active_particles.analysis.css._r_max_c44
SMOOTH [not('real') and PLOT mode] : float
	C44 Gaussian smoothing length scale in units of average particle
	separation.
	DEFAULT: 0
R_CUT_FOURIER [not('real') and PLOT mode] : float
    Initial wave length Gaussian cut-off radius in units of average particle
	separation.
    DEFAULT: active_particles.analysis.css._r_cut_fourier
SLOPE_C44 [not('real'), PLOT and FITTING_LINE mode] : slope
	Initial slope for fitting line in C44 figure.
	DEFAULT: active_particles.analysis.css._slope0_c44
SLOPE_MIN_C44 [not('real'), PLOT and FITTING_LINE mode] : float
	Minimum slope for fitting line in C44 figure.
	DEFAULT: active_particles.analysis.css._slope_min_c44
SLOPE_MAX_C44 [not('real'), PLOT and FITTING_LINE mode] : float
	Maximum slope for fitting line in C44 figure.
	DEFAULT: active_particles.analysis.css._slope_max_c44
THETA [not('real') and PLOT mode] : float
	Angle at which to evaluate Css(r, theta) (in multiple of pi).
	DEFAULT:  active_particles.analysis.css._theta
POINTS_X_THETA [not('real') and PLOT mode] : int
	Number of radii at which to evaluate Css(r, theta).
	DEFAULT: POINTS_X_C44
Y_MIN_THETA [not('real') and PLOT mode] : float
	Minimum plot value for Css(r, theta).
	DEFAULT: active_particles.analysis.css._y_min_theta
Y_MAX_THETA [not('real') and PLOT mode] : float
	Maximum plot value for Css(r, theta).
	DEFAULT: active_particles.analysis.css._y_max_theta
R_MIN_THETA [not('real') and PLOT mode] : float
	Minimum radius in average particle separation for Css(r, theta)
	calculation.
	DEFAULT: active_particles.analysis.css._r_min_theta
R_MAX_THETA [not('real') and PLOT mode] : float
	Maximum radius in average particle separation for Css(r, theta)
	calculation.
	DEFAULT: active_particles.analysis.css._r_max_theta
SLOPE_THETA [not('real'), PLOT and FITTING_LINE mode] : slope
	Initial slope for fitting line in Css(r, theta) figure.
	DEFAULT: SLOPE_C44
SLOPE_MIN_THETA [not('real'), PLOT and FITTING_LINE mode] : float
	Minimum slope for fitting line in Css(r, theta) figure.
	DEFAULT: SLOPE_MIN_C44
SLOPE_MAX_THETA [not('real'), PLOT and FITTING_LINE mode] : float
	Maximum slope for fitting line in Css(r, theta) figure.
	DEFAULT: SLOPE_MAX_C44

Output
------
[COMPUTE and 'real' mode]
> Prints neigbours grid computation time and execution time.
> Saves a computed map of shear strain and the averaged shear strain
correlation according to active_particles.naming.Css standards in
DATA_DIRECTORY.
> Saves a computed map of displacement vorticity and the averaged displacement
vorticity correlations according to active_particles.naming.Ccc standards
in DATA_DIRECTORY.
[COMPUTE and 'fourier' mode]
> Saves average square norm of shear strain Fourier transforms according to
active_particles.naming.Css standards in DATA_DIRECTORY.
> Saves average square norm of displacement vorticity Fourier transforms
according to active_particles.naming.Ccc standards in DATA_DIRECTORY.
[COMPUTE and 'ovito' mode]
> Saves average square norm of shear strain Fourier transforms according to
active_particles.naming.Css standards in DATA_DIRECTORY.
[SHOW or PLOT mode]
> Plots data map and/or correlation for shear strain and/or displacement
vorticity.
[SAVE and 'real' mode]
> Saves data map and correlation figures in DATA_DIRECTORY.
[SAVE and not('real') mode]
> Saves shear strain correlation figure in DATA_DIRECTORY.
"""

import active_particles.naming as naming

from active_particles.init import get_env, slurm_output, linframes
from active_particles.dat import Dat, Gsd
from active_particles.maths import relative_positions, wave_vectors_2D,\
	FFT2Dfilter, gaussian_smooth_1D
from active_particles.quantities import nD0_active

from active_particles.analysis.neighbours import NeighboursGrid
from active_particles.analysis.correlations import corField2D_scalar_average
from active_particles.analysis.correlations import CorGrid
from active_particles.analysis.coarse_graining import GaussianCG,\
	CoarseGraining
from active_particles.analysis.cuu import displacement_grid, Cnn
from active_particles.analysis.number import count_particles

from os import getcwd
from os import environ as envvar
from os.path import join as joinpath

from math import ceil

import numpy as np

import pickle

from operator import itemgetter

from datetime import datetime

import matplotlib as mpl
if not(get_env('SHOW', default=False, vartype=bool)):
	mpl.use('Agg')	# avoids crash if launching without display
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.widgets import Slider
from matplotlib.gridspec import GridSpec
from matplotlib.lines import Line2D

from active_particles.plot.mpl_tools import GridCircle, FittingLine

# DEFAULT VARIABLES

_r_cut = 2	# default cut-off radius for coarse graining function
_r_max = 20	# default half size of the box showed for 2D correlation

_c_min = -0.2	# default minimum value for correlations
_c_max = 1		# default maximum value for correlations

_font_size = 10	# default plot font size

_slope0_c44 = -2    # default initial slope for fitting line slider in C44 figure
_slope_min_c44 = -5 # default minimum slope for fitting line slider in C44 figure
_slope_max_c44 = 0  # default maximum slope for fitting line slider in C44 figure

_points_x_c44 = 100     # default number of radii at which to evaluate integrated strain correlation
_points_theta_c44 = 100 # default number of angles to evaluate integrated strain correlation

_y_min_c44 = 1e-4   # default minimum C44 value
_y_max_c44 = 2e-1   # default maximum C44 value
_r_min_c44 = 1      # default minimum radius over average particle separation value for C44 figure
_r_max_c44 = 20     # default maximum radius over average particle separation value for C44 figure

_theta = 0	# default angle at which to evaluate Css(r, theta) (in multiple of pi)

_y_min_theta = 1e-4	# default minimum Css(r, theta) value
_y_max_theta = 2e-1	# default maximum Css(r, theta) value
_r_min_theta = 1	# default minimum radius over average particle separation value for Css(r, theta) figure
_r_max_theta = 20	# default maximum radius over average particle separation value for Css(r, theta) figure

_r_cut_fourier = 0	# default initial wave length Gaussian cut-off radius

# FUNCTIONS AND CLASSES

def strain_vorticity(point, time, dt, positions, displacements,
	sigma, r_cut, box_size, neighbours_grid):
	"""
	From coarse-grained displacement field u, calculates
	> linearised shear strain, \\epsilon_{xy} =
		1/2 (\\frac{\\partial u_y}{\\partial x}
		     + \\frac{\\partial u_x}{\\partial y})
	> and displacement vorticity, \\omega =
		\\frac{\\partial u_y}{\\partial x} - \\frac{\\partial u_x}{\\partial y}

	Resorts to neighbours grids to speed up calculations.

	Parameters
	----------
	point : array like
		Coordinates of point at which to calculate shear strain and
		displacement vorticity.
	time : int
		Frame at which shear strain will be calculated.
	dt : int
		Length of the interval of time for which the displacements are
		calculated.
	positions : (N, 2) shaped array like
		Array of wrapped particle positions.
	displacements : (N, 2) shaped array like
		Array of particle displacements.
	sigma : float
		Length scale of the spatial extent of the coarse graining function.
	r_cut : float
		Cut-off radius for coarse graining function.
		Also cut-off radius with which neighbours grid has been computed.
	box_size : float
		Length of the system's square box.
	neighbours_grid : active_particles.analysis.neighbours.NeighboursGrid
		Neighbours grid.

	Returns
	-------
	strain : float
		Linearised shear strain at point.
	vorticity : float
		Displacement vorticity at point.
	"""

	wrcut = neighbours_grid.get_neighbours(point)	# particles indexes within r_cut of point
	Nwrcut = len(wrcut)								# number of neighbouring particles
	if Nwrcut == 0:									# if there is no particles within r_cut of point
		return 0, 0

	pos_wrcut = relative_positions(
		np.array(itemgetter(*wrcut)(positions), ndmin=2), point, box_size)	# position at time time (with boundary conditions) with point as the centre of the frame
	dis_wrcut = np.array(itemgetter(*wrcut)(displacements), ndmin=2)		# displacements of the particles between time and time + dt

	coarse_graining = CoarseGraining(GaussianCG(sigma, r_cut).factors,
		pos_wrcut)															# coarse graining object
	rho = coarse_graining.average(np.full(len(pos_wrcut), fill_value=1))	# averaged density
	if rho == 0: return 0, 0												# if there is no particles within r_cut

	Ax, Ay, Aux, Auy, Auxy, Auyx = tuple(map(coarse_graining.average,
		[pos_wrcut[:, 0], pos_wrcut[:, 1], dis_wrcut[:, 0], dis_wrcut[:, 1],
		dis_wrcut[:, 0]*pos_wrcut[:, 1], dis_wrcut[:, 1]*pos_wrcut[:, 0]]))	# coarse grained x, y, u_x, u_y, u_x * y, u_y * x
	strain = 0.5*((Ay*Aux + Ax*Auy)/((rho*sigma)**2)
		- (Auxy + Auyx)/(rho*(sigma**2)))									# linearised shear strain
	vorticity = (Ax*Auy - Ay*Aux)/((rho*sigma)**2)\
		- (Auyx - Auxy)/(rho*(sigma**2))									# displacement vorticity

	return strain, vorticity

def strain_vorticity_grid(box_size, Ncases, grid_points, time, dt, w_traj,
	u_traj, sigma, r_cut):
	"""
	Calculates grids of (linearised) shear strain and displacement vorticity
	from coarse-grained displacement field.

	Resorts to neighbours grids to speed up calculations.

	Parameters
	----------
	box_size : float
		Length of the considered system's square box.
	Ncases : int
		Number of boxes in each direction to compute the shear strain and
		displacement vorticity grid.
	grid_points : array like of coordinates
		Grid points at which shear strain will be evaluated.
	time : int
		Frame at which shear strain and displacement vorticity will be
		calculated.
	dt : int
		Length of the interval of time for which the displacements are
		calculated.
	w_traj : active_particles.dat.Gsd
		Wrapped trajectory object.
	u_traj : active_particles.dat.Dat
		Unwrapped trajectory object.
	sigma : float
		Length scale of the spatial extent of the coarse graining function.
	r_cut : float
		Cut-off radius for coarse graining function.

	Returns
	-------
	sgrid : 2D array like
		Shear strain grid.
	cgrid : 2D array like
		Displacement vorticity grid.

	Output
	------
	Prints neighbours grid computation time.
	"""

	# NEIGHBOURS GRID

	startTime0 = datetime.now()	# start time for neighbours grid computation

	positions = w_traj.position(time +
		dt*get_env('ENDPOINT', default=False, vartype=bool))		# array of wrapped particle positions
	neighbours_grid = NeighboursGrid(positions, box_size, r_cut)	# neighbours grid

	print("Neighbours grid computation time (time = %e): %s" %
		(time, datetime.now() - startTime0))	# neighbours grid computation time

	# SHEAR STRAIN AND DISPLACEMENT VORTICITY GRIDS CALCULATION

	displacements = u_traj.displacement(time, time + dt)	# array of particle displacements

	sgrid, cgrid = tuple(np.transpose(list(map(
		lambda point: strain_vorticity(point, time, dt, positions,
		displacements, sigma, r_cut, box_size, neighbours_grid),
		grid_points))))	# shear strain and displacement vorticity lists

	correct_grid = lambda grid: np.transpose(
		np.reshape(grid, (Ncases, Ncases)))[::-1]	# get grids with the same orientation as positions
	return correct_grid(sgrid), correct_grid(cgrid)	# shear strain and displacement vorticity grids

def strain_vorticity_fftsqnorm_grid(box_size, centre, Ncases, time, dt,
	w_traj, u_traj):
	"""
	Calculates grids of square norm of fast Fourier transforms of (linearised)
	shear strain and displacement vorticity from fast Fourier transform of
	displacement grid.

	Parameters
	----------
	box_size : float
		Length of the considered system's square box.
	centre : float array
		Centre of the box.
	Ncases : int
		Number of boxes in each direction to compute the displacements.
	time : int
		Frame at which displacements will be calculated.
	dt : int
		Length of the interval of time for which the displacements are
		calculated.
	w_traj : active_particles.dat.Gsd
		Wrapped trajectory object.
	u_traj : active_particles.dat.Dat
		Unwrapped trajectory object.

	Returns
	-------
	FFTsgridsqnorm : 2D Numpy array
		Square norm of shear strain fast Fourier transform grid.
	FFTcgridsqnorm : 2D Numpy array
		Square norm of displacement vorticity fast Fourier transform grid.
	"""

	# DISPLACEMENT GRID FOURIER TRANSFORM

	ugrid = displacement_grid(box_size, centre, Ncases, time, dt,
		w_traj, u_traj)							# displacment grid
	FFTugrid = np.fft.fft2(ugrid, axes=(0, 1))	# displacement grid Fourier transform

	# SHEAR STRAIN AND DISPLACEMENT VORTICITY FOURIER TRANSFORM CALCULATION

	wave_vectors = wave_vectors_2D(Ncases, Ncases, d=box_size/Ncases)	# wave vectors corresponding to displacement grid Fourier transform
	eikxm1 = np.exp(1j*wave_vectors[:, :, 0]) - 1						# grid exponential of i times wave vectors' x-coordinates - 1
	eikym1 = np.exp(1j*wave_vectors[:, :, 1]) - 1						# grid exponential of i times wave vectors' y-coordinates - 1

	FFTsgrid = (eikxm1*FFTugrid[:, :, 1] + eikym1*FFTugrid[:, :, 0])/2	# shear strain Fourier transform grid
	FFTcgrid = eikxm1*FFTugrid[:, :, 1] - eikym1*FFTugrid[:, :, 0]		# displacement vorticity Fourier transform grid

	return np.conj(FFTsgrid)*FFTsgrid, np.conj(FFTcgrid)*FFTcgrid

def strain_OVITO_fftsqnorm_grid(box_size, centre, Ncases, time, dt, w_traj):
	"""
	Calculates grid of square norm of fast Fourier transforms of (linearised)
	shear strain with OVITO.

	Parameters
	----------
	box_size : float
		Length of the considered system's square box.
	centre : float array
		Centre of the box.
	Ncases : int
		Number of boxes in each direction to compute the displacements.
	time : int
		Frame at which displacements will be calculated.
	dt : int
		Length of the interval of time for which the displacements are
		calculated.
	w_traj : active_particles.dat.Gsd
		Wrapped trajectory object.

	Returns
	-------
	FFTsgridsqnorm : 2D Numpy array
		Square norm of shear strain fast Fourier transform grid.
	"""

	sgrid = w_traj.to_grid(
		time + dt*get_env('ENDPOINT', default=False, vartype=bool),
		w_traj.xy_strain(time, time + dt),
		Ncases=Ncases, box_size=box_size, centre=centre)	# shear strain grid
	FFTsgrid = np.fft.fft2(sgrid, axes=(0, 1))				# shear strain Fourier transform grid

	return np.conj(FFTsgrid)*FFTsgrid

def suptitle():
	"""
	Returns suptitles for figures.
	"""

	return (
		r'$N=%.2e, \phi=%1.2f, \tilde{v}=%.2e, \tilde{\nu}_r=%.2e,$'
		% (parameters['N'], parameters['density'], parameters['vzero'],
		parameters['dr'])
		+ r'$\bar{N}=%.2e, \Delta t=%.2e, nD_0 \Delta t=%.2e, L=%.2e$'
		% (Nmean, dt*parameters['period_dump']*parameters['time_step'],
		nD0*dt*parameters['period_dump']*parameters['time_step'], box_size)
		+ '\n' + r'$x_0=%.2e, y_0=%.2e,$' % centre
		+ r'$S_{init}=%.2e$' % init_frame
		+ r'$, S_{max}=%.2e, N_{cases}=%.2e$' % (int_max, Ncases)
		+ r'$, dL = %.2e a$' % (np.sqrt(Nmean)/Ncases))

def plot(grid, corr, box_size, var, naming_standard):
	"""
	Plots variable grid and correlations.

	Parameters
	----------
	grid : 2D array-like
		Variable grid.
	corr : 2D array-like
		Variable correlation grid.
	box_size : float
		Length of the box in one dimension.
	var : string
		Name of variable.
	naming_standard : active_particles.naming standard
		Standard naming object.

	Returns
	-------
	fig : matplotlib figure
		Figure.
	ax : matplotlib axis
		Axis.
	gridcircle : active_particles.plot.mpl_tools.GridCircle
		Grid circle object.
	"""

	# GRID AND CORRELATION FIGURE

	cmap = plt.cm.jet

	fig, ax = plt.subplots(1, 2)	# 1 x 2 figure

	fig.set_size_inches(16, 16)		# figure size
	fig.subplots_adjust(wspace=0.4)	# width space
	fig.subplots_adjust(hspace=0.3)	# height space

	fig.suptitle(suptitle()
		+ r'$, r_{cut}=%.2e, \sigma=%.2e$' % (r_cut, sigma))

	# VARIABLE GRID

	Smin = -2*np.std(grid)
	Smax = 2*np.std(grid)

	SvNorm = colors.Normalize(vmin=Smin, vmax=Smax)
	SscalarMap = cmx.ScalarMappable(norm=SvNorm, cmap=cmap)

	ax[0].imshow(grid, cmap=cmap, norm=SvNorm,
		extent=[-box_size/2, box_size/2, -box_size/2, box_size/2])

	ax[0].set_xlabel(r'$x$')
	ax[0].set_ylabel(r'$y$')
	ax[0].set_title('2D ' + r'$%s$' % var)

	divider0 = make_axes_locatable(ax[0])
	cax0 = divider0.append_axes('right', size='5%', pad=0.05)
	cb0 = mpl.colorbar.ColorbarBase(cax0, cmap=cmap, norm=SvNorm,
		orientation='vertical')
	cb0.set_label(r'$%s$' % var, labelpad=20, rotation=270)

	# 2D CORRELATIONS

	C = 'C_{%s%s}' % (var, var)

	Cmin = np.max((np.min(corr), _c_min))
	Cmax = np.max((np.min(corr), _c_max))

	CvNorm = colors.Normalize(vmin=Cmin, vmax=Cmax)

	cgrid = CorGrid(corr, box_size, display_size=2*r_max)

	ax[1].imshow(cgrid.display_grid.grid, cmap=cmap, norm=CvNorm,
		extent=[-r_max, r_max, -r_max, r_max])

	ax[1].set_xlabel(r'$x$')
	ax[1].set_ylabel(r'$y$')
	ax[1].set_title('2D ' + r'$%s$' % C)

	divider1 = make_axes_locatable(ax[1])
	cax1 = divider1.append_axes("right", size="5%", pad=0.05)
	cb1 = mpl.colorbar.ColorbarBase(cax1, cmap=cmap, norm=CvNorm,
		orientation='vertical')
	cb1.set_label(r'$%s$' % C, labelpad=20, rotation=270)

	if get_env('SAVE', default=False, vartype=bool):	# SAVE mode
		image_name, = naming_standard.image().filename(**attributes)
		fig.savefig(joinpath(data_dir, image_name))

	# GRID CIRCLE FIGURE

	gridcircle = GridCircle(cgrid.display_grid.grid,
		extent=[-r_max, r_max, -r_max, r_max])
	fig_gc, (ax_grid, ax_plot), cb_gc = gridcircle.get_fig_ax_cmap()

	fig_gc.set_size_inches(16, 16)		# figure size
	fig_gc.subplots_adjust(wspace=0.4)	# width space
	fig_gc.subplots_adjust(hspace=0.3)	# height space

	fig_gc.suptitle(suptitle()
		+ r'$, r_{cut}=%.2e, \sigma=%.2e$' % (r_cut, sigma))

	ax_grid.set_xlabel(r'$x$')
	ax_grid.set_ylabel(r'$y$')
	ax_grid.set_title('2D ' + r'$%s$' % C)
	cb_gc.set_label(r'$%s$' % C, labelpad=20, rotation=270)

	ax_plot.set_xlabel(r'$\theta$')
	ax_plot.set_ylabel(r'$%s(r, \theta)$' % C)

	# RETURN

	return fig, ax, gridcircle

class StrainCorrelations:
	"""
	Manipulate and plot strain correlations from its Fourier transform.
	"""

	def __init__(self, wave_vectors, FFTCss):
		"""
		Parameters
		----------
		wave_vectors : (_, _, 2) array-like
			Wave vectors at which Fourier transform was calculated.
		FFTCss : 2D Numpy array
			Strain correlations Fourrier transform.
		"""

		self.wave_vectors = np.array(wave_vectors)
		self.strain_correlations_FFT = np.array(FFTCss)

	def strain_correlations(self, r_cut=0):
		"""
		Computes strain correlations from inverse fast Fourier transform of
		self.strain_correlations_FFT with low wave lengths Gaussian cut at
		r_cut.

		Parameters
		----------
		r_cut : float
			Wave length Gaussian cut-off radius, equivalent to coarse-graining
			cut-off radius. (default: 0)

		Returns
		-------
		sc : Numpy array
			Strain correlations.
		"""

		sc = (FFT2Dfilter(self.strain_correlations_FFT,
			wave_vectors=self.wave_vectors)
			.gaussian_filter(np.sqrt(2)*r_cut)
			.get_signal()).real
		return sc/sc[0, 0]	# correlation normalisation

	def plot(self, box_size, r_max_css, av_p_sep,
		points_x_c44=_points_x_c44, points_theta_c44=_points_theta_c44,
		y_min_c44=_y_min_c44, y_max_c44=_y_max_c44,
		r_min_c44=_r_min_c44, r_max_c44=_r_max_c44,
		smooth=0, r_cut=0, linear_interpolation=False,
		theta=0, points_x_theta=None,
		y_min_theta=_y_min_theta, y_max_theta=_y_max_theta,
		r_min_theta=_r_min_theta, r_max_theta=_r_max_theta,
		superimpose_c44=False):
		"""
		Plots strain correlations with slider for cut-off radius.

		Parameters
		----------
		box_size : float
			System length.
		r_max_css : float
			Maximum radius in infinite norm for strain correlations plots.
		av_p_sep : float
			Average particle separation.
		points_x_c44 : int
			Number of radii at which to evaluate integrated strain correlation.
			(default: active_particles.plot.c44._points_x)
		points_theta_c44 : int
			Number of angles to evaluate integrated strain correlation.
			(default: active_particles.plot.c44._points_theta)
		y_min_c44 : float
			Minimum plot value for C44.
			(default: active_particles.analysis.css._y_min_c44)
		y_max_c44 : float
			Maximum plot value for C44.
			(default: active_particles.analysis.css._y_max_c44)
		r_min_c44 : float
			Minimum radius in average particle separation for C44 calculation.
			(default: active_particles.analysis.css._r_min_c44)
		r_max_c44 : float
			Maximum radius in average particle separation for C44 calculation.
			(default: active_particles.analysis.css._r_max_c44)
		smooth : float
			C44 Gaussian smoothing length scale in units of average particle
			separation. (default: 0)
        linear_interpolation : bool
            Get value on grid by linear interpolation of neighbouring grid
			boxes. (default: False)
		r_cut : float
			Initial wave length Gaussian cut-off radius in units of average
			particles separation. (default: 0)
		theta : float
			Angle at which to plot Css(r, theta). (default: 0)
		points_x_theta : int
			Number of radii at which to evaluate Css(r, theta). (default: None)
			NOTE: if points_x_theta == None, then it is taken equal to
			      points_x_c44.
		y_min_theta : float
			Minimum plot value for Css(r, theta).
			(default: active_particles.analysis.css._y_min_theta)
		y_max_theta : float
			Maximum plot value for Css(r, theta).
			(default: active_particles.analysis.css._y_max_theta)
		r_min_theta : float
			Minimum radius in average particle separation for Css(r, theta)
			calculation.
			(default: active_particles.analysis.css._r_min_theta)
		r_max_theta : float
			Maximum radius in average particle separation for Css(r, theta)
			calculation.
			(default: active_particles.analysis.css._r_max_theta)
		superimpose_c44 : bool
			Superimpose C44 curve to Css(r, theta) curve. (default: False)
		"""

		self.box_size = box_size
		self.av_p_sep = av_p_sep

		self.r_max_css = r_max_css

		self.r_cut = r_cut	# Gaussian cut-off radius in units of average particle separation

		self.cor_name = 'C_{\\varepsilon_{xy}\\varepsilon_{xy}}'	# name of plotted correlation

		grid = self.strain_correlations_corgrid()	# correlation grid

		# CSS FIGURE

		self.fig_css, self.ax_css = plt.subplots()
		self.ax_css.set_title(
			'2D ' + r'$%s(r_{cut}/a = %.2e)$' % (self.cor_name, self.r_cut))

		Cmin = np.max((np.min(grid.grid), _c_min))
		Cmax = np.max((np.min(grid.grid), _c_max))
		cmap = plt.cm.jet
		CvNorm = colors.Normalize(vmin=Cmin, vmax=Cmax)

		self.grid_plot = self.ax_css.imshow(grid.display_grid.grid,
			cmap=cmap, norm=CvNorm, extent=
			[-self.r_max_css, self.r_max_css, -self.r_max_css, self.r_max_css])

		self.colormap = plt.colorbar(self.grid_plot)	# color bar

		self.ax_slider = make_axes_locatable(self.ax_css).append_axes('bottom',
			size='5%', pad=0.6)	# slider Axes

		self.slider = Slider(self.ax_slider, r'$r_{cut}/a$', 0,
			self.r_max_css/self.av_p_sep, valinit=self.r_cut)	# slider
		self.slider.on_changed(self.update_r_cut)				# call update_r_cut when slider value is changed

		# GRID CIRCLE FIGURE

		self.linear_interpolation = linear_interpolation

		self.grid_circle = GridCircle(grid.display_grid.grid, extent=
			[-self.r_max_css/self.av_p_sep, self.r_max_css/self.av_p_sep,
				-self.r_max_css/self.av_p_sep, self.r_max_css/self.av_p_sep],
			min=Cmin, max=Cmax, linear_interpolation=self.linear_interpolation)
		self.grid_circle.ax_grid.set_title(
			'2D ' + r'$%s(r_{cut}/a = %.2e)$' % (self.cor_name, self.r_cut))

		# C44 FIGURE

		self.points_x_c44 = points_x_c44
		self.points_theta_c44 = points_theta_c44
		self.y_min_c44 = y_min_c44
		self.y_max_c44 = y_max_c44
		self.r_min_c44 = r_min_c44
		self.r_max_c44 = r_max_c44
		self.smooth = smooth	# C44 Gaussian smoothing length scale in units of average particle separation

		self.toC44 = Css2DtoC44(self.box_size,
			self.points_x_c44, self.points_theta_c44,
			self.av_p_sep*self.r_min_c44, self.av_p_sep*self.r_max_c44,
			linear_interpolation=self.linear_interpolation)

		self.fig_c44, self.ax_c44 = plt.subplots()
		self.ax_c44.set_title(r'$r_{cut}/a = %.2e$' % self.r_cut)
		self.ax_c44.set_xlim([self.r_min_c44, self.r_max_c44])
		self.ax_c44.set_ylim([self.y_min_c44, self.y_max_c44])

		self.c44 = self.css2Dtoc44(grid.grid)				# list of [r, C44(r)]
		self.line_c44, = self.ax_c44.plot(
			self.c44[:, 0]/self.av_p_sep, self.c44[:, 1])	# C44 plotting line

		# CSS(r, theta) FIGURE

		self.theta = theta
		self.points_x_theta = (points_x_theta if points_x_theta != None
			else self.points_x_c44)
		self.y_min_theta = y_min_theta
		self.y_max_theta = y_max_theta
		self.r_min_theta = r_min_theta
		self.r_max_theta = r_max_theta
		self.superimpose_c44 = superimpose_c44

		self.toCsstheta = Css2DtoCsstheta(self.box_size, self.points_x_theta,
			self.av_p_sep*self.r_min_theta, self.av_p_sep*self.r_max_theta,
			linear_interpolation=self.linear_interpolation)

		self.fig_theta, self.ax_theta = plt.subplots()
		self.ax_theta.set_title(r'$r_{cut}/a = %.2e$' % self.r_cut)
		self.ax_theta.set_xlim([self.r_min_theta, self.r_max_theta])
		self.ax_theta.set_ylim([self.y_min_theta, self.y_max_theta])

		self.csstheta = self.css2Dtocsstheta(grid.grid)				# list of [r, Css(r, theta)]
		self.line_csstheta, = self.ax_theta.plot(
			self.csstheta[:, 0]/self.av_p_sep, self.csstheta[:, 1])	# Css(r, theta) plotting line

		if self.superimpose_c44:								# SUPERIMPOSE_C44 mode
			self.line_c44_superimposed, = self.ax_theta.plot(
				self.c44[:, 0]/self.av_p_sep, self.c44[:, 1])	# C44 plotting line superimposed on Css(r, theta)

	# METHODS CALLED BY SELF.PLOT()

	def strain_correlations_corgrid(self):
		"""
		Define strain correlations grid, with wave length cut-off radius
		self.r_cut equivalent to coarse-graining cut-off radius.

		Returns
		-------
		corgrid : active_particles.analysis.correlations.CorGrid
			Strain correlations CorGrid object.
		"""

		sc = self.strain_correlations(r_cut=self.r_cut*self.av_p_sep)	# strain correlations grid
		return CorGrid(sc, self.box_size, display_size=2*self.r_max_css)

	def css2Dtoc44(self, css2D):
		"""
		Returns strain correlations projected on cos(4 \\theta) from strain
		correlations grid, smoothed over self.smooth*self.av_p_sep.

		Parameters
		----------
		css2D : array-like
			Strain correlations grid.

		Returns
		-------
		c44 : Numpy array
			Array of position and projection of strain correlations on
			cos(4 \\theta) at this position.
		"""

		return self.toC44.get_C44(css2D, smooth=self.av_p_sep*self.smooth)	# list of [r, C44(r)]

	def css2Dtocsstheta(self, css2D):
		"""
		Returns strain correlations along half-line of origin (0, 0) and
		inclination self.theta.

		Parameters
		----------
		css2D : array-like
			Strain correlations grid.

		Returns
		-------
		csstheta : Numpy array
			Array of radii r and values of strain correlations (r, theta).
		"""

		return self.toCsstheta.get_Csstheta(css2D, self.theta)

	def update_r_cut(self, val):
		"""
		Updates cut-off radius on slider change.
		"""

		self.r_cut = self.slider.val	# new cut-off radius
		self.draw()						# updates figure

	def draw(self):
		"""
		Updates line position and strain correlations grid according to cut-off
		radius self.r_cut.
		"""

		self.ax_css.set_title(
			'2D ' + r'$%s(r_{cut}/a = %.2e)$' % (self.cor_name, self.r_cut))	# update title
		self.grid_circle.ax_grid.set_title(
			'2D ' + r'$%s(r_{cut} = %.2e)$' % (self.cor_name, self.r_cut))		# update title
		self.ax_c44.set_title(r'$r_{cut}/a = %.2e$' % self.r_cut)				# update title
		self.ax_theta.set_title(r'$r_{cut}/a = %.2e$' % self.r_cut)				# update title

		grid = self.strain_correlations_corgrid()		# new grid
		self.grid_plot.set_data(grid.display_grid.grid)	# plot grid
		self.grid_plot.figure.canvas.draw()				# update plot

		self.grid_circle.update_grid_plot(grid.display_grid.grid)	# plot grid

		self.c44 = self.css2Dtoc44(grid.grid)		# update C44 values
		self.line_c44.set_ydata(self.c44[:, 1])		# plot C44
		self.line_c44.figure.canvas.draw()			# update plot

		self.csstheta = self.css2Dtocsstheta(grid.grid)		# update Css(r, theta) values
		self.line_csstheta.set_ydata(self.csstheta[:, 1])	# plot Css(r, theta)
		self.line_csstheta.figure.canvas.draw()				# update plot
		if self.superimpose_c44: self.line_c44_superimposed.set_ydata(
			self.c44[:, 1])									# update plot

class Css2DtoC44:
    """
    Calculates C44 as projection of 2D shear strain correlations on
    cos(4\\theta).
    """

    def __init__(self, box_size, points_x, points_theta,
        r_min, r_max, linear_interpolation=False):
        """
        Sets parameters for C44 integration.

        Parameters
        ----------
        box_size : float
            Size of the square box to consider.
        points_x : int
            Number of radii at which to evaluate integrated strain correlation.
        points_theta : int
            Number of angles to evaluate integrated strain correlation.
        r_min : float
            Minimum radius.
        r_max : float
            Maximum radius.
        linear_interpolation : bool
            Get value on grid by linear interpolation of neighbouring grid
			boxes. (default: False)
        """

        self.box_size = box_size

        self.points_x = points_x
        self.points_theta = points_theta

        self.r_min = r_min
        self.r_max = r_max
        self.c44_x = np.linspace(self.r_min, self.r_max, self.points_x)

        self.linear_interpolation = linear_interpolation

    def get_C44(self, Css2D, smooth=0):
        """
        From 2D strain correlations Css2D, returns values of C44 at
        self.points_x radii between r_max and r_min.

        Parameters
        ----------
        Css2D : 2D array-like
            Shear strain correlation grid.
		smooth : float
			C44 Gaussian smoothing length scale. (default: 0)

        Returns
        -------
        C44 : 2D Numpy array
            List of [r, C44(r)].
        """

        self.css2Dgrid = CorGrid(Css2D, self.box_size)	# shear strain 2D CorGrid object

        self.c44 = np.array(list(map(
            lambda r: self.css2Dgrid.integrate_over_angles(r,
            	projection=lambda theta: np.cos(4*theta)/np.pi,
            	points_theta=self.points_theta,
				linear_interpolation=self.linear_interpolation),
            self.c44_x)))

        return np.array(list(map(lambda x, y: [x, y],
			*(self.c44_x, gaussian_smooth_1D(self.c44_x, self.c44, smooth)))))

class Css2DtoCsstheta:
    """
    Calculates Css(r, theta) as values of strain correlations along the
	half-line of origin (0, 0) and inclination of angle theta.
    """

    def __init__(self, box_size, points_x, r_min, r_max,
		linear_interpolation=False):
        """
        Sets parameters for Css(r, theta) calculation.

        Parameters
        ----------
        box_size : float
            Size of the square box to consider.
        points_x : int
            Number of radii at which to evaluate strain correlations.
        r_min : float
            Minimum radius.
        r_max : float
            Maximum radius.
        linear_interpolation : bool
            Get value on grid by linear interpolation of neighbouring grid
			boxes. (default: False)
        """

        self.box_size = box_size

        self.points_x = points_x

        self.r_min = r_min
        self.r_max = r_max
        self.csstheta_x = np.linspace(self.r_min, self.r_max, self.points_x)

        self.linear_interpolation = linear_interpolation

    def get_Csstheta(self, Css2D, theta=0):
        """
        From 2D strain correlations Css2D, returns values of Css(r, theta) at
        self.points_x radii between r_max and r_min.

        Parameters
        ----------
        Css2D : 2D array-like
            Shear strain correlation grid.
		theta : float
			Inclination of half-line in Css2D grid. (default: 0)

        Returns
        -------
        Csstheta : 2D Numpy array
            List of [r, Css(r, theta)].
        """

        self.css2Dgrid = CorGrid(Css2D, self.box_size)	# shear strain 2D CorGrid object
        self.csstheta = np.array(list(map(
			lambda r: [r, self.css2Dgrid.get_value_polar(r, theta,
				linear_interpolation=self.linear_interpolation)],
            self.csstheta_x)))

        return self.csstheta

def plot_fft():
	"""
	Plots shear strain correlations.

	Returns
	-------
	sc : active_particles.analysis.css.StrainCorrelations
		Strain correlations object.
	"""

	cor_name = r'$C_{\varepsilon_{xy}\varepsilon_{xy}}$'	# name of plotted correlation

	wave_vectors = wave_vectors_2D(*FFTsgridsqnorm.shape, d=box_size/Ncases)	# wave vectors at which Fourier transform was calculated

	sc = StrainCorrelations(wave_vectors, FFTsgridsqnorm)

	av_p_sep = box_size/np.sqrt(Nmean)
	sc.plot(parameters['box_size'], r_max, av_p_sep,
		points_x_c44=points_x_c44, points_theta_c44=points_theta_c44,
		y_min_c44=y_min_c44, y_max_c44=y_max_c44,
		r_min_c44=r_min_c44, r_max_c44=r_max_c44,
		smooth=smooth, r_cut=r_cut_fourier,
		linear_interpolation=linear_interpolation,
		theta=theta, points_x_theta=points_x_theta,
		y_min_theta=y_min_theta, y_max_theta=y_max_theta,
		r_min_theta=r_min_theta, r_max_theta=r_max_theta,
		superimpose_c44=superimpose_c44)

	# CSS FIGURE

	sc.fig_css.set_size_inches(16, 16)

	sc.fig_css.suptitle(suptitle())

	sc.ax_css.set_xlabel(r'$x$')
	sc.ax_css.set_ylabel(r'$y$')
	sc.colormap.set_label(cor_name, labelpad=20, rotation=270)

	if get_env('SAVE', default=False, vartype=bool):	# SAVE mode
		sc.ax_slider.set_visible(False)	# hide slider
		image_name, = naming_Css.image().filename(
			**{**attributes, 'r_cut': r_cut_fourier})
		sc.fig_css.savefig(joinpath(data_dir, image_name))
		sc.ax_slider.set_visible(True)	# show slider

	# GRID CIRCLE FIGURE

	fig_gc, (ax_grid, ax_plot), cb_gc = sc.grid_circle.get_fig_ax_cmap()

	sc.grid_circle.fig.suptitle(suptitle())

	sc.grid_circle.fig.set_size_inches(16, 16)		# figure size
	sc.grid_circle.fig.subplots_adjust(wspace=0.4)	# width space
	sc.grid_circle.fig.subplots_adjust(hspace=0.3)	# height space

	sc.grid_circle.ax_grid.set_xlabel(r'$x/a$')
	sc.grid_circle.ax_grid.set_ylabel(r'$y/a$')
	sc.grid_circle.colormap.set_label(cor_name, labelpad=20, rotation=270)

	sc.grid_circle.ax_plot.set_xlabel(r'$\theta$')
	sc.grid_circle.ax_plot.set_ylabel(cor_name + r'$(r, \theta)$')

	# C44 FIGURE

	sc.fig_c44.suptitle(suptitle())

	sc.fig_c44.set_size_inches(16, 16)	# figure size

	sc.ax_c44.set_xlabel(r'$r/a$' + ' ' + r'$(a = L/\sqrt{\bar{N}})$')
	sc.ax_c44.set_ylabel(r'$C_4^4(r) = \frac{1}{\pi}\int_0^{2\pi}d\theta$'
		+ ' ' + cor_name + r'$(r, \theta)$'
		+ ' ' + r'$\cos4\theta$'
		+ ' ' + (r'($\sigma_{smooth}/a=%.2e$)' % smooth if smooth!=0 else ''))
	sc.ax_c44.set_xlim([r_min_c44, r_max_c44])
	sc.ax_c44.set_ylim([y_min_c44, y_max_c44])
	sc.ax_c44.set_yscale('log')
	sc.ax_c44.set_xscale('log')

	if get_env('FITTING_LINE', default=False, vartype=bool):	# FITTING_LINE mode
		sc.fl_c44 = FittingLine(sc.ax_c44, slope0_c44, slope_min_c44,
			slope_max_c44, x_fit='(r/a)', y_fit='C_4^4(r)')		# add fitting line to plot

	# CSS(r, theta) FIGURE

	sc.fig_theta.suptitle(suptitle())

	sc.fig_theta.set_size_inches(16, 16)	# figure size

	sc.ax_theta.set_xlabel(r'$r/a$' + ' ' + r'$(a = L/\sqrt{\bar{N}})$')
	sc.ax_theta.set_xlim([r_min_theta, r_max_theta])
	sc.ax_theta.set_ylim([y_min_theta, y_max_theta])
	sc.ax_theta.set_yscale('log')
	sc.ax_theta.set_xscale('log')

	if sc.superimpose_c44:	# SUPERIMPOSE_C44 mode
		sc.ax_theta.set_ylabel(r'$C(r)$')
		plt.sca(sc.ax_theta)
		sc.ax_theta.add_artist(plt.legend(handles=[
			Line2D([0], [0], color=sc.line_csstheta.get_color(),
				label='%s' % cor_name + r'$(r, \theta=%.2e)$' % theta),
			Line2D([0], [0], color=sc.line_c44_superimposed.get_color(),
				label=r'$C_4^4(r)$')]))
	else:
		sc.ax_theta.set_ylabel('%s' % cor_name + r'$(r, \theta=%.2e)$' % theta)

	if get_env('FITTING_LINE', default=False, vartype=bool):					# FITTING_LINE mode
		sc.fl_theta = FittingLine(sc.ax_theta, slope0_theta, slope_min_theta,
			slope_max_theta, x_fit='(r/a)', y_fit='%s(r, \\theta)' % cor_name)	# add fitting line to plot

	# RETURN

	return sc

# SCRIPT

if __name__ == '__main__':	# executing as script

	# VARIABLE DEFINITIONS

	data_dir = get_env('DATA_DIRECTORY', default=getcwd())	# data directory

	wrap_file_name = get_env('WRAPPED_FILE',
		default=joinpath(data_dir, naming.wrapped_trajectory_file))		# wrapped trajectory file (.gsd)
	unwrap_file_name = get_env('UNWRAPPED_FILE',
		default=joinpath(data_dir, naming.unwrapped_trajectory_file))	# unwrapped trajectory file (.dat)

	dt = get_env('TIME', default=-1, vartype=int)	# lag time for displacement

	init_frame = get_env('INITIAL_FRAME', default=-1, vartype=int)	# frame to consider as initial
	int_max = get_env('INTERVAL_MAXIMUM', default=1, vartype=int)	# maximum number of intervals of length dt considered in correlations calculations

	parameters_file = get_env('PARAMETERS_FILE',
		default=joinpath(data_dir, naming.parameters_file))	# simulation parameters file
	with open(parameters_file, 'rb') as param_file:
		parameters = pickle.load(param_file)				# parameters hash table

	r_cut = parameters['a']*get_env('R_CUT', default=_r_cut, vartype=float)	# cut-off radius for coarse graining function
	sigma = get_env('SIGMA', default=r_cut, vartype=float)					# length scale of the spatial extent of the coarse graining function

	box_size = get_env('BOX_SIZE', default=parameters['box_size'],
		vartype=float)									# size of the square box to consider
	centre = (get_env('X_ZERO', default=0, vartype=float),
		get_env('Y_ZERO', default=0, vartype=float))	# centre of the box

	prep_frames = ceil(parameters['prep_steps']/parameters['period_dump'])	# number of preparation frames (FIRE energy minimisation)

	Ncases = get_env('N_CASES', default=ceil(np.sqrt(parameters['N'])),
		vartype=int)	# number of boxes in each direction to compute the shear strain and displacement vorticity grid

	Nentries = parameters['N_steps']//parameters['period_dump']		# number of time snapshots in unwrapped trajectory file
	init_frame = int(Nentries/2) if init_frame < 0 else init_frame	# initial frame
	Nframes = Nentries - init_frame									# number of frames available for the calculation

	dt = Nframes + dt if dt <= 0 else dt	# length of the interval of time for which displacements are calculated

	times = linframes(init_frame, Nentries - dt, int_max)	# frames at which shear strain is calculated

	mode = get_env('MODE', default='fourier')				# computation mode
	if not(mode in ('real', 'fourier', 'ovito')):
		raise ValueError('Mode %s is not known.' % mode)	# mode is not known

	# NAMING

	attributes = {'density': parameters['density'],
		'vzero': parameters['vzero'], 'dr': parameters['dr'],
		'N': parameters['N'], 'init_frame': init_frame, 'dt': dt,
		'int_max': int_max, 'Ncases': Ncases, 'r_cut': r_cut, 'sigma': sigma,
		'box_size': box_size, 'x_zero': centre[0], 'y_zero': centre[1]}	# attributes displayed in filenames
	naming_Css = naming.Css(mode=mode)									# Css naming object
	Css_filename, = naming_Css.filename(**attributes)					# Css filename
	naming_Ccc = naming.Ccc(mode=mode)									# Ccc naming object
	Ccc_filename, = naming_Ccc.filename(**attributes)					# Css filename

	# STANDARD OUTPUT

	if 'SLURM_JOB_ID' in envvar:	# script executed from Slurm job scheduler
		slurm_output(joinpath(data_dir, 'out'), naming_Css, attributes)

	# MODE SELECTION

	if get_env('COMPUTE', default=False, vartype=bool):	# COMPUTE mode

		startTime = datetime.now()

		if mode == 'real':	# calculation of shear strain and vorticity in real space

			grid_points = np.array([(x, y) for x in\
				relative_positions(np.linspace(- box_size*(1 - 1./Ncases)/2,
				box_size*(1 - 1./Ncases)/2, Ncases, endpoint=True) + centre[0],
				0, parameters['box_size'])\
				for y in\
				relative_positions(np.linspace(- box_size*(1 - 1./Ncases)/2,
				box_size*(1 - 1./Ncases)/2, Ncases, endpoint=True) + centre[1],
				0, parameters['box_size'])\
				])	# grid points at which shear strain will be evaluated

			display_grid = get_env('DISPLAY_GRID', default=0, vartype=int)	# index of map in list of variable maps to display

			# SHEAR STRAIN, DISPLACEMENT VORTICITY AND THEIR CORRELATIONS

			with open(wrap_file_name, 'rb') as wrap_file,\
				open(unwrap_file_name, 'rb') as unwrap_file:	# opens wrapped and unwrapped trajectory files

				w_traj = Gsd(wrap_file, prep_frames=prep_frames)	# wrapped trajectory object
				u_traj = Dat(unwrap_file, parameters['N'])			# unwrapped trajectory object
				Sgrid, Cgrid = tuple(np.transpose(list(map(lambda time:
					strain_vorticity_grid(parameters['box_size'], Ncases,
					grid_points, time, dt, w_traj, u_traj, sigma, r_cut),
					times)), (1, 0, 2, 3)))							# lists of shear strain and displacement vorticity correlations

			Css2D, Ccc2D = tuple(map(corField2D_scalar_average,
				[Sgrid, Cgrid]))	# shear strain and displacement vorticity fields correlations

			# SAVING

			sgrid = Sgrid[display_grid]
			cgrid = Cgrid[display_grid]

			with open(joinpath(data_dir, Css_filename),
				'wb') as Css_dump_file,\
				open(joinpath(data_dir, Ccc_filename),
				'wb') as Ccc_dump_file:
				pickle.dump([sgrid, Css2D], Css_dump_file)
				pickle.dump([cgrid, Ccc2D], Ccc_dump_file)

		elif mode == 'fourier':	# calculation of shear strain and vorticity in Fourier space

			# SHEAR STRAIN AND DISPLACEMENT VORTICITY FAST FOURIER TRANSFORMS

			with open(wrap_file_name, 'rb') as wrap_file,\
				open(unwrap_file_name, 'rb') as unwrap_file:	# opens wrapped and unwrapped trajectory files

				w_traj = Gsd(wrap_file, prep_frames=prep_frames)	# wrapped trajectory object
				u_traj = Dat(unwrap_file, parameters['N'])			# unwrapped trajectory object
				FFTsgridsqnorm, FFTcgridsqnorm = tuple(np.mean(np.transpose(
					list(map(lambda time:
					strain_vorticity_fftsqnorm_grid(
						box_size, centre, Ncases, time, dt, w_traj, u_traj),
					times)), (1, 0, 2, 3)), axis=1))				# average square norm of shear strain and displacement vorticity Fourier transforms

			# SAVING

			with open(joinpath(data_dir, Css_filename),
				'wb') as Css_dump_file,\
				open(joinpath(data_dir, Ccc_filename),
				'wb') as Ccc_dump_file:
				pickle.dump(FFTsgridsqnorm, Css_dump_file)
				pickle.dump(FFTcgridsqnorm, Ccc_dump_file)

		elif mode == 'ovito': # calculation of shear strain from OVITO

			# SHEAR STRAIN FAST FOURIER TRANSFORT

			with open(wrap_file_name, 'rb') as wrap_file:	# opens wrapped trajectory file

				w_traj = Gsd(wrap_file, prep_frames=prep_frames)	# wrapped trajectory object
				FFTsgridsqnorm = np.mean(list(map(lambda time:
					strain_OVITO_fftsqnorm_grid(
						box_size, centre, Ncases, time, dt, w_traj),
					times)), axis=0)								# average square norm of shear strain Fourier transform

			# SAVING

			with open(joinpath(data_dir, Css_filename), 'wb') as Css_dump_file:
				pickle.dump(FFTsgridsqnorm, Css_dump_file)

		# EXECUTION TIME

		print("Execution time: %s" % (datetime.now() - startTime))

	if get_env('PLOT', default=False, vartype=bool):	# PLOT mode

		# DATA

		if mode == 'real':	# shear strain and vorticity in real space
			with open(joinpath(data_dir, Css_filename),
				'rb') as Css_dump_file,\
				open(joinpath(data_dir, Ccc_filename),
				'rb') as Ccc_dump_file:
				sgrid, Css2D = pickle.load(Css_dump_file)
				cgrid, Ccc2D = pickle.load(Ccc_dump_file)
		else:				# shear strain in Fourier space
			with open(joinpath(data_dir, Css_filename), 'rb') as Css_dump_file:
				FFTsgridsqnorm = pickle.load(Css_dump_file)

	if get_env('PLOT', default=False, vartype=bool) or\
		get_env('SHOW', default=False, vartype=bool):	# PLOT or SHOW mode

		# PARAMETERS

		with open(wrap_file_name, 'rb') as wrap_file:
			w_traj = Gsd(wrap_file, prep_frames=prep_frames)		# wrapped trajectory object
			Nmean = np.mean(count_particles(
				w_traj, *times, box_size=box_size, centre=centre))	# average number of particles in box
		nD0 = nD0_active(
			Nmean, parameters['vzero'], parameters['dr'], box_size)	# product of particle density and active diffusion constant
		av_p_sep = parameters['box_size']/np.sqrt(Nmean)			# average particle separation

		# PLOT

		linear_interpolation = get_env('LINEAR_INTERPOLATION', default=False,
			vartype=bool)	# LINEAR_INTERPOLATION mode

		font_size = get_env('FONT_SIZE', default=_font_size, vartype=int)       # plot font size
		mpl.rcParams.update({'font.size': font_size})

		r_max = get_env('R_MAX', default=_r_max, vartype=float)	# half size of the box showed for 2D correlation
		r_max = box_size/2 if r_max < 0 else r_max

		if mode == 'real':	# shear strain and vorticity in real space

			fig_Css, ax_Css, gridcircle_Css = plot(sgrid, Css2D, box_size,
				'\epsilon_{xy}', naming_Css)	# plotting shear strain map and correlation
			fig_Ccc, ax_Ccc, gridcircle_Ccc = plot(cgrid, Ccc2D, box_size,
				'\omega', naming_Ccc)			# plotting displacement vorticity map and correlation

		else:				# shear strain in Fourier space

			points_x_c44 = get_env('POINTS_X_C44', default=_points_x_c44,
				vartype=int)							# number of radii at which to evaluate integrated strain correlation
			points_theta_c44 = get_env('POINTS_THETA_C44',
				default=_points_theta_c44, vartype=int)	# number of angles to evaluate integrated strain correlation

			y_min_c44 = get_env('Y_MIN_C44', default=_y_min_c44, vartype=float)	# minimum plot value for C44
			y_max_c44 = get_env('Y_MAX_C44', default=_y_max_c44, vartype=float)	# maximum plot value for C44

			r_min_c44 = get_env('R_MIN_C44', default=_r_min_c44, vartype=float)	# minimum radius in average particle separation for C44 calculation
			r_max_c44 = get_env('R_MAX_C44', default=_r_max_c44, vartype=float)	# maximum radius in average particle separation for C44 calculation

			smooth = get_env('SMOOTH', default=0, vartype=float)	# C44 Gaussian smoothing length scale in units of average particle separation

			r_cut_fourier = get_env('R_CUT_FOURIER', default=_r_cut_fourier,
				vartype=float)	# initial wave length Gaussian cut-off radius in units of average particle separation

			theta = np.pi*get_env('THETA', default=_theta, vartype=float)	# angle at which to evaluate Css(r, theta)

			points_x_theta = get_env('POINTS_X_THETA', default=points_x_c44,
				vartype=int)	# number of radii at which to evaluate Css(r, theta)

			y_min_theta = get_env('Y_MIN_THETA', default=_y_min_theta,
				vartype=float)	# minimum plot value for Css(r, theta)
			y_max_theta = get_env('Y_MAX_THETA', default=_y_max_theta,
				vartype=float)	# maximum plot value for Css(r, theta)

			r_min_theta = get_env('R_MIN_THETA', default=_r_min_theta,
				vartype=float)	# minimum radius in average particle separation for Css(r, theta) calculation
			r_max_theta = get_env('R_MAX_THETA', default=_r_max_theta,
				vartype=float)	# maximum radius in average particle separation for Css(r, theta) calculation

			superimpose_c44 = get_env('SUPERIMPOSE_C44',
				default=False, vartype=bool)	# superimpose C44 curve to Css(r, theta) curve

			if get_env('FITTING_LINE', default=False, vartype=bool):	# FITTING_LINE mode

				slope0_c44 = get_env('SLOPE_C44',
					default=_slope0_c44, vartype=float)		# initial slope for fitting line in C44 figure
				slope_min_c44 = get_env('SLOPE_MIN_C44',
					default=_slope_min_c44, vartype=float)	# minimum slope for fitting line in C44 figure
				slope_max_c44 = get_env('SLOPE_MAX_C44',
					default=_slope_max_c44, vartype=float)	# maximum slope for fitting line in C44 figure

				slope0_theta = get_env('SLOPE_THETA',
					default=slope0_c44, vartype=float)		# initial slope for fitting line in Css(r, theta) figure
				slope_min_theta = get_env('SLOPE_MIN_THETA',
					default=slope_min_c44, vartype=float)	# minimum slope for fitting line in Css(r, theta) figure
				slope_max_theta = get_env('SLOPE_MAX_THETA',
					default=slope_max_c44, vartype=float)	# maximum slope for fitting line in Css(r, theta) figure

			plot = plot_fft()	# plotting shear strain correlations with slider for cut-off radius

		if get_env('SHOW', default=False, vartype=bool):	# SHOW mode
			plt.show()


text:
"""
This script compares the png test image given in the second argument against
the known good baseline given in the third. Alternate known good images
(with _number.png in the filename) are also compared against.
"""

import math
import numpy
import vtk
import os
import os.path
import sys
import logging
import shutil

defaultThreshold=10.0

def image_compare(testImage, baselineImage):
    imageDiff = vtk.vtkImageDifference()
    imageDiff.SetInputData(testImage)
    imageDiff.SetImageData(baselineImage)
    imageDiff.Update()
    return (imageDiff.GetThresholdedError(), imageDiff.GetOutput())

def dump_image_to_file(fname,img):
    wr = vtk.vtkPNGWriter()
    wr.SetFileName(fname)
    wr.SetInputData(img)
    wr.Write()

def image_from_file(fname):
    try:
        rd = vtk.vtkPNGReader()
        rd.SetFileName(fname)
        removeAlpha = vtk.vtkImageExtractComponents()
        removeAlpha.SetComponents(0, 1, 2)
        removeAlpha.SetInputConnection(rd.GetOutputPort())
        removeAlpha.Update()
        return removeAlpha.GetOutput();
    except Exception as err:
        print(("Problem opening file '%s': %s"%(fname,err)))
        return None

def find_alternates(fname):
    dirname = os.path.dirname(fname)
    prefix, ext = os.path.splitext(os.path.split(fname)[1])
    files = os.listdir(dirname)
    results = [fname]
    for i in files:
        if i.startswith(prefix) and i.endswith(ext) and i != prefix+ext:
            results.append(os.path.join(dirname, i))
    return results

def check_result_image(fname, baselinefname, threshold = defaultThreshold,
                       baseline = True, cleanup=True, update_baseline=False):
    testImage = image_from_file(fname)
    if testImage is None:
        print("Testing image missing, test failed.")
        return -1

    if baseline:
        baselinefnames = find_alternates(baselinefname)
    else:
        baselinefnames = [baselinefname,]

    print("Found Baselines:")
    for baselinefname in baselinefnames:
        print(("- %s"%baselinefname))

    bestImage = None
    bestFilename = None
    bestDiff = None
    bestDiffImage = None
    for baselineFilename in baselinefnames:
        sys.stdout.write("Comparing %s %s ... "%(fname, baselineFilename))
        baselineImage = image_from_file(baselineFilename)
        if baselineImage is None:
            continue

        diff, diffImage = image_compare(testImage, baselineImage)

        sys.stdout.write("diff=%f"%diff)

        if bestDiff is None or diff < bestDiff:
            sys.stdout.write(", New best!")
            bestDiff = diff
            bestFilename = baselineFilename
            bestImage = baselineImage
            bestDiffImage = diffImage
            if bestDiff < threshold:
                break
        sys.stdout.write("\n")

    if bestImage is None:
      print("No baseline images found. Test failed.")
      return -1

    if update_baseline is True:
      print(("UPDATE MODE updating %s with %s" % (bestFilename,fname)))
      shutil.copy(fname,bestFilename)
      return 0

    if bestDiff < threshold:
        print(("Baseline '%s' is the best match with a difference of %f."%(bestFilename, bestDiff)))
        if cleanup:
            print(("Deleting test image '%s'..."%(fname)))
            os.remove(fname)
        return 0

    print(("All baselines failed! Lowest error (%f) exceeds threshold (%f)."%(bestDiff, threshold)))

    sp = fname.split(".")
    diffFilename = ".".join(sp[:-1])+"_diff."+sp[-1]
    print(("Saving image diff at %s"%diffFilename))
    dump_image_to_file(diffFilename, bestDiffImage)

    # Print metadata for CDash image upload:
    def printDart(name, type, value, suff=""):
      print(('<DartMeasurement%s name="%s" type="%s">%s</DartMeasurement%s>'%(
        suff, name, type, value, suff)))
    printDart("ImageError", "numeric/double", "%f"%bestDiff)
    printDart("TestImage", "image/png", os.path.abspath(fname), "File")
    printDart("DifferenceImage", "image/png", os.path.abspath(diffFilename), "File")
    printDart("ValidImage", "image/png", os.path.abspath(bestFilename), "File")
    return -1

def main():
    if len(sys.argv) != 4:
        print("Error:")
        print(("Called with: " + str(sys.argv)))
        print(("Call with " + sys.argv[0] + " test_directory/test_filename.png  baseline_directory/baseline_filename.png threshold"))
        sys.exit(-1)

    ret = check_result_image(sys.argv[1], sys.argv[2], float(sys.argv[3]))
    sys.exit(ret)

if __name__ == "__main__":
    main()

ext: py

max_stars_repo_name: scottwittenburg/thermo

max_issues_repo_name: scottwittenburg/thermo

max_forks_repo_name: scottwittenburg/thermo

reason: import numpy


text:
# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the shot adaptive optimizer"""
import pytest
from scipy.stats import multinomial

import pennylane as qml
from pennylane import numpy as np

class TestExceptions:
    """Test exceptions are raised for incorrect usage"""

    def test_analytic_device_error(self):
        """Test that an exception is raised if an analytic device is used"""
        H = qml.Hamiltonian([0.3, 0.1], [qml.PauliX(0), qml.PauliZ(0)])
        dev = qml.device("default.qubit", wires=1, shots=None)

        def ansatz(x, **kwargs):
            qml.RX(x, wires=0)

        expval_cost = qml.ExpvalCost(ansatz, H, dev)

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        # test expval cost
        with pytest.raises(ValueError, match="can only be used with devices that"):
            opt.step(expval_cost, 0.5)

        # test qnode cost
        with pytest.raises(ValueError, match="can only be used with devices that"):
            opt.step(expval_cost.qnodes[0], 0.5)

    def test_learning_error(self):
        """Test that an exception is raised if the learning rate is beyond the
        lipschitz bound"""
        coeffs = [0.3, 0.1]
        H = qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)])
        dev = qml.device("default.qubit", wires=1, shots=100)

        def ansatz(x, **kwargs):
            qml.RX(x, wires=0)

        expval_cost = qml.ExpvalCost(ansatz, H, dev)

        opt = qml.ShotAdaptiveOptimizer(min_shots=10, stepsize=100.0)

        # lipschitz constant is given by sum(|coeffs|)
        lipschitz = np.sum(np.abs(coeffs))

        assert opt.stepsize > 2 / lipschitz

        with pytest.raises(
            ValueError, match=f"The learning rate must be less than {2 / lipschitz}"
        ):
            opt.step(expval_cost, 0.5)

        # for a single QNode, the lipschitz constant is simply 1
        opt = qml.ShotAdaptiveOptimizer(min_shots=10, stepsize=100.0)
        with pytest.raises(ValueError, match=f"The learning rate must be less than {2 / 1}"):
            opt.step(expval_cost.qnodes[0], 0.5)

    def test_unknown_objective_function(self):
        """Test that an exception is raised if an unknown objective function is passed"""
        dev = qml.device("default.qubit", wires=1, shots=100)

        @qml.qnode(dev)
        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        def cost(x):
            return np.sin(circuit(x))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        # test expval cost
        with pytest.raises(
            ValueError, match="The objective function must either be encoded as a single QNode"
        ):
            opt.step(cost, 0.5)

        # defining the device attribute allows it to proceed
        cost.device = circuit.device
        new_x = opt.step(cost, 0.5)

        assert isinstance(new_x, float)

class TestSingleShotGradientIntegration:
    """Integration tests to ensure that the single shot gradient is correctly computed
    for a variety of argument types."""

    dev = qml.device("default.qubit", wires=1, shots=100)
    H = qml.Hamiltonian([1.0], [qml.PauliZ(0)])

    def ansatz(x, **kwargs):
        qml.RX(x, wires=0)

    expval_cost = qml.ExpvalCost(ansatz, H, dev)

    @qml.qnode(dev)
    def qnode(x):
        qml.RX(x, wires=0)
        return qml.expval(qml.PauliZ(0))

    @pytest.mark.parametrize("cost_fn", [qnode, expval_cost])
    def test_single_argument_step(self, cost_fn, mocker, monkeypatch):
        """Test that a simple QNode with a single argument correctly performs an optimization step,
        and that the single-shot gradients generated have the correct shape"""

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy_single_shot_expval = mocker.spy(opt, "_single_shot_expval_gradients")
        spy_single_shot_qnodes = mocker.spy(opt, "_single_shot_qnode_gradients")
        spy_grad = mocker.spy(opt, "compute_grad")

        x_init = 0.5
        new_x = opt.step(cost_fn, x_init)

        assert isinstance(new_x, float)
        assert new_x != x_init

        spy_grad.assert_called_once()

        if isinstance(cost_fn, qml.ExpvalCost):
            spy_single_shot_expval.assert_called_once()
            single_shot_grads = opt._single_shot_expval_gradients(cost_fn, [x_init], {})
        else:
            spy_single_shot_qnodes.assert_called_once()
            single_shot_grads = opt._single_shot_qnode_gradients(cost_fn, [x_init], {})

        # assert single shot gradients are computed correctly
        assert len(single_shot_grads) == 1
        assert single_shot_grads[0].shape == (10,)

        # monkeypatch the optimizer to use the same single shot gradients
        # as previously
        monkeypatch.setattr(
            opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads
        )
        monkeypatch.setattr(
            opt, "_single_shot_expval_gradients", lambda *args, **kwargs: single_shot_grads
        )

        # reset the shot budget
        opt.s = [np.array(10)]

        # check that the gradient and variance are computed correctly
        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
        assert len(grad) == 1
        assert len(grad_variance) == 1
        assert np.allclose(grad, np.mean(single_shot_grads))
        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1))

        # check that the gradient and variance are computed correctly
        # with a different shot budget
        opt.s = [np.array(5)]
        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
        assert len(grad) == 1
        assert len(grad_variance) == 1
        assert np.allclose(grad, np.mean(single_shot_grads[0][:5]))
        assert np.allclose(grad_variance, np.var(single_shot_grads[0][:5], ddof=1))

    def ansatz(x, **kwargs):
        qml.RX(x[0, 0], wires=0)
        qml.RY(x[0, 1], wires=0)
        qml.RZ(x[0, 2], wires=0)
        qml.RX(x[1, 0], wires=0)
        qml.RY(x[1, 1], wires=0)
        qml.RZ(x[1, 2], wires=0)

    expval_cost = qml.ExpvalCost(ansatz, H, dev)
    qnode = expval_cost.qnodes[0]

    @pytest.mark.parametrize("cost_fn", [qnode, expval_cost])
    def test_single_array_argument_step(self, cost_fn, mocker, monkeypatch):
        """Test that a simple QNode with a single array argument correctly performs an optimization step,
        and that the single-shot gradients generated have the correct shape"""
        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy_single_shot_expval = mocker.spy(opt, "_single_shot_expval_gradients")
        spy_single_shot_qnodes = mocker.spy(opt, "_single_shot_qnode_gradients")
        spy_grad = mocker.spy(opt, "compute_grad")

        x_init = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
        new_x = opt.step(cost_fn, x_init)

        assert isinstance(new_x, np.ndarray)
        assert not np.allclose(new_x, x_init)

        if isinstance(cost_fn, qml.ExpvalCost):
            spy_single_shot_expval.assert_called_once()
            single_shot_grads = opt._single_shot_expval_gradients(cost_fn, [x_init], {})
        else:
            spy_single_shot_qnodes.assert_called_once()
            single_shot_grads = opt._single_shot_qnode_gradients(cost_fn, [x_init], {})

        spy_grad.assert_called_once()

        # assert single shot gradients are computed correctly
        assert len(single_shot_grads) == 1
        assert single_shot_grads[0].shape == (10, 2, 3)

        # monkeypatch the optimizer to use the same single shot gradients
        # as previously
        monkeypatch.setattr(
            opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads
        )
        monkeypatch.setattr(
            opt, "_single_shot_expval_gradients", lambda *args, **kwargs: single_shot_grads
        )

        # reset the shot budget
        opt.s = [10 * np.ones([2, 3], dtype=np.int64)]

        # check that the gradient and variance are computed correctly
        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
        assert len(grad) == 1
        assert len(grad_variance) == 1
        assert grad[0].shape == x_init.shape
        assert grad_variance[0].shape == x_init.shape

        assert np.allclose(grad, np.mean(single_shot_grads, axis=1))
        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1, axis=1))

        # check that the gradient and variance are computed correctly
        # with a different shot budget
        opt.s[0] = opt.s[0] // 2  # all array elements have a shot budget of 5
        opt.s[0][0, 0] = 8  # set the shot budget of the zeroth element to 8

        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
        assert len(grad) == 1
        assert len(grad_variance) == 1

        # check zeroth element
        assert np.allclose(grad[0][0, 0], np.mean(single_shot_grads[0][:8, 0, 0]))
        assert np.allclose(grad_variance[0][0, 0], np.var(single_shot_grads[0][:8, 0, 0], ddof=1))

        # check other elements
        assert np.allclose(grad[0][0, 1], np.mean(single_shot_grads[0][:5, 0, 1]))
        assert np.allclose(grad_variance[0][0, 1], np.var(single_shot_grads[0][:5, 0, 1], ddof=1))

    def test_multiple_argument_step(self, mocker, monkeypatch):
        """Test that a simple QNode with multiple scalar arguments correctly performs an optimization step,
        and that the single-shot gradients generated have the correct shape"""
        dev = qml.device("default.qubit", wires=1, shots=100)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            return qml.expval(qml.PauliZ(0))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy_single_shot = mocker.spy(opt, "_single_shot_qnode_gradients")
        spy_grad = mocker.spy(opt, "compute_grad")

        args = [0.1, 0.2]
        new_x = opt.step(circuit, *args)

        assert isinstance(new_x, list)
        assert len(new_x) == 2

        spy_single_shot.assert_called_once()
        spy_grad.assert_called_once()

        # assert single shot gradients are computed correctly
        single_shot_grads = opt._single_shot_qnode_gradients(circuit, args, {})
        assert len(single_shot_grads) == 2
        assert single_shot_grads[0].shape == (10,)

        # monkeypatch the optimizer to use the same single shot gradients
        # as previously
        monkeypatch.setattr(
            opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads
        )

        # reset the shot budget
        opt.s = [np.array(10), np.array(10)]

        # check that the gradient and variance are computed correctly
        grad, grad_variance = opt.compute_grad(circuit, args, {})
        assert len(grad) == 2
        assert len(grad_variance) == 2
        assert np.allclose(grad, np.mean(single_shot_grads, axis=1))
        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1, axis=1))

        # check that the gradient and variance are computed correctly
        # with a different shot budget
        opt.s = [np.array(5), np.array(7)]
        grad, grad_variance = opt.compute_grad(circuit, args, {})
        assert len(grad) == 2
        assert len(grad_variance) == 2

        for p, s in zip(range(2), opt.s):
            assert np.allclose(grad[p], np.mean(single_shot_grads[p][:s]))
            assert np.allclose(grad_variance[p], np.var(single_shot_grads[p][:s], ddof=1))

    def test_multiple_array_argument_step(self, mocker, monkeypatch):
        """Test that a simple QNode with multiple array arguments correctly performs an optimization step,
        and that the single-shot gradients generated have the correct shape"""
        dev = qml.device("default.qubit", wires=1, shots=100)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x[0, 0], wires=0)
            qml.RY(x[0, 1], wires=0)
            qml.RZ(x[0, 2], wires=0)
            qml.RX(x[1, 0], wires=0)
            qml.RY(x[1, 1], wires=0)
            qml.RZ(x[1, 2], wires=0)
            qml.RX(y[0], wires=0)
            qml.RY(y[1], wires=0)
            qml.RZ(y[2], wires=0)
            return qml.expval(qml.PauliZ(0))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy_single_shot = mocker.spy(opt, "_single_shot_qnode_gradients")
        spy_grad = mocker.spy(opt, "compute_grad")

        args = [np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), np.array([1.0, 2.0, 3.0])]
        new_x = opt.step(circuit, *args)

        assert isinstance(new_x, list)
        assert len(new_x) == 2

        spy_single_shot.assert_called_once()
        spy_grad.assert_called_once()

        # assert single shot gradients are computed correctly
        single_shot_grads = opt._single_shot_qnode_gradients(circuit, args, {})
        assert len(single_shot_grads) == 2
        assert single_shot_grads[0].shape == (10, 2, 3)
        assert single_shot_grads[1].shape == (10, 3)

        # monkeypatch the optimizer to use the same single shot gradients
        # as previously
        monkeypatch.setattr(
            opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads
        )

        # reset the shot budget
        opt.s = [10 * np.ones([2, 3], dtype=np.int64), 10 * np.ones([3], dtype=np.int64)]

        # check that the gradient and variance are computed correctly
        grad, grad_variance = opt.compute_grad(circuit, args, {})
        assert len(grad) == 2
        assert len(grad_variance) == 2

        for i in range(2):
            assert grad[i].shape == args[i].shape
            assert grad_variance[i].shape == args[i].shape

        assert np.allclose(grad[0], np.mean(single_shot_grads[0], axis=0))
        assert np.allclose(grad_variance[0], np.var(single_shot_grads[0], ddof=1, axis=0))

        assert np.allclose(grad[1], np.mean(single_shot_grads[1], axis=0))
        assert np.allclose(grad_variance[1], np.var(single_shot_grads[1], ddof=1, axis=0))

        # check that the gradient and variance are computed correctly
        # with a different shot budget
        opt.s[0] = opt.s[0] // 2  # all array elements have a shot budget of 5
        opt.s[0][0, 0] = 8  # set the shot budget of the zeroth element to 8

        opt.s[1] = opt.s[1] // 5  # all array elements have a shot budget of 2
        opt.s[1][0] = 7  # set the shot budget of the zeroth element to 7

        grad, grad_variance = opt.compute_grad(circuit, args, {})
        assert len(grad) == 2
        assert len(grad_variance) == 2

        # check zeroth element of arg 0
        assert np.allclose(grad[0][0, 0], np.mean(single_shot_grads[0][:8, 0, 0]))
        assert np.allclose(grad_variance[0][0, 0], np.var(single_shot_grads[0][:8, 0, 0], ddof=1))

        # check other elements of arg 0
        assert np.allclose(grad[0][0, 1], np.mean(single_shot_grads[0][:5, 0, 1]))
        assert np.allclose(grad_variance[0][0, 1], np.var(single_shot_grads[0][:5, 0, 1], ddof=1))

        # check zeroth element of arg 1
        assert np.allclose(grad[1][0], np.mean(single_shot_grads[1][:7, 0]))
        assert np.allclose(grad_variance[1][0], np.var(single_shot_grads[1][:7, 0], ddof=1))

        # check other elements of arg 1
        assert np.allclose(grad[1][1], np.mean(single_shot_grads[1][:2, 1]))
        assert np.allclose(grad_variance[1][1], np.var(single_shot_grads[1][:2, 1], ddof=1))

class TestWeightedRandomSampling:
    """Tests for weighted random Hamiltonian term sampling"""

    def test_wrs_expval_cost(self, mocker):
        """Checks that cost functions that are expval costs can
        make use of weighted random sampling"""
        coeffs = [0.2, 0.1]
        dev = qml.device("default.qubit", wires=2, shots=100)
        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])

        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy = mocker.spy(opt, "weighted_random_sampling")

        new_weights = opt.step(expval_cost, weights)
        spy.assert_called_once()

        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)
        assert len(grads) == 1
        assert grads[0].shape == (10, *weights.shape)

    def test_wrs_disabled(self, mocker):
        """Checks that cost functions that are expval costs can
        disable use of weighted random sampling"""
        coeffs = [0.2, 0.1]
        dev = qml.device("default.qubit", wires=2, shots=100)
        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])

        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling=None)
        spy = mocker.spy(opt, "weighted_random_sampling")

        new_weights = opt.step(expval_cost, weights)
        spy.assert_not_called()

    def test_unknown_term_sampling_method(self):
        """Checks that an exception is raised if the term sampling method is unknown"""
        coeffs = [0.2, 0.1]
        dev = qml.device("default.qubit", wires=2, shots=100)
        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])

        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="uniform_random_sampling")

        with pytest.raises(ValueError, match="Unknown Hamiltonian term sampling method"):
            opt.step(expval_cost, weights)

    def test_zero_shots(self, mocker, monkeypatch):
        """Test that, if the shot budget for a single term is 0,
        that the jacobian computation is skipped"""
        coeffs = [0.2, 0.1, 0.1]
        dev = qml.device("default.qubit", wires=2, shots=100)
        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliX(1), qml.PauliZ(0) @ qml.PauliZ(1)])

        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
        spy = mocker.spy(qml, "jacobian")
        mocker.patch(
            "scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 0, 6]])
        )
        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)

        assert len(spy.call_args_list) == 2
        assert len(grads) == 1
        assert grads[0].shape == (10, *weights.shape)

    def test_single_shots(self, mocker, monkeypatch):
        """Test that, if the shot budget for a single term is 1,
        that the number of dimensions for the returned Jacobian is expanded"""
        coeffs = [0.2, 0.1, 0.1]
        dev = qml.device("default.qubit", wires=2, shots=100)
        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliX(1), qml.PauliZ(0) @ qml.PauliZ(1)])

        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        spy = mocker.spy(qml, "jacobian")
        spy_dims = mocker.spy(np, "expand_dims")
        mocker.patch(
            "scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 1, 5]])
        )
        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)

        spy_dims.assert_called_once()
        assert len(spy.call_args_list) == 3
        assert len(grads) == 1
        assert grads[0].shape == (10, *weights.shape)

class TestOptimization:
    """Integration test to ensure that the optimizer
    minimizes simple examples"""

    def test_multi_qubit_rotation(self):
        """Test that multiple qubit rotation can be optimized"""
        dev = qml.device("default.qubit", wires=2, shots=100)

        @qml.qnode(dev)
        def circuit(x):
            qml.RX(x[0], wires=0)
            qml.RY(x[1], wires=1)
            qml.RZ(x[2], wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliX(0) @ qml.PauliY(1))

        params = np.array([0.1, 0.3, 0.4], requires_grad=True)
        initial_loss = circuit(params)

        min_shots = 10
        loss = initial_loss
        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        for i in range(100):
            params = opt.step(circuit, params)
            loss = circuit(params)

        assert loss < initial_loss
        assert np.allclose(circuit(params), -1, atol=0.1, rtol=0.2)
        assert opt.shots_used > min_shots

    @pytest.mark.slow
    def test_vqe_optimization(self):
        """Test that a simple VQE circuit can be optimized"""
        dev = qml.device("default.qubit", wires=2, shots=100)
        coeffs = [0.1, 0.2]
        obs = [qml.PauliZ(0), qml.PauliX(0)]
        H = qml.Hamiltonian(coeffs, obs)

        def ansatz(x, **kwargs):
            qml.Rot(*x[0], wires=0)
            qml.Rot(*x[1], wires=1)
            qml.CNOT(wires=[0, 1])
            qml.Rot(*x[2], wires=0)
            qml.Rot(*x[3], wires=1)
            qml.CNOT(wires=[0, 1])

        cost = qml.ExpvalCost(ansatz, H, dev)
        params = np.random.random((4, 3), requires_grad=True)
        initial_loss = cost(params)

        min_shots = 10
        loss = initial_loss
        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        for i in range(100):
            params = opt.step(cost, params)
            loss = cost(params)

        assert loss < initial_loss
        assert np.allclose(loss, -1 / (2 * np.sqrt(5)), atol=0.1, rtol=0.2)
        assert opt.shots_used > min_shots

class TestStepAndCost:
    """Tests for the step_and_cost method"""

    def test_qnode_cost(self, tol):
        """Test that the cost is correctly returned
        when using a QNode as the cost function"""
        dev = qml.device("default.qubit", wires=1, shots=10)

        @qml.qnode(dev)
        def circuit(x):
            qml.RX(x[0], wires=0)
            qml.RY(x[1], wires=0)
            return qml.expval(qml.PauliZ(0))

        params = np.array([0.1, 0.3], requires_grad=True)

        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        for i in range(100):
            params, res = opt.step_and_cost(circuit, params)

        assert np.allclose(res, -1, atol=tol, rtol=0)

    def test_expval_cost(self, tol):
        """Test that the cost is correctly returned
        when using a QNode as the cost function"""
        dev = qml.device("default.qubit", wires=1, shots=10)

        def ansatz(x, **kwargs):
            qml.RX(x[0], wires=0)
            qml.RY(x[1], wires=0)

        H = qml.Hamiltonian([1.0], [qml.PauliZ(0)])
        circuit = qml.ExpvalCost(ansatz, H, dev)
        params = np.array([0.1, 0.3], requires_grad=True)
        opt = qml.ShotAdaptiveOptimizer(min_shots=10)

        for i in range(100):
            params, res = opt.step_and_cost(circuit, params)

        assert np.allclose(res, -1, atol=tol, rtol=0)

reason: import numpy,from scipy


text:
"""
Get specified number of random tweets from CSV with over 100K entries

"""

#-------------------------------------------#
# Import dependencies                       #
#-------------------------------------------#

import time
import re
import string
import numpy as np
import pandas as pd
import random as rnd
import multiprocessing as mp
import csv

from datetime import datetime, timedelta

#-------------------------------------------#
# Global variables                          #
#-------------------------------------------#

CHUNKSIZE = 10000                                                                          # processing 100,000 rows at a time
WORKER_PROCESSES = 4
SAMPLE_SIZE = 62200
USERID = '2980186012'
INPUT = '../datasets/top_users_csv/{}.csv'.format(USERID)
OUTPUT = '../datasets/top_users_csv/{}_{}.csv'.format(SAMPLE_SIZE, USERID)

#-------------------------------------------#
# Set random seed                           #
#-------------------------------------------#
rnd.seed(8)

#-------------------------------------------#
# Columns of CSV file                       #
#-------------------------------------------#

ibm_cols = [
    'content',
    'id',
    'created',
    'contenttype',
    'language',
    'parentid',
    'reply',
    'forward'
]

#-------------------------------------------#
# D-types                                   #
#-------------------------------------------#

ibm_d_types = {
    'content':                  'str',                                                      # tweet_text
    'id':                       'str',                                                      # tweetid
    'created':                  'int',                                                      # tweet_time (needs to be converted to UNIX)
    'contenttype':              'str',                                                      # set default value as 'text/plain'
    'language':                 'str',                                                      # set default value as 'en' since we check this in cleaning step
    'parentid':                 'str',                                                      # set default value as userid
    'reply':                    'str',                                                     # check if in_reply_to_tweetid/userid is filled or not, return bool
    'forward':                  'str',                                                     # check if is_retweet is True or False, return bool

}

#-------------------------------------------#
# Function                                  #
#-------------------------------------------#

def process_frame(df):
    # print(df.tweet_text)
    return len(df)                                                                          # process data frame

#-------------------------------------------#
# Main                                      #
#-------------------------------------------#

def main():
    start_time = time.time()                                                                # program start time
    print("Processing file: ",INPUT)

    reader = pd.read_csv(
                         INPUT,
                         header=0,
                         dtype=ibm_d_types,
                         chunksize=CHUNKSIZE,
                         keep_default_na=False,                                  # do not parse rows with default NaN values
                        )
    pool = mp.Pool(4)                                                               # use 4 processes

    funclist = []
    reader_count = 1
    for df in reader:
        print("Processing %d df of reader..."%(reader_count))
        f = pool.apply_async(process_frame,[df])
        funclist.append(f)
        reader_count += 1

    n = 0
    for f in funclist:
        n += f.get(timeout=10)                                                      # timeout in 10 seconds

    print("There are ",n," lines in ",INPUT)
    skip = sorted(rnd.sample(range(1,n+1),n-SAMPLE_SIZE))                           # the 0-indexed header will not be included in the skip list
    df = pd.read_csv(
                    INPUT,
                    skiprows=skip
                    )
    df.to_csv(
                OUTPUT,
                columns=ibm_cols,
                encoding='utf-8',
                index=False,
                header=True                                                           # include header for first chunk
            )

    runtime = "Execution: %s seconds" % timedelta(seconds=time.time() - start_time)         # calculate runtime
    print(runtime)

#-------------------------------------------#
# Call Main()                               #
#-------------------------------------------#

if __name__ == "__main__":
    main()

max_stars_repo_name: TheDestroyerOfSomeWorlds/msc

max_issues_repo_name: TheDestroyerOfSomeWorlds/msc

max_forks_repo_name: TheDestroyerOfSomeWorlds/msc


text:
import pandas as pd
import numpy as np

df_train = pd.read_csv("trainset2-15.csv")

X = df_train.iloc[:, 0:-1].values
y = df_train.iloc[:, -1].values

w = np.array([8.7519,5.3491,-1.4437,14.0686,-3.5229 ])

y_pred = np.dot(X, w)
print(y_pred)
# Visualizing
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.scatter(X[:,0], X[:,1], y)
ax.plot(X[:,0], X[:,1], y_pred)

ax.set_xlabel('X Label')
ax.set_ylabel('Y Label')
ax.set_zlabel('Z Label')

plt.xticks(())
plt.yticks(())

plt.show()


text:
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm

def title(constraint):
    increasing = constraint['increasing']
    convex = constraint['convex']
    bounded = len(constraint['bounded']) > 0

    if increasing and not convex and not bounded:
        title = 'Monotonicity'
    elif increasing and not convex and bounded:
        title = 'Monotonicity and boundedness'
    elif increasing and convex and not bounded:
        title = 'Monotonicity and convexity'
    elif increasing and convex and bounded:
        title = 'Monotonicity, boundedness and convexity'
    elif not increasing and not convex and bounded:
        title = 'Boundedness'
    elif not increasing and convex and not bounded:
        title = 'Convexity'
    elif not increasing and convex and bounded:
        title = 'Convexity and boundedness'
    else:
        title = 'Unconstrained'
    return title

def plot_fig(m, constraint, method, n, burn_in):
    Gp = ConstrainedGP(m, constraints=constraint)
    rv = norm()

    def f(x):
        return rv.cdf((x - 0.5) / 0.2)

    x_train = np.array([0.25, 0.5, 0.75])
    y_train = f(x_train)
    Gp.fit_gp(x_train, y_train, n=n, burn_in=burn_in, alpha=0.0000001,
              verbose=False, method=method)

    t = np.arange(0, 1 + 0.01, 0.01)
    y_true = f(t)

    Gp.mean_var(t)

    y_pred = Gp.mean
    ci_l, ci_u = Gp.confidence_interval()

    fig, axis = plt.subplots()
    axis.plot(t, y_true, 'r', label="true function")
    axis.plot(t, y_pred, 'b', label="sample")
    axis.fill_between(t, ci_l, ci_u, color='lightgrey')
    axis.plot(x_train, y_train, 'ko', label="training points")
    axis.set_xlim(0., 1.)
    axis.set_ylim(0., 1.)
    axis.set_xlabel('x')
    axis.set_ylabel('y(x)')
    axis.legend(loc='upper left')
    axis.set_title(title(constraint) + ',method: ' + method)

    plt.show()

def condition_num(const, x, alpha=None):
    print("{} {} {} {} {} {} {} {}".format("  m", "rank_Gamma", "cond_Gamma", "PSD_sigma", "rank_sigma", "cond_sigma",
                                           "PSD_Rinv", "cond_Rinv"))
    for m in range(10, 105, 5):
        Gp = ConstrainedGP(m, constraints=const)
        Gamma = Gp.covariance()
        if alpha is not None:
            Gamma = Gamma + alpha * np.eye(len(Gamma))
        Phi = Gp.interpolation_constraints(x)
        l, Lambda, u = Gp.inequality_constraints()
        Sigma = Gamma - Gamma @ Phi.T @ np.linalg.solve(Phi @ Gamma @ Phi.T, Phi, assume_a='sym') @ Gamma
        if alpha is not None:
            Sigma = Sigma + alpha * np.eye(len(Sigma))

        R = Lambda @ Sigma @ Lambda.T

        if alpha is not None:
            R = R + alpha * np.eye(len(R))
        Rinv = np.linalg.inv(R)

        cond_Gamma = np.linalg.cond(Gamma)
        rank_Gamma = np.linalg.matrix_rank(Gamma)
        cond_Sigma = np.linalg.cond(Sigma)
        rank_Sigma = np.linalg.matrix_rank(Sigma)
        PSD_Sigma = np.all(np.linalg.eigvals(Sigma) > 0)
        cond_Rinv = np.linalg.cond(Rinv)
        PSD_Rinv = np.all(np.linalg.eigvals(Rinv) > 0)
        print("%3d %5d      %5.4E %5d       %5d    %5.4E %4d     %5.4E"
              % (m, rank_Gamma, cond_Gamma, PSD_Sigma, rank_Sigma, cond_Sigma, PSD_Rinv, cond_Rinv))

def plot_fig2(m, method, n, burn_in):
    const = [{'increasing': False, 'bounded': [], 'convex': False},
             {'increasing': False, 'bounded': [0, 1], 'convex': False},
             {'increasing': True, 'bounded': [], 'convex': False},
             {'increasing': True, 'bounded': [0, 1], 'convex': False}]

    rv = norm()

    def f(x):
        return rv.cdf((x - 0.5) / 0.2)

    x_train = np.array([0.25, 0.5, 0.75])
    y_train = f(x_train)
    t = np.arange(0, 1 + 0.01, 0.01)
    y_true = f(t)

    fig, axs = plt.subplots(2, 2)
    for i in range(4):
        Gp = ConstrainedGP(m, constraints=const[i])
        Gp.fit_gp(x_train, y_train, n=n, burn_in=burn_in, alpha=0.0000001,
                  verbose=False, method=method)
        Gp.mean_var(t)

        y_pred = Gp.mean
        ci_l, ci_u = Gp.confidence_interval()

        axs[i // 2, i % 2].plot(t, y_true, 'r', label="true function")
        axs[i // 2, i % 2].plot(t, y_pred, 'b', label="sample")
        axs[i // 2, i % 2].fill_between(t, ci_l, ci_u, color='lightgrey')
        axs[i // 2, i % 2].plot(x_train, y_train, 'ko', label="training points")
        axs[i // 2, i % 2].set_xlim(0., 1.)
        axs[i // 2, i % 2].set_ylim(0., 1.)
        axs[i // 2, i % 2].set_xlabel('x')
        axs[i // 2, i % 2].set_ylabel('y(x)')
        axs[i // 2, i % 2].legend(loc='upper left')
        axs[i // 2, i % 2].set_title(title(const[i]) + ', method: ' + method)

    plt.show()