diff --git a/lcls_tools/common/image/fit.py b/lcls_tools/common/image/fit.py index 225eea8f..c0a0246b 100644 --- a/lcls_tools/common/image/fit.py +++ b/lcls_tools/common/image/fit.py @@ -1,10 +1,11 @@ from abc import ABC, abstractmethod import importlib -from typing import List +from typing import List, Optional import numpy as np from numpy import ndarray from pydantic import PositiveFloat, Field, ConfigDict +from enum import Enum from lcls_tools.common.measurements.utils import NDArrayAnnotatedType import lcls_tools import warnings @@ -17,6 +18,24 @@ class ImageFitResult(lcls_tools.common.BaseModel): image: NDArrayAnnotatedType +class ImageProjectionFitFailureMode(Enum): + """ + NONE : 0 + Flag for when the projection fitting had no validation errors + BEAM_SIZE_TOO_SMALL : 1 + Flag for when the beam size from the fit is below the min_beam_size threshold + BEAM_EXTENT_OFF_SCREEN : 2 + FLag for when the beam extends off the screen + LOW_SIGNAL_TO_NOISE_RATIO : 3 + Flag for when the signal_to_noise_ratio is below the signal_to_noise_threshold + """ + + NONE = 0 + BEAM_SIZE_TOO_SMALL = 1 + BEAM_EXTENT_OFF_SCREEN = 2 + LOW_SIGNAL_TO_NOISE_RATIO = 3 + + class ImageProjectionFitResult(ImageFitResult): projection_fit_module: str projection_fit_parameters: List[dict[str, float]] @@ -26,6 +45,10 @@ class ImageProjectionFitResult(ImageFitResult): beam_extent: NDArrayAnnotatedType = Field( description="Extent of the beam in the data, defined as mean +/- 2*sigma" ) + failure_mode: Optional[List[ImageProjectionFitFailureMode]] = Field( + default=None, + description="Reason for failure if either the x/y projection fit was rejected during fit validation", + ) class ImageFit(lcls_tools.common.BaseModel, ABC): @@ -63,8 +86,9 @@ def _validate_parameters(self, parameters: list[float]): class ImageProjectionFit(ImageFit): """ Image fitting class that gets the beam size and location by independently fitting - the x/y projections. The default configuration uses a Gaussian fitting of the - profile with prior distributions placed on the model parameters. + the x/y projections of the input image, given as a 2d array of intensity per pixel. + The default configuration uses a Gaussian fitting of the profile with prior distributions + placed on the model parameters. Attributes ---------- @@ -79,6 +103,9 @@ class ImageProjectionFit(ImageFit): Fit amplitude to noise threshold for the fit. If the amplitude of the fit is below this threshold times the noise standard deviation and `validate_fit` is True, the fit parameters will be set to NaN. + min_beam_size : PositiveFloat + Minimum beam size in pixels for the fit. If the beam size of the fit is below this threshold, + and 'validate_fit' is True, the fit parameters will be set to NaN. beam_extent_n_stds : PositiveFloat Number of standard deviations on either side to use for the beam extent. If the beam extent is outside the image and `validate_fit` is True, the fit parameters will be set to NaN. @@ -94,6 +121,9 @@ class ImageProjectionFit(ImageFit): signal_to_noise_threshold: PositiveFloat = Field( 2.0, description="Fit amplitude to noise threshold for the fit" ) + min_beam_size: PositiveFloat = Field( + 1.0, description="Minimum beam size in pixels to accept for validation" + ) beam_extent_n_stds: PositiveFloat = Field( 2.0, description="Number of standard deviations on either side to use for the beam extent", @@ -107,6 +137,7 @@ def _fit_image(self, image: ndarray) -> ImageProjectionFitResult: fit_parameters = [] signal_to_noise_ratios = [] beam_extent = [] + failure_mode = [] module = importlib.import_module(f"lcls_tools.common.model.{self.fit_module}") @@ -122,7 +153,12 @@ def _fit_image(self, image: ndarray) -> ImageProjectionFitResult: # perform validation checks, modify parameters if checks fail if self.validate_fit: - self._validate_parameters(parameters, snr, extent, projection, dim) + fmode = self._validate_parameters( + parameters, snr, extent, projection, dim + ) + failure_mode.append(fmode) + else: + failure_mode = None fit_parameters.append(parameters) signal_to_noise_ratios.append(snr) @@ -137,6 +173,7 @@ def _fit_image(self, image: ndarray) -> ImageProjectionFitResult: projection_fit_module=self.fit_module, signal_to_noise_ratio=signal_to_noise_ratios, beam_extent=beam_extent, + failure_mode=failure_mode, ) return result @@ -144,18 +181,18 @@ def _fit_image(self, image: ndarray) -> ImageProjectionFitResult: def _validate_parameters( self, parameters, signal_to_noise_ratios, beam_extent, projection, dim ): - # if the amplitude of the the fit is smaller than noise then reject - # moving this into a validate function to clean it up. - if self.signal_to_noise_threshold is not None: - if signal_to_noise_ratios < self.signal_to_noise_threshold: + # if the beam size is smaller than the specified minimum number of pixels, fits cannot be trusted + if self.min_beam_size is not None: + if parameters["sigma"] < self.min_beam_size: for name in parameters.keys(): parameters[name] = np.nan warnings.warn( - f"Projection in {dim} had a low amplitude relative to noise" + f"Projection in {dim} had too small a beam size, fit cannot be trusted" ) + return ImageProjectionFitFailureMode.BEAM_SIZE_TOO_SMALL - # if the beam extent is outside the image then its off the screen etc. and fits cannot be trusted + # if the beam extent is outside the image then it's off the screen and fits cannot be trusted if self.beam_extent_n_stds is not None: if beam_extent[0] < 0 or beam_extent[1] > len(projection): for name in parameters.keys(): @@ -164,3 +201,17 @@ def _validate_parameters( warnings.warn( f"Projection in {dim} was off the screen, fit cannot be trusted" ) + return ImageProjectionFitFailureMode.BEAM_EXTENT_OFF_SCREEN + + # if the amplitude of the the fit is smaller than noise then reject + if self.signal_to_noise_threshold is not None: + if signal_to_noise_ratios < self.signal_to_noise_threshold: + for name in parameters.keys(): + parameters[name] = np.nan + + warnings.warn( + f"Projection in {dim} had a low amplitude relative to noise" + ) + return ImageProjectionFitFailureMode.LOW_SIGNAL_TO_NOISE_RATIO + + return ImageProjectionFitFailureMode.NONE