aimbat.io

def clear_seismogram_cache() -> None:
    """Clear the in-memory seismogram data cache."""
    _cache.clear()

def create_event(datasource: str | PathLike, datatype: DataType) -> AimbatEvent:
    """Create an `AimbatEvent` from a data source.

    Args:
        datasource: Data source path or name.
        datatype: Data type of the source.

    Returns:
        A new `AimbatEvent` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered event creator.
    """
    logger.debug(f"Creating AimbatEvent from {datasource}.")
    creator = _event_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support event creation.")
    return creator(datasource)

def create_event_from_sacfile(sacfile: str | PathLike) -> AimbatEvent:
    """Create an `AimbatEvent` instance from a SAC file.

    Args:
        sacfile: Name of the SAC file.

    Returns:
        A new `AimbatEvent` instance.
    """

    from aimbat.models import AimbatEvent, AimbatEventParameters

    logger.debug(f"Reading event data from {sacfile}.")

    event = SAC.from_file(sacfile).event
    aimbat_event = AimbatEvent.model_validate(
        event, update={"parameters": AimbatEventParameters()}
    )
    return aimbat_event

def create_seismogram(
    datasource: str | PathLike, datatype: DataType
) -> AimbatSeismogram:
    """Create an `AimbatSeismogram` from a data source.

    Args:
        datasource: Data source path or name.
        datatype: Data type of the source.

    Returns:
        A new `AimbatSeismogram` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered seismogram creator.
    """
    logger.debug(f"Creating AimbatSeismogram from {datasource}.")
    creator = _seismogram_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support seismogram creation.")
    return creator(datasource)

def create_seismogram_from_sacfile_and_pick_header(
    sacfile: str | PathLike, sac_pick_header: str
) -> AimbatSeismogram:
    """Create an AimbatSeismogram instance from a SAC file.

    Args:
        sacfile: Name of the SAC file.
        sac_pick_header: SAC header to use as t0 in AIMBAT.
    """

    from aimbat.models import AimbatSeismogram, AimbatSeismogramParameters

    logger.debug(f"Reading seismogram metadata from {sacfile}.")

    sac = SAC.from_file(sacfile)
    t0 = getattr(sac.timestamps, sac_pick_header)
    seismogram = sac.seismogram
    aimbat_seismogram = AimbatSeismogram.model_validate(
        seismogram, update={"t0": t0, "parameters": AimbatSeismogramParameters()}
    )
    return aimbat_seismogram

def create_station(datasource: str | PathLike, datatype: DataType) -> AimbatStation:
    """Create an `AimbatStation` from a data source.

    Args:
        datasource: Data source path or name.
        datatype: Data type of the source.

    Returns:
        A new `AimbatStation` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered station creator.
    """
    logger.debug(f"Creating AimbatStation from {datasource}.")
    creator = _station_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support station creation.")
    return creator(datasource)

def create_station_from_sacfile(sacfile: str | PathLike) -> AimbatStation:
    """Create an AimbatStation instance from a SAC file.

    Args:
        sacfile: Name of the SAC file.

    Returns:
        A new AimbatStation instance.
    """

    from aimbat.models import AimbatStation

    logger.debug(f"Reading station data from {sacfile}.")

    station = SAC.from_file(sacfile).station
    aimbat_station = AimbatStation.model_validate(station)
    return aimbat_station

def read_seismogram_data(
    datasource: str | PathLike, datatype: DataType
) -> npt.NDArray[np.float64]:
    """Read seismogram waveform data from a data source.

    Results are cached in memory by `(datasource, datatype)` key. The cache
    entry is invalidated when `write_seismogram_data` is called for the same
    key, and can be cleared manually with `clear_seismogram_cache`.

    Args:
        datasource: Data source path or name.
        datatype: Data type of the source.

    Returns:
        Seismogram waveform data as a NumPy array.

    Raises:
        NotImplementedError: If `datatype` has no registered data reader.
    """
    logger.debug(f"Reading seismogram data from {datasource}.")
    reader = _seismogram_data_readers.get(datatype)
    if reader is None:
        raise NotImplementedError(
            f"{datatype} does not support reading seismogram data."
        )
    key = (str(datasource), datatype)
    if key not in _cache:
        _cache[key] = reader(datasource)
    return _cache[key]

