Implement WedgeSegmentedDeformableMirror, and related refactoring of …

…WedgeSegmentedCircularAperture

poppy/dms.py

@@ -20,7 +20,8 @@
 if accel_math._NUMEXPR_AVAILABLE:
     import numexpr as ne
 
-__all__ = ['ContinuousDeformableMirror', 'HexSegmentedDeformableMirror', 'CircularSegmentedDeformableMirror']
+__all__ = ['ContinuousDeformableMirror', 'HexSegmentedDeformableMirror', 'CircularSegmentedDeformableMirror',
+           'WedgeSegmentedDeformableMirror']
 
 
 # noinspection PyUnresolvedReferences
@@ -350,7 +351,7 @@ def get_opd(self, wave):
                     pixscale_m = wave.pixelscale.to(u.m/u.pixel).value
                     shift_y_pix = int(np.round(self.shift_y /pixscale_m))
                     interpolated_surface = xp.roll(interpolated_surface, shift_y_pix, axis=0)
-        
+
         # account for DM being reflective (optional, governed by include_factor_of_two parameter)
         coefficient = 2 if self.include_factor_of_two else 1
 
@@ -466,7 +467,7 @@ def _get_surface_via_convolution(self, wave):
                         self._surface_trace_flat[self._act_ind_flat[0] + ix + iy*wave.shape[0]]=(xweight*yweight).flat*target_val
                 except:
                     pass # Ignore any actuators outside the FoV
-        
+
         # Now we can convolve with the influence function to get the full continuous surface.
         influence_rescaled = self._get_rescaled_influence_func(wave.pixelscale)
         dm_surface = _scipy.signal.fftconvolve(self._surface_trace_flat.reshape(wave.shape), influence_rescaled, mode='same')
@@ -685,6 +686,14 @@ def display_influence_fn(self):
 class SegmentedDeformableMirror(ABC):
     """ Abstract class for segmented DMs.
     See below for subclasses for hexagonal and circular apertures.
+
+    Classes to turn into a SegmentedDeformableMirror must implement the
+    poppy.MultiSegmentAperture ABC, including methods for
+       _n_aper_inside_ring
+       _one_aperture
+       _aper_center
+    and attributes:
+       segmentlist
     """
     def __init__(self, rings=1, include_factor_of_two=False):
         self._surface = xp.zeros((self._n_aper_inside_ring(rings + 1), 3))
@@ -833,10 +842,47 @@ class CircularSegmentedDeformableMirror(SegmentedDeformableMirror, optics.MultiC
                 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, rings=1, segment_radius=1.0 * u.m, gap=0.01 * u.m,
                  name='CircSegDM', center=True, include_factor_of_two=False, **kwargs):
         #FIXME ? using grey pixel does not work. something in the geometry module generate a true divide error
         optics.MultiCircularAperture.__init__(self, name=name, rings=rings, segment_radius=segment_radius,
                                               gap=gap, center=center, gray_pixel = False, **kwargs)
         SegmentedDeformableMirror.__init__(self, rings=rings, include_factor_of_two=include_factor_of_two)
+
+
+# note, must inherit first from SegmentedDeformableMirror to get correct method resolution order
+class WedgeSegmentedDeformableMirror(SegmentedDeformableMirror, optics.WedgeSegmentedCircularAperture):
+    """ Circularly segmented DM. Each actuator is controllable in piston, tip, and tilt (and any zernike term)
+
+            Parameters
+            ----------
+            rings, segment_radius, gap, center : various
+                All keywords for defining the segmented aperture geometry are inherited from
+                the MultiCircularAperture class. See that class for details.
+
+             include_factor_of_two : Bool
+                include the factor of two due to reflection in the OPD function (optional, default False).
+                If this is set False (default), actuator commands are interpreted as being in units of
+                desired wavefront error directly; the returned WFE will be directly proportional to the requested
+                values (convolved with the actuator response function etc).
+                If this is set to True, then the actuator commands are interpreted as being in physical surface
+                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
+        optics.WedgeSegmentedCircularAperture.__init__(self, name=name, radius=radius, rings=rings,
+                                              nsections=nsections, gap_radii=gap_radii,
+                                              gap=gap,   **kwargs)
+        SegmentedDeformableMirror.__init__(self, rings=rings, include_factor_of_two=include_factor_of_two)
+
+
+    def _setup_arrays(self, npix, pixelscale, wave=None):
+        # Small tweak to the superclass function, to allow invoking slightly better handling for pixels near
+        # edges of segments. This approach results in the DM segment maps covering the segment gaps better, to
+        # accomodate 'gray' pixels in the transmission map
+        super()._setup_arrays(npix, pixelscale, wave=wave)
+        self._transmission = optics.WedgeSegmentedCircularAperture.get_transmission(self, wave)

poppy/optics.py

@@ -1482,7 +1482,7 @@ class NgonAperture(AnalyticOpticalElement):
         radius to the vertices, meters. Default is 1.
     rotation : float
         Rotation angle to first vertex, in degrees counterclockwise from the +X axis. Default is 0.
-        
+
     TODO: get_transmission() extremely slow when using CuPy, find better solution
     """
 
