Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Trying out linting with Black #103

Merged
merged 1 commit into from
Dec 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 93 additions & 69 deletions poke/interfaces.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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




91 changes: 50 additions & 41 deletions poke/materials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
return index_model
Loading