add pie wedge segmented aperture

poppy/optics.py

@@ -24,7 +24,7 @@
            'BandLimitedCoron', 'BandLimitedCoronagraph', 'IdealFQPM', 'CircularPhaseMask', 'RectangularFieldStop', 'SquareFieldStop',
            'AnnularFieldStop', 'HexagonFieldStop',
            'CircularOcculter', 'BarOcculter', 'FQPM_FFT_aligner', 'CircularAperture',
-           'HexagonAperture', 'MultiHexagonAperture', 'NgonAperture', 'MultiCircularAperture', 'RectangleAperture',
+           'HexagonAperture', 'MultiHexagonAperture', 'NgonAperture', 'MultiCircularAperture', 'WedgeSegmentedCircularAperture', 'RectangleAperture',
            'SquareAperture', 'SecondaryObscuration', 'LetterFAperture', 'AsymmetricSecondaryObscuration',
            'ThinLens',  'GaussianAperture', 'KnifeEdge', 'TiltOpticalPathDifference', 'CompoundAnalyticOptic', 'fixed_sampling_optic']
 
@@ -1588,6 +1588,88 @@ def _one_aperture(self, wave, index, value=1):
         return self.transmission
 
 
+class WedgeSegmentedCircularAperture(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):
+        """ Define a circular aperture made of pie-wedge or keystone shaped segments.
+
+        Parameters
+        ----------
+        name : string
+            Descriptive name
+        radius : float
+            Radius of the pupil, in meters.
+        rings : int
+            Number of rings of segments
+        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.
+        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
+        kwargs : other kwargs are passed to CircularAperture
+
+        Potential TODO: also have this inherit from MultiSegmentedAperture and subclass
+        some of those functions as appropriate. Consider refactoring from gap_radii to instead
+        provide the widths of each segment. Add option for including the center segment or having
+        a missing one in the middle for on-axis apertures. Use grayscale approximation for rasterizing
+        the circular gaps between the rings.
+        """
+
+        if name is None:
+            name = "Circle of Wedge Sections, radius={}".format(radius)
+        super().__init__(name=name, radius=radius, **kwargs)
+
+        self._default_display_size = 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
+        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.
+        """
+        self.transmission = super().get_transmission(wave)
+
+        y, x = self.get_coordinates(wave)
+        r = np.sqrt(x ** 2 + y ** 2)
+
+        halfgapwidth = self.gap.to_value(u.m) / 2
+        for iring in range(self.rings):
+
+            # Draw the radial gaps around the azimuth in the Nth ring
+            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}")
+
+            # 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
+
+        return self.transmission
+
+
 class RectangleAperture(AnalyticOpticalElement):
     """ Defines an ideal rectangular pupil aperture
 

poppy/tests/test_optics.py

@@ -5,6 +5,7 @@
 
 from .. import poppy_core
 from .. import optics
+from .. import utils
 
 from .. import accel_math
 import numpy as np
@@ -390,6 +391,24 @@ def test_NgonAperture(display=False):
         optic.display()
 
 
+def test_WedgeSegmentedCircularAperture(plot=False):
+    """ test WedgeSegmentedCircularAperture """
+
+    ap_circ = optics.CircularAperture()
+    ap_wedge = optics.WedgeSegmentedCircularAperture(rings=3, nsections=[0, 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
+
+    wave1 *= ap_circ
+    wave2 *= ap_wedge
+
+    assert wave1.total_intensity * 0.95 < wave2.total_intensity < wave1.total_intensity, 'wedge segmented circle should have slightly less clear area than monolith circle'
+
+    if plot:
+        fig, axes = plt.subplots(figsize=(16, 9), ncols=2)
+        wave1.display(ax=axes[0])
+        wave2.display(ax=axes[1])
+
 
 def test_ObscuredCircularAperture_Airy(display=False):
     """ Compare analytic 2d Airy function with the results of a POPPY

poppy/tests/test_utils.py

@@ -140,6 +140,10 @@ def test_radial_profile(plot=False):
     assert prof_stds.shape == prof.shape, "Radial profile and stddev array output sizes should match"
     assert prof_mads.shape == prof.shape, "Radial profile and median absolute deviation array output sizes should match"
 
+    # Test computing some arbitrary function
+    rad4, prof_max = poppy.radial_profile(psf, custom_function=np.max)
+    # compare, but ignore the 0th bin and last bin that have no pixels within that bin etc
+    assert np.all(prof_max[1:-1] >= prof[1:-1]), "Radial profile using max function should be >= profile using mean"
 
 def test_radial_profile_of_offset_source():
     """Test that we can compute the radial profile for a source slightly outside the FOV,

poppy/utils.py

@@ -540,7 +540,8 @@ def display_profiles(hdulist_or_filename=None, ext=0, overplot=False, title=None
 
 
 def radial_profile(hdulist_or_filename=None, ext=0, ee=False, center=None, stddev=False, mad=False,
-                   binsize=None, maxradius=None, normalize='None', pa_range=None, slice=0):
+                   binsize=None, maxradius=None, normalize='None', pa_range=None, slice=0,
+                   custom_function=None):
     """ Compute a radial profile of the image.
 
     This computes a discrete radial profile evaluated on the provided binsize. For a version
@@ -579,6 +580,10 @@ def radial_profile(hdulist_or_filename=None, ext=0, ee=False, center=None, stdde
     slice: integer, optional
         Slice into a datacube, for use on cubes computed by calc_datacube. Default 0 if a
         cube is provided with no slice specified.
+    custom_function : function
+        To evaluate an arbitrary function in each radial bin, provide some callable function that
+        takes a list of pixels and returns one float. For instance custome_function=np.min to compute
+        the minimum in each radial bin.
 
     Returns
     --------
@@ -696,6 +701,21 @@ def radial_profile(hdulist_or_filename=None, ext=0, ee=False, center=None, stdde
                 wg = np.where((r_pix >= (radius - binsize / 2)) & (r_pix < (radius + binsize / 2)))
             mads[i] = np.nanmedian(np.absolute(image[wg]-np.nanmedian(image[wg])))
         return rr, mads
+    elif custom_function is not None:
+        # Compute some custom function in each radial bin
+        results = np.zeros_like(radialprofile2)
+        r_pix = r * binsize
+        for i, radius in enumerate(rr):
+            if i == 0:
+                wg = np.where(r < radius + binsize / 2)
+            else:
+                wg = np.where((r_pix >= (radius - binsize / 2)) & (r_pix < (radius + binsize / 2)))
+            if len(wg[0])==0: # Zero elements in this bin
+                results[i] = np.nan
+            else:
+                results[i] = custom_function(image[wg])
+        return rr, results
+
 
     elif not ee:
         # (Default behavior) Compute average in each radial bin