add some plots, move GS to threadripper

photonics/__init__.py

@@ -1,11 +1,9 @@
-from .utils import *
-from .linearity import *
+from matplotlib import pyplot as plt
 
-from .experiments.calibrate_dm import *
-from .experiments.lantern_reader import *
-from .experiments.seal import *
-
-from .simulations.lantern_optics import *
-from .simulations.optics import *
-from .simulations.pyramid_optics import *
-from .simulations.second_stage import *
+plt.rc('font', family='serif',size=12)
+plt.rcParams.update({
+    "text.usetex": False,
+    "font.family": "serif",
+    "font.serif" : "cmr10",
+    "axes.formatter.use_mathtext" : True
+})

photonics/experiments/lantern_fits.py

@@ -0,0 +1,27 @@
+from datetime import datetime
+from astropy.io import fits
+
+# Copilot experiment to help me get over procrastinating on making a standard format for lantern outputs.
+def create_fits_file(commands, science_images, exposure_time=10.0, dark_frame=None, output_path='/path/to/output.fits'):
+    # Create a new FITS file
+    fits_file = fits.HDUList()
+
+    # Add a primary header
+    primary_header = fits.Header()
+    primary_header['DATE-OBS'] = datetime.now().isoformat()
+    fits_file.append(fits.PrimaryHDU(header=primary_header))
+
+    # Add camera exposure time
+    fits_file[0].header['EXPTIME'] = exposure_time
+
+    # Add dark frame image
+    if dark_frame is not None:
+        fits_file.append(fits.ImageHDU(data=dark_frame, name='DARK_FRAME'))
+
+    # Add deformable mirror commands and science images
+    for command, science_image in zip(commands, science_images):
+        fits_file.append(fits.ImageHDU(data=command, name='COMMAND'))
+        fits_file.append(fits.ImageHDU(data=science_image, name='SCIENCE_IMAGE'))
+
+    # Save the FITS file
+    fits_file.writeto(output_path, overwrite=True)

photonics/simulations/lantern_optics.py

@@ -1,15 +1,19 @@
-import os
 from os.path import join
 import numpy as np
 import hcipy as hc
 import lightbeam as lb
 from hcipy import imshow_field
 from matplotlib import pyplot as plt
 from tqdm import trange
-from ..utils import PROJECT_ROOT, date_now, zernike_names
+from ..utils import PROJECT_ROOT, date_now, zernike_names, nanify
 from .command_matrix import make_command_matrix
 
 class LanternOptics:
+	"""
+	Implements the optics around a photonic lantern and wraps propagations from Lightbeam to determine the lantern's behaviour.
+ 
+	Need to ask Emiel about the canonical hcipy way of organizing this vs. the parent "optics" class.
+	"""
 	def __init__(self, optics, nports=19, nmodes=9):
 		self.name = f"lantern_ports{nports}_modes{nmodes}"
 		self.nports = nports # update this later
@@ -47,6 +51,7 @@ def __init__(self, optics, nports=19, nmodes=9):
 		self.extent_y = (np.min(self.input_footprint[1]), np.max(self.input_footprint[1]))
 		self.load_outputs()
 		self.focal_propagator = optics.focal_propagator
+		self.focal_grid = optics.focal_grid
 		make_command_matrix(optics.deformable_mirror, self, optics.wf, rerun=True)
 
 	def input_to_2d(self, input_efield, zoomed=True):
@@ -172,55 +177,68 @@ def show_linearity(self, amplitudes, linearity_responses):
 				axs[r][c].plot(amplitudes, linearity_responses[i,:,j], alpha=alpha)
 		plt.show()
 
-	def forward(self, pupil):
-		focal = self.prop.forward(pupil)
-		_, _, reading = self.lantern_output(focal)
-		return hc.Wavefront(hc.Field(reading.ravel(), self.focal_grid), wavelength=self.wl)
+	def forward(self, optics, pupil):
+		focal = self.focal_propagator.forward(pupil)
+		_, reading = self.lantern_output(focal)
+		return hc.Wavefront(hc.Field(reading.ravel(), optics.focal_grid), wavelength=self.wl)
 
