Merge pull request #292 from sot/performance

Improve performance

proseco/acq.py

@@ -89,8 +89,15 @@ def get_acq_catalog(obsid=0, **kwargs):
     acqs.update_p_acq_column(acqs)
 
     # Sort to make order match the original candidate list order (by
-    # increasing mag), and assign a slot.
+    # increasing mag), and assign a slot.  Sadly astropy 3.1 has a real
+    # performance bug here and doing the sort makes 6 deepcopy's of the
+    # meta, which in this case is substantial (mostly stars).  So temporarily
+    # clear out the meta before sorting and then restore from a (light) copy.
+    acqs_meta_copy = acqs.meta.copy()
+    acqs.meta.clear()
     acqs.sort('idx')
+    acqs.meta.update(acqs_meta_copy)
+
     acqs['slot'] = np.arange(len(acqs), dtype=np.int64)
 
     # Add slot to cand_acqs table, putting in -99 if not selected as acq.
@@ -491,9 +498,11 @@ def select_best_p_acqs(self, cand_acqs, min_p_acq, acq_indices, box_sizes):
             # probability to zero.  This happens for small maneuver angles where
             # acq.box_sizes might be only [60] or [60, 80].
             p_acqs_for_box = np.zeros(len(cand_acqs))
-            for idx, acq in enumerate(cand_acqs):
-                if box_size in acq['box_sizes']:
-                    p_acqs_for_box[idx] = acq['probs'].p_acq_marg(box_size, self)
+            my_box_sizes = cand_acqs['box_sizes']
+            my_probs = cand_acqs['probs']
+            for idx in range(len(cand_acqs)):
+                if box_size in my_box_sizes[idx]:
+                    p_acqs_for_box[idx] = my_probs[idx].p_acq_marg(box_size, self)
 
             self.log(f'Trying search box size {box_size} arcsec', level=1)
 
@@ -573,8 +582,7 @@ def get_initial_catalog(self):
         acqs_init = cand_acqs[acq_indices]
 
         # Transfer to acqs (which at this point is an empty table)
-        for name, col in acqs_init.columns.items():
-            self[name] = col
+        self.add_columns(acqs_init.columns.values())
 
     def calc_p_brightest(self, acq, box_size, man_err=0, bgd=0):
         """
@@ -639,11 +647,15 @@ def calc_p_safe(self, verbose=False):
 
         p_no_safe = 1.0
 
+        self_halfws = self['halfw']
+        self_probs = self['probs']
+
         for man_err, p_man_err in zip(ACQ.man_errs, self.p_man_errs):
             if p_man_err == 0.0:
                 continue
 
-            p_acqs = [acq['probs'].p_acqs(acq['halfw'], man_err, self) for acq in self]
+            p_acqs = [prob.p_acqs(halfw, man_err, self)
+                      for halfw, prob in zip(self_halfws, self_probs)]
 
             p_n_cum = prob_n_acq(p_acqs)[1]  # This returns (p_n, p_n_cum)
 
@@ -734,7 +746,7 @@ def optimize_acqs_halfw(self, verbose=False):
         :param verbose: include additional information in the run log
         """
         p_safe = self.calc_p_safe()
-        idxs = self.argsort('p_acq')
+        idxs = self['p_acq'].argsort()
 
         # Any updates made?
         any_improved = False
@@ -782,13 +794,18 @@ def optimize_catalog(self, verbose=False):
             # Get the index of the worst p_acq in the catalog, excluding acq stars
             # that are in include_ids (since they are not to be replaced).
             ok = [acq['id'] not in self.include_ids for acq in self]
-            acqs = self[ok]
+            # acqs = self[ok]
+            acqs_probs_ok = self['probs'][ok]
+            acqs_halfw_ok = self['halfw'][ok]
+            acqs_id_ok = self['id'][ok]
 
             # Sort by the marginalized acq probability for the current box size
-            p_acqs = [acq['probs'].p_acq_marg(acq['halfw'], acqs) for acq in acqs]
+            p_acqs = [acq_probs.p_acq_marg(acq_halfw, self)
+                      for acq_probs, acq_halfw in zip(acqs_probs_ok, acqs_halfw_ok)]
+            # TODO: performance?
             idx_worst = np.argsort(p_acqs, kind='mergesort')[0]
 
-            idx = self.get_id_idx(acqs[idx_worst]['id'])
+            idx = self.get_id_idx(acqs_id_ok[idx_worst])
 
             self.log('Trying to use {} mag={:.2f} to replace idx={} with p_acq={:.3f}'
                      .format(cand_id, cand_acq['mag'], idx, p_acqs[idx_worst]), id=cand_id)
