Python 時間波形を時間方向にシフトさせるコード

import numpy as np
import matplotlib.pyplot as plt

from scipy.ndimage import shift
from typing import List


def shift_waveform_interpolation(
    original_waveform: np.ndarray, shift_time: float, sampling_time: float
) -> np.ndarray:
    """
    Shift a waveform using interpolation.

    Parameters
    ----------
    original_waveform : np.ndarray
        The original waveform.
    shift_time : float
        The amount of time by which the waveform should be shifted.
    sampling_time : float
        The sampling time of the waveform.

    Returns
    -------
    out : np.ndarray
        The shifted waveform.
    """
    t = np.arange(0, len(original_waveform) * sampling_time, sampling_time)
    t_shifted = t + shift_time
    return np.interp(t, t_shifted, original_waveform, left=0, right=0)


def shift_waveform_ndimage(
    waveform: np.ndarray, shift_time: float, sampling_time: float
) -> np.ndarray:
    """
    Shift a waveform using scipy.ndimage's shift function.

    Parameters
    ----------
    waveform : np.ndarray
        The original waveform.
    shift_time : float
        The amount of time by which the waveform should be shifted.
    sampling_time : float
        The sampling time of the waveform.

    Returns
    -------
    out : np.ndarray
        The shifted waveform.
    """
    shift_amount = int(shift_time / sampling_time)
    return shift(waveform, shift_amount, cval=0)


def shift_waveform(
    waveform: np.ndarray, shift_time: float, sampling_time: float
) -> np.ndarray:
    """
    Shift a waveform by a specified time.

    Parameters
    ----------
    waveform : np.ndarray
        The original waveform.
    shift_time : float
        The amount of time by which the waveform should be shifted.
    sampling_time : float
        The sampling time of the waveform.

    Returns
    -------
    out : np.ndarray
        The shifted waveform.
    """
    # シフト量をサンプル数に変換
    shift_amount = int(shift_time / sampling_time)

    # 0で埋められた配列を用意
    shifted_waveform = np.zeros_like(waveform)

    # シフト量が正なら右に、負なら左にシフト
    if shift_amount > 0:
        shifted_waveform[shift_amount:] = waveform[:-shift_amount]
    elif shift_amount < 0:
        shifted_waveform[:shift_amount] = waveform[-shift_amount:]

    return shifted_waveform


def shift_waveform_concatenation(
    original_waveform: np.ndarray, shift_time: float, sampling_time: float
) -> np.ndarray:
    """
    Shift a waveform by concatenating zeros.

    Parameters
    ----------
    original_waveform : np.ndarray
        The original waveform.
    shift_time : float
        The amount of time by which the waveform should be shifted.
    sampling_time : float
        The sampling time of the waveform.

    Returns
    -------
    out : np.ndarray
        The shifted waveform.
    """
    # シフト量をサンプル数に変換
    shift_amount = int(shift_time / sampling_time)

    # 0で埋められた配列を用意し、元の波形を連結
    shifted_waveform = np.concatenate(
        [np.zeros(shift_amount), original_waveform[:-shift_amount]]
    )

    return shifted_waveform


def plot_shifted_waveforms(
    original_waveform: np.ndarray,
    shifted_waveforms: List[np.ndarray],
    shift_times: List[float],
    method_names: List[str],
) -> None:
    """
    Plot the original and shifted waveforms obtained using different methods.

    Parameters
    ----------
    original_waveform : np.ndarray
        The original waveform.
    shifted_waveforms : List of np.ndarray
        List of shifted waveforms obtained using different methods.
    shift_times : List of float
        List of shift times for each method.
    method_names : List of str
        List of names for each method.
    """
    t = np.arange(0, len(original_waveform) * 0.1, 0.1)

    plt.plot(t, original_waveform, label="Original Waveform")

    for shifted_waveform, shift_time, method_name in zip(
        shifted_waveforms, shift_times, method_names
    ):
        plt.plot(
            t,
            shifted_waveform,
            label=f"{method_name} Shift by {shift_time}s",
        )

    plt.xlabel("Time")
    plt.ylabel("Amplitude")
    plt.legend()
    plt.show()


# オリジナルのガウス波形を生成
t = np.arange(0, 100, 0.1)
original_gaussian = np.exp(-((t - 50) ** 2) / (2 * 10**2))

# サンプリング時間
sampling_time = 0.1

# シフトする時間 T[s]
shift_time = 50

# 各関数で波形をシフト
shifted_waveform_interp = shift_waveform_interpolation(
    original_gaussian, shift_time, sampling_time
)
shifted_waveform_ndimage = shift_waveform_ndimage(
    original_gaussian, shift_time, sampling_time
)
shifted_waveform_custom = shift_waveform(original_gaussian, shift_time, sampling_time)
shifted_gaussian_concatenation = shift_waveform_concatenation(
    original_gaussian, shift_time, sampling_time
)

# 波形を比較してプロット
plot_shifted_waveforms(
    original_gaussian,
    [
        shifted_waveform_interp,
        shifted_waveform_ndimage,
        shifted_waveform_custom,
        shifted_gaussian_concatenation,
    ],
    [shift_time, shift_time, shift_time, shift_time],
    ["Interpolation", "ndimage", "Custom", "Concatennation"],
)

print(np.all(shifted_waveform_interp == shifted_waveform_ndimage))  # False
print(np.all(shifted_waveform_interp == shifted_waveform_custom))  # False
print(np.all(shifted_waveform_interp == shifted_gaussian_concatenation))  # False
print(np.all(shifted_waveform_ndimage == shifted_waveform_custom))  # False
print(np.all(shifted_waveform_ndimage == shifted_gaussian_concatenation))  # False
print(np.all(shifted_waveform_custom == shifted_gaussian_concatenation))  # True