def test_tritools():
# Tests TriAnalyzer.scale_factors on masked triangulation
# Tests circle_ratios on equilateral and right-angled triangle.
x = np.array([0., 1., 0.5, 0., 2.])
y = np.array([0., 0., 0.5 * np.sqrt(3.), -1., 1.])
triangles = np.array([[0, 1, 2], [0, 1, 3], [1, 2, 4]], dtype=np.int32)
mask = np.array([False, False, True], dtype=np.bool)
triang = mtri.Triangulation(x, y, triangles, mask=mask)
analyser = mtri.TriAnalyzer(triang)
assert_array_almost_equal(analyser.scale_factors,
np.array([1., 1. / (1. + 0.5 * np.sqrt(3.))]))
assert_array_almost_equal(
analyser.circle_ratios(rescale=False),
np.ma.masked_array([0.5, 1. / (1. + np.sqrt(2.)), np.nan], mask))
# Tests circle ratio of a flat triangle
x = np.array([0., 1., 2.])
y = np.array([1., 1. + 3., 1. + 6.])
triangles = np.array([[0, 1, 2]], dtype=np.int32)
triang = mtri.Triangulation(x, y, triangles)
analyser = mtri.TriAnalyzer(triang)
assert_array_almost_equal(analyser.circle_ratios(), np.array([0.]))
# Tests TriAnalyzer.get_flat_tri_mask
# Creates a triangulation of [-1, 1] x [-1, 1] with contiguous groups of
# 'flat' triangles at the 4 corners and at the center. Checks that only
# those at the borders are eliminated by TriAnalyzer.get_flat_tri_mask
n = 9
def power(x, a):
return np.abs(x)**a * np.sign(x)
x = np.linspace(-1., 1., n + 1)
x, y = np.meshgrid(power(x, 2.), power(x, 0.25))
x = x.ravel()
y = y.ravel()
triang = mtri.Triangulation(x, y, triangles=meshgrid_triangles(n + 1))
analyser = mtri.TriAnalyzer(triang)
mask_flat = analyser.get_flat_tri_mask(0.2)
verif_mask = np.zeros(162, dtype=np.bool)
corners_index = [
0, 1, 2, 3, 14, 15, 16, 17, 18, 19, 34, 35, 126, 127, 142, 143, 144,
145, 146, 147, 158, 159, 160, 161
]
verif_mask[corners_index] = True
assert_array_equal(mask_flat, verif_mask)
# Now including a hole (masked triangle) at the center. The center also
# shall be eliminated by get_flat_tri_mask.
mask = np.zeros(162, dtype=np.bool)
mask[80] = True
triang.set_mask(mask)
mask_flat = analyser.get_flat_tri_mask(0.2)
center_index = [44, 45, 62, 63, 78, 79, 80, 81, 82, 83, 98, 99, 116, 117]
verif_mask[center_index] = True
assert_array_equal(mask_flat, verif_mask)
def test_trianalyzer_mismatched_indices():
# github issue 4999.
x = np.array([0., 1., 0.5, 0., 2.])
y = np.array([0., 0., 0.5 * np.sqrt(3.), -1., 1.])
triangles = np.array([[0, 1, 2], [0, 1, 3], [1, 2, 4]], dtype=np.int32)
mask = np.array([False, False, True], dtype=bool)
triang = mtri.Triangulation(x, y, triangles, mask=mask)
analyser = mtri.TriAnalyzer(triang)
# numpy >= 1.10 raises a VisibleDeprecationWarning in the following line
# prior to the fix.
analyser._get_compressed_triangulation()
def plot_surface(cube, model='', unstructure=False, **kw):
projection = kw.pop('projection', ccrs.PlateCarree())
figsize = kw.pop('figsize', (8, 6))
cmap = kw.pop('cmap', plt.cm.rainbow)
fig, ax = plt.subplots(figsize=figsize,
subplot_kw=dict(projection=projection))
ax.set_extent(get_bbox(cube))
ax.add_feature(LAND)
ax.coastlines(resolution='10m')
gl = ax.gridlines(draw_labels=True)
gl.xlabels_top = gl.ylabels_right = False
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER
z = z_coord(cube)
if z:
positive = z.attributes.get('positive', None)
if positive == 'up':
idx = np.unique(z.points.argmax(axis=0))[0]
else:
idx = np.unique(z.points.argmin(axis=0))[0]
c = cube[idx, ...].copy()
else:
idx = None
c = cube.copy()
c.data = ma.masked_invalid(c.data)
t = time_coord(cube)
t = t.units.num2date(t.points)[0]
if unstructure:
# The following lines would work if the cube is note bbox-sliced.
# lon = cube.mesh.nodes[:, 0]
# lat = cube.mesh.nodes[:, 1]
# nv = cube.mesh.faces
lon = cube.coord(axis='X').points
lat = cube.coord(axis='Y').points
nv = Delaunay(np.c_[lon, lat]).vertices
triang = tri.Triangulation(lon, lat, triangles=nv)
# http://matplotlib.org/examples/pylab_examples/ tricontour_smooth_delaunay.html
if False: # TODO: Test this.
subdiv = 3
min_circle_ratio = 0.01
mask = tri.TriAnalyzer(triang).get_flat_tri_mask(min_circle_ratio)
triang.set_mask(mask)
refiner = tri.UniformTriRefiner(triang)
tri_ref, data_ref = refiner.refine_field(cube.data, subdiv=subdiv)
cs = ax.tricontourf(triang, c.data, cmap=cmap, **kw)
else:
cs = ax.pcolormesh(c.coord(axis='X').points,
c.coord(axis='Y').points,
c.data, cmap=cmap, **kw)
title = (model, t, c.name(), idx)
ax.set_title('{}: {}\nVariable: {} level: {}'.format(*title))
return fig, ax, cs
def setup(ds, whichgrids=None):
'''Set up for using ll2xe().
Set up Delaunay triangulation by calculating triangulation and functions for
calculating grid coords from lon/lat pairs and save into and return ds object.
Create a separate triangulation setup for each grid since otherwise it impacts
the performance, especially at edges. Can input keyword whichgrids to only
calculate for particular grids — this is intended for testing purposes to save time.
Usage is demonstrated in ll2xe().
'''
tris = {}
# Set up Delaunay triangulation of grid space in lon/lat coordinates
if whichgrids is None:
whichgrids = ['rho', 'u', 'v', 'psi']
for whichgrid in whichgrids:
lonkey = 'lon_' + whichgrid
latkey = 'lat_' + whichgrid
# Triangulation for curvilinear space to grid space
# Have to use SciPy's Triangulation to be more robust.
# http://matveichev.blogspot.com/2014/02/matplotlibs-tricontour-interesting.html
lon = ds[lonkey].values.flatten()
lat = ds[latkey].values.flatten()
pts = np.column_stack((lon, lat))
tess = Delaunay(pts)
tri = mtri.Triangulation(lon, lat, tess.simplices.copy())
# For the triangulation, need to preprocess the mask to get rid of potential
# flat triangles at the boundaries.
# http://matplotlib.org/1.3.1/api/tri_api.html#matplotlib.tri.TriAnalyzer
mask = mtri.TriAnalyzer(tri).get_flat_tri_mask(0.01, rescale=False)
tri.set_mask(mask)
# Next set up the grid (or index) space grids: integer arrays that are the
# shape of the horizontal model grid.
J, I = ds[lonkey].shape
X, Y = np.meshgrid(np.arange(I), np.arange(J))
# these are the functions for interpolating X and Y (the grid arrays) onto
# lon/lat coordinates. That is, the functions for calculating grid space coords
# corresponding to input lon/lat coordinates.
fx = mtri.LinearTriInterpolator(tri, X.flatten())
fy = mtri.LinearTriInterpolator(tri, Y.flatten())
tris[whichgrid] = {'name': whichgrid, 'tri': tri, 'fx': fx, 'fy': fy}
return tris
def _plot_function2D(funcarr, coordsarr, name):
plt.close('all')
fig = plt.figure(figsize=(10, 10))
x = np.ravel(coordsarr[..., 0])
y = np.ravel(coordsarr[..., 1])
if len(funcarr.shape) == len(coordsarr.shape): # vector quantity
datu = np.ravel(funcarr[..., 0])
datv = np.ravel(funcarr[..., 1])
mag = np.sqrt(np.square(datu) + np.square(datv))
plt.quiver(x, y, datu, datv, mag)
else: # scalar quantity
triang = tri.Triangulation(x, y)
mask = tri.TriAnalyzer(triang).get_flat_tri_mask(min_circle_ratio=0.05)
triang.set_mask(mask)
dat = np.ravel(funcarr)
plt.tripcolor(triang, dat, cmap='jet', shading='gouraud')
# plt.tricontour(x, y, dat, cmap='jet')
# plt.tricontourf(x, y, dat, cmap='jet')
plt.colorbar()
fig.savefig(name + '.png')
plt.close('all')
def fwhm_psf_detector(zosapi, sys_mode, ao_modes, spaxel_scale, grating, N_configs, N_waves, N_rays, mode, files_path, results_path):
"""
Calculate the FWHM PSF (including diffraction effects)
in both directions, acrosss and along the Slices
for a given spaxel scale and grating configuration, across all 4 IFU channels
Because of the nature of HARMONI, we split the calculation for the X and Y direction
The FWHM X is calculated at the Detector plane, in the spatial direction
The FWHM Y is calculated at the Image Slicer plane, across the slices
Details on the methodology can be found in e2e_analysis.py, at "FWHM_PSF_FastAnalysis"
:param zosapi: the Zemax API
:param sys_mode: [str] the system mode, either 'HARMONI', 'ELT' or 'IFS
:param ao_modes: [list] containing the AO mode we want. Only 1 should be used
:param spaxel_scale: [str] the spaxel scale, ex. '60x30'
:param grating: [str] the spectral band to analyze, ex. 'HK'
:param N_waves: how many wavelengths to use. If 3, we will use the min, central and max (typically 5-7)
:param N_rays: how many pupil rays to trace to estimate the geometric PSF
:param mode: whether to use 'geometric' or 'diffraction' PSF for the FWHM calculation
:param files_path: path to the E2E files to load for the analysis
:param results_path: path to save the results
:return:
"""
# We will create a separate folder within results_path to save the FWHM results
analysis_dir = os.path.join(results_path, 'FWHM')
print("Analysis Results will be saved in folder: ", analysis_dir)
if not os.path.exists(analysis_dir):
os.mkdir(analysis_dir)
analysis = e2e.FWHM_PSF_FastAnalysis(zosapi=zosapi)
wavelength_idx = np.linspace(1, 23, N_waves).astype(int)
# Select a subsect of configuration indices. Using all of them takes forever...
# We jump every N_configs
odd_configs = np.arange(1, 76, 2)[::N_configs]
even_configs = odd_configs + 1
# We use both even and odd configurations to cover both A and B paths of the IFU
configs = list(np.concatenate([odd_configs, even_configs]))
N = len(configs)
print("Total Configurations: ", len(configs))
fx, fy = [], []
focal_coord, object_coord = [], []
ifu_sections = ['AB', 'CD', 'EF', 'GH']
for ifu_section in ifu_sections:
options = {'which_system': sys_mode, 'AO_modes': ao_modes, 'scales': [spaxel_scale], 'IFUs': [ifu_section],
'grating': [grating]}
list_results = analysis.loop_over_files(files_dir=files_path, files_opt=options, results_path=results_path,
wavelength_idx=wavelength_idx, configuration_idx=configs, N_rays=N_rays,
mode=mode)
# Only 1 list, no Monte Carlo
fwhm, obj_xy, foc_xy, wavelengths = list_results[0]
focal_coord.append(foc_xy)
object_coord.append(obj_xy)
fwhm_x, fwhm_y = fwhm[:, :, 0], fwhm[:, :, 1]
# Convert to milli-arcseconds
plate_x = plate_scales['DET'][spaxel_scale][0] # mm / arcsec
plate_y = plate_scales['IS'][spaxel_scale][1]
fwhm_x /= plate_x
fwhm_y /= plate_y
print("IFU-%s: %.1f mas" % (ifu_section, np.mean(fwhm_y)))
fx.append(fwhm_x)
fy.append(fwhm_y)
fx = np.array(fx)
fy = np.array(fy)
# (1) DETECTOR plane plot
# Get a separate figure for each direction X and Y
spax = int(spaxel_scale.split('x')[0])
for fwhm, label in zip([fx, fy], ['FWHM_X', 'FWHM_Y']):
# min_fwhm = np.nanmin(fwhm)
# max_fwhm = np.nanmax(fwhm)
min_fwhm = 0
max_fwhm = spax if spaxel_scale == "60x30" else 4 * spax
fig, axes = plt.subplots(2, 2, figsize=(10, 10))
# Loop over the IFU channels: AB, CD, EF, GH
for i in range(2):
for j in range(2):
k = 2 * i + j
ifu_section = ifu_sections[k]
ax = axes[i][j]
_foc_xy = focal_coord[k]
_fwhm = fwhm[k]
# If we don't split the results between "even" and "odd" configurations (A and B paths),
# pyplot will triangulate over the gap between the two detector footprints (where there is no data)
# so it's better to split the results and generate 2 separate triangulations
x_odd, y_odd = _foc_xy[:N//2, :, 0].flatten(), _foc_xy[:N//2, :, 1].flatten()
x_even, y_even = _foc_xy[N//2:, :, 0].flatten(), _foc_xy[N//2:, :, 1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
# This thing here makes sure that we don't get weird triangles if the detector footprint
# has a funny shape (with a lot of spectrograph smile)
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(min_circle_ratio)
triang_even.set_mask(mask_even)
tpc_odd = ax.tripcolor(triang_odd, _fwhm[:N//2].flatten(), shading='flat', cmap='jet')
tpc_odd.set_clim(vmin=min_fwhm, vmax=max_fwhm)
tpc_even = ax.tripcolor(triang_even, _fwhm[N//2:].flatten(), shading='flat', cmap='jet')
tpc_even.set_clim(vmin=min_fwhm, vmax=max_fwhm)
# Draw the boundaries of the detector for reference
draw_detector_boundary(ax)
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
# ax.scatter(_foc_xy[:, :, 0].flatten(), _foc_xy[:, :, 1].flatten(), s=2, color='black')
ax.set_aspect('equal')
cbar = plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
# cbar.ax.set_xlabel('[$\mu$m]')
cbar.ax.set_xlabel('[mas]')
title = r'%s IFU-%s | %s | %s | %s' % (label, ifu_section, spaxel_scale, grating, ao_modes[0])
ax.set_title(title)
fig_name = "%s_%s_%s_DETECTOR_SPEC_%s_MODE_%s_%s" % (label, mode, spaxel_scale, grating, sys_mode, ao_modes[0])
save_path = os.path.join(results_path, analysis_dir)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig.savefig(os.path.join(save_path, fig_name))
# # (2) Stitch the Object space coordinates
# object_coord = np.array(object_coord)
# x_obj, y_obj = object_coord[:, :, 0].flatten(), object_coord[:, :, 1].flatten()
# triang = tri.Triangulation(x_obj, y_obj)
#
# for k_wave, wave_str in zip([0, N_waves // 2, -1], ['MIN', 'CENT', 'MAX']):
#
# for fwhm, label in zip([fx, fy], ['FWHM_X', 'FWHM_Y']):
#
# min_fwhm = np.nanmin(fwhm)
# max_fwhm = np.nanmax(fwhm)
#
# fig_obj, ax = plt.subplots(1, 1)
#
# _fwhm = fwhm[:, :, k_wave].flatten()
# tpc = ax.tripcolor(triang, _fwhm, shading='flat', cmap='jet')
# tpc.set_clim(vmin=min_fwhm, vmax=max_fwhm)
# ax.scatter(x_obj, y_obj, s=2, color='black')
#
# axis_label = 'Object'
# ax.set_xlabel(axis_label + r' X [mm]')
# ax.set_ylabel(axis_label + r' Y [mm]')
# ax.set_aspect('equal')
# cbar = plt.colorbar(tpc, ax=ax, orientation='horizontal')
# cbar.ax.set_xlabel('[$\mu$m]')
#
# wave = wavelengths[k_wave]
#
# title = r'%s | %s | %s SPEC %.3f $\mu$m | %s %s' % (label, spaxel_scale, grating, wave, sys_mode, ao_modes[0])
# ax.set_title(title)
# fig_name = "%s_OBJECT_%s_SPEC_%s_MODE_%s_%s_WAVE_%s" % (label, spaxel_scale, grating, sys_mode, ao_modes[0], wave_str)
#
# save_path = os.path.join(results_path, analysis_dir)
# if os.path.isfile(os.path.join(save_path, fig_name)):
# os.remove(os.path.join(save_path, fig_name))
# fig_obj.savefig(os.path.join(save_path, fig_name))
return fx, fy
import matplotlib.tri as tri
normscatter = np.nanmax(fluor_int)
mask = fluor_int / normscatter > 1e-5
# Now create the Triangulation.
# (Creating a Triangulation without specifying the triangles results in the
# Delaunay triangulation of the points.)
x = jday_int[mask]
y = zcom_t[mask]
z = np.log10(fluor_int[mask] / normscatter)
triang = tri.Triangulation(jday_int[mask], zcom_t[mask])
#-----------------------------------------------------------------------------
# Improving the triangulation before high-res plots: removing flat triangles
#-----------------------------------------------------------------------------
# masking badly shaped triangles at the border of the triangular mesh.
maskt = tri.TriAnalyzer(triang).get_flat_tri_mask(0.1)
#triang.set_mask(maskt)
# refining the data
refiner = tri.UniformTriRefiner(triang)
tri_refi, z_test_refi = refiner.refine_field(z, subdiv=3)
# Temperature
cl = [-3, 0]
axsurf = plt.subplot2grid((3, 4), (1, 0), rowspan=1, colspan=4)
axsurf.grid()
ix = axsurf.tricontourf(jday_int[mask],
zcom_ti[mask],
np.log10(fluor_int[mask] / normscatter),
def detector_rms_wfe(zosapi, sys_mode, ao_modes, spaxel_scale,
spaxels_per_slice, grating, files_path, results_path):
analysis_dir = os.path.join(results_path, 'RMS_WFE')
print("Analysis Results will be saved in folder: ", analysis_dir)
if not os.path.exists(analysis_dir):
os.mkdir(analysis_dir)
ifu_sections = ['AB', 'CD', 'EF', 'GH']
analysis = e2e.RMS_WFE_Analysis(zosapi=zosapi)
rms_maps, object_coord, focal_coord = [], [], []
for ifu_section in ifu_sections:
options = {
'which_system': sys_mode,
'AO_modes': ao_modes,
'scales': [spaxel_scale],
'IFUs': [ifu_section],
'grating': [grating]
}
list_results = analysis.loop_over_files(
files_dir=files_path,
files_opt=options,
results_path=results_path,
wavelength_idx=None,
configuration_idx=None,
surface=None,
spaxels_per_slice=spaxels_per_slice)
# Only 1 list, no Monte Carlo
rms_wfe, obj_xy, foc_xy, global_xy, waves = list_results[0]
# print a summary:
print("\nFor %s scale, IFU-%s, SPEC-%s: " %
(spaxel_scale, ifu_section, grating))
print("RMS: min %.2f | mean %.2f | max %.2f nm " %
(np.min(rms_wfe), np.mean(rms_wfe), np.max(rms_wfe)))
rms_maps.append(rms_wfe)
focal_coord.append(foc_xy)
object_coord.append(obj_xy)
# Stitch the different IFU sections
rms_field = np.concatenate(rms_maps, axis=1)
min_rms = np.min(rms_field)
max_rms = np.max(rms_field)
fig, axes = plt.subplots(2, 2, figsize=(10, 10))
# Loop over the IFU channels: AB, CD, EF, GH
for i in range(2):
for j in range(2):
k = 2 * i + j
ifu_section = ifu_sections[k]
ax = axes[i][j]
_foc_xy = focal_coord[k]
_rms_field = rms_maps[k]
# Separate Odd and Even configs to avoid triangulating over the detector gap
x_odd, y_odd = _foc_xy[:, ::2, :,
0].flatten(), _foc_xy[:, ::2, :,
1].flatten()
x_even, y_even = _foc_xy[:, 1::2, :,
0].flatten(), _foc_xy[:, 1::2, :,
1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
# Remove the flat triangles at the detector edges that appear because of the field curvature
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(
min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(
min_circle_ratio)
triang_even.set_mask(mask_even)
tpc_odd = ax.tripcolor(triang_odd,
_rms_field[:, ::2].flatten(),
shading='flat',
cmap='jet')
tpc_odd.set_clim(vmin=min_rms, vmax=max_rms)
tpc_even = ax.tripcolor(triang_even,
_rms_field[:, 1::2].flatten(),
shading='flat',
cmap='jet')
tpc_even.set_clim(vmin=min_rms, vmax=max_rms)
# ax.scatter(x_odd, y_odd, s=2, color='black')
# ax.scatter(x_even, y_even, s=2, color='black')
draw_detector_boundary(ax)
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
cbar = plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
cbar.ax.set_xlabel('[nm]')
title = r'IFU-%s | %s mas | %s SPEC | %s' % (
ifu_section, spaxel_scale, grating, sys_mode)
ax.set_title(title)
fig_name = "RMSMAP_%s_DETECTOR_SPEC_%s_MODE_%s" % (
spaxel_scale, grating, sys_mode)
# save_path = os.path.join(results_path, analysis_dir)
# if os.path.isfile(os.path.join(save_path, fig_name)):
# os.remove(os.path.join(save_path, fig_name))
# fig.savefig(os.path.join(save_path, fig_name))
# Object space coordinates
fig_obj, ax = plt.subplots(1, 1)
rms_maps = np.array(rms_maps)
object_coord = np.array(object_coord)
N_waves = rms_maps.shape[1]
k_wave = int(N_waves // 2)
x_obj, y_obj = object_coord[:, k_wave, :, :,
0].flatten(), object_coord[:, k_wave, :, :,
1].flatten()
triang = tri.Triangulation(x_obj, y_obj)
tpc = ax.tripcolor(triang,
rms_maps[:, k_wave].flatten(),
shading='flat',
cmap='jet')
tpc.set_clim(vmin=min_rms, vmax=max_rms)
axis_label = 'Object'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
cbar = plt.colorbar(tpc, ax=ax, orientation='horizontal')
cbar.ax.set_xlabel('[nm]')
wave = waves[k_wave]
title = r'%s mas | %s SPEC %.3f $\mu$m | %s' % (spaxel_scale, grating,
wave, sys_mode)
ax.set_title(title)
fig_name = "OBJ_RMSMAP_%s_DETECTOR_SPEC_%s_MODE_%s" % (spaxel_scale,
grating, sys_mode)
save_path = os.path.join(results_path, analysis_dir)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig_obj.savefig(os.path.join(save_path, fig_name))
return rms_field
def readgrid(grid_filename, proj, vert_filename=None, usespherical=True):
"""
readgrid(loc)
Kristen Thyng, March 2013
This function should be read in at the beginnind of a run.py call.
It reads in all necessary grid information that won't change in time
and stores it in a dictionary called grid.
All arrays are changed to Fortran ordering (from Python ordering)
and to tracmass variables ordering from ROMS ordering
i.e. from [t,k,j,i] to [i,j,k,t]
right away after reading in.
Args:
grid_filename: File name (with extension) where grid information is
stored
vert_filename (optional): File name (with extension) where vertical
grid information is stored, if not in grid_loc. Can also skip
this if don't need vertical grid info. also optional prjection
box parameters. Default is None.
proj: Projection object.
usespherical: Use spherical geometric coordinates (lat/lon) or not.
Returns:
* grid - Dictionary containing all necessary time-independent grid fields
grid dictionary contains: (array sizing is for tracmass ordering)
* imt,jmt,km: Grid index sizing constants in (x,y,z), are for horizontal
rho grid [scalar]
* dxv: Horizontal grid cell walls areas in x direction [imt,jmt-1]
* dyu: Horizontal grid cell walls areas in y direction [imt-1,jmt]
* dxdy: Horizontal area of cells defined at cell centers [imt,jmt]
* mask: Land/sea mask [imt,jmt]
* pm,pn: Difference in horizontal grid spacing in x and y [imt,jmt]
* kmt: Number of vertical levels in horizontal space [imt,jmt]
* dzt0: Thickness in meters of grid at each k-level with
time-independent free surface. Surface is at km [imt,jmt,km].
* zrt0: Depth in meters of grid at each k-level on vertical rho grid
with time-independent free surface. Surface is at km [imt,jmt,km]
* zwt0: Depth in meters of grid at each k-level on vertical w grid with
time-independent free surface. Surface is at km [imt,jmt,km]
* xr, yr: Rho grid zonal (x) and meriodional (y) coordinates [imt,jmt]
* xu, yu: U grid zonal (x) and meriodional (y) coordinates [imt,jmt]
* xv, yv: V grid zonal (x) and meriodional (y) coordinates [imt,jmt]
* xpsi, ypsi: Psi grid zonal (x) and meriodional (y) coordinates
[imt, jmt]
* X, Y: Grid index arrays
* tri, trir: Delaunay triangulations
* Cs_r, sc_r: Vertical grid streching paramters [km-1]
* hc: Critical depth [scalar]
* h: Depths [imt,jmt]
* theta_s: Vertical stretching parameter [scalar]. A parameter
(typically 0.0 <= theta_s < 5.0) that defines the amount of grid
focusing. A higher value for theta_s will focus the grid more.
* theta_b: Vertical stretching parameter [scalar]. A parameter (0.0 <
theta_b < 1.0) that says whether the coordinate will be focused at the
surface (theta_b -> 1.0) or split evenly between surface and bottom
(theta_b -> 0)
* basemap: Basemap object
Note: all are in fortran ordering and tracmass ordering except for X, Y,
tri, and tric
To test: [array].flags['F_CONTIGUOUS'] will return true if it is fortran
ordering
"""
# Read in grid parameters and find x and y in domain on different grids
# use full dataset to get grid information
gridfile = netCDF.Dataset(grid_filename)
if usespherical:
try:
lon_vert = gridfile.variables['lon_vert'][:]
lat_vert = gridfile.variables['lat_vert'][:]
except:
lon_rho = gridfile.variables['lon_rho'][:]
lat_rho = gridfile.variables['lat_rho'][:]
x_rho, y_rho = proj(lon_rho, lat_rho)
# get vertex locations
try:
angle = gridfile.variables['angle'][:]
except:
angle = np.zeros(x_rho.shape)
x_vert, y_vert = octant.grid.rho_to_vert(
x_rho, y_rho, gridfile.variables['pm'][:],
gridfile.variables['pn'][:], angle)
lon_vert, lat_vert = proj(x_vert, y_vert, inverse=True)
try:
mask_rho = gridfile.variables['mask'][:]
except:
mask_rho = gridfile.variables['mask_rho'][:]
grid = octant.grid.CGrid_geo(lon_vert, lat_vert, proj)
grid.mask_rho = mask_rho
else: # read cartesian data
try:
x_vert = gridfile.variables['x_vert'][:]
y_vert = gridfile.variables['y_vert'][:]
except:
x_rho = gridfile.variables['x_rho'][:]
y_rho = gridfile.variables['y_rho'][:]
# get vertex locations
try:
angle = gridfile.variables['angle'][:]
except:
angle = np.zeros(x_rho.shape)
x_vert, y_vert = octant.grid.rho_to_vert(
x_rho, y_rho, gridfile.variables['pm'][:],
gridfile.variables['pn'][:], angle)
grid = octant.grid.CGrid(x_vert, y_vert)
try:
mask_rho = gridfile.variables['mask'][:]
grid.mask_rho = mask_rho
# except KeyError as 'mask':
# mask_rho = gridfile.variables['mask_rho'][:]
# grid.mask_rho = mask_rho
except KeyError:
print('No mask.')
# Add into grid spherical coord variables so they are avaiable as
# expected for the code but set them equal to the projected coords.
# Make this better in the future.
grid.lon_rho = grid.x_rho
grid.lat_rho = grid.y_rho
grid.lon_psi = grid.x_psi
grid.lat_psi = grid.y_psi
grid.lon_u = grid.x_u
grid.lat_u = grid.y_u
grid.lon_v = grid.x_v
grid.lat_v = grid.y_v
# vertical grid info
if (vert_filename is not None) or ('s_w' in gridfile.variables):
if 's_w' in gridfile.variables: # test for presence of vertical info
nc = gridfile
else:
nc = netCDF.Dataset(vert_filename)
if 's_w' in nc.variables:
grid.sc_r = nc.variables['s_w'][:] # sigma coords, 31 layers
else:
grid.c_r = nc.variables['sc_w'][:] # sigma coords, 31 layers
# stretching curve in sigma coords, 31 layers
grid.Cs_r = nc.variables['Cs_w'][:]
grid.hc = nc.variables['hc'][:]
grid.theta_s = nc.variables['theta_s'][:]
grid.theta_b = nc.variables['theta_b'][:]
if 'Vtransform' in nc.variables:
grid.Vtransform = nc.variables['Vtransform'][:]
grid.Vstretching = nc.variables['Vstretching'][:]
else:
grid.Vtransform = 1
grid.Vstretching = 1
# Basing this on setupgrid.f95 for rutgersNWA example project from Bror
grid.h = gridfile.variables['h'][:]
# Grid sizes
grid.imt = grid.h.shape[1] # 191
grid.jmt = grid.h.shape[0] # 671
if hasattr(grid, 'sc_r'):
grid.km = grid.sc_r.shape[0] - 1 # 30 NOT SURE ON THIS ONE YET
# Index grid, for interpolation between real and grid space
# This is for rho
# X goes from 0 to imt-1 and Y goes from 0 to jmt-1
# grid in index coordinates, without ghost cells
grid.X, grid.Y = np.meshgrid(np.arange(grid.imt), np.arange(grid.jmt))
# Triangulation for grid space to curvilinear space
pts = np.column_stack((grid.X.flatten(), grid.Y.flatten()))
tess = Delaunay(pts)
grid.tri = mtri.Triangulation(grid.X.flatten(), grid.Y.flatten(),
tess.simplices.copy())
# Triangulation for curvilinear space to grid space
# Have to use SciPy's Triangulation to be more robust.
# http://matveichev.blogspot.com/2014/02/matplotlibs-tricontour-interesting.html
if isinstance(grid.x_rho, np.ma.MaskedArray):
pts = np.column_stack(
(grid.x_rho.data.flatten(), grid.y_rho.data.flatten()))
else:
pts = np.column_stack((grid.x_rho.flatten(), grid.y_rho.flatten()))
tess = Delaunay(pts)
grid.trir = mtri.Triangulation(grid.x_rho.flatten(), grid.y_rho.flatten(),
tess.simplices.copy())
# For the two triangulations that are not integer based, need to
# preprocess the mask to get rid of potential flat triangles at the
# boundaries
# http://matplotlib.org/1.3.1/api/tri_api.html#matplotlib.tri.TriAnalyzer
# Hopefully these will work for other cases too: for the xy spherical
# unit test cases, I needed these both for the triangulation to be valid.
mask = mtri.TriAnalyzer(grid.trir).get_flat_tri_mask(0.01, rescale=True)
grid.trir.set_mask(mask)
mask = mtri.TriAnalyzer(grid.trir).get_flat_tri_mask(0.01, rescale=False)
grid.trir.set_mask(mask)
if isinstance(grid.x_rho, np.ma.MaskedArray):
pts = np.column_stack(
(grid.lon_rho.data.flatten(), grid.lat_rho.data.flatten()))
else:
pts = np.column_stack((grid.lon_rho.flatten(), grid.lat_rho.flatten()))
tess = Delaunay(pts)
grid.trirllrho = mtri.Triangulation(grid.lon_rho.flatten(),
grid.lat_rho.flatten(),
tess.simplices.copy())
mask = mtri.TriAnalyzer(grid.trirllrho).get_flat_tri_mask(0.01,
rescale=True)
grid.trirllrho.set_mask(mask)
mask = mtri.TriAnalyzer(grid.trirllrho).get_flat_tri_mask(0.01,
rescale=False)
grid.trirllrho.set_mask(mask)
# tracmass ordering.
# Not sure how to convert this to pm, pn with appropriate shift
grid.dxv = 1 / grid.pm # pm is 1/\Delta x at cell centers
grid.dyu = 1 / grid.pn # pn is 1/\Delta y at cell centers
grid.dxdy = grid.dyu * grid.dxv
# Change dxv,dyu to be correct u and v grid size after having
# them be too big for dxdy calculation. This is not in the
# rutgersNWA example and I am not sure why. [i,j]
grid.dxv = 0.5 * (grid.dxv[:-1, :] + grid.dxv[1:, :])
grid.dyu = 0.5 * (grid.dyu[:, :-1] + grid.dyu[:, 1:])
# Adjust masking according to setupgrid.f95 for rutgersNWA example
# project from Bror
if hasattr(grid, 'sc_r'):
mask2 = grid.mask.copy()
grid.kmt = np.ones((grid.jmt, grid.imt)) * grid.km
ind = (mask2 == 1)
ind[0:grid.jmt - 1, :] = ind[1:grid.jmt, :]
mask2[ind] = 1
ind = (mask2 == 1)
ind[:, 0:grid.imt - 1] = ind[:, 1:grid.imt]
mask2[ind] = 1
ind = (mask2 == 0)
grid.kmt[ind] = 0
# Use octant to calculate depths/thicknesses for the appropriate
# vertical grid parameters have to transform a few back to ROMS
# coordinates and python ordering for this
grid.zwt0 = octant.depths.get_zw(grid.Vtransform,
grid.Vstretching,
grid.km + 1,
grid.theta_s,
grid.theta_b,
grid.h,
grid.hc,
zeta=0,
Hscale=3)
grid.zrt0 = octant.depths.get_zrho(grid.Vtransform,
grid.Vstretching,
grid.km,
grid.theta_s,
grid.theta_b,
grid.h,
grid.hc,
zeta=0,
Hscale=3)
# this should be the base grid layer thickness that doesn't change in
# time because it is for the reference vertical level
grid.dzt0 = grid.zwt0[1:, :, :] - grid.zwt0[:-1, :, :]
gridfile.close()
return grid
def detector_ensquared_energy(zosapi, file_options, N_rays, box_size,
files_path, results_path):
"""
Calculate the Ensquared Energy at the detector plane for all 4 IFU channels
The Ensquared Energy is calculated for the Field Point at the centre of the slice
for all slices and all wavelengths
The results are shown at the DETECTOR Plane for all 4 IFU channels
:param zosapi: Zemax API
:param sys_mode: [str] system mode, either 'HARMONI', 'ELT' or 'IFS'
:param ao_modes: [list] containing the AO mode we want to analyze (if running 'HARMONI' system mode')
:param spaxel_scale: [str] the spaxel scale, ex. '60x30'
:param grating: [str] spectral band
:param N_rays: how many rays to trace in the Pupil to calculate the Ensquared Energy (typically 500 - 1000)
:param box_size: size in [spaxels] for the EE (default is 2 spaxels, HARMONI requirement)
:param files_path: path to the E2E files we will load for the analysis
:param results_path: path where we want to save the results
:return:
"""
spaxel_scale = file_options['SPAX_SCALE']
grating = file_options['GRATING']
monte_carlo = file_options['MONTE_CARLO']
ao_mode = file_options['AO_MODE']
# We will create a separate folder within results_path to save the results
analysis_dir = os.path.join(results_path, 'ENSQUARED ENERGY')
print("Analysis Results will be saved in folder: ", analysis_dir)
if not os.path.exists(analysis_dir):
os.mkdir(analysis_dir)
analysis = e2e.EnsquaredEnergyFastAnalysis(
zosapi=zosapi) # The analysis object
# For speed, we only use a fraction of the available wavelengths
# at least, 3 wavelengths (min, central, max), typically 5 or 7 to get decent sampling at the detector plane
waves = np.linspace(1, 23, 7).astype(int)
ifu_sections = ['AB', 'CD', 'EF', 'GH']
# ifu_sections = ['AB', 'CD']
fig, axes = plt.subplots(2, 2, figsize=(15, 10))
for i, ifu in enumerate(ifu_sections): # Loop over the IFU channels
# change the IFU path option
file_options['IFU_PATH'] = ifu
if monte_carlo is True:
# read the MC instance number of the specific IFU path
ifu_mc = file_options['IFU_PATHS_MC'][ifu]
file_options['IFU_MC'] = ifu_mc
list_results = analysis.loop_over_files(files_dir=files_path,
files_opt=file_options,
results_path=results_path,
wavelength_idx=waves,
configuration_idx=None,
N_rays=N_rays,
box_size=box_size,
monte_carlo=monte_carlo)
# No Monte Carlo so the list_results only contains 1 entry, for the IFU channel
energy, obj_xy, slicer_xy, detector_xy, wavelengths = list_results[0]
# Separate the Odd / Even configurations to avoid triangulating over the gap on the detector plane
ener_ogg, ener_even = energy[::2].flatten(), energy[1::2].flatten()
x, y = detector_xy[:, :, 0].flatten(), detector_xy[:, :, 1].flatten()
x_odd, y_odd = detector_xy[::2, :,
0].flatten(), detector_xy[::2, :,
1].flatten()
x_even, y_even = detector_xy[1::2, :,
0].flatten(), detector_xy[1::2, :,
1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
# Remove the flat triangles at the detector edges that appear because of the field curvature
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(
min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(
min_circle_ratio)
triang_even.set_mask(mask_even)
min_ener, max_ener = np.min(energy), np.max(energy)
ax = axes.flatten()[i]
tpc_odd = ax.tripcolor(triang_odd,
ener_ogg,
shading='flat',
cmap='Blues',
vmin=min_ener,
vmax=1.0)
tpc_odd.set_clim(vmin=min_ener, vmax=1.0)
tpc_even = ax.tripcolor(triang_even,
ener_even,
shading='flat',
cmap='Blues',
vmin=min_ener,
vmax=1.0)
tpc_even.set_clim(vmin=min_ener, vmax=1.0)
ax.scatter(x, y, s=2, color='black')
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
title = r'IFU-%s | %s | %s | %s | EE min:%.2f max:%.2f' % (
ifu, spaxel_scale, grating, ao_mode, min_ener, max_ener)
ax.set_title(title)
save_path = os.path.join(results_path, analysis_dir)
fig_name = "ENSQ_ENERGY_DETECTOR_%s_SPEC_%s_MODE_%s" % (spaxel_scale,
grating, ao_mode)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig.savefig(os.path.join(save_path, fig_name))
return
def detector_rms_wfe(zosapi, file_options, spaxels_per_slice, pupil_sampling,
files_path, results_path):
"""
Calculate the RMS WFE for a given system mode (typically HARMONI), AO mode, spaxel scale
and spectral band, at the DETECTOR PLANE
We loop over the IFU channels and calculate, for each E2E file, the RMS WFE for all wavelengths and configurations
We can specify the number of points per slice [spaxels per slice] to use for the calculation (typically 3)
The results are plotted:
(1) at the DETECTOR plane for all wavelengths
(2) at the OBJECT plane (stitching all IFU channels) for the Min, Central and Max wavelength
In order to accommodate both Nominal Design and Monte Carlo files, we have abstracted most of the parameters
into the "file_options" dictionary
:param zosapi: the Zemax API
:param file_options: dictionary containing the different parameters
:param spaxels_per_slice: number of field points per slice to use for the calculation
:param pupil_sampling: the pupil sampling for the RWRE operand, representing a grid of N x N per pupil quadrant
:param files_path: path to the E2E files to load for the analysis
:param results_path: path to save the results
:return:
"""
spaxel_scale = file_options['SPAX_SCALE']
grating = file_options['GRATING']
monte_carlo = file_options['MONTE_CARLO']
ao_mode = file_options['AO_MODE']
# We will create a separate folder within results_path to save the RMS WFE results
analysis_dir = os.path.join(results_path, 'RMS_WFE')
print("Analysis Results will be saved in folder: ", analysis_dir)
if not os.path.exists(analysis_dir):
os.mkdir(analysis_dir)
ifu_sections = ['AB', 'CD', 'EF', 'GH']
# ifu_sections = ['AB', 'CD']
analysis = e2e.RMS_WFE_FastAnalysis(zosapi=zosapi) # The analysis object
rms_maps, object_coord, focal_coord = [], [], [
] # Lists to save the results for each IFU channel
for ifu_section in ifu_sections:
# change the IFU path option to the current IFU section
file_options['IFU_PATH'] = ifu_section
if monte_carlo is True:
# read the MC instance number of the specific IFU path
ifu_mc = file_options['IFU_PATHS_MC'][ifu_section]
file_options['IFU_MC'] = ifu_mc
# select the proper ISP MC instance, according to the IFU path
isp_mc = file_options['IFU_ISP_MC'][ifu_section]
file_options['ISP_MC'] = isp_mc
# Run the Analysis for a given E2E file, corresponding to a single IFU path and ISP
list_results = analysis.loop_over_files(
files_dir=files_path,
files_opt=file_options,
results_path=results_path,
wavelength_idx=None,
configuration_idx=None,
surface=None,
spaxels_per_slice=spaxels_per_slice,
pupil_sampling=pupil_sampling,
remove_slicer_aperture=True,
monte_carlo=monte_carlo)
rms_wfe, obj_xy, foc_xy, waves = list_results[
0] # Only 1 item on the list, no Monte Carlo files
# print a summary to spot any issues:
print("\nFor %s scale, IFU-%s, SPEC-%s: " %
(spaxel_scale, ifu_section, grating))
print("RMS: min %.2f | mean %.2f | max %.2f nm " %
(np.min(rms_wfe), np.mean(rms_wfe), np.max(rms_wfe)))
# Add the results to a list that covers all IFU paths
rms_maps.append(rms_wfe)
focal_coord.append(foc_xy)
object_coord.append(obj_xy)
# Stitch the different IFU sections
rms_field = np.concatenate(rms_maps, axis=1)
# Read the RMS WFE Requirement for the particular Spaxel Scale
req = RMS_WFE[spaxel_scale]
min_rms = 0
max_rms = req
# min_rms = np.min(rms_field)
# max_rms = np.max(rms_field)
# (1) DETECTOR plane plot
fig, axes = plt.subplots(2, 2, figsize=(10, 10))
# Loop over the IFU channels: AB, CD, EF, GH
for i in range(2):
# i = 0
for j in range(2):
k = 2 * i + j
ifu_section = ifu_sections[k]
ax = axes[i][j]
_foc_xy = focal_coord[k]
_rms_field = rms_maps[k]
# This is where we create the Detector plane maps
# We separate the Odd and Even configurations as these represent the two sides of the detector
x_odd, y_odd = _foc_xy[::2, :, :,
0].flatten(), _foc_xy[::2, :, :,
1].flatten()
x_even, y_even = _foc_xy[1::2, :, :,
0].flatten(), _foc_xy[1::2, :, :,
1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
# We fix the triangulation to avoid getting weird artifacts
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(
min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(
min_circle_ratio)
triang_even.set_mask(mask_even)
tpc_odd = ax.tripcolor(triang_odd,
_rms_field[::2].flatten(),
shading='flat',
cmap='jet')
tpc_odd.set_clim(vmin=min_rms, vmax=max_rms)
tpc_even = ax.tripcolor(triang_even,
_rms_field[1::2].flatten(),
shading='flat',
cmap='jet')
tpc_even.set_clim(vmin=min_rms, vmax=max_rms)
draw_detector_boundary(ax)
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
cbar = plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
cbar.ax.set_xlabel('[nm]')
if monte_carlo is False:
title = r'IFU-%s | %s | %s | %s | Req: %d nm' % (
ifu_section, spaxel_scale, grating, ao_mode, req)
else:
title = r'IFU-%s | %s | %s | %s MC | Req: %d nm' % (
ifu_section, spaxel_scale, grating, ao_mode, req)
ax.set_title(title)
fig_name = "RMSWFE_%s_DETECTOR_SPEC_%s_MODE_%s" % (spaxel_scale, grating,
ao_mode)
save_path = os.path.join(results_path, analysis_dir)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig.savefig(os.path.join(save_path, fig_name))
N_waves = len(waves)
rms_maps = np.array(rms_maps)
object_coord = np.array(object_coord)
# (2) Stitch the Object space coordinates
# This is a new routine, as suggested by Fraser. How does it work
# First, it interpolates along each slice based on N=spaxels_per_slice values, evaluating the interpolation at
# 76 pixels per slice. Then, it defines a grid of 204 x 152 pixels covering the field of view in spaxels
# It uses the results of the along-slice interpolation to interpolate onto that grid
for k_wave, wave_str in zip([0, N_waves // 2, -1], ['MIN', 'CENT', 'MAX']):
data = rms_maps[:, :, k_wave].flatten()
x_obj, y_obj = object_coord[:, :, :,
0].flatten(), object_coord[:, :, :,
1].flatten()
# we have to transform the object coordinates [mm] into arcseconds using the plate scales
x_obj /= plate_scale
y_obj /= plate_scale
N_samples = x_obj.shape[0]
Nx = 76 # Number of pixels along the slice to interpolate to.
xhi, yhi, dhi = [], [], []
for i in np.arange(0, N_samples, spaxels_per_slice):
xlo, ylo = x_obj[i:i + spaxels_per_slice], y_obj[i:i +
spaxels_per_slice]
dlo = data[i:i + spaxels_per_slice]
xmin, xmax = np.min(xlo), np.max(xlo)
# ymin, ymax = np.min(ylo), np.max(ylo)
yi = np.linspace(ylo.mean(), ylo.mean(),
Nx) # Collapse Y (across slice)
xi = np.linspace(xmin, xmax, Nx) # Linear sample along slice
# Need to use scipy interpolation version, as some slices have decreasing Y values.
f = interp1d(xlo, dlo, assume_sorted=False)
di = f(xi) # Interpolate WFE values along slice
xhi.append(xi)
yhi.append(yi)
dhi.append(di)
xhi = np.array(xhi).flatten()
yhi = np.array(yhi).flatten()
dhi = np.array(dhi).flatten()
xmin, xmax = np.min(xhi), np.max(xhi)
ymin, ymax = np.min(yhi), np.max(yhi)
# Define a regular grid onto which we will interpolate the data
# Used 204 x 152 pixels, which is the total field of view in spaxels.
xi, yi = np.mgrid[xmin:xmax:204j,
ymin:ymax:152j] # 204 pixels by 152 slices
# zi = griddata((xhi[good], yhi[good]), dhi[good], (xi, yi), method='linear') # Nearest or linear.
zi = griddata(points=(xhi, yhi),
values=dhi,
xi=(xi, yi),
method='nearest') # Nearest or linear.
fig_obj, ax = plt.subplots(1, 1)
img = ax.imshow(zi.T,
interpolation='none',
cmap='jet',
origin='lower left',
extent=[xmin, xmax, ymin, ymax])
img.set_clim(vmin=min_rms, vmax=max_rms)
cbar = plt.colorbar(img, ax=ax, orientation='horizontal')
cbar.ax.set_xlabel('[nm]')
axis_label = 'Object'
ax.set_xlabel(axis_label + r' X [arcsec]')
ax.set_ylabel(axis_label + r' Y [arcsec]')
wave = waves[k_wave]
if monte_carlo is False:
title = r'%s mas | %s SPEC %.3f $\mu$m | %s' % (
spaxel_scale, grating, wave, ao_mode)
else:
title = r'%s mas | %s SPEC %.3f $\mu$m | %s MC' % (
spaxel_scale, grating, wave, ao_mode)
ax.set_title(title)
fig_name = "RMSWFE_OBJECT_INTERP_%s_SPEC_%s_MODE_%s_WAVE_%s" % (
spaxel_scale, grating, ao_mode, wave_str)
save_path = os.path.join(results_path, analysis_dir)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig_obj.savefig(os.path.join(save_path, fig_name))
# (2) Stitch the Object space coordinates
# # We create a plot for each case of Minimum, Central and Maximum Wavelength in the spectral band
# for k_wave, wave_str in zip([0, N_waves // 2, -1], ['MIN', 'CENT', 'MAX']):
# fig_obj, ax = plt.subplots(1, 1)
# x_obj, y_obj = object_coord[:, :, :, 0].flatten(), object_coord[:, :, :, 1].flatten()
#
# # we have to transform the object coordinates [mm] into arcseconds using the plate scales
# x_obj /= plate_scale
# y_obj /= plate_scale
#
# triang = tri.Triangulation(x_obj, y_obj)
# rms_ = rms_maps[:, :, k_wave].flatten()
# tpc = ax.tripcolor(triang, rms_, shading='flat', cmap='jet')
# tpc.set_clim(vmin=min_rms, vmax=max_rms)
# # ax.scatter(x_obj, y_obj, s=3, color='black')
#
# axis_label = 'Object'
# ax.set_xlabel(axis_label + r' X [arcsec]')
# ax.set_ylabel(axis_label + r' Y [arcsec]')
# ax.set_aspect('equal')
# cbar = plt.colorbar(tpc, ax=ax, orientation='horizontal')
# cbar.ax.set_xlabel('[nm]')
#
# wave = waves[k_wave]
#
# if monte_carlo is False:
# title = r'%s mas | %s SPEC %.3f $\mu$m | %s' % (spaxel_scale, grating, wave, ao_mode)
# else:
# title = r'%s mas | %s SPEC %.3f $\mu$m | %s MC' % (spaxel_scale, grating, wave, ao_mode)
# ax.set_title(title)
#
# fig_name = "RMSWFE_OBJECT_%s_SPEC_%s_MODE_%s_WAVE_%s" % (spaxel_scale, grating, ao_mode, wave_str)
# save_path = os.path.join(results_path, analysis_dir)
# if os.path.isfile(os.path.join(save_path, fig_name)):
# os.remove(os.path.join(save_path, fig_name))
# fig_obj.savefig(os.path.join(save_path, fig_name))
return rms_field
def stitch_fields(zosapi,
mode,
spaxel_scale,
spaxels_per_slice,
wavelength_idx,
spectral_band,
plane,
files_path,
results_path,
show='focal'):
"""
Loop over the IFU sections AB, CD, EF, GH and calculate the RMS WFE map for each case
Gather the results and stitch them together to get a complete Field Map of the IFU
We can specify the Focal Plane at which we calculate the RMS. For example:
'PO': preoptics, 'IS': Image Slicer, 'DET': detector
We have to be careful with the plot format because the coordinate systems vary depending on
the focal plane we choose. Moreover, we can choose whether to refer the results to 'object' or
'focal' coordinates, which also depends on the surface. For instance, if we go to the DETECTOR plane,
it makes no sense to use 'object' coordinates because the wavelengths would overlap, we must use 'focal'
:param zosapi: the Python Standalone Application
:param mode: E2E file mode: IFS, HARMONI or ELT
:param spaxel_scale: string, spaxel scale like '60x30'
:param spaxels_per_slice: how many points to sample each slice.
:param wavelength_idx: list of Zemax wavelength indices to analyse
:param spectral_band: string, the grating we use
:param plane: string, codename for the focal plane surface 'PO': PreOptics, 'IS': Image Slicer, 'DET': Detector
:param files_path, path where the E2E model Zemax files are stored
:param results_path, path where we want to save the plots
:param show: string, either 'focal' for focal plane coordinates or 'object' for object space coordinates
:return:
"""
if mode == 'IFS':
ao_modes = []
elif mode == 'HARMONI':
ao_modes = ['NOAO']
ifu_sections = ['AB', 'CD', 'EF', 'GH']
analysis = e2e.RMS_WFE_Analysis(zosapi=zosapi)
rms_maps, object_coord, focal_coord = [], [], []
for ifu_section in ifu_sections:
options = {
'which_system': mode,
'AO_modes': ao_modes,
'scales': [spaxel_scale],
'IFUs': [ifu_section],
'grating': [spectral_band]
}
# We need to know the Surface Number for the focal plane of interest in the Zemax files
# this is something varies with mode, scale, ifu channel so we save those numbers on e2e.focal_planes dictionary
focal_plane = e2e.focal_planes[mode][spaxel_scale][ifu_section][plane]
list_results = analysis.loop_over_files(
files_dir=files_path,
files_opt=options,
results_path=results_path,
wavelength_idx=wavelength_idx,
configuration_idx=None,
surface=focal_plane,
spaxels_per_slice=spaxels_per_slice)
# Only 1 list, no Monte Carlo
rms_wfe, obj_xy, foc_xy, global_xy, waves = list_results[0]
rms_maps.append(rms_wfe)
focal_coord.append(foc_xy)
object_coord.append(obj_xy)
# Stitch the different IFU sections
rms_field = np.concatenate(rms_maps, axis=1)
obj_xy = np.concatenate(object_coord, axis=1)
foc_xy = np.concatenate(focal_coord, axis=1)
if plane != 'DET':
# For all surfaces except the detector plane we want to separate the plots for each wavelength
min_rms = np.min(rms_field)
max_rms = np.max(rms_field)
for j_wave, wavelength in enumerate(waves):
if show == 'object':
x, y = obj_xy[j_wave, :, :, 0].flatten(), obj_xy[j_wave, :, :,
1].flatten()
axis_label = 'Object Plane'
ref = 'OBJ'
elif show == 'focal':
x, y = foc_xy[j_wave, :, :, 0].flatten(), foc_xy[j_wave, :, :,
1].flatten()
axis_label = 'Focal Plane'
ref = 'FOC'
if show == 'object' or show == 'focal':
triang = tri.Triangulation(x, y)
fig, ax = plt.subplots(1, 1)
ax.set_aspect('equal')
tpc = ax.tripcolor(triang,
rms_field[j_wave, :, :].flatten(),
shading='flat',
cmap='jet')
tpc.set_clim(vmin=min_rms, vmax=max_rms)
# ax.scatter(x, y, s=2, color='black')
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
title = r'RMS Field Map Stitched IFU | %s mas %s %.3f $\mu$m' % (
spaxel_scale, spectral_band, wavelength)
ax.set_title(title)
plt.colorbar(tpc, ax=ax)
fig_name = "STITCHED_RMSMAP_%s_SPEC_%s_%s_SURF_%s_WAVE%d" % (
spaxel_scale, spectral_band, ref, plane,
wavelength_idx[j_wave])
if os.path.isfile(os.path.join(results_path, fig_name)):
os.remove(os.path.join(results_path, fig_name))
fig.savefig(os.path.join(results_path, fig_name))
elif show == 'both':
x_obj, y_obj = obj_xy[j_wave, :, :,
0].flatten(), obj_xy[j_wave, :, :,
1].flatten()
x_foc, y_foc = foc_xy[j_wave, :, :,
0].flatten(), foc_xy[j_wave, :, :,
1].flatten()
xy = [[x_obj, y_obj], [x_foc, y_foc]]
axis_labels = ['Object Plane', 'Focal Plane']
refs = ['OBJ', 'FOC']
for i in range(2):
x, y = xy[i]
triang = tri.Triangulation(x, y)
fig, ax = plt.subplots(1, 1)
ax.set_aspect('equal')
tpc = ax.tripcolor(triang,
rms_field[j_wave, :, :].flatten(),
shading='flat',
cmap='jet')
tpc.set_clim(vmin=min_rms, vmax=max_rms)
ax.set_xlabel(axis_labels[i] + r' X [mm]')
ax.set_ylabel(axis_labels[i] + r' Y [mm]')
title = r'RMS Field Map Stitched IFU | %s mas %s %.3f $\mu$m' % (
spaxel_scale, spectral_band, wavelength)
ax.set_title(title)
plt.colorbar(tpc, ax=ax)
fig_name = "STITCHED_RMSMAP_%s_SPEC_%s_%s_SURF_%s_WAVE%d" % (
spaxel_scale, spectral_band, refs[i], plane,
wavelength_idx[j_wave])
if os.path.isfile(os.path.join(results_path, fig_name)):
os.remove(os.path.join(results_path, fig_name))
fig.savefig(os.path.join(results_path, fig_name))
else:
raise ValueError('show should be "object" / "focal" / "both"')
if plane == 'DET':
# For the detector plane we want to plot ALL wavelengths at the same time!
# And using the FOCAL PLANE COORDINATES as reference
# at the Object plane there is an overlap in wavelength!!
min_rms = np.min(rms_field)
max_rms = np.max(rms_field)
fig, axes = plt.subplots(2, 2, figsize=(10, 10))
# Loop over the IFU channels: AB, CD, EF, GH
for i in range(2):
for j in range(2):
k = 2 * i + j
ifu_section = ifu_sections[k]
ax = axes[i][j]
_foc_xy = focal_coord[k]
_rms_field = rms_maps[k]
x_odd, y_odd = _foc_xy[:, ::2, :,
0].flatten(), _foc_xy[:, ::2, :,
1].flatten()
x_even, y_even = _foc_xy[:, 1::2, :,
0].flatten(), _foc_xy[:, 1::2, :,
1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
# Remove the flat triangles at the detector edges that appear because of the field curvature
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(
min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(
min_circle_ratio)
triang_even.set_mask(mask_even)
tpc_odd = ax.tripcolor(triang_odd,
_rms_field[:, ::2].flatten(),
shading='flat',
cmap='jet')
tpc_odd.set_clim(vmin=min_rms, vmax=max_rms)
tpc_even = ax.tripcolor(triang_even,
_rms_field[:, 1::2].flatten(),
shading='flat',
cmap='jet')
tpc_even.set_clim(vmin=min_rms, vmax=max_rms)
# ax.scatter(x, y, s=2, color='black')
draw_detector_boundary(ax)
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
cbar = plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
cbar.ax.set_xlabel('[nm]')
title = r'IFU-%s | %s mas | %s SPEC | %s' % (
ifu_section, spaxel_scale, spectral_band, mode)
ax.set_title(title)
fig_name = "RMSMAP_%s_DETECTOR_SPEC_%s_MODE_%s" % (
spaxel_scale, spectral_band, mode)
if os.path.isfile(os.path.join(results_path, fig_name)):
os.remove(os.path.join(results_path, fig_name))
fig.savefig(os.path.join(results_path, fig_name))
# plt.close(fig)
# # Plot the histograms
# fig_hist, axes = plt.subplots(2, 2, figsize=(10, 10))
# for i in range(4):
# ax = axes.flatten()[i]
# ax.hist(rms_maps[i].flatten(), bins=np.arange(0, max_rms, 2), histtype='step', color='red')
# ax.set_title(r'IFU-%s' % ifu_sections[i])
# ax.set_xlabel(r'RMS WFE [nm]')
# ax.set_ylabel(r'Frequency')
# ax.set_xlim([min_rms, max_rms])
# ax.set_xticks(np.arange(0, max_rms, 5))
#
# fig_name = "RMSMAP_%s_DETECTOR_SPEC_%s_HISTOGRAM" % (spaxel_scale, spectral_band)
# if os.path.isfile(os.path.join(results_path, fig_name)):
# os.remove(os.path.join(results_path, fig_name))
# fig_hist.savefig(os.path.join(results_path, fig_name))
# plt.close(fig_hist)
# Box and Whisker plots
rms_flats = [array.flatten() for array in rms_maps]
rms_flats = np.array(rms_flats).T
data = pd.DataFrame(rms_flats, columns=ifu_sections)
fig_box, ax = plt.subplots(1, 1)
sns.boxplot(data=data, ax=ax)
ax.set_ylabel(r'RMS WFE [nm]')
ax.set_title(r'RMS WFE Detector | %s scale | %s band | %s' %
(spaxel_scale, spectral_band, mode))
fig_name = "RMSMAP_%s_DETECTOR_SPEC_%s_BOXWHISKER_MODE_%s" % (
spaxel_scale, spectral_band, mode)
if os.path.isfile(os.path.join(results_path, fig_name)):
os.remove(os.path.join(results_path, fig_name))
fig_box.savefig(os.path.join(results_path, fig_name))
plt.close(fig_box)
return rms_field, obj_xy, foc_xy, rms_maps, object_coord, focal_coord
def detector_ensquared_energy(zosapi, sys_mode, ao_modes, spaxel_scale, grating, N_rays, alpha, files_path, results_path):
"""
Calculate the Ensquared Energy at the detector plane for all 4 IFU channels
The Ensquared Energy is calculated for the Field Point at the centre of the slice
for all slices and all wavelengths
:param zosapi: the Python Standalone Application
:param spaxel_scale: str, spaxel scale to analyze
:param grating: str, grating to use
:param N_rays: number of rays to sample the pupil with
:param files_path: path to the Zemax E2E model files
:param results_path: path to save the results
:return:
"""
analysis_dir = os.path.join(results_path, 'ENSQUARED ENERGY')
print("Analysis Results will be saved in folder: ", analysis_dir)
if not os.path.exists(analysis_dir):
os.mkdir(analysis_dir)
analysis = e2e.EnsquaredEnergyAnalysis(zosapi=zosapi)
waves = np.linspace(1, 23, 5).astype(int)
ifu_sections = ['AB', 'CD', 'EF', 'GH']
fig, axes = plt.subplots(2, 2, figsize=(15, 10))
for i, ifu in enumerate(ifu_sections):
options = {'which_system': sys_mode, 'AO_modes': ao_modes, 'scales': [spaxel_scale], 'IFUs': [ifu],
'grating': [grating]}
list_results = analysis.loop_over_files(files_dir=files_path, files_opt=options, results_path=results_path,
wavelength_idx=waves, configuration_idx=None, N_rays=N_rays,
alpha=alpha)
# No Monte Carlo so the list_results only contains 1 entry, for the IFU channel
energy, obj_xy, slicer_xy, detector_xy, wavelengths = list_results[0]
# Separate the Odd / Even configurations to avoid triangulating over the gap on the detector plane
ener_ogg, ener_even = energy[:, ::2].flatten(), energy[:, 1::2].flatten()
x, y = detector_xy[:, :, 0].flatten(), detector_xy[:, :, 1].flatten()
x_odd, y_odd = detector_xy[:, ::2, 0].flatten(), detector_xy[:, ::2, 1].flatten()
x_even, y_even = detector_xy[:, 1::2, 0].flatten(), detector_xy[:, 1::2, 1].flatten()
triang_odd = tri.Triangulation(x_odd, y_odd)
triang_even = tri.Triangulation(x_even, y_even)
#Remove the flat triangles at the detector edges that appear because of the field curvature
min_circle_ratio = .05
mask_odd = tri.TriAnalyzer(triang_odd).get_flat_tri_mask(min_circle_ratio)
triang_odd.set_mask(mask_odd)
mask_even = tri.TriAnalyzer(triang_even).get_flat_tri_mask(min_circle_ratio)
triang_even.set_mask(mask_even)
min_ener, max_ener = np.min(energy), np.max(energy)
ax = axes.flatten()[i]
tpc_odd = ax.tripcolor(triang_odd, ener_ogg, shading='flat', cmap='Blues', vmin=min_ener, vmax=1.0)
tpc_odd.set_clim(vmin=min_ener, vmax=1.0)
tpc_even = ax.tripcolor(triang_even, ener_even, shading='flat', cmap='Blues', vmin=min_ener, vmax=1.0)
tpc_even.set_clim(vmin=min_ener, vmax=1.0)
ax.scatter(x, y, s=2, color='black')
axis_label = 'Detector'
ax.set_xlabel(axis_label + r' X [mm]')
ax.set_ylabel(axis_label + r' Y [mm]')
ax.set_aspect('equal')
plt.colorbar(tpc_odd, ax=ax, orientation='horizontal')
title = r'IFU-%s | %s mas | %s SPEC | EE min:%.2f max:%.2f' % (ifu, spaxel_scale, grating, min_ener, max_ener)
ax.set_title(title)
save_path = os.path.join(results_path, analysis_dir)
fig_name = "ENSQ_ENERGY_DETECTOR_%s_SPEC_%s_MODE_%s" % (spaxel_scale, grating, sys_mode)
if os.path.isfile(os.path.join(save_path, fig_name)):
os.remove(os.path.join(save_path, fig_name))
fig.savefig(os.path.join(save_path, fig_name))
return
