archeo 0.1.0__tar.gz

archeo-0.1.0/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.1
+Name: archeo
+Version: 0.1.0
+Summary: 
+License: MIT
+Author: wyhwong
+Author-email: wyhwong@link.cuhk.edu.hk
+Requires-Python: >=3.11,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: corner (>=2.2.2,<3.0.0)
+Requires-Dist: gwsurrogate (==0.5.0)
+Requires-Dist: matplotlib (>=3.8.2,<4.0.0)
+Requires-Dist: mpl-scatter-density (>=0.7,<0.8)
+Requires-Dist: numpy (>=1.26.3,<2.0.0)
+Requires-Dist: pandas (>=2.1.4,<3.0.0)
+Requires-Dist: pyarrow (>=15.0.0,<16.0.0)
+Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
+Requires-Dist: seaborn (>=0.13.1,<0.14.0)
+Requires-Dist: setuptools (>=70.3.0,<71.0.0)
+Requires-Dist: surfinbh (>=1.2.3,<2.0.0)
+Requires-Dist: tqdm (>=4.66.4,<5.0.0)
+Description-Content-Type: text/markdown
+
+## ARCHEO - Inferring the natal kick and parental masses posterior of black holes in Pair-instability Supernova (PISN) gap.
+
+## Basic Usage
+
+The following example demonstrates how to use the package to visualize the prior and posterior distributions of a single event.
+
+```python
+import archeo
+
+# Load the mass/spin samples from a file
+# They are expected to be a list of floats
+mass_posterior = [68.0, 71.4, ..., 91.4]
+spin_posterior = [0.31, 0.54, ..., 0.64]
+
+# Create a prior (preset priors are "precessing" and "aligned_spin")
+prior = archeo.Prior.from_prior_config("precessing")
+# Create a posterior from the samples and the prior
+posterior = prior.to_posterior(mass_posterior, spin_posterior)
+
+# Visualize the prior and the posterior
+archeo.visualize_prior_distribution(prior, output_dir="./")
+archeo.visualize_posterior_estimation({"GW190521": posterior}, output_dir="./")
+```
+
+## Configure your own prior
+
+Check out the preset priors in [precessing.py](./archeo/preset/precessing.py) and [aligned_spin.py](./archeo/preset/aligned_spin.py). From that, one should be able to create their own prior by following the same structure.
+

archeo-0.1.0/README.md

@@ -0,0 +1,27 @@
+## ARCHEO - Inferring the natal kick and parental masses posterior of black holes in Pair-instability Supernova (PISN) gap.
+
+## Basic Usage
+
+The following example demonstrates how to use the package to visualize the prior and posterior distributions of a single event.
+
+```python
+import archeo
+
+# Load the mass/spin samples from a file
+# They are expected to be a list of floats
+mass_posterior = [68.0, 71.4, ..., 91.4]
+spin_posterior = [0.31, 0.54, ..., 0.64]
+
+# Create a prior (preset priors are "precessing" and "aligned_spin")
+prior = archeo.Prior.from_prior_config("precessing")
+# Create a posterior from the samples and the prior
+posterior = prior.to_posterior(mass_posterior, spin_posterior)
+
+# Visualize the prior and the posterior
+archeo.visualize_prior_distribution(prior, output_dir="./")
+archeo.visualize_posterior_estimation({"GW190521": posterior}, output_dir="./")
+```
+
+## Configure your own prior
+
+Check out the preset priors in [precessing.py](./archeo/preset/precessing.py) and [aligned_spin.py](./archeo/preset/aligned_spin.py). From that, one should be able to create their own prior by following the same structure.

archeo-0.1.0/archeo/__init__.py

@@ -0,0 +1,5 @@
+from archeo.constants import Fits
+from archeo.core.prior import Prior
+from archeo.preset import get_prior_config
+from archeo.schema import PriorConfig
+from archeo.visualization import visualize_posterior_estimation, visualize_prior_distribution

archeo-0.1.0/archeo/constants.py

