diff --git a/regularizepsf/fitter.py b/regularizepsf/fitter.py index d204f4c..2ba6bd3 100644 --- a/regularizepsf/fitter.py +++ b/regularizepsf/fitter.py @@ -270,6 +270,7 @@ def find_stars_and_average(cls, patch_size: int, interpolation_scale: int = 1, average_mode: str = "median", + percentile: float = 10, star_threshold: int = 3, hdu_choice: int=0) -> CoordinatePatchCollection: """Loads a series of images, finds stars in each, and builds a CoordinatePatchCollection with averaged stars @@ -288,7 +289,11 @@ def find_stars_and_average(cls, if >1, the image are first scaled by this factor. This results in stars being aligned at a subpixel scale average_mode : str - "median" or "mean" determines how patches are combined + "median", "percentile", or "mean": determines how patches are + combined + percentile : float + If `average_mode` is `"percentile"`, use this percentile value + (from 0 to 100) star_threshold : int SEP's threshold for finding stars. See `threshold` in https://sep.readthedocs.io/en/v1.1.x/api/sep.extract.html#sep-extract @@ -385,7 +390,7 @@ def generator(): patch_size * interpolation_scale) averaged = this_collection.average(corners, patch_size * interpolation_scale, psf_size * interpolation_scale, - mode=average_mode) + mode=average_mode, percentile=percentile) if interpolation_scale != 1: for coordinate, _ in averaged.items(): @@ -403,15 +408,20 @@ def generator(): return output def average(self, corners: np.ndarray, patch_size: int, psf_size: int, # noqa: ARG002, kept for consistency - mode: str = "median") -> PatchCollectionABC: - CoordinatePatchCollection._validate_average_mode(mode) + mode: str = "median", percentile: float = 10) -> PatchCollectionABC: + CoordinatePatchCollection._validate_average_mode(mode, percentile) if mode == "mean": mean_stack = {tuple(corner): np.zeros((psf_size, psf_size)) for corner in corners} - counts = {tuple(corner): 0 for corner in corners} - elif mode == "median": - median_stack = {tuple(corner): [] for corner in corners} + counts = {tuple(corner): np.zeros((psf_size, psf_size)) + for corner in corners} + else: + # n.b. If mode is 'median', we could set mode='percentile' + # and percentile=50 to simplify parts of this function, but + # np.nanpercentile(x, 50) seems to be about half as fast as + # np.nanmedian(x), so let's keep a speedy special case for medians. + stack = {tuple(corner): [] for corner in corners} corners_x, corners_y = corners[:, 0], corners[:, 1] x_bounds = np.stack([corners_x, corners_x + patch_size], axis=-1) @@ -434,9 +444,9 @@ def average(self, corners: np.ndarray, patch_size: int, psf_size: int, # noqa: if mode == "mean": mean_stack[match_corner] = np.nansum([mean_stack[match_corner], patch], axis=0) - counts[match_corner] += 1 - elif mode == "median": - median_stack[match_corner].append(patch) + counts[match_corner] += np.isfinite(patch) + else: + stack[match_corner].append(patch) if mode == "mean": averages = {CoordinateIdentifier(None, corner[0], corner[1]): @@ -444,10 +454,18 @@ def average(self, corners: np.ndarray, patch_size: int, psf_size: int, # noqa: for corner in mean_stack} elif mode == "median": averages = {CoordinateIdentifier(None, corner[0], corner[1]): - np.nanmedian(median_stack[corner], axis=0) - if len(median_stack[corner]) > 0 else + np.nanmedian(stack[corner], axis=0) + if len(stack[corner]) > 0 else + np.zeros((psf_size, psf_size)) + for corner in stack} + elif mode == "percentile": + averages = {CoordinateIdentifier(None, corner[0], corner[1]): + np.nanpercentile(stack[corner], + percentile, + axis=0) + if len(stack[corner]) > 0 else np.zeros((psf_size, psf_size)) - for corner in median_stack} + for corner in stack} # Now that we have our combined patches, pad them as appropriate pad_shape = self._calculate_pad_shape(patch_size) for key, patch in averages.items(): @@ -455,12 +473,14 @@ def average(self, corners: np.ndarray, patch_size: int, psf_size: int, # noqa: return CoordinatePatchCollection(averages) @staticmethod - def _validate_average_mode(mode: str) -> None: + def _validate_average_mode(mode: str, percentile: float) -> None: """Determine if the average_mode is a valid kind""" - valid_modes = ["median", "mean"] + valid_modes = ["median", "mean", "percentile"] if mode not in valid_modes: msg = f"Found a mode of {mode} but it must be in the list {valid_modes}." raise ValueError(msg) + if mode == "percentile" and not (0 <= percentile <= 100): + raise ValueError("`percentile` must be between 0 and 100, inclusive") def _calculate_pad_shape(self, size: int) -> Tuple[int, int]: pad_amount = size - self.size diff --git a/tests/test_fitter.py b/tests/test_fitter.py index 4549862..ba43978 100644 --- a/tests/test_fitter.py +++ b/tests/test_fitter.py @@ -72,10 +72,31 @@ def test_saving_and_loading(): def test_coordinate_patch_average(): - collection = CoordinatePatchCollection({CoordinateIdentifier(0, 0, 0): np.zeros((10, 10)), - CoordinateIdentifier(0, 0, 0): np.ones((10, 10))*2}) - averaged_collection = collection.average(np.array([[0, 0]]), 10, 10, mode='median') - assert averaged_collection[CoordinateIdentifier(None, 0, 0)][1, 1] == 1 + collection = CoordinatePatchCollection( + {CoordinateIdentifier(0, 0, 0): np.full((10, 10), .3), + CoordinateIdentifier(1, 0, 0): np.full((10, 10), .5), + CoordinateIdentifier(2, 0, 0): np.full((10, 10), .9), + # Exercise the nan-rejection in CoordinatePatchCollection.average() + CoordinateIdentifier(3, 0, 0): np.full((10, 10), np.nan), + }) + for patch in collection.values(): + # Make the normalization of each patch a no-op + patch[-1, -1] = 1 + + averaged_collection = collection.average( + np.array([[0, 0]]), 10, 10, mode='median') + expected = np.nanmedian([.3, .5, .9]) + assert averaged_collection[CoordinateIdentifier(None, 0, 0)][1, 1] == expected + + averaged_collection = collection.average( + np.array([[0, 0]]), 10, 10, mode='mean') + expected = np.nanmean([.3, .5, .9]) + assert averaged_collection[CoordinateIdentifier(None, 0, 0)][1, 1] == expected + + averaged_collection = collection.average( + np.array([[0, 0]]), 10, 10, mode='percentile', percentile=20) + expected = np.nanpercentile([.3, .5, .9], 20) + assert averaged_collection[CoordinateIdentifier(None, 0, 0)][1, 1] == expected def test_calculate_pad_shape(): @@ -98,7 +119,20 @@ def test_odd_pad_shape_errors(): def test_validate_average_mode(): with pytest.raises(ValueError): - CoordinatePatchCollection._validate_average_mode("nonexistent_method") + CoordinatePatchCollection._validate_average_mode( + "nonexistent_method", 1) + + # Ensure valid modes are accepted + for mode in ('mean', 'median', 'percentile'): + CoordinatePatchCollection._validate_average_mode(mode, 1) + + # Check invalid percentile values + with pytest.raises(ValueError): + CoordinatePatchCollection._validate_average_mode( + "percentile", -1) + with pytest.raises(ValueError): + CoordinatePatchCollection._validate_average_mode( + "percentile", 101) def test_find_stars_and_average():