gwsim 0.1.0__py3-none-any.whl

gwsim/noise/base.py

@@ -0,0 +1,134 @@
+"""Base class for noise simulators."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import cast
+
+import numpy as np
+
+from gwsim.cli.utils.utils import get_file_name_from_template_with_dict
+from gwsim.mixin.detector import DetectorMixin
+from gwsim.mixin.gwf import GWFOutputMixin
+from gwsim.mixin.randomness import RandomnessMixin
+from gwsim.mixin.time_series import TimeSeriesMixin
+from gwsim.simulator.base import Simulator
+from gwsim.simulator.state import StateAttribute
+from gwsim.utils.random import get_state
+
+
+class NoiseSimulator(
+    RandomnessMixin, DetectorMixin, TimeSeriesMixin, GWFOutputMixin, Simulator
+):  # pylint: disable=duplicate-code
+    """Base class for noise simulators."""
+
+    start_time = StateAttribute(0)
+
+    def __init__(
+        self,
+        sampling_frequency: float,
+        duration: float,
+        start_time: float = 0,
+        max_samples: int | None = None,
+        seed: int | None = None,
+        detectors: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        """Initialize the base noise simulator.
+
+        Args:
+            sampling_frequency: Sampling frequency of the noise in Hz.
+            duration: Duration of each noise segment in seconds.
+            start_time: Start time of the first noise segment in GPS seconds. Default is 0
+            max_samples: Maximum number of samples to generate. None means infinite.
+            seed: Seed for the random number generator. If None, the RNG is not initialized.
+            detectors: List of detector names. Default is None.
+            **kwargs: Additional arguments absorbed by subclasses and mixins.
+        """
+        super().__init__(
+            sampling_frequency=sampling_frequency,
+            duration=duration,
+            start_time=start_time,
+            max_samples=max_samples,
+            seed=seed,
+            detectors=detectors,
+            **kwargs,
+        )
+
+    def save_batch(self, batch: np.ndarray, file_name: str | Path, overwrite: bool = False, **kwargs) -> None:
+        """Save a batch of noise data to a file.
+
+        Args:
+            batch: Batch of noise data to save.
+            file_name: Name of the output file.
+            overwrite: Whether to overwrite existing files. Default is False.
+            **kwargs: Additional arguments for the output mixin.
+
+        Raises:
+            NotImplementedError: If the output mixin does not implement this method.
+        """
+        suffix = Path(file_name).suffix.lower()
+        if suffix == ".gwf":
+            save_function = self.save_batch_to_gwf
+        else:
+            raise NotImplementedError(f"Output format {suffix} not supported by the output mixin.")
+
+        # Check whether the file_name contains the {detector} placeholder
+        if "{detector}" in str(file_name).replace(" ", ""):
+            # Check whether self.detectors is set
+            if self.detectors is None:
+                raise ValueError(
+                    "The file_name contains the {detector} placeholder, but the simulator does not have detectors set."
+                )
+            # Check whether the dimension of batch matches number of detectors
+            if len(batch.shape) == 1:
+                batch = batch[None, :]
+            # Check whether the length of batch matches number of detectors
+            if batch.shape[0] != len(self.detectors):
+                raise ValueError(
+                    f"The batch has {batch.shape[0]} channels, but the simulator has {len(self.detectors)} detectors."
+                )
+            # Save each detector's data separately
+            for i, detector in enumerate(self.detectors):
+                detector_file_name = get_file_name_from_template_with_dict(
+                    template=str(file_name),
+                    values={
+                        "detector": detector,
+                    },
+                )
+                self.save_batch_to_gwf(
+                    batch=batch[i, :],
+                    file_path=detector_file_name,
+                    overwrite=overwrite,
+                    **kwargs,
+                )
+        else:
+            save_function(
+                batch=batch,
+                file_path=file_name,
+                overwrite=overwrite,
+                **kwargs,
+            )
+
+    @property
+    def metadata(self) -> dict:
+        """Get a dictionary of metadata.
+        This can be overridden by the subclass.
+
+        Returns:
+            dict: A dictionary of metadata.
+        """
+        # Get metadata from all parent classes using cooperative inheritance
+        metadata = super().metadata
+
+        return metadata
+
+    def update_state(self) -> None:
+        """Update internal state after each sample generation.
+
+        This method can be overridden by subclasses to update any internal state
+        after generating a sample. The default implementation does nothing.
+        """
+        self.counter = cast(int, self.counter) + 1
+        self.start_time += self.duration
+        self.rng_state = get_state()

gwsim/noise/bilby_stationary_gaussian.py

@@ -0,0 +1,117 @@
+"""Stationary Gaussian noise simulator using Bilby."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from bilby.core.utils.random import Generator as BilbyGenerator
+from bilby.core.utils.series import infft
+from bilby.gw.detector.psd import PowerSpectralDensity
+from numpy.random import Generator
+
+from gwsim.noise.stationary_gaussian import StationaryGaussianNoiseSimulator
+
+
+class BilbyStationaryGaussianNoiseSimulator(
+    StationaryGaussianNoiseSimulator
+):  # pylint: disable=too-many-ancestors, duplicate-code
+    """Stationary Gaussian noise simulator using Bilby."""
+
+    def __init__(
+        self,
+        frequency_array: np.ndarray[Any, np.dtype[Any]] | None = None,
+        psd_array: np.ndarray[Any, np.dtype[Any]] | None = None,
+        psd_file: str | None = None,
+        sampling_frequency: float = 4096,
+        duration: float = 4,
+        start_time: float = 0,
+        max_samples: int | None = None,
+        seed: int | None = None,
+        detectors: list[str] | None = None,
+        **kwargs,
+    ):
+        """Initialize Bilby stationary Gaussian noise simulator.
+
+        Args:
+            frequency_array: Frequency array for the PSD.
+            psd_array: PSD values corresponding to the frequency array.
+            psd_file: Path to a file containing the PSD.
+            psd: Path to PSD file or numpy array with PSD values, or label of PSD
+            sampling_frequency: Sampling frequency in Hz. Default is 4096.
+            duration: Duration of each segment in seconds. Default is 4.
+            start_time: Start time in GPS seconds. Default is 0.
+            max_samples: Maximum number of samples. None means infinite.
+            seed: Random seed. If None, RNG is not initialized.
+            detectors: List of detector names. Default is None.
+            **kwargs: Additional arguments.
+        """
+        super().__init__(
+            sampling_frequency=sampling_frequency,
+            duration=duration,
+            start_time=start_time,
+            max_samples=max_samples,
+            seed=seed,
+            detectors=detectors,
+            **kwargs,
+        )
+        self.frequency_array = frequency_array
+        self.psd_array = psd_array
+        self.psd_file = psd_file
+        self._setup_psd()
+
+    @property
+    def frequency_array(self) -> np.ndarray[Any, np.dtype[Any]] | None:
+        """Get the frequency array."""
+        return self._frequency_array
+
+    @frequency_array.setter
+    def frequency_array(self, value: np.ndarray[Any, np.dtype[Any]] | None) -> None:
+        """Set the frequency array."""
+        if value is None:
+            self._frequency_array = np.arange(int(self.sampling_frequency * self.duration // 2) + 1) / self.duration
+        else:
+            self._frequency_array = value
+
+    def _setup_psd(self) -> None:
+        if self.frequency_array is not None and self.psd_array is not None:
+            self.psd = PowerSpectralDensity(frequency_array=self.frequency_array, psd_array=self.psd_array)
+        elif self.psd_file is not None:
+            self.psd = PowerSpectralDensity.from_power_spectral_density_file(self.psd_file)
+        else:
+            raise ValueError("Either frequency_array and psd_array or psd_file must be provided.")
+
+    @property
+    def rng(self) -> Generator | None:
+        """Get the random number generator.
+
+        Returns:
+            Random number generator instance or None if no seed was set.
+        """
+        return self._rng
+
+    @rng.setter
+    def rng(self, value: Generator | None) -> None:
+        """Set the random number generator.
+
+        Args:
+            value: Random number generator instance.
+        """
+        self._rng = value
+        # Override the bilby RNG
+        if value is not None:
+            BilbyGenerator.rng = value
+
+    def simulate(self, *args, **kwargs) -> np.ndarray:
+        """Simulate a noise segment.
+
+        Returns:
+            np.ndarray: Simulated noise segment as a numpy array.
+        """
+        if self.rng is None:
+            raise RuntimeError("Random number generator not initialized. Set seed in constructor.")
+        # Placeholder implementation; replace with actual Bilby PSD-based noise generation
+        frequency_domain_strain, _frequencies = self.psd.get_noise_realisation(
+            sampling_frequency=self.sampling_frequency, duration=self.duration
+        )
+        return infft(frequency_domain_strain=frequency_domain_strain, sampling_frequency=self.sampling_frequency)

gwsim/noise/colored_noise.py

@@ -0,0 +1,275 @@
+"""Colored noise simulator for gravitational wave detectors."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import numpy as np
+from scipy.interpolate import interp1d
+from scipy.signal.windows import tukey
+
+from gwsim.data.time_series.time_series import TimeSeries
+from gwsim.data.time_series.time_series_list import TimeSeriesList
+from gwsim.noise.base import NoiseSimulator
+from gwsim.simulator.state import StateAttribute
+
+logger = logging.getLogger("gwsim")
+
+# The default base path for PSD files
+DEFAULT_PSD_PATH = Path(__file__).parent.parent / "detector/noise_curves"
+
+
+class ColoredNoiseSimulator(NoiseSimulator):  # pylint: disable=too-many-instance-attributes
+    """Colored noise simulator for gravitational wave detectors.
+
+    This class generates noise time series with a specified power spectral density (PSD).
+    It uses an overlap-add method with windowing to produce smooth, continuous time series
+    across segment boundaries.
+
+    The simulator maintains state between batches to ensure continuity of the noise
+    time series across multiple calls to simulate().
+    """
+
+    # State attribute to track the previous strain buffer for continuity
+    previous_strain = StateAttribute(default=None)
+
+    def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments,duplicate-code
+        self,
+        psd_file: str | Path,
+        detectors: list[str],
+        sampling_frequency: float = 4096,
+        duration: float = 4,
+        start_time: float = 0,
+        max_samples: int | None = None,
+        seed: int | None = None,
+        low_frequency_cutoff: float = 2.0,
+        high_frequency_cutoff: float | None = None,
+        **kwargs,
+    ):
+        """Initialize the colored noise simulator.
+
+        Args:
+            psd_file: Path to file containing Power Spectral Density array with shape (N, 2),
+                where the first column is frequency (Hz) and the second is PSD values.
+            detectors: List of detector names (e.g., ['H1', 'L1']).
+            sampling_frequency: Sampling frequency in Hz. Default is 4096.
+            duration: Duration of each noise segment in seconds. Default is 4.
+            start_time: GPS start time for the time series. Default is 0.
+            max_samples: Maximum number of samples to generate. None means infinite.
+            seed: Seed for random number generation. If None, RNG is not initialized.
+            low_frequency_cutoff: Lower frequency cutoff in Hz. Default is 2.0.
+            high_frequency_cutoff: Upper frequency cutoff in Hz. Default is Nyquist frequency.
+            **kwargs: Additional arguments passed to parent classes.
+
+        Raises:
+            ValueError: If detectors list is empty.
+            ValueError: If duration is too short for proper noise generation.
+        """
+        if not detectors or len(detectors) == 0:
+            raise ValueError("detectors must contain at least one detector.")
+
+        super().__init__(  # pylint: disable=duplicate-code
+            sampling_frequency=sampling_frequency,
+            duration=duration,
+            start_time=start_time,
+            max_samples=max_samples,
+            seed=seed,
+            detectors=detectors,
+            **kwargs,
+        )
+
+        self.psd_file = psd_file
+        self.low_frequency_cutoff = low_frequency_cutoff
+        self.high_frequency_cutoff = (
+            high_frequency_cutoff
+            if (high_frequency_cutoff is not None and high_frequency_cutoff <= sampling_frequency / 2)
+            else sampling_frequency // 2
+        )
+
+        # Initialize noise generation properties
+        self._n_det = len(detectors)
+        self._initialize_window_properties()
+        self._initialize_frequency_properties()
+        self._initialize_psd()
+
+        # Initialize the previous strain buffer (will be populated on first simulate call)
+        self.previous_strain = np.zeros((self._n_det, self._n_chunk))
+        self._temp_strain_buffer: np.ndarray | None = None
+
+    def _initialize_window_properties(self) -> None:
+        """Initialize window properties for connecting noise realizations.
+
+        Raises:
+            ValueError: If the duration is too short for proper noise generation.
+        """
+        self._t_window = 2048
+        self._f_window = 1.0 / self._t_window
+        self._t_overlap = self._t_window / 2.0
+        self._n_overlap = int(self._t_overlap * self.sampling_frequency.value)
+
+        # Create overlap windows for smooth transitions
+        t_overlap_array = np.linspace(0, self._t_overlap, self._n_overlap)
+        self._w0 = 0.5 + np.cos(2 * np.pi * self._f_window * t_overlap_array) / 2
+        self._w1 = 0.5 + np.sin(2 * np.pi * self._f_window * t_overlap_array - np.pi / 2) / 2
+
+        # Safety check to ensure proper noise generation
+        if self.duration.value < self._t_window / 2:
+            raise ValueError(
+                f"Duration ({self.duration.value:.1f} seconds) must be at least "
+                f"{self._t_window / 2:.1f} seconds to ensure noise continuity."
+            )
+
+    def _initialize_frequency_properties(self) -> None:
+        """Initialize frequency and time properties for noise generation."""
+        self._t_chunk = self._t_window
+        self._df_chunk = 1.0 / self._t_chunk
+        self._n_chunk = int(self._t_chunk * self.sampling_frequency.value)
+        self._k_min_chunk = int(self.low_frequency_cutoff / self._df_chunk)
+        self._k_max_chunk = int(self.high_frequency_cutoff / self._df_chunk) + 1
+        self._frequency_chunk = np.arange(0.0, self._n_chunk / 2.0 + 1) * self._df_chunk
+        self._n_freq_chunk = len(self._frequency_chunk[self._k_min_chunk : self._k_max_chunk])
+        self._dt = 1.0 / self.sampling_frequency.value
+
+    def _load_spectral_data(self, file_path: str | Path) -> np.ndarray:  # pylint: disable=duplicate-code
+        """Load spectral data from file.
+
+        Args:
+            file_path: Path to file containing spectral data.
+
+        Returns:
+            Loaded array.
+
+        Raises:
+            ValueError: If file format is not supported.
+            TypeError: If file_path is not a string or Path.
+        """
+        if not isinstance(file_path, (str, Path)):
+            raise TypeError("file_path must be a string or Path.")
+
+        path = Path(file_path)
+        if not path.exists():
+            psd_dir = DEFAULT_PSD_PATH
+            path = next(iter(psd_dir.rglob(path.name)))
+
+        if path.suffix == ".npy":
+            return np.load(path)
+        if path.suffix == ".txt":
+            return np.loadtxt(path)
+        if path.suffix == ".csv":
+            return np.loadtxt(path, delimiter=",")
+        raise ValueError(f"Unsupported file format: {path.suffix}. Use .npy, .txt, or .csv.")
+
+    def _initialize_psd(self) -> None:
+        """Initialize PSD interpolation for the frequency range.
+
+        Raises:
+            ValueError: If PSD array doesn't have shape (N, 2).
+        """
+        psd_data = self._load_spectral_data(self.psd_file)
+
+        if psd_data.shape[1] != 2:
+            raise ValueError("PSD file must have shape (N, 2).")
+
+        # Interpolate the PSD to the relevant frequencies
+        freqs = self._frequency_chunk[self._k_min_chunk : self._k_max_chunk]
+        psd_interp = interp1d(psd_data[:, 0], psd_data[:, 1], bounds_error=False, fill_value="extrapolate")(freqs)
+
+        # Add a roll-off at the edges using a Tukey window
+        window = tukey(self._n_freq_chunk, alpha=1e-3)
+        self._psd = psd_interp * window
+
+    def _generate_single_realization(self) -> np.ndarray:
+        """Generate a single noise realization in the time domain.
+
+        Returns:
+            Time series array with shape (n_detectors, n_samples).
+        """
+        if self.rng is None:
+            raise RuntimeError("Random number generator not initialized. Set seed in constructor.")
+
+        freq_series = np.zeros((self._n_det, self._frequency_chunk.size), dtype=np.complex128)
+
+        # Generate white noise and color it with the PSD
+        white_strain = (
+            self.rng.standard_normal((self._n_det, self._n_freq_chunk))
+            + 1j * self.rng.standard_normal((self._n_det, self._n_freq_chunk))
+        ) / np.sqrt(2)
+        colored_strain = white_strain * np.sqrt(self._psd * 0.5 / self._df_chunk)
+        freq_series[:, self._k_min_chunk : self._k_max_chunk] += colored_strain
+
+        # Transform to time domain
+        time_series = np.fft.irfft(freq_series, n=self._n_chunk, axis=1) * self._df_chunk * self._n_chunk
+
+        return time_series
+
+    def _simulate(self, *args, **kwargs) -> TimeSeriesList:
+        """Simulate colored noise for all detectors.
+
+        Returns:
+            TimeSeriesList containing a single TimeSeries with shape (n_detectors, n_samples).
+        """
+        n_frame = int(self.duration.value * self.sampling_frequency.value)
+
+        # Load previous strain, or generate new if all zeros
+        if self.previous_strain.shape[-1] < self._n_overlap:
+            raise ValueError(
+                f"previous_strain has only {self.previous_strain.shape[-1]} samples per detector, "
+                f"but expected at least {self._n_overlap}."
+            )
+
+        strain_buffer = self.previous_strain[:, -self._n_chunk :]
+        if np.all(strain_buffer == 0):
+            strain_buffer = self._generate_single_realization()
+
+        # Apply the final part of the window
+        strain_buffer[:, -self._n_overlap :] *= self._w0
+
+        # Extend the strain buffer until it has more valid data than a single frame
+        while strain_buffer.shape[-1] - self._n_chunk - self._n_overlap < n_frame:
+            new_strain = self._generate_single_realization()
+            new_strain[:, : self._n_overlap] *= self._w1
+            new_strain[:, -self._n_overlap :] *= self._w0
+            strain_buffer[:, -self._n_overlap :] += new_strain[:, : self._n_overlap]
+            strain_buffer[:, -self._n_overlap :] *= 1 / np.sqrt(self._w0**2 + self._w1**2)
+            strain_buffer = np.concatenate((strain_buffer, new_strain[:, self._n_overlap :]), axis=1)
+
+        # Extract the frame data
+        output_strain = strain_buffer[:, self._n_chunk : (self._n_chunk + n_frame)]
+
+        # Store the output strain temporarily for state update
+        self._temp_strain_buffer = output_strain
+
+        return TimeSeriesList(
+            [TimeSeries(data=output_strain, start_time=self.start_time, sampling_frequency=self.sampling_frequency)]
+        )
+
+    def update_state(self) -> None:
+        """Update internal state after each sample generation.
+
+        Updates the previous_strain buffer to ensure continuity across batches.
+        """
+        # Call parent's update_state first (increments counter, advances start_time, saves rng_state)
+        super().update_state()
+
+        # Update the previous strain buffer for continuity
+        if self._temp_strain_buffer is not None:
+            self.previous_strain = self._temp_strain_buffer
+            self._temp_strain_buffer = None
+
+    @property
+    def metadata(self) -> dict:
+        """Get metadata including colored noise configuration.
+
+        Returns:
+            Dictionary containing metadata.
+        """
+        meta = super().metadata
+        meta["colored_noise"] = {
+            "arguments": {
+                "psd_file": str(self.psd_file),
+                "low_frequency_cutoff": self.low_frequency_cutoff,
+                "high_frequency_cutoff": self.high_frequency_cutoff,
+            }
+        }
+        return meta