Merge pull request #104 from ajeldorado/algodiff_aj

Add regularization to AD-EFC. Add dead actuators.

Author: ajeldorado

examples/try_running_falco/EXAMPLE_try_running_FALCO_AD_EFC.py

@@ -52,10 +52,9 @@
 mp.ctrl.ad.iprint = 10
 mp.ctrl.ad.maxfun = 1000000
 
-
-mp.ctrl.log10regVec = np.array([-6, ])
-
-
+mp.ctrl.log10regVec = [0, ]
+mp.ctrl.ad.maxiter = 30  # 50
+mp.ctrl.ad.iprint = 5
 
 # %% Perform the Wavefront Sensing and Control
 

falco/ctrl.py

@@ -954,7 +954,10 @@ def _ad_efc(ni, vals_list, mp, cvar):
         bounds[cvar.uLegend == 2, 0] = dm2lb[mp.dm2.act_ele]
         bounds[cvar.uLegend == 2, 1] = dm2ub[mp.dm2.act_ele]
 
-    EFendPrev = []
+    EFend_list = []
+    Edm1post_list = []
+    Edm2pre_list = []
+    DM2surf_list = []
     for iMode in range(mp.jac.Nmode):
 
         modvar = falco.config.ModelVariables()
@@ -964,19 +967,23 @@ def _ad_efc(ni, vals_list, mp, cvar):
         modvar.starIndex = mp.jac.star_inds[iMode]
 
         # Calculate E-Field for previous EFC iteration