-	def backward(self, reading):
+	def backward(self, optics, reading):
 		coeffs = self.lantern_reverse @ reading.electric_field
 		lantern_input = self.input_to_2d(coeffs @ self.outputs, zoomed=False)
 		focal_wf = hc.Wavefront(hc.Field(lantern_input.ravel(), self.focal_grid), wavelength=self.wl)
-		pupil_field = self.prop.backward(focal_wf)
+		pupil_field = self.focal_propagator.backward(focal_wf)
 		return pupil_field
-
-	def GS(self, img, niter=10):
-		"""
-		Gerchberg-Saxton algorithm
-		"""        
+
+	def GS_init(self, optics, img):
 		# Normalization
 		img /= img.sum()
 
 		# Amplitude in - measured (or assumed to be known)
-		measuredAmplitude_in = self.pupil_wf_ref
+		measuredAmplitude_in = optics.wf
 		# Amplitude out - measured
 		measuredAmplitude_out = np.sqrt(img) 
-
 		# Forward Propagation in WFS assuming 0 phase
-		EM_out = self.forward(measuredAmplitude_in)
+		EM_out = self.forward(optics, measuredAmplitude_in)
 		# replacing by known amplitude
 		phase_out_0 = EM_out.phase
 		EM_out = measuredAmplitude_out * np.exp(1j*phase_out_0)
 		# Back Propagation in PWFS
-		EM_in = self.backward(hc.Wavefront(EM_out, wavelength=self.wl))
+		EM_in = self.backward(optics, hc.Wavefront(EM_out, wavelength=self.wl))
+		return EM_in, measuredAmplitude_in, measuredAmplitude_out
+
+	def GS_iteration(self, optics, EM_in, measuredAmplitude_in, measuredAmplitude_out):
+		# replacing by known amplitude
+		phase_in_k = EM_in.phase
+		EM_in = hc.Wavefront(measuredAmplitude_in.electric_field*np.exp(1j*phase_in_k),wavelength=self.wl)
+		# Lantern forward propagation
+		EM_out = self.forward(optics, EM_in)
+		# replacing by known amplitude
+		phase_out_k = EM_out.phase
+		EM_out = hc.Wavefront(measuredAmplitude_out*np.exp(1j*phase_out_k), wavelength=self.wl)
+		# Lantern backward propagation    
+		EM_in = self.backward(optics, EM_out)
+
+		return EM_in
+
+	def GS(self, optics, img, niter=10):
+		"""
+		Gerchberg-Saxton algorithm
+		"""
+		EM_in, measuredAmplitude_in, measuredAmplitude_out = self.GS_init(optics, img)
+		print(EM_in)
 		for _ in trange(niter):
-			# replacing by known amplitude
-			phase_in_k = EM_in.phase
-			EM_in = hc.Wavefront(measuredAmplitude_in.electric_field*np.exp(1j*phase_in_k),wavelength=self.wl)
-			# Lantern forward propagation
-			EM_out = self.forward(EM_in)
-			# replacing by known amplitude
-			phase_out_k = EM_out.phase
-			EM_out = hc.Wavefront(measuredAmplitude_out*np.exp(1j*phase_out_k), wavelength=self.wl)
-			# Lantern backward propagation    
-			EM_in = self.backward(EM_out)
+			EM_in = self.GS_iteration(optics, EM_in, measuredAmplitude_in, measuredAmplitude_out)
 
 		return EM_in
 
-	def show_GS(self, zernike, amplitude, niter=10):
-		input_phase = self.zernike_to_phase(zernike, amplitude)
-		reading = self.forward(self.phase_to_pupil(input_phase))
-		retrieved = self.GS(reading.intensity, niter=niter)
+	def show_GS(self, optics, zernike, amplitude, niter=10):
+		input_phase = optics.zernike_to_phase(zernike, amplitude)
+		input_pupil = optics.phase_to_pupil(input_phase)
+		reading = self.forward(optics, input_pupil)
+		input_focal = self.focal_propagator.forward(input_pupil)
+		coeffs = self.lantern_coeffs(input_focal)
+		out_to_plot = np.abs(sum(c * lf for (c, lf) in zip(coeffs, self.plotting_launch_fields))) ** 2
+		retrieved = self.GS(optics, reading.intensity, niter=niter)
 		retphase = retrieved.phase # (retrieved.phase % np.pi) - np.pi / 2
 		fig, axs = plt.subplots(1, 3)
 		fig.suptitle(f"Lantern phase retrieval, Zernike {zernike}, amplitude {amplitude}")
