Skip to content

Commit

Permalink
Merge pull request #103 from Jashcraf/linting
Browse files Browse the repository at this point in the history 
Trying out linting with Black
  • Loading branch information
@Jashcraf
Jashcraf authored Dec 26, 2023
2 parents 3bc0390 + 43d6264 commit cc46565
Show file tree
Hide file tree
Showing 9 changed files with 1,170 additions and 1,099 deletions.
162 changes: 93 additions & 69 deletions poke/interfaces.py
View file
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
View file
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

0 comments on commit cc46565

Please sign in to comment.