ml4gw 0.7.4__py3-none-any.whl → 0.7.5__py3-none-any.whl

ml4gw/distributions.py

@@ -6,13 +6,18 @@ from the corresponding distribution.
 """
 
 import math
-from typing import Optional
+from typing import Callable, Optional
 
 import torch
 import torch.distributions as dist
 from jaxtyping import Float
 from torch import Tensor
 
+from ml4gw.constants import C
+
+_PLANCK18_H0 = 67.66  # Hubble constant in km/s/Mpc
+_PLANCK18_OMEGA_M = 0.30966  # Matter density parameter
+
 
 class Cosine(dist.Distribution):
     """
@@ -173,3 +178,202 @@ class DeltaFunction(dist.Distribution):
         return self.peak * torch.ones(
             sample_shape, device=self.peak.device, dtype=torch.float32
         )
+
+
+class UniformComovingVolume(dist.Distribution):
+    """
+    Sample either redshift, comoving distance, or luminosity distance
+    such that they are uniform in comoving volume, assuming a flat
+    lambda-CDM cosmology. Default H0 and Omega_M values match
+    astropy.cosmology.Planck18
+
+    Args:
+        minimum: Minimum distance in the specified distance type
+        maximum: Maximum distance in the specified distance type
+        distance_type:
+            Type of distance to sample from. Can be 'redshift',
+            'comoving_distance', or 'luminosity_distance'
+        h0: Hubble constant in km/s/Mpc
+        omega_m: Matter density parameter
+        z_max: Maximum redshift for the grid
+        grid_size: Number of points in the grid for interpolation
+        validate_args: Whether to validate arguments
+    """
+
+    arg_constraints = {}
+    support = dist.constraints.nonnegative
+
+    def __init__(
+        self,
+        minimum: float,
+        maximum: float,
+        distance_type: str = "redshift",
+        h0: float = _PLANCK18_H0,
+        omega_m: float = _PLANCK18_OMEGA_M,
+        z_grid_max: float = 5,
+        grid_size: int = 10000,
+        validate_args: bool = None,
+    ):
+        super().__init__(validate_args=validate_args)
+        if distance_type not in [
+            "redshift",
+            "comoving_distance",
+            "luminosity_distance",
+        ]:
+            raise ValueError(
+                "Distance type must be 'redshift', 'comoving_distance', "
+                f"or 'luminosity_distance'; got {distance_type}"
+            )
+
+        self.minimum = minimum
+        self.maximum = maximum
+        self.distance_type = distance_type
+        self.grid_size = grid_size
+        self.z_grid_max = z_grid_max
+        self.h0 = h0
+        self.omega_m = omega_m
+
+        # Compute redshift range based on the given min and max distances
+        z_min, z_max = self._get_z_bounds()
+        if z_max > z_grid_max:
+            raise ValueError(
+                f"Maximum {distance_type} {maximum} "
+                f"exceeds given z_max {z_grid_max}."
+            )
+
+        # Restrict distance grids to the specified redshift range
+        mask = (self.z_grid >= z_min) & (self.z_grid <= z_max)
+        self.distance_grid = self.distance_grid[mask]
+        self.z_grid = self.z_grid[mask]
+        self.comoving_dist_grid = self.comoving_dist_grid[mask]
+        self.luminosity_dist_grid = self.luminosity_dist_grid[mask]
+        # Compute probability arrays from those grids
+        self._generate_probability_grids()
+
+    def _hubble_function(self):
+        """
+        Compute H(z) assuming a flat lambda-CDM cosmology.
+        """
+        omega_l = 1 - self.omega_m
+        return self.h0 * torch.sqrt(
+            self.omega_m * (1 + self.z_grid) ** 3 + omega_l
+        )
+
+    def _get_z_bounds(self):
+        """
+        Compute the bounds on redshift based on the given minimum and maximum
+        distances, using the specified distance type.
+        """
+        self._generate_distance_grids()
+        bounds = torch.tensor([self.minimum, self.maximum])
+        z_min, z_max = self._linear_interp_1d(
+            self.distance_grid, self.z_grid, bounds
+        )
+
+        return z_min, z_max
+
+    def _generate_distance_grids(self):
+        """
+        Generate distance grids based on the specified redshift range.
+        """
+        self.z_grid = torch.linspace(0, self.z_grid_max, self.grid_size)
+        self.dz = self.z_grid[1] - self.z_grid[0]
+        # C is specfied in m/s, h0 in km/s/Mpc, so divide by 1000 to convert
+        comoving_dist_grid = (
+            torch.cumulative_trapezoid(
+                (C / self._hubble_function()), self.z_grid
+            )
+            / 1000
+        )
+        zero_prefix = torch.zeros(1, dtype=comoving_dist_grid.dtype)
+        self.comoving_dist_grid = torch.cat([zero_prefix, comoving_dist_grid])
+        self.luminosity_dist_grid = self.comoving_dist_grid * (1 + self.z_grid)
+
+        if self.distance_type == "redshift":
+            self.distance_grid = self.z_grid
+        elif self.distance_type == "comoving_distance":
+            self.distance_grid = self.comoving_dist_grid
+        else:  # luminosity_distance
+            self.distance_grid = self.luminosity_dist_grid
+
+    def _p_of_distance(self):
+        """
+        Compute the unnormalized probability as a function of distance
+        """
+        dV_dz = self.comoving_dist_grid**2 / self._hubble_function()
+        # This is a tensor of ones if the distance type is redshift
+        jacobian = torch.gradient(self.distance_grid, spacing=self.dz)[0]
+        return dV_dz / jacobian
+
+    def _generate_probability_grids(self):
+        """
+        Compute the pdf, cdf, and log pdf based on the
+        comoving volume differential and distance grid.
+        """
+        p_of_distance = self._p_of_distance()
+        self.pdf = p_of_distance / torch.trapz(
+            p_of_distance, self.distance_grid
+        )
+        cdf = torch.cumulative_trapezoid(self.pdf, self.distance_grid)
+        zero_prefix = torch.zeros(1, dtype=cdf.dtype)
+        self.cdf = torch.cat([zero_prefix, cdf])
+        self.log_pdf = torch.log(self.pdf)
+
+    def _linear_interp_1d(self, x_grid, y_grid, x_query):
+        idx = torch.bucketize(x_query, x_grid, right=True)
+        idx = idx.clamp(min=1, max=len(x_grid) - 1)
+
+        x0 = x_grid[idx - 1]
+        x1 = x_grid[idx]
+        y0 = y_grid[idx - 1]
+        y1 = y_grid[idx]
+
+        t = (x_query - x0) / (x1 - x0)
+        return y0 + t * (y1 - y0)
+
+    def rsample(self, sample_shape: torch.Size = None) -> Tensor:
+        sample_shape = sample_shape or torch.Size()
+        u = torch.rand(sample_shape)
+        return self._linear_interp_1d(self.cdf, self.distance_grid, u)
+
+    def log_prob(self, value: Tensor) -> Tensor:
+        log_prob = self._linear_interp_1d(
+            self.distance_grid, self.log_pdf, value
+        )
+        inside_range = (value >= self.minimum) & (value <= self.maximum)
+        log_prob[~inside_range] = float("-inf")
+        return log_prob
+
+
+class RateEvolution(UniformComovingVolume):
+    """
+    Wrapper around `UniformComovingVolume` to allow for
+    arbitrary rate evolution functions. E.g., if
+    `rate_function = 1 / (1 + z)`, then the distribution
+    will sample values such that they occur uniform in
+    source frame time.
+
+    Args:
+        rate_function: Callable that takes redshift as input
+            and returns the rate evolution factor.
+        *args, **kwargs: Arguments passed to `UniformComovingVolume`
+            constructor.
+    """
+
+    def __init__(
+        self,
+        rate_function: Callable,
+        *args,
+        **kwargs,
+    ):
+        self.rate_function = rate_function
+        super().__init__(*args, **kwargs)
+
+    def _p_of_distance(self):
+        """
+        Compute the unnormalized probability as a function of distance
+        """
+        dV_dz = self.comoving_dist_grid**2 / self._hubble_function()
+        # This is a tensor of ones if the distance type is redshift
+        jacobian = torch.gradient(self.distance_grid, spacing=self.dz)[0]
+        return dV_dz / jacobian * self.rate_function(self.z_grid)