@@ -230,20 +248,21 @@ def show_GS(self, zernike, amplitude, niter=10):
 			ax.set_yticks([])
 		# vmin = np.minimum(np.min(input_phase), np.min(retphase))
 		# vmax = np.maximum(np.max(input_phase), np.max(retphase))
-		axs[0].imshow(input_phase.shaped)
+		axs[0].imshow(nanify(input_phase.shaped, optics.aperture.shaped), cmap="RdBu")
 		axs[0].set_title("Phase screen")
-		axs[1].imshow(np.abs(reading.intensity.shaped))
+		axs[1].imshow(out_to_plot)
 		axs[1].set_title("Lantern output")
-		hc.imshow_field(retphase * self.aperture, ax=axs[2])
+		axs[2].imshow(nanify(retphase.shaped, optics.aperture.shaped), cmap="RdBu")
 		axs[2].set_title("Retrieved phase")
 		plt.show()
 
 	def plot_outputs(self):
-		fig, axs = plt.subplots(5, 4)
+		rm, cm = 4, 5
+		fig, axs = plt.subplots(rm, cm)
 		plt.suptitle("Photonic lantern entrance modes")
-		plt.subplots_adjust(wspace=-0.7, hspace=0.1)
+		plt.subplots_adjust(wspace=0.05, hspace=0.05)
 		for (i, o) in enumerate(self.outputs):
-			r, c = i // 4, i % 4
+			r, c = i // cm, i % cm
 			axs[r][c].imshow(np.abs(self.input_to_2d(o)))
 			axs[r][c].set_xticks([])
 			axs[r][c].set_yticks([])

photonics/utils.py

@@ -25,7 +25,7 @@ def time_now():
     return datetime.now().strftime("%H%M")
 
 def datetime_now():
-    return date_now() + "_" + time_now()
+    return datetime.now().isoformat()
 
 def rms(x):
     return np.sqrt(np.mean((x - np.mean(x)) ** 2))

scripts_jl/plotting/lantern_static_aberration.jl

@@ -32,7 +32,7 @@ heatmap(log10.(psf_after)[m...], c=:grays)
 
 begin
     p = plot(vcat(rms_before, rms_during, rms_after) / (2pi / λ), xlabel="Number of iterations", ylabel="Wavefront error measured on the PL (nm)", label=nothing, ylim=(0, 170), title="Correcting a static aberration with the photonic lantern", titlefontsize=13, dpi=800)
-    vspan!([20, 60], color=RGBA(0.76, 1, 0.76, 1), alpha=0.3, label="Loop closed")
+    vspan!([20, 60], color=RGBA(0.2, 1, 0.2, 1), alpha=0.3, label="Loop closed")
     BB1 = bbox(0.01, 0.2, 0.3, 0.3)
     BB2 = bbox(0.65, 0.2, 0.3, 0.3)
     BB3 = bbox(0.01, 0.6, 0.2, 0.3)
@@ -41,6 +41,6 @@ begin
     im_show!(phase_screen_final * (λ / (2pi)), inset=(1, BB2), clim=(-λ/2, λ/2), subplot=3, title="Final phase (nm)", titlefontsize=7, background_color_subplot=RGBA(0, 0, 0, 0), background_alpha_subplot=0.0, ytickfontsize=6, topmargin=-2mm)
     heatmap!(log10.(psf_before[m...]), c=:grays, showaxis=false, grid=false, inset=(1, BB3), subplot=4, title="Initial image", titlefontsize=7, xlim=(0, length(m[1])), ylim=(0, length(m[2])), legend=nothing, yflip=true, topmargin=-2mm)
     heatmap!(log10.(psf_after[m...]), c=:grays, showaxis=false, grid=false, inset=(1, BB4), subplot=5, title="Final image", titlefontsize=7, xlim=(0, length(m[1])), ylim=(0, length(m[2])), legend=nothing, yflip=true, topmargin=-2mm)
