49 finalize jones pupil to fits file conversion

poke/interfaces.py

@@ -166,7 +166,7 @@ def zernike(rho, phi, J):
             m=m+2
     return values
 
-def regularly_space_jones(rayfront,nmodes,npix,which=-1):
+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
@@ -179,6 +179,8 @@ def regularly_space_jones(rayfront,nmodes,npix,which=-1):
         number of samples along the side of the output array
     which : int, optional
         which jones pupil in the rf.jones_pupil list to use, by default -1
+    return_residuals : bool, optional
+        Whether to return the full np.linalg.lstsq residuals, by default False
 
     Returns
     -------
@@ -211,7 +213,10 @@ def regularly_space_jones(rayfront,nmodes,npix,which=-1):
     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])
 
-    return return_jones
+    if return_residuals:
+        return return_jones, [cxx,cxy,cyx,cyy]
+    else:
+        return return_jones
 
 
 

poke/writing.py

@@ -2,6 +2,7 @@
 import msgpack
 import msgpack_numpy as m
 from poke.poke_core import Rayfront
+from poke.interfaces import regularly_space_jones
 
 m.patch()
 
@@ -94,18 +95,24 @@ def read_serial_to_rayfront(filename):
     return rayfront
 
 
-def jones_to_fits(data,filename,realimag=True):
+def jones_to_fits(rayfront, filename, realimag=True, which=-1, nmodes=11, npix=128):
     """Write a N x N x 2 x 2 x 2 Jones Pupil to a FITS file
 
     Parameters
     ----------
-    data : numpy.ndarray complex128
+    rayfront : poke.Rayfront
         Jones pupil data generated by poke. Generally the first two dimensions are square, and correspond to the
         exit pupil coordinate. The remainder are matrix row, column, and then real/imag
     filename : str
         name of file to save the jones pupil as
     realimag : bool
         whether to save in real/imaginary parts. Defaults to True. If False, saves as absolute value and phase.
+    which : int
+        Which index of the poke.Rayfront.jones_pupil list to save, by default -1 or the most recent
+    nmodes : int
+        Number of Noll-ordered Zernike polynomials to use in the decomposition, by default 11
+    npix : int
+        Number of samples along one dimension of the square jones pupil, by default 128
     """
 
     # we don't require astropy, so the user needs to install it
@@ -114,21 +121,44 @@ def jones_to_fits(data,filename,realimag=True):
     except Exception as 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)
+
+    # The FITS headers that we want
+    # feel free to just use the variable names as the headers
+    wavelength = rayfront.wavelength
+    field_of_view_x = rayfront.fov[0]
+    field_of_view_y = rayfront.fov[1]
+    residuals_jxx = residual_list[0][1]
+    residuals_jxy = residual_list[1][1]
+    residuals_jyx = residual_list[2][1]
+    residuals_jyy = residual_list[3][1]
+
     if realimag:
-        realpart = np.real(data)
-        imagpart = np.imag(data)
+        realpart = np.real(jones)
+        imagpart = np.imag(jones)
     else:
-        realpart = np.abs(data)
-        imagpart = np.angle(data)
+        realpart = np.abs(jones)
+        imagpart = np.angle(jones)
 
-    box = np.empty([data.shape[0],data.shape[1],2,2,2])
+    box = np.empty([*jones.shape,2])
     box[...,0] = realpart
     box[...,1] = imagpart
 
-    hdu = fits.PrimaryHDU(box)
-    hdul = fits.HDUList([hdu])
-
-    hdul['wavelength'] = 1.6
-    hdul['pixelscale'] = 1
-
-    pass
+    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]))
+    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)