@@ -971,23 +988,26 @@ def get_imposter_stars(dark, star_row, star_col, thresh=None,
 
         yang, zang = pixels_to_yagzag(row, col, allow_bad=True)
 
-        out = {'row': row,
-               'col': col,
-               'd_row': row - star_row,
-               'd_col': col - star_col,
-               'yang': yang,
-               'zang': zang,
-               'row0': c_row - 4,
-               'col0': c_col - 4,
-               'img': img,
-               'img_sum': img_sum,
-               'mag': mag,
-               'mag_err': get_mag_std(mag).item(),
-               }
+        out = (row,
+               col,
+               row - star_row,
+               col - star_col,
+               yang,
+               zang,
+               c_row - 4,
+               c_col - 4,
+               img,
+               img_sum,
+               mag,
+               get_mag_std(mag).item(),
+               )
         outs.append(out)
 
     if len(outs) > 0:
-        outs = Table(outs).as_array()
+        dtype = [('row', '<f8'), ('col', '<f8'), ('d_row', '<f8'), ('d_col', '<f8'),
+                 ('yang', '<f8'), ('zang', '<f8'), ('row0', '<i8'), ('col0', '<i8'),
+                 ('img', 'f8', (8, 8)), ('img_sum', '<f8'), ('mag', '<f8'), ('mag_err', '<f8')]
+        outs = np.rec.fromrecords(outs, dtype=dtype)
         outs.sort(order=['mag'])
 
     return outs

proseco/core.py

@@ -13,7 +13,7 @@
 
 from chandra_aca.transform import (yagzag_to_pixels, pixels_to_yagzag,
                                    count_rate_to_mag, mag_to_count_rate)
-from chandra_aca.aca_image import ACAImage, AcaPsfLibrary
+from chandra_aca.aca_image import AcaPsfLibrary
 from Ska.quatutil import radec2yagzag, yagzag2radec
 import agasc
 from Quaternion import Quat
@@ -565,21 +565,18 @@ def dark(self):
         self.log(f'Getting dark cal image at date={self.date} t_ccd={self.t_ccd:.1f}')
         self.dark = get_dark_cal_image(date=self.date, select='before',
                                        t_ccd_ref=self.t_ccd,
-                                       aca_image=True)
+                                       aca_image=False)
         return self._dark
 
     @dark.setter
     def dark(self, value):
-        if not isinstance(value, ACAImage):
-            assert value.shape == (1024, 1024)
-            value = ACAImage(value, row0=-512, col0=-512, copy=False)
-
+        assert value.shape == (1024, 1024)
         self._dark = value
 
         # Set pixel regions from ACA.bad_pixels to have acqs.dark=700000 (5.0 mag
         # star) per pixel.
         for r0, r1, c0, c1 in ACA.bad_pixels:
-            self._dark.aca[r0:r1 + 1, c0:c1 + 1] = ACA.bad_pixel_dark_current
+            self._dark[r0 + 512:r1 + 513, c0 + 512:c1 + 513] = ACA.bad_pixel_dark_current
 
     def set_attrs_from_kwargs(self, **kwargs):
         for name, val in kwargs.items():

proseco/fid.py

@@ -358,7 +358,7 @@ def near_hot_or_bad_pixel(self, fid):
         c0 = int(fid['col'] - dp)
         r1 = int(fid['row'] + dp) + 1
         c1 = int(fid['col'] + dp) + 1
-        dark = self.dark.aca[r0:r1, c0:c1]
+        dark = self.dark[r0 + 512:r1 + 512, c0 + 512:c1 + 512]
 
         bad = dark > FID.hot_pixel_spoiler_limit
         if np.any(bad):

proseco/guide.py

@@ -15,6 +15,7 @@
 from .core import bin2x2, ACACatalogTable, MetaAttribute, AliasAttribute
 
 CCD = ACA.CCD
+APL = AcaPsfLibrary()
 
 STAR_PAIR_DIST_CACHE = {}
 
@@ -55,8 +56,7 @@ def get_guide_catalog(obsid=0, **kwargs):
     selected = guides.run_search_stages()
 
     # Transfer to table (which at this point is an empty table)
-    for name, col in selected.columns.items():
-        guides[name] = col
+    guides.add_columns(selected.columns.values())
 
     guides['idx'] = np.arange(len(guides))
 
