gwsim 0.1.0__py3-none-any.whl

gwsim/simulator/base.py

@@ -0,0 +1,315 @@
+"""Refactored base simulator with clean separation of concerns."""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+import yaml
+
+from gwsim import __version__
+from gwsim.simulator.state import StateAttribute
+from gwsim.utils.io import get_file_name_from_template
+
+logger = logging.getLogger("gwsim")
+
+
+class Simulator(ABC):
+    """Base simulator class providing core interface and iteration capabilities.
+
+    This class provides the minimal common interface that all simulators share:
+    - State management and persistence
+    - Iterator protocol for data generation
+    - Metadata handling
+    - File I/O operations
+
+    Specialized functionality (randomness, timing, etc.) should be added
+    via mixins to avoid bloating the base interface.
+
+    Args:
+        max_samples: Maximum number of samples to generate. None means infinite.
+        **kwargs: Additional arguments absorbed by subclasses and mixins.
+    """
+
+    # State attributes using StateAttribute descriptor
+    counter = StateAttribute(default=0)
+
+    def __init__(self, max_samples: int | float | None = None, **kwargs):
+        """Initialize the base simulator.
+
+        Args:
+            max_samples: Maximum number of samples to generate.
+            **kwargs: Additional arguments for subclasses and mixins.
+        """
+        # Absorb unused kwargs to enable flexible parameter passing
+        if kwargs:
+            logger.debug("Unused kwargs in Simulator.__init__: %s", kwargs)
+
+        # Non-state attributes
+        if max_samples is None:
+            self.max_samples = np.inf
+            logger.debug("max_samples set to None, interpreted as infinite.")
+        else:
+            self.max_samples = max_samples
+
+    @property
+    def max_samples(self) -> int | float:
+        """Get the maximum number of samples.
+
+        Returns:
+            Maximum number of samples (np.inf for unlimited).
+        """
+        return self._max_samples
+
+    @max_samples.setter
+    def max_samples(self, value: int | float) -> None:
+        """Set the maximum number of samples.
+
+        Args:
+            value: Maximum number of samples. None interpreted as infinite.
+
+        Raises:
+            ValueError: If value is negative.
+        """
+        if value < 0:
+            raise ValueError("Max samples cannot be negative.")
+        self._max_samples = value
+
+    @property
+    def state(self) -> dict:
+        """Get the current simulator state.
+
+        Returns:
+            Dictionary containing all state attributes.
+        """
+        # Get state attributes from all classes in MRO (set by StateAttribute descriptors)
+        state_attrs: list[Any] = []
+        for cls in self.__class__.__mro__:
+            state_attrs.extend(getattr(cls, "_state_attributes", []))
+        # Remove duplicates while preserving order
+        seen = set()
+        state_attrs = [x for x in state_attrs if not (x in seen or seen.add(x))]
+        return {key: getattr(self, key) for key in state_attrs}
+
+    @state.setter
+    def state(self, state: dict) -> None:
+        """Set the simulator state.
+
+        Args:
+            state: Dictionary of state values.
+
+        Raises:
+            ValueError: If state contains unregistered attributes.
+        """
+        # Get state attributes from all classes in MRO (set by StateAttribute descriptors)
+        state_attrs: list[Any] = []
+        for cls in self.__class__.__mro__:
+            state_attrs.extend(getattr(cls, "_state_attributes", []))
+        # Remove duplicates while preserving order
+        seen = set()
+        state_attrs = [x for x in state_attrs if not (x in seen or seen.add(x))]
+
+        for key, value in state.items():
+            if key not in state_attrs:
+                raise ValueError(f"Attribute {key} is not registered as a state attribute.")
+            setattr(self, key, value)
+
+    @property
+    def metadata(self) -> dict:
+        """Get simulator metadata.
+
+        This can be overridden by subclasses to include additional metadata.
+        Mixins should call super().metadata and update the returned dictionary.
+
+        Returns:
+            Dictionary containing metadata.
+        """
+        metadata = {
+            "max_samples": self.max_samples,
+            "counter": self.counter,
+            "version": __version__,
+            "state": {},
+        }
+        metadata["state"].update(self.state)
+        return metadata
+
+    # Iterator protocol
+    def __iter__(self):
+        """Return iterator interface."""
+        return self
+
+    def __next__(self):
+        """Generate next sample.
+
+        Returns:
+            Next generated sample.
+
+        Raises:
+            StopIteration: When max_samples is reached.
+        """
+        if self.counter >= self.max_samples:
+            raise StopIteration("Maximum number of samples reached.")
+
+        sample = self.simulate()
+        self.update_state()
+        return sample
+
+    def save_state(self, file_name: str | Path, overwrite: bool = False, encoding: str = "utf-8") -> None:
+        """Save simulator state to file.
+
+        Args:
+            file_name: Output file path (must have .yml or .yaml extension).
+            overwrite: Whether to overwrite existing files.
+            encoding: File encoding to use when writing the file.
+
+        Raises:
+            ValueError: If file extension is not .yml or .yaml.
+            FileExistsError: If file exists and overwrite=False.
+        """
+        file_name = Path(file_name)
+
+        if file_name.suffix.lower() not in [".yml", ".yaml"]:
+            raise ValueError(f"Unsupported file format: {file_name.suffix}. Supported: [.yml, .yaml]")
+
+        if not overwrite and file_name.exists():
+            raise FileExistsError(f"File '{file_name}' already exists. Use overwrite=True to overwrite it.")
+
+        with file_name.open("w", encoding=encoding) as f:
+            yaml.safe_dump(self.state, f)
+
+    def load_state(self, file_name: str | Path, encoding: str = "utf-8") -> None:
+        """Load simulator state from file.
+
+        Args:
+            file_name: Input file path (must have .yml or .yaml extension).
+            encoding: File encoding to use when reading the file.
+
+        Raises:
+            FileNotFoundError: If file doesn't exist.
+            ValueError: If file extension is not .yml or .yaml.
+        """
+        file_name = Path(file_name)
+
+        if file_name.suffix.lower() not in [".yml", ".yaml"]:
+            raise ValueError(f"Unsupported file format: {file_name.suffix}. Supported: [.yml, .yaml]")
+
+        if not file_name.exists():
+            raise FileNotFoundError(f"File '{file_name}' does not exist.")
+
+        with file_name.open("r", encoding=encoding) as f:
+            state = yaml.safe_load(f)
+
+        self.state = state
+
+    def save_metadata(
+        self,
+        file_name: str | Path,
+        output_directory: str | Path | None = None,
+        overwrite: bool = False,
+        encoding: str = "utf-8",
+    ) -> None:
+        """Save simulator metadata to file.
+
+        Args:
+            file_name: Output file path (must have .yml or .yaml extension).
+            output_directory: Optional output directory to prepend to the file name.
+            overwrite: Whether to overwrite existing files.
+            encoding: File encoding to use when writing the file.
+
+        Raises:
+            ValueError: If file extension is not .yml or .yaml.
+            FileExistsError: If file exists and overwrite=False.
+        """
+        file_name_resolved = get_file_name_from_template(
+            template=str(file_name),
+            instance=self,
+            output_directory=output_directory,
+        )
+        if not isinstance(file_name_resolved, Path):
+            raise ValueError("Resolved file name for metadata must be a single Path.")
+
+        if file_name_resolved.suffix.lower() not in [".yml", ".yaml"]:
+            raise ValueError(f"Unsupported file format: {file_name_resolved.suffix}. Supported: [.yml, .yaml]")
+
+        if not overwrite and file_name_resolved.exists():
+            raise FileExistsError(f"File '{file_name_resolved}' already exists. Use overwrite=True to overwrite it.")
+
+        with file_name_resolved.open("w", encoding=encoding) as f:
+            yaml.safe_dump(self.metadata, f)
+
+    def update_state(self) -> None:
+        """Update internal state after each sample generation.
+
+        This method must be implemented by all simulator subclasses.
+        """
+        self.counter = cast(int, self.counter) + 1
+
+    # Abstract methods that subclasses must implement
+    @abstractmethod
+    def simulate(self, *args, **kwargs) -> Any:
+        """Generate a single sample.
+
+        This method must be implemented by all simulator subclasses.
+
+        Returns:
+            A single generated sample.
+        """
+
+    @abstractmethod
+    def _save_data(self, data: Any, file_name: str | Path | np.ndarray[Any, np.dtype[np.object_]], **kwargs) -> None:
+        """Internal method to save data to a file.
+
+        This method must be implemented by all simulator subclasses.
+
+        Args:
+            batch: Batch of generated samples.
+            file_name: Output file path.
+            overwrite: Whether to overwrite existing files.
+            **kwargs: Additional arguments for specific file formats.
+        """
+
+    def save_data(
+        self,
+        data: Any,
+        file_name: str | Path,
+        output_directory: str | Path | None = None,
+        overwrite: bool = False,
+        **kwargs,
+    ) -> None:
+        """Save data to a file.
+
+        This method must be implemented by all simulator subclasses.
+
+        Args:
+            batch: Batch of generated samples.
+            file_name: Output file path.
+                If the file_name contains placeholders (e.g., {{detector}}, {{duration}}),
+                they are filled by the attributes of the simulator.
+            output_directory: Optional output directory to prepend to the file name.
+            overwrite: Whether to overwrite existing files.
+            save_metadata: Whether to save metadata alongside the data.
+            **kwargs: Additional arguments for specific file formats.
+        """
+        file_name_resolved = get_file_name_from_template(
+            template=str(file_name),
+            instance=self,
+            output_directory=output_directory,
+        )
+
+        if not overwrite:
+            if isinstance(file_name_resolved, Path):
+                if file_name_resolved.exists():
+                    raise FileExistsError(
+                        f"File '{file_name_resolved}' already exists. " f"Use overwrite=True to overwrite it."
+                    )
+            else:
+                for single_file in file_name_resolved.flatten():
+                    if single_file.exists():
+                        raise FileExistsError(
+                            f"File '{single_file}' already exists. " f"Use overwrite=True to overwrite it."
+                        )
+
+        self._save_data(data=data, file_name=file_name_resolved, **kwargs)

