DRAFT : Add MOS functionality to ScopeSim

scopesim/effects/apertures.py

@@ -7,11 +7,11 @@
 from matplotlib.path import Path as MPLPath  # rename to avoid conflict with pathlib
 from astropy.io import fits
 from astropy import units as u
-from astropy.table import Table
+from astropy.table import Table, Column
 
 from .effects import Effect
 from ..optics import image_plane_utils as imp_utils
-from ..base_classes import FOVSetupBase
+from ..base_classes import FOVSetupBase, FieldOfViewBase
 
 from ..utils import (quantify, quantity_from_table, from_currsys, check_keys,
                      figure_factory, get_logger)
@@ -70,6 +70,10 @@ class ApertureMask(Effect):
 
     Other Parameters
     ----------------
+    extend_fov_beyond_slit : float
+        [arcsec] Default 0. Increases the on-sky size of the FOV by this
+        amount in all directions when extracting flux from the Source object.
+
     pixel_scale : float
         [arcsec] Defaults to ``"!INST.pixel_scale"`` from the config
 
@@ -89,6 +93,7 @@ def __init__(self, **kwargs):
 
         super().__init__(**kwargs)
         params = {
+            "extend_fov_beyond_slit": 0,
             "pixel_scale": "!INST.pixel_scale",
             "no_mask": True,
             "angle": 0,
@@ -118,9 +123,13 @@ def apply_to(self, obj, **kwargs):
                                     u.arcsec).to(u.arcsec).value
             y = quantity_from_table("y", self.table,
                                     u.arcsec).to(u.arcsec).value
-            obj.shrink(["x", "y"], ([min(x), max(x)], [min(y), max(y)]))
+            dr = self.meta["extend_fov_beyond_slit"]
+            x_min, x_max = min(x) - dr, max(x) + dr
+            y_min, y_max = min(y) - dr, max(y) + dr
+            obj.shrink(["x", "y"], ([x_min, x_max], [y_min, y_max]))
 
             # ..todo: HUGE HACK - Get rid of this!
+            # Assume the slit coord 'xi' is along the x axis
             for vol in obj.volumes:
                 vol["meta"]["xi_min"] = min(x) * u.arcsec
                 vol["meta"]["xi_max"] = max(x) * u.arcsec
@@ -232,6 +241,25 @@ class ApertureList(Effect):
 
     Examples
     --------
+    Defining the apertures of an image slicer IFU::
+
+        - name: image_slicer
+          description: collection of slits corresponding to the image slicer apertures
+          class: ApertureList
+          kwargs:
+            filename: "INS_ifu_apertures.dat"
+            fov_for_each_aperture : True
+            extend_fov_beyond_slit : 2            # [arcsec]
+
+    Where the file ``INS_ifu_apertures.dat`` includes this table::
+
+        id  left    right   bottom  top     angle   conserve_image  shape
+        0   -14     14      -5      -3      0       True            rect
+        1   -14     14      -3      -1      0       True            rect
+        2   -14     14      -1      1       0       True            rect
+        3   -14     14      1       3       0       True            rect
+        4   -14     14      3       5       0       True            rect
+
 
     File format
     -----------
@@ -261,7 +289,7 @@ class ApertureList(Effect):
     .. note:: ``shape`` values ``"rect"`` and ``4`` do not produce equal results
 
        Both use 4 equidistant points around an ellipse constrained by
-       [``left``, ..., etc]. However ``"rect"`` aims to fill the contraining
+       [``left``, ..., etc]. However ``"rect"`` aims to fill the constraining
        area, while ``4`` simply uses 4 points on the ellipse.
        Consequently, ``4`` results in a diamond shaped mask covering only
        half of the constraining area filled by ``"rect"``.
@@ -272,12 +300,14 @@ class ApertureList(Effect):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         params = {
+            "fov_for_each_aperture": True,
+            "extend_fov_beyond_slit": 0,
             "pixel_scale": "!INST.pixel_scale",
             "n_round_corners": 32,        # number of corners use to estimate ellipse
             "no_mask": False,             # .. todo:: is this necessary when we have conserve_image?
             "report_plot_include": True,
             "report_table_include": True,
-            "report_table_rounding": 4,
+            "report_table_rounding": 4
         }
         self.meta["z_order"] = [81, 281]
         self.meta.update(params)
@@ -291,23 +321,49 @@ def __init__(self, **kwargs):
     def apply_to(self, obj, **kwargs):
         """See parent docstring."""
         if isinstance(obj, FOVSetupBase):