@@ -540,7 +540,6 @@ def check_spoil_contrib(cand_stars, ok, stars, regfrac, bgthresh):
     :returns: reg_spoiled, bg_spoiled, rej - two masks on cand_stars and a list of reject debug dicts
     """
     fraction = regfrac
-    APL = AcaPsfLibrary()
     bg_spoiled = np.zeros_like(ok)
     reg_spoiled = np.zeros_like(ok)
     bgpix = CCD['bgpix']
@@ -623,17 +622,20 @@ def check_mag_spoilers(cand_stars, ok, stars, n_sigma):
 
     mag_spoiled = np.zeros(len(ok)).astype(bool)
     rej = []
-    for cand in cand_stars[ok]:
-        pix_dist = np.sqrt(((cand['row'] - stars['row']) ** 2) +
-                           ((cand['col'] - stars['col']) ** 2))
-        spoilers = ((np.abs(cand['row'] - stars['row']) < 10) &
-                    (np.abs(cand['col'] - stars['col']) < 10))
+    cand_idxs = np.flatnonzero(ok)
+
+    for cand_idx in cand_idxs:
+        cand = cand_stars[cand_idx]
+        spoil_idxs = np.flatnonzero(
+            (np.abs(cand['row'] - stars['row']) < 10) &
+            (np.abs(cand['col'] - stars['col']) < 10))
 
         # If there is only one match, it is the candidate so there's nothing to do
-        if np.count_nonzero(spoilers) == 1:
+        if len(spoil_idxs) == 1:
             continue
 
-        for spoil, dist in zip(stars[spoilers], pix_dist[spoilers]):
+        for spoil_idx in spoil_idxs:
+            spoil = stars[spoil_idx]
             if spoil['id'] == cand['id']:
                 continue
             if (cand['mag'] - spoil['mag']) < magdifflim:
@@ -642,6 +644,8 @@ def check_mag_spoilers(cand_stars, ok, stars, n_sigma):
                                   (spoil['MAG_ACA_ERR'] * 0.01) ** 2)
             delmag = cand['mag'] - spoil['mag'] + n_sigma * mag_err_sum
             thsep = intercept + delmag * spoilslope
+            dist = np.sqrt(((cand['row'] - spoil['row']) ** 2) +
+                           ((cand['col'] - spoil['col']) ** 2))
             if dist < thsep:
                 rej.append({'id': cand['id'],
                             'spoiler': spoil['id'],
@@ -751,7 +755,7 @@ def get_imposter_mags(cand_stars, dark, dither):
     for cand in cand_stars:
         rminus, rplus = get_ax_range(cand['row'], row_extent)
         cminus, cplus = get_ax_range(cand['col'], col_extent)
-        pix = np.array(dark.aca[rminus:rplus, cminus:cplus])
+        pix = np.array(dark[rminus + 512:rplus + 512, cminus + 512:cplus + 512])
         pixmax = 0
         max_r = None
         max_c = None

proseco/tests/test_acq.py

@@ -776,7 +776,7 @@ def test_acq_fid_catalog_probs_low_level():
     aca = get_aca_catalog(**kwargs)
     acqs = aca.acqs
 
-    assert np.all(acqs.dark.aca[-512:0, -512:0] == 40)
+    assert np.all(acqs.dark[0:512, 0:512] == 40)
 
     # Initial fid set is empty () and we check baseline p_safe
     assert acqs.fid_set == ()

proseco/tests/test_catalog.py

@@ -405,7 +405,7 @@ def test_bad_pixel_dark_current():
     aca = get_aca_catalog(**kwargs)
 
     # Make sure bad pixels have expected value
-    assert np.all(aca.acqs.dark.aca[-245:0, 454] == ACA.bad_pixel_dark_current)
+    assert np.all(aca.acqs.dark[-245 + 512:512, 454 + 512] == ACA.bad_pixel_dark_current)
 
     exp_ids = [2, 100, 101, 102, 103]
     assert sorted(aca.guides['id']) == exp_ids

proseco/tests/test_common.py

@@ -20,7 +20,7 @@ def mod_std_info(**kwargs):
 
 
 # Flat dark current map
-DARK40 = ACAImage(np.full(shape=(1024, 1024), fill_value=40), row0=-512, col0=-512)
+DARK40 = np.full(shape=(1024, 1024), fill_value=40)
 
 
 # Parameters for test cases (to avoid starcheck.db3 dependence)

proseco/tests/test_fid.py

@@ -211,7 +211,7 @@ def test_fid_hot_pixel_reject():
         fid = FIDS.cand_fids.get_id(fid_id)
         r = int(round(fid['row'] + off))
         c = int(round(fid['col'] + off))
-        dark.aca[r, c] = dc
+        dark[r + 512, c + 512] = dc
 
     fids = get_fid_catalog(stars=StarsTable.empty(), dark=dark, **STD_INFO)
     assert fids['id'].tolist() == [2, 3, 6]

proseco/tests/test_guide.py

@@ -340,15 +340,15 @@ def test_pix_spoiler(case):
     stars = StarsTable.empty()
     stars.add_fake_star(row=0, col=0, mag=7.0, id=1, ASPQ1=0)
     stars.add_fake_constellation(n_stars=4)
-    dark = ACAImage(np.zeros((1024, 1024)), row0=-512, col0=-512)
+    dark = np.zeros((1024, 1024))
     pix_config = {'att': (0, 0, 0),
                   'date': '2018:001',
                   't_ccd': -10,
                   'n_guide': 5,
                   'stars': stars}
     # Use the "case" to try to spoil the first star with a bad pixel
-    dark.aca[case['offset_row'] + int(stars[0]['row']),
-             case['offset_col'] + int(stars[0]['col'])] = mag_to_count_rate(stars[0]['mag'])
+    dark[case['offset_row'] + int(stars[0]['row']) + 512,
+         case['offset_col'] + int(stars[0]['col']) + 512] = mag_to_count_rate(stars[0]['mag'])
     selected = get_guide_catalog(**pix_config, dither=case['dither'], dark=dark)
     assert (1 not in selected['id']) == case['spoils']