-    Plots.savefig("./figures/seal_static_cl.png")
+    Plots.savefig("./figures/seal_static_cl.pdf")
     p
 end

scripts_py/sim/fnumber_tests.py

@@ -3,9 +3,6 @@
 from matplotlib import pyplot as plt
 from matplotlib import animation
 from IPython.display import HTML
-import hcipy as hc
-from hcipy import imshow_field
-from scipy.optimize import minimize
 from photonics.simulations.lantern_optics import LanternOptics
 from itertools import product, repeat
 from photonics import PROJECT_ROOT

scripts_py/sim/gs.py

@@ -1,7 +1,64 @@
 # %%
 from photonics.simulations.lantern_optics import LanternOptics
+from photonics.simulations.optics import Optics
+from hcipy import imshow_field
+from photonics.utils import nanify
+import numpy as np
+from matplotlib import pyplot as plt
+from tqdm import trange
 
-lo = LanternOptics(f_number=10.0)
 # %%
-lo.show_GS(3, -0.5, niter=10)
+optics = Optics(lantern_fnumber=6.5)
+lo = LanternOptics(optics)
 # %%
+lo.show_GS(optics, 3, 0.5, niter=10)
+# %%
+gs_ret = lo.GS(
+    optics,
+    lo.forward(
+        optics,
+        optics.zernike_to_pupil(3, 0.1)
+    ).intensity,
+    niter=20
+)
+# %%
+optics.zernike_basis.coefficients_for(optics.zernike_to_pupil(3, 0.1).phase)[:lo.nmodes]
+# %%
+optics.zernike_basis.coefficients_for(gs_ret.phase)[:lo.nmodes]
+# %%
+imshow_field(nanify(gs_ret.phase - np.mean(gs_ret.phase), optics.aperture), cmap="RdBu")
+plt.gca().set_axis_off()
+# %%
+z_applied = 2
+a_applied = -0.5
+zdecomps = []
+EM_in, measured_in, measured_out = lo.GS_init(
+    optics,
+    lo.forward(
+        optics,
+        optics.zernike_to_pupil(z_applied, a_applied)
+    ).intensity
+)
+zdecomps.append(
+    optics.zernike_basis.coefficients_for(EM_in.phase)[:lo.nmodes]
+)
+for i in trange(4):
+    EM_in = lo.GS_iteration(optics, EM_in, measured_in, measured_out)
+    zdecomps.append(
+        optics.zernike_basis.coefficients_for(EM_in.phase)[:lo.nmodes]
+    )
+
+zdecomps = np.array(zdecomps).T
+
+# %%
+plt.hlines(a_applied, 0, zdecomps.shape[1] - 1, label="Target")
+for (i, r) in enumerate(zdecomps):
+    if i == z_applied:
+        plt.plot(r, alpha=1, label="Injected mode", color="k")
+    else:
+        plt.plot(r, alpha=0.1, color="r")
+plt.xticks(np.arange(0, zdecomps.shape[1], 4))
+plt.xlabel("Gerchberg-Saxton iteration")
+plt.ylabel("Amplitude (rad)")
+plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0.)
+plt.subplots_adjust(right=0.8)

src/PhotonicLantern.jl

@@ -64,11 +64,11 @@ module PhotonicLantern
     end
 
     function im_show(x; kwargs...)
-        heatmap(nanify(x), aspect_ratio=1, showaxis=false, grid=false, xlim=(0, size(x, 1) + 2), c=:thermal; kwargs...)
+        heatmap(nanify(x), aspect_ratio=1, showaxis=false, grid=false, xlim=(0, size(x, 1) + 2), c=:RdBu; kwargs...)
     end
 
     function im_show!(x; kwargs...)
