Merge pull request #460 from desihub/dstn/speedup-hardware

Speed up the Hardware() constructor, and other optimizations

doc/changes.rst

@@ -11,12 +11,14 @@ fiberassign change log
 * Update ReadTheDocs configuration (PR `#459`_).
 * Add Assignment.check_avail_collisions(tile, all_matches=False) to check potential assignments for collisions (PR `#454`_).
 * uses the focalplane radius from DESIMODEL instead of the hardcoded version in fiberassign (PR `#453`_).
+* speeds up the Hardware constructor's search for neighboring positioners by pre-sorting in one dimension first (PR `#460`_).
 
 .. _`#453`: https://github.com/desihub/fiberassign/pull/453
 .. _`#454`: https://github.com/desihub/fiberassign/pull/454
 .. _`#459`: https://github.com/desihub/fiberassign/pull/459
 .. _`#457`: https://github.com/desihub/fiberassign/pull/457
 .. _`#458`: https://github.com/desihub/fiberassign/pull/458
+.. _`#460`: https://github.com/desihub/fiberassign/pull/460
 
 5.7.1 (2023-02-06)
 ------------------

py/fiberassign/hardware.py

@@ -14,6 +14,7 @@
 import numpy as np
 
 from scipy.interpolate import interp1d
+import astropy.table
 
 import desimodel.io as dmio
 
@@ -31,6 +32,18 @@
     Shape,
 )
 