@@ -0,0 +1,98 @@
+import enum
+
+import archeo.logger
+
+
+local_logger = archeo.logger.get_logger(__name__)
+
+SPEED_OF_LIGHT = 299792.458  # km/s
+
+
+class Fits(enum.Enum):
+    """Surrogate models for binary black hole merger simulations.
+
+    Attributes:
+        NRSur3dq8Remnant: non precessing BHs with mass ratio<=8, anti-/aligned spin <= 0.8.
+        NRSur7dq4Remnant: precessing BHs with mass ratio<=4, generic spin <= 0.8.
+        surfinBH7dq2: precessing BHs with mass ratio <= 2, generic spin <= 0.8.
+
+    Details please refer to https://pypi.org/project/surfinBH/.
+    """
+
+    NRSUR3DQ8REMNANT = "NRSur3dq8Remnant"
+    NRSUR7DQ4REMNANT = "NRSur7dq4Remnant"
+    SURFINBH7DQ2 = "surfinBH7dq2"
+
+    def load(self):
+        """Load a surfinBH fits.
+
+        Returns:
+            fits (surfinBH.surfinBH.SurFinBH): The loaded fits.
+        """
+
+        import surfinBH  # pylint: disable=import-outside-toplevel
+
+        local_logger.info(
+            "Loading surfinBH %s, description: %s.",
+            self.value,
+            surfinBH.fits_collection[self.value].desc,
+        )
+        return surfinBH.LoadFits(self.value)
+
+
+class EscapeVelocity(enum.Enum):
+    """Escape velocity (Unit in km s^-1)"""
+
+    GLOBULAR_CLUSTER = 50.0
+    MILKY_WAY = 600.0
+    NUCLEAR_STAR_CLUSTER = 1500.0
+    ELLIPTICAL_GALAXY = 2500.0
+
+    def label(self):
+        """Return the escape velocity label"""
+
+        if self is EscapeVelocity.GLOBULAR_CLUSTER:
+            return "$v_{esc, GC}$"
+
+        if self is EscapeVelocity.MILKY_WAY:
+            return "$v_{esc, MW}$"
+
+        if self is EscapeVelocity.NUCLEAR_STAR_CLUSTER:
+            return "$v_{esc, NSC}$"
+
+        if self is EscapeVelocity.ELLIPTICAL_GALAXY:
+            return "$v_{esc, EG}$"
+
+        raise ValueError(f"Unknown escape velocity {self}")
+
+    @classmethod
+    def to_vlines(cls) -> dict[str, float]:
+        """Return a dictionary for vlines plotting
+
+        Returns:
+            vlines (Dict[str, float]): The escape velocity vlines.
+        """
+
+        return {esc_vel.label(): esc_vel.value for esc_vel in cls}
+
+
+class Columns(str, enum.Enum):
+    """Columns in the prior dataframe"""
+
+    HEAVIER_BH_MASS = "m_1"
+    HEAVIER_BH_SPIN = "a_1"
+    HEAVIER_BH_CHI = "chi_1"
+    LIGHTER_BH_MASS = "m_2"
+    LIGHTER_BH_SPIN = "a_2"
+    LIGHTER_BH_CHI = "chi_2"
+    MASS_RATIO = "q"
+    RETAINED_MASS = "m_ret"
+    LIKELIHOOD = "l"
+    RECOVERY_RATE = "r_rec"
+    BH_MASS = "m_f"
+    BH_KICK = "k_f"
+    BH_VEL = "v_f"
+    BH_CHI = "chi_f"
+    BH_SPIN = "a_f"
+    BH_EFF_SPIN = "a_eff"
+    BH_PREC_SPIN = "a_prec"

archeo-0.1.0/archeo/core/mahapatra.py

