gwsim 0.1.0__py3-none-any.whl

gwsim/mixin/population_reader.py

@@ -0,0 +1,175 @@
+"""Population reader mixin for simulators."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import h5py
+import numpy as np
+import pandas as pd
+import yaml
+
+logger = logging.getLogger("gwsim")
+
+
+class PopulationIterationState:  # pylint: disable=too-few-public-methods
+    """Manages state for population file iteration with checkpoint support."""
+
+    def __init__(self, checkpoint_file: str | Path | None = None, encoding: str = "utf-8") -> None:
+        self.checkpoint_file = checkpoint_file
+        self.encoding = encoding
+        self.current_index = 0
+        self.injected_indices: list[int] = []
+        self.segment_map: dict[int, list[int]] = {}
+        self._load_checkpoint()
+
+    @property
+    def checkpoint_file(self) -> Path | None:
+        """Get the checkpoint file path.
+
+        Returns:
+            Path to the checkpoint file or None if not set.
+        """
+        return self._checkpoint_file
+
+    @checkpoint_file.setter
+    def checkpoint_file(self, value: str | Path | None) -> None:
+        """Set the checkpoint file path.
+
+        Args:
+            value: Path to the checkpoint file or None to unset.
+        """
+        if value is None:
+            self._checkpoint_file = None
+        else:
+            self._checkpoint_file = Path(value)
+
+    def _load_checkpoint(self) -> None:
+        if self.checkpoint_file and self.checkpoint_file.is_file():
+            try:
+                with open(self.checkpoint_file, encoding=self.encoding) as f:
+                    data = yaml.safe_load(f)["population"]
+                    self.current_index = data.get("current_index", 0)
+                    self.injected_indices = data.get("injected_indices", [])
+                    self.segment_map = data.get("segment_map", {})
+                logger.info(
+                    "Loaded checkpoint: current_index=%s, injected=%s",
+                    self.current_index,
+                    self.injected_indices,
+                )
+            except (OSError, yaml.YAMLError, KeyError) as e:
+                logger.warning("Failed to load checkpoint %s: %s. Starting fresh.", self.checkpoint_file, e)
+
+
+PARAMETER_NAME_MAPPER = {
+    "pycbc": {
+        "ra": "right_ascension",
+        "dec": "declination",
+        "polarization": "polarization_angle",
+    }
+}
+
+
+class PopulationReaderMixin:  # pylint: disable=too-few-public-methods
+    """A mixin class to read population files for GW signal simulators."""
+
+    population_counter = 0
+
+    def __init__(self, population_file: str | Path, population_file_type: str = "pycbc", **kwargs):
+        """Initialize the PopulationReaderMixin.
+
+        Args:
+            population_file: Path to the population file.
+            population_file_type: Type of the population file. Default is 'pycbc'.
+        """
+        super().__init__(**kwargs)
+        self.population_file = Path(population_file)
+        self.population_file_type = population_file_type
+        if not self.population_file.is_file():
+            raise FileNotFoundError(f"Population file {self.population_file} does not exist.")
+
+        if population_file_type == "pycbc":
+            self.population_data = self._read_pycbc_population_file(self.population_file, **kwargs)
+        else:
+            raise ValueError(f"Unsupported population file type: {population_file_type}")
+
+    def _read_pycbc_population_file(  # pylint: disable=unused-argument
+        self, file_name: str | Path, **kwargs
+    ) -> pd.DataFrame:
+        """Read a pycbc population file in HDF5 format.
+
+        Args:
+            file_name: Path to the pycbc population file.
+            **kwargs: Additional arguments (not used currently).
+
+        Returns:
+            A pandas DataFrame containing the population data.
+        """
+        # Load the pycbc population file and create a pandas DataFrame
+
+        with h5py.File(file_name, "r") as f:
+            data = {key: value[()] for key, value in f.items()}
+
+            # Create a DataFrame
+            population_data = pd.DataFrame(data)
+
+            # Save the attributes to metadata
+            attrs = dict(f.attrs.items())
+            # If there is any numpy array in attrs, convert it to list
+            for key, value in attrs.items():
+                if isinstance(value, np.ndarray):
+                    attrs[key] = value.tolist()
+            self._population_metadata = attrs
+
+        # Order the DataFrame by the coalescence time 'tc'
+        return (
+            population_data.sort_values(by="tc")
+            .reset_index(drop=True)
+            .rename(columns=PARAMETER_NAME_MAPPER.get("pycbc", {}))
+        )
+
+    def get_next_injection_parameters(self) -> dict[str, float | int] | None:
+        """Get the next set of injection parameters from the population.
+
+        Returns:
+            A dictionary of injection parameters for the next event,
+                or None if all events have been used.
+        """
+        if self.population_counter < len(self.population_data):
+            output = self.population_data.iloc[self.population_counter].to_dict()
+            self.population_counter += 1
+        else:
+            output = None
+        return output
+
+    def get_injection_parameter_keys(self) -> list[str]:
+        """Get the list of injection parameter keys from the population data.
+
+        Returns:
+            A list of strings representing the injection parameter keys.
+        """
+        if not self.population_data.empty:
+            output = list(self.population_data.columns)
+        else:
+            output = []
+        return output
+
+    @property
+    def metadata(self) -> dict:
+        """Get metadata including population file information.
+
+        Returns:
+            Dictionary containing metadata.
+        """
+        metadata = {
+            "population_reader": {
+                "arguments": {
+                    "population_file": str(self.population_file),
+                    "population_file_type": self.population_file_type,
+                }
+            }
+        }
+        if hasattr(self, "_population_metadata"):
+            metadata["population_reader"].update(self._population_metadata)
+        return metadata

gwsim/mixin/randomness.py

@@ -0,0 +1,107 @@
+"""Mixins for GW simulation classes for randomness handling."""
+
+from __future__ import annotations
+
+import logging
+
+from gwsim.simulator.state import StateAttribute
+from gwsim.utils.random import Generator, get_rng, get_state, set_seed, set_state
+
+logger = logging.getLogger("gwsim")
+
+
+class RandomnessMixin:
+    """Mixin providing random number generation capabilities.
+
+    This mixin should be used by simulators that require randomness.
+    It provides RNG state management, seed handling, and state persistence.
+
+    Example:
+        >>> class MySimulator(RandomnessMixin, Simulator):
+        ...     def __init__(self, seed=None, **kwargs):
+        ...         super().__init__(**kwargs)
+        ...         self.seed = seed
+        ...
+        ...     def simulate(self):
+        ...         # Use self.rng for random operations
+        ...         return self.rng.random()
+    """
+
+    # State attributes for RNG persistence
+    rng_state = StateAttribute(
+        lambda self: get_state() if self.rng else None, post_set_hook=lambda self, state: self.init_rng(state)
+    )
+
+    def __init__(self, seed: int | None = None, **kwargs):
+        """Initialize the randomness mixin.
+
+        Args:
+            seed: Random seed for reproducibility. If None, no RNG is created.
+            **kwargs: Additional arguments passed to parent classes.
+        """
+        super().__init__(**kwargs)
+        self.seed = seed
+
+    @property
+    def seed(self) -> int | None:
+        """Get the random seed.
+
+        Returns:
+            The random seed used for initialization.
+        """
+        return self._seed
+
+    @seed.setter
+    def seed(self, value: int | None) -> None:
+        """Set the random seed and reinitialize RNG.
+
+        Args:
+            value: New random seed. If None, RNG is disabled.
+        """
+        self._seed = value
+        if value is not None:
+            set_seed(value)
+            self.rng = get_rng()
+            self.rng_state = get_state()
+        else:
+            self.rng = None
+
+    @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
+
+    def init_rng(self, state: dict | None) -> None:
+        """Initialize RNG from saved state.
+
+        Args:
+            state: Saved RNG state dictionary.
+        """
+        if state is not None and self.rng is not None:
+            set_state(state)
+            self._rng = get_rng()
+        else:
+            logger.debug("init_rng called but state is %s and self.rng is %s.", state, self.rng)
+
+    @property
+    def metadata(self) -> dict:
+        """Get metadata including seed information.
+
+        Returns:
+            Dictionary containing seed and other metadata.
+        """
+        metadata = {"seed": self.seed}
+        return metadata

gwsim/mixin/time_series.py

@@ -0,0 +1,295 @@
+"""Mixins for simulator classes providing optional functionality."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+from astropy.units.quantity import Quantity
+from gwpy.timeseries import TimeSeries as GWPyTimeSeries
+
+from gwsim.data.time_series.time_series import TimeSeries
+from gwsim.data.time_series.time_series_list import TimeSeriesList
+from gwsim.simulator.state import StateAttribute
+from gwsim.utils.datetime_parser import parse_duration_to_seconds
+
+logger = logging.getLogger("gwsim")
+
+
+class TimeSeriesMixin:  # pylint: disable=too-few-public-methods,too-many-instance-attributes
+    """Mixin providing timing and duration management.
+
+    This mixin adds time-based parameters commonly used
+    in gravitational wave simulations.
+    """
+
+    start_time = StateAttribute(Quantity(0, unit="s"))
+    cached_data_chunks = TimeSeriesList()
+
+    def __init__(
+        self,
+        start_time: int = 0,
+        duration: float = 4,
+        total_duration: float | str | None = None,
+        sampling_frequency: float = 4096,
+        num_of_channels: int | None = None,
+        dtype: type = np.float64,
+        **kwargs,
+    ):
+        """Initialize timing parameters.
+
+        Args:
+            start_time: Start time in GPS seconds. Default is 0.
+            duration: Duration of simulation in seconds. Default is 4.
+            total_duration
+            sampling_frequency: Sampling frequency in Hz. Default is 4096.
+            dtype: Data type for the time series data. Default is np.float64.
+            **kwargs: Additional arguments passed to parent classes.
+        """
+        super().__init__(**kwargs)
+        # TimeSeriesMixin is the last mixin in the hierarchy, so no super().__init__() call needed
+        self.start_time = Quantity(start_time, unit="s")
+        self.duration = duration
+        self.total_duration = total_duration
+        self.sampling_frequency = sampling_frequency
+        self.dtype = dtype
+
+        # Get the number of channels.
+        if num_of_channels is not None:
+            self.num_of_channels = num_of_channels
+            if "detectors" in kwargs and kwargs["detectors"] is not None:
+                if len(kwargs["detectors"]) != num_of_channels:
+                    raise ValueError("Number of detectors does not match num_of_channels.")
+        elif "detectors" in kwargs and kwargs["detectors"] is not None:
+            self.num_of_channels = len(kwargs["detectors"])
+        else:
+            self.num_of_channels = 1
+
+    @property
+    def duration(self) -> Quantity:
+        """Get the duration of each simulation segment.
+
+        Returns:
+            Duration in seconds.
+        """
+        return self._duration
+
+    @duration.setter
+    def duration(self, value: float) -> None:
+        """Set the duration of each simulation segment.
+
+        Args:
+            value: Duration in seconds.
+        """
+        if value <= 0:
+            raise ValueError("duration must be positive.")
+        self._duration = Quantity(value, unit="s")
+
+    @property
+    def sampling_frequency(self) -> Quantity:
+        """Get the sampling frequency.
+
+        Returns:
+            Sampling frequency in Hz.
+        """
+        return self._sampling_frequency
+
+    @sampling_frequency.setter
+    def sampling_frequency(self, value: float) -> None:
+        """Set the sampling frequency.
+
+        Args:
+            value: Sampling frequency in Hz.
+        """
+        if value <= 0:
+            raise ValueError("sampling_frequency must be positive.")
+        self._sampling_frequency = Quantity(value, unit="Hz")
+
+    @property
+    def total_duration(self) -> Quantity:
+        """Get the total duration of the simulation.
+
+        Returns:
+            Total duration in seconds.
+        """
+        return self._total_duration
+
+    @total_duration.setter
+    def total_duration(self, value: int | float | str | None) -> None:
+        """Set the total duration of the simulation.
+
+        Args:
+            value: Total duration in seconds.
+        """
+        if value is not None:
+            if isinstance(value, (float, int)):
+                self._total_duration = Quantity(value, unit="s")
+            elif isinstance(value, str):
+                self._total_duration = Quantity(parse_duration_to_seconds(value), unit="s")
+            else:
+                raise ValueError("total_duration must be a float, int, or str representing duration.")
+
+            if self.total_duration < 0:
+                raise ValueError("total_duration must be non-negative.")
+
+            if self.total_duration < self.duration:
+                raise ValueError("total_duration must be greater than or equal to duration.")
+
+            # Round the total_duration to the nearest multiple of duration
+            num_segments = round(self.total_duration.value / self.duration.value)
+            self._total_duration = Quantity(num_segments * self.duration, unit="s")
+
+            logger.info("Total duration set to %s seconds.", self.total_duration)
+
+            # Set the max_samples based on total_duration and duration
+            self.max_samples = int(self.total_duration.value / self.duration.value)
+            logger.info("Setting max_samples to %s based on total_duration and duration.", self.max_samples)
+        else:
+            self._total_duration = self.duration * self.max_samples
+            logger.info("total_duration not set, using duration * max_samples = %s seconds.", self.total_duration.value)
+
+    @property
+    def end_time(self) -> Quantity:
+        """Calculate the end time of the current segment.
+
+        Returns:
+            End time in GPS seconds.
+        """
+        return cast(Quantity, self.start_time + self.duration)
+
+    @property
+    def final_end_time(self) -> Quantity:
+        """Calculate the final end time of the entire simulation.
+
+        Returns:
+            Final end time in GPS seconds.
+        """
+        return cast(Quantity, self.start_time + self.total_duration)
+
+    def _simulate(self, *args, **kwargs) -> TimeSeriesList:
+        """Generate time series data chunks.
+
+        This method should be implemented by subclasses to generate
+        the actual time series data.
+        """
+        raise NotImplementedError("Subclasses must implement the _simulate method.")
+
+    def simulate(self, *args, **kwargs) -> TimeSeries:
+        """
+        Simulate a segment of time series data.
+
+        Args:
+            *args: Positional arguments for the _simulate method.
+            **kwargs: Keyword arguments for the _simulate method.
+
+        Returns:
+            TimeSeries: Simulated time series segment.
+        """
+        # First create a new segment
+        segment = TimeSeries(
+            data=np.zeros(
+                (self.num_of_channels, int(self.duration.value * self.sampling_frequency.value)), dtype=self.dtype
+            ),
+            start_time=self.start_time,
+            sampling_frequency=self.sampling_frequency,
+        )
+
+        # Inject cached data chunks into the segment
+        self.cached_data_chunks = segment.inject_from_list(self.cached_data_chunks)
+
+        # Generate new chunks of data
+        new_chunks = self._simulate(*args, **kwargs)
+
+        # Add the new chunks to the segment
+        remaining_chunks = segment.inject_from_list(new_chunks)
+
+        # Add the remaining chunks to the cache
+        self.cached_data_chunks.extend(remaining_chunks)
+
+        # Check whether there are chunks that are outside the whole dataset duration
+        # Remove the chunks that are outside the total duration
+        for i in reversed(range(len(self.cached_data_chunks))):
+            chunk = self.cached_data_chunks[i]
+            if chunk.start_time >= self.final_end_time:
+                logger.info(
+                    "Removing cached chunk starting at %s which is outside the total duration ending at %s.",
+                    chunk.start_time,
+                    self.final_end_time,
+                )
+                self.cached_data_chunks.pop(i)
+            elif chunk.end_time <= self.start_time:
+                logger.info(
+                    "Removing cached chunk ending at %s which is before the current segment starting at %s.",
+                    chunk.end_time,
+                    self.start_time,
+                )
+                self.cached_data_chunks.pop(i)
+
+        return segment
+
+    @property
+    def metadata(self) -> dict:
+        """Get metadata including timing information.
+
+        Returns:
+            Dictionary containing timing parameters and other metadata.
+        """
+        metadata = {
+            "time_series": {
+                "arguments": {
+                    "start_time": self.start_time,
+                    "duration": self.duration,
+                    "sampling_frequency": self.sampling_frequency,
+                    "num_of_channels": self.num_of_channels,
+                    "dtype": str(self.dtype),
+                }
+            }
+        }
+        return metadata
+
+    def _save_data(  # pylint: disable=unused-argument
+        self,
+        data: TimeSeries,
+        file_name: str | Path | np.ndarray[Any, np.dtype[np.object_]],
+        channel_names: str | None = None,
+        **kwargs,
+    ) -> None:
+        """Save time series data to a file.
+
+        Args:
+            data: Time series data to save.
+            file_name: Path to the output file.
+            **kwargs: Additional arguments for the saving function.
+        """
+        if data.num_of_channels == 1 and isinstance(file_name, (str, Path)):
+            self._save_gwf_data(data=data[0], file_name=file_name, **kwargs)
+        elif (
+            data.num_of_channels > 1
+            and isinstance(file_name, np.ndarray)
+            and len(file_name.shape) == 1
+            and file_name.shape[0] == data.num_of_channels
+        ):
+            for i in range(data.num_of_channels):
+                single_file_name = cast(Path, file_name[i])
+                self._save_gwf_data(data=data[i], file_name=single_file_name, **kwargs)
+        else:
+            raise ValueError(
+                "file_name must be a single path for single-channel data or an array of paths for multi-channel data."
+            )
+
+    def _save_gwf_data(  # pylint: disable=unused-argument
+        self, data: GWPyTimeSeries, file_name: str | Path, channel: str | None = None, **kwargs
+    ) -> None:
+        """Save GWPy TimeSeries data to a GWF file.
+
+        Args:
+            data: GWPy TimeSeries data to save.
+            file_name: Path to the output GWF file.
+            channel: Optional channel name to set in the data.
+            **kwargs: Additional arguments for the write function.
+        """
+        if channel is not None:
+            data.channel = channel
+        data.write(str(file_name))