{ml4gw-0.7.4.dist-info → ml4gw-0.7.5.dist-info}/METADATA

@@ -1,9 +1,15 @@
 Metadata-Version: 2.4
 Name: ml4gw
-Version: 0.7.4
+Version: 0.7.5
 Summary: Tools for training torch models on gravitational wave data
 Author-email: Ethan Marx <emarx@mit.edu>, Will Benoit <benoi090@umn.edu>, Deep Chatterjee <deep1018@mit.edu>, Alec Gunny <alec.gunny@ligo.org>
 License-File: LICENSE
+Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.13,>=3.9
 Requires-Dist: jaxtyping<0.3,>=0.2
 Requires-Dist: numpy<2.0.0

{ml4gw-0.7.4.dist-info → ml4gw-0.7.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
 ml4gw/__init__.py,sha256=81quoggCuIypZjZs3bbf1Ty70KHdva5RGEJxi0oC57E,25
 ml4gw/augmentations.py,sha256=pZH9tjEpXV0AIqvHHDkpUE-BorG02beOz2pmSipw2EY,1232
 ml4gw/constants.py,sha256=RQPXwavlw_cWu3ByltvTejPsi6EWXHDJQ1HaV9iE3Lg,850
-ml4gw/distributions.py,sha256=6UOgq8W-Bs-9170Jor_0hyeRnmC74zwbUrwAcJEz1jI,5082
+ml4gw/distributions.py,sha256=YbkPqeYBDC91aM59R7-n6NpBMgvMZcZoOAW_U-Jgrdo,12420
 ml4gw/gw.py,sha256=0ovW_HJ3j2b5Yq3mduYtGLSl2RrvFyNNcOsZFf7koHY,19794
 ml4gw/spectral.py,sha256=sao_D0ceeMEatABfiabpqb-xxRfQO8Tz7yk9N7ciOAU,19858
 ml4gw/types.py,sha256=CcctqDcNajR7khGT6BD-WYsfRKpiP0udoSAB0k1qcFw,863
@@ -49,7 +49,7 @@ ml4gw/waveforms/cbc/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9
 ml4gw/waveforms/cbc/phenom_p.py,sha256=RZzzKQzqZW3rQuWZ41htTZOwwulYP61ow87HRRrel5A,27612
 ml4gw/waveforms/cbc/taylorf2.py,sha256=cmYrVL29dwX2Icp7I6SXqRIjtPmoljK5DP_ofx2heiM,10505
 ml4gw/waveforms/cbc/utils.py,sha256=LT1ky10_6ZrbwTcxIrWP1O75GUEuU5q2ZE2yYDhadQE,3037
-ml4gw-0.7.4.dist-info/METADATA,sha256=5nM8sBFDpqrKHQNqtssbhFYFkwze3IE_HLM8Zb8qXQU,3049
-ml4gw-0.7.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ml4gw-0.7.4.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-ml4gw-0.7.4.dist-info/RECORD,,
+ml4gw-0.7.5.dist-info/METADATA,sha256=KI_oVRLxKfCFsBESasNrDh2QfMR4lsriRZwCCAYFt4c,3380
+ml4gw-0.7.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ml4gw-0.7.5.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+ml4gw-0.7.5.dist-info/RECORD,,