@@ -1524,7 +1524,7 @@ def get_transmission(self, wave):
 
 class MultiCircularAperture(MultiSegmentAperture):
     """ Defines a circularly segmented aperture in close compact configuration
-    
+
     Parameters
     ----------
     name : string
@@ -1546,9 +1546,9 @@ class MultiCircularAperture(MultiSegmentAperture):
     gray_pixel : bool, optional
         Apply gray pixel approximation to return fractional transmission for
         edge pixels that are only partially within this aperture? default : True
-    
+
     """
-    
+
     @utils.quantity_input(segment_radius=u.meter, gap=u.meter)
     def __init__(self, name="multiCirc", rings=1, segment_radius=1.0, gap=0.01,
                  segmentlist=None, center=True, gray_pixel=True, **kwargs):
@@ -1558,9 +1558,9 @@ def __init__(self, name="multiCirc", rings=1, segment_radius=1.0, gap=0.01,
         super().__init__(name=name, segment_size=segment_diameter,
                          gap=gap, rings=rings, segmentlist=segmentlist, center=center, **kwargs)
         self.pupil_diam = (segment_diameter) * (2 * self.rings + 1) + gap * (2*rings)
-        
+
         self._use_gray_pixel = bool(gray_pixel)
-        
+
     def _one_aperture(self, wave, index, value=1):
         """ Draw one circular aperture into the self.transmission array """
 
@@ -1571,11 +1571,11 @@ def _one_aperture(self, wave, index, value=1):
 
         y -= ceny
         x -= cenx
-        
+
         if self._use_gray_pixel:
             pixscale = wave.pixelscale.to(u.meter/u.pixel).value
             tmpTransmission = geometry.filled_circle_aa(wave.shape, 0, 0, segRadius/pixscale, x/pixscale, y/pixscale)
-            self.transmission += tmpTransmission 
+            self.transmission += tmpTransmission
         else:
             r = _r(x, y)
             del x
@@ -1588,10 +1588,12 @@ def _one_aperture(self, wave, index, value=1):
         return self.transmission
 
 
-class WedgeSegmentedCircularAperture(CircularAperture):
+class WedgeSegmentedCircularAperture(MultiSegmentAperture, CircularAperture):
     @utils.quantity_input(radius=u.meter, gap=u.meter)
     def __init__(self, name=None, radius=1.0 * u.meter,
-                 rings=2, nsections=4, gap_radii=None, gap=0.01 * u.meter, **kwargs):
+                 rings=2, nsections=4, gap_radii=None, gap=0.01 * u.meter,
+                 gray_pixel=False,
+                 rotation=0, **kwargs):
         """ Define a circular aperture made of pie-wedge or keystone shaped segments.
 
         Parameters
@@ -1605,10 +1607,17 @@ def __init__(self, name=None, radius=1.0 * u.meter,
         nsections : int or list of ints
             Number of segments per ring. If one int, same number of segments in each ring.
             Or provide a list of ints to set different numbers per ring.
+            To exclude the center for an on-axis aperture, provide a 0 as the first
+            element of nsections to indicate 0 segments in the first ring.
         gap_radii : quantity length
             Radii from the center for the gaps between rings
         gap : quantity length
             Width of gaps between segments, in both radial and azimuthal directions
+        gray_pixel : bool, optional
+            Apply gray pixel approximation to return fractional transmission for
+            edge pixels that are only partially within this aperture?
+            (Note, currently this gives a warning; disabled by default)
+
         kwargs : other kwargs are passed to CircularAperture
 
         Potential TODO: also have this inherit from MultiSegmentedAperture and subclass
@@ -1620,26 +1629,42 @@ def __init__(self, name=None, radius=1.0 * u.meter,
 
         if name is None:
             name = "Circle of Wedge Sections, radius={}".format(radius)
-        super().__init__(name=name, radius=radius, **kwargs)
+        CircularAperture.__init__(self, name=name, radius=radius, rotation=rotation,
+                                  gray_pixel=gray_pixel, **kwargs)
+
+        # This class inherits from MultiSegmentAperture, but intentionally
+        # does **not** call MultiSegmentAperture.__init__, because some of the
+        # assumptions made there for a regular geometry do not apply. We instead
+        # perform the necessary initialization steps here directly.
+        self.nsections = [nsections, ] * rings if np.isscalar(nsections) else nsections
+        self.segmentlist = np.arange(np.sum(self.nsections))
+        self._include_center = self.nsections[0] != 0
+        if not self._include_center:
+            self.segmentlist = self.segmentlist[1:]  # remove center segment 0
 
         self._default_display_size = 2 * self.radius
+        self.pupil_diam = 2*self.radius
 
         self.rings = rings
-        self.nsections = [nsections, ] * rings if np.isscalar(nsections) else nsections
         self.gap = gap
         self.gap_radii = gap_radii if gap_radii is not None else ((np.arange(
             self.rings) + 1) / self.rings) * self.radius
 
         # determine angles per each section gap
+        # Note, this starts with angle 0 = +Y in the array, and
+        # increases counterclockwise around the aperture.
         self.gap_angles = []
         for iring in range(self.rings):
             nsec = self.nsections[iring]
             self.gap_angles.append(np.arange(nsec) / nsec * 2 * np.pi)
 
     def get_transmission(self, wave):
-        """ Compute the transmission inside/outside of the occulter.
+        """ Compute the transmission inside/outside of the aperture.
+
+        Note, this implementation draws the whole circular aperture then draws in
+        the individual gaps, rather than drawing the aperture one segment at a time.
         """
