gwsim 0.1.0__py3-none-any.whl

gwsim/repository/zenodo.py

@@ -0,0 +1,269 @@
+"""Zenodo publishing client."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import requests
+from requests import Response
+
+from gwsim.utils.retry import retry_on_failure
+
+
+def get_deposition_id_from_doi(doi: str) -> tuple[str, bool]:
+    """Extract deposition ID from a Zenodo DOI.
+
+    Args:
+        doi: The DOI string.
+
+    returns:
+        A tuple containing the deposition ID and a boolean indicating if it's from sandbox.
+    """
+    parts = doi.split(".")
+    deposition_id = parts[1]
+    if parts[0] == "10.5072/zenodo":
+        sandbox = True
+    elif parts[0] == "10.5281/zenodo":
+        sandbox = False
+    else:
+        raise ValueError(f"Invalid Zenodo DOI: {doi}")
+
+    return deposition_id, sandbox
+
+
+class ZenodoClient:
+    """Client for interacting with Zenodo API (production and sandbox)."""
+
+    def __init__(self, access_token: str, sandbox: bool = False):
+        """Initialize the Zenodo client.
+
+        Args:
+            access_token: Zenodo API access token.
+            sandbox: Whether to use the sandbox environment. Default is False.
+        """
+        self.access_token = access_token
+        self.sandbox = sandbox
+        self.base_url = "https://sandbox.zenodo.org/api/" if sandbox else "https://zenodo.org/api/"
+
+        self.headers = {
+            "Authorization": f"Bearer {self.access_token}",
+        }
+
+    @retry_on_failure()
+    def _request(self, method: str, url: str, headers: dict, timeout: int = 60, **kwargs) -> Response:
+        """Make a request to the Zenodo API.
+
+        Args:
+            method: HTTP method (GET, POST, PUT, DELETE).
+            url: URL to make the request to.
+            headers: Headers to include in the request.
+            timeout: Timeout for the request in seconds. Default is 60.
+            **kwargs: Additional arguments to pass to requests.
+
+        Returns:
+            Response object.
+        """
+        response = requests.request(method, url, headers=headers, timeout=timeout, **kwargs)
+        response.raise_for_status()
+        return response
+
+    def create_deposition(self, metadata: dict[str, Any] | None = None, timeout=60) -> dict[str, Any]:
+        """Create a new deposition.
+
+        Args:
+            metadata: Optional metadata dictionary for the deposition.
+            timeout: Timeout for the request in seconds. Default is 60.
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        data = {"metadata": metadata} if metadata else {}
+
+        response: Response = self._request(
+            "POST",
+            f"{self.base_url}deposit/depositions",
+            headers={"Content-Type": "application/json", **self.headers},
+            timeout=timeout,
+            json=data,
+        )
+        return response.json()
+
+    def upload_file(
+        self, deposition_id: str, file_path: Path, timeout=300, auto_timeout: bool = True
+    ) -> dict[str, Any]:
+        """Upload a file to a deposition.
+
+        Args:
+            deposition_id: ID of the deposition to upload the file to.
+            file_path: Path to the file to upload.
+            timeout: Timeout for the request in seconds. Default is 300.
+            auto_timeout: Whether to automatically adjust timeout based on file size. Default is True.
+                If True, the timeout is set to max(timeout, file_size_in_MB * 10).
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        deposition = self.get_deposition(deposition_id)
+
+        bucket = deposition["links"]["bucket"]
+
+        if auto_timeout:
+            # Get the size of the file in MB
+            file_size = file_path.stat().st_size / (1024 * 1024)
+
+            timeout = max(timeout, int(file_size * 10))  # 10 seconds per MB, minimum 300 seconds
+
+        with file_path.open("rb") as f:
+            response: Response = self._request(
+                "PUT",
+                f"{bucket}/{file_path.name}",
+                headers=self.headers,
+                timeout=timeout,
+                data=f,
+            )
+
+        return response.json()
+
+    def update_metadata(self, deposition_id: str, metadata: dict[str, Any], timeout: int = 60) -> dict[str, Any]:
+        """Update metadata for a deposition.
+
+        Args:
+            deposition_id: ID of the deposition to update.
+            metadata: Metadata dictionary to update.
+            timeout: Timeout for the request in seconds. Default is 60.
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        data = {"metadata": metadata}
+        response: Response = self._request(
+            "PUT",
+            f"{self.base_url}deposit/depositions/{deposition_id}",
+            headers={"Content-Type": "application/json", **self.headers},
+            timeout=timeout,
+            data=json.dumps(data),
+        )
+
+        return response.json()
+
+    def publish_deposition(self, deposition_id: str, timeout: int = 300) -> dict[str, Any]:
+        """Publish a deposition.
+
+        Args:
+            deposition_id: ID of the deposition to publish.
+            timeout: Timeout for the request in seconds. Default is 300.
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        response: Response = self._request(
+            "POST",
+            f"{self.base_url}deposit/depositions/{deposition_id}/actions/publish",
+            headers=self.headers,
+            timeout=timeout,
+        )
+        return response.json()
+
+    def get_deposition(self, deposition_id: str, timeout: int = 60) -> dict[str, Any]:
+        """Retrieve a deposition's details.
+
+        Args:
+            deposition_id: ID of the deposition to retrieve.
+            timeout: Timeout for the request in seconds. Default is 60.
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        response: Response = self._request(
+            "GET",
+            f"{self.base_url}deposit/depositions/{deposition_id}",
+            headers={"Content-Type": "application/json", **self.headers},
+            timeout=timeout,
+        )
+        return response.json()
+
+    def download_file(
+        self,
+        deposition_id: str,
+        filename: str,
+        output_path: Path,
+        is_draft: bool = False,
+        timeout: int = 300,
+        file_size_in_mb: int | None = None,
+    ) -> dict[str, Any]:
+        """Download a file from Zenodo.
+
+        Args:
+            deposition_id: ID of the deposition.
+            filename: Name of the file to download.
+            output_path: Path to save the downloaded file.
+            is_draft: Whether the file is in a draft deposition. Default is False.
+            timeout: Timeout for the request in seconds. Default is 300.
+            file_size_in_mb: Optional size of the file in MB to adjust timeout. If provided
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        if is_draft:
+            deposition_id += "/draft"
+
+        file_url = f"{self.base_url}records/{deposition_id}/files/{filename}"
+
+        if file_size_in_mb is not None:
+            timeout = max(timeout, int(file_size_in_mb * 10))  # 10 seconds per MB
+
+        response: Response = self._request(
+            "GET",
+            file_url,
+            headers={"Content-Type": "application/json", **self.headers},
+            timeout=timeout,
+            stream=True,
+        )
+
+        # Atomic write to avoid incomplete files
+        output_path_tmp = output_path.with_suffix(".tmp")
+        with output_path_tmp.open("wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        output_path_tmp.rename(output_path)
+        return response.json()
+
+    def list_depositions(self, status: str = "published", timeout: int = 60) -> list[dict[str, Any]]:
+        """List all depositions for the authenticated user.
+
+        Args:
+            status: Filter by deposition status ('draft', 'unsubmitted', 'published'). Default is 'published'.
+            timeout: Timeout for the request in seconds. Default is 60.
+
+        Returns:
+            List of deposition dictionaries.
+        """
+        params = {"status": status}
+        response = self._request(
+            "GET",
+            f"{self.base_url}deposit/depositions",
+            headers={"Content-Type": "application/json", **self.headers},
+            timeout=timeout,
+            params=params,
+        )
+        return response.json()
+
+    def delete_deposition(self, deposition_id: str, timeout: int = 60) -> dict[str, Any]:
+        """Delete a deposition.
+
+        Args:
+            deposition_id: ID of the deposition to delete.
+            timeout: Timeout for the request in seconds. Default is 60.
+
+        Returns:
+            Response JSON as a dictionary.
+        """
+        response = self._request(
+            "DELETE",
+            f"{self.base_url}deposit/depositions/{deposition_id}",
+            headers=self.headers,
+            timeout=timeout,
+        )
+        return response

gwsim/signal/__init__.py

@@ -0,0 +1,11 @@
+"""Initialization for the signal module."""
+
+from __future__ import annotations
+
+from gwsim.signal.base import SignalSimulator
+from gwsim.signal.cbc import CBCSignalSimulator
+
+__all__ = [
+    "CBCSignalSimulator",
+    "SignalSimulator",
+]

gwsim/signal/base.py

@@ -0,0 +1,137 @@
+"""Base class for signal simulators."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+
+from gwsim.data.time_series.time_series_list import TimeSeriesList
+from gwsim.mixin.detector import DetectorMixin
+from gwsim.mixin.population_reader import PopulationReaderMixin
+from gwsim.mixin.time_series import TimeSeriesMixin
+from gwsim.mixin.waveform import WaveformMixin
+from gwsim.simulator.base import Simulator
+
+logger = logging.getLogger("gwsim")
+
+
+class SignalSimulator(PopulationReaderMixin, WaveformMixin, TimeSeriesMixin, DetectorMixin, Simulator):
+    """Base class for signal simulators."""
+
+    def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
+        self,
+        population_file: str | Path,
+        population_file_type: str = "pycbc",
+        waveform_model: str | Callable = "IMRPhenomXPHM",
+        waveform_arguments: dict[str, Any] | None = None,
+        start_time: int = 0,
+        duration: float = 1024,
+        sampling_frequency: float = 4096,
+        max_samples: int | None = None,
+        dtype: type = np.float64,
+        detectors: list[str] | None = None,
+        minimum_frequency: float = 5,
+        **kwargs,
+    ) -> None:
+        """Initialize the base signal simulator.
+
+        Args:
+            population_file: Path to the population file.
+            population_file_type: Type of the population file (e.g., 'pycbc').
+            waveform_model: Name (from registry) or callable for waveform generation.
+            waveform_arguments: Fixed parameters to pass to waveform model.
+            start_time: Start time of the first signal segment in GPS seconds. Default is 0.
+            duration: Duration of each signal segment in seconds. Default is 1024.
+            sampling_frequency: Sampling frequency of the signals in Hz. Default is 4096.
+            max_samples: Maximum number of samples to generate. None means infinite.
+            dtype: Data type for the time series data. Default is np.float64.
+            detectors: List of detector names. Default is None.
+            minimum_frequency: Minimum GW frequency for waveform generation. Default is 5 Hz.
+            **kwargs: Additional arguments absorbed by subclasses and mixins.
+        """
+        waveform_arguments = waveform_arguments or {}
+        required_waveform_arguments = {
+            "minimum_frequency": minimum_frequency,
+            "sampling_frequency": sampling_frequency,
+        }
+        for key, value in required_waveform_arguments.items():
+            if key not in waveform_arguments:
+                logger.info("%s not specified in waveform_arguments; setting to %s", key, value)
+                waveform_arguments[key] = value
+
+        super().__init__(
+            population_file=population_file,
+            population_file_type=population_file_type,
+            waveform_model=waveform_model,
+            waveform_arguments=waveform_arguments,
+            detectors=detectors,
+            start_time=start_time,
+            duration=duration,
+            sampling_frequency=sampling_frequency,
+            max_samples=max_samples,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    def _simulate(self, *args, **kwargs) -> TimeSeriesList:
+        """Simulate signals for the current segment.
+
+        Returns:
+            TimeSeriesList: List of simulated signals.
+        """
+        output = []
+
+        while True:
+            # Get the next injection parameters
+            parameters = self.get_next_injection_parameters()
+
+            # If the parameters are None, break the loop
+            if parameters is None:
+                break
+
+            # Get the polarizations
+            polarizations = self.waveform_factory.generate(
+                waveform_model=self.waveform_model, parameters=parameters, **self.waveform_arguments
+            )
+
+            # Project onto detectors
+            strain = self.project_polarizations(
+                polarizations=polarizations,
+                right_ascension=parameters["right_ascension"],
+                declination=parameters["declination"],
+                polarization_angle=parameters["polarization_angle"],
+                **self.waveform_arguments,
+            )
+
+            # Register the parameters
+            strain.metadata.update({"injection_parameters": parameters})
+
+            output.append(strain)
+
+            # Check whether the start time of the strain is at or after the end time of the current segment
+            if strain.start_time >= self.end_time:
+                break
+        return TimeSeriesList(output)
+
+    @property
+    def metadata(self) -> dict:
+        """Get the metadata of the simulator.
+
+        Returns:
+            Metadata dictionary.
+        """
+        meta = super().metadata
+        return meta
+
+    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

gwsim/signal/cbc.py

@@ -0,0 +1,61 @@
+"""Compact Binary Coalescence (CBC) signal simulation module."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from gwsim.signal.base import SignalSimulator
+
+
+class CBCSignalSimulator(SignalSimulator):
+    """CBC Signal Simulator class."""
+
+    def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments
+        self,
+        population_file: str | Path,
+        population_file_type: str = "pycbc",
+        waveform_model: str | Callable = "IMRPhenomXPHM",
+        waveform_arguments: dict[str, Any] | None = None,
+        start_time: int = 0,
+        duration: float = 1024,
+        sampling_frequency: float = 4096,
+        max_samples: int | None = None,
+        dtype: type = np.float64,
+        detectors: list[str] | None = None,
+        minimum_frequency: float = 5,
+        **kwargs,
+    ) -> None:
+        """Initialize the CBC signal simulator.
+
+        Args:
+            population_file: Path to the population file.
+            population_file_type: Type of the population file (e.g., 'pycbc').
+            waveform_model: Name (from registry) or callable for waveform generation.
+            waveform_arguments: Fixed parameters to pass to waveform model.
+            start_time: Start time of the first signal segment in GPS seconds. Default is 0.
+            duration: Duration of each signal segment in seconds. Default is 1024.
+            sampling_frequency: Sampling frequency of the signals in Hz. Default is 4096.
+            max_samples: Maximum number of samples to generate. None means infinite.
+            dtype: Data type for the time series data. Default is np.float64.
+            detectors: List of detector names. Default is None.
+            minimum_frequency: Minimum GW frequency for waveform generation. Default is 5 Hz.
+            **kwargs: Additional arguments absorbed by subclasses and mixins.
+        """
+        super().__init__(
+            population_file=population_file,
+            population_file_type=population_file_type,
+            waveform_model=waveform_model,
+            waveform_arguments=waveform_arguments,
+            start_time=start_time,
+            duration=duration,
+            sampling_frequency=sampling_frequency,
+            max_samples=max_samples,
+            dtype=dtype,
+            detectors=detectors,
+            minimum_frequency=minimum_frequency,
+            **kwargs,
+        )

gwsim/simulator/__init__.py

@@ -0,0 +1,7 @@
+"""Simulator module for GWSim."""
+
+from __future__ import annotations
+
+from .base import Simulator
+
+__all__ = ["Simulator"]