Merge pull request #743 from desihub/secondsubprio

override_subpriority support duplicate TARGETID inputs

bin/select_secondary

@@ -5,6 +5,7 @@ import numpy as np
 from desitarget.secondary import select_secondary, _get_scxdir
 from desitarget.brightmask import is_in_bright_mask, get_recent_mask_dir
 from desitarget import io
+from desitarget.subpriority import override_subpriority
 import os
 from glob import glob
 from time import time

doc/changes.rst

@@ -5,6 +5,8 @@ desitarget Change Log
 1.1.0 (unreleased)
 ------------------
 
+* override_subpriority support duplicate input TARGETID from secondaries
+  [`PR #743`_].
 * Support reading mtl 1.0.0 format with different columns [`PR #742`_].
 * Enable optional subpriority overrides [`PR #740`_, `PR #741`_].
 * Allow initial ledgers to use a preordained timestamp [`PR #739`_].
@@ -38,6 +40,7 @@ desitarget Change Log
 .. _`PR #740`: https://github.com/desihub/desitarget/pull/740
 .. _`PR #741`: https://github.com/desihub/desitarget/pull/741
 .. _`PR #742`: https://github.com/desihub/desitarget/pull/742
+.. _`PR #743`: https://github.com/desihub/desitarget/pull/743
 
 1.0.1 (2021-05-14)
 ------------------

py/desitarget/subpriority.py

@@ -26,12 +26,29 @@ def override_subpriority(targets, override):
     Rows in ``targets`` that aren't in ``override`` are unchanged.
     Rows in ``override`` that aren't in ``targets`` are ignored.
     """
-    ii_targ, ii_over = match(targets['TARGETID'], override['TARGETID'])
-    if len(ii_targ) > 0:
-        targets['SUBPRIORITY'][ii_targ] = override['SUBPRIORITY'][ii_over]
-
-    return np.sort(ii_targ)
-
+    # SB Use geomask.match if input targets are unique,
+    # SB otherwise use slower code that supports duplicates (e.g. secondaries)
+    if len(np.unique(targets['TARGETID'])) == len(targets['TARGETID']):
+        ii_targ, ii_over = match(targets['TARGETID'], override['TARGETID'])
+        if len(ii_targ) > 0:
+            targets['SUBPRIORITY'][ii_targ] = override['SUBPRIORITY'][ii_over]
+        return np.sort(ii_targ)
+    else:
+        ii = np.where(np.isin(targets['TARGETID'], override['TARGETID']))[0]
+        n = len(ii)
+        if n > 0:
+            # SB create TARGETID->SUBPRIORITY dict only for TARGETID in targets
+            subprio_dict = dict()
+            jj = np.where(np.isin(override['TARGETID'], targets['TARGETID']))[0]
+            for tid, subprio in zip(
+                    override['TARGETID'][jj], override['SUBPRIORITY'][jj]):
+                subprio_dict[tid] = subprio
+
+            for i in ii:
+                tid = targets['TARGETID'][i]
+                targets['SUBPRIORITY'][i] = subprio_dict[tid]
+
+        return ii
 
 def get_fiberassign_subpriorities(fiberassignfiles):
     """

py/desitarget/test/test_subpriority.py

@@ -41,6 +41,30 @@ def test_override(self):
             else:
                 self.assertEqual(targets['SUBPRIORITY'][i], orig_subpriority[i])
 
+    def test_override_duplicates(self):
+        """Test subpriority override with duplicate input TARGETIDs"""
+        targets = Table()
+        targets['TARGETID'] = [1,2,3,2,1,5]
+        n = len(targets['TARGETID'])
+        orig_subpriority = np.random.random(n)
+        targets['SUBPRIORITY'] = orig_subpriority.copy()
+
+        override = Table()
+        override['TARGETID'] = np.array([3,2,20])
+        override['SUBPRIORITY'] = np.array([10.0, 20.0, 30.0])
+
+        desitarget.subpriority.override_subpriority(targets, override)
+
+        #- Check that we overrode correctly; don't juse geomask.match
+        #- to avoid circularity of code and test
+        for i, tid in enumerate(targets['TARGETID']):
+            in_override = np.where(override['TARGETID'] == tid)[0]
+            if len(in_override) > 0:
+                j = in_override[0]
+                self.assertEqual(targets['SUBPRIORITY'][i], override['SUBPRIORITY'][j])
+            else:
+                self.assertEqual(targets['SUBPRIORITY'][i], orig_subpriority[i])
+
 
 if __name__ == '__main__':
     unittest.main()