@@ -0,0 +1,51 @@
+from typing import Callable
+
+import numpy as np
+import pandas as pd
+
+import archeo.logger
+from archeo.schema import Domain
+
+
+local_logger = archeo.logger.get_logger(__name__)
+
+
+def get_mahapatra_mass_fn(mass: Domain, alpha=2.3, dm=4.83, n_samples=500000) -> Callable:
+    """Get a mass function from Mahapatra's mass distribution.
+    NOTE: For details, see https://arxiv.org/abs/2209.05766.
+
+    Args:
+        mass (Domain): Mass domain.
+        alpha (float): Power law index.
+        dm (float): Tapering parameter.
+        n_samples (int): Number of samples to generate.
+
+    Returns:
+        mass_fn (Callable): Mass function.
+    """
+
+    def _f(ds: pd.Series) -> pd.Series:
+        """Calculate the function f in Mahapatra's paper"""
+
+        mp = ds - mass.low
+        return np.exp(dm / mp + dm / (mp - dm))
+
+    def smoothing_func(ds: pd.Series) -> pd.Series:
+        """Smoothing function."""
+
+        probis = ds.copy()
+        probis[ds < mass.low + dm] = 1 / (_f(ds[ds < mass.low + dm]) + 1)  # type: ignore
+        probis[ds > mass.low + dm] = 1  # type: ignore
+        probis *= ds ** (-alpha)
+        return probis
+
+    masses = pd.Series(np.random.uniform(mass.low, mass.high, size=n_samples))
+    probis = smoothing_func(masses)
+    probis /= probis.sum()
+
+    def mass_from_mahapatra() -> float:
+        """Generate a mass from Mahapatra's mass distribution."""
+
+        return np.random.choice(masses, p=probis)
+
+    return mass_from_mahapatra

archeo-0.1.0/archeo/core/prior.py