+def get_default_exclusion_margins():
+    '''
+    We add an additional margin around the exclusion zones in the hardware descriptions.
+    This function returns the defaults used by the fiberassign scripts.
+
+    Returns:
+    margins (dict): with keys "pos", "gfa" and "petal", to floating-point margins in millimeters.
+    '''
+    return dict(pos=0.05,
+                petal=0.4,
+                gfa=0.4)
+
 def expand_closed_curve(xx, yy, margin):
     '''
     For the --margin-{pos,petal,gfa} options, we can add a buffer zone
@@ -121,9 +134,74 @@ def expand_closed_curve(xx, yy, margin):
 
     return ex, ey
 
+# Global cache of exclusions.  (exclusions are geometric regions describing the keep-out
+# areas around positioners, GFAs, and petal edges)
+_cache_shape_to_excl = dict()
+
+def get_exclusions(exclude, exclname, margins, local_cache=None):
+    '''
+    Fetches an exclusion region from a data file, applying additional margins as required.
+
+    A "local_cache", valid for a single "exclude" file, can be passed in.
+
+    Args:
+    * exclude: the exclusions file, eg from desimodel.io.load_focalplane.
+    * exclname: the string name to look up
+    * margins: dict of additional margins to apply
+    * local_cache: dict
+    '''
+    global _cache_shape_to_excl
+
+    if local_cache is not None:
+        # check the local cache, where names are unique
+        e = local_cache.get(exclname)
+        if e is not None:
+            return e
 
-def load_hardware(focalplane=None, rundate=None,
-                  add_margins={}):
+    shp = exclude[exclname]
+
+    # NOTE that different exclusions files can assign
+    # different values to the same key!  ie, code 0000fbb5303ebbc2 appears in
+    #   desi-exclusion_2021-03-17T23:20:01.json and
+    #   desi-exclusion_2021-06-25T22:38:59+00:00.json
+    # with totally different values.
+    # We therefore cannot use the "exclname" as the (global) cache key.
+    #
+    # I tried caching the "expand_closed_curve" computation using a hash of the
+    # shape itself as the cache key, but that turned out not to help -- computing
+    # the key was too expensive?
+    #
+    # For further speedups, I suspect that moving expand_closed_curve to C++ would be
+    # very effective.
+
+    rtn = dict()
+    for obj in shp.keys():
+        cr = list()
+        for crc in shp[obj]["circles"]:
+            cr.append(Circle(crc[0], crc[1]))
+        sg = list()
+        for sgm in shp[obj]["segments"]:
+            if obj in margins:
+                # Create a hashable version of the line-segment list
+                key = (obj, tuple(tuple(xy) for xy in sgm))
+                cached = _cache_shape_to_excl.get(key)
+                if cached is not None:
+                    sgm = cached
+                else:
+                    sx = [x for x,y in sgm]
+                    sy = [y for x,y in sgm]
+                    ex,ey = expand_closed_curve(sx, sy, margins[obj])
+                    sgm = list(zip(ex, ey))
+                    _cache_shape_to_excl[key] = sgm
+            sg.append(Segments(sgm))
+        fshp = Shape((0.0, 0.0), cr, sg)
+        rtn[obj] = fshp
+    if local_cache is not None:
+        local_cache[exclname] = rtn
+    return rtn
+
+def load_hardware(focalplane=None, rundate=None, add_margins={},
+                  get_time_range=False):
     """Create a hardware class representing properties of the telescope.
 
     Args:
@@ -132,10 +210,45 @@ def load_hardware(focalplane=None, rundate=None,
             desimodel.io.load_focalplane()
         rundate (str):  ISO 8601 format time stamp as a string in the
             format YYYY-MM-DDTHH:MM:SS+-zz:zz.  If None, uses current time.
+        add_margins (dict): additional margins to add around positioners, GFAs,
+            and petals.  Dict with keys "pos", "gfa", "petal" and values of
+            millimeters.
+        get_time_range (bool): if True, return (hw, time_lo, time_hi), where
+            time_lo and time_hi are datetime objects corresponding to the first and
+            last dates when the hardware was in this state.
 
     Returns:
+        if get_time_range is True: (hardware, time_lo, time_hi)
+        else: hardware
+
         (Hardware):  The hardware object.
+    """
+    args = load_hardware_args(focalplane=focalplane, rundate=rundate, add_margins=add_margins)
+    args, time_lo, time_hi = args
+    hw = Hardware(*args)
+    if get_time_range:
+        return hw, time_lo, time_hi
+    return hw
+
+def load_hardware_args(focalplane=None, rundate=None, add_margins={}):
+    """Reads the arguments needed to create a Hardware object representing the
+    properties of the instrument at a given date.  Also returns the range of
+    dates when this hardware configuration is valid.
+
+    Args:
+        focalplane (tuple):  Override the focalplane model.  If not None, this
+            should be a tuple of the same data types returned by
+            desimodel.io.load_focalplane()
+        rundate (str):  ISO 8601 format time stamp as a string in the
+            format YYYY-MM-DDTHH:MM:SS+-zz:zz.  If None, uses current time.
+        add_margins (dict): additional margins to add around positioners, GFAs,
+            and petals.  Dict with keys "pos", "gfa", "petal" and values of
+            millimeters.
 
+    Returns:
+        args (tuple): The arguments to be passed to the Hardware constructor
+        time_lo (datetime): The earliest date when this hardware configuration is valid
+        time_hi (datetime): The latest date when this hardware configuration is valid
     """
     log = Logger.get()
 
@@ -162,8 +275,9 @@ def load_hardware(focalplane=None, rundate=None,
     exclude = None
     state = None
     create_time = "UNKNOWN"
+    time_lo = time_hi = None
     if focalplane is None:
-        fp, exclude, state, create_time = dmio.load_focalplane(runtime)
+        fp, exclude, state, create_time, time_lo, time_hi = dmio.load_focalplane(runtime, get_time_range=True)
     else:
         fp, exclude, state = focalplane
 
@@ -232,44 +346,14 @@ def load_hardware(focalplane=None, rundate=None,
     if 'petal' in add_margins:
         margins['petal'] = -add_margins['petal']
 
-    # Convert the exclusion polygons into shapes (as required)
-    excl = dict()
-    # cache expanded polygons (because many of the polygons are actually duplicates)
-    expanded = {}
-    def get_exclusions(exclname):
-        e = excl.get(exclname)
-        if e is not None:
-            return e
-        shp = exclude[exclname]
-        excl[exclname] = dict()
-        for obj in shp.keys():
-            cr = list()
-            for crc in shp[obj]["circles"]:
-                cr.append(Circle(crc[0], crc[1]))
-            sg = list()
-            for sgm in shp[obj]["segments"]:
-                if obj in margins:
-                    key = (obj, tuple(tuple(xy) for xy in sgm))
-                    if key in expanded:
-                        sgm = expanded[key]
-                    else:
-                        sx = [x for x,y in sgm]
-                        sy = [y for x,y in sgm]
-                        ex,ey = expand_closed_curve(sx, sy, margins[obj])
-                        sgm = list(zip(ex, ey))
-                        expanded[key] = sgm
-                sg.append(Segments(sgm))
-            fshp = Shape((0.0, 0.0), cr, sg)
-            excl[exclname][obj] = fshp
-        return excl[exclname]
 
     # For each positioner, select the exclusion polynomials.
     positioners = dict()
-
+    excl_cache = dict()
     for loc in locations:
         exclname = state["EXCLUSION"][loc_to_state[loc]]
         positioners[loc] = dict()
-        posexcl = get_exclusions(exclname)
+        posexcl = get_exclusions(exclude, exclname, margins, local_cache=excl_cache)
         positioners[loc]["theta"] = Shape(posexcl["theta"])
         positioners[loc]["phi"] = Shape(posexcl["phi"])
         if "gfa" in posexcl:
@@ -280,11 +364,13 @@ def get_exclusions(exclname):
             positioners[loc]["petal"] = Shape(posexcl["petal"])
         else:
             positioners[loc]["petal"] = Shape()
+    del excl_cache
 
-    hw = None
     if "MIN_P" in state.colnames:
         # This is a new-format focalplane model (after desimodel PR #143)
-        hw = Hardware(
+        Istate = np.array([loc_to_state[x] for x in locations])
+        Ifp    = np.array([loc_to_fp   [x] for x in locations])
+        args = (
             runtimestr,
             locations,
             fp["PETAL"][keep_rows],
@@ -295,23 +381,23 @@ def get_exclusions(exclname):
             device_type,
             fp["OFFSET_X"][keep_rows],
             fp["OFFSET_Y"][keep_rows],
-            np.array([state["STATE"][loc_to_state[x]] for x in locations]),
-            np.array([fp["OFFSET_T"][loc_to_fp[x]] for x in locations]),
-            np.array([state["MIN_T"][loc_to_state[x]] for x in locations]),
-            np.array([state["MAX_T"][loc_to_state[x]] for x in locations]),
-            np.array([state["POS_T"][loc_to_state[x]] for x in locations]),
-            np.array([fp["LENGTH_R1"][loc_to_fp[x]] for x in locations]),
-            np.array([fp["OFFSET_P"][loc_to_fp[x]] for x in locations]),
-            np.array([state["MIN_P"][loc_to_state[x]] for x in locations]),
-            np.array([state["MAX_P"][loc_to_state[x]] for x in locations]),
-            np.array([state["POS_P"][loc_to_state[x]] for x in locations]),
-            np.array([fp["LENGTH_R2"][loc_to_fp[x]] for x in locations]),
+            state['STATE'][Istate],
+            fp['OFFSET_T'][Ifp],
+            state['MIN_T'][Istate],
+            state['MAX_T'][Istate],
+            state['POS_T'][Istate],
+            fp['LENGTH_R1'][Ifp],
+            fp['OFFSET_P'][Ifp],
+            state['MIN_P'][Istate],
+            state['MAX_P'][Istate],
+            state['POS_P'][Istate],
+            fp['LENGTH_R2'][Ifp],
             fine_radius,
             fine_theta,
             fine_arc,
             [positioners[x]["theta"] for x in locations],
-            [positioners[x]["phi"] for x in locations],
-            [positioners[x]["gfa"] for x in locations],
+            [positioners[x]["phi"]   for x in locations],
+            [positioners[x]["gfa"]   for x in locations],
             [positioners[x]["petal"] for x in locations],
             add_margins
         )
@@ -330,7 +416,8 @@ def get_exclusions(exclname):
                 pp = 180.0
             fake_pos_p[ilid] = pp
             fake_pos_t[ilid] = pt
-        hw = Hardware(
+
+        args = (
             runtimestr,
             locations,
             fp["PETAL"][keep_rows],
@@ -361,7 +448,14 @@ def get_exclusions(exclname):
             [positioners[x]["petal"] for x in locations],
             add_margins
         )
-    return hw
+
+    aa = []
+    for a in args:
+        if type(a) == astropy.table.Column:
+            a = a.data
+        aa.append(a)
+
+    return aa, time_lo, time_hi
 
 def radec2xy(hw, tile_ra, tile_dec, tile_obstime, tile_obstheta, tile_obsha,
              ra, dec, use_cs5, threads=0):

py/fiberassign/scripts/assign.py

@@ -16,7 +16,7 @@
 
 from ..utils import GlobalTimers, Logger
 
-from ..hardware import load_hardware
+from ..hardware import load_hardware, get_default_exclusion_margins
 
 from ..tiles import load_tiles
 
@@ -53,6 +53,8 @@ def parse_assign(optlist=None):
     log = Logger.get()
     parser = argparse.ArgumentParser()
 
+    margins = get_default_exclusion_margins()
+
     parser.add_argument("--targets", type=str, required=True, nargs="+",
                         help="Input file with targets of any type.  This "
                         "argument can be specified multiple times (for "
@@ -127,11 +129,11 @@ def parse_assign(optlist=None):
                         default=1, help="Required number of sky targets per"
                         " fiber slitblock")
 
-    parser.add_argument("--margin-pos", "--margin_pos", type=float, required=False, default=0.05,
+    parser.add_argument("--margin-pos", "--margin_pos", type=float, required=False, default=margins['pos'],
                         help="Add margin (in mm) around positioner keep-out polygons")
-    parser.add_argument("--margin-petal", "--margin_petal", type=float, required=False, default=0.4,
+    parser.add_argument("--margin-petal", "--margin_petal", type=float, required=False, default=margins['petal'],
                         help="Add margin (in mm) around petal-boundary keep-out polygons")
-    parser.add_argument("--margin-gfa", "--margin_gfa", type=float, required=False, default=0.4,
+    parser.add_argument("--margin-gfa", "--margin_gfa", type=float, required=False, default=margins['gfa'],
                         help="Add margin (in mm) around GFA keep-out polygons")
 
     parser.add_argument("--write_all_targets", required=False, default=False,

py/fiberassign/test/collide.py

@@ -6,11 +6,10 @@
 
 from astropy.table import Table
 
-from fiberassign.hardware import load_hardware
+from fiberassign.hardware import load_hardware, get_default_exclusion_margins
 from fiberassign.tiles import load_tiles
 from fiberassign.targets import Targets, TargetsAvailable, LocationsAvailable, create_tagalong, load_target_file, targets_in_tiles
 from fiberassign.assign import Assignment
-
 from fiberassign.utils import Logger
 
 import fitsio
@@ -22,10 +21,7 @@ def main():
     # return
     log = Logger.get()
 
-    margins = dict(pos=0.05,
-                   petal=0.4,
-                   gfa=0.4)
-
+    margins = get_default_exclusion_margins()
     hw = load_hardware(rundate=None, add_margins=margins)
 
     t = Table.read('/global/cfs/cdirs/desi/survey/catalogs/Y1/LSS/tiles-DARK.fits')