diff --git a/bin/select_secondary b/bin/select_secondary index 8c714728f..727160b87 100755 --- a/bin/select_secondary +++ b/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 diff --git a/doc/changes.rst b/doc/changes.rst index 3107e8e39..e6cf2d012 100644 --- a/doc/changes.rst +++ b/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) ------------------ diff --git a/py/desitarget/subpriority.py b/py/desitarget/subpriority.py index 8db4d16eb..9f6ca092f 100644 --- a/py/desitarget/subpriority.py +++ b/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): """ diff --git a/py/desitarget/test/test_subpriority.py b/py/desitarget/test/test_subpriority.py index 3d07196a5..aa05d1dd6 100644 --- a/py/desitarget/test/test_subpriority.py +++ b/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()