-        self.transmission = super().get_transmission(wave)
+        self.transmission = CircularAperture.get_transmission(self, wave)
 
         y, x = self.get_coordinates(wave)
         r = np.sqrt(x ** 2 + y ** 2)
@@ -1667,8 +1692,94 @@ def get_transmission(self, wave):
                 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
+
         return self.transmission
 
+    def _n_aper_in_ring(self, n):
+        """ How many hexagons or circles in ring N? """
+        return self.nsections[n] if (n < len(self.nsections)) else 0
+
+    def _one_aperture(self, wave, index, value=1):
+        """ Draw one wedge aperture into the existing self.transmission array
+        """
+
+        #self.transmission = CircularAperture.get_transmission(self, wave)
+
+        y, x = self.get_coordinates(wave)
+        r = np.sqrt(x ** 2 + y ** 2)
+        theta = np.arctan2(y, x)
+        theta[theta<0] += 2*np.pi  # we want angles between 0 and 2 pi, below
+
+        halfgapwidth = self.gap.to_value(u.m) / 2
+
+        # which ring is this?
+        iring = self._aper_in_ring(index)
+        # which segment within this ring?
+        iseg_in_ring = index - self._n_aper_inside_ring(iring)
+
+        # Determine the inner and outer radii of the Nth ring
+        # (Not counting the gap width yet here)
+        r_ring_inner = 0 if iring == 0 else self.gap_radii[iring - 1].to_value(u.m)
+        r_ring_outer = self.radius.to_value(u.m) if iring == self.rings - 1 else self.gap_radii[iring].to_value( u.m)
+        # print(f"{iring}: gap from inner: {r_ring_inner} to outer: {r_ring_outer}")
+
+        gap_angles_this_ring = self.gap_angles[iring]
+        theta_min = gap_angles_this_ring[iseg_in_ring]
+        theta_max = gap_angles_this_ring[iseg_in_ring+1] if (iseg_in_ring < self._n_aper_in_ring(iring)-1) else (gap_angles_this_ring[0] + 2*np.pi)
+
+        self.transmission[(r_ring_inner < r) &
+                          (r < r_ring_outer) &
+                          (theta_min < theta) &
+                          (theta < theta_max)] = value
+        return
+
+        # Draw the azimuthal gap after the ring
+        if iring > 0:
+            # 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
+
+
+    def _aper_center(self, aper_index):
+        """ Center coordinates of a given wedge aperture
+        counting counter clockwise around each ring
+
+        Returns y, x coords
+
+        """
+        # which ring is this?
+        iring = self._aper_in_ring(aper_index)
+        # which segment within this ring?
+        iseg_in_ring = aper_index - self._n_aper_inside_ring(iring)
+
+        # Determine the inner and outer radii of the Nth ring
+        # (Not counting the gap width yet here)
+        r_ring_inner = 0 if iring == 0 else self.gap_radii[iring - 1].to_value(u.m)
+        r_ring_outer = self.radius.to_value(u.m) if iring == self.rings - 1 else self.gap_radii[iring].to_value( u.m)
+        r_center = (r_ring_inner + r_ring_outer) / 2
+
+        gap_angles_this_ring = self.gap_angles[iring]
+        theta_min = gap_angles_this_ring[iseg_in_ring]
+        theta_max = gap_angles_this_ring[iseg_in_ring+1] if (iseg_in_ring < self._n_aper_in_ring(iring)-1) else (gap_angles_this_ring[0] + 2*np.pi)
+        theta_center = (theta_min + theta_max) / 2
+
+        xpos = r_center * np.cos(theta_center)
+        ypos = r_center * np.sin(theta_center)
+
+        return ypos, xpos
+
+
 
 class RectangleAperture(AnalyticOpticalElement):
     """ Defines an ideal rectangular pupil aperture