pysmo.tools.iccs

@beartype
@define(slots=True)
class ICCS:
    """Class to store a list of [`ICCSSeismograms`][pysmo.tools.iccs.ICCSSeismogram] and run the ICCS algorithm.

    The [`ICCS`][pysmo.tools.iccs.ICCS] class serves as a container to store a
    list of seismograms (typically recordings of the same event at different
    stations), and to then run the ICCS algorithm when an instance of this
    class is called. Processing parameters that are common to all seismograms
    are stored as attributes (e.g. time window limits).

    The seismograms stored in an instance are prepared in two distinct ways:

    - **For cross-correlation:** uses [`ramp_width`][pysmo.tools.iccs.ICCS.ramp_width]
      to define how much of the seismogram should be used as *taper* before and
      after the time window of interest. These are the seismograms that form
      the stack and are used in the cross-correlation.
    - **For added context:** uses [`context_width`][pysmo.tools.iccs.ICCS.context_width]
      to define how much of the seismogram should be used as *extra context*
      before and after the time window of interest.
      [`context_width`][pysmo.tools.iccs.ICCS.context_width] should be chosen
      such that a large enough portion of the seismogram is shown to e.g.
      interactively pick new time window boundaries.

    Apart from tapering the two types are processed the same. For performance,
    the prepared seismograms are cached and only calculated on a first call or
    if relevant parameters are updated.

    Examples:
        We begin with a set of SAC files of the same event, recorded at different
        stations. All files have a preliminary phase arrival estimate saved in the
        `T0` SAC header, so we can use these files to create instances of the
        [`MiniICCSSeismogram`][pysmo.tools.iccs.MiniICCSSeismogram] class for use
        with the [`ICCS`][pysmo.tools.iccs.ICCS] class:

        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.functions import clone_to_mini
        >>> from pysmo.tools.iccs import MiniICCSSeismogram
        >>> from pathlib import Path
        >>>
        >>> sacfiles = sorted(Path("iccs-example/").glob("*.bhz"))
        >>>
        >>> seismograms = []
        >>> for sacfile in sacfiles:
        ...     sac = SAC.from_file(sacfile)
        ...     update = {"t0": sac.timestamps.t0}
        ...     iccs_seismogram = clone_to_mini(MiniICCSSeismogram, sac.seismogram, update=update)
        ...     seismograms.append(iccs_seismogram)
        ...
        >>>
        ```

        To better illustrate the different modes of running the ICCS algorithm,
        we modify the data and picks in the seismograms to make them **worse**
        than they actually are:

        ```python
        >>> from pandas import Timedelta
        >>> from copy import deepcopy
        >>> import numpy as np
        >>>
        >>> # change the sign of the data in the first seismogram
        >>> seismograms[0].data *= -1
        >>>
        >>> # move the initial pick 2 seconds earlier in second seismogram
        >>> seismograms[1].t0 += Timedelta(seconds=-2)
        >>>
        >>> # move the initial pick 2 seconds later in third seismogram
        >>> seismograms[2].t0 += Timedelta(seconds=2)
        >>>
        >>> # create a seismogram with completely random data
        >>> iccs_random: MiniICCSSeismogram = deepcopy(seismograms[-1])
        >>> np.random.seed(42)  # set this for consistent results during testing
        >>> iccs_random.data = np.random.rand(len(iccs_random.data))
        >>> seismograms.append(iccs_random)
        >>>
        ```

        We can now create an instance of the [`ICCS`][pysmo.tools.iccs.ICCS]
        class and plot the initial [`stack`][pysmo.tools.iccs.ICCS.stack] and
        [`cc_seismograms`][pysmo.tools.iccs.ICCS.cc_seismograms]:

        ```python
        >>> from pysmo.tools.iccs import ICCS, plot_stack
        >>> iccs = ICCS(seismograms)
        >>> plot_stack(iccs, context=False)
        >>>
        ```

        <!-- invisible-code-block: python
        ```
        >>> import matplotlib.pyplot as plt
        >>> plt.close("all")
        >>> if savedir:
        ...     plt.style.use("dark_background")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_initial-dark.png", transparent=True)
        ...
        ...     plt.style.use("default")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_initial.png", transparent=True)
        >>>
        ```
        -->

        ![Initial stack](../../../images/tools/iccs/iccs_stack_initial.png#only-light){ loading=lazy }
        ![Initial stack](../../../images/tools/iccs/iccs_stack_initial-dark.png#only-dark){ loading=lazy }

        The phase emergence is not visible in the stack, and the (absolute)
        correlation coefficients of the seismograms are not very high. This
        shows the initial picks are not very good and/or that the data are of
        low quality. To run the ICCS algorithm, we simply call (execute) the
        ICCS instance:

        ```python
        >>> convergence_list = iccs()  # this runs the ICCS algorithm and returns
        >>>                            # a list of the convergence value after each
        >>>                            # iteration.
        >>> plot_stack(iccs, context=False)
        >>>
        ```

        <!-- invisible-code-block: python
        ```
        >>> import matplotlib.pyplot as plt
        >>> plt.close("all")
        >>> if savedir:
        ...     plt.style.use("dark_background")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_first_run-dark.png", transparent=True)
        ...
        ...     plt.style.use("default")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_first_run.png", transparent=True)
        >>>
        ```
        -->

        ![Stack after first run](../../../images/tools/iccs/iccs_stack_first_run.png#only-light){ loading=lazy }
        ![Stack after first run](../../../images/tools/iccs/iccs_stack_first_run-dark.png#only-dark){ loading=lazy }

        Despite the random noise seismogram, the phase arrival is now visible in
        the stack. Seismograms with low correlation coefficients can automatically
        be deselected from the calculations by running ICCS again with
        `autoselect=True`:


        ```python
        >>> _ = iccs(autoselect=True)
        >>> plot_stack(iccs, context=False)
        >>>
        ```

        <!-- invisible-code-block: python
        ```
        >>> import matplotlib.pyplot as plt
        >>> plt.close("all")
        >>> if savedir:
        ...     plt.style.use("dark_background")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_autoselect-dark.png", transparent=True)
        ...
        ...     plt.style.use("default")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_autoselect.png", transparent=True)
        >>>
        ```
        -->

        ![Stack after run with autoselect](../../../images/tools/iccs/iccs_stack_autoselect.png#only-light){ loading=lazy }
        ![Stack after run with autoselect](../../../images/tools/iccs/iccs_stack_autoselect-dark.png#only-dark){ loading=lazy }


        Seismograms that fit better with their polarity reversed can be flipped
        automatically by setting `autoflip=True`:

        ```python
        >>> _ = iccs(autoflip=True)
        >>> plot_stack(iccs, context=False)
        >>>
        ```

        <!-- invisible-code-block: python
        ```
        >>> import matplotlib.pyplot as plt
        >>> plt.close("all")
        >>> if savedir:
        ...     plt.style.use("dark_background")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_autoflip-dark.png", transparent=True)
        ...
        ...     plt.style.use("default")
        ...     fig, _ = plot_stack(iccs, context=False, return_fig=True)
        ...     fig.savefig(savedir / "iccs_stack_autoflip.png", transparent=True)
        >>>
        ```
        -->

        ![Stack after run with autoflip](../../../images/tools/iccs/iccs_stack_autoflip.png#only-light){ loading=lazy }
        ![Stack after run with autoflip](../../../images/tools/iccs/iccs_stack_autoflip-dark.png#only-dark){ loading=lazy }


        To further improve results, you can interactively update the picks,
        time window, and minimum correlation coefficient using
        [`update_pick`][pysmo.tools.iccs.update_pick],
        [`update_timewindow`][pysmo.tools.iccs.update_timewindow], and
        [`update_min_ccnorm`][pysmo.tools.iccs.update_min_ccnorm],
        respectively, and then run the ICCS algorithm again.
    """

    seismograms: Sequence[ICCSSeismogram] = field(
        factory=lambda: list[ICCSSeismogram](), validator=_clear_cache_on_update
    )
    """Input seismograms.

    These seismograms provide the input data for ICCS. They are used to store
    processing parameters and create shorter seismograms (based on pick and
    time window) that are then used for cross-correlation. The shorter
    seismograms are created on the fly and then cached within an [`ICCS`]
    [pysmo.tools.iccs.ICCS] instance.
    """

    window_pre: Timedelta = field(
        default=ICCS_DEFAULTS.window_pre,
        validator=[
            validators.lt(Timedelta(seconds=0)),
            _validate_window_pre,
            _clear_cache_on_update,
        ],
    )
    """Beginning of the time window relative to the pick."""

    window_post: Timedelta = field(
        default=ICCS_DEFAULTS.window_post,
        validator=[
            validators.gt(Timedelta(seconds=0)),
            _validate_window_post,
            _clear_cache_on_update,
        ],
    )
    """End of the time window relative to the pick."""

    context_width: Timedelta = field(
        default=ICCS_DEFAULTS.context_width,
        validator=[validators.gt(Timedelta(seconds=0)), _clear_cache_on_update],
    )
    """Context padding to apply before and after the time window.

    This padding is *not* used for the cross-correlation."""

    ramp_width: NonNegativeTimedelta | NonNegativeNumber = field(
        default=ICCS_DEFAULTS.ramp_width, validator=_clear_cache_on_update
    )
    """Width of taper ramp up and down.

    Can be either a timedelta or a float - see [`pysmo.functions.window()`][pysmo.functions.window]
    for details.
    """

    bandpass_apply: bool = field(
        default=ICCS_DEFAULTS.bandpass_apply, validator=_clear_cache_on_update
    )
    """Filter seismograms with a bandpass filter before running ICCS.

    Setting this to [`True`][] will apply a
    [`bandpass`][pysmo.tools.signal.bandpass] filter (with `zerophase` set to
    [`True`][]) to the [`cc_seismograms`][pysmo.tools.iccs.ICCS.cc_seismograms]
    and [`context_seismograms`][pysmo.tools.iccs.ICCS.context_seismograms].

    As the [`seismograms`][pysmo.tools.iccs.ICCS.seismograms] may have already
    been pre-processed (i.e. already filtered) the default value for this
    parameter is [`False`][].
    """

    bandpass_fmin: float = field(
        default=ICCS_DEFAULTS.bandpass_fmin, validator=_clear_cache_on_update
    )
    """Bandpass filter minimum frequency (Hz). Only used if [`bandpass_apply`][pysmo.tools.iccs.ICCS.bandpass_apply] is `True`."""

    bandpass_fmax: float = field(
        default=ICCS_DEFAULTS.bandpass_fmax, validator=_clear_cache_on_update
    )
    """Bandpass filter maximum frequency (Hz). Only used if [`bandpass_apply`][pysmo.tools.iccs.ICCS.bandpass_apply] is `True`."""

    min_ccnorm: np.floating | float = ICCS_DEFAULTS.min_ccnorm
    """Minimum normalised cross-correlation coefficient for seismograms.

    When the ICCS algorithm is [executed][pysmo.tools.iccs.ICCS.__call__],
    the cross-correlation coefficient for each seismogram is calculated after
    each iteration. If `autoselect` is set to `True`, the
    [`select`][pysmo.tools.iccs.ICCSSeismogram.select] attribute of seismograms
    with with correlation coefficients below this value is set to `False`, and
    they are no longer used for the [`stack`][pysmo.tools.iccs.ICCS.stack].
    """

    # The following attributes are cached to prevent unnecessary processing.
    # Setting the caches to None will force a new calculation when they are
    # requested.
    _cc_seismograms: list[_EphemeralSeismogram] | None = field(init=False)
    _context_seismograms: list[_EphemeralSeismogram] | None = field(init=False)
    _ccnorms: np.ndarray | None = field(init=False)
    _cc_stack: MiniSeismogram | None = field(init=False)
    _context_stack: MiniSeismogram | None = field(init=False)
    _valid_pick_range: tuple[timedelta, timedelta] | None = field(init=False)
    _valid_time_window_range: tuple[Timedelta, Timedelta] | None = field(init=False)

    def __attrs_post_init__(self) -> None:
        self._clear_caches()

    def _clear_caches(self) -> None:
        """Clear all cached attributes."""
        self._cc_seismograms = None
        self._context_seismograms = None
        self._ccnorms = None
        self._cc_stack = None
        self._context_stack = None
        self._valid_pick_range = None
        self._valid_time_window_range = None

    def _prepare_seismograms(
        self,
        add_context: bool = False,
    ) -> list[_EphemeralSeismogram]:
        """Prepare cc_seismograms or context_seismograms."""

        ephemeral_seismograms: list[_EphemeralSeismogram] = []

        min_delta = min((s.delta for s in self.seismograms))

        for seismogram in self.seismograms:
            pick = seismogram.t1 or seismogram.t0
            window_start = pick + self.window_pre
            window_end = pick + self.window_post

            ephemeral_seismogram = _EphemeralSeismogram(parent_seismogram=seismogram)

            if self.bandpass_apply:
                bandpass(
                    ephemeral_seismogram,
                    self.bandpass_fmin,
                    self.bandpass_fmax,
                    zerophase=True,
                )

            if not np.isclose(
                ephemeral_seismogram.delta.total_seconds(), min_delta.total_seconds()
            ):
                resample(ephemeral_seismogram, min_delta)

            if add_context:
                context_window_start = window_start - self.context_width
                context_window_end = window_end + self.context_width

                if (
                    context_window_start < seismogram.begin_time
                    or context_window_end > seismogram.end_time
                ):
                    pad(
                        ephemeral_seismogram,
                        context_window_start,
                        context_window_end,
                        mode="linear_ramp",
                        end_values=(0, 0),
                    )

                crop(ephemeral_seismogram, context_window_start, context_window_end)
                detrend(ephemeral_seismogram)
                normalize(ephemeral_seismogram, window_start, window_end)
            else:
                window(ephemeral_seismogram, window_start, window_end, self.ramp_width)
                normalize(ephemeral_seismogram)

            if seismogram.flip:
                ephemeral_seismogram.data *= -1

            ephemeral_seismograms.append(ephemeral_seismogram)

        # If all seismograms have the same length, return them now.
        if len(lengths := set(len(s.data) for s in ephemeral_seismograms)) == 1:
            return ephemeral_seismograms

        # Shorten seismograms if necessary and return (floating-point precision
        # can cause small differences in length after resampling). We cut off
        # at the end so the `begin_time` will not change.
        for s in ephemeral_seismograms:
            s.data = s.data[: min(lengths)]
        return ephemeral_seismograms

    @property
    def cc_seismograms(self) -> list[_EphemeralSeismogram]:
        """Returns the seismograms as used for the cross-correlation.

        These seismograms are derived from the input seismograms and used for
        the cross-correlation steps. Starting with the input seismograms, they
        are processed as follows:

        1. Bandpass filtered if
           [`bandpass_apply`][pysmo.tools.iccs.ICCS.bandpass_apply] is
           [`True`][].
        2. Resampled to the minimum sampling interval of all input seismograms
           (only if it is not equal in all seismograms).
        3. Cropped to `ramp_width` +  current time window + `ramp_width`.
        4. Detrended.
        5. Tapered using [`ramp_width`][pysmo.tools.iccs.ICCS.ramp_width]
           (tapered sections are *outside* time window).
        6. Normalised based on the highest absolute value within the cropped
           window. This step is done slightly differently in
           [`context_seismograms`][pysmo.tools.iccs.ICCS.context_seismograms]
           --see the documentation of that property for details.
        """

        if self._cc_seismograms is None:
            self._cc_seismograms = self._prepare_seismograms(add_context=False)
        return self._cc_seismograms

    @property
    def context_seismograms(self) -> list[_EphemeralSeismogram]:
        """Returns the seismograms with extra context for plotting.

        These seismograms are derived from the input seismograms and used
        primarily for plotting with extra context (e.g. when selecting new
        time window boundaries). Starting with the input seismograms, they are
        processed as follows:

        1. Bandpass filtered if
           [`bandpass_apply`][pysmo.tools.iccs.ICCS.bandpass_apply] is
           [`True`][].
        2. Resampled to the minimum sampling interval of all input seismograms
           (only if it is not equal in all seismograms).
        3. Cropped and/or padded to `context_width` +  current time window +
           `context_width`.
        4. Detrended.
        5. Normalised based on the highest absolute value within the selected
           time window (i.e. without the context).
        """

        if self._context_seismograms is None:
            self._context_seismograms = self._prepare_seismograms(add_context=True)
        return self._context_seismograms

    @property
    def ccnorms(self) -> np.ndarray:
        """Returns an array of the normalised cross-correlation coefficients."""

        if self._ccnorms is None:
            matrix = np.array([s.data for s in self.cc_seismograms])
            self._ccnorms = pearson_matrix_vector(matrix, self.stack.data)
        return self._ccnorms

    @property
    def stack(self) -> MiniSeismogram:
        """Returns the stacked [`cc_seismograms`][pysmo.tools.iccs.ICCS.cc_seismograms]).

        The stack is calculated as the average of all seismograms with the
        attribute [`select`][pysmo.tools.iccs.ICCSSeismogram.select] set to
        [`True`][True]. The [`begin_time`][pysmo.MiniSeismogram.begin_time] of
        the returned stack is the average of the [`begin_time`]
        [pysmo.tools.iccs.ICCSSeismogram.begin_time] of the input seismograms.

        Returns:
            Stacked input seismograms.
        """
        if self._cc_stack is None:
            self._cc_stack = _create_stack(self.cc_seismograms)
        return self._cc_stack

    @property
    def context_stack(self) -> MiniSeismogram:
        """Returns the stacked [`context_seismograms`][pysmo.tools.iccs.ICCS.context_seismograms].

        Returns:
            Stacked input seismograms with context padding.
        """
        if self._context_stack is None:
            self._context_stack = _create_stack(self.context_seismograms)
        return self._context_stack

    def validate_pick(self, pick: Timedelta) -> bool:
        """Check if a new pick is valid.

        This checks if a new manual pick is valid for all selected seismograms.

        Args:
            pick: New pick to validate.

        Returns:
            Whether the new pick is valid.
        """

        # first calculate valid pick range if not done yet
        if self._valid_pick_range is None:
            self._valid_pick_range = _calc_valid_pick_range(self)

        return self._valid_pick_range[0] <= pick <= self._valid_pick_range[1]

    def validate_time_window(
        self, window_pre: Timedelta, window_post: Timedelta
    ) -> bool:
        """Check if a new time window (relative to pick) is valid.

        Args:
            window_pre: New window start time to validate.
            window_post: New window end time to validate.

        Returns:
            Whether the new time window is valid.
        """

        if window_pre >= window_post:
            return False

        if window_pre > -self.stack.delta:
            return False

        if window_post < self.stack.delta:
            return False

        # Calculate valid time window range if not done yet
        if self._valid_time_window_range is None:
            self._valid_time_window_range = _calc_valid_time_window_range(self)

        return (
            self._valid_time_window_range[0] <= window_pre
            and window_post <= self._valid_time_window_range[1]
        )

    def __call__(
        self,
        autoflip: bool = False,
        autoselect: bool = False,
        convergence_limit: float = ICCS_DEFAULTS.convergence_limit,
        convergence_method: ConvergenceMethod = ICCS_DEFAULTS.convergence_method,
        max_iter: int = ICCS_DEFAULTS.max_iter,
        max_shift: Timedelta | None = None,
    ) -> np.ndarray:
        """Run the iccs algorithm.

        Args:
            autoflip: Automatically toggle [`flip`][pysmo.tools.iccs.ICCSSeismogram.flip] attribute of seismograms.
            autoselect: Automatically set `select` attribute to `False` for poor quality seismograms.
            convergence_limit: Convergence limit at which the algorithm stops.
            convergence_method: Method to calculate convergence criterion.
            max_iter: Maximum number of iterations.
            max_shift: Maximum shift in seconds (see [`delay()`][pysmo.tools.signal.delay]).

        Returns:
            convergence: Array of convergence criterion values.
        """
        convergence_list = []

        for _ in range(max_iter):
            # Save the previous stack to calculate convergence criterion after updating the seismograms.
            prev_stack = clone_to_mini(MiniSeismogram, self.stack)

            # Get delays and correlation coefficients for all seismograms in one go
            delays, ccnorms = multi_delay(
                self.stack, self.cc_seismograms, abs_max=autoflip
            )

            # Update seismograms based on results and settings.
            for delay, ccnorm, cc_seismogram in zip(
                delays, ccnorms, self.cc_seismograms
            ):
                _update_seismogram(
                    delay,
                    ccnorm,
                    cc_seismogram.parent_seismogram,
                    autoflip,
                    autoselect,
                    self.min_ccnorm,
                    (self.window_pre, self.window_post),
                )

            self._clear_caches()

            convergence = _calc_convergence(self.stack, prev_stack, convergence_method)
            convergence_list.append(convergence)
            if convergence <= convergence_limit:
                break

        return np.array(convergence_list)

