archeo 2.0.0.dev6__tar.gz → 2.0.1__tar.gz

{archeo-2.0.0.dev6 → archeo-2.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archeo
-Version: 2.0.0.dev6
+Version: 2.0.1
 Summary: Bayesian framework for inferring natal kick, ancestral masses and spins of black holes.
 Author-email: wyhwong <wyhwong@link.cuhk.edu.hk>
 License-Expression: MIT

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo/bayesian/importance_sampling/resampler/generic.py

@@ -48,8 +48,9 @@ class ISDataGeneric(ImportanceSamplingDataBase):
         weights_matrix = self._safe_divide(1.0, prior_hist)
 
         def _get_pdf(row: pd.Series):
-            idx = tuple(np.searchsorted(edges[col], row[col], side="right") - 1 for col in self.common_columns)
-            return weights_matrix[idx]
+            indices = tuple(np.searchsorted(edges[col], row[col], side="right") - 1 for col in self.common_columns)
+            indices = tuple(max(0, min(weights_matrix.shape[i] - 1, idx)) for i, idx in enumerate(indices))
+            return weights_matrix[indices]
 
         weights = self.posterior_samples.apply(_get_pdf, axis=1)
 
@@ -73,8 +74,9 @@ class ISDataGeneric(ImportanceSamplingDataBase):
         weights_matrix = self._safe_divide(new_prior_hist, prior_hist)
 
         def _get_pdf(row: pd.Series):
-            idx = tuple(np.searchsorted(edges[col], row[col], side="right") - 1 for col in self.common_columns)
-            return weights_matrix[idx]
+            indices = tuple(np.searchsorted(edges[col], row[col], side="right") - 1 for col in self.common_columns)
+            indices = tuple(max(0, min(weights_matrix.shape[i] - 1, idx)) for i, idx in enumerate(indices))
+            return weights_matrix[indices]
 
         return self.posterior_samples.apply(_get_pdf, axis=1)
 

archeo-2.0.1/archeo/constants/bayesian.py

@@ -0,0 +1,3 @@
+DEFAULT_BINSIZE_MASS: float = 2.0
+DEFAULT_BINSIZE_SPIN: float = 0.1
+ASSUME_PARAMETER_INDEPENDENCE: bool = False

archeo-2.0.1/archeo/constants/physics.py

@@ -0,0 +1,61 @@
+import enum
+import os
+
+import pandas as pd
+from pydantic import BaseModel, NonNegativeFloat
+
+
+SPEED_OF_LIGHT = 299792.458  # km/s
+BH_MASS_LB: float = float(os.environ.get("BH_MASS_LB", 5.0))  # Solar masses
+PISN_LB: float = float(os.environ.get("PISN_LB", 65.0))  # Solar masses
+PISN_UB: float = float(os.environ.get("PISN_UB", 130.0))  # Solar masses
+
+
+class _TypicalHostEscapeVelocityMeta(BaseModel):
+    v_esc: NonNegativeFloat  # km s^-1
+    short: str
+    latex: str
+
+
+class TypicalHostEscapeVelocity(enum.Enum):
+    """Escape velocity (units: km s^-1)."""
+
+    GLOBULAR_CLUSTER = _TypicalHostEscapeVelocityMeta(v_esc=50.0, short="GC", latex=r"$v_{esc, GC}$")
+    MILKY_WAY = _TypicalHostEscapeVelocityMeta(v_esc=600.0, short="MW", latex=r"$v_{esc, MW}$")
+    NUCLEAR_STAR_CLUSTER = _TypicalHostEscapeVelocityMeta(v_esc=1500.0, short="NSC", latex=r"$v_{esc, NSC}$")
+    ELLIPTICAL_GALAXY = _TypicalHostEscapeVelocityMeta(v_esc=2500.0, short="EG", latex=r"$v_{esc, EG}$")
+
+    @property
+    def v_esc(self) -> float:
+        return self.value.v_esc
+
+    @property
+    def short(self) -> str:
+        return self.value.short
+
+    @property
+    def latex(self) -> str:
+        return self.value.latex
+
+    def compute_p2g(
+        self,
+        df: pd.DataFrame,
+        kf_col: str = "k_f",
+        m1_col: str = "m_1",
+        m2_col: str = "m_2",
+    ) -> float:
+        """Probability of being 2nd-generation under this escape velocity."""
+
+        if df.empty:
+            return 0.0
+
+        mask = (df[kf_col] <= self.v_esc) & (df[m1_col] <= PISN_LB) & (df[m2_col] <= PISN_LB)
+        return mask.mean() * 100.0
+
+    @classmethod
+    def latex_to_values(cls) -> dict[str, float]:
+        return {m.latex: m.v_esc for m in cls}
+
+    @classmethod
+    def short_to_values(cls) -> dict[str, float]:
+        return {m.short: m.v_esc for m in cls}

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo/data_structures/physics/black_hole.py

@@ -4,6 +4,7 @@ import numpy as np
 import pandas as pd
 from pydantic import BaseModel, NonNegativeFloat, PositiveFloat, field_validator
 
+from archeo.constants.physics import BH_MASS_LB, PISN_LB
 from archeo.data_structures.annotation import Distribution
 from archeo.data_structures.distribution import Uniform
 
@@ -35,7 +36,7 @@ BlackHoles: TypeAlias = list[BlackHole]
 class BlackHoleGenerator(BaseModel, frozen=True):
     """Black hole generator data class."""
 
-    mass_distribution: Distribution = Uniform(low=5, high=65)
+    mass_distribution: Distribution = Uniform(low=BH_MASS_LB, high=PISN_LB)
     spin_magnitude_distribution: Distribution = Uniform(low=0, high=1)
     phi_distribution: Distribution = Uniform(low=0, high=2 * np.pi)
     theta_distribution: Distribution = Uniform(low=0, high=np.pi)

archeo-2.0.1/archeo/postprocessing/eval_utils/bias.py

@@ -0,0 +1,13 @@
+import pandas as pd
+
+
+def compute_bias_for_remnant_spin(df_samples: pd.DataFrame) -> float:
+    """Compute the bias in remnant spin for a given set of samples."""
+
+    return (df_samples["a_f"] - df_samples["spin_measure"]).sum() / df_samples.shape[0]
+
+
+def compute_bias_for_remnant_mass(df_samples: pd.DataFrame) -> float:
+    """Compute the bias in remnant mass for a given set of samples."""
+
+    return (df_samples["m_f"] - df_samples["mass_measure"]).sum() / df_samples.shape[0]

archeo-2.0.1/archeo/postprocessing/eval_utils/kl.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+
+def compute_kl_divergence_from_samples(X: np.ndarray, Y: np.ndarray) -> float:
+    """
+    Compute KL(N0 || N1) where N0 and N1 are multivariate Gaussians.
+    X and Y should be 2D arrays of shape (n, d) and (m, d) respectively.
+    """
+
+    if X.shape[1] != Y.shape[1]:
+        raise ValueError("X and Y must have the same dimension.")
+
+    k = X.shape[1]
+    mu0 = X.mean(axis=0)
+    mu1 = Y.mean(axis=0)
+    S0 = np.cov(X.T, bias=False)
+    S1 = np.cov(Y.T, bias=False)
+
+    invS1 = np.linalg.inv(S1)
+    diff = (mu1 - mu0).reshape(k, 1)
+
+    term_trace = np.trace(invS1 @ S0)
+    term_quad = float((diff.T @ invS1 @ diff)[0, 0])
+    term_logdet = np.log(np.linalg.det(S1) / np.linalg.det(S0))
+
+    return 0.5 * (term_trace + term_quad - k + term_logdet)

archeo-2.0.1/archeo/postprocessing/evaluation.py

@@ -0,0 +1,29 @@
+import pandas as pd
+
+from archeo.postprocessing.eval_utils.bias import compute_bias_for_remnant_mass, compute_bias_for_remnant_spin
+from archeo.postprocessing.eval_utils.kl import compute_kl_divergence_from_samples
+
+
+def evaluate_ancestral_inference(df_samples: pd.DataFrame) -> dict[str, float]:
+    """Evaluate the validity of the estimated posterior samples by computing:
+    1. Estimation bias: the average difference between the estimated parameters and the true parameters
+    2. KL divergence between the Gaussian fitted to the estimated samples and the Gaussian fitted to
+       the true parameters.
+    """
+
+    bias_spin = compute_bias_for_remnant_spin(df_samples)
+    bias_mass = compute_bias_for_remnant_mass(df_samples)
+
+    mask = df_samples[["a_f", "m_f"]].notna().all(axis=1)
+    # Here we need to apply the mask to the inferred samples.
+    # The reason is that there could be no similar samples
+    # in the ancestral prior, resulting in NaN values.
+    # Those NaN values would cause the KL divergence computation to fail.
+    kl_div = compute_kl_divergence_from_samples(
+        df_samples.loc[mask, ["a_f", "m_f"]].values, df_samples[["spin_measure", "mass_measure"]].values
+    )
+    return {
+        "bias_spin": bias_spin,
+        "bias_mass": bias_mass,
+        "kl_divergence": kl_div,
+    }

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo/visualization/base.py

@@ -5,7 +5,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
-from archeo.constants.physics import EscapeVelocity
+from archeo.constants.physics import TypicalHostEscapeVelocity
 from archeo.data_structures.visualization import Labels, Padding
 from archeo.utils.fs import check_and_create_dir
 from archeo.utils.logger import get_logger
@@ -112,7 +112,7 @@ def add_escape_velocity(ax, v_max: float, y_max: float, log_xscale: bool = False
 
     colors = iter(mcolors.TABLEAU_COLORS.keys())
     # Plot vertical lines and labels (escape velocities)
-    for label, v_esc in EscapeVelocity.to_vlines().items():
+    for label, v_esc in TypicalHostEscapeVelocity.latex_to_values().items():
         # Skip if out of scope
         if v_esc > v_max:
             return

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo/visualization/estimation.py

@@ -8,7 +8,7 @@ import numpy as np
 import pandas as pd
 import seaborn as sns
 
-from archeo.constants.physics import EscapeVelocity
+from archeo.constants.physics import PISN_LB, PISN_UB, TypicalHostEscapeVelocity
 from archeo.data_structures.visualization import Labels, Padding
 from archeo.utils.logger import get_logger
 from archeo.visualization import base
@@ -87,8 +87,8 @@ def _add_pisn_gap(ax, color: str) -> None:
         color (str): Color.
     """
 
-    ax.axvline(65.0, color=color, linewidth=0.9, linestyle="--", label="PISN Gap")
-    ax.axvline(130.0, color=color, linewidth=0.9, linestyle="--")
+    ax.axvline(PISN_LB, color=color, linewidth=0.9, linestyle="--", label="PISN Gap")
+    ax.axvline(PISN_UB, color=color, linewidth=0.9, linestyle="--")
 
 
 def corner_estimates(  # pylint: disable=dangerous-default-value
@@ -235,7 +235,7 @@ def second_generation_probability_curve(
         # Calculate the CDF
         y = []
         for kick in x:
-            df_samples = df.loc[(df["k_f"] <= kick) & (df["m_1"] <= 65) & (df["m_2"] <= 65)]
+            df_samples = df.loc[(df["k_f"] <= kick) & (df["m_1"] <= PISN_LB) & (df["m_2"] <= PISN_LB)]
             if df_samples.empty:
                 y.append(0.0)
             else:
@@ -363,7 +363,7 @@ def table_estimates(
     data = {
         "": dfs.keys(),
         "Recovery Rate": [df["m_1"].notna().sum() / df.shape[0] for df in dfs.values()],
-        **{f"p2g_{v_esc.short()}": [v_esc.compute_p2g(df) for df in dfs.values()] for v_esc in EscapeVelocity},
+        **{f"p2g_{v_esc.short}": [v_esc.compute_p2g(df) for df in dfs.values()] for v_esc in TypicalHostEscapeVelocity},
     }
 
     for col, name in col_to_names.items():

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archeo
-Version: 2.0.0.dev6
+Version: 2.0.1
 Summary: Bayesian framework for inferring natal kick, ancestral masses and spins of black holes.
 Author-email: wyhwong <wyhwong@link.cuhk.edu.hk>
 License-Expression: MIT

{archeo-2.0.0.dev6 → archeo-2.0.1}/archeo.egg-info/SOURCES.txt

@@ -36,6 +36,10 @@ archeo/data_structures/physics/mahapatra.py
 archeo/data_structures/physics/simulation.py
 archeo/postprocessing/__init__.py
 archeo/postprocessing/dataframe.py
+archeo/postprocessing/evaluation.py
+archeo/postprocessing/eval_utils/__init__.py
+archeo/postprocessing/eval_utils/bias.py
+archeo/postprocessing/eval_utils/kl.py
 archeo/preset/__init__.py
 archeo/preset/cli.py
 archeo/preset/forward/__init__.py

{archeo-2.0.0.dev6 → archeo-2.0.1}/pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
     "Operating System :: OS Independent"
 ]
 keywords = ["black-holes", "gravitational-waves", "black-hole-archeology"]
-version = "2.0.0.dev6"
+version = "2.0.1"
 readme = "README.md"
 requires-python = ">=3.11,<3.14"
 dependencies = [
@@ -43,7 +43,7 @@ documentation = "https://wyhwong.github.io/archeo/"
 
 [dependency-groups]
 dev = [
-    "ipykernel>=6.9.2, <7.0.0",
+    "ipykernel>=7.2.0",
     "pre-commit>=4.5.1",
     "pyarrow>=23.0.1",
     "pyinstrument>=5.1.2",

archeo-2.0.0.dev6/archeo/constants/bayesian.py

@@ -1,3 +0,0 @@
-DEFAULT_BINSIZE_MASS = 2.0
-DEFAULT_BINSIZE_SPIN = 0.1
-ASSUME_PARAMETER_INDEPENDENCE = False

archeo-2.0.0.dev6/archeo/constants/physics.py

@@ -1,70 +0,0 @@
-import enum
-
-import pandas as pd
-
-
-SPEED_OF_LIGHT = 299792.458  # km/s
-
-
-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) -> str:
-        """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}")
-
-    def short(self) -> str:
-        """Return the short name of the escape velocity"""
-
-        if self is EscapeVelocity.GLOBULAR_CLUSTER:
-            return "GC"
-
-        if self is EscapeVelocity.MILKY_WAY:
-            return "MW"
-
-        if self is EscapeVelocity.NUCLEAR_STAR_CLUSTER:
-            return "NSC"
-
-        if self is EscapeVelocity.ELLIPTICAL_GALAXY:
-            return "EG"
-
-        raise ValueError(f"Unknown escape velocity {self}")
-
-    def compute_p2g(self, df: pd.DataFrame) -> float:
-        """Return the probability of
-        the black hole being a 2nd generation black hole
-        under different escape velocity conditions."""
-
-        if df.empty:
-            return 0.0
-
-        mask = (df["k_f"] <= self.value) & (df["m_1"] <= 65) & (df["m_2"] <= 65)
-        return mask.sum() / len(df) * 100.0
-
-    @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}