aimbat.io

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_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_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 event_creator(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike], AimbatEvent]], Callable[[str | PathLike], AimbatEvent]
]:
    """Decorator that registers a function as an event creator for `datatype`.

    Example:
        ```python
        @event_creator(DataType.SAC)
        def create_event_from_sacfile(sacfile: str | PathLike) -> AimbatEvent:
            ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike], AimbatEvent],
    ) -> Callable[[str | PathLike], AimbatEvent]:
        register_event_creator(datatype, fn)
        return fn

    return decorator

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 returned
    array is read-only; to write new data use `write_seismogram_data`. The
    cache entry is invalidated when `write_seismogram_data` is called for the
    same key.

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

    Returns:
        Read-only 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:
        arr = reader(datasource)
        arr.flags.writeable = False
        _cache[key] = arr
    else:
        logger.debug(f"Retrieved seismogram data from cache for {datasource}.")
    return _cache[key]

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.
    """
    logger.debug(f"Registering event creator for {datatype}.")
    _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.
    """
    logger.debug(f"Registering seismogram creator for {datatype}.")
    _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.
    """
    logger.debug(f"Registering seismogram data reader for {datatype}.")
    _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.
    """
    logger.debug(f"Registering seismogram data writer for {datatype}.")
    _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.
    """
    logger.debug(f"Registering station creator for {datatype}.")
    _station_creators[datatype] = fn

def seismogram_creator(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike], AimbatSeismogram]],
    Callable[[str | PathLike], AimbatSeismogram],
]:
    """Decorator that registers a function as a seismogram creator for `datatype`.

    Example:
        ```python
        @seismogram_creator(DataType.SAC)
        def create_seismogram_from_sacfile(sacfile: str | PathLike) -> AimbatSeismogram:
            ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike], AimbatSeismogram],
    ) -> Callable[[str | PathLike], AimbatSeismogram]:
        register_seismogram_creator(datatype, fn)
        return fn

    return decorator

def seismogram_data_reader(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike], npt.NDArray[np.float64]]],
    Callable[[str | PathLike], npt.NDArray[np.float64]],
]:
    """Decorator that registers a function as a seismogram data reader for `datatype`.

    Example:
        ```python
        @seismogram_data_reader(DataType.SAC)
        def read_seismogram_data_from_sacfile(sacfile: str | PathLike) -> npt.NDArray[np.float64]:
            ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike], npt.NDArray[np.float64]],
    ) -> Callable[[str | PathLike], npt.NDArray[np.float64]]:
        register_seismogram_data_reader(datatype, fn)
        return fn

    return decorator

def seismogram_data_writer(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike, npt.NDArray[np.float64]], None]],
    Callable[[str | PathLike, npt.NDArray[np.float64]], None],
]:
    """Decorator that registers a function as a seismogram data writer for `datatype`.

    Example:
        ```python
        @seismogram_data_writer(DataType.SAC)
        def write_seismogram_data_to_sacfile(
            sacfile: str | PathLike, data: npt.NDArray[np.float64]
        ) -> None:
            ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike, npt.NDArray[np.float64]], None],
    ) -> Callable[[str | PathLike, npt.NDArray[np.float64]], None]:
        register_seismogram_data_writer(datatype, fn)
        return fn

    return decorator

def station_creator(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike], AimbatStation]],
    Callable[[str | PathLike], AimbatStation],
]:
    """Decorator that registers a function as a station creator for `datatype`.

    Example:
        ```python
        @station_creator(DataType.SAC)
        def create_station_from_sacfile(sacfile: str | PathLike) -> AimbatStation:
            ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike], AimbatStation],
    ) -> Callable[[str | PathLike], AimbatStation]:
        register_station_creator(datatype, fn)
        return fn

    return decorator

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)