archeo 2.0.0.dev7__tar.gz → 2.0.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archeo
-Version: 2.0.0.dev7
+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.dev7 → 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.0.dev7 → archeo-2.0.1}/archeo/constants/physics.py

@@ -1,10 +1,14 @@
 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):
@@ -36,7 +40,6 @@ class TypicalHostEscapeVelocity(enum.Enum):
     def compute_p2g(
         self,
         df: pd.DataFrame,
-        max_mass: float = 65.0,
         kf_col: str = "k_f",
         m1_col: str = "m_1",
         m2_col: str = "m_2",
@@ -46,7 +49,7 @@ class TypicalHostEscapeVelocity(enum.Enum):
         if df.empty:
             return 0.0
 
-        mask = (df[kf_col] <= self.v_esc) & (df[m1_col] <= max_mass) & (df[m2_col] <= max_mass)
+        mask = (df[kf_col] <= self.v_esc) & (df[m1_col] <= PISN_LB) & (df[m2_col] <= PISN_LB)
         return mask.mean() * 100.0
 
     @classmethod

{archeo-2.0.0.dev7 → 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.dev7 → 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 TypicalHostEscapeVelocity
+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:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archeo
-Version: 2.0.0.dev7
+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.dev7 → 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.dev7 → 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.dev7"
+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",