gwsim/mixin/waveform.py

@@ -0,0 +1,47 @@
+"""Mixin for waveform generation in signal simulators."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from gwsim.waveform.factory import WaveformFactory
+
+logger = logging.getLogger("gwsim")
+
+
+class WaveformMixin:  # pylint: disable=too-few-public-methods
+    """Mixin class for waveform generation in signal simulators."""
+
+    def __init__(
+        self,
+        waveform_model: str | Callable = "IMRPhenomXPHM",
+        waveform_arguments: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        """Initialize the WaveformMixin.
+
+        Args:
+            waveform_model: Name (from registry) or callable for waveform generation.
+            waveform_arguments: Fixed parameters to pass to waveform model.
+            parameter_mapping: Dict mapping population column names to waveform parameter names.
+            minimum_frequency: Minimum GW frequency for waveform generation.
+            **kwargs: Additional arguments for other mixins.
+        """
+        super().__init__(**kwargs)
+        self.waveform_factory = WaveformFactory()
+        if waveform_model not in self.waveform_factory.list_models():
+            # Register the model if not already registered
+            self.waveform_factory.register_model(name=str(waveform_model), factory_func=waveform_model)
+        self.waveform_model = str(waveform_model)
+        self.waveform_arguments = waveform_arguments or {}
+
+    @property
+    def metadata(self) -> dict:
+        """Include waveform metadata."""
+        metadata = {
+            "waveform_model": self.waveform_model,
+            "waveform_arguments": self.waveform_arguments,
+        }
+        return metadata

gwsim/noise/__init__.py

@@ -0,0 +1,19 @@
+"""
+Noise models for gravitational wave detector simulations.
+"""
+
+from __future__ import annotations
+
+from gwsim.noise.base import NoiseSimulator
+from gwsim.noise.bilby_stationary_gaussian import BilbyStationaryGaussianNoiseSimulator
+from gwsim.noise.colored_noise import ColoredNoiseSimulator
+from gwsim.noise.correlated_noise import CorrelatedNoiseSimulator
+from gwsim.noise.pycbc_stationary_gaussian import PyCBCStationaryGaussianNoiseSimulator
+
+__all__ = [
+    "BilbyStationaryGaussianNoiseSimulator",
+    "ColoredNoiseSimulator",
+    "CorrelatedNoiseSimulator",
+    "NoiseSimulator",
+    "PyCBCStationaryGaussianNoiseSimulator",
+]