-            new_vols = []
-            for row in self.table:
-                vols = obj.extract(["x", "y"], ([row["left"], row["right"]],
-                                                [row["bottom"], row["top"]]))
-                for vol in vols:
-                    vol["meta"]["aperture_id"] = row["id"]
-
-                    # ..todo: HUGE HACK - Get rid of this!
-                    vol["meta"]["xi_min"] = row["left"] * u.arcsec
-                    vol["meta"]["xi_max"] = row["right"] * u.arcsec
-                    vol["conserve_image"] = row["conserve_image"]
-                    vol["shape"] = row["shape"]
-                    vol["angle"] = row["angle"]
-
-                new_vols += vols
-
-            obj.volumes = new_vols
+            dr = self.meta["extend_fov_beyond_slit"]
+            if self.meta["fov_for_each_aperture"] is True:
+                new_vols = []
+                for row in self.table:
+                    x_min, x_max = row["left"] - dr, row["right"] + dr
+                    y_min, y_max = row["bottom"] - dr, row["top"] + dr
+                    vols = obj.extract(["x", "y"], ([x_min, x_max],
+                                                    [y_min, y_max]))
+                    for vol in vols:
+                        vol["meta"]["aperture_id"] = row["id"]
+                        vol["meta"]["extend_fov_beyond_slit"] = dr
+                        vol["meta"]["sub_apertures"] = Table(rows=row)
+
+                        # ..todo: HUGE HACK - Get rid of this!
+                        # Assume the slit coord 'xi' is along the x axis
+                        vol["meta"]["xi_min"] = row["left"] * u.arcsec
+                        vol["meta"]["xi_max"] = row["right"] * u.arcsec
+                        vol["conserve_image"] = row["conserve_image"]
+                        vol["shape"] = row["shape"]
+                        vol["angle"] = row["angle"]
+
+                    new_vols += vols
+
+                obj.volumes = new_vols
+
+            else:
+                x_min = min(self.table["left"]) - dr
+                x_max = max(self.table["right"]) + dr
+                y_min = min(self.table["bottom"]) - dr
+                y_max = max(self.table["top"]) + dr
+
+                obj.shrink(["x", "y"], ([x_min, x_max], [y_min, y_max]))
+                vol = obj.volumes[0]
+                vol["meta"]["aperture_id"] = list(self.table["id"])
+                vol["meta"]["sub_apertures"] = self.table
+
+                # ..todo: HUGE HACK - Get rid of this!
+                # Assume the slit coord 'xi' is along the x axis
+                vol["meta"]["xi_min"] = x_min * u.arcsec
+                vol["meta"]["xi_max"] = x_max * u.arcsec
+                vol["conserve_image"] = True
+                vol["shape"] = "rect"
+                vol["angle"] = 0
 
         return obj
 
@@ -373,8 +429,63 @@ def __add__(self, other):
             raise ValueError("Secondary argument not of type ApertureList: "
                              f"{type(other) = }")
 
-    # def __getitem__(self, item):
-    #     return self.get_apertures(item)[0]
+
+class FibreApertureList(ApertureList):
+    """
+    A subclass of ApertureList for fibre apertures
+
+    File format
+    -----------
+    Much like an ApertureList, a FibreApertureList can be initialised by either
+    of the three standard DataContainer methods.
+    The easiest is however to make an ASCII file with the following columns::
+
+        id  x       y       r       angle  conserve_image  shape
+            arcsec  arcsec  arcsec  deg
+        int float   float   float   float  bool            str/int
+
+    where:
+    - x,y : centres of the fibres
+    - r : radius (NOT diameter) of the fibre head.
+    - angle : rotation of the first vertex from x=0
+    - conserve_image : flag for maintaining spatial information of flux
+    - shape : string or int. See ApertureList for allowed strings
+
+    .. note:: FibreApertureList does not support Elliptical fibres.
+
+
+    """
+    def __init__(self, **kwargs):
+        # Initialise with the super-super class, i.e. Effect, not ApertureList
+        super(ApertureList, self).__init__(**kwargs)
+        params = {"fov_for_each_aperture": True,
+                  "extend_fov_beyond_slit": 0,
+                  "pixel_scale": "!INST.pixel_scale",
+                  "n_round_corners": 32,        # number of corners use to estimate ellipse
+                  "no_mask": False,             # .. todo:: is this necessary when we have conserve_image?
+                  "report_plot_include": True,
+                  "report_table_include": True,
+                  "report_table_rounding": 4}
+        self.meta["z_order"] = [81, 281]
+        self.meta.update(params)
+        self.meta.update(kwargs)
+
+        if self.table is not None:
+            required_keys = ["id", "x", "y", "r", "angle",
+                             "conserve_image", "shape"]
+            check_keys(self.table.colnames, required_keys)
+
+        left = self.table["x"] - self.table["r"]
+        right = self.table["x"] + self.table["r"]
+        bottom = self.table["y"] - self.table["r"]
+        top = self.table["y"] + self.table["r"]
+        left.name = "left"
+        right.name = "right"
+        bottom.name = "bottom"
+        top.name = "top"
+        self.table.add_columns([left, right, bottom, top])
+
+
 
 
 class SlitWheel(Effect):