gwsim/simulator/state.py

@@ -0,0 +1,85 @@
+"""
+A descriptor class to handle a state attribute.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar, overload
+
+T = TypeVar("T")
+
+
+class StateAttribute(Generic[T]):
+    """A state attribute."""
+
+    def __init__(
+        self,
+        default: T | None = None,
+        default_factory: Callable[[], T] | None = None,
+        post_set_hook: Callable[[Any, T], None] | None = None,
+    ) -> None:
+        """A state attribute.
+
+        Args:
+            default (T | None, optional): The default value of the attribute. Defaults to None.
+            default_factory (Callable[[], T] | None, optional): A factory to create the default value.
+                This is useful when you want to have a list as a state attribute, and
+                do not want to share the list across instances. Defaults to None.
+            post_set_hook (Callable[[Any, T], None] | None): Called after the value of the attribute is set.
+        """
+        self.default = default
+        self.default_factory = default_factory
+        self.post_set_hook = post_set_hook
+        self.name = None
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        """Set the name of the attribute.
+
+        Args:
+            owner (type): Owner of the attribute.
+            name (str): Name.
+        """
+        self.name = name
+        # Ensure each class has its OWN _state_attributes list (not inherited from parent).
+        # We check __dict__ directly to avoid inheriting a parent's list.
+        if "_state_attributes" not in owner.__dict__:
+            owner._state_attributes = []
+        if name not in owner._state_attributes:
+            owner._state_attributes.append(name)
+
+    @overload
+    def __get__(self, instance: None, owner: type) -> StateAttribute[T]: ...
+
+    @overload
+    def __get__(self, instance: Any, owner: type) -> T: ...
+
+    def __get__(self, instance: Any, owner: type) -> T | StateAttribute[T]:
+        """Get the value of the attribute.
+
+        Args:
+            instance (Any): Instance of the owner.
+            owner (type): Placeholder.
+
+        Returns:
+            T | StateAttribute: Value of the attribute.
+        """
+        if instance is None:
+            return self
+        if self.name not in instance.__dict__:
+            if self.default_factory is not None:
+                instance.__dict__[self.name] = self.default_factory()
+            else:
+                instance.__dict__[self.name] = self.default
+        return instance.__dict__[self.name]
+
+    def __set__(self, instance: Any, value: T) -> None:
+        """Set the value of the attribute.
+
+        Args:
+            instance (Any): An instance of the owner.
+            value (T): Value to set.
+        """
+        instance.__dict__[self.name] = value
+        if self.post_set_hook is not None:
+            self.post_set_hook(instance, value)

gwsim/utils/__init__.py

@@ -0,0 +1,11 @@
+"""Utility functions."""
+
+from __future__ import annotations
+
+from .random import generate_seeds, get_rng, set_seed
+
+__all__ = [
+    "generate_seeds",
+    "get_rng",
+    "set_seed",
+]

gwsim/utils/datetime_parser.py

@@ -0,0 +1,44 @@
+"""Utilities for parsing human-readable time durations."""
+
+from __future__ import annotations
+
+import re
+
+SECOND_UNITS = {
+    "second": 1,
+    "minute": 60,
+    "hour": 3_600,
+    "day": 86_400,
+    "week": 604_800,
+    # approximate values for longer units
+    "month": 2_592_000,  # 30 days
+    "year": 31_536_000,  # 365 days
+}
+
+
+def parse_duration_to_seconds(duration: str) -> float:
+    """Convert a human-friendly duration like "1 day" into seconds.
+
+    Args:
+        duration: A string such as "1 week", "2 days", "1.5 hours".
+
+    Returns:
+        Number of seconds represented by the duration.
+
+    Raises:
+        ValueError: If the string cannot be parsed or the unit is unsupported.
+    """
+
+    pattern = re.compile(r"^\s*(?P<value>\d+(?:\.\d+)?)\s*(?P<unit>\w+)s?\s*$", re.IGNORECASE)
+    match = pattern.match(duration)
+    if not match:
+        raise ValueError("Duration must be of the form '<number> <unit>' (e.g. '1 week').")
+
+    value = float(match.group("value"))
+    unit = match.group("unit").lower().rstrip("s")
+
+    seconds = SECOND_UNITS.get(unit)
+    if seconds is None:
+        raise ValueError(f"Unsupported duration unit '{unit}'. Supported units: {', '.join(SECOND_UNITS)}.")
+
+    return value * seconds

gwsim/utils/et_2l_geometry.py

@@ -0,0 +1,165 @@
+"""Script to compute the ET 2L aligned/misaligned geometry at a given location"""
+
+from __future__ import annotations
+
+import numpy as np
+import pymap3d as pm
+from pycbc.detector import Detector, add_detector_on_earth
+
+
+def get_unit_vector_angles(unit_vector: np.ndarray, ellipsoid_position: np.ndarray) -> np.ndarray:
+    """
+    Compute the azimuthal angle and altitude (elevation) of a given unit vector
+    relative to the local tangent plane at the specified ellipsoid position.
+
+    Args:
+        unit vector (np.ndarray): A 3-element array representing the unit vector in
+            geocentric (ECEF) coordinates.
+        ellipsoid_position (np.ndarray): A 3-element array specifying the reference position
+            [latitude (rad), longitude (rad), height (meters)] on the Earth's ellipsoid
+
+    Returns:
+        (np.ndarray): A 2-element array [azimuth (rad), altitude (rad)], where:
+            - azimuth is the angle from local north (0 to 2π, increasing eastward),
+            - altitude is the elevation angle from the local horizontal plane (-π/2 to π/2).
+    """
+    lat, lon, _ = ellipsoid_position
+    normal_vector = np.array([np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)])
+    north_vector = np.array([-np.sin(lat) * np.cos(lon), -np.sin(lat) * np.sin(lon), np.cos(lat)])
+    east_vector = np.array([-np.sin(lon), np.cos(lon), 0])
+    altitude = np.arcsin(np.dot(unit_vector, normal_vector))
+    azimuth = np.mod(np.arctan2(np.dot(unit_vector, east_vector), np.dot(unit_vector, north_vector)), 2 * np.pi)
+
+    return np.array([azimuth, altitude])
+
+
+def add_et_2l_detectors_at_location(  # pylint: disable=too-many-locals
+    e1_latitude: float,
+    e1_longitude: float,
+    e1_height: float,
+    e2_latitude: float,
+    e2_longitude: float,
+    e2_height: float,
+    alpha: float,
+    e1_location_name: str,
+    e2_location_name: str,
+    et_arm_l: float = 15000,
+) -> tuple[Detector, Detector]:
+    """
+    Add the 2L Einstein Telescope detectors with PyCBC at two given locations and heights,
+    for a given relative angle alpha.
+    The arms of the detectors are defined on the tangent plane at their vertex position.
+    The arm 1 of E1 has the same azimuth angle and altitude of the Virgo arm 1 in the
+    local horizontal coordinate system center at the E1 vertex.
+    All the angles are measured clockwise in the local horizontal coordinate system (North to East).
+
+    Args:
+        E1_latitude (float): E1 vertex latitude (rad)
+        E1_longitude (float): E1 vertex longitude (rad)
+        E1_height (float): E1 vertex height above the standard reference ellipsoidal earth (meters)
+        E2_latitude (float): E2 vertex latitude (rad)
+        E2_longitude (float): E2 vertex longitude (rad)
+        E2_height (float): E2 vertex height above the standard reference ellipsoidal earth (meters)
+        alpha (float): Relative orientation angle alpha in radians. Alpha is defined
+            as the relative angle between the two detectors, oriented w.r.t their local North.
+        E1_location_name (str): Name of the E1 location (e.g., Sardinia, EMR, Cascina, ...)
+            for detector naming convention
+        E2_location_name (str): Name of the E1 location (e.g., Sardinia, EMR, Cascina, ...)
+            for detector naming convention
+        ETArmL (float, optional): ET arm length (meters). Default to 10000 meters.
+
+    Returns:
+        (Detector, Detector): pycbc.detector.Detector objects for E1, E2.
+    """
+
+    if alpha == 0:
+        config = "Aligned"
+    elif alpha == np.pi / 4:
+        config = "Misaligned"
+    else:
+        raise ValueError("Only alpha = 0 (aligned configuration) and π/4 (misaligned configuration) are supported.")
+
+    # === Detector E1 ===
+
+    e1_ellipsoid = [e1_latitude, e1_longitude, e1_height]
+
+    # E1 vertex location in geocentric (ECEF) coordinates
+    e1 = np.array(pm.geodetic2ecef(*e1_ellipsoid, deg=False))
+
+    # Normal vector to the tangent plane at the E1 vertex (ECEF coordinates)
+    e1_norm_vec = np.array(
+        [np.cos(e1_latitude) * np.cos(e1_longitude), np.cos(e1_latitude) * np.sin(e1_longitude), np.sin(e1_latitude)]
+    )
+
+    # Azimuth and altitude of Virgo arm 1 from LAL
+    v1_arm1_az = 0.3391628563404083
+    v1_arm1_alt = 0.0
+
+    # Define the arm 1 of E1 with the same azimuth and altitude of the Virgo arm 1 (ECEF coordinates)
+    e1_arm1 = np.array(
+        pm.aer2ecef(
+            az=v1_arm1_az, el=v1_arm1_alt, srange=1, lat0=e1_latitude, lon0=e1_longitude, alt0=e1_height, deg=False
+        )
+        - e1
+    )
+
+    # Vector perpendicular to E1Arm1 on the same plane
+    e1_arm2 = np.cross(e1_arm1, e1_norm_vec)
+
+    e1_arm1_angles = get_unit_vector_angles(e1_arm1, e1_ellipsoid)
+    e1_arm2_angles = get_unit_vector_angles(e1_arm2, e1_ellipsoid)
+
+    add_detector_on_earth(  # pylint: disable=duplicate-code
+        name=f"E1_{config}_" + e1_location_name,
+        latitude=e1_ellipsoid[0],
+        longitude=e1_ellipsoid[1],
+        height=e1_ellipsoid[2],
+        xangle=e1_arm1_angles[0],
+        yangle=e1_arm2_angles[0],
+        xaltitude=e1_arm1_angles[1],
+        yaltitude=e1_arm2_angles[1],
+        xlength=et_arm_l,
+        ylength=et_arm_l,
+    )
+
+    # === Detector E2 ===
+
+    e2_ellipsoid = [e2_latitude, e2_longitude, e2_height]
+
+    # E2 vertex location in geocentric (ECEF) coordinates
+    e2 = np.array(pm.geodetic2ecef(*e2_ellipsoid, deg=False))
+
+    # Normal vector to the tangent plane at the E2 vertex (ECEF coordinates)
+    e2_norm_vec = np.array(
+        [np.cos(e2_latitude) * np.cos(e2_longitude), np.cos(e2_latitude) * np.sin(e2_longitude), np.sin(e2_latitude)]
+    )
+
+    # Define the arm 1 of E2 with the same azimuth and altitude of the Virgo arm 1 + alpha (ECEF coordinates)
+    e2_arm1_az = v1_arm1_az + alpha
+    e2_arm1 = np.array(
+        pm.aer2ecef(
+            az=e2_arm1_az, el=v1_arm1_alt, srange=1, lat0=e2_latitude, lon0=e2_longitude, alt0=e2_height, deg=False
+        )
+        - e2
+    )
+
+    # Vector perpendicular to E2Arm1 on the same plane
+    e2_arm2 = np.cross(e2_arm1, e2_norm_vec)
+
+    e2_arm1_angles = get_unit_vector_angles(e2_arm1, e2_ellipsoid)
+    e2_arm2_angles = get_unit_vector_angles(e2_arm2, e2_ellipsoid)
+
+    add_detector_on_earth(
+        name=f"E2_{config}_" + e2_location_name,
+        latitude=e2_ellipsoid[0],
+        longitude=e2_ellipsoid[1],
+        height=e2_ellipsoid[2],
+        xangle=e2_arm1_angles[0],
+        yangle=e2_arm2_angles[0],
+        xaltitude=e2_arm1_angles[1],
+        yaltitude=e2_arm2_angles[1],
+        xlength=et_arm_l,
+        ylength=et_arm_l,
+    )
+
+    return Detector(f"E1_{config}_" + e1_location_name), Detector(f"E2_{config}_" + e2_location_name)