From 43d62649fa99fb887ab43795d42a3fe98539efd7 Mon Sep 17 00:00:00 2001 From: Jaren Ashcraft Date: Tue, 26 Dec 2023 11:31:17 -1000 Subject: [PATCH] black linting --- poke/interfaces.py | 162 ++++++----- poke/materials.py | 91 +++--- poke/plotting.py | 97 +++---- poke/poke_core.py | 342 ++++++++++++++-------- poke/poke_math.py | 233 +++++++++------ poke/polarization.py | 471 +++++++++++------------------- poke/raytrace.py | 666 +++++++++++++++++++++++-------------------- poke/thinfilms.py | 157 +++++----- poke/writing.py | 50 ++-- 9 files changed, 1170 insertions(+), 1099 deletions(-) diff --git a/poke/interfaces.py b/poke/interfaces.py index cae2210..93afdaf 100644 --- a/poke/interfaces.py +++ b/poke/interfaces.py @@ -1,7 +1,10 @@ """interfaces with POPPY, HCIPy""" from poke.poke_math import np -def jones_pupil_to_hcipy_wavefront(jones_pupil,pupil_grid,input_stokes_vector=[1,0,0,0],shape=None): + +def jones_pupil_to_hcipy_wavefront( + jones_pupil, pupil_grid, input_stokes_vector=[1, 0, 0, 0], shape=None +): """converts a poke jones pupil to an HCIPy partially polarized wavefront, only works on square jones pupils @@ -23,31 +26,37 @@ def jones_pupil_to_hcipy_wavefront(jones_pupil,pupil_grid,input_stokes_vector=[1 """ # First test the import try: - + import hcipy - except Exception as e: - - print(f'Error trying to import HCIPy \n {e}') + except Exception as e: + + print(f"Error trying to import HCIPy \n {e}") # Next test the ability to reshape the jones pupil # TODO: Add option to fit data to Zernike polynomials if shape is None: - size = jones_pupil[-1][...,0,0].shape[0] + size = jones_pupil[-1][..., 0, 0].shape[0] shape = np.sqrt(size) shape = int(np.sqrt(size)) - jones_reshaped = jones_pupil[-1][...,:2,:2] - field = hcipy.Field([[jones_reshaped[...,0,0],jones_reshaped[...,0,1]], - [jones_reshaped[...,1,0],jones_reshaped[...,1,1]]],pupil_grid) - - wavefront = hcipy.Wavefront(field,input_stokes_vector=input_stokes_vector) + jones_reshaped = jones_pupil[-1][..., :2, :2] + field = hcipy.Field( + [ + [jones_reshaped[..., 0, 0], jones_reshaped[..., 0, 1]], + [jones_reshaped[..., 1, 0], jones_reshaped[..., 1, 1]], + ], + pupil_grid, + ) + + wavefront = hcipy.Wavefront(field, input_stokes_vector=input_stokes_vector) return wavefront -def jones_pupil_to_poppy_wavefronts(jones_pupil,wavelength=1e-6,shape=None): + +def jones_pupil_to_poppy_wavefronts(jones_pupil, wavelength=1e-6, shape=None): """converts a Poke jones pupil to a POPPY wavefront list Parameters @@ -70,23 +79,31 @@ def jones_pupil_to_poppy_wavefronts(jones_pupil,wavelength=1e-6,shape=None): import astropy.units as u except Exception as e: - print(f'Error trying to import POPPY and/or astropy \n {e}') + print(f"Error trying to import POPPY and/or astropy \n {e}") if shape is None: - size = jones_pupil[-1][...,0,0].shape[0] + size = jones_pupil[-1][..., 0, 0].shape[0] shape = int(np.sqrt(size)) - jones_reshaped = jones_pupil[-1][...,:2,:2].reshape([shape,shape,2,2]) - jones_pupils = [jones_reshaped[...,0,0],jones_reshaped[...,0,1],jones_reshaped[...,1,0],jones_reshaped[...,1,1]] + jones_reshaped = jones_pupil[-1][..., :2, :2].reshape([shape, shape, 2, 2]) + jones_pupils = [ + jones_reshaped[..., 0, 0], + jones_reshaped[..., 0, 1], + jones_reshaped[..., 1, 0], + jones_reshaped[..., 1, 1], + ] wflist = [] for jones in jones_pupils: - wf = poppy.Wavefront(wavelength=wavelength*u.m,npix=shape,diam=1*u.m,oversample=1) + wf = poppy.Wavefront(wavelength=wavelength * u.m, npix=shape, diam=1 * u.m, oversample=1) wf.wavefront = jones wflist.append(wf) return wflist -def rayfront_to_hcipy_wavefront(rayfront,npix,pupil_grid,nmodes=11,input_stokes_vector=[1,0,0,0],which=-1): + +def rayfront_to_hcipy_wavefront( + rayfront, npix, pupil_grid, nmodes=11, input_stokes_vector=[1, 0, 0, 0], which=-1 +): """convert rayfront to an hcipy wavefront using zernike decomposition Parameters @@ -112,19 +129,25 @@ def rayfront_to_hcipy_wavefront(rayfront,npix,pupil_grid,nmodes=11,input_stokes_ # First test the import try: - + import hcipy except Exception as e: - print(f'Error trying to import HCIPy \n {e}') - - jones_pupil = regularly_space_jones(rayfront,nmodes,npix,which=which) - field = hcipy.Field([[jones_pupil[...,0,0].ravel(),jones_pupil[...,0,1].ravel()], - [jones_pupil[...,1,0].ravel(),jones_pupil[...,1,1].ravel()]],pupil_grid) - wavefront = hcipy.Wavefront(field,input_stokes_vector=input_stokes_vector) + print(f"Error trying to import HCIPy \n {e}") + + jones_pupil = regularly_space_jones(rayfront, nmodes, npix, which=which) + field = hcipy.Field( + [ + [jones_pupil[..., 0, 0].ravel(), jones_pupil[..., 0, 1].ravel()], + [jones_pupil[..., 1, 0].ravel(), jones_pupil[..., 1, 1].ravel()], + ], + pupil_grid, + ) + wavefront = hcipy.Wavefront(field, input_stokes_vector=input_stokes_vector) return wavefront + def zernike(rho, phi, J): """Generates an array containing Zernike polynomials contributed by Emory Jenkins with edits made by Jaren Ashcraft @@ -143,30 +166,35 @@ def zernike(rho, phi, J): values : numpy.ndarray array containing the Zernike modes """ - N=int(np.ceil(np.sqrt(2*J + 0.25)-0.5)) # N = number of rows on the zernike pyramid - values = np.zeros([rho.size, J+1]) - j=0 # ANSI index of zernike - for n in range(0,N): - m=-n - while m<=n: + N = int(np.ceil(np.sqrt(2 * J + 0.25) - 0.5)) # N = number of rows on the zernike pyramid + values = np.zeros([rho.size, J + 1]) + j = 0 # ANSI index of zernike + for n in range(0, N): + m = -n + while m <= n: R = 0 - for k in range(0,1+(n-abs(m))//2): - c_k = ((-1)**k * np.math.factorial(n-k))/(np.math.factorial(k) * np.math.factorial((n+abs(m))//2 - k) * np.math.factorial((n-abs(m))//2 - k)) - R += c_k * rho**(n-2*k) - if m>0: - Z = R*np.cos(m*phi) - elif m<0: - Z = R*np.sin(abs(m)*phi) + for k in range(0, 1 + (n - abs(m)) // 2): + c_k = ((-1) ** k * np.math.factorial(n - k)) / ( + np.math.factorial(k) + * np.math.factorial((n + abs(m)) // 2 - k) + * np.math.factorial((n - abs(m)) // 2 - k) + ) + R += c_k * rho ** (n - 2 * k) + if m > 0: + Z = R * np.cos(m * phi) + elif m < 0: + Z = R * np.sin(abs(m) * phi) else: Z = R - values[:,j] = Z - j=j+1 - if j>J: + values[:, j] = Z + j = j + 1 + if j > J: break - m=m+2 + m = m + 2 return values -def regularly_space_jones(rayfront,nmodes,npix,which=-1,return_residuals=False): + +def regularly_space_jones(rayfront, nmodes, npix, which=-1, return_residuals=False): """converts a jones pupil from a rayfront to a regularly-spaced array with zernike decomposition Parameters @@ -191,33 +219,29 @@ def regularly_space_jones(rayfront,nmodes,npix,which=-1,return_residuals=False): jones_pupil = rayfront.jones_pupil[which] # TODO: This breaks for decentered pupils, need to implement an offset - x,y = rayfront.xData[0,0],rayfront.yData[0,0] - x = x/np.max(x) - y = y/np.max(y) - r,t = np.sqrt(x**2 + y**2),np.arctan2(y,x) - irregularly_spaced_basis = zernike(r,t,nmodes) - - cxx = np.linalg.lstsq(irregularly_spaced_basis,jones_pupil[...,0,0],rcond=None) - cxy = np.linalg.lstsq(irregularly_spaced_basis,jones_pupil[...,0,1],rcond=None) - cyx = np.linalg.lstsq(irregularly_spaced_basis,jones_pupil[...,1,0],rcond=None) - cyy = np.linalg.lstsq(irregularly_spaced_basis,jones_pupil[...,1,1],rcond=None) - - x = np.linspace(-1,1,npix) - x,y = np.meshgrid(x,x) - r,t = np.sqrt(x**2 + y**2),np.arctan2(y,x) - regularly_spaced_basis = zernike(r.ravel(),t.ravel(),nmodes) - - return_jones = np.empty([npix,npix,2,2],dtype=np.complex128) - return_jones[...,0,0] = np.sum(regularly_spaced_basis*cxx[0],axis=-1).reshape([npix,npix]) - return_jones[...,0,1] = np.sum(regularly_spaced_basis*cxy[0],axis=-1).reshape([npix,npix]) - return_jones[...,1,0] = np.sum(regularly_spaced_basis*cyx[0],axis=-1).reshape([npix,npix]) - return_jones[...,1,1] = np.sum(regularly_spaced_basis*cyy[0],axis=-1).reshape([npix,npix]) + x, y = rayfront.xData[0, 0], rayfront.yData[0, 0] + x = x / np.max(x) + y = y / np.max(y) + r, t = np.sqrt(x ** 2 + y ** 2), np.arctan2(y, x) + irregularly_spaced_basis = zernike(r, t, nmodes) + + cxx = np.linalg.lstsq(irregularly_spaced_basis, jones_pupil[..., 0, 0], rcond=None) + cxy = np.linalg.lstsq(irregularly_spaced_basis, jones_pupil[..., 0, 1], rcond=None) + cyx = np.linalg.lstsq(irregularly_spaced_basis, jones_pupil[..., 1, 0], rcond=None) + cyy = np.linalg.lstsq(irregularly_spaced_basis, jones_pupil[..., 1, 1], rcond=None) + + x = np.linspace(-1, 1, npix) + x, y = np.meshgrid(x, x) + r, t = np.sqrt(x ** 2 + y ** 2), np.arctan2(y, x) + regularly_spaced_basis = zernike(r.ravel(), t.ravel(), nmodes) + + return_jones = np.empty([npix, npix, 2, 2], dtype=np.complex128) + return_jones[..., 0, 0] = np.sum(regularly_spaced_basis * cxx[0], axis=-1).reshape([npix, npix]) + return_jones[..., 0, 1] = np.sum(regularly_spaced_basis * cxy[0], axis=-1).reshape([npix, npix]) + return_jones[..., 1, 0] = np.sum(regularly_spaced_basis * cyx[0], axis=-1).reshape([npix, npix]) + return_jones[..., 1, 1] = np.sum(regularly_spaced_basis * cyy[0], axis=-1).reshape([npix, npix]) if return_residuals: - return return_jones, [cxx,cxy,cyx,cyy] + return return_jones, [cxx, cxy, cyx, cyy] else: return return_jones - - - - \ No newline at end of file diff --git a/poke/materials.py b/poke/materials.py index 88187a7..a364603 100644 --- a/poke/materials.py +++ b/poke/materials.py @@ -3,19 +3,28 @@ from pathlib import Path # Silica == Fused Silica -avail_materials = ['Al','Ag', # metals - 'HfO2','SiO2','Ta2O5','TiO2','Nb2O5', # Oxides - 'SiN', # Nitrides - 'MgF2','CaF2','LiF', # Fluorides - 'Silica' # Glasses - ] +avail_materials = [ + "Al", + "Ag", # metals + "HfO2", + "SiO2", + "Ta2O5", + "TiO2", + "Nb2O5", # Oxides + "SiN", # Nitrides + "MgF2", + "CaF2", + "LiF", # Fluorides + "Silica", # Glasses +] + def get_abs_path(file): - fullpth = Path(__file__).parent/'material_data'/file + fullpth = Path(__file__).parent / "material_data" / file return fullpth -def create_index_model(material,verbose=False): +def create_index_model(material, verbose=False): """creates an interpolated material based on data available from refractiveindex.info "[refractiveindex.info] is made freely available under the CC0 1.0 Universal Public Domain Dedication. @@ -42,65 +51,65 @@ def create_index_model(material,verbose=False): k_start = None if material not in avail_materials: - print(f'Material {material} not recognized') - print('Materials supported:') + print(f"Material {material} not recognized") + print("Materials supported:") print(avail_materials) # Load materials - this will be a lot - if material == 'Al': - pth = get_abs_path('Cheng_Al.csv') + if material == "Al": + pth = get_abs_path("Cheng_Al.csv") n_end = 427 k_start = n_end + 3 - elif material == 'Ag': - pth = get_abs_path('Ciesielski_Ag.csv') + elif material == "Ag": + pth = get_abs_path("Ciesielski_Ag.csv") n_end = 333 k_start = n_end + 3 - elif material == 'HfO2': - pth = get_abs_path('Kuhaili_HfO2.csv') - elif material == 'SiO2': - pth = get_abs_path('Lemarchand_SiO2.csv') + elif material == "HfO2": + pth = get_abs_path("Kuhaili_HfO2.csv") + elif material == "SiO2": + pth = get_abs_path("Lemarchand_SiO2.csv") n_end = 451 k_start = n_end + 3 - elif material == 'SiN': - pth = get_abs_path('Philipp_SiN.csv') - elif material == 'MgF2': - pth = get_abs_path('Rodriguez-de Marcos_MgF2.csv') + elif material == "SiN": + pth = get_abs_path("Philipp_SiN.csv") + elif material == "MgF2": + pth = get_abs_path("Rodriguez-de Marcos_MgF2.csv") n_end = 960 k_start = n_end + 3 - elif material == 'CaF2': - pth = get_abs_path('Daimon_CaF2.csv') - elif material == 'LiF': - pth = get_abs_path('Li_LiF.csv') - elif material == 'Ta2O5': - pth = get_abs_path('Gao_Ta2O5.csv') + elif material == "CaF2": + pth = get_abs_path("Daimon_CaF2.csv") + elif material == "LiF": + pth = get_abs_path("Li_LiF.csv") + elif material == "Ta2O5": + pth = get_abs_path("Gao_Ta2O5.csv") n_end = 726 k_start = n_end + 3 - elif material == 'TiO2': - pth = get_abs_path('Sarkar_TiO2.csv') + elif material == "TiO2": + pth = get_abs_path("Sarkar_TiO2.csv") n_end = 977 k_start = n_end + 3 - elif material == 'Nb2O5': - pth = get_abs_path('Lemarchand_Nb2O5.csv') + elif material == "Nb2O5": + pth = get_abs_path("Lemarchand_Nb2O5.csv") n_end = 451 k_start = n_end + 3 - elif material == 'Silica': - pth = get_abs_path('Malitson_Silica.csv') + elif material == "Silica": + pth = get_abs_path("Malitson_Silica.csv") # bunch of conditionals on if we have extinction coefficients if n_end is not None: - n_data = np.genfromtxt(pth,skip_header=1,delimiter=',')[:n_end].T + n_data = np.genfromtxt(pth, skip_header=1, delimiter=",")[:n_end].T if k_start is not None: - k_data = np.genfromtxt(pth,skip_header=1,delimiter=',')[k_start:].T + k_data = np.genfromtxt(pth, skip_header=1, delimiter=",")[k_start:].T else: - n_data = np.genfromtxt(pth,skip_header=1,delimiter=',').T + n_data = np.genfromtxt(pth, skip_header=1, delimiter=",").T # create the index splines - n_model = interp1d(n_data[0],n_data[1]) + n_model = interp1d(n_data[0], n_data[1]) if k_start is not None: - k_model = interp1d(k_data[0],k_data[1]) - index_model = lambda wvl: n_model(wvl) + 1j*k_model(wvl) + k_model = interp1d(k_data[0], k_data[1]) + index_model = lambda wvl: n_model(wvl) + 1j * k_model(wvl) else: index_model = n_model - return index_model \ No newline at end of file + return index_model diff --git a/poke/plotting.py b/poke/plotting.py index d901004..84295ea 100644 --- a/poke/plotting.py +++ b/poke/plotting.py @@ -2,25 +2,8 @@ import matplotlib as mpl from poke.poke_math import np -# Goal is to set-up publication-ready plots for PSF's and Jones Pupils - -# Default params -params = { - 'image.origin':'lower', - 'image.interpolation':'nearest', - 'image.cmap':'magma', - 'axes.labelsize':18, - 'axes.titlesize':18, - 'font.size':14, - 'xtick.labelsize':10, - 'ytick.labelsize':10, - 'figure.figsize':[3.39,2.10], - 'font.family':'sans serif', -} - -mpl.rcParams.update(params) - -def jones_pupil(raybundle,which=-1): + +def jones_pupil(raybundle, which=-1): """plot the jones pupil Parameters @@ -34,16 +17,16 @@ def jones_pupil(raybundle,which=-1): x = raybundle.xData[0][0] y = raybundle.yData[0][0] Jmat = raybundle.jones_pupil[which] - - fig,axs = plt.subplots(figsize=[12,5],nrows=2,ncols=4) - plt.suptitle('Jones Pupil') + + fig, axs = plt.subplots(figsize=[12, 5], nrows=2, ncols=4) + plt.suptitle("Jones Pupil") for j in range(2): for k in range(2): - ax = axs[j,k] - ax.set_title('|J{j}{k}|'.format(j=j,k=k)) - sca = ax.scatter(x,y,c=np.abs(Jmat[...,j,k]),cmap='inferno') - fig.colorbar(sca,ax=ax) - + ax = axs[j, k] + ax.set_title("|J{j}{k}|".format(j=j, k=k)) + sca = ax.scatter(x, y, c=np.abs(Jmat[..., j, k]), cmap="inferno") + fig.colorbar(sca, ax=ax) + # turn off the ticks if j != 1: ax.xaxis.set_visible(False) @@ -52,29 +35,30 @@ def jones_pupil(raybundle,which=-1): for j in range(2): for k in range(2): - + # Offset the p coefficient if j == 1: if k == 1: - offset = 0#-np.pi + offset = 0 # -np.pi else: offset = 0 else: offset = 0 - ax = axs[j,k+2] - ax.set_title(r'$\angle$' + 'J{j}{k}'.format(j=j,k=k)) - sca = ax.scatter(x,y,c=np.angle(Jmat[...,j,k])+offset,cmap='coolwarm') - fig.colorbar(sca,ax=ax) - + ax = axs[j, k + 2] + ax.set_title(r"$\angle$" + "J{j}{k}".format(j=j, k=k)) + sca = ax.scatter(x, y, c=np.angle(Jmat[..., j, k]) + offset, cmap="coolwarm") + fig.colorbar(sca, ax=ax) + # turn off the ticks if j != 1: ax.xaxis.set_visible(False) - + ax.yaxis.set_visible(False) plt.show() -def ray_opd(raybundle,which=-1): + +def ray_opd(raybundle, which=-1): """plot the OPD of the ray trace Parameters @@ -85,41 +69,44 @@ def ray_opd(raybundle,which=-1): Which index of the jones pupil list to plot, by default -1 """ - x = raybundle.xData[0,0] - y = raybundle.yData[0,0] - opd = raybundle.opd[0,-1] + x = raybundle.xData[0, 0] + y = raybundle.yData[0, 0] + opd = raybundle.opd[0, -1] opd -= np.mean(opd) - plt.figure(figsize=[5,5]) - plt.title('OPD for raybundle [m]') - plt.scatter(x,y,c=opd,cmap='coolwarm') + plt.figure(figsize=[5, 5]) + plt.title("OPD for raybundle [m]") + plt.scatter(x, y, c=opd, cmap="coolwarm") plt.colorbar() plt.show() + def mueller_pupil(M): - fig,axs = plt.subplots(figsize=[12,12],nrows=4,ncols=4) - plt.suptitle('Mueller Pupil') + fig, axs = plt.subplots(figsize=[12, 12], nrows=4, ncols=4) + plt.suptitle("Mueller Pupil") for i in range(4): for j in range(4): - ax = axs[i,j] - ax.set_title('J{i}{j}'.format(i=i,j=j)) - sca = ax.imshow(M[i,j,:,:]) - fig.colorbar(sca,ax=ax) + ax = axs[i, j] + ax.set_title("J{i}{j}".format(i=i, j=j)) + sca = ax.imshow(M[i, j, :, :]) + fig.colorbar(sca, ax=ax) ax.axes.xaxis.set_visible(False) ax.axes.yaxis.set_visible(False) plt.show() + def point_spread_matrix(PSM): from matplotlib.colors import LogNorm - fig,axs = plt.subplots(figsize=[12,12],nrows=4,ncols=4) - plt.suptitle('Point-spread Matrix') + + fig, axs = plt.subplots(figsize=[12, 12], nrows=4, ncols=4) + plt.suptitle("Point-spread Matrix") for i in range(4): for j in range(4): - ax = axs[i,j] - ax.set_title('M{i}{j}'.format(i=i,j=j)) - sca = ax.imshow(PSM[...,i,j],cmap='coolwarm') - fig.colorbar(sca,ax=ax) - + ax = axs[i, j] + ax.set_title("M{i}{j}".format(i=i, j=j)) + sca = ax.imshow(PSM[..., i, j], cmap="coolwarm") + fig.colorbar(sca, ax=ax) + if i != 3: ax.axes.xaxis.set_visible(False) if j != 0: diff --git a/poke/poke_core.py b/poke/poke_core.py index 8faf3a8..aff4cb4 100644 --- a/poke/poke_core.py +++ b/poke/poke_core.py @@ -4,6 +4,7 @@ from poke.poke_math import np import poke.plotting as plot import poke.polarization as pol + # import poke.gbd as gbd import poke.beamlets as beam import poke.raytrace as rt @@ -16,12 +17,22 @@ 3) No plotting/writing here, call other functions """ -GOLDEN = (1 + np.sqrt(5))/2 - -class Rayfront: +GOLDEN = (1 + np.sqrt(5)) / 2 - def __init__(self,nrays,wavelength,pupil_radius,max_fov,normalized_pupil_radius=1,fov=[0.,0.],waist_pad=None,circle=True,grid='even'): +class Rayfront: + def __init__( + self, + nrays, + wavelength, + pupil_radius, + max_fov, + normalized_pupil_radius=1, + fov=[0.0, 0.0], + waist_pad=None, + circle=True, + grid="even", + ): """class for the Rayfront object that 1) traces rays with the zosapi @@ -54,56 +65,58 @@ def __init__(self,nrays,wavelength,pupil_radius,max_fov,normalized_pupil_radius= """ - self.nrays = nrays # rays across a square pupil + self.nrays = nrays # rays across a square pupil self.wavelength = wavelength self.pupil_radius = pupil_radius - self.normalized_pupil_radius = normalized_pupil_radius # normalized radius + self.normalized_pupil_radius = normalized_pupil_radius # normalized radius self.max_fov = max_fov self.fov = np.array(fov) - self.normalized_fov = self.fov/max_fov - self.raybundle_extent = pupil_radius*normalized_pupil_radius # the actual extent of the raybundle + self.normalized_fov = self.fov / max_fov + self.raybundle_extent = ( + pupil_radius * normalized_pupil_radius + ) # the actual extent of the raybundle # init rayset # init raysets - x = np.linspace(-self.raybundle_extent,self.raybundle_extent,nrays) - x,y = np.meshgrid(x,x) + x = np.linspace(-self.raybundle_extent, self.raybundle_extent, nrays) + x, y = np.meshgrid(x, x) X = x Y = y - + if circle == True: if waist_pad: wo = waist_pad else: wo = 0 - x = x[np.sqrt(X**2 + Y**2) < self.raybundle_extent-wo/4] - y = y[np.sqrt(X**2 + Y**2) < self.raybundle_extent-wo/4] - - if grid == 'fib': - i = len(x) # use however many rays are in a circular aperture with even sampling - n = np.arange(1,i) - Rn = np.sqrt(n/i) - Tn = 2*np.pi/GOLDEN**2 * n - x_fib = Rn*np.cos(Tn) - y_fib = Rn*np.sin(Tn) - x = x_fib + x = x[np.sqrt(X ** 2 + Y ** 2) < self.raybundle_extent - wo / 4] + y = y[np.sqrt(X ** 2 + Y ** 2) < self.raybundle_extent - wo / 4] + + if grid == "fib": + i = len(x) # use however many rays are in a circular aperture with even sampling + n = np.arange(1, i) + Rn = np.sqrt(n / i) + Tn = 2 * np.pi / GOLDEN ** 2 * n + x_fib = Rn * np.cos(Tn) + y_fib = Rn * np.sin(Tn) + x = x_fib y = y_fib - x = np.ravel(x)/pupil_radius - y = np.ravel(y)/pupil_radius + x = np.ravel(x) / pupil_radius + y = np.ravel(y) / pupil_radius - print('norm fov = ',self.normalized_fov) + print("norm fov = ", self.normalized_fov) # in normalized pupil and field coords for an on-axis field - self.base_rays = np.array([x, - y, - 0*x + self.normalized_fov[0], - 0*y + self.normalized_fov[1]]) - print('base ray shape ',self.base_rays.shape) + self.base_rays = np.array( + [x, y, 0 * x + self.normalized_fov[0], 0 * y + self.normalized_fov[1]] + ) + print("base ray shape ", self.base_rays.shape) + # First optional constructors of our core physics modules - #@classmethod - def as_gaussianbeamlets(self,wo): + # @classmethod + def as_gaussianbeamlets(self, wo): """optional constructor to init the rayfront for GBD, comes with additional args @@ -115,18 +128,18 @@ def as_gaussianbeamlets(self,wo): # gaussian beam parameters self.wo = wo - self.div = self.wavelength/(np.pi*self.wo) * 180 / np.pi # beam divergence in deg + self.div = self.wavelength / (np.pi * self.wo) * 180 / np.pi # beam divergence in deg # ray differentials in normalized coords - dPx = self.wo/self.pupil_radius - dPy = self.wo/self.pupil_radius - dHx = self.div/self.max_fov - dHy = self.div/self.max_fov + dPx = self.wo / self.pupil_radius + dPy = self.wo / self.pupil_radius + dHx = self.div / self.max_fov + dHy = self.div / self.max_fov # differential ray bundles from base rays self.Px_rays = np.copy(self.base_rays) - if np.__name__ == 'jax.numpy': + if np.__name__ == "jax.numpy": self.Px_rays = self.Px_rays.at[0].set(dPx) self.Py_rays = np.copy(self.base_rays) @@ -151,7 +164,7 @@ def as_gaussianbeamlets(self,wo): self.Hy_rays[3] += dHy # total set of rays - self.raysets = [self.base_rays,self.Px_rays,self.Py_rays,self.Hx_rays,self.Hy_rays] + self.raysets = [self.base_rays, self.Px_rays, self.Py_rays, self.Hx_rays, self.Hy_rays] # Will force the transverse coords to be x and y self.global_coords = False @@ -162,9 +175,8 @@ def as_gaussianbeamlets(self,wo): self.dHx = dHx self.dHy = dHy - - #@classmethod - def as_polarized(self,surfaces): + # @classmethod + def as_polarized(self, surfaces): """optional constructor to init the rayfront for PRT, comes with additional args @@ -178,7 +190,7 @@ def as_polarized(self,surfaces): } """ - self._surfaces = surfaces # a list of dictionaries + self._surfaces = surfaces # a list of dictionaries self.raysets = [self.base_rays] self.global_coords = True self.P_total = [] @@ -188,7 +200,7 @@ def as_polarized(self,surfaces): ########################### GENERAL RAY TRACING METHODS ########################### """ - def trace_rayset(self,pth,wave=1,surfaces=None,_experimental=True): + def trace_rayset(self, pth, wave=1, surfaces=None, _experimental=True): """ Parameters ---------- @@ -204,14 +216,19 @@ def trace_rayset(self,pth,wave=1,surfaces=None,_experimental=True): if surfaces != None: self._surfaces = surfaces - if (pth[-3:] == 'zmx') or (pth[-3:] == 'zos'): - positions,directions,normals,self.opd,self.vignetted = rt.trace_through_zos(self.raysets,pth,self._surfaces,self.nrays,wave,self.global_coords) - elif (pth[-3:] == 'seq') or (pth[-3:] == 'len'): + if (pth[-3:] == "zmx") or (pth[-3:] == "zos"): + positions, directions, normals, self.opd, self.vignetted = rt.trace_through_zos( + self.raysets, pth, self._surfaces, self.nrays, wave, self.global_coords + ) + elif (pth[-3:] == "seq") or (pth[-3:] == "len"): if _experimental: - positions,directions,normals,self.opd = rt.trace_through_cv(self.raysets,pth,self._surfaces,self.nrays,wave,self.global_coords) + positions, directions, normals, self.opd = rt.trace_through_cv( + self.raysets, pth, self._surfaces, self.nrays, wave, self.global_coords + ) else: - positions,directions,normals,self.opd = rt.TraceThroughCV(self.raysets,pth,self._surfaces,self.nrays,wave,self.global_coords) - + positions, directions, normals, self.opd = rt.TraceThroughCV( + self.raysets, pth, self._surfaces, self.nrays, wave, self.global_coords + ) self.xData = positions[0] self.yData = positions[1] @@ -220,7 +237,7 @@ def trace_rayset(self,pth,wave=1,surfaces=None,_experimental=True): self.lData = directions[0] self.mData = directions[1] self.nData = directions[2] - + # Keep sign in raytracer coordinate system self.l2Data = normals[0] self.m2Data = normals[1] @@ -230,7 +247,14 @@ def trace_rayset(self,pth,wave=1,surfaces=None,_experimental=True): ########################### GAUSSIAN BEAMLET TRACING METHODS ########################### """ - def beamlet_decomposition_field(self,dcoords,dnorms=np.array([0.,0.,1.]),memory_avail=4,misaligned=True,vignette=True): + def beamlet_decomposition_field( + self, + dcoords, + dnorms=np.array([0.0, 0.0, 1.0]), + memory_avail=4, + misaligned=True, + vignette=True, + ): """computes the coherent field by decomposing the entrance pupil into gaussian beams and propagating them to the final surface @@ -246,40 +270,95 @@ def beamlet_decomposition_field(self,dcoords,dnorms=np.array([0.,0.,1.]),memory_ """ # converting memory - nrays = self.nData[:,-1].shape[1] - npix = dcoords.shape[-1] # need to have coords in first dimension and be raveled - print('pixels = ',npix) - print('rays = ',nrays) - total_size = nrays*npix*128 * 1e-9 # complex128, 4 is a fudge factor to account for intermediate variables - total_blocks = total_size/memory_avail + nrays = self.nData[:, -1].shape[1] + npix = dcoords.shape[-1] # need to have coords in first dimension and be raveled + print("pixels = ", npix) + print("rays = ", nrays) + total_size = ( + nrays * npix * 128 * 1e-9 + ) # complex128, 4 is a fudge factor to account for intermediate variables + total_blocks = total_size / memory_avail nloops = np.ceil(total_blocks) if nloops < 1: nloops = 1 - print(f'beamlet field at wavelength = {self.wavelength}') + print(f"beamlet field at wavelength = {self.wavelength}") if misaligned: if vignette: - field = beam.misaligned_beamlet_field(self.xData,self.yData,self.zData,self.lData,self.mData,self.nData,self.opd, - self.wo,self.wo,self.div*np.pi/180,self.div*np.pi/180, dcoords,dnorms, - wavelength=self.wavelength,nloops=nloops,use_centroid=True,vignetting=self.vignetted) + field = beam.misaligned_beamlet_field( + self.xData, + self.yData, + self.zData, + self.lData, + self.mData, + self.nData, + self.opd, + self.wo, + self.wo, + self.div * np.pi / 180, + self.div * np.pi / 180, + dcoords, + dnorms, + wavelength=self.wavelength, + nloops=nloops, + use_centroid=True, + vignetting=self.vignetted, + ) else: - field = beam.misaligned_beamlet_field(self.xData,self.yData,self.zData,self.lData,self.mData,self.nData,self.opd, - self.wo,self.wo,self.div*np.pi/180,self.div*np.pi/180, dcoords,dnorms, - wavelength=self.wavelength,nloops=nloops,use_centroid=True) + field = beam.misaligned_beamlet_field( + self.xData, + self.yData, + self.zData, + self.lData, + self.mData, + self.nData, + self.opd, + self.wo, + self.wo, + self.div * np.pi / 180, + self.div * np.pi / 180, + dcoords, + dnorms, + wavelength=self.wavelength, + nloops=nloops, + use_centroid=True, + ) else: - field = beam.beamlet_decomposition_field(self.xData,self.yData,self.zData,self.lData,self.mData,self.nData,self.opd, - self.wo,self.wo,self.div*np.pi/180,self.div*np.pi/180, dcoords,dnorms, - wavelength=self.wavelength,nloops=int(nloops),use_centroid=True,vignetting=self.vignetted) - + field = beam.beamlet_decomposition_field( + self.xData, + self.yData, + self.zData, + self.lData, + self.mData, + self.nData, + self.opd, + self.wo, + self.wo, + self.div * np.pi / 180, + self.div * np.pi / 180, + dcoords, + dnorms, + wavelength=self.wavelength, + nloops=int(nloops), + use_centroid=True, + vignetting=self.vignetted, + ) + return field - """ ########################### POLARIZATION RAY TRACING METHODS ########################### """ - def compute_jones_pupil(self,ambient_index=1,aloc=np.array([0.,0.,1.]),entrance_x=np.array([1.,0.,0.]),exit_x=np.array([1.,0.,0.]),proper_retardance=False): + def compute_jones_pupil( + self, + ambient_index=1, + aloc=np.array([0.0, 0.0, 1.0]), + entrance_x=np.array([1.0, 0.0, 0.0]), + exit_x=np.array([1.0, 0.0, 0.0]), + proper_retardance=False, + ): """compute jones pupil from ray data using the double pole coordinate system Parameters @@ -297,76 +376,91 @@ def compute_jones_pupil(self,ambient_index=1,aloc=np.array([0.,0.,1.]),entrance_ """ if proper_retardance: - warnings.warn('The proper retardance calculation is prone to unphysical results and requires further testing') - - for rayset_ind,rayset in enumerate(self.raysets): - - - aoi,kin,kout,norm = rt.convert_ray_data_to_prt_data(self.lData[rayset_ind],self.mData[rayset_ind],self.nData[rayset_ind], - self.l2Data[rayset_ind],self.m2Data[rayset_ind],self.n2Data[rayset_ind], - self._surfaces) - - Psys,Jsys,Qsys = pol.system_prt_matrices(aoi,kin,kout,norm,self._surfaces,self.wavelength,ambient_index) - P,Q = pol.total_prt_matrix(Psys,Qsys) + warnings.warn( + "The proper retardance calculation is prone to unphysical results and requires further testing" + ) + + for rayset_ind, rayset in enumerate(self.raysets): + + aoi, kin, kout, norm = rt.convert_ray_data_to_prt_data( + self.lData[rayset_ind], + self.mData[rayset_ind], + self.nData[rayset_ind], + self.l2Data[rayset_ind], + self.m2Data[rayset_ind], + self.n2Data[rayset_ind], + self._surfaces, + ) + + Psys, Jsys, Qsys = pol.system_prt_matrices( + aoi, kin, kout, norm, self._surfaces, self.wavelength, ambient_index + ) + P, Q = pol.total_prt_matrix(Psys, Qsys) if proper_retardance: - Jpupil = pol.global_to_local_coordinates(P,kin[0],kout[-1],aloc,entrance_x,exit_x,Q=Q) + Jpupil = pol.global_to_local_coordinates( + P, kin[0], kout[-1], aloc, entrance_x, exit_x, Q=Q + ) else: - Jpupil = pol.global_to_local_coordinates(P,kin[0],kout[-1],aloc,entrance_x,exit_x) + Jpupil = pol.global_to_local_coordinates( + P, kin[0], kout[-1], aloc, entrance_x, exit_x + ) self.jones_pupil.append(Jpupil) self.P_total.append(P) - def compute_arm(self,pad=2,circle=True,is_square=True): + def compute_arm(self, pad=2, circle=True, is_square=True): """Computes the amplitude response matrix from the Jones Pupil, requires a square array """ if is_square: - - J = self.JonesPupil[-1][:,:2,:2] + + J = self.JonesPupil[-1][:, :2, :2] J_dim = int(np.sqrt(J.shape[0])) - J = np.reshape(J,[J_dim,J_dim,2,2]) + J = np.reshape(J, [J_dim, J_dim, 2, 2]) else: # Expand into a polynomial basis - J = inter.regularly_space_jones(self,11,self.nrays) - - A = np.empty([J_dim*pad,J_dim*pad,2,2],dtype='complex128') - + J = inter.regularly_space_jones(self, 11, self.nrays) + + A = np.empty([J_dim * pad, J_dim * pad, 2, 2], dtype="complex128") + # Create a circular aperture - x = np.linspace(-1,1,J.shape[0]) - x,y = np.meshgrid(x,x) - mask = np.ones([J.shape[0],J.shape[0]]) - + x = np.linspace(-1, 1, J.shape[0]) + x, y = np.meshgrid(x, x) + mask = np.ones([J.shape[0], J.shape[0]]) + if circle: - mask[x**2 + y**2 > 1] = 0 + mask[x ** 2 + y ** 2 > 1] = 0 for i in range(2): for j in range(2): - A[...,i,j] = np.fft.fftshift(np.fft.fft2(np.pad(J[...,i,j]*mask,int(J_dim*pad/2-(J_dim/2))))) - + A[..., i, j] = np.fft.fftshift( + np.fft.fft2(np.pad(J[..., i, j] * mask, int(J_dim * pad / 2 - (J_dim / 2)))) + ) + self.ARM = A return A - - def compute_psm(self,cut=128,stokes=np.array([1.,0.,0.,0.])): - + + def compute_psm(self, cut=128, stokes=np.array([1.0, 0.0, 0.0, 0.0])): + """ We regrettably need to loop over this because we use numpy.kron() """ - + # cut out the center - size = self.ARM.shape[0]/2 - A = self.ARM[int(size-cut):int(size+cut),int(size-cut):int(size+cut)] - P = np.empty([A.shape[0],A.shape[1],4,4]) + size = self.ARM.shape[0] / 2 + A = self.ARM[int(size - cut) : int(size + cut), int(size - cut) : int(size + cut)] + P = np.empty([A.shape[0], A.shape[1], 4, 4]) for i in range(A.shape[0]): for j in range(A.shape[1]): - - P[i,j] = pol.jones_to_mueller(A[i,j]) - + + P[i, j] = pol.jones_to_mueller(A[i, j]) + img = P @ stokes self.PSM = P - return img[...,0] - + return img[..., 0] + """ ########################### PROPERTIES ########################### """ @@ -374,39 +468,43 @@ def compute_psm(self,cut=128,stokes=np.array([1.,0.,0.,0.])): @property def surfaces(self): return self._surfaces - + @surfaces.setter - def surfaces(self,surflist): + def surfaces(self, surflist): self._surfaces = surflist """ ########################### Source Module Conversions ########################### """ - def convert_data_sourcemodule(self,new_backend='numpy'): + def convert_data_sourcemodule(self, new_backend="numpy"): """This is a bit cursed, but in the case where data is initialized in numpy, but we want to use it in Jax/Cupy, then we have to convert it and vice versa """ - from poke.poke_math import np,set_backend_to_cupy,set_backend_to_jax,set_backend_to_numpy # make sure we have the current source module loaded + from poke.poke_math import ( + np, + set_backend_to_cupy, + set_backend_to_jax, + set_backend_to_numpy, + ) # make sure we have the current source module loaded - if new_backend == 'numpy': + if new_backend == "numpy": set_backend_to_numpy() - elif new_backend == 'jax': - + elif new_backend == "jax": + set_backend_to_jax() - elif new_backend == 'cupy': + elif new_backend == "cupy": set_backend_to_cupy() else: - print('Did not recognize module, defaulting to numpy') + print("Did not recognize module, defaulting to numpy") set_backend_to_numpy() - # Ray Data self.xData = np.asarray(self.xData) self.yData = np.asarray(self.yData) @@ -420,4 +518,4 @@ def convert_data_sourcemodule(self,new_backend='numpy'): self.m2Data = np.asarray(self.m2Data) self.n2Data = np.asarray(self.n2Data) - self.opd = np.asarray(self.opd) \ No newline at end of file + self.opd = np.asarray(self.opd) diff --git a/poke/poke_math.py b/poke/poke_math.py index bbe0b78..f76af2f 100644 --- a/poke/poke_math.py +++ b/poke/poke_math.py @@ -1,35 +1,43 @@ import numpy as np + class BackendShim: """A shim that allows a backend to be swapped at runtime. Taken from prysm.mathops with permission from Brandon Dube """ + def __init__(self, src): self._srcmodule = src def __getattr__(self, key): - if key == '_srcmodule': + if key == "_srcmodule": return self._srcmodule return getattr(self._srcmodule, key) + _np = np np = BackendShim(_np) + def set_backend_to_numpy(): """Convenience method to automatically configure poke's backend to cupy.""" import numpy as cp + np._srcmodule = cp - + return + def set_backend_to_cupy(): """Convenience method to automatically configure poke's backend to cupy.""" import cupy as cp + np._srcmodule = cp return + def set_backend_to_jax(): """Convenience method to automatically configure poke's backend to cupy.""" @@ -38,14 +46,16 @@ def set_backend_to_jax(): # jax defaults to 32 bit but we need 64bit from jax.config import config + config.update("jax_enable_x64", True) np._srcmodule = cp - print('source module switched to ',np.__name__) + print("source module switched to ", np.__name__) return + def det_2x2(array): """compute determinant of 2x2 matrix, broadcasted @@ -59,15 +69,16 @@ def det_2x2(array): det determinant array """ - a = array[...,0,0] - b = array[...,0,1] - c = array[...,1,0] - d = array[...,1,1] + a = array[..., 0, 0] + b = array[..., 0, 1] + c = array[..., 1, 0] + d = array[..., 1, 1] - det = (a*d - b*c) + det = a * d - b * c return det + def mat_inv_2x2(array): """compute inverse of 2x2 matrix, broadcasted @@ -82,20 +93,21 @@ def mat_inv_2x2(array): matrix inverse array """ - a = array[...,0,0] - b = array[...,0,1] - c = array[...,1,0] - d = array[...,1,1] + a = array[..., 0, 0] + b = array[..., 0, 1] + c = array[..., 1, 0] + d = array[..., 1, 1] - det = (a*d - b*c) + det = a * d - b * c - matinv = np.array([[d,-b],[-c,a]]) / det + matinv = np.array([[d, -b], [-c, a]]) / det if matinv.ndim > 2: - for i in range(matinv.ndim-2): - matinv = np.moveaxis(matinv,-1,0) - + for i in range(matinv.ndim - 2): + matinv = np.moveaxis(matinv, -1, 0) + return matinv + def mat_inv_3x3(array): """compute inverse of 3x3 matrix, broadcasted @@ -110,40 +122,39 @@ def mat_inv_3x3(array): matrix inverse array """ - a = array[...,0,0] # row 1 - b = array[...,0,1] - c = array[...,0,2] - - d = array[...,1,0] # row 2 - e = array[...,1,1] - f = array[...,1,2] - - g = array[...,2,0] # row 3 - h = array[...,2,1] - i = array[...,2,2] + a = array[..., 0, 0] # row 1 + b = array[..., 0, 1] + c = array[..., 0, 2] + + d = array[..., 1, 0] # row 2 + e = array[..., 1, 1] + f = array[..., 1, 2] + + g = array[..., 2, 0] # row 3 + h = array[..., 2, 1] + i = array[..., 2, 2] # determine cofactor elements - ac = e*i - f*h - bc = -(d*i - f*g) - cc = d*h - e*g - dc = -(b*i - c*h) - ec = a*i - c*g - fc = -(a*h - b*g) - gc = b*f - c*e - hc = -(a*f - c*d) - ic = a*e - b*d + ac = e * i - f * h + bc = -(d * i - f * g) + cc = d * h - e * g + dc = -(b * i - c * h) + ec = a * i - c * g + fc = -(a * h - b * g) + gc = b * f - c * e + hc = -(a * f - c * d) + ic = a * e - b * d # get determinant - det = a*ac + b*bc + c*cc # second term's negative is included in cofactor term - det = det[...,np.newaxis,np.newaxis] + det = a * ac + b * bc + c * cc # second term's negative is included in cofactor term + det = det[..., np.newaxis, np.newaxis] # Assemble adjucate matrix (transpose of cofactor) - arrayinv = np.asarray([[ac,bc,cc], - [dc,ec,fc], - [gc,hc,ic]]).T / det - + arrayinv = np.asarray([[ac, bc, cc], [dc, ec, fc], [gc, hc, ic]]).T / det + return arrayinv + def eigenvalues_2x2(array): """ Computes the eigenvalues of a 2x2 matrix using a trick @@ -157,17 +168,18 @@ def eigenvalues_2x2(array): The eigenvalues of the array """ - a = array[...,0,0] - b = array[...,0,1] - c = array[...,1,0] - d = array[...,1,1] + a = array[..., 0, 0] + b = array[..., 0, 1] + c = array[..., 1, 0] + d = array[..., 1, 1] - determinant = (a*d - b*c) - mean_ondiag = (a+d)/2 - e1 = mean_ondiag + np.sqrt(mean_ondiag**2 - determinant) - e2 = mean_ondiag - np.sqrt(mean_ondiag**2 - determinant) + determinant = a * d - b * c + mean_ondiag = (a + d) / 2 + e1 = mean_ondiag + np.sqrt(mean_ondiag ** 2 - determinant) + e2 = mean_ondiag - np.sqrt(mean_ondiag ** 2 - determinant) + + return e1, e2 - return e1,e2 def vector_norm(vector): """computes the magnitude of a vector @@ -182,13 +194,14 @@ def vector_norm(vector): numpy.ndarray magnitude of the vector """ - vx = vector[...,0] * vector[...,0] - vy = vector[...,1] * vector[...,1] - vz = vector[...,2] * vector[...,2] + vx = vector[..., 0] * vector[..., 0] + vy = vector[..., 1] * vector[..., 1] + vz = vector[..., 2] * vector[..., 2] return np.sqrt(vx + vy + vz) -def vector_angle(u,v): + +def vector_angle(u, v): """computes the vector angle between two vectors Parameters @@ -203,28 +216,31 @@ def vector_angle(u,v): ndarray vector of angle between u and v in x, y, z in radians """ - u = u/(vector_norm(u)[...,np.newaxis]) - v = v/(vector_norm(v)[...,np.newaxis]) + u = u / (vector_norm(u)[..., np.newaxis]) + v = v / (vector_norm(v)[..., np.newaxis]) - dot = np.sum(u*v,axis=-1) + dot = np.sum(u * v, axis=-1) angles = np.zeros_like(dot) # Make exceptions for angles turning around if dot.any() < 0: if np.__name__ == "jax.numpy": - angles = angles.at[dot < 0].set((np.pi - 2*np.arcsin(vector_norm(-v-u)/2))[dot < 0]) + angles = angles.at[dot < 0].set( + (np.pi - 2 * np.arcsin(vector_norm(-v - u) / 2))[dot < 0] + ) else: - angles[dot < 0] = (np.pi - 2*np.arcsin(vector_norm(-v-u)/2))[dot < 0] + angles[dot < 0] = (np.pi - 2 * np.arcsin(vector_norm(-v - u) / 2))[dot < 0] elif dot.any() >= 0: if np.__name__ == "jax.numpy": - angles = angles.at[dot >= 0].set((2*np.arcsin(vector_norm(v-u)/2))[dot >= 0]) + angles = angles.at[dot >= 0].set((2 * np.arcsin(vector_norm(v - u) / 2))[dot >= 0]) else: - angles[dot >= 0] = (2*np.arcsin(vector_norm(v-u)/2))[dot >= 0] - + angles[dot >= 0] = (2 * np.arcsin(vector_norm(v - u) / 2))[dot >= 0] + return angles -def rotation_3d(angle,axis): + +def rotation_3d(angle, axis): """Rotation matrix about an axis by an angle Parameters @@ -241,17 +257,35 @@ def rotation_3d(angle,axis): """ c = np.cos(angle) s = np.sin(angle) - mat = np.array([[(1-c)*axis[...,0]**2 + c, (1-c)*axis[...,0]*axis[...,1] - s*axis[...,2], (1-c)*axis[...,0]*axis[...,2] + s*axis[...,1]], - [(1-c)*axis[...,1]*axis[...,0] + s*axis[...,2], (1-c)*axis[...,1]**2 + c, (1-c)*axis[...,1]*axis[...,2] - s*axis[...,0]], - [(1-c)*axis[...,2]*axis[...,0] - s*axis[...,1], (1-c)*axis[...,1]*axis[...,2] + s*axis[...,0], (1-c)*axis[...,2]**2 + c]]) + mat = np.array( + [ + [ + (1 - c) * axis[..., 0] ** 2 + c, + (1 - c) * axis[..., 0] * axis[..., 1] - s * axis[..., 2], + (1 - c) * axis[..., 0] * axis[..., 2] + s * axis[..., 1], + ], + [ + (1 - c) * axis[..., 1] * axis[..., 0] + s * axis[..., 2], + (1 - c) * axis[..., 1] ** 2 + c, + (1 - c) * axis[..., 1] * axis[..., 2] - s * axis[..., 0], + ], + [ + (1 - c) * axis[..., 2] * axis[..., 0] - s * axis[..., 1], + (1 - c) * axis[..., 1] * axis[..., 2] + s * axis[..., 0], + (1 - c) * axis[..., 2] ** 2 + c, + ], + ] + ) if mat.ndim > 2: - for i in range(mat.ndim-2): - mat = np.moveaxis(mat,-1,0) + for i in range(mat.ndim - 2): + mat = np.moveaxis(mat, -1, 0) return mat + "Vector Operations from Quinn Jarecki" -def rotation3D(angle,axis): + +def rotation3D(angle, axis): """Rotation matrix about an axis by an angle Parameters @@ -268,12 +302,29 @@ def rotation3D(angle,axis): """ c = np.cos(angle) s = np.sin(angle) - mat = np.array([[(1-c)*axis[0]**2 + c, (1-c)*axis[0]*axis[1] - s*axis[2], (1-c)*axis[0]*axis[2] + s*axis[1]], - [(1-c)*axis[1]*axis[0] + s*axis[2], (1-c)*axis[1]**2 + c, (1-c)*axis[1]*axis[2] - s*axis[0]], - [(1-c)*axis[2]*axis[0] - s*axis[1], (1-c)*axis[1]*axis[2] + s*axis[0], (1-c)*axis[2]**2 + c]]) + mat = np.array( + [ + [ + (1 - c) * axis[0] ** 2 + c, + (1 - c) * axis[0] * axis[1] - s * axis[2], + (1 - c) * axis[0] * axis[2] + s * axis[1], + ], + [ + (1 - c) * axis[1] * axis[0] + s * axis[2], + (1 - c) * axis[1] ** 2 + c, + (1 - c) * axis[1] * axis[2] - s * axis[0], + ], + [ + (1 - c) * axis[2] * axis[0] - s * axis[1], + (1 - c) * axis[1] * axis[2] + s * axis[0], + (1 - c) * axis[2] ** 2 + c, + ], + ] + ) return mat -def vectorAngle(u,v): + +def vectorAngle(u, v): """computes the vector angle between two vectors Parameters @@ -288,13 +339,14 @@ def vectorAngle(u,v): ndarray vector of angle between u and v in x, y, z in radians """ - u = u/vector_norm(u) - v = v/vector_norm(v) + u = u / vector_norm(u) + v = v / vector_norm(v) - if u@v<0: # dot product - return np.pi - 2*np.arcsin(np.linalg.norm(-v-u)/2) + if u @ v < 0: # dot product + return np.pi - 2 * np.arcsin(np.linalg.norm(-v - u) / 2) else: - return 2*np.arcsin(np.linalg.norm(v-u)/2) + return 2 * np.arcsin(np.linalg.norm(v - u) / 2) + def pauli_spin_matrix(i): @@ -312,17 +364,16 @@ def pauli_spin_matrix(i): """ if i == 0: - - return np.array([[1,0],[0,1]]) - + + return np.array([[1, 0], [0, 1]]) + if i == 1: - - return np.array([[1,0],[0,-1]]) - + + return np.array([[1, 0], [0, -1]]) + if i == 2: - return np.array([[0,1],[1,0]]) - + return np.array([[0, 1], [1, 0]]) + if i == 3: - - return np.array([[0,-1j],[1j,0]]) - \ No newline at end of file + + return np.array([[0, -1j], [1j, 0]]) diff --git a/poke/polarization.py b/poke/polarization.py index f8c8b7f..7e9cef8 100644 --- a/poke/polarization.py +++ b/poke/polarization.py @@ -1,5 +1,13 @@ # dependencies -from poke.poke_math import np,mat_inv_3x3,vector_norm,vector_angle,rotation_3d,vectorAngle,rotation3D +from poke.poke_math import ( + np, + mat_inv_3x3, + vector_norm, + vector_angle, + rotation_3d, + vectorAngle, + rotation3D, +) import poke.thinfilms as tf import poke.poke_math as math import matplotlib.pyplot as plt @@ -17,7 +25,8 @@ # ax[row,column].scatter(x,y,c=op(raybundle.P_total[0][...,row,column])) # plt.show() -def fresnel_coefficients(aoi,n1,n2,mode='reflect'): + +def fresnel_coefficients(aoi, n1, n2, mode="reflect"): """Computes Fresnel Coefficients for a single surface interaction Parameters @@ -35,26 +44,31 @@ def fresnel_coefficients(aoi,n1,n2,mode='reflect'): the Fresnel s- and p- coefficients of the surface interaction """ - if (mode != 'reflect') and (mode != 'transmit'): - print('not a valid mode, please use reflect, transmit, or both. Defaulting to reflect') - mode = 'reflect' + if (mode != "reflect") and (mode != "transmit"): + print("not a valid mode, please use reflect, transmit, or both. Defaulting to reflect") + mode = "reflect" # ratio of refractive indices - n = n2/n1 + n = n2 / n1 + + if mode == "reflect": - if mode == 'reflect': + fs = (np.cos(aoi) - np.sqrt(n ** 2 - np.sin(aoi) ** 2)) / ( + np.cos(aoi) + np.sqrt(n ** 2 - np.sin(aoi) ** 2) + ) + fp = (n ** 2 * np.cos(aoi) - np.sqrt(n ** 2 - np.sin(aoi) ** 2)) / ( + n ** 2 * np.cos(aoi) + np.sqrt(n ** 2 - np.sin(aoi) ** 2) + ) # * np.exp(-1j*np.pi) - fs = (np.cos(aoi) - np.sqrt(n**2 - np.sin(aoi)**2))/(np.cos(aoi) + np.sqrt(n**2 - np.sin(aoi)**2)) - fp = (n**2 * np.cos(aoi) - np.sqrt(n**2 - np.sin(aoi)**2))/(n**2 * np.cos(aoi) + np.sqrt(n**2 - np.sin(aoi)**2)) # * np.exp(-1j*np.pi) + elif mode == "transmit": - elif mode == 'transmit': + fs = (2 * np.cos(aoi)) / (np.cos(aoi) + np.sqrt(n ** 2 - np.sin(aoi) ** 2)) + fp = (2 * n * np.cos(aoi)) / (n ** 2 * np.cos(aoi) + np.sqrt(n ** 2 - np.sin(aoi) ** 2)) - fs = (2*np.cos(aoi))/(np.cos(aoi) + np.sqrt(n**2 - np.sin(aoi)**2)) - fp = (2*n*np.cos(aoi))/(n**2 * np.cos(aoi) + np.sqrt(n**2 - np.sin(aoi)**2)) + return fs, fp - return fs,fp -def orthogonal_transofrmation_matrices(kin,kout,normal): +def orthogonal_transofrmation_matrices(kin, kout, normal): """compute the orthogonal transformation matrices that rotate into and out of the local coordinates of a surface @@ -74,39 +88,40 @@ def orthogonal_transofrmation_matrices(kin,kout,normal): """ # ensure wave vectors are normalized - kin = kin / vector_norm(kin)[...,np.newaxis] - kout = kout / vector_norm(kout)[...,np.newaxis] + kin = kin / vector_norm(kin)[..., np.newaxis] + kout = kout / vector_norm(kout)[..., np.newaxis] # get s-basis vector - sin = np.cross(kin,normal) - sin = sin / vector_norm(sin)[...,np.newaxis] + sin = np.cross(kin, normal) + sin = sin / vector_norm(sin)[..., np.newaxis] # get p-basis vector - pin = np.cross(kin,sin) - pin = pin / vector_norm(pin)[...,np.newaxis] + pin = np.cross(kin, sin) + pin = pin / vector_norm(pin)[..., np.newaxis] # Assemble Oinv - Oinv = np.array([sin,pin,kin]) - Oinv = np.moveaxis(Oinv,-1,0) - if Oinv.ndim >2: + Oinv = np.array([sin, pin, kin]) + Oinv = np.moveaxis(Oinv, -1, 0) + if Oinv.ndim > 2: for i in range(Oinv.ndim - 2): - Oinv = np.moveaxis(Oinv,-1,0) - Oinv = np.swapaxes(Oinv,-1,-2) # take the transpose/inverse + Oinv = np.moveaxis(Oinv, -1, 0) + Oinv = np.swapaxes(Oinv, -1, -2) # take the transpose/inverse # outgoing basis vectors sout = sin - pout = np.cross(kout,sout) - pout = pout / vector_norm(pout)[...,np.newaxis] - Oout = np.array([sout,pout,kout]) - Oout = np.moveaxis(Oout,-1,0) - if Oout.ndim >2: + pout = np.cross(kout, sout) + pout = pout / vector_norm(pout)[..., np.newaxis] + Oout = np.array([sout, pout, kout]) + Oout = np.moveaxis(Oout, -1, 0) + if Oout.ndim > 2: for i in range(Oout.ndim - 2): - Oout = np.moveaxis(Oout,-1,0) + Oout = np.moveaxis(Oout, -1, 0) # Oout = np.moveaxis(Oout,0,-1) - return Oinv,Oout + return Oinv, Oout -def prt_matrix(kin,kout,normal,aoi,surfdict,wavelength,ambient_index): + +def prt_matrix(kin, kout, normal, aoi, surfdict, wavelength, ambient_index): """prt matrix for a single surface Parameters @@ -136,65 +151,80 @@ def prt_matrix(kin,kout,normal,aoi,surfdict,wavelength,ambient_index): offdiagbool = False # A surface decision tree - TODO: it is worth trying to make this more robust - if type(surfdict['coating']) == list: + if type(surfdict["coating"]) == list: # prysm likes films in degress, wavelength in microns, thickness in microns - rs,ts = tf.compute_thin_films_broadcasted(surfdict['coating'],aoi,wavelength,substrate_index=surfdict['coating'][-1],polarization='s') - rp,tp = tf.compute_thin_films_broadcasted(surfdict['coating'],aoi,wavelength,substrate_index=surfdict['coating'][-1],polarization='p') - - if surfdict['mode'] == 'reflect': + rs, ts = tf.compute_thin_films_broadcasted( + surfdict["coating"], + aoi, + wavelength, + substrate_index=surfdict["coating"][-1], + polarization="s", + ) + rp, tp = tf.compute_thin_films_broadcasted( + surfdict["coating"], + aoi, + wavelength, + substrate_index=surfdict["coating"][-1], + polarization="p", + ) + + if surfdict["mode"] == "reflect": fss = rs - fpp = rp * np.exp(-1j*np.pi) # The Thin Film Correction + fpp = rp * np.exp(-1j * np.pi) # The Thin Film Correction - if surfdict['mode'] == 'transmit': + if surfdict["mode"] == "transmit": fss = ts fpp = tp - elif type(surfdict['coating']) == np.ndarray: # assumes the film is defined with first index as fs,fp - - fss = surfdict['coating'][0,0] - fsp = surfdict['coating'][0,1] - fps = surfdict['coating'][1,0] - fpp = surfdict['coating'][1,1] + elif ( + type(surfdict["coating"]) == np.ndarray + ): # assumes the film is defined with first index as fs,fp + + fss = surfdict["coating"][0, 0] + fsp = surfdict["coating"][0, 1] + fps = surfdict["coating"][1, 0] + fpp = surfdict["coating"][1, 1] offdiagbool = True - elif callable(surfdict['coating']): # check if a function - fss,fps = surfdict['coating'](aoi) + elif callable(surfdict["coating"]): # check if a function + fss, fps = surfdict["coating"](aoi) else: - fss,fpp = fresnel_coefficients(aoi,ambient_index,surfdict['coating'],mode=surfdict['mode']) - if np.imag(surfdict['coating']) < 0: # TODO: This is a correction for the n - ik configuration, need to investigate if physical - fss *= np.exp(-1j*np.pi) - fpp *= np.exp(1j*np.pi) + fss, fpp = fresnel_coefficients( + aoi, ambient_index, surfdict["coating"], mode=surfdict["mode"] + ) + if ( + np.imag(surfdict["coating"]) < 0 + ): # TODO: This is a correction for the n - ik configuration, need to investigate if physical + fss *= np.exp(-1j * np.pi) + fpp *= np.exp(1j * np.pi) - Oinv,Oout = orthogonal_transofrmation_matrices(kin,kout,normal) + Oinv, Oout = orthogonal_transofrmation_matrices(kin, kout, normal) # Compute the Jones matrix and parallel transport matrix zeros = np.zeros(fss.shape) ones = np.ones(fss.shape) if offdiagbool: - J = np.asarray([[fss,fsp,zeros], - [fps,fpp,zeros], - [zeros,zeros,ones]]) + J = np.asarray([[fss, fsp, zeros], [fps, fpp, zeros], [zeros, zeros, ones]]) else: - J = np.asarray([[fss,zeros,zeros], - [zeros,fpp,zeros], - [zeros,zeros,ones]]) - B = np.asarray([[1,0,0],[0,1,0],[0,0,1]]) + J = np.asarray([[fss, zeros, zeros], [zeros, fpp, zeros], [zeros, zeros, ones]]) + B = np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) # dimensions need to be appropriate if J.ndim > 2: - for _ in range(J.ndim-2): - J = np.moveaxis(J,-1,0) + for _ in range(J.ndim - 2): + J = np.moveaxis(J, -1, 0) # compute PRT matrix and orthogonal transformation Pmat = Oout @ J @ Oinv - Qmat = Oout @ B @ Oinv # test if this broadcasts + Qmat = Oout @ B @ Oinv # test if this broadcasts - return Pmat,J,Qmat + return Pmat, J, Qmat -def system_prt_matrices(aoi,kin,kout,norm,surfaces,wavelength,ambient_index): + +def system_prt_matrices(aoi, kin, kout, norm, surfaces, wavelength, ambient_index): """computes the PRT matrices for each surface in the optical system Parameters @@ -224,22 +254,24 @@ def system_prt_matrices(aoi,kin,kout,norm,surfaces,wavelength,ambient_index): J = [] Q = [] - - for i,surfdict in enumerate(surfaces): + for i, surfdict in enumerate(surfaces): + + kisurf = np.moveaxis(kin[i], -1, 0) + kosurf = np.moveaxis(kout[i], -1, 0) + normsurf = np.moveaxis(norm[i], -1, 0) + aoisurf = np.moveaxis(aoi[i], -1, 0) - kisurf = np.moveaxis(kin[i],-1,0) - kosurf = np.moveaxis(kout[i],-1,0) - normsurf = np.moveaxis(norm[i],-1,0) - aoisurf = np.moveaxis(aoi[i],-1,0) - - Pmat,Jmat,Qmat = prt_matrix(kisurf,kosurf,normsurf,aoisurf,surfdict,wavelength,ambient_index) + Pmat, Jmat, Qmat = prt_matrix( + kisurf, kosurf, normsurf, aoisurf, surfdict, wavelength, ambient_index + ) P.append(Pmat) J.append(Jmat) Q.append(Qmat) - return P,J,Q + return P, J, Q -def total_prt_matrix(P,Q): + +def total_prt_matrix(P, Q): """computes the total PRT matrix for the optical system Parameters @@ -255,19 +287,20 @@ def total_prt_matrix(P,Q): the total PRT and Parallel transport matrices """ - for i,(p,q) in enumerate(zip(P,Q)): + for i, (p, q) in enumerate(zip(P, Q)): if i == 0: Ptot = p Qtot = q - + else: Ptot = p @ Ptot Qtot = q @ Qtot - - return Ptot,Qtot -def global_to_local_coordinates(P,kin,k,a,xin,exit_x,Q=None): + return Ptot, Qtot + + +def global_to_local_coordinates(P, kin, k, a, xin, exit_x, Q=None): """Use the double pole basis to compute the local coordinate system of the Jones pupil. Vectorized to perform on arrays of arbitrary shape, assuming the PRT matrix is in the last two dimensions. @@ -304,28 +337,32 @@ def global_to_local_coordinates(P,kin,k,a,xin,exit_x,Q=None): # Double Pole Coordinate System, requires a rotation about an axis # Wikipedia article seems to disagree with CLY Example 11.4 # Default entrance pupil in Zemax. Note that this assumes the stop is at the first surface - kin = np.moveaxis(kin,-1,0) - k = np.moveaxis(k,-1,0) - xin = xin / vector_norm(xin)[...,np.newaxis] - xin = np.broadcast_to(xin,kin.shape) - yin = np.cross(kin,xin) - yin = yin / vector_norm(yin)[...,np.newaxis] - yin = np.broadcast_to(yin,kin.shape) - O_e = np.array([[xin[...,0],yin[...,0],kin[...,0]], - [xin[...,1],yin[...,1],kin[...,1]], - [xin[...,2],yin[...,2],kin[...,2]]]) - O_e = np.moveaxis(O_e,-1,0) + kin = np.moveaxis(kin, -1, 0) + k = np.moveaxis(k, -1, 0) + xin = xin / vector_norm(xin)[..., np.newaxis] + xin = np.broadcast_to(xin, kin.shape) + yin = np.cross(kin, xin) + yin = yin / vector_norm(yin)[..., np.newaxis] + yin = np.broadcast_to(yin, kin.shape) + O_e = np.array( + [ + [xin[..., 0], yin[..., 0], kin[..., 0]], + [xin[..., 1], yin[..., 1], kin[..., 1]], + [xin[..., 2], yin[..., 2], kin[..., 2]], + ] + ) + O_e = np.moveaxis(O_e, -1, 0) # Compute Exit Pupil Basis Vectors # For arbitrary k each ray will have it's own pair of basis vectors - r = np.cross(k,a) - r = r / vector_norm(r)[...,np.newaxis] # match shapes - th = -vector_angle(k,a) - R = rotation_3d(th,r) + r = np.cross(k, a) + r = r / vector_norm(r)[..., np.newaxis] # match shapes + th = -vector_angle(k, a) + R = rotation_3d(th, r) # Local basis vectors xout = exit_x - yout = np.cross(a,xout) + yout = np.cross(a, xout) yout /= vector_norm(yout) # add axes to match shapes @@ -334,20 +371,25 @@ def global_to_local_coordinates(P,kin,k,a,xin,exit_x,Q=None): x = R @ xout y = R @ yout - O_x = np.array([[x[...,0],y[...,0],k[...,0]], - [x[...,1],y[...,1],k[...,1]], - [x[...,2],y[...,2],k[...,2]]]) - - O_x = np.moveaxis(O_x,-1,0) + O_x = np.array( + [ + [x[..., 0], y[..., 0], k[..., 0]], + [x[..., 1], y[..., 1], k[..., 1]], + [x[..., 2], y[..., 2], k[..., 2]], + ] + ) + + O_x = np.moveaxis(O_x, -1, 0) # apply proper retardance correction if type(Q) == np.ndarray: - P = mat_inv_3x3(Q) @ P + P = mat_inv_3x3(Q) @ P J = mat_inv_3x3(O_x) @ P @ O_e return J + def jones_to_mueller(Jones): """Converts a Jones matrix to a Mueller matrix @@ -362,17 +404,15 @@ def jones_to_mueller(Jones): Mueller matrix from Jones matrix """ - U = np.array([[1,0,0,1], - [1,0,0,-1], - [0,1,1,0], - [0,1j,-1j,0]]) + U = np.array([[1, 0, 0, 1], [1, 0, 0, -1], [0, 1, 1, 0], [0, 1j, -1j, 0]]) - U *= np.sqrt(1/2) + U *= np.sqrt(1 / 2) - M = np.real(U @ (np.kron(np.conj(Jones),Jones)) @ np.linalg.inv(U)) + M = np.real(U @ (np.kron(np.conj(Jones), Jones)) @ np.linalg.inv(U)) return M + def mueller_to_jones(M): """Converts Mueller matrix to a relative Jones matrix. Phase aberration is relative to the Pxx component. @@ -385,206 +425,23 @@ def mueller_to_jones(M): "CLY Eq. 6.112" "Untested" - print('warning : This operation looses global phase') - - pxx = np.sqrt((M[0,0] + M[0,1] + M[1,0] + M[1,1])/2) - pxy = np.sqrt((M[0,0] - M[0,1] + M[1,0] - M[1,1])/2) - pyx = np.sqrt((M[0,0] + M[0,1] - M[1,0] - M[1,1])/2) - pyy = np.sqrt((M[0,0] - M[0,1] - M[1,0] + M[1,1])/2) + print("warning : This operation looses global phase") - txx = 0 # This phase is not determined - txy = -np.arctan((M[0,3]+M[1,3])/(M[0,2]+M[1,2])) - tyx = np.arctan((M[3,0]+M[3,1])/(M[2,0]+M[2,1])) - tyy = np.arctan((M[3,2]-M[2,3])/(M[2,2]+M[3,3])) + pxx = np.sqrt((M[0, 0] + M[0, 1] + M[1, 0] + M[1, 1]) / 2) + pxy = np.sqrt((M[0, 0] - M[0, 1] + M[1, 0] - M[1, 1]) / 2) + pyx = np.sqrt((M[0, 0] + M[0, 1] - M[1, 0] - M[1, 1]) / 2) + pyy = np.sqrt((M[0, 0] - M[0, 1] - M[1, 0] + M[1, 1]) / 2) - J = np.array([[pxx*np.exp(-1j*txx),pxy*np.exp(-1j*txy)], - [pyx*np.exp(-1j*tyx),pyy*np.exp(-1j*tyy)]]) - - return J + txx = 0 # This phase is not determined + txy = -np.arctan((M[0, 3] + M[1, 3]) / (M[0, 2] + M[1, 2])) + tyx = np.arctan((M[3, 0] + M[3, 1]) / (M[2, 0] + M[2, 1])) + tyy = np.arctan((M[3, 2] - M[2, 3]) / (M[2, 2] + M[3, 3])) + J = np.array( + [ + [pxx * np.exp(-1j * txx), pxy * np.exp(-1j * txy)], + [pyx * np.exp(-1j * tyx), pyy * np.exp(-1j * tyy)], + ] + ) -"""Beware weary traverler, below is a mueseum containing first-generation Poke code. Browse if you dare""" - -# # Step 2) Construct Orthogonal Transfer Matrices -# def ConstructOrthogonalTransferMatrices(kin,kout,normal,check_orthogonal=False): -# print('This function has been depreciated, please use orthogonal_transformation_matrices') -# """Construct the Orthogonal transformations to rotate from global to local coordinates and back again - -# Parameters -# ---------- -# kin : ndarray -# incident direction cosine vector -# kout : ndarray -# exiting direction cosine vector -# normal : ndarray -# direction cosine vector of the surface normal -# check_orthogonal : bool, optional -# prints the difference of the inverse(O) and transpose(O), should be apprx 0. by default False - -# Returns -# ------- -# Oinv,Oout : ndarrays -# orthogonal transformation matrices to rotate into the surface local coords (Oinv) and back into global coords (Oout) -# """ -# # PL&OS Page 326 Eq 9.5 - 9.7 -# # Construct Oin-1 with incident ray, say vectors are row vectors -# kin /= np.linalg.norm(kin) # these were not in chippman and lam - added 03/30/2022 -# kout /= np.linalg.norm(kout) - -# sin = np.cross(kin,normal) -# sin /= np.linalg.norm(sin) # normalize the s-vector -# pin = np.cross(kin,sin) -# pin /= np.linalg.norm(pin) -# Oinv = np.array([sin,pin,kin]) - -# sout = sin #np.cross(kout,normal) -# pout = np.cross(kout,sout) -# pout /= np.linalg.norm(pout) -# Oout = np.transpose(np.array([sout,pout,kout])) - -# if check_orthogonal == True: -# print('Oinv orthogonality : ',Oinv.transpose() == np.linalg.inv(Oinv)) -# print('Oout orthogonality : ',Oout.transpose() == np.linalg.inv(Oout)) - -# return Oinv,Oout - - -# # Step 3) Create Polarization Ray Trace matrix -# def ConstructPRTMatrix(kin,kout,normal,aoi,surfdict,wavelength,ambient_index): -# print('This function has been depreciated, please use system_prt_matrices') - -# """Assembles the PRT matrix, relies on the previous two functions - -# Parameters -# ---------- -# kin : ndarray -# incident direction cosine vector -# kout : ndarray -# exiting direction cosine vector -# normal : ndarray -# direction cosine vector of the surface normal -# aoi : float or array of floats -# angle of incidence in radians on the interface -# surfdict : dict -# dictionary that describe surfaces. Including surface number in raytrace, -# interaction mode, coating, etc. -# wavelength : float -# wavelength of light in meters -# ambient_index : float -# index of the medium that the optical system exists in - -# Returns -# ------- -# Pmat,J : ndarrays -# Pmat is the polarization ray tracing matrix, J is the same matrix without the orthogonal transformations -# """ - -# # negate to get to chipman sign convention from zemax sign convention -# normal = -normal - -# # Compute the Fresnel coefficients for either transmission OR reflection -# if type(surfdict['coating']) == list: - -# # prysm likes films in degress, wavelength in microns, thickness in microns -# rs,rp = tf.compute_thin_films_macleod(surfdict['coating'][:-1],aoi,wavelength,substrate_index=surfdict['coating'][-1]) -# ts,tp = 0,0 - -# if surfdict['mode'] == 'reflect': -# fs = rs -# fp = rp # * np.exp(-1j*np.pi) # The Thin Film Correction -# if surfdict['mode'] == 'transmit': -# fs = ts -# fp = tp - -# # Single surface coefficients -# else: - -# fs,fp = FresnelCoefficients(aoi,ambient_index,surfdict['coating'],mode=surfdict['mode']) - - -# # Compute the orthogonal transfer matrices -# Oinv,Oout = ConstructOrthogonalTransferMatrices(kin,kout,normal) - -# # Compute the Jones matrix -# J = np.array([[fs,0,0],[0,fp,0],[0,0,1]]) -# B = np.array([[1,0,0],[0,1,0],[0,0,1]]) - -# # Compute the Polarization Ray Tracing Matrix -# Pmat = Oout @ J @ Oinv -# Omat = Oout @ B @ Oinv # The parallel transport matrix, return when ready to implement. This will matter for berry phase - -# # This returns the polarization ray tracing matrix but I'm not 100% sure its in the coordinate system of the Jones Pupil -# return Pmat,J#,Omat - - -# def GlobalToLocalCoordinates(Pmat,kin,k,a,exit_x,check_orthogonal=False): -# print('This function has been depreciated, please use global_to_local_coordinates') -# """Use the double pole basis to compute the local coordinate system of the Jones pupil -# Chipman, Lam, Young, from Ch 11 : The Jones Pupil - -# Parameters -# ---------- -# Pmat : ndarray -# Pmat is the polarization ray tracing matrix -# kin : ndarray -# incident direction cosine vector at the entrance pupil -# kout : ndarray -# exiting direction cosine vector at the exit pupil -# a : ndarray -# vector in global coordinates describing the antipole direction -# exit_x : ndarray -# vector in global coordinates describing the direction that should be the -# "local x" direction -# check_orthogonal : bool, optional -# prints the difference of the inverse(O) and transpose(O), should be apprx 0. by default False - -# Returns -# ------- -# J : ndarray -# shape 3 x 3 ndarray containing the Jones pupil of the optical system. The elements -# Jtot[0,2], Jtot[1,2], Jtot[2,0], Jtot[2,1] should be zero. -# Jtot[-1,-1] should be 1 -# """ - - -# # Double Pole Coordinate System, requires a rotation about an axis -# # Wikipedia article seems to disagree with CLY Example 11.4 -# # Default entrance pupil in Zemax. Note that this assumes the stop is at the first surface -# xin = np.array([1.,0.,0.]) -# xin /= np.linalg.norm(xin) -# yin = np.cross(kin,xin) -# yin /= np.linalg.norm(yin) -# O_e = np.array([[xin[0],yin[0],kin[0]], -# [xin[1],yin[1],kin[1]], -# [xin[2],yin[2],kin[2]]]) - -# # Compute Exit Pupil Basis Vectors -# # For arbitrary k each ray will have it's own pair of basis vectors -# r = np.cross(k,a) -# r /= np.linalg.norm(r) -# th = -math.vectorAngle(k,a) -# R = math.rotation3D(th,r) - -# # Local basis vectors -# xout = exit_x -# yout = np.cross(a,xout) -# yout /= np.linalg.norm(yout) -# x = R @ xout -# x /= np.linalg.norm(x) -# y = R @ yout -# y /= np.linalg.norm(y) - -# O_x = np.array([[x[0],y[0],k[0]], -# [x[1],y[1],k[1]], -# [x[2],y[2],k[2]]]) - -# # Check orthogonality -# if check_orthogonal == True: -# print('O_x difference = ') -# print(O_x.transpose() - np.linalg.inv(O_x)) -# print('O_e difference = ') -# print(O_e.transpose() - np.linalg.inv(O_e)) - -# J = np.linalg.inv(O_x) @ Pmat @ O_e - -# return J + return J \ No newline at end of file diff --git a/poke/raytrace.py b/poke/raytrace.py index 2f64656..b36e0cc 100644 --- a/poke/raytrace.py +++ b/poke/raytrace.py @@ -4,7 +4,8 @@ import poke.thinfilms as tf import os -def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): + +def trace_through_zos(raysets, pth, surflist, nrays, wave, global_coords): """Traces initialized rays through a zemax opticstudio file Parameters @@ -56,11 +57,12 @@ def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): """ import zosapi - from System import Enum,Int32,Double,Array - import clr,os + from System import Enum, Int32, Double, Array + import clr, os + # known directory # dll = os.path.join(os.path.dirname(os.path.realpath(__file__)),r'RayTrace.dll') - dll = os.path.dirname(__file__)+'\RayTrace.dll' + dll = os.path.dirname(__file__) + "\RayTrace.dll" clr.AddReference(dll) import BatchRayTrace @@ -68,39 +70,39 @@ def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): zos = zosapi.App() TheSystem = zos.TheSystem ZOSAPI = zos.ZOSAPI - TheSystem.LoadFile(pth,False) + TheSystem.LoadFile(pth, False) # Check to make sure the ZOSAPI is working if TheSystem.LDE.NumberOfSurfaces < 4: - print('File was not loaded correctly') + print("File was not loaded correctly") exit() - - if surflist[-1]['surf'] > TheSystem.LDE.NumberOfSurfaces: - print('last surface > num surfaces, setting last surface to num surfaces') - surflist[-1]['surf'] = TheSystem.LDE.NumberOfSurfaces + + if surflist[-1]["surf"] > TheSystem.LDE.NumberOfSurfaces: + print("last surface > num surfaces, setting last surface to num surfaces") + surflist[-1]["surf"] = TheSystem.LDE.NumberOfSurfaces maxrays = raysets[0].shape[-1] # Dimension 0 is ray set, Dimension 1 is surface, dimension 2 is coordinate # Satisfies broadcasting rules! - xData = np.empty([len(raysets),len(surflist),maxrays]) - yData = np.empty([len(raysets),len(surflist),maxrays]) - zData = np.empty([len(raysets),len(surflist),maxrays]) + xData = np.empty([len(raysets), len(surflist), maxrays]) + yData = np.empty([len(raysets), len(surflist), maxrays]) + zData = np.empty([len(raysets), len(surflist), maxrays]) - lData = np.empty([len(raysets),len(surflist),maxrays]) - mData = np.empty([len(raysets),len(surflist),maxrays]) - nData = np.empty([len(raysets),len(surflist),maxrays]) + lData = np.empty([len(raysets), len(surflist), maxrays]) + mData = np.empty([len(raysets), len(surflist), maxrays]) + nData = np.empty([len(raysets), len(surflist), maxrays]) - l2Data = np.empty([len(raysets),len(surflist),maxrays]) - m2Data = np.empty([len(raysets),len(surflist),maxrays]) - n2Data = np.empty([len(raysets),len(surflist),maxrays]) + l2Data = np.empty([len(raysets), len(surflist), maxrays]) + m2Data = np.empty([len(raysets), len(surflist), maxrays]) + n2Data = np.empty([len(raysets), len(surflist), maxrays]) - mask = np.empty([len(raysets),len(surflist),maxrays]) + mask = np.empty([len(raysets), len(surflist), maxrays]) # Necessary for GBD calculations, might help PRT calculations - opd = np.empty([len(raysets),len(surflist),maxrays]) + opd = np.empty([len(raysets), len(surflist), maxrays]) - for rayset_ind,rayset in enumerate(raysets): + for rayset_ind, rayset in enumerate(raysets): # Get the normalized coordinates Px = rayset[0] @@ -108,24 +110,23 @@ def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): Hx = rayset[2] Hy = rayset[3] - for surf_ind,surfdict in enumerate(surflist): + for surf_ind, surfdict in enumerate(surflist): + + surf = surfdict["surf"] - surf = surfdict['surf'] - # Some ZOS-API setup tool = TheSystem.Tools.OpenBatchRayTrace() normUnpol = tool.CreateNormUnpol(maxrays, ZOSAPI.Tools.RayTrace.RaysType.Real, surf) reader = BatchRayTrace.ReadNormUnpolData(tool, normUnpol) reader.ClearData() - # THIS OVER-INITIALIZES THE NUMBER OF RAYS. + # THIS OVER-INITIALIZES THE NUMBER OF RAYS. # This initialization is weird because it requires allocating space for a square of rays - # so there will be extra which we remove later. + # so there will be extra which we remove later. rays = reader.InitializeOutput(nrays) # Add rays to reader - reader.AddRay(wave,Hx,Hy,Px,Py, - Enum.Parse(ZOSAPI.Tools.RayTrace.OPDMode,'None')) + reader.AddRay(wave, Hx, Hy, Px, Py, Enum.Parse(ZOSAPI.Tools.RayTrace.OPDMode, "None")) isfinished = False @@ -138,38 +139,60 @@ def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): # Global Coordinate Conversion # Have to pre-allocate a sysDbl for this method to execute sysDbl = Double(1.0) - success,R11,R12,R13,R21,R22,R23,R31,R32,R33,XO,YO,ZO = TheSystem.LDE.GetGlobalMatrix(int(surf), - sysDbl,sysDbl,sysDbl, - sysDbl,sysDbl,sysDbl, - sysDbl,sysDbl,sysDbl, - sysDbl,sysDbl,sysDbl) + ( + success, + R11, + R12, + R13, + R21, + R22, + R23, + R31, + R32, + R33, + XO, + YO, + ZO, + ) = TheSystem.LDE.GetGlobalMatrix( + int(surf), + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + sysDbl, + ) # Did the raytrace succeed? if success != 1: - print('Ray Failure at surface {}'.format(surf)) + print("Ray Failure at surface {}".format(surf)) # Global Rotation Matrix - Rmat = np.array([[R11,R12,R13], - [R21,R22,R23], - [R31,R32,R33]]) + Rmat = np.array([[R11, R12, R13], [R21, R22, R23], [R31, R32, R33]]) - position = np.array([np.array(list(rays.X)), - np.array(list(rays.Y)), - np.array(list(rays.Z))]) + position = np.array( + [np.array(list(rays.X)), np.array(list(rays.Y)), np.array(list(rays.Z))] + ) # I think this is just per-surface so it doesn't really need to be a big list, just a single surface. # TODO: Change later when cleaning up the code offset = np.zeros(position.shape) - offset[0,:] = XO - offset[1,:] = YO - offset[2,:] = ZO + offset[0, :] = XO + offset[1, :] = YO + offset[2, :] = ZO - angle = np.array([np.array(list(rays.L)), - np.array(list(rays.M)), - np.array(list(rays.N))]) + angle = np.array( + [np.array(list(rays.L)), np.array(list(rays.M)), np.array(list(rays.N))] + ) - normal = np.array([np.array(list(rays.l2)), - np.array(list(rays.m2)), - np.array(list(rays.n2))]) + normal = np.array( + [np.array(list(rays.l2)), np.array(list(rays.m2)), np.array(list(rays.n2))] + ) OPD = np.array(list(rays.opd)) @@ -178,51 +201,58 @@ def trace_through_zos(raysets,pth,surflist,nrays,wave,global_coords): # rotate into global coordinates - necessary for PRT if global_coords == True: - print('tracing with global coordinates') + print("tracing with global coordinates") position = offset + Rmat @ position angle = Rmat @ angle normal = Rmat @ normal # Filter the values at the end because ZOS allocates extra space - position = position[:,:maxrays] - angle = angle[:,:maxrays] - normal = normal[:,:maxrays] + position = position[:, :maxrays] + angle = angle[:, :maxrays] + normal = normal[:, :maxrays] OPD = OPD[:maxrays] # Append data to lists along the surface dimension - xData[rayset_ind,surf_ind] = position[0] - yData[rayset_ind,surf_ind] = position[1] - zData[rayset_ind,surf_ind] = position[2] + xData[rayset_ind, surf_ind] = position[0] + yData[rayset_ind, surf_ind] = position[1] + zData[rayset_ind, surf_ind] = position[2] - lData[rayset_ind,surf_ind] = angle[0] - mData[rayset_ind,surf_ind] = angle[1] - nData[rayset_ind,surf_ind] = angle[2] + lData[rayset_ind, surf_ind] = angle[0] + mData[rayset_ind, surf_ind] = angle[1] + nData[rayset_ind, surf_ind] = angle[2] - l2Data[rayset_ind,surf_ind] = normal[0] - m2Data[rayset_ind,surf_ind] = normal[1] - n2Data[rayset_ind,surf_ind] = normal[2] + l2Data[rayset_ind, surf_ind] = normal[0] + m2Data[rayset_ind, surf_ind] = normal[1] + n2Data[rayset_ind, surf_ind] = normal[2] # I don't think we need R and O, but might be useful to store just in case. Commenting out for now # R.append(Rmat) # O.append(offset) - opd[rayset_ind,surf_ind] = OPD - mask[rayset_ind,surf_ind] = rays_that_passed + opd[rayset_ind, surf_ind] = OPD + mask[rayset_ind, surf_ind] = rays_that_passed # always close your tools tool.Close() # This isn't necessary but makes the code more readable - positions = [xData,yData,zData] - directions = [lData,mData,nData] - normals = [l2Data,m2Data,n2Data] + positions = [xData, yData, zData] + directions = [lData, mData, nData] + normals = [l2Data, m2Data, n2Data] # Just a bit of celebration - print('{nrays} Raysets traced through {nsurf} surfaces'.format(nrays=rayset_ind+1,nsurf=surf_ind+1)) - + print( + "{nrays} Raysets traced through {nsurf} surfaces".format( + nrays=rayset_ind + 1, nsurf=surf_ind + 1 + ) + ) + # And finally return everything - return positions,directions,normals,opd,mask + return positions, directions, normals, opd, mask -def trace_through_cv(raysets,pth,surflist,nrays,wave,global_coords,global_coord_reference='1'): + +def trace_through_cv( + raysets, pth, surflist, nrays, wave, global_coords, global_coord_reference="1" +): """trace raysets through a sequential code v optical system Parameters @@ -249,13 +279,14 @@ def trace_through_cv(raysets,pth,surflist,nrays,wave,global_coords,global_coord_ position, direction, surface normal, and OPD data """ - # Code V Imports for com interface import sys - from pythoncom import (CoInitializeEx, CoUninitialize, COINIT_MULTITHREADED, com_error ) + from pythoncom import CoInitializeEx, CoUninitialize, COINIT_MULTITHREADED, com_error from pythoncom import VT_VARIANT, VT_BYREF, VT_ARRAY, VT_R8 from win32com.client import DispatchWithEvents, Dispatch, gencache, VARIANT - from win32com.client import constants as c # To use enumerated constants from the COM object typelib + from win32com.client import ( + constants as c, + ) # To use enumerated constants from the COM object typelib from win32api import FormatMessage sys.coinit_flags = COINIT_MULTITHREADED @@ -264,219 +295,225 @@ def trace_through_cv(raysets,pth,surflist,nrays,wave,global_coords,global_coord_ # Class to instantiate CV interface class ICVApplicationEvents: def OnLicenseError(self, error): - # This event handler is called when a licensing error is + # This event handler is called when a licensing error is # detected in the CODE V application. - print ("License error: %s " % error) + print("License error: %s " % error) def OnCodeVError(self, error): # This event handler is called when a CODE V error message is issued - print ("CODE V error: %s " % error) + print("CODE V error: %s " % error) def OnCodeVWarning(self, warning): # This event handler is called when a CODE V warning message is issued - print ("CODE V warning: %s " % warning) + print("CODE V warning: %s " % warning) def OnPlotReady(self, filename, plotwindow): # This event handler is called when a plot file, refered to by filename, # is ready to be displayed. # The event handler is responsible for saving/copying the # plot data out of the file specified by filename - print ("CODE V Plot: %s in plot window %d" % (filename ,plotwindow) ) + print("CODE V Plot: %s in plot window %d" % (filename, plotwindow)) zoompos = 1 wavelen = 1 fieldno = 0 refsurf = 0 ray_ind = 0 - + cv = DispatchWithEvents("CodeV.Application", ICVApplicationEvents) cv.StartCodeV() # Load the file - if pth[-3:] == 'len': - print(f'res {pth}') - cv.Command(f'res {pth}') - elif pth[-3:] == 'seq': - print(f'in {pth}') - cv.Command(f'in '+pth) + if pth[-3:] == "len": + print(f"res {pth}") + cv.Command(f"res {pth}") + elif pth[-3:] == "seq": + print(f"in {pth}") + cv.Command(f"in " + pth) # configure the file - cv.Command('cd '+dir) - cv.Command('dim m') + cv.Command("cd " + dir) + cv.Command("dim m") # clear any existing buffer data - cv.Command('buf n') # turn off buffer saving if it exists - cv.Command('buf del b0') # clear the buffer + cv.Command("buf n") # turn off buffer saving if it exists + cv.Command("buf del b0") # clear the buffer # Set wavelength to 1um so OPD are in units of um TODO: This breaks refractive element tracing - cv.Command('wl w1 1000') - + cv.Command("wl w1 1000") + # Set up global coordinate reference if global_coords: - cv.Command(f'glo s{global_coord_reference} 0 0 0') + cv.Command(f"glo s{global_coord_reference} 0 0 0") # cv.Command(f'pol y') - print(f'global coordinate reference set to surface {global_coord_reference}') + print(f"global coordinate reference set to surface {global_coord_reference}") else: - cv.Command('glo n') + cv.Command("glo n") # Configure ray output format to get everything we need for PRT/GBD - cv.Command('rof x y z l m n srl srm srn aoi aor') + cv.Command("rof x y z l m n srl srm srn aoi aor") # How many surfaces do we have? - numsurf = int(cv.EvaluateExpression('(NUM S)')) - assert numsurf >= 3, f'number of surfaces = {numsurf}' + numsurf = int(cv.EvaluateExpression("(NUM S)")) + assert numsurf >= 3, f"number of surfaces = {numsurf}" maxrays = raysets[0].shape[-1] - print('maxrays = ',maxrays) + print("maxrays = ", maxrays) # Dimension 0 is ray set, Dimension 1 is surface, dimension 2 is coordinate # Satisfies broadcasting rules! - xData = np.empty([len(raysets),len(surflist),maxrays]) - yData = np.empty([len(raysets),len(surflist),maxrays]) - zData = np.empty([len(raysets),len(surflist),maxrays]) + xData = np.empty([len(raysets), len(surflist), maxrays]) + yData = np.empty([len(raysets), len(surflist), maxrays]) + zData = np.empty([len(raysets), len(surflist), maxrays]) - lData = np.empty([len(raysets),len(surflist),maxrays]) - mData = np.empty([len(raysets),len(surflist),maxrays]) - nData = np.empty([len(raysets),len(surflist),maxrays]) + lData = np.empty([len(raysets), len(surflist), maxrays]) + mData = np.empty([len(raysets), len(surflist), maxrays]) + nData = np.empty([len(raysets), len(surflist), maxrays]) - l2Data = np.empty([len(raysets),len(surflist),maxrays]) - m2Data = np.empty([len(raysets),len(surflist),maxrays]) - n2Data = np.empty([len(raysets),len(surflist),maxrays]) + l2Data = np.empty([len(raysets), len(surflist), maxrays]) + m2Data = np.empty([len(raysets), len(surflist), maxrays]) + n2Data = np.empty([len(raysets), len(surflist), maxrays]) # Necessary for GBD calculations, might help PRT calculations - opd = np.empty([len(raysets),len(surflist),maxrays]) + opd = np.empty([len(raysets), len(surflist), maxrays]) # Open an intermediate raytrace file - for rayset_ind,rayset in enumerate(raysets): + for rayset_ind, rayset in enumerate(raysets): - Hx,Hy = rayset[2,0],rayset[3,0] + Hx, Hy = rayset[2, 0], rayset[3, 0] fac = -1 - for surf_ind,surfdict in enumerate(surflist): + for surf_ind, surfdict in enumerate(surflist): - surf = surfdict['surf'] - file = open(dir+'intermediate_raytrace.seq','w') + surf = surfdict["surf"] + file = open(dir + "intermediate_raytrace.seq", "w") # Begin file construction # x,y,z,l,m,n,l2,m2,n2,aoi,aor - file.write('! Define input variables\n') - file.write(f'num ^input_array(2,{int(nrays**2)}) ^output_array(12,{int(nrays**2)}) ^input_ray(4)\n') - file.write('num ^success \n') + file.write("! Define input variables\n") + file.write( + f"num ^input_array(2,{int(nrays**2)}) ^output_array(12,{int(nrays**2)}) ^input_ray(4)\n" + ) + file.write("num ^success \n") # file.write('! set up global coordinates\n') # file.write(f'glo s{global_coord_reference} 0 0 0\n') # file.write('glo y\n') - file.write(f'^nrays == {nrays**2} \n') - file.write(f'^nrays_across == sqrt(^nrays) \n') - file.write('^x_start == -1\n') - file.write('^x_end == 1\n') - file.write('^y_start == -1\n') - file.write('^y_end == 1\n') - file.write('\n') - file.write('rof x y z l m n srl srm srn aoi aor \n') - - file.write('! compute step size and construct input array\n') - file.write('^step_size == absf(^y_start-^y_end)/(^nrays_across-1)\n') - file.write('^x_cord == ^x_start\n') - file.write('^y_cord == ^y_start\n') - file.write('^next_row == 0\n') - file.write('FOR ^iter 1 ^nrays\n') - file.write(' ^input_array(1,^iter) == ^x_cord\n') - file.write(' ^input_array(2,^iter) == ^y_cord\n') - file.write(' ! update the x_cord\n') - file.write(' ^x_cord == ^x_cord + ^step_size\n') - file.write(' IF modf(^iter,^nrays_across) = 0\n') - file.write(' ^next_row == 1\n') - file.write(' END IF\n') - file.write(' IF ^next_row = 1\n') - file.write(' ! update y_cord\n') - file.write(' ^y_cord == ^y_cord + ^step_size\n') - file.write(' ! reset x_cord\n') - file.write(' ^x_cord == ^x_start\n') - file.write(' ! reset next_row\n') - file.write(' ^next_row == 0\n') - file.write(' END IF\n') - file.write('END FOR\n') - file.write('\n') - - file.write('! Run RAYRSI and pass data to output_array\n') - file.write('FOR ^iter 1 ^nrays\n') - file.write(' ^input_ray(1) == ^input_array(1,^iter)\n') - file.write(' ^input_ray(2) == ^input_array(2,^iter)\n') - file.write(f' ^input_ray(3) == {Hx}\n') - file.write(f' ^input_ray(4) == {Hy}\n') - file.write('\n') - file.write(' ! Execute RAYRSI\n') - file.write(' ^success == RAYRSI(0,0,0,0,^input_ray)\n') + file.write(f"^nrays == {nrays**2} \n") + file.write(f"^nrays_across == sqrt(^nrays) \n") + file.write("^x_start == -1\n") + file.write("^x_end == 1\n") + file.write("^y_start == -1\n") + file.write("^y_end == 1\n") + file.write("\n") + file.write("rof x y z l m n srl srm srn aoi aor \n") + + file.write("! compute step size and construct input array\n") + file.write("^step_size == absf(^y_start-^y_end)/(^nrays_across-1)\n") + file.write("^x_cord == ^x_start\n") + file.write("^y_cord == ^y_start\n") + file.write("^next_row == 0\n") + file.write("FOR ^iter 1 ^nrays\n") + file.write(" ^input_array(1,^iter) == ^x_cord\n") + file.write(" ^input_array(2,^iter) == ^y_cord\n") + file.write(" ! update the x_cord\n") + file.write(" ^x_cord == ^x_cord + ^step_size\n") + file.write(" IF modf(^iter,^nrays_across) = 0\n") + file.write(" ^next_row == 1\n") + file.write(" END IF\n") + file.write(" IF ^next_row = 1\n") + file.write(" ! update y_cord\n") + file.write(" ^y_cord == ^y_cord + ^step_size\n") + file.write(" ! reset x_cord\n") + file.write(" ^x_cord == ^x_start\n") + file.write(" ! reset next_row\n") + file.write(" ^next_row == 0\n") + file.write(" END IF\n") + file.write("END FOR\n") + file.write("\n") + + file.write("! Run RAYRSI and pass data to output_array\n") + file.write("FOR ^iter 1 ^nrays\n") + file.write(" ^input_ray(1) == ^input_array(1,^iter)\n") + file.write(" ^input_ray(2) == ^input_array(2,^iter)\n") + file.write(f" ^input_ray(3) == {Hx}\n") + file.write(f" ^input_ray(4) == {Hy}\n") + file.write("\n") + file.write(" ! Execute RAYRSI\n") + file.write(" ^success == RAYRSI(0,0,0,0,^input_ray)\n") # file.write(' rsi ^input_array(1), ^input_array(2), ^input_array(3), ^input_array(4),\n') - file.write('\n') - file.write(' ! Read ray data from lens database into output_array\n') - file.write(f' ^output_array(1,^iter) == (X S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(2,^iter) == (Y S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(3,^iter) == (Z S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(4,^iter) == (L S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(5,^iter) == (M S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(6,^iter) == (N S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(7,^iter) == (SRL S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(8,^iter) == (SRM S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(9,^iter) == (SRN S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(10,^iter) == (AOI S{surf} G{global_coord_reference})\n') - file.write(f' ^output_array(11,^iter) == (AOR S{surf} G{global_coord_reference})\n') + file.write("\n") + file.write(" ! Read ray data from lens database into output_array\n") + file.write(f" ^output_array(1,^iter) == (X S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(2,^iter) == (Y S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(3,^iter) == (Z S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(4,^iter) == (L S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(5,^iter) == (M S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(6,^iter) == (N S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(7,^iter) == (SRL S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(8,^iter) == (SRM S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(9,^iter) == (SRN S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(10,^iter) == (AOI S{surf} G{global_coord_reference})\n") + file.write(f" ^output_array(11,^iter) == (AOR S{surf} G{global_coord_reference})\n") # TODO: This defaults to the total OPD, not per-surface - file.write(f' ^output_array(12,^iter) == (OPD)\n') - file.write('END FOR\n') - file.write('\n') - file.write('! Write output_array to text file\n') - file.write('BUF Y\n') - file.write('BUF DEL B1\n') - file.write(f'BUF LEN {maxrays} \n') - file.write('^result == ARRAY_TO_BUFFER(^output_array,1,1)\n') + file.write(f" ^output_array(12,^iter) == (OPD)\n") + file.write("END FOR\n") + file.write("\n") + file.write("! Write output_array to text file\n") + file.write("BUF Y\n") + file.write("BUF DEL B1\n") + file.write(f"BUF LEN {maxrays} \n") + file.write("^result == ARRAY_TO_BUFFER(^output_array,1,1)\n") file.write('BUF EXP B1 SEP "intermediate_output.txt"\n') file.close() # Now execute the macro - cv.Command(f'in intermediate_raytrace.seq') + cv.Command(f"in intermediate_raytrace.seq") # And read the raydata - raydata = np.genfromtxt(dir+"intermediate_output.txt",delimiter=' ') + raydata = np.genfromtxt(dir + "intermediate_output.txt", delimiter=" ") # raydata = np.genfromtxt(dir+"intermediate_output.txt",delimiter=' ',usecols=np.arange(0,int(nrays**2))) # Load into numpy arrays - xData[rayset_ind,surf_ind] = raydata[:,0] - yData[rayset_ind,surf_ind] = raydata[:,1] - zData[rayset_ind,surf_ind] = raydata[:,2] + xData[rayset_ind, surf_ind] = raydata[:, 0] + yData[rayset_ind, surf_ind] = raydata[:, 1] + zData[rayset_ind, surf_ind] = raydata[:, 2] - lData[rayset_ind,surf_ind] = raydata[:,3]*fac - mData[rayset_ind,surf_ind] = raydata[:,4]*fac - nData[rayset_ind,surf_ind] = raydata[:,5]*fac + lData[rayset_ind, surf_ind] = raydata[:, 3] * fac + mData[rayset_ind, surf_ind] = raydata[:, 4] * fac + nData[rayset_ind, surf_ind] = raydata[:, 5] * fac - l2Data[rayset_ind,surf_ind] = raydata[:,6] - m2Data[rayset_ind,surf_ind] = raydata[:,7] - n2Data[rayset_ind,surf_ind] = raydata[:,8] + l2Data[rayset_ind, surf_ind] = raydata[:, 6] + m2Data[rayset_ind, surf_ind] = raydata[:, 7] + n2Data[rayset_ind, surf_ind] = raydata[:, 8] - opd[rayset_ind,surf_ind] = raydata[:,11] + opd[rayset_ind, surf_ind] = raydata[:, 11] fac *= -1 # delete the files made - os.remove(dir+'intermediate_raytrace.seq') - os.remove(dir+'intermediate_output.txt') + os.remove(dir + "intermediate_raytrace.seq") + os.remove(dir + "intermediate_output.txt") - positions = [xData*1e-3,yData*1e-3,zData*1e-3] # correct for default to mm - norm = np.sqrt(lData**2 + mData**2 + nData**2) + positions = [xData * 1e-3, yData * 1e-3, zData * 1e-3] # correct for default to mm + norm = np.sqrt(lData ** 2 + mData ** 2 + nData ** 2) lData /= norm mData /= norm nData /= norm - norm = np.sqrt(l2Data**2 + m2Data**2 + n2Data**2) + norm = np.sqrt(l2Data ** 2 + m2Data ** 2 + n2Data ** 2) l2Data /= norm m2Data /= norm n2Data /= norm - directions = [lData,mData,nData] - normals = [l2Data,m2Data,n2Data] + directions = [lData, mData, nData] + normals = [l2Data, m2Data, n2Data] # Just a bit of celebration - print('{nrays} Raysets traced through {nsurf} surfaces'.format(nrays=rayset_ind+1,nsurf=surf_ind+1)) + print( + "{nrays} Raysets traced through {nsurf} surfaces".format( + nrays=rayset_ind + 1, nsurf=surf_ind + 1 + ) + ) # Close up cv.StopCodeV() @@ -484,20 +521,22 @@ def OnPlotReady(self, filename, plotwindow): # Delete the intermediate files - # And finally return everything, OPD needs to be converted to meters - return positions,directions,normals,opd*1e-6 + return positions, directions, normals, opd * 1e-6 + -def TraceThroughCV(raysets,pth,surflist,nrays,wave,global_coords,global_coord_reference='1'): +def TraceThroughCV(raysets, pth, surflist, nrays, wave, global_coords, global_coord_reference="1"): - print('This function is depreciated, please use trace_through_cv') + print("This function is depreciated, please use trace_through_cv") # Code V Imports for com interface import sys - from pythoncom import (CoInitializeEx, CoUninitialize, COINIT_MULTITHREADED, com_error ) + from pythoncom import CoInitializeEx, CoUninitialize, COINIT_MULTITHREADED, com_error from pythoncom import VT_VARIANT, VT_BYREF, VT_ARRAY, VT_R8 from win32com.client import DispatchWithEvents, Dispatch, gencache, VARIANT - from win32com.client import constants as c # To use enumerated constants from the COM object typelib + from win32com.client import ( + constants as c, + ) # To use enumerated constants from the COM object typelib from win32api import FormatMessage sys.coinit_flags = COINIT_MULTITHREADED @@ -506,85 +545,84 @@ def TraceThroughCV(raysets,pth,surflist,nrays,wave,global_coords,global_coord_re # Class to instantiate CV interface class ICVApplicationEvents: def OnLicenseError(self, error): - # This event handler is called when a licensing error is + # This event handler is called when a licensing error is # detected in the CODE V application. - print ("License error: %s " % error) + print("License error: %s " % error) def OnCodeVError(self, error): # This event handler is called when a CODE V error message is issued - print ("CODE V error: %s " % error) + print("CODE V error: %s " % error) def OnCodeVWarning(self, warning): # This event handler is called when a CODE V warning message is issued - print ("CODE V warning: %s " % warning) + print("CODE V warning: %s " % warning) def OnPlotReady(self, filename, plotwindow): # This event handler is called when a plot file, refered to by filename, # is ready to be displayed. # The event handler is responsible for saving/copying the # plot data out of the file specified by filename - print ("CODE V Plot: %s in plot window %d" % (filename ,plotwindow) ) + print("CODE V Plot: %s in plot window %d" % (filename, plotwindow)) zoompos = 1 wavelen = 1 fieldno = 0 refsurf = 0 ray_ind = 0 - + cv = DispatchWithEvents("CodeV.Application", ICVApplicationEvents) cv.StartCodeV() # Load the file - print(f'res {pth}') - cv.Command(f'res {pth}') + print(f"res {pth}") + cv.Command(f"res {pth}") # clear any existing buffer data - cv.Command('buf n') # turn off buffer saving if it exists - cv.Command('buf del b0') # clear the buffer + cv.Command("buf n") # turn off buffer saving if it exists + cv.Command("buf del b0") # clear the buffer # Set wavelength to 1um so OPD are in units of um - cv.Command('wl w1 1000') - + cv.Command("wl w1 1000") + # Set up global coordinate reference if global_coords: - cv.Command(f'glo s{global_coord_reference} 0 0 0') + cv.Command(f"glo s{global_coord_reference} 0 0 0") # cv.Command(f'pol y') - print(f'global coordinate reference set to surface {global_coord_reference}') + print(f"global coordinate reference set to surface {global_coord_reference}") offset_rows = 7 - offset_columns = 0 # for apertured systems + offset_columns = 0 # for apertured systems else: - cv.Command('glo n') - offset_rows = -1+7 + cv.Command("glo n") + offset_rows = -1 + 7 # Configure ray output format to get everything we need for PRT/GBD - cv.Command('rof x y z l m n srl srm srn aoi aor') + cv.Command("rof x y z l m n srl srm srn aoi aor") # How many surfaces do we have? - numsurf = int(cv.EvaluateExpression('(NUM S)')) + numsurf = int(cv.EvaluateExpression("(NUM S)")) if numsurf < 3: - raise Exception('File was not loaded correctly') + raise Exception("File was not loaded correctly") maxrays = raysets[0].shape[-1] # Dimension 0 is ray set, Dimension 1 is surface, dimension 2 is coordinate # Satisfies broadcasting rules! - xData = np.empty([len(raysets),len(surflist),maxrays]) - yData = np.empty([len(raysets),len(surflist),maxrays]) - zData = np.empty([len(raysets),len(surflist),maxrays]) + xData = np.empty([len(raysets), len(surflist), maxrays]) + yData = np.empty([len(raysets), len(surflist), maxrays]) + zData = np.empty([len(raysets), len(surflist), maxrays]) - lData = np.empty([len(raysets),len(surflist),maxrays]) - mData = np.empty([len(raysets),len(surflist),maxrays]) - nData = np.empty([len(raysets),len(surflist),maxrays]) + lData = np.empty([len(raysets), len(surflist), maxrays]) + mData = np.empty([len(raysets), len(surflist), maxrays]) + nData = np.empty([len(raysets), len(surflist), maxrays]) - l2Data = np.empty([len(raysets),len(surflist),maxrays]) - m2Data = np.empty([len(raysets),len(surflist),maxrays]) - n2Data = np.empty([len(raysets),len(surflist),maxrays]) + l2Data = np.empty([len(raysets), len(surflist), maxrays]) + m2Data = np.empty([len(raysets), len(surflist), maxrays]) + n2Data = np.empty([len(raysets), len(surflist), maxrays]) # Necessary for GBD calculations, might help PRT calculations - opd = np.empty([len(raysets),len(surflist),maxrays]) - + opd = np.empty([len(raysets), len(surflist), maxrays]) - for rayset_ind,rayset in enumerate(raysets): + for rayset_ind, rayset in enumerate(raysets): # Get the normalized coordinates Px = rayset[0] @@ -592,56 +630,55 @@ def OnPlotReady(self, filename, plotwindow): Hx = rayset[2] Hy = rayset[3] - for ray_ind,(px,py,hx,hy) in enumerate(zip(Px,Py,Hx,Hy)): + for ray_ind, (px, py, hx, hy) in enumerate(zip(Px, Py, Hx, Hy)): # TODO : Make compatible with different wavelengths and zooms # out = cv.RAYRSI(zoompos,wavelen,fieldno,refsurf,[px,py,hx,hy]) - - cv.Command('buf y') - cv.Command(f'RSI {px} {py} {hx} {hy}') - cv.Command('buf n') + + cv.Command("buf y") + cv.Command(f"RSI {px} {py} {hx} {hy}") + cv.Command("buf n") # if out != 0: # print('raytrace failure') - - # TODO figure out why negated + + # TODO figure out why negated fac = -1 - for surf_ind,surfdict in enumerate(surflist): + for surf_ind, surfdict in enumerate(surflist): # print(rayset_ind) # print('surf ind = ',surf_ind) # print(ray_ind) # print('-------') - surf = surfdict['surf'] # surface in LDE + surf = surfdict["surf"] # surface in LDE # print(surf) # Do this the buffer scraping way - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{2+offset_columns}') - xData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression("(BUF.NUM)") - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{3+offset_columns}') - yData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{4+offset_columns}') - zData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - - - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{5+offset_columns}') - lData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{6+offset_columns}') - mData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{7+offset_columns}') - nData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{2+offset_columns}") + xData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{3+offset_columns}") + yData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{4+offset_columns}") + zData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{5+offset_columns}") + lData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{6+offset_columns}") + mData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{7+offset_columns}") + nData[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + # apply the factor - lData[rayset_ind,surf_ind,ray_ind] *= fac - mData[rayset_ind,surf_ind,ray_ind] *= fac - nData[rayset_ind,surf_ind,ray_ind] *= fac - - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{8+offset_columns}') - l2Data[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{9+offset_columns}') - m2Data[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') - cv.Command(f'BUF MOV B0 i{surf+offset_rows} j{10+offset_columns}') - n2Data[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') + lData[rayset_ind, surf_ind, ray_ind] *= fac + mData[rayset_ind, surf_ind, ray_ind] *= fac + nData[rayset_ind, surf_ind, ray_ind] *= fac + + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{8+offset_columns}") + l2Data[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{9+offset_columns}") + m2Data[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") + cv.Command(f"BUF MOV B0 i{surf+offset_rows} j{10+offset_columns}") + n2Data[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") # xData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression(f'(x s{surf})') # yData[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression(f'(y s{surf})') @@ -656,33 +693,40 @@ def OnPlotReady(self, filename, plotwindow): # n2Data[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression(f'(srn s{surf})') # print(f'{surf}') # TODO: Check that this is returning the correct OPD - cv.Command(f'BUF MOV B0 i{1+numsurf+offset_rows} j{2+offset_columns}') - opd[rayset_ind,surf_ind,ray_ind] = cv.EvaluateExpression('(BUF.NUM)') + cv.Command(f"BUF MOV B0 i{1+numsurf+offset_rows} j{2+offset_columns}") + opd[rayset_ind, surf_ind, ray_ind] = cv.EvaluateExpression("(BUF.NUM)") fac *= -1 - cv.Command('buf del b0') + cv.Command("buf del b0") # This isn't necessary but makes the code more readable # CODE V will default to mm, so we need to scale back to meters - positions = [xData*1e-3,yData*1e-3,zData*1e-3] - norm = np.sqrt(lData**2 + mData**2 + nData**2) + positions = [xData * 1e-3, yData * 1e-3, zData * 1e-3] + norm = np.sqrt(lData ** 2 + mData ** 2 + nData ** 2) lData /= norm mData /= norm nData /= norm - directions = [lData,mData,nData] - normals = [l2Data,m2Data,n2Data] + directions = [lData, mData, nData] + normals = [l2Data, m2Data, n2Data] # Just a bit of celebration - print('{nrays} Raysets traced through {nsurf} surfaces'.format(nrays=rayset_ind+1,nsurf=surf_ind+1)) + print( + "{nrays} Raysets traced through {nsurf} surfaces".format( + nrays=rayset_ind + 1, nsurf=surf_ind + 1 + ) + ) # Close up cv.StopCodeV() del cv - + # And finally return everything, OPD needs to be converted to meters - return positions,directions,normals,opd*1e-6 + return positions, directions, normals, opd * 1e-6 + -def convert_ray_data_to_prt_data(LData,MData,NData,L2Data,M2Data,N2Data,surflist,ambient_index=1): +def convert_ray_data_to_prt_data( + LData, MData, NData, L2Data, M2Data, N2Data, surflist, ambient_index=1 +): """Function that computes the PRT-relevant data from ray and material data Mathematics principally from Polarized Light and Optical Systems by Chipman, Lam, Young 2018 @@ -739,7 +783,7 @@ def convert_ray_data_to_prt_data(LData,MData,NData,L2Data,M2Data,N2Data,surflist n1 = ambient_index # Loop over surfaces - for surf_ind,surfdict in enumerate(surflist): + for surf_ind, surfdict in enumerate(surflist): # Do some digesting # directions @@ -753,19 +797,17 @@ def convert_ray_data_to_prt_data(LData,MData,NData,L2Data,M2Data,N2Data,surflist n2Data = N2Data[surf_ind] # grab the index - if type(surfdict['coating']) == list: + if type(surfdict["coating"]) == list: # compute coeff from last film, first location - n2 = surfdict['coating'][-1] + n2 = surfdict["coating"][-1] - else: + else: # assume scalar - n2 = surfdict['coating'] - - + n2 = surfdict["coating"] # Maintain right handed coords to stay with Chipman sign convention - norm = -np.array([l2Data,m2Data,n2Data]) + norm = -np.array([l2Data, m2Data, n2Data]) # # Compute number of rays # total_rays_in_both_axes = int(lData.shape[0]) @@ -774,41 +816,37 @@ def convert_ray_data_to_prt_data(LData,MData,NData,L2Data,M2Data,N2Data,surflist # calculates angle of exitance from direction cosine # the LMN direction cosines are for AFTER refraction # need to calculate via Snell's Law the angle of incidence - numerator = (lData*l2Data + mData*m2Data + nData*n2Data) - denominator = ((lData**2 + mData**2 + nData**2)**0.5)*(l2Data**2 + m2Data**2 + n2Data**2)**0.5 - aoe_data = np.arccos(-numerator/denominator) # now in radians + numerator = lData * l2Data + mData * m2Data + nData * n2Data + denominator = ((lData ** 2 + mData ** 2 + nData ** 2) ** 0.5) * ( + l2Data ** 2 + m2Data ** 2 + n2Data ** 2 + ) ** 0.5 + aoe_data = np.arccos(-numerator / denominator) # now in radians # aoe = aoe_data - (aoe_data[0:total_rays_in_both_axes] > np.pi/2) * np.pi aoe = aoe_data # Compute kin with Snell's Law: https://en.wikipedia.org/wiki/Snell%27s_law#Vector_form - kout.append(np.array([lData,mData,nData])/np.sqrt(lData**2 + mData**2 + nData**2)) + kout.append(np.array([lData, mData, nData]) / np.sqrt(lData ** 2 + mData ** 2 + nData ** 2)) - if surfdict['mode'] == 'transmit': + if surfdict["mode"] == "transmit": # Snell's Law - aoi.append((np.arcsin(n2/n1 * np.sin(aoe)))) - kin.append(np.cos(np.arcsin(n2*np.sin(np.arccos(kout[surf_ind]))/n1))) + aoi.append((np.arcsin(n2 / n1 * np.sin(aoe)))) + kin.append(np.cos(np.arcsin(n2 * np.sin(np.arccos(kout[surf_ind])) / n1))) - elif surfdict['mode'] == 'reflect': + elif surfdict["mode"] == "reflect": # Snell's Law aoi.append(-aoe) - kin_norm = kout[surf_ind] - 2*np.cos(aoi[surf_ind])*norm - kin_norm /= np.sqrt(kin_norm[0]**2 + kin_norm[1]**2 + kin_norm[2]**2) + kin_norm = kout[surf_ind] - 2 * np.cos(aoi[surf_ind]) * norm + kin_norm /= np.sqrt(kin_norm[0] ** 2 + kin_norm[1] ** 2 + kin_norm[2] ** 2) kin.append(kin_norm) else: - print('Interaction mode not recognized') - - # saves normal in zemax sign convention - normal.append(np.array([l2Data,m2Data,n2Data])/np.sqrt(l2Data**2 + m2Data**2 + n2Data**2)) - - return aoi,kin,kout,normal - - - - - - + print("Interaction mode not recognized") + # saves normal in zemax sign convention + normal.append( + np.array([l2Data, m2Data, n2Data]) / np.sqrt(l2Data ** 2 + m2Data ** 2 + n2Data ** 2) + ) + return aoi, kin, kout, normal diff --git a/poke/thinfilms.py b/poke/thinfilms.py index 0592fd1..ccf53d7 100644 --- a/poke/thinfilms.py +++ b/poke/thinfilms.py @@ -1,14 +1,17 @@ from poke.poke_math import np # Inputs are list of index, distance, and the wavelength -VACUUM_PERMITTIVITY = 8.8541878128e-12 # Farad / M -VACUUM_PERMEABILITY = 1.25663706212e-6 # Henry / M -FREESPACE_IMPEDANCE = 376.730313668 # ohms -FREESPACE_IMPEDANCE_INV = 1/FREESPACE_IMPEDANCE # Chipman includes this, Macleod ignores this -ONE_COMPLEX = 1 + 0*1j -ZERO_COMPLEX = 0 + 0*1j - -def compute_thin_films_broadcasted(stack, aoi, wavelength, ambient_index=1, substrate_index=1.5, polarization='s'): +VACUUM_PERMITTIVITY = 8.8541878128e-12 # Farad / M +VACUUM_PERMEABILITY = 1.25663706212e-6 # Henry / M +FREESPACE_IMPEDANCE = 376.730313668 # ohms +FREESPACE_IMPEDANCE_INV = 1 / FREESPACE_IMPEDANCE # Chipman includes this, Macleod ignores this +ONE_COMPLEX = 1 + 0 * 1j +ZERO_COMPLEX = 0 + 0 * 1j + + +def compute_thin_films_broadcasted( + stack, aoi, wavelength, ambient_index=1, substrate_index=1.5, polarization="s" +): """compute fresnel coefficients for a multilayer stack using the BYU Optics Book method Parameters @@ -33,84 +36,81 @@ def compute_thin_films_broadcasted(stack, aoi, wavelength, ambient_index=1, subs """ # Do some digesting - stack = stack[:-1] # ignore the last element, which contains the substrate index + stack = stack[:-1] # ignore the last element, which contains the substrate index # Consider the incident media system_matrix = np.array([[1, 0], [0, 1]], dtype=np.complex128) if len(aoi.shape) > 0: - system_matrix = np.broadcast_to(system_matrix,[*aoi.shape,*system_matrix.shape]) + system_matrix = np.broadcast_to(system_matrix, [*aoi.shape, *system_matrix.shape]) # Consider the terminating media - aor = np.arcsin(ambient_index/substrate_index*np.sin(aoi)) - ncosAOR = substrate_index*np.cos(aor) + aor = np.arcsin(ambient_index / substrate_index * np.sin(aoi)) + ncosAOR = substrate_index * np.cos(aor) cosAOI = np.cos(aoi) sinAOI = np.sin(aoi) ncosAOI = ambient_index * cosAOI - - n0 = np.full_like(aoi,ambient_index,dtype=np.complex128) - nM = np.full_like(aoi,substrate_index,dtype=np.complex128) - zeros = np.full_like(aoi,0.,dtype=np.complex128) - ones = np.full_like(aoi,1.,dtype=np.complex128) + + n0 = np.full_like(aoi, ambient_index, dtype=np.complex128) + nM = np.full_like(aoi, substrate_index, dtype=np.complex128) + zeros = np.full_like(aoi, 0.0, dtype=np.complex128) + ones = np.full_like(aoi, 1.0, dtype=np.complex128) for layer in stack: ni = layer[0] - di = layer[1] # has some dimension + di = layer[1] # has some dimension - angle_in_film = np.arcsin(ambient_index/ni*sinAOI) + angle_in_film = np.arcsin(ambient_index / ni * sinAOI) Beta = 2 * np.pi * ni * di * np.cos(angle_in_film) / wavelength - cosB = np.full_like(aoi,np.cos(Beta),dtype=np.complex128) - sinB = np.full_like(aoi,np.sin(Beta),dtype=np.complex128) - cosT = np.full_like(aoi,np.cos(angle_in_film),dtype=np.complex128) + cosB = np.full_like(aoi, np.cos(Beta), dtype=np.complex128) + sinB = np.full_like(aoi, np.sin(Beta), dtype=np.complex128) + cosT = np.full_like(aoi, np.cos(angle_in_film), dtype=np.complex128) - if polarization == 'p': - newfilm = np.array([[cosB, -1j*sinB*cosT/ni], - [-1j*ni*sinB/cosT, cosB]]) + if polarization == "p": + newfilm = np.array([[cosB, -1j * sinB * cosT / ni], [-1j * ni * sinB / cosT, cosB]]) - elif polarization == 's': - newfilm = np.array([[cosB, -1j*sinB/(cosT* ni)], - [-1j*ni*sinB*cosT, cosB]]) - if newfilm.ndim > 2: - for i in range(newfilm.ndim-2): + elif polarization == "s": + newfilm = np.array([[cosB, -1j * sinB / (cosT * ni)], [-1j * ni * sinB * cosT, cosB]]) + if newfilm.ndim > 2: + for i in range(newfilm.ndim - 2): newfilm = np.moveaxis(newfilm, -1, 0) - + system_matrix = system_matrix @ newfilm # Final matrix - coeff = 1/(2*ncosAOI) + coeff = 1 / (2 * ncosAOI) if system_matrix.ndim > 2: - coeff = coeff[...,np.newaxis,np.newaxis] - - if polarization == 'p': - front_matrix = np.array([[n0, cosAOI], - [n0, -cosAOI]]) - back_matrix = np.array([[np.cos(aor), zeros], - [nM, zeros]]) - - elif polarization == 's': - front_matrix = np.array([[ncosAOI, ones], - [ncosAOI, -1*ones]]) - back_matrix = np.array([[ones, zeros], - [ncosAOR, zeros]]) - + coeff = coeff[..., np.newaxis, np.newaxis] + + if polarization == "p": + front_matrix = np.array([[n0, cosAOI], [n0, -cosAOI]]) + back_matrix = np.array([[np.cos(aor), zeros], [nM, zeros]]) + + elif polarization == "s": + front_matrix = np.array([[ncosAOI, ones], [ncosAOI, -1 * ones]]) + back_matrix = np.array([[ones, zeros], [ncosAOR, zeros]]) + if front_matrix.ndim > 2: - for i in range(front_matrix.ndim-2): - front_matrix = np.moveaxis(front_matrix,-1,0) - + for i in range(front_matrix.ndim - 2): + front_matrix = np.moveaxis(front_matrix, -1, 0) + if back_matrix.ndim > 2: - for i in range(back_matrix.ndim-2): - back_matrix = np.moveaxis(back_matrix,-1,0) + for i in range(back_matrix.ndim - 2): + back_matrix = np.moveaxis(back_matrix, -1, 0) characteristic_matrix = coeff * (front_matrix @ system_matrix @ back_matrix) - ttot = 1/characteristic_matrix[..., 0, 0] - rtot = characteristic_matrix[..., 1, 0]/characteristic_matrix[..., 0, 0] + ttot = 1 / characteristic_matrix[..., 0, 0] + rtot = characteristic_matrix[..., 1, 0] / characteristic_matrix[..., 0, 0] return rtot, ttot -def compute_thin_films_macleod(stack, aoi, wavelength, ambient_index=1, substrate_index=1.5, polarization='s'): + +def compute_thin_films_macleod( + stack, aoi, wavelength, ambient_index=1, substrate_index=1.5, polarization="s" +): """compute fresnel coefficients for a multilayer stack using the Macleod 1969 method Parameters @@ -135,65 +135,64 @@ def compute_thin_films_macleod(stack, aoi, wavelength, ambient_index=1, substrat """ # Do some digesting - stack = stack[:-1] # ignore the last element, which contains the substrate index + stack = stack[:-1] # ignore the last element, which contains the substrate index # Consider the incident media system_matrix = np.array([[1, 0], [0, 1]], dtype=np.complex128) if len(aoi.shape) > 0: - system_matrix = np.broadcast_to(system_matrix,[*aoi.shape,*system_matrix.shape]) - - aor = np.arcsin(ambient_index/substrate_index*np.sin(aoi)) + system_matrix = np.broadcast_to(system_matrix, [*aoi.shape, *system_matrix.shape]) + + aor = np.arcsin(ambient_index / substrate_index * np.sin(aoi)) cosAOR = np.cos(aor) sinAOI = np.sin(aoi) cosAOI = np.cos(aoi) ones = np.ones_like(aoi) # compute the substrate admittance - if polarization == 's': + if polarization == "s": n_substrate = substrate_index * cosAOR eta_medium = ambient_index * cosAOI - elif polarization == 'p': + elif polarization == "p": n_substrate = substrate_index / cosAOR eta_medium = ambient_index / cosAOI else: - print('polarization not recognized, defaulting to s') + print("polarization not recognized, defaulting to s") n_substrate = substrate_index * cosAOR for layer in stack: ni = layer[0] - di = layer[1] # has some dimension + di = layer[1] # has some dimension - angle_in_film = np.arcsin(ambient_index/ni*sinAOI) + angle_in_film = np.arcsin(ambient_index / ni * sinAOI) # phase thickness Beta = 2 * np.pi * ni * di * np.cos(angle_in_film) / wavelength cosB = np.cos(Beta) - isinB = 1j*np.sin(Beta) + isinB = 1j * np.sin(Beta) # film admittance - if polarization == 'p': - eta_film = ni/np.cos(angle_in_film) + if polarization == "p": + eta_film = ni / np.cos(angle_in_film) else: - eta_film = ni*np.cos(angle_in_film) + eta_film = ni * np.cos(angle_in_film) # assemble the characteristic matrix - newfilm = np.array([[cosB,isinB/eta_film], - [isinB*eta_film,cosB]]) - - if newfilm.ndim > 2: - for i in range(newfilm.ndim-2): + newfilm = np.array([[cosB, isinB / eta_film], [isinB * eta_film, cosB]]) + + if newfilm.ndim > 2: + for i in range(newfilm.ndim - 2): newfilm = np.moveaxis(newfilm, -1, 0) - + # apply the matrix system_matrix = system_matrix @ newfilm # layer finished - substrate_vec = np.array([ones,n_substrate]) - substrate_vec = np.swapaxes(substrate_vec,-1,0) - substrate_vec = substrate_vec[...,np.newaxis] + substrate_vec = np.array([ones, n_substrate]) + substrate_vec = np.swapaxes(substrate_vec, -1, 0) + substrate_vec = substrate_vec[..., np.newaxis] BC = system_matrix @ substrate_vec - Y = BC[...,1,0]/BC[...,0,0] + Y = BC[..., 1, 0] / BC[..., 0, 0] - rtot = (eta_medium - Y)/(eta_medium + Y) - return rtot \ No newline at end of file + rtot = (eta_medium - Y) / (eta_medium + Y) + return rtot diff --git a/poke/writing.py b/poke/writing.py index 781673d..bc23e77 100644 --- a/poke/writing.py +++ b/poke/writing.py @@ -6,6 +6,7 @@ m.patch() + def serialize(T): """serializes a class using msgpack written by Brandon Dube, docstring by Jaren Ashcraft @@ -23,18 +24,21 @@ class to convert to hex code. Used for rayfronts glb = globals() Tname = T.__class__.__name__ # assert Tname in glb, 'class must exist in globals in order to be re-hydrateable, with the same constraint' - + # now we make our storage format. It will be: # 1) a header with the class name # 2) the content of vars(T) serdat = (Tname, vars(T)) return msgpack.packb(serdat) + class MsgpackTrickerEmpty: """dummy class to trick msgpack """ + pass + def deserialize(buf): """deserializes a class using msgpack written by Brandon Dube, docstring by Jaren Ashcraft @@ -53,11 +57,12 @@ def deserialize(buf): Tname, varzzz = msgpack.unpackb(buf, use_list=True) for k, v in varzzz.items(): setattr(e, k, v) - + e.__class__ = globals()[Tname] return e -def write_rayfront_to_serial(rayfront,filename): + +def write_rayfront_to_serial(rayfront, filename): """writes rayfront to serial file using msgpack Parameters @@ -70,9 +75,10 @@ def write_rayfront_to_serial(rayfront,filename): serdata = serialize(rayfront) - with open(filename+'.msgpack','wb') as outfile: + with open(filename + ".msgpack", "wb") as outfile: outfile.write(serdata) + def read_serial_to_rayfront(filename): """reads serial data containing Rayfront into a Rayfront object @@ -87,9 +93,9 @@ def read_serial_to_rayfront(filename): the saved poke.Rayfront object """ - with open(filename,'rb') as infile: + with open(filename, "rb") as infile: serdata = infile.read() - + rayfront = deserialize(serdata) return rayfront @@ -119,10 +125,12 @@ def jones_to_fits(rayfront, filename, realimag=True, which=-1, nmodes=11, npix=1 try: from astropy.io import fits except Exception as e: - print(f'Error in importing astropy \n {e}') + print(f"Error in importing astropy \n {e}") # convert the jones pupil data into regularly spaced data - jones,residual_list = regularly_space_jones(rayfront,nmodes,npix,which=which,return_residuals=True) + jones, residual_list = regularly_space_jones( + rayfront, nmodes, npix, which=which, return_residuals=True + ) # The FITS headers that we want # feel free to just use the variable names as the headers @@ -141,24 +149,24 @@ def jones_to_fits(rayfront, filename, realimag=True, which=-1, nmodes=11, npix=1 realpart = np.abs(jones) imagpart = np.angle(jones) - box = np.empty([*jones.shape,2]) - box[...,0] = realpart - box[...,1] = imagpart + box = np.empty([*jones.shape, 2]) + box[..., 0] = realpart + box[..., 1] = imagpart - hdu_primary = fits.PrimaryHDU(box) + hdu_primary = fits.PrimaryHDU(box) # TODO: Add exit pupil calculation # c2 = fits.Column(name='pixelscale', format='K', array= np.array([1])) - c1 = fits.Column(name='wavelength', format='E', array= np.array([wavelength])) - c3 = fits.Column(name='field_of_view_x', format='D', array=np.array([field_of_view_x])) - c4 = fits.Column(name='field_of_view_y', format='D', array=np.array([field_of_view_y])) - c5 = fits.Column(name='residuals_jxx', format='D', array=np.array([residuals_jxx])) - c6 = fits.Column(name='residuals_jxy', format='D', array=np.array([residuals_jxy])) - c7 = fits.Column(name='residuals_jyx', format='D', array=np.array([residuals_jyx])) - c8 = fits.Column(name='residuals_jyy', format='D', array=np.array([residuals_jyy])) + c1 = fits.Column(name="wavelength", format="E", array=np.array([wavelength])) + c3 = fits.Column(name="field_of_view_x", format="D", array=np.array([field_of_view_x])) + c4 = fits.Column(name="field_of_view_y", format="D", array=np.array([field_of_view_y])) + c5 = fits.Column(name="residuals_jxx", format="D", array=np.array([residuals_jxx])) + c6 = fits.Column(name="residuals_jxy", format="D", array=np.array([residuals_jxy])) + c7 = fits.Column(name="residuals_jyx", format="D", array=np.array([residuals_jyx])) + c8 = fits.Column(name="residuals_jyy", format="D", array=np.array([residuals_jyy])) cols = fits.ColDefs([c1, c3, c4, c5, c6, c7, c8]) hdu_table = fits.BinTableHDU.from_columns(cols) hdul = fits.HDUList([hdu_primary, hdu_table]) - - hdul.writeto(filename + '.fits', overwrite=True) \ No newline at end of file + + hdul.writeto(filename + ".fits", overwrite=True)