-        EFend = falco.model.compact(mp, modvar, isNorm=True, isEvalMode=False,
-                                    useFPM=True, forRevGradModel=False)
-        EFendPrev.append(EFend)
+        EFend, Edm1post, Edm2pre, DM1surf, DM2surf = falco.model.compact(
+            mp, modvar, isNorm=True, isEvalMode=False, useFPM=True, forRevGradModel=True)
+        EFend_list.append(EFend)
+        Edm1post_list.append(Edm1post)
+        Edm2pre_list.append(Edm2pre)
+        DM2surf_list.append(DM2surf)
 
     t0 = time.time()
 
     u_sol = scipy.optimize.minimize(
-        falco.model.compact_reverse_gradient, dm0, args=(mp, cvar.Eest, EFendPrev, log10reg),
+        falco.model.compact_reverse_gradient, dm0, args=(
+            log10reg, mp, cvar.Eest, EFend_list, Edm1post_list, Edm2pre_list, DM2surf_list),
         method='L-BFGS-B', jac=True, bounds=bounds,
         tol=None, callback=None,
         options={'disp': None,
                  'ftol': 1e-99,
-                 'gtol': 1e-10,
+                 'gtol': 1e-99,
                  'maxls': 20,
                  'maxiter': mp.ctrl.ad.maxiter,
                  'maxfun': mp.ctrl.ad.maxfun,
@@ -987,7 +994,6 @@ def _ad_efc(ni, vals_list, mp, cvar):
     duVec = u_sol.x
     print(u_sol.success)
     print(u_sol.nit)
-
     '''
     def optim_cost_func(Vdm):
         #cost function for L-BFGS

falco/dm.py

@@ -77,7 +77,7 @@ def gen_surf_from_act(dm, dx, Nout):
         dm = discretize_surf(dm, dm.HminStepMethod)
         heightMap = dm.VtoH*dm.Vquantized
     else:  # Quantization not desired; send raw, continuous voltages
-        heightMap = dm.VtoH*dm.V
+        heightMap = dm.VtoH * dm.V
 
     if hasattr(dm, 'orientation'):
         if dm.orientation.lower() == 'rot0':
@@ -904,6 +904,15 @@ def enforce_constraints(dm):
     # 1) Find actuators that exceed min and max values. Any actuators reaching
     # those limits are added to the pinned actuator list.
 
+    # Create dead actuator map if it doesn't exist already
+    if not hasattr(dm, 'dead_map'):
+        dm.dead_map = np.zeros((dm.Nact, dm.Nact), dtype=bool)  # initialize
+        for ii in dm.dead:
+            dm.dead_map.ravel()[ii] = True
+
+    # Apply dead actuator map
+    dm.V = dm.V * ~dm.dead_map
+
     # Min voltage limit
     Vtotal = dm.V + dm.biasMap
     new_inds = np.nonzero(Vtotal.flatten() < dm.Vmin)[0]  # linear indices of new actuators breaking their bounds

falco/model/jacobians.py

@@ -229,6 +229,18 @@ def lyot(mp, imode, idm, iact):
         specified Zernike mode, DM number, subband, and star.
 
     """
+    # Return zeros immediately if it's a dead actuator
+    if idm == 1:
+        if iact in mp.dm1.dead:
+            EFend = np.zeros_like(mp.Fend.corr.maskBool, dtype=complex)
+            Gchunk = EFend[mp.Fend.corr.maskBool]
+            return Gchunk
+    if idm == 2:
+        if iact in mp.dm2.dead:
+            EFend = np.zeros_like(mp.Fend.corr.maskBool, dtype=complex)
+            Gchunk = EFend[mp.Fend.corr.maskBool]
+            return Gchunk
+
     modvar = falco.config.ModelVariables()
     modvar.sbpIndex = mp.jac.sbp_inds[imode]
     modvar.zernIndex = mp.jac.zern_inds[imode]
@@ -657,6 +669,18 @@ def vortex(mp, imode, idm, iact):
         specified Zernike mode, DM number, subband, and star.
 
     """
+    # Return zeros immediately if it's a dead actuator
+    if idm == 1:
+        if iact in mp.dm1.dead:
+            EFend = np.zeros_like(mp.Fend.corr.maskBool, dtype=complex)
+            Gchunk = EFend[mp.Fend.corr.maskBool]
+            return Gchunk
+    if idm == 2:
+        if iact in mp.dm2.dead:
+            EFend = np.zeros_like(mp.Fend.corr.maskBool, dtype=complex)
+            Gchunk = EFend[mp.Fend.corr.maskBool]
+            return Gchunk
+
     modvar = falco.config.ModelVariables()
     modvar.sbpIndex = mp.jac.sbp_inds[imode]
     modvar.zernIndex = mp.jac.zern_inds[imode]

falco/model/models.py

@@ -422,7 +422,7 @@ def full_Fourier(mp, wvl, Ein, normFac, flagScaleFPM=False):
     return Eout
 
 
-def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
+def compact_reverse_gradient(command_vec, log10reg, mp, EestAll, EFend_list, Edm1post_list, Edm2pre_list, DM2surf_list):
     """
     Simplified (aka compact) model used by estimator and controller.
 
@@ -488,11 +488,10 @@ def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
     dDM2Vvec[mp.dm2.act_ele] = command_vec[mp.ctrl.uLegend == 2]
     dv_dm1 = dDM1Vvec.reshape((mp.dm1.Nact, mp.dm1.Nact))
     dv_dm2 = dDM2Vvec.reshape((mp.dm2.Nact, mp.dm2.Nact))
-    # mp.dm1.V += command_vec[0:mp.dm1.NactTotal].reshape([mp.dm1.Nact, mp.dm1.Nact])
-    # mp.dm2.V += command_vec[mp.dm2.NactTotal::].reshape([mp.dm2.Nact, mp.dm2.Nact])
 
-    total_cost = 0  # initialize
-    normFacADweightedSum = 0
+    # initialize
+    total_cost = 0
+    # normFacADweightedSum = 0
 
     for imode in range(mp.jac.Nmode):
         modvar = falco.config.ModelVariables()
@@ -505,17 +504,18 @@ def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
         kk = mirrorFac*2*np.pi/wvl
         I00 = mp.Fend.compact.I00[modvar.sbpIndex]
         # normFac = mp.Fend.compact.I00[modvar.sbpIndex]
+        # print(f'normFacAD = {normFacAD}')
         # normFacFull = np.mean(mp.Fend.full.I00[modvar.sbpIndex, :])
         EestVec = EestAll[:, imode]
         Eest2D = np.zeros_like(mp.Fend.corr.maskBool, dtype=complex)
-        Eest2D[mp.Fend.corr.maskBool] = EestVec  # * np.sqrt(normFacFull)  # Remove normalization
-        normFacAD = np.sum(np.abs(EestVec)**2)
-        # print(f'normFacAD = {normFacAD}')
+        Eest2D[mp.Fend.corr.maskBool] = EestVec
+        # normFacAD = np.sum(np.abs(EestVec)**2)
 
-        # Get model-based E-field before deltas
-        EFendA, Edm1post, Edm2pre, DM1surf, DM2surf = compact(
-           mp, modvar, isNorm=True, isEvalMode=isEvalMode, useFPM=useFPM,
-           forRevGradModel=True)
+        # Get model-based E-field before deltas. Should be pre-computed for speed.
+        EFendA = EFend_list[imode]
+        Edm1post = Edm1post_list[imode]
+        Edm2pre = Edm2pre_list[imode]
+        DM2surf = DM2surf_list[imode]
 
         # Get model-based E-field With delta DM commands applied.
         mp.dm1.V = mp.dm1.V0 + dv_dm1
@@ -528,17 +528,14 @@ def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
         mp.dm2.V = mp.dm2.V0.copy()
 
         # Compute the delta E-field from the latest commands (model new - model old).
-        # EFendA = EFendPrev[imode]
         dEend = EFendB - EFendA
 
         # DH = EFend[mp.Fend.corr.maskBool]
-        # Eest2D = EFendA  # DEBUGGING
         EdhNew = Eest2D + dEend
         DH = EdhNew[mp.Fend.corr.maskBool]
         int_in_dh = np.sum(np.abs(DH)**2)
-        # print(f'int_in_dh = {int_in_dh}')
         total_cost += mp.jac.weights[imode] * int_in_dh
-        normFacADweightedSum += mp.jac.weights[imode] / normFacAD
+        # normFacADweightedSum += mp.jac.weights[imode] / normFacAD
 
         # Gradient
         Fend_masked = mp.jac.weights[imode]*2/np.sqrt(I00)*EdhNew*mp.Fend.corr.maskBool
@@ -612,14 +609,14 @@ def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
         Edm1_grad = falco.prop.ptp(Edm2_grad, mp.P2.compact.dx*Edm2_grad.shape[0], wvl, -mp.d_dm1_dm2)
         surf_dm1_bar = -kk*np.imag(Edm1_grad * np.conj(Edm1post))
 
-        surf_dm2_bar_total += mp.jac.weights[imode] * surf_dm2_bar
-        surf_dm1_bar_total += mp.jac.weights[imode] * surf_dm1_bar
+        surf_dm2_bar_total += surf_dm2_bar
+        surf_dm1_bar_total += surf_dm1_bar
 
-    # # Calculate DM penalty term component of cost function
-    # utu_coefs = normFacADweightedSum * mp.ctrl.ad.utu_scale_fac * 10.0**(log10reg)
-    # total_cost += utu_coefs * np.sum(command_vec**2)
-    # # print('normFacADweightedSum = %.4g' % normFacADweightedSum)
-    # # print('utu_coefs = %.4g' % utu_coefs)
+    # Calculate DM penalty term component of cost function
+    utu_coef = mp.ctrl.ad.utu_scale_fac * 10.0**(log10reg)
+    total_cost += utu_coef * np.sum(command_vec**2)
+    # print('normFacADweightedSum = %.4g' % normFacADweightedSum)
+    # print('utu_coef = %.4g' % utu_coef)
 
     if mp.dm1.useDifferentiableModel and mp.dm2.useDifferentiableModel:
         Vout1 = mp.dm1.differentiableModel.render_backprop(
@@ -630,11 +627,12 @@ def compact_reverse_gradient(command_vec, mp, EestAll, EFendPrev, log10reg):
     else:
         raise ValueError('mp.dm1.useDifferentiableModel and mp.dm2.useDifferentiableModel must be True for AD-EFC.')
 
-    Vout1 *= mp.dm1.VtoH
-    Vout2 *= mp.dm2.VtoH
+    Vout1 *= mp.dm1.VtoH * np.conj(~mp.dm1.dead_map)
+    Vout2 *= mp.dm2.VtoH * np.conj(~mp.dm2.dead_map)
     gradient = np.concatenate((Vout1.reshape([mp.dm1.NactTotal])[mp.dm1.act_ele],
                                Vout2.reshape([mp.dm2.NactTotal])[mp.dm2.act_ele]),
                               axis=None)
+    gradient += 2 * utu_coef * gradient
 
     return total_cost, gradient
 

falco/setup.py

@@ -312,6 +312,8 @@ def set_optional_variables(mp):
         mp.dm1.orientation = 'rot0'  # Change to mp.dm1.V orientation before generating DM surface. Options: rot0, rot90, rot180, rot270, flipxrot0, flipxrot90, flipxrot180, flipxrot270
     if not hasattr(mp.dm1, 'fitType'):
         mp.dm1.fitType = 'linear'  # Type of response for displacement vs voltage. Options are 'linear', 'quadratic', and 'fourier2'.
+    if not hasattr(mp.dm1, 'dead'):
+        mp.dm1.dead = []  # Linear indices of dead actuators
     if not hasattr(mp.dm1, 'pinned'):
         mp.dm1.pinned = np.array([])  # Indices of pinned actuators
     if not hasattr(mp.dm1, 'Vpinned'):
@@ -346,6 +348,8 @@ def set_optional_variables(mp):
         mp.dm2.orientation = 'rot0'  # Change to mp.dm2.V orientation before generating DM surface. Options: rot0, rot90, rot180, rot270, flipxrot0, flipxrot90, flipxrot180, flipxrot270
     if not hasattr(mp.dm2, 'fitType'):
         mp.dm2.fitType = 'linear'  # Type of response for displacement vs voltage. Options are 'linear', 'quadratic', and 'fourier2'.
+    if not hasattr(mp.dm2, 'dead'):
+        mp.dm2.dead = []  # Linear indices of dead actuators
     if not hasattr(mp.dm2, 'pinned'):
         mp.dm2.pinned = np.array([])  # Indices of pinned actuators
     if not hasattr(mp.dm2, 'Vpinned'):
@@ -1260,7 +1264,14 @@ def falco_configure_dm1_and_dm2(mp):
         mp.dm1.V = np.zeros((mp.dm1.Nact, mp.dm1.Nact))
     if not hasattr(mp.dm2, 'V'):
         mp.dm2.V = np.zeros((mp.dm2.Nact, mp.dm2.Nact))
-    pass
+
+    # Dead actuator 2-D maps
+    mp.dm1.dead_map = np.zeros((mp.dm1.Nact, mp.dm1.Nact), dtype=bool)
+    mp.dm2.dead_map = np.zeros((mp.dm2.Nact, mp.dm2.Nact), dtype=bool)
+    for ii in mp.dm1.dead:
+        mp.dm1.dead_map.ravel()[ii] = True
+    for ii in mp.dm2.dead:
+        mp.dm2.dead_map.ravel()[ii] = True
 
     # Initialize the number of elements used per DM
     if np.any(mp.dm_ind == 1):
@@ -1296,7 +1307,7 @@ def falco_gen_dm_stops(mp):
     if mp.flagDM2stop:
         mp.dm2.full.mask = falco.mask.falco_gen_DM_stop(mp.P2.full.dx, mp.dm2.Dstop, mp.centering)
         mp.dm2.compact.mask = falco.mask.falco_gen_DM_stop(mp.P2.compact.dx, mp.dm2.Dstop, mp.centering)
-    
+
     return None
 
 

tests/functional/test_ad_efc_gradient_lc.py

@@ -131,7 +131,11 @@ def test_adjoint_model_lc():
         # Compute gradient
         EestAll = ev.Eest
         log10reg = -np.inf
-        EFendPrev = []
+        mp.ctrl.ad.utu_scale_fac = 1  # not used when log10reg = -np.inf
+        EFend_list = []
+        Edm1post_list = []
+        Edm2pre_list = []
+        DM2surf_list = []
 
         falco.ctrl.init(mp, cvar)
 
@@ -144,10 +148,15 @@ def test_adjoint_model_lc():
             modvar.starIndex = mp.jac.star_inds[iMode]
 
             # Calculate E-Field for previous EFC iteration
-            EFend = falco.model.compact(mp, modvar, isNorm=True, isEvalMode=False,
-                                        useFPM=True, forRevGradModel=False)
-            EFendPrev.append(EFend)
-        total_cost, u_bar_out = falco.model.compact_reverse_gradient(du, mp, EestAll, EFendPrev, log10reg)
+            EFend, Edm1post, Edm2pre, DM1surf, DM2surf = falco.model.compact(
+                mp, modvar, isNorm=True, isEvalMode=False, useFPM=True, forRevGradModel=True)
+            EFend_list.append(EFend)
+            Edm1post_list.append(Edm1post)
+            Edm2pre_list.append(Edm2pre)
+            DM2surf_list.append(DM2surf)
+
+        total_cost, u_bar_out = falco.model.compact_reverse_gradient(
+            du, log10reg, mp, EestAll, EFend_list, Edm1post_list, Edm2pre_list, DM2surf_list)
 
         u1_bar_out = u_bar_out[0:mp.dm1.NactTotal]
         u2_bar_out = u_bar_out[mp.dm1.NactTotal::]

tests/functional/test_ad_efc_gradient_vc.py

@@ -187,7 +187,11 @@ def test_adjoint_model_vc():
         # Compute gradient
         EestAll = ev.Eest
         log10reg = -np.inf
-        EFendPrev = []
+        mp.ctrl.ad.utu_scale_fac = 1  # not used when log10reg = -np.inf
+        EFend_list = []
+        Edm1post_list = []
+        Edm2pre_list = []
+        DM2surf_list = []
 
         falco.ctrl.init(mp, cvar)
 
@@ -200,10 +204,15 @@ def test_adjoint_model_vc():
             modvar.starIndex = mp.jac.star_inds[iMode]
 
             # Calculate E-Field for previous EFC iteration
-            EFend = falco.model.compact(mp, modvar, isNorm=True, isEvalMode=False,
-                                        useFPM=True, forRevGradModel=False)
-            EFendPrev.append(EFend)
-        total_cost, u_bar_out = falco.model.compact_reverse_gradient(du, mp, EestAll, EFendPrev, log10reg)
+            EFend, Edm1post, Edm2pre, DM1surf, DM2surf = falco.model.compact(
+                mp, modvar, isNorm=True, isEvalMode=False, useFPM=True, forRevGradModel=True)
+            EFend_list.append(EFend)
+            Edm1post_list.append(Edm1post)
+            Edm2pre_list.append(Edm2pre)
+            DM2surf_list.append(DM2surf)
+            
+        total_cost, u_bar_out = falco.model.compact_reverse_gradient(
+            du, log10reg, mp, EestAll, EFend_list, Edm1post_list, Edm2pre_list, DM2surf_list)
 
         u1_bar_out = u_bar_out[0:mp.dm1.NactTotal]
         u2_bar_out = u_bar_out[mp.dm1.NactTotal::]

tests/functional/test_diff_dm_model.py

@@ -41,6 +41,7 @@ def test_diff_dm_model():
     mp.dm1.edgeBuffer = 1  # max radius (in actuator spacings) outside of beam on DM surface to compute influence functions for. [actuator widths]
 
     mp.dm1.fitType = 'linear'
+    mp.dm1.dead = []
     mp.dm1.pinned = np.array([])
     mp.dm1.Vpinned = np.zeros_like(mp.dm1.pinned)
     mp.dm1.tied = np.zeros((0, 2))

tests/functional/test_dm_orientation.py

@@ -32,6 +32,7 @@ def test_surface_orientation():
     mp.dm1.edgeBuffer = 1  # max radius (in actuator spacings) outside of beam on DM surface to compute influence functions for. [actuator widths]
 
     mp.dm1.fitType = 'linear'
+    mp.dm1.dead = []
     mp.dm1.pinned = np.array([])
     mp.dm1.Vpinned = np.zeros_like(mp.dm1.pinned)
     mp.dm1.tied = np.zeros((0, 2))
@@ -126,6 +127,7 @@ def test_surface_orientation_from_cube():
     mp.dm1.edgeBuffer = 1  # max radius (in actuator spacings) outside of beam on DM surface to compute influence functions for. [actuator widths]
 
     mp.dm1.fitType = 'linear'
+    mp.dm1.dead = []
     mp.dm1.pinned = np.array([])
     mp.dm1.Vpinned = np.zeros_like(mp.dm1.pinned)
     mp.dm1.tied = np.zeros((0, 2))