adding in expmap from sas to change, and putting sources into a list

xga/generate/esass/phot.py

@@ -30,7 +30,11 @@ def evtool_image(sources: Union[BaseSource, NullSource, BaseSample], lo_en: Quan
     stack = False
     execute = True
 
-        # Checking user's choice of energy limit parameters
+    # This function supports passing both individual sources and sets of sources
+    if isinstance(sources, (BaseSource, NullSource)):
+        sources = [sources]
+
+    # Checking user's choice of energy limit parameters
     if not isinstance(lo_en, Quantity) or not isinstance(hi_en, Quantity):
         raise TypeError("The lo_en and hi_en arguments must be astropy quantities in units "
                         "that can be converted to keV.")
@@ -110,3 +114,105 @@ def evtool_image(sources: Union[BaseSource, NullSource, BaseSample], lo_en: Quan
     # I only return num_cores here so it has a reason to be passed to this function, really
     # it could just be picked up in the decorator.
     return sources_cmds, stack, execute, num_cores, sources_types, sources_paths, sources_extras, disable_progress
+
+@esass_call
+def expmap(sources: Union[BaseSource, NullSource, BaseSample], lo_en: Quantity = Quantity(0.2, 'keV'),
+            hi_en: Quantity = Quantity(10, 'keV'), num_cores: int = NUM_CORES, disable_progress: bool = False):
+    """
+    A convenient Python wrapper for the eSASS expmap command.
+    Expmaps will be generated for every observation associated with every source passed to this function.
+    If expmaps in the requested energy band are already associated with the source,
+    they will not be generated again.
+
+    :param BaseSource/NullSource/BaseSample sources: A single source object, or sample of sources.
+    :param Quantity lo_en: The lower energy limit for the expmap, in astropy energy units.
+    :param Quantity hi_en: The upper energy limit for the expmap, in astropy energy units.
+    :param int num_cores: The number of cores to use (if running locally), default is set to
+        90% of available.
+    :param bool disable_progress: Setting this to true will turn off the eSASS generation progress bar.
+    """
+    # I know that a lot of this code is the same as the evselect_image code, but its 1am so please don't
+    #  judge me too much.
+
+    # This function supports passing both individual sources and sets of sources
+    if isinstance(sources, (BaseSource, NullSource)):
+        sources = [sources]
+
+    # Don't do much value checking in this module, but this one is so fundamental that I will do it
+    if lo_en > hi_en:
+        raise ValueError("lo_en cannot be greater than hi_en")
+    else:
+        # Calls a useful little function that takes an astropy energy quantity to the XMM channels
+        # required by SAS commands
+        lo_chan = energy_to_channel(lo_en)
+        hi_chan = energy_to_channel(hi_en)
+
+    # These are crucial, to generate an exposure map one must have a ccf.cif calibration file, and a reference
+    # image. If they do not already exist, these commands should generate them.
+    cifbuild(sources, disable_progress=disable_progress, num_cores=num_cores)
+    sources = evselect_image(sources, lo_en, hi_en)
+    # This is necessary because the decorator will reduce a one element list of source objects to a single
+    # source object. Useful for the user, not so much here where the code expects an iterable.
+    if not isinstance(sources, (list, BaseSample)):
+        sources = [sources]
+
+    # These lists are to contain the lists of commands/paths/etc for each of the individual sources passed
+    # to this function
+    sources_cmds = []
+    sources_paths = []
+    sources_extras = []
+    sources_types = []
+    for source in sources:
+        cmds = []
+        final_paths = []
+        extra_info = []
+        # Check which event lists are associated with each individual source
+        for pack in source.get_products("events", just_obj=False):
+            obs_id = pack[0]
+            inst = pack[1]
+
+            if not os.path.exists(OUTPUT + obs_id):
+                os.mkdir(OUTPUT + obs_id)
+
+            en_id = "bound_{l}-{u}".format(l=lo_en.value, u=hi_en.value)
+            exists = [match for match in source.get_products("expmap", obs_id, inst, just_obj=False)
+                      if en_id in match]
+            if len(exists) == 1 and exists[0][-1].usable:
+                continue
+            # Generating an exposure map requires a reference image.
+            ref_im = [match for match in source.get_products("image", obs_id, inst, just_obj=False)
+                      if en_id in match][0][-1]
+            # It also requires an attitude file
+            att = source.get_att_file(obs_id)
+            # Set up the paths and names of files
+            evt_list = pack[-1]
+            dest_dir = OUTPUT + "{o}/{i}_{l}-{u}_{n}_temp/".format(o=obs_id, i=inst, l=lo_en.value, u=hi_en.value,
+                                                                   n=source.name)
+            exp_map = "{o}_{i}_{l}-{u}keVexpmap.fits".format(o=obs_id, i=inst, l=lo_en.value, u=hi_en.value)
+
+            # If something got interrupted and the temp directory still exists, this will remove it
+            if os.path.exists(dest_dir):
+                rmtree(dest_dir)
+
+            os.makedirs(dest_dir)
+            cmds.append("cd {d}; cp ../ccf.cif .; export SAS_CCF={ccf}; eexpmap eventset={e} "
+                        "imageset={im} expimageset={eim} withdetcoords=no withvignetting=yes "
+                        "attitudeset={att} pimin={l} pimax={u}; mv * ../; cd ..; "
+                        "rm -r {d}".format(e=evt_list.path, im=ref_im.path, eim=exp_map, att=att, l=lo_chan,
+                                           u=hi_chan, d=dest_dir, ccf=dest_dir + "ccf.cif"))
+
+            # This is the products final resting place, if it exists at the end of this command
+            final_paths.append(os.path.join(OUTPUT, obs_id, exp_map))
+            extra_info.append({"lo_en": lo_en, "hi_en": hi_en, "obs_id": obs_id, "instrument": inst})
+        sources_cmds.append(np.array(cmds))
+        sources_paths.append(np.array(final_paths))
+        # This contains any other information that will be needed to instantiate the class
+        # once the SAS cmd has run
+        sources_extras.append(np.array(extra_info))
+        sources_types.append(np.full(sources_cmds[-1].shape, fill_value="expmap"))
+
+    stack = False  # This tells the sas_call routine that this command won't be part of a stack
+    execute = True  # This should be executed immediately
+    # I only return num_cores here so it has a reason to be passed to this function, really
+    # it could just be picked up in the decorator.
+    return sources_cmds, stack, execute, num_cores, sources_types, sources_paths, sources_extras, disable_progress