-        heatmap!(nanify(x), aspect_ratio=1, showaxis=false, grid=false, xlim=(0, size(x, 1) + 2), c=:thermal; kwargs...)
+        heatmap!(nanify(x), aspect_ratio=1, showaxis=false, grid=false, xlim=(0, size(x, 1) + 2), c=:RdBu; kwargs...)
     end
 
     function phasewrap(x)

writing/ncpa.tex

@@ -0,0 +1,53 @@
+\documentclass{article}
+
+\usepackage{custom}
+\usepackage{tikz}
+\usetikzlibrary{shapes,arrows,positioning,decorations.pathmorphing}
+
+\begin{document}
+\tikzstyle{block} = [draw, rectangle, 
+    minimum height=3em, minimum width=6em]
+\tikzstyle{sum} = [draw, fill=white, circle, node distance=1cm]
+\tikzstyle{input} = [coordinate]
+\tikzstyle{output} = [coordinate]
+\tikzstyle{pinstyle} = [pin edge={to-,thin,black}]
+\tikzstyle{gain} = [draw, fill=white, isosceles triangle, isosceles triangle apex angle = 60, shape border rotate=#1]
+
+\tikzset{snake it/.style={decorate, decoration=snake}}
+
+\begin{tikzpicture}[auto, node distance=3cm,>=latex', text width=2.2cm, align=center]
+    \node [name=disturbance] {Atmospheric disturbance};
+    \node [sum, below of=disturbance, node distance=2cm, text width=0.5cm] (dm) {+};
+    \node [draw, rectangle, below of=dm, node distance=2cm, minimum height=2em, minimum width=2em, text width=0.1cm, color=gray] (beamsplitter) {};
+    \node [input, below right=1em and 1em of beamsplitter.center, anchor=center] (beamsplitterbr) {};
+    \node [input, above left=1em and 1em of beamsplitter.center, anchor=center] (beamsplittertl) {};
+    \node [input, below of=beamsplitter, node distance=2cm] (turn) {};
+    \node [name=beamsplittertext, left of=beamsplitter, node distance=1.0cm] {Beam\\splitter};
+    \node [block, minimum height=5em] (dmblock) at (dm) {};
+    \node [above right=1.5em and 2em of dmblock.center, anchor=center] (dmtext) {DM};
+    \node [block, right of=turn] (pl) {Science image};
+    \node [block, right of=beamsplitter] (pywfs) {Wavefront sensor};
+    \node [block, right of=pywfs,] (pyrecon) {Wavefront reconstructor};
+    \node [input, right of=pyrecon, node distance=1.5cm] (pyswitch) {};
+    \node [input, above of=pyswitch, node distance=0.9cm] (turn4) {};
+    \node [input, right of=pyrecon, text width=0.5cm, node distance=2cm] (sum2) {};
+    \node [input, above of=sum2, node distance=2cm] (turn3) {};
+    \node [gain, left of=turn3, shape border rotate=-180, text width=0.8cm, node distance=1.5cm] (gain) {Gain};
+    \node [block, left of=gain, node distance=3cm] (integ) {Integrator};
+    \node [input, left of=turn, node distance=1.05cm] (wave) {};
+
+    \path[draw,->] (disturbance) -- (dm);
+    \path[draw,->] (dm) -- (beamsplitter.center);
+    \path[draw,-,color=gray] (beamsplittertl) -- (beamsplitterbr);
+    \path[draw,->] (beamsplitter.center) -- (pywfs);
+    \path[draw,->] (pywfs) -- (pyrecon);
+    \path[draw,-] (pyrecon) -- (sum2);
+    \path[draw,-] (beamsplitter.center) -- (turn);
+    \path[draw,->] (turn) -- (pl);
+    \path[draw,-] (sum2) -- (turn3);
+    \path[draw,->] (turn3) -- (gain);
+    \path[draw,->] (gain) -- (integ);
+    \path[draw,->] (integ) -- (dm) node[above, pos=0.95] {--};
+    \path[draw,->, snake it, color=red] (wave) -- (turn) node[left, pos=0.7] {NCPA};
+\end{tikzpicture}
+\end{document}