S2Generator 0.0.7__tar.gz → 0.0.9__tar.gz

{s2generator-0.0.7 → s2generator-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: S2Generator
-Version: 0.0.7
+Version: 0.0.9
 Summary: A series-symbol (S2) dual-modality data generation mechanism, enabling the unrestricted creation of high-quality time series data paired with corresponding symbolic representations.
 Home-page: https://github.com/wwhenxuan/S2Generator
 Author: whenxuan, johnfan12, changewam

{s2generator-0.0.7 → s2generator-0.0.9}/S2Generator.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: S2Generator
-Version: 0.0.7
+Version: 0.0.9
 Summary: A series-symbol (S2) dual-modality data generation mechanism, enabling the unrestricted creation of high-quality time series data paired with corresponding symbolic representations.
 Home-page: https://github.com/wwhenxuan/S2Generator
 Author: whenxuan, johnfan12, changewam

{s2generator-0.0.7 → s2generator-0.0.9}/S2Generator.egg-info/SOURCES.txt

@@ -10,6 +10,14 @@ s2generator/__init__.py
 s2generator/base.py
 s2generator/encoders.py
 s2generator/generators.py
+s2generator/augmentation/__init__.py
+s2generator/augmentation/_amplitude_modulation.py
+s2generator/augmentation/_censor_augmentation.py
+s2generator/augmentation/_empirical_mode_modulation.py
+s2generator/augmentation/_frequency_perturbation.py
+s2generator/augmentation/_spike_injection.py
+s2generator/augmentation/_time_transformation.py
+s2generator/augmentation/_wiener_filter.py
 s2generator/excitation/__init__.py
 s2generator/excitation/_interface.py
 s2generator/excitation/autoregressive_moving_average.py
@@ -36,6 +44,7 @@ s2generator/utils/visualization.py
 tests/__init__.py
 tests/test_all.py
 tests/test_arma.py
+tests/test_augmentation.py
 tests/test_base_excitation.py
 tests/test_forecast_pfn.py
 tests/test_intrinsic_mode_functions.py
@@ -45,4 +54,5 @@ tests/test_series_params.py
 tests/test_stl.py
 tests/test_symbol_params.py
 tests/test_tools.py
-tests/test_wasserstein.py
+tests/test_wasserstein.py
+tests/test_wiener_filter_simulator.py

{s2generator-0.0.7 → s2generator-0.0.9}/s2generator/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "0.0.7"
+__version__ = "0.0.9"
 
 __all__ = [
     "Node",
@@ -17,6 +17,7 @@ __all__ = [
     "print_hello",
     "excitation",
     "simulator",
+    "augmentation",
     "utils",
     "params",
 ]

s2generator-0.0.9/s2generator/augmentation/__init__.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/02 12:15:45
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+
+__all__ = [
+    "amplitude_modulation",
+    "censor_augmentation",
+    "empirical_mode_modulation",
+    "frequency_perturbation",
+    "spike_injection",
+    "wiener_filter",
+    "add_linear_trend",
+    "time_series_mixup",
+]
+
+# Import the amplitude modulation function
+from ._amplitude_modulation import amplitude_modulation
+
+# Import the censoring augmentation function
+from ._censor_augmentation import censor_augmentation
+
+# Import the empirical mode modulation function
+from ._empirical_mode_modulation import empirical_mode_modulation
+
+# Import the frequency perturbation function
+from ._frequency_perturbation import frequency_perturbation
+
+# Import the spike injection function
+from ._spike_injection import spike_injection
+
+# Import the wiener filter function
+from ._wiener_filter import wiener_filter
+
+# Import the time transformation functions
+from ._time_transformation import add_linear_trend, time_series_mixup

s2generator-0.0.9/s2generator/augmentation/_amplitude_modulation.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/04 22:52:40
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+
+import numpy as np
+
+from s2generator.utils._tools import (
+    linear_interpolation,
+    cubic_spline_interpolation,
+    lagrange_interpolation,
+)
+
+
+def amplitude_modulation(
+    time_series: np.ndarray,
+    num_changepoints: int = 5,
+    amplitude_mean: float = 1.0,
+    amplitude_variation: float = 1.0,
+    interpolation_method: str = "linear",
+    rng: np.random.RandomState = None,
+    seed: int = 42,
+) -> np.ndarray:
+    """
+    Perform amplitude modulation on the input time series.
+    This augmentation introduces scale trends and change points into the time series
+    by multiplying the signal with a piecewise linear trend.
+    The modulation trend is generated by sampling change points and interpolating amplitudes between them.
+
+    :param time_series: Input time series, a 1D numpy array
+    :param num_changepoints: Number of change points to introduce in the modulation trend.
+    :param amplitude_mean: The mean amplitude of the modulation trend.
+    :param amplitude_variation: The variation of the amplitude around the mean.
+    :param interpolation_method: The method to interpolate the modulation trend.
+                                 Options are "linear", "cubic", or "lagrange".
+    :param rng: Optional random number generator for reproducibility. If None, a new RNG will be created using the provided seed.
+    :param seed: Random seed for reproducibility if rng is not provided.
+
+    :return: Amplitude modulated time series, a 1D numpy array of the same length as the input series.
+    """
+    # Validate the input time series
+    time_series = np.asarray(time_series)
+    if time_series.ndim != 1:
+        raise ValueError("Input time_series must be a 1D array.")
+
+    # Validate interpolation method
+    if interpolation_method not in ["linear", "cubic", "lagrange"]:
+        raise ValueError(
+            "interpolation_method must be one of 'linear', 'cubic', or 'lagrange'."
+        )
+
+    # Get the length of the time series
+    n = len(time_series)
+
+    # Validate num_changepoints
+    if num_changepoints < 2:
+        raise ValueError(
+            "num_changepoints must be at least 2 to create a modulation trend."
+        )
+    if num_changepoints > n:
+        raise ValueError(
+            "num_changepoints cannot exceed the length of the time series."
+        )
+
+    # Initialize random number generator
+    if rng is None:
+        rng = np.random.RandomState(seed)
+
+    # Sample change points and their corresponding amplitudes
+    changepoints = np.hstack(
+        [
+            0,
+            np.sort(rng.choice(n - 2, size=num_changepoints - 2, replace=False) + 1),
+            n - 1,
+        ]
+    )
+
+    # Generate random amplitudes for each change point
+    amplitude = rng.normal(
+        loc=amplitude_mean, scale=amplitude_variation, size=num_changepoints
+    )
+
+    # Interpolate the modulation trend across the entire time series
+    if interpolation_method == "linear":
+        modulation_trend = linear_interpolation(
+            x_known=changepoints, y_known=amplitude, x_new=np.arange(n)
+        )
+    elif interpolation_method == "cubic":
+        modulation_trend = cubic_spline_interpolation(
+            x_known=changepoints, y_known=amplitude, x_new=np.arange(n)
+        )
+    elif interpolation_method == "lagrange":
+        modulation_trend = lagrange_interpolation(
+            x_known=changepoints, y_known=amplitude, x_new=np.arange(n)
+        )
+
+    # Apply the modulation trend to the original time series
+    modulated_series = time_series * modulation_trend
+
+    return modulated_series, np.array(modulation_trend)

s2generator-0.0.9/s2generator/augmentation/_censor_augmentation.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/05 00:42:51
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+
+import numpy as np
+
+
+def censor_augmentation(
+    time_series: np.ndarray,
+    upper_quantile: float = 0.65,
+    lower_quantile: float = 0.35,
+    bernoulli_p: float = 0.8,
+    rng: np.random.RandomState = None,
+    seed: int = 42,
+) -> np.ndarray:
+    """
+    Perform censoring augmentation on the input time series.
+
+    This augmentation censors (clips) the input signal either from below or
+    above, depending on a randomly sampled direction. The clipping threshold is determined by drawing
+    a quantile uniformly from the empirical distribution of the signal.
+
+    :param time_series: Input time series, a 1D numpy array
+    :param upper_quantile: Upper quantile threshold for censoring, default is 0.65
+    :param lower_quantile: Lower quantile threshold for censoring, default is 0.35
+    :param bernoulli_p: Probability of censoring direction (0 for lower censoring, 1 for upper censoring),
+                        the default is 0.5, meaning equal probability for both directions.
+    :param rng: Optional random number generator for reproducibility.
+                If None, a new RNG will be created using the provided seed.
+    :param seed: Random seed for reproducibility if rng is not provided.
+
+    :return: Censored time series, a 1D numpy array of the same length as the input series.
+    """
+
+    # Validate the input time series
+    time_series = np.asarray(time_series)
+    if time_series.ndim != 1:
+        raise ValueError("Input time_series must be a 1D array.")
+
+    # Validate bernoulli_p
+    if not (0 <= bernoulli_p <= 1):
+        raise ValueError("bernoulli_p must be in the range [0, 1].")
+
+    # Get the length of the time series
+    length = time_series.shape[0]
+
+    # Set random seed for reproducibility
+    if rng is None:
+        rng = np.random.RandomState(seed)
+
+    # Randomly sample quantile thresholds for each time step
+    quantile_threshold = rng.uniform(lower_quantile, upper_quantile, size=length)
+
+    # Compute the threshold value based on the quantile of the time series
+    threshold_value = np.quantile(time_series, quantile_threshold)
+
+    # Sample the censor direction from bernoulli distribution (0.5)
+    censor_direction = rng.binomial(
+        n=1, p=bernoulli_p, size=length
+    )  # 0 for lower censoring, 1 for upper censoring
+
+    for t in range(length):
+        if censor_direction[t] == 1:
+            # Lower censoring
+            time_series[t] = max(time_series[t], threshold_value[t])
+        else:
+            # Upper censoring
+            time_series[t] = min(time_series[t], threshold_value[t])
+
+    return time_series

s2generator-0.0.9/s2generator/augmentation/_empirical_mode_modulation.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/05 11:05:45
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+from typing import Optional
+
+import numpy as np
+
+from pysdkit import EMD
+
+
+def empirical_mode_modulation(
+    time_series: np.ndarray,
+    min_scale_factor: float = 0.5,
+    max_scale_factor: float = 2.0,
+    low_frequency_enhancement: bool = True,
+    spline_kind: str = "cubic",
+    extrema_detection: str = "parabol",
+    max_imfs: Optional[int] = None,
+    rng: Optional[np.random.RandomState] = None,
+    seed: int = 42,
+) -> np.ndarray:
+    """
+    Perform empirical mode modulation on the input time series.
+    This augmentation decomposes the time series into intrinsic mode functions (IMFs)
+    using Empirical Mode Decomposition (EMD) and then reconstructs the signal
+    by randomly modifying the IMFs to introduce non-linear trends and variations.
+
+    :param time_series: Input time series, a 1D numpy array.
+    :param min_scale_factor: The minimum scaling factor to apply to each IMF.
+    :param max_scale_factor: The maximum scaling factor to apply to each IMF.
+    :param low_frequency_enhancement: Whether to enhance the low-frequency components.
+                                      If True, the method will suppress the high-frequency noise in the
+                                      empirical mode decomposition results and focus on enhancing
+                                      the characterization of the low-frequency part.
+    :param spline_kind: The kind of spline to use for interpolation,
+                        options are "akima", "cubic", "pchip", "cubic_hermite", "slinear", "quadratic", "linear"
+    :param extrema_detection: The method for detecting extrema in the EMD process, options are "parabol" or "simple"
+    :param max_imfs: The maximum number of IMFs to extract,
+                     if None, it will extract all possible IMFs until the residue is a monotonic function.
+
+    :return: Empirical mode modulated time series, a 1D numpy array of the same length as the input series.
+    """
+    # Validate the input time series
+    time_series = np.asarray(time_series)
+    if time_series.ndim != 1:
+        raise ValueError("Input time_series must be a 1D array.")
+
+    # Normalize the input time series to have zero mean and unit variance
+    mean, std = np.mean(time_series), np.std(time_series)
+    time_series = (time_series - mean) / (
+        std + 1e-8
+    )  # Add a small value to avoid division by zero
+
+    # Check if max_imfs is valid
+    if max_imfs is None:
+        # 表示会完整的分解所有的IMF分量
+        # 直到剩余的分量不再满足IMF的定义为止
+        max_imfs = -1
+
+    # Validate the spline kind
+    assert spline_kind in [
+        "akima",
+        "cubic",
+        "pchip",
+        "cubic_hermite",
+        "slinear",
+        "quadratic",
+        "linear",
+    ], "spline_kind must be one of 'akima', 'cubic', 'pchip', 'cubic_hermite', 'slinear', 'quadratic', 'linear'."
+
+    # Validate the extrema detection method
+    assert extrema_detection in [
+        "parabol",
+        "simple",
+    ], "extrema_detection must be one of 'parabol' or 'simple'."
+
+    # Initialize random number generator
+    if rng is None:
+        rng = np.random.RandomState(seed=seed)
+
+    # Perform Empirical Mode Decomposition
+    emd = EMD(
+        max_imfs=max_imfs, spline_kind=spline_kind, extrema_detection=extrema_detection
+    )
+    print(type(time_series))
+    imfs = emd.fit_transform(signal=time_series)
+
+    # Get the number of IMFs extracted
+    num_imfs = imfs.shape[0]
+
+    # Randomly select a scaling factor for modulation
+    scale_factor = rng.uniform(
+        low=min_scale_factor, high=max_scale_factor, size=num_imfs
+    )
+
+    # Validate the low_frequency_enhancement parameter
+    assert isinstance(
+        low_frequency_enhancement, bool
+    ), "low_frequency_enhancement must be a boolean value."
+    if low_frequency_enhancement is True:
+        # If low-frequency enhancement is enabled, we can apply a stronger scaling to the lower frequency IMFs
+        scale_factor = np.sort(
+            scale_factor
+        )  # Sort the scale factors to enhance low-frequency components more than high-frequency ones
+
+    # Randomly modify the IMFs to create modulation
+    modified_imfs = []
+    for index, imf in enumerate(imfs):
+        # Randomly scale each IMF by a factor between min_scale_factor and max_scale_factor
+        scaled_imf = imf * scale_factor[index]
+        modified_imfs.append(scaled_imf)
+
+    # Reconstruct the signal from the modified IMFs
+    modulated_time_series = np.sum(modified_imfs, axis=0)
+
+    # Denormalize the modulated time series to restore the original scale
+    modulated_time_series = (modulated_time_series - np.mean(modulated_time_series)) / (
+        np.std(modulated_time_series) + 1e-8
+    )
+    modulated_time_series = modulated_time_series * (std + 1e-8) + mean
+
+    return modulated_time_series

s2generator-0.0.9/s2generator/augmentation/_frequency_perturbation.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/02 12:16:05
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+
+import numpy as np
+from numpy import fft
+
+
+def sample_random_perturbation(
+    K: int, min_alpha: float, max_alpha: float, rng: np.random.RandomState = None
+) -> np.ndarray:
+    """
+    Randomly sample K numbers in the interval [-alpha_max, -alpha_min] U [alpha_min, alpha_max]
+    The purpose of this sampling is to construct random perturbations in the frequency domain.
+
+    :param K: Number of random numbers to sample
+    :param min_alpha: Minimum absolute value of the random numbers
+    :param max_alpha: Maximum absolute value of the random numbers
+    :param rng: Optional random number generator, if not provided, the global numpy random number generator will be used
+
+    :return: A numpy array containing K random numbers, which are uniformly distributed in the interval [-alpha_max, -alpha_min] U [alpha_min, alpha_max]
+    """
+
+    # First generate random numbers in [alpha_min, alpha_max]
+    if rng is not None:
+        positive_rand = rng.uniform(min_alpha, max_alpha, K)
+
+        # Randomly generate sign (-1 or 1)
+        signs = rng.choice([-1, 1], size=K)
+
+    else:
+        # When the random number generator is not passed in, use the global numpy random number generator
+        positive_rand = np.random.uniform(min_alpha, max_alpha, K)
+        signs = np.random.choice([-1, 1], size=K)
+
+    # Combine to get the final result
+    final_random_nums = positive_rand * signs
+
+    return final_random_nums
+
+
+def frequency_perturbation(
+    time_series: np.ndarray,
+    min_alpha: float,
+    max_alpha: float,
+    r: float = 0.5,
+    rng: np.random.RandomState = None,
+) -> np.ndarray:
+    """
+    Perform frequency domain perturbation on the input time series.
+    This method adds random perturbations to the frequency components of the time series,
+    which can help to enhance the diversity of the data and improve the robustness of models trained on it.
+
+    :param time_series: Input time series, a 1D numpy array
+    :param min_alpha: Minimum absolute value of the random perturbation added to the frequency components
+    :param max_alpha: Maximum absolute value of the random perturbation added to the frequency components
+    :param r: Proportion of frequency components to perturb (default is 0.5, meaning 50% of the frequency components will be perturbed)
+    :param rng: Optional random number generator, if not provided, the global numpy random number generator will be used.
+
+    :return: Perturbed time series, a 1D numpy array of the same length as the input series.
+    """
+    # Validate the input parameters
+    assert 0 <= r <= 1, "The proportion r must be between 0 and 1."
+    assert min_alpha >= 0, "min_alpha must be non-negative."
+    assert (
+        max_alpha >= min_alpha
+    ), "max_alpha must be greater than or equal to min_alpha."
+
+    # Validate that the input time series is ndarray
+    if isinstance(time_series, list):
+        time_series = np.array(time_series)
+
+    # Validate that the input time series is 1D
+    if time_series.ndim != 1:
+        raise ValueError("Input time series must be a 1D array.")
+
+    # Perform Fast Fourier Transform to convert the time series to the frequency domain
+    f = fft.rfft(time_series)
+    f_perturbed = f.copy()
+    frequencies = fft.fftfreq(len(time_series))
+
+    # Calculate the number of frequency domain components that can be perturbed
+    K = int(len(frequencies) * r)
+
+    # Sample random perturbations for the real and imaginary parts of the frequency components
+    alpha_real = sample_random_perturbation(
+        K=K, min_alpha=min_alpha, max_alpha=max_alpha, rng=rng
+    )
+    alpha_imag = sample_random_perturbation(
+        K=K, min_alpha=min_alpha, max_alpha=max_alpha, rng=rng
+    )
+
+    # Randomly select K frequency domain components for perturbation
+    indices = np.random.choice(len(f_perturbed), size=K, replace=False)
+    f_perturbed[indices] += alpha_real + 1j * alpha_imag
+
+    # Perform inverse Fourier transform to restore the original time-domain signal
+    perturbed_series = fft.irfft(f_perturbed).real
+
+    return perturbed_series

s2generator-0.0.9/s2generator/augmentation/_spike_injection.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/05 15:54:19
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+import numpy as np
+
+
+def spike_injection(time_series: np.ndarray) -> np.ndarray:
+    """
+    Perform spike injection augmentation on the input time series.
+    This augmentation randomly injects spikes into the input time series to simulate sudden and extreme events, which can help models learn to handle such anomalies.
+    """
+    pass

s2generator-0.0.9/s2generator/augmentation/_time_transformation.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/05 16:19:59
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+import numpy as np
+
+
+def add_linear_trend(
+    time_series: np.ndarray, trend_strength: float = 1.0, direction: str = "upward"
+) -> np.ndarray:
+    """
+    Perform linear trend augmentation on the input time series.
+    This augmentation adds a linear trend to the input time series,
+    which can help models learn to handle non-stationary data and improve their robustness to trends.
+
+    :param time_series: Input time series, a 1D numpy array
+    :param trend_strength: The strength of the linear trend to be added, default is 1.0.
+    :param direction: The direction of the linear trend, either "upward" or "downward", default is "upward".
+
+    :return: Augmented time series with a linear trend, a 1D numpy array of the same length as the input series.
+    """
+
+    # Get the length of the time series
+    seq_length = len(time_series)
+
+    # Calculate the the energy of the original time series
+    original_energy = np.mean(time_series**2)
+
+    # Create a linear trend
+    if direction == "upward":
+        trend = np.linspace(0, trend_strength * seq_length, seq_length)
+    elif direction == "downward":
+        trend = np.linspace(0, -trend_strength * seq_length, seq_length)
+    else:
+        raise ValueError("direction must be either 'upward' or 'downward'")
+
+    # Scale the trend to have the same energy as the original time series
+    trend_energy = np.mean(trend**2)
+
+    if trend_energy > 0:
+        trend = trend * np.sqrt(original_energy / trend_energy)
+
+    # Average the original signal and the trend to maintain the overall scale
+    return (time_series + trend) / 2
+
+
+def time_series_mixup(a: np.ndarray, b: np.ndarray, alpha: float = 0.7) -> np.ndarray:
+    """
+    Mixup Enhancement: Weighted mixing of two time series to create a new augmented signal.
+    This method combines two time series by taking a weighted average of them,
+    where the weights are determined by a mixing parameter alpha.
+    This can help models learn to generalize better by exposing them to a wider variety of signal combinations.
+
+    :param a: First input time series, a 1D numpy array.
+    :param b: Second input time series, a 1D numpy array of the same length as a.
+    :param alpha: The mixing parameter that controls the weight of each time series in the mixup,
+    default is 0.7. A value of alpha close to 1 gives more weight to the first time series (a),
+    while a value close to 0 gives more weight to the second time series (b).
+
+    :return: Mixed time series, a 1D numpy array of the same length as the input series.
+    """
+    assert a.shape == b.shape, "Input time series must have the same shape"
+
+    # Calculate the mixed signal as a weighted average of the two input signals
+    return alpha * a + (1 - alpha) * b

s2generator-0.0.9/s2generator/augmentation/_wiener_filter.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 2026/03/05 16:12:41
+@author: Whenxuan Wang
+@email: wwhenxuan@gmail.com
+@url: https://github.com/wwhenxuan/S2Generator
+"""
+
+from typing import Optional
+
+import numpy as np
+
+
+def wiener_filter(
+    time_series: np.ndarray,
+    noise_variance: float = 1.0,
+    window_size: Optional[int] = None,
+) -> np.ndarray:
+    """
+    Implement Wiener filter to remove white noise from signals.
+
+    :param time_series: The input time series, a 1D numpy array representing the noisy signal.
+    :param noise_variance: The known variance of the white noise to be removed, default is 1.0
+    :param window_size: The window size for power spectral density estimation, optional
+        if None, it will use the entire signal length as the window size.
+
+    :return: The filtered time series, a 1D numpy array of the same length as the input series.
+    """
+    # Validate the input time series
+    time_series = np.asarray(time_series, dtype=np.float64)
+
+    # If window_size is not provided, use the length of the time series
+    if window_size is None:
+        window_size = len(time_series)
+
+    # Calculate the Fourier transform of the signal
+    signal_fft = np.fft.fft(time_series)
+
+    # Calculate the power spectral density (PSD) of the signal.
+    signal_psd = np.abs(signal_fft) ** 2 / len(time_series)
+
+    # Estimate the power spectrum of the original signal (noisy PSD - noise PSD)
+    # The power spectral density of noise is constant for white noise: noise variance
+    noise_psd = noise_variance
+    original_signal_psd_estimate = np.maximum(
+        signal_psd - noise_psd, 1e-10
+    )  # Avoid negative values
+
+    # Calculate the frequency domain response of the Wiener filter.
+    wiener_filter_freq = original_signal_psd_estimate / (
+        original_signal_psd_estimate + noise_psd
+    )
+
+    # Apply a filter and perform an inverse Fourier transform to return to the time domain.
+    filtered_fft = signal_fft * wiener_filter_freq
+
+    # Remove numerical error by taking the real part
+    filtered_signal = np.fft.ifft(filtered_fft).real
+
+    return filtered_signal

{s2generator-0.0.7 → s2generator-0.0.9}/s2generator/base.py

@@ -376,9 +376,9 @@ class NodeList(object):
         if xs.ndim > 1:
             # For multivariate case, keep other dimensions constant
             x_uniform_input = np.tile(np.mean(xs, axis=0), (n_integration_points, 1))
-            x_uniform_input[:, 0] = (
-                x_uniform  # Replace first dimension with uniform grid
-            )
+            x_uniform_input[
+                :, 0
+            ] = x_uniform  # Replace first dimension with uniform grid
         else:
             x_uniform_input = x_uniform.reshape(-1, 1)  # Ensure 2D array for val method
 

{s2generator-0.0.7 → s2generator-0.0.9}/s2generator/simulator/__init__.py

@@ -6,6 +6,8 @@ Created on 2026/02/13 13:04:42
 @url: https://github.com/wwhenxuan/S2Generator
 """
 
+__all__ = ["ARIMASimulator", "WienerFilterSimulator"]
+
 from .arima import ARIMASimulator
 
 from .wiener_filter import WienerFilterSimulator