From 0dd9d49610c46b42c48c1ad65d6d896f9619dd19 Mon Sep 17 00:00:00 2001
From: Marshall Perrin <mperrin@stsci.edu>
Date: Tue, 5 Mar 2024 14:35:30 -0500
Subject: [PATCH] fix tests

---
 poppy/dms.py               |  2 +-
 poppy/optics.py            | 20 +++++++++++---------
 poppy/tests/test_dms.py    | 35 +++++++++++++++++++++++++++++++++++
 poppy/tests/test_optics.py |  2 +-
 4 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/poppy/dms.py b/poppy/dms.py
index d4ffed46..f5eea52d 100644
--- a/poppy/dms.py
+++ b/poppy/dms.py
@@ -870,7 +870,7 @@ class WedgeSegmentedDeformableMirror(SegmentedDeformableMirror, optics.WedgeSegm
                 units, and the WFE is therefore a factor of two larger. The returned WFE will be twice the
                 amplitude of the requested values (convolved with the actuator response function etc.)
     """
-   
+
     def __init__(self, name='WedgeSegDM', radius=1.0 * u.m, rings=1, nsections=4, gap_radii=None, gap=0.01 * u.m,
                  include_factor_of_two=False, **kwargs):
         #FIXME ? using grey pixel does not work. something in the geometry module generate a true divide error
diff --git a/poppy/optics.py b/poppy/optics.py
index 9b84f14c..0c3c4383 100644
--- a/poppy/optics.py
+++ b/poppy/optics.py
@@ -1682,15 +1682,17 @@ def get_transmission(self, wave):
                 # print(f"drawing ring gap {iring} at {r_ring_inner}")
                 self.transmission[np.abs(r - r_ring_inner) < halfgapwidth] = 0
 
-            for igap in range(self.nsections[iring]):
-                angle = self.gap_angles[iring][igap]
-                # print(f"  linear gap {igap} at {angle} radians")
-                # calculate rotated x' and y' coordinates after rotation by that angle.
-                x_p = np.cos(angle) * x + np.sin(angle) * y
-                y_p = -np.sin(angle) * x + np.cos(angle) * y
-
-                self.transmission[(0 < x_p) & (r_ring_inner < r) & (r < r_ring_outer) &
-                                  (np.abs(y_p) < halfgapwidth)] = 0
+            if self.nsections[iring] > 1:
+                # If we have more than 1 segment in this ring, draw the gaps
+                for igap in range(self.nsections[iring]):
+                    angle = self.gap_angles[iring][igap]
+                    # print(f"  linear gap {igap} at {angle} radians")
+                    # calculate rotated x' and y' coordinates after rotation by that angle.
+                    x_p = np.cos(angle) * x + np.sin(angle) * y
+                    y_p = -np.sin(angle) * x + np.cos(angle) * y
+
+                    self.transmission[(0 < x_p) & (r_ring_inner < r) & (r < r_ring_outer) &
+                                      (np.abs(y_p) < halfgapwidth)] = 0
 
         if not self._include_center: # mask out the center ring / center zeroth segment
             self.transmission[r < self.gap_radii[0].to_value(u.m)] = 0
diff --git a/poppy/tests/test_dms.py b/poppy/tests/test_dms.py
index c5e8235d..a971965d 100644
--- a/poppy/tests/test_dms.py
+++ b/poppy/tests/test_dms.py
@@ -194,3 +194,38 @@ def test_basic_circular_dm():
     assert peak_aberrated < peak_perf, "Adding nonzero WFE did not decrease the Strehl as expected."
 
     return psf_aberrated, psf_perf, osys
+
+
+
+def test_basic_wedge_dm():
+    """ A simple test for the circularly segmented deformable mirror code -
+    can we move actuators, and does adding nonzero WFE result in decreased Strehl?"""
+
+    dm = dms.WedgeSegmentedDeformableMirror(rings=2, nsections=[1,6])
+
+    osys = poppy_core.OpticalSystem(npix=256)
+    osys.add_pupil(dm)
+    osys.add_detector(0.010, fov_pixels=128)
+
+    psf_perf = osys.calc_psf()
+
+    for act in ( 3,6):
+        dm.set_actuator(act, 1e-6, 0, 1e-7) # 1000 nm = 1 micron
+        assert np.allclose(dm.surface[act,0],  1e-6), "Segment {} did not move as expected using bare floats".format(actx,acty)
+        assert np.allclose(dm.surface[act,2],  1e-7), "Segment {} did not move as expected using bare floats".format(actx,acty)
+
+    for act in ( 5,2):
+        dm.set_actuator(act, 1*u.nm, 0, 0) # 1 nm
+        assert np.allclose(dm.surface[act,0],  1e-9), "Segment {} did not move as expected in piston using astropy quantities".format(actx,acty)
+
+    for act in ( 1,4):
+        dm.set_actuator(act,  0, 1*u.microradian, 0) # 1 nm
+        assert np.allclose(dm.surface[act,1],  1e-6), "Segment {} did not move as expected in tilt using astropy quantities".format(actx,acty)
+
+    psf_aberrated = osys.calc_psf()
+
+    peak_perf = psf_perf[0].data.max()
+    peak_aberrated = psf_aberrated[0].data.max()
+    assert peak_aberrated < peak_perf, "Adding nonzero WFE did not decrease the Strehl as expected."
+
+    return psf_aberrated, psf_perf, osys
diff --git a/poppy/tests/test_optics.py b/poppy/tests/test_optics.py
index 9ff1a07b..3c32a568 100644
--- a/poppy/tests/test_optics.py
+++ b/poppy/tests/test_optics.py
@@ -395,7 +395,7 @@ def test_WedgeSegmentedCircularAperture(plot=False):
     """ test WedgeSegmentedCircularAperture """
 
     ap_circ = optics.CircularAperture()
-    ap_wedge = optics.WedgeSegmentedCircularAperture(rings=3, nsections=[0, 6, 8])
+    ap_wedge = optics.WedgeSegmentedCircularAperture(rings=3, nsections=[1, 6, 8])
     wave1 = poppy_core.Wavefront(npix=256, diam=2, wavelength=1e-6)  # 10x10 meter square
     wave2 = poppy_core.Wavefront(npix=256, diam=2, wavelength=1e-6)  # 10x10 meter square