@@ -486,6 +597,11 @@ def current_slit(self):
             return False
         return self.slits[currslit]
 
+    @property
+    def display_name(self):
+        return f'{self.meta["name"]} : ' \
+               f'[{from_currsys(self.meta["current_slit"])}]'
+
     def __getattr__(self, item):
         return getattr(self.current_slit, item)
 

scopesim/effects/psfs.py

@@ -1,6 +1,6 @@
 import numpy as np
 from scipy.signal import convolve
-from scipy.interpolate import RectBivariateSpline
+from scipy.interpolate import RectBivariateSpline, interp1d
 
 from astropy import units as u
 from astropy.io import fits
@@ -189,6 +189,98 @@ def get_kernel(self, obj):
         return self.kernel.astype(float)
 
 
+class RadialProfilePSF(AnalyticalPSF):
+    """
+    Creates a wavelength independent kernel image
+    """
+    def __init__(self, **kwargs):
+        params = {"unit": "pixel",
+                  "pixel_scale": "!INST.pixel_scale",
+                  "plate_scale": "!INST.plate_scale"}
+        params.update(kwargs)
+        super().__init__(**params)
+        self.meta["z_order"] = [244, 744]
+
+        self.convolution_classes = ImagePlaneBase
+
+        self.required_keys = ["r", "z", "unit"]
+        if self.meta["unit"] == "mm":
+            self.required_keys += ["pixel_scale", "plate_scale"]
+        check_keys(self.meta, self.required_keys, action="error")
+
+        self.kernel = None
+
+    def get_kernel(self, obj):
+        if self.kernel is None:
+            dr = 1
+            if self.meta["unit"] == "mm":
+                dr = from_currsys(self.meta["pixel_scale"], self.cmds) / \
+                     from_currsys(self.meta["plate_scale"], self.cmds)
+
+            rpix = np.array(self.meta["r"]) / dr
+            z = np.array(self.meta["z"])
+            fn = interp1d(rpix, z, "linear", fill_value=0, bounds_error=False)
+            rmax = np.ceil(max(rpix))
+            xx, yy = np.mgrid[-rmax:rmax+1, -rmax:rmax+1]
+            rr = (xx**2 + yy**2)**0.5
+            kernel = fn(rr)
+            kernel[kernel<0] = 0
+            kernel /= np.sum(kernel)
+            self.kernel = kernel
+
+        return self.kernel.astype(float)
+
+
+class MosaicBundleTracePSF(RadialProfilePSF):
+    """
+    A hack to get the mosaic bundle traces onto the detector, assuming full ADC
+
+    Basically it creates a 5x5 RadialProfilePSF (defined by r and z) and then
+    copies and pastes the kernel side-by-side for each fibre in the mosaic
+    bundle. Each fibre PSF kernel can be weighted by the filling factor of the
+    PSF in said fibre.
+
+    The main problem here is that it doesn't take into account any wavelength
+    dependent effects on the PSF. This could be an issue down the line.
+
+    """
+
+    def __init__(self, **kwargs):
+        params = {"unit": "pixel",
+                  "pixel_scale": "!INST.pixel_scale",
+                  "plate_scale": "!INST.plate_scale",
+                  "r": [0, 1, 2],               # distance from centre of trace in [unit]
+                  "z": [0.6, 0.2, 0],           # relative flux of PSF for values in "r"
+                  "trace_spacing": 2,           # [unit]
+                  "trace_flux_weights": [0.7] + [0.05] * 6      # Number of traces taken from the length of this list
+                  }
+        params.update(kwargs)
+        super().__init__(**params)
+        self.meta["z_order"] = [244, 744]
+
+        self.convolution_classes = ImagePlaneBase
+        self.kernel = None
+        self.single_kernel = None
+
+    def get_kernel(self, obj):
+        self.single_kernel = super().get_kernel(obj)
+        kernel_list = [self.single_kernel * weight
+                       for weight in self.meta["trace_flux_weights"]]
+
+        dr = 1
+        if self.meta["unit"] == "mm":
+            dr = from_currsys(self.meta["pixel_scale"], self.cmds) / \
+                 from_currsys(self.meta["plate_scale"], self.cmds)
+        pad_arr = np.zeros([self.single_kernel.shape[0],
+                            int(dr * self.meta["trace_spacing"])])
+        pad_list = [pad_arr] * len(kernel_list)
+
+        zipped_list = [j for i in zip(kernel_list, pad_list) for j in i][:-1]     # this zips together the psf and pad lists
+        self.kernel = np.concatenate(zipped_list, axis=1)
+
+        return self.kernel.astype(float)
+
+
 class NonCommonPathAberration(AnalyticalPSF):
     """
     TBA.