@runtime_checkable
class ICCSSeismogram(Seismogram, Protocol):
    """Protocol class to define the `ICCSSeismogram` type.

    The `ICCSSeismogram` type extends the [`Seismogram`][pysmo.Seismogram] type
    with the addition of parameters that are required for ICCS.
    """

    t0: Timestamp
    """Initial pick."""

    t1: Timestamp | None
    """Updated pick."""

    flip: bool
    """Data in seismogram should be flipped for ICCS."""

    select: bool
    """Use seismogram to create stack."""

@beartype
@define(kw_only=True, slots=True)
class MiniICCSSeismogram(SeismogramEndtimeMixin, ICCSSeismogram):
    """Minimal implementation of the [`ICCSSeismogram`][pysmo.tools.iccs.ICCSSeismogram] type.

    The [`MiniICCSSeismogram`][pysmo.tools.iccs.ICCSSeismogram] class provides
    a minimal implementation of a class that is compatible with the
    [`ICCSSeismogram`][pysmo.tools.iccs.ICCSSeismogram] protocol.

    Examples:
        Because [`ICCSSeismogram`][pysmo.tools.iccs.ICCSSeismogram] inherits
        from [`Seismogram`][pysmo.Seismogram], we can easily create
        [`MiniICCSSeismogram`][pysmo.tools.iccs.MiniICCSSeismogram] instances
        from existing seismograms using the
        [`clone_to_mini()`][pysmo.functions.clone_to_mini] function, whereby
        the `update` parameter is used to provide the extra information needed:

        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.functions import clone_to_mini
        >>> from pysmo.tools.iccs import MiniICCSSeismogram
        >>> from pandas import Timedelta
        >>> sac = SAC.from_file("example.sac")
        >>> sac_seis = sac.seismogram
        >>> # Use existing pick or set a new one 10 seconds after begin time
        >>> update = {"t0": sac.timestamps.t0 or sac_seis.begin_time + Timedelta(seconds=10)}
        >>> mini_iccs_seis = clone_to_mini(MiniICCSSeismogram, sac_seis, update=update)
        >>>
        ```
    """

    begin_time: Timestamp = field(
        default=SEISMOGRAM_DEFAULTS.begin_time.value, validator=datetime_is_utc
    )
    """Seismogram begin time."""

    delta: PositiveTimedelta = SEISMOGRAM_DEFAULTS.delta.value
    """Seismogram sampling interval."""

    data: np.ndarray = field(factory=lambda: np.array([]))
    """Seismogram data."""

    t0: Timestamp = field(validator=datetime_is_utc)
    """Initial pick."""

    t1: Timestamp | None = field(
        default=None, validator=validators.optional(datetime_is_utc)
    )
    """Updated pick."""

    flip: bool = False
    """Data in seismogram should be flipped for ICCS."""

    select: bool = True
    """Use seismogram to create stack."""