def read_seismogram_data_from_sacfile(
    sacfile: str | PathLike,
) -> npt.NDArray[np.float64]:
    """Read seismogram data from a SAC file.

    Args:
        sacfile: Name of the SAC file.

    Returns:
        Seismogram data.
    """

    logger.debug(f"Reading seismogram data from {sacfile}.")

    return SAC.from_file(sacfile).seismogram.data

def register_event_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike], AimbatEvent],
) -> None:
    """Register a function that creates an `AimbatEvent` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a datasource path or name and returns an
            `AimbatEvent` instance.
    """
    _event_creators[datatype] = fn

def register_seismogram_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike], AimbatSeismogram],
) -> None:
    """Register a function that creates an `AimbatSeismogram` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a datasource path or name and returns an
            `AimbatSeismogram` instance.
    """
    _seismogram_creators[datatype] = fn

def register_seismogram_data_reader(
    datatype: DataType,
    fn: Callable[[str | PathLike], npt.NDArray[np.float64]],
) -> None:
    """Register a function that reads seismogram waveform data from a data source.

    Args:
        datatype: The data type this reader handles.
        fn: Callable that accepts a datasource path or name and returns the
            waveform data as a NumPy array.
    """
    _seismogram_data_readers[datatype] = fn

def register_seismogram_data_writer(
    datatype: DataType,
    fn: Callable[[str | PathLike, npt.NDArray[np.float64]], None],
) -> None:
    """Register a function that writes seismogram waveform data to a data source.

    Args:
        datatype: The data type this writer handles.
        fn: Callable that accepts a datasource path or name and a NumPy array,
            and writes the data to the source.
    """
    _seismogram_data_writers[datatype] = fn

def register_station_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike], AimbatStation],
) -> None:
    """Register a function that creates an `AimbatStation` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a datasource path or name and returns an
            `AimbatStation` instance.
    """
    _station_creators[datatype] = fn

def supports_event_creation(datatype: DataType) -> bool:
    """Return whether `datatype` has a registered event creator."""
    return datatype in _event_creators

def supports_seismogram_creation(datatype: DataType) -> bool:
    """Return whether `datatype` has a registered seismogram creator."""
    return datatype in _seismogram_creators

def supports_seismogram_data_reading(datatype: DataType) -> bool:
    """Return whether `datatype` has a registered seismogram data reader."""
    return datatype in _seismogram_data_readers

def supports_seismogram_data_writing(datatype: DataType) -> bool:
    """Return whether `datatype` has a registered seismogram data writer."""
    return datatype in _seismogram_data_writers

def supports_station_creation(datatype: DataType) -> bool:
    """Return whether `datatype` has a registered station creator."""
    return datatype in _station_creators

def write_seismogram_data(
    datasource: str | PathLike,
    datatype: DataType,
    data: npt.NDArray[np.float64],
) -> None:
    """Write seismogram waveform data to a data source.

    Invalidates the cache entry for `(datasource, datatype)` after writing.

    Args:
        datasource: Data source path or name.
        datatype: Data type of the source.
        data: Seismogram waveform data to write.

    Raises:
        NotImplementedError: If `datatype` has no registered data writer.
    """
    logger.debug(f"Writing seismogram data to {datasource}.")
    writer = _seismogram_data_writers.get(datatype)
    if writer is None:
        raise NotImplementedError(
            f"{datatype} does not support writing seismogram data."
        )
    writer(datasource, data)
    _cache.pop((str(datasource), datatype), None)

def write_seismogram_data_to_sacfile(
    sacfile: str | PathLike, data: npt.NDArray[np.float64]
) -> None:
    """Write seismogram data to a SAC file.

    Args:
        sacfile: Name of the SAC file.
        data: Seismogram data.
    """

    logger.debug(f"Writing seismogram data to {sacfile}.")

    sac = SAC.from_file(sacfile)
    sac.seismogram.data = data
    sac.write(sacfile)