@@ -0,0 +1,257 @@
+from typing import Union
+
+import numpy as np
+import pandas as pd
+
+import archeo.logger
+from archeo.constants import SPEED_OF_LIGHT
+from archeo.constants import Columns as C
+from archeo.core.simulator import Simulator
+from archeo.preset import get_prior_config
+from archeo.schema import PriorConfig
+from archeo.utils.executor import MultiThreadExecutor
+
+
+local_logger = archeo.logger.get_logger(__name__)
+
+
+class Prior(pd.DataFrame):
+    """A class to represent the prior distribution."""
+
+    def __init__(
+        self,
+        *args,
+        ignore_simulated_mass: bool = False,
+        sample_ratio: int = 1,
+        spin_tolerance: float = 0.05,  # unit: dimensionless
+        mass_tolerance: float = 1.0,  # unit: solar mass
+        **kwargs,
+    ) -> None:
+        """Construct a prior dataframe.
+
+        Args:
+            ignore_simulated_mass: Whether to ignore the simulated mass
+            sample_ratio (int): The number of samples to be sampled each time
+            spin_tolerance (float): The tolerance of the spin
+            mass_tolerance (float): The tolerance of the mass
+        """
+
+        super().__init__(*args, **kwargs)
+
+        self._ignore_simulated_mass = ignore_simulated_mass
+        self._sample_ratio = sample_ratio
+        self._spin_tolerance = spin_tolerance
+        self._mass_tolerance = mass_tolerance
+
+    def _sample_from_possible_samples(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Sample from a dataframe.
+
+        Args:
+            df (pd.DataFrame): The dataframe to sample from.
+
+        Returns:
+            df (pd.DataFrame): The sampled dataframe.
+        """
+
+        if df.empty:
+            local_logger.warning("No similar samples in the prior.")
+        else:
+            df = df.sample(self._sample_ratio, replace=True)
+        return df
+
+    def retrieve_samples(self, spin_measure: float, mass_measure: float) -> pd.DataFrame:
+        """Retrieve the samples from prior.
+
+        Args:
+            spin_measure (float): The measured spin
+            mass_measure (float): The measured mass
+
+        Returns:
+            pd.DataFrame: The sampled dataframe
+        """
+
+        if not self._ignore_simulated_mass:
+            # Find the possible samples in the prior
+            # Based on:
+            #    1. mass_prior - tol < mass_measure < mass_prior + tol
+            #    2. spin_prior - tol < spin_measure < spin_prior + tol
+            possible_samples = self.loc[
+                ((self[C.BH_MASS] - mass_measure).abs() < self._mass_tolerance)
+                & ((self[C.BH_SPIN] - spin_measure).abs() < self._spin_tolerance)
+            ]
+            likelihood = len(possible_samples) / len(self)
+
+            # Sample n_sample samples from the possible samples
+            samples = self._sample_from_possible_samples(possible_samples)
+            samples[C.LIKELIHOOD] = likelihood
+            return samples
+
+        # Find the possible samples in the prior
+        # Based on:
+        #    1. spin_prior - tol < spin_measure < spin_prior + tol
+        possible_samples = self.loc[(self[C.BH_SPIN] - spin_measure).abs() < self._spin_tolerance]
+        likelihood = len(possible_samples) / len(self)
+
+        # Sample n_sample samples from the possible samples
+        samples = self._sample_from_possible_samples(possible_samples)
+
+        # Calculate the mass parameters (for mass not injected case)
+        samples[C.HEAVIER_BH_MASS] = (
+            mass_measure / samples[C.RETAINED_MASS] * samples[C.MASS_RATIO] / (1 + samples[C.MASS_RATIO])
+        )
+        samples[C.LIGHTER_BH_MASS] = mass_measure / samples[C.RETAINED_MASS] / (1 + samples[C.MASS_RATIO])
+        samples[C.BH_MASS] = mass_measure
+        samples[C.LIKELIHOOD] = likelihood
+        return samples
+
+    @property
+    def _constructor(self):
+        """Return the constructor of the class."""
+
+        return pd.DataFrame
+
+    @classmethod
+    def from_feather(cls, path: str, **kwargs) -> "Prior":
+        """Read the feather file.
+
+        Args:
+            path (str): The path to the feather file.
+            **kwargs: The keyword arguments for parental class (pd.DataFrame).
+
+        Returns:
+            Prior: The prior distribution.
+        """
+
+        return cls(pd.read_feather(path), **kwargs)
+
+    @classmethod
+    def from_csv(cls, path: str, **kwargs) -> "Prior":
+        """Read the csv file.
+
+        Args:
+            path (str): The path to the csv file.
+            **kwargs: The keyword arguments for parental class (pd.DataFrame).
+
+        Returns:
+            Prior: The prior distribution
+        """
+
+        return cls(pd.read_csv(path), **kwargs)
+
+    @classmethod
+    def from_parquet(cls, path: str, **kwargs) -> "Prior":
+        """Read the parquet file.
+
+        Args:
+            path (str): The path to the parquet file.
+            **kwargs: The keyword arguments for parental class (pd.DataFrame).
+
+        Returns:
+            Prior: The prior distribution.
+        """
+
+        return cls(pd.read_parquet(path), **kwargs)
+
+    @classmethod
+    def from_prior_config(cls, prior_config: Union[PriorConfig, str], use_threads=True, **kwargs) -> "Prior":
+        """Generate the prior from the prior config.
+
+        Args:
+            prior_config (PriorConfig): The prior configuration.
+            use_threads (bool): Whether to use threads.
+            **kwargs: The keyword arguments for the class.
+
+        Returns:
+            Prior: The prior distribution.
+        """
+
+        if isinstance(prior_config, str):
+            prior_config = get_prior_config(prior_config)
+
+        simulator = Simulator(prior_config)
+        return cls.from_simulator(simulator, use_threads=use_threads, **kwargs)
+
+    @classmethod
+    def from_simulator(cls, simulator: Simulator, use_threads=True, **kwargs) -> "Prior":
+        """Generate the prior from the simulator.
+
+        Args:
+            simulator (Simulator): The simulator.
+            use_threads (bool): Whether to use threads.
+            **kwargs: The keyword arguments for the class.
+
+        Returns:
+            Prior: The prior distribution.
+        """
+
+        df = cls(simulator.simulate(use_threads=use_threads), **kwargs)
+
+        # Extract more information from the samples
+
+        # Define nan recovery rate
+        df[C.RECOVERY_RATE] = float("nan")
+
+        # Calculate the mass ratio
+        m1, m2 = df[C.HEAVIER_BH_MASS], df[C.LIGHTER_BH_MASS]
+        df[C.MASS_RATIO] = q = m1 / m2
+
+        # Calculate the remnant mass
+        df[C.BH_MASS] = df[C.RETAINED_MASS] * (m1 + m2)
+
+        # Calculate the BH kick velocity
+        df[C.BH_KICK] = df[C.BH_VEL].apply(lambda vf: np.sqrt(np.dot(vf, vf)) * SPEED_OF_LIGHT)
+
+        # Calculate the BH spin
+        df[C.BH_SPIN] = df[C.BH_CHI].apply(lambda vf: np.sqrt(np.dot(vf, vf)))
+
+        # Calculate the parental spins
+        df[C.HEAVIER_BH_SPIN] = df[C.HEAVIER_BH_CHI].apply(lambda chi: np.sqrt(np.dot(chi, chi)))
+        df[C.LIGHTER_BH_SPIN] = df[C.LIGHTER_BH_CHI].apply(lambda chi: np.sqrt(np.dot(chi, chi)))
+
+        # Calculate the effective spin
+        a1z = df[C.HEAVIER_BH_CHI].apply(lambda chi: chi[-1])
+        a2z = df[C.LIGHTER_BH_CHI].apply(lambda chi: chi[-1])
+        df[C.BH_EFF_SPIN] = (m1 * a1z + m2 * a2z) / (m1 + m2)
+
+        # Calculate the precession spin
+        a1h = df[C.HEAVIER_BH_CHI].apply(lambda chi: np.sqrt(chi[0] ** 2 + chi[1] ** 2))
+        a2h = df[C.LIGHTER_BH_CHI].apply(lambda chi: np.sqrt(chi[0] ** 2 + chi[1] ** 2))
+        df[C.BH_PREC_SPIN] = np.maximum(a1h, (4 / q + 3) / (3 / q + 4) / q * a2h)
+
+        return df
+
+    def to_posterior(
+        self,
+        mass_posterior: list[float],
+        spin_posterior: list[float],
+        use_threads=True,
+    ) -> pd.DataFrame:
+        """Convert the prior to the posterior.
+
+        Args:
+            mass_posterior (list[float]): The posterior mass.
+            spin_posterior (list[float]): The posterior spin.
+            use_threads (bool): Whether to use threads.
+
+        Returns:
+            pd.DataFrame: The posterior distribution.
+        """
+
+        if use_threads:
+            exc = MultiThreadExecutor()
+            input_kwargs = [
+                dict(spin_measure=spin_measure, mass_measure=mass_measure)
+                for spin_measure, mass_measure in zip(spin_posterior, mass_posterior)
+            ]
+            samples = exc.run(func=self.retrieve_samples, input_kwargs=input_kwargs)
+
+        else:
+            samples = [
+                self.retrieve_samples(spin_measure=spin_measure, mass_measure=mass_measure)
+                for spin_measure, mass_measure in zip(spin_posterior, mass_posterior)
+            ]
+
+        df_posterior = pd.concat(samples)
+        df_posterior[C.RECOVERY_RATE] = len(df_posterior) / (len(mass_posterior) * self._sample_ratio)
+
+        return df_posterior

archeo-0.1.0/archeo/core/simulator.py

@@ -0,0 +1,127 @@
+from dataclasses import asdict
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from archeo.core.mahapatra import get_mahapatra_mass_fn
+from archeo.schema import Binary, Event, PriorConfig
+from archeo.utils.executor import MultiThreadExecutor
+from archeo.utils.math import sph2cart
+
+
+class Simulator:
+    """Simulator class to simulate binary black hole mergers"""
+
+    def __init__(self, prior_config: PriorConfig):
+        """Initializes the simulator with the given prior configuration
+
+        Args:
+            prior_config (PriorConfig): The prior configuration to use
+        """
+
+        self._prior_config = prior_config
+
+        self._fits = self._prior_config.fits.load()
+        self._n_samples = self._prior_config.n_samples
+
+        if self._prior_config.is_mahapatra:
+            self._mass_fn = get_mahapatra_mass_fn(mass=self._prior_config.mass)
+        else:
+            self._mass_fn = self._prior_config.mass.draw
+
+        self._theta_fn = self._prior_config.theta.draw
+        self._phi_fn = self._prior_config.phi.draw
+        self._spin_fn = self._prior_config.spin.draw
+
+        self._q_bounds = self._prior_config.mass_ratio
+
+        self._is_spin_aligned = self._prior_config.is_spin_aligned
+        self._only_up_aligned_spin = self._prior_config.is_only_up_aligned_spin
+
+    def __call__(self) -> Event:
+        """Simulates a binary black hole merger event
+
+        Returns:
+            Event: The simulated event
+        """
+
+        b = self._get_binary()
+
+        q = b.m_1 / b.m_2
+        v_f, _ = self._fits.vf(q, b.chi_1, b.chi_2)
+        chi_f, _ = self._fits.chif(q, b.chi_1, b.chi_2)
+        m_ret, _ = self._fits.mf(q, b.chi_1, b.chi_2)
+
+        return Event(
+            m_1=b.m_1,
+            m_2=b.m_2,
+            m_ret=m_ret,
+            v_f=v_f,
+            chi_1=b.chi_1,
+            chi_2=b.chi_2,
+            chi_f=chi_f,
+        )
+
+    def simulate(self, use_threads=True) -> pd.DataFrame:
+        """Simulates multiple binary black hole merger events"""
+
+        if use_threads:
+            exc = MultiThreadExecutor()
+            events = exc.run(self, [{} for _ in range(self._n_samples)])
+        else:
+            events = [self() for _ in tqdm(range(self._n_samples))]
+
+        df = pd.DataFrame([asdict(event) for event in events])
+        return df
+
+    def _get_binary(self) -> Binary:
+        """Draws a binary from the prior distribution
+
+        Returns:
+            Binary: The drawn binary
+        """
+
+        m_1, m_2 = self._get_masses()
+        chi_1, chi_2 = self._get_spin(), self._get_spin()
+
+        return Binary(m_1=m_1, m_2=m_2, chi_1=chi_1, chi_2=chi_2)
+
+    def _get_spin(self) -> tuple[float, float, float]:
+        """Draws the spin of the binary
+
+        Returns:
+            tuple[float, float, float]: The drawn spin
+        """
+
+        spin = self._spin_fn()
+
+        if self._is_spin_aligned:
+            if self._only_up_aligned_spin:
+                return (0, 0, spin)
+
+            direction = np.random.choice([-1, 1])
+            return (0, 0, direction * spin)
+
+        theta = np.arccos(-1 + 2 * self._theta_fn())
+        phi = self._phi_fn() * np.pi
+        univ = sph2cart(theta=theta, phi=phi)
+        return tuple(spin * univ)
+
+    def _get_masses(self) -> tuple[float, float]:
+        """Draws the masses of the binary
+
+        Returns:
+            tuple[float, float]: The drawn masses
+        """
+
+        masses = (self._mass_fn(), self._mass_fn())
+        m_1, m_2 = sorted(masses, reverse=True)
+
+        # Check whether the mass ratio is in the domain
+        # If not, resample the masses (recursion)
+        q = m_1 / m_2
+        if not self._q_bounds.contain(q):
+            return self._get_masses()
+
+        return (m_1, m_2)

archeo-0.1.0/archeo/env.py

@@ -0,0 +1,14 @@
+import os
+import zoneinfo
+
+
+# For logging
+STREAMING_LOG_LEVEL = int(os.getenv("STREAMING_LOG_LEVEL", "30"))
+FILE_LOG_LEVEL = int(os.getenv("FILE_LOG_LEVEL", "10"))
+LOG_FILEPATH = os.getenv("LOG_FILEPATH", "./runtime.log")
+LOG_FMT = os.getenv("LOG_FMT", "%(asctime)s [%(name)s | %(levelname)s]: %(message)s")
+LOG_DATEFMT = os.getenv("LOG_DATEFMT", "%Y-%m-%dT%H:%M:%SZ")
+TZ = zoneinfo.ZoneInfo(os.getenv("TZ") or "UTC")
+
+# For simulation
+MAX_MULTITHREAD_WORKER = int(os.getenv("MAX_MULTITHREAD_WORKER", "20"))