diff --git a/python/lsst/analysis/tools/actions/plot/plotUtils.py b/python/lsst/analysis/tools/actions/plot/plotUtils.py index 305cdcb8c..02f83082f 100644 --- a/python/lsst/analysis/tools/actions/plot/plotUtils.py +++ b/python/lsst/analysis/tools/actions/plot/plotUtils.py @@ -27,6 +27,7 @@ import matplotlib import matplotlib.pyplot as plt import numpy as np +import pandas as pd from lsst.geom import Box2D, SpherePoint, degrees from lsst.pex.config import Config, Field from matplotlib import colors @@ -384,21 +385,26 @@ def sortAllArrays(arrsToSort, sortArrayIndex=0): Parameters ---------- - arrsToSort : `list` [`np.array`] - A list of arrays to be simultaneously sorted based on the array in - the list position given by ``sortArrayIndex`` (defaults to be the - first array in the list). + arrsToSort : `list` [`np.array`] | `list` [`pd.Series`] + A list of arrays or Series to be simultaneously sorted based on the + array in the list position given by ``sortArrayIndex`` (defaults to be + the first array in the list). sortArrayIndex : `int`, optional Zero-based index indicating the array on which to base the sorting. Returns ------- - arrsToSort : `list` [`np.array`] - The list of arrays sorted on array in list index ``sortArrayIndex``. + arrsToSort : `list` [`np.array`] | `list` [`pd.Series`] + The list of arrays or Series (same type as the input) sorted on array + in list index ``sortArrayIndex``. """ ids = extremaSort(arrsToSort[sortArrayIndex]) for i, arr in enumerate(arrsToSort): - arrsToSort[i] = arr[ids] + if isinstance(arr, pd.Series): + arrsToSort[i] = arr.iloc[ids] + else: + arrsToSort[i] = arr[ids] + return arrsToSort