Skip to content

torchwave API Reference

This section provides the auto-generated reference documentation for PyTorch-based waveform generation, approximants, detector projections, and inference modules in torchwave.

For physical coverage and execution restrictions, see model support. For numerical qualification, see accuracy and validation.

Core Waveform Functions

torchwave.get_fd_waveform(approximant, mass1=None, mass2=None, chi1=_UNSET, chi2=_UNSET, spin1x=0.0, spin1y=0.0, spin1z=None, spin2x=0.0, spin2y=0.0, spin2z=None, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, lambda1=None, lambda2=None, f_ref=None, phasing_convention=None, reference_frame=None, f_lower=None, f_final=None, delta_f=None, sample_frequencies=None, modes=None, device=None, dtype=torch.float64, frequency_support=None, **kwargs)

Generate frequency-domain gravitational-wave polarizations h_plus(f) and h_cross(f).

Parameters

approximant : str Waveform model name (e.g., 'IMRPhenomD', 'IMRPhenomXHM', 'IMRPhenomXP', 'IMRPhenomXPHM', 'TaylorF2', 'SEOBNRv4_ROM', etc.). mass1, mass2 : float or torch.Tensor Component masses in solar masses (M_sun). Standard aliases (mass_1, m1, etc.) accepted. chi1, chi2 : float or torch.Tensor, optional Aligned dimensionless spins in [-1, 1]. Deprecated in favor of spin1z, spin2z. spin1x, spin1y, spin1z : float or torch.Tensor, optional Dimensionless Cartesian spin components for body 1. spin2x, spin2y, spin2z : float or torch.Tensor, optional Dimensionless Cartesian spin components for body 2. distance : float or torch.Tensor, default 100.0 Luminosity distance in Megaparsecs (Mpc). inclination : float or torch.Tensor, default 0.0 Inclination angle between line of sight and orbital angular momentum in radians. phic : float or torch.Tensor, default 0.0 Coalescence phase in radians. tc : float or torch.Tensor, default 0.0 Time of coalescence in seconds. lambda1, lambda2 : float or torch.Tensor, optional Dimensionless tidal deformabilities Lambda1, Lambda2. f_ref : float or torch.Tensor, optional Reference frequency in Hz. f_lower : float, optional Low-frequency cutoff in Hz. Required if sample_frequencies is not provided. f_final : float, optional High-frequency cutoff in Hz. Defaults to ISCO or Nyquist limit if omitted. delta_f : float, optional Frequency bin width in Hz for uniform grids. sample_frequencies : torch.Tensor, optional Explicit arbitrary frequency grid in Hz. modes : sequence of (l, m) tuples, optional Spherical harmonic modes to include for higher-order approximants. device : str or torch.device, optional Target torch device for waveform tensors. dtype : torch.dtype, default torch.float64 Floating-point dtype for real parameters and grid. **kwargs : dict Approximant-specific keyword arguments (e.g. calibration='sovereign').

Returns

h_plus, h_cross : torch.Tensor Complex Fourier amplitudes of the plus and cross gravitational wave polarizations.

Source code in src/torchwave/waveforms.py 
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
def get_fd_waveform(
    approximant: str,
    mass1: Optional[Union[float, torch.Tensor]] = None,
    mass2: Optional[Union[float, torch.Tensor]] = None,
    chi1: Union[float, torch.Tensor, object] = _UNSET,
    chi2: Union[float, torch.Tensor, object] = _UNSET,
    spin1x: Union[float, torch.Tensor] = 0.0,
    spin1y: Union[float, torch.Tensor] = 0.0,
    spin1z: Optional[Union[float, torch.Tensor]] = None,
    spin2x: Union[float, torch.Tensor] = 0.0,
    spin2y: Union[float, torch.Tensor] = 0.0,
    spin2z: Optional[Union[float, torch.Tensor]] = None,
    distance: Union[float, torch.Tensor] = 100.0,
    inclination: Union[float, torch.Tensor] = 0.0,
    phic: Union[float, torch.Tensor] = 0.0,
    tc: Union[float, torch.Tensor] = 0.0,
    lambda1: Optional[Union[float, torch.Tensor]] = None,
    lambda2: Optional[Union[float, torch.Tensor]] = None,
    f_ref: Optional[Union[float, torch.Tensor]] = None,
    phasing_convention: Optional[str] = None,
    reference_frame: Optional[str] = None,
    f_lower: Optional[float] = None,
    f_final: Optional[float] = None,
    delta_f: Optional[float] = None,
    sample_frequencies: Optional[torch.Tensor] = None,
    modes: Optional[Sequence[Tuple[int, int]]] = None,
    device: Optional[Union[str, torch.device]] = None,
    dtype: torch.dtype = torch.float64,
    frequency_support: Optional[str] = None,
    **kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Generate frequency-domain gravitational-wave polarizations h_plus(f) and h_cross(f).

    Parameters
    ----------
    approximant : str
        Waveform model name (e.g., 'IMRPhenomD', 'IMRPhenomXHM', 'IMRPhenomXP',
        'IMRPhenomXPHM', 'TaylorF2', 'SEOBNRv4_ROM', etc.).
    mass1, mass2 : float or torch.Tensor
        Component masses in solar masses (M_sun). Standard aliases (mass_1, m1, etc.) accepted.
    chi1, chi2 : float or torch.Tensor, optional
        Aligned dimensionless spins in [-1, 1]. Deprecated in favor of spin1z, spin2z.
    spin1x, spin1y, spin1z : float or torch.Tensor, optional
        Dimensionless Cartesian spin components for body 1.
    spin2x, spin2y, spin2z : float or torch.Tensor, optional
        Dimensionless Cartesian spin components for body 2.
    distance : float or torch.Tensor, default 100.0
        Luminosity distance in Megaparsecs (Mpc).
    inclination : float or torch.Tensor, default 0.0
        Inclination angle between line of sight and orbital angular momentum in radians.
    phic : float or torch.Tensor, default 0.0
        Coalescence phase in radians.
    tc : float or torch.Tensor, default 0.0
        Time of coalescence in seconds.
    lambda1, lambda2 : float or torch.Tensor, optional
        Dimensionless tidal deformabilities Lambda1, Lambda2.
    f_ref : float or torch.Tensor, optional
        Reference frequency in Hz.
    f_lower : float, optional
        Low-frequency cutoff in Hz. Required if sample_frequencies is not provided.
    f_final : float, optional
        High-frequency cutoff in Hz. Defaults to ISCO or Nyquist limit if omitted.
    delta_f : float, optional
        Frequency bin width in Hz for uniform grids.
    sample_frequencies : torch.Tensor, optional
        Explicit arbitrary frequency grid in Hz.
    modes : sequence of (l, m) tuples, optional
        Spherical harmonic modes to include for higher-order approximants.
    device : str or torch.device, optional
        Target torch device for waveform tensors.
    dtype : torch.dtype, default torch.float64
        Floating-point dtype for real parameters and grid.
    **kwargs : dict
        Approximant-specific keyword arguments (e.g. calibration='sovereign').

    Returns
    -------
    h_plus, h_cross : torch.Tensor
        Complex Fourier amplitudes of the plus and cross gravitational wave polarizations.
    """
    if kwargs:
        norm_kwargs = normalize_waveform_params(kwargs)
        if mass1 is None and "mass1" in norm_kwargs:
            mass1 = norm_kwargs.pop("mass1")
        if mass2 is None and "mass2" in norm_kwargs:
            mass2 = norm_kwargs.pop("mass2")
        if chi1 is _UNSET and "chi1" in norm_kwargs:
            chi1 = norm_kwargs.pop("chi1")
        if chi2 is _UNSET and "chi2" in norm_kwargs:
            chi2 = norm_kwargs.pop("chi2")
        if spin1z is None and "spin1z" in norm_kwargs:
            spin1z = norm_kwargs.pop("spin1z")
        if spin2z is None and "spin2z" in norm_kwargs:
            spin2z = norm_kwargs.pop("spin2z")
        if spin1x == 0.0 and "spin1x" in norm_kwargs:
            spin1x = norm_kwargs.pop("spin1x")
        if spin1y == 0.0 and "spin1y" in norm_kwargs:
            spin1y = norm_kwargs.pop("spin1y")
        if spin2x == 0.0 and "spin2x" in norm_kwargs:
            spin2x = norm_kwargs.pop("spin2x")
        if spin2y == 0.0 and "spin2y" in norm_kwargs:
            spin2y = norm_kwargs.pop("spin2y")
        if distance == 100.0 and "distance" in norm_kwargs:
            distance = norm_kwargs.pop("distance")
        if inclination == 0.0 and "inclination" in norm_kwargs:
            inclination = norm_kwargs.pop("inclination")
        if phic == 0.0 and "phic" in norm_kwargs:
            phic = norm_kwargs.pop("phic")
        if tc == 0.0 and "tc" in norm_kwargs:
            tc = norm_kwargs.pop("tc")
        if lambda1 is None and "lambda1" in norm_kwargs:
            lambda1 = norm_kwargs.pop("lambda1")
        if lambda2 is None and "lambda2" in norm_kwargs:
            lambda2 = norm_kwargs.pop("lambda2")
        if f_ref is None and "f_ref" in norm_kwargs:
            f_ref = norm_kwargs.pop("f_ref")
        if f_lower is None and "f_lower" in norm_kwargs:
            f_lower = norm_kwargs.pop("f_lower")
        if f_final is None and "f_final" in norm_kwargs:
            f_final = norm_kwargs.pop("f_final")
        if delta_f is None and "delta_f" in norm_kwargs:
            delta_f = norm_kwargs.pop("delta_f")
        if sample_frequencies is None and "sample_frequencies" in norm_kwargs:
            sample_frequencies = norm_kwargs.pop("sample_frequencies")
        kwargs = norm_kwargs

    if mass1 is None or mass2 is None:
        raise TypeError("get_fd_waveform() requires both 'mass1' and 'mass2' (or aliases 'mass_1', 'm1', etc.)")

    out = _get_fd_waveform_raw(
        approximant=approximant,
        mass1=mass1,
        mass2=mass2,
        chi1=chi1,
        chi2=chi2,
        spin1x=spin1x,
        spin1y=spin1y,
        spin1z=spin1z,
        spin2x=spin2x,
        spin2y=spin2y,
        spin2z=spin2z,
        distance=distance,
        inclination=inclination,
        phic=phic,
        tc=tc,
        lambda1=lambda1,
        lambda2=lambda2,
        f_ref=f_ref,
        phasing_convention=phasing_convention,
        reference_frame=reference_frame,
        f_lower=f_lower,
        f_final=f_final,
        delta_f=delta_f,
        sample_frequencies=sample_frequencies,
        frequency_support=frequency_support,
        modes=modes,
        device=device,
        dtype=dtype,
        **kwargs,
    )
    if isinstance(out, tuple) and len(out) >= 2 and torch.is_tensor(out[0]) and torch.is_tensor(out[1]):
        hp, hc = out[0], out[1]
        f = sample_frequencies
        if f is None and delta_f is not None and f_lower is not None and f_final is not None:
            f = torch.arange(f_lower, f_final + delta_f / 2.0, delta_f, dtype=dtype, device=device)
        if f is not None and (f_lower is not None or f_final is not None):
            # Preserve frequency batch axes; ordinary broadcasting prepends
            # any missing parameter axes. Flattening here merges rows/samples.
            f_eval = torch.as_tensor(f, device=hp.device, dtype=hp.real.dtype)
            zero_mask = f_eval <= 0.0
            if f_lower is not None:
                zero_mask = zero_mask | (f_eval < f_lower)
            key = approximant.lower().replace("_", "").replace("-", "")
            pv2_sequence = key == "imrphenompv2" and sample_frequencies is not None and frequency_support != "uniform"
            if f_final is not None and not pv2_sequence:
                zero_mask = zero_mask | (f_eval >= f_final)
            hp_re = torch.where(zero_mask, torch.zeros_like(hp.real), hp.real)
            hp_im = torch.where(zero_mask, torch.zeros_like(hp.imag), hp.imag)
            hc_re = torch.where(zero_mask, torch.zeros_like(hc.real), hc.real)
            hc_im = torch.where(zero_mask, torch.zeros_like(hc.imag), hc.imag)
            hp = torch.view_as_complex(torch.stack([hp_re, hp_im], dim=-1))
            hc = torch.view_as_complex(torch.stack([hc_re, hc_im], dim=-1))
        if len(out) == 2:
            return hp, hc
        return (hp, hc) + out[2:]
    return out

torchwave.get_td_waveform(approximant, mass1=None, mass2=None, chi1=_UNSET, chi2=_UNSET, spin1x=0.0, spin1y=0.0, spin1z=None, spin2x=0.0, spin2y=0.0, spin2z=None, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, lambda1=None, lambda2=None, delta_t=1.0 / 2048.0, f_lower=20.0, f_final=None, duration=None, t_pad=0.2, center_coalescence=True, device=None, dtype=torch.float64, return_times=False, **kwargs)

Generate conditioned time-domain gravitational-wave polarizations h_plus(t) and h_cross(t).

Literature: - Allen et al. (2012), Phys. Rev. D 85, 122006 (arXiv:gr-qc/0509116) - Buonanno et al. (2009), Phys. Rev. D 80, 084043 (arXiv:0907.0700)

The IMRPhenomT family is evaluated natively in time; f_ref fixes its orbital phase. Other models use the inverse Fourier path, which conserves signal energy: int |h(t)|^2 dt = 2 * int |h_tilde(f)|^2 df

Parameters:

Name Type Description Default
approximant str

Waveform approximant name

 required
mass1 Optional[Union[float, Tensor]]

Component mass of primary in solar masses

 None
mass2 Optional[Union[float, Tensor]]

Component mass of secondary in solar masses

 None
chi1 Union[float, Tensor, object]

Dimensionless aligned spin of body 1

 _UNSET
chi2 Union[float, Tensor, object]

Dimensionless aligned spin of body 2

 _UNSET
spin1x Union[float, Tensor]

Primary dimensionless spin x-component

 0.0
spin1y Union[float, Tensor]

Primary dimensionless spin y-component

 0.0
spin1z Optional[Union[float, Tensor]]

Primary dimensionless spin z-component

 None
spin2x Union[float, Tensor]

Secondary dimensionless spin x-component

 0.0
spin2y Union[float, Tensor]

Secondary dimensionless spin y-component

 0.0
spin2z Optional[Union[float, Tensor]]

Secondary dimensionless spin z-component

 None
distance Union[float, Tensor]

Luminosity distance in Mpc

 100.0
inclination Union[float, Tensor]

Orbital inclination in radians

 0.0
phic Union[float, Tensor]

Coalescence phase in radians

 0.0
tc Union[float, Tensor]

Requested coalescence peak time offset in seconds

 0.0
lambda1 Optional[Union[float, Tensor]]

Dimensionless tidal deformability of body 1

 None
lambda2 Optional[Union[float, Tensor]]

Dimensionless tidal deformability of body 2

 None
delta_t float

Sampling interval in seconds (default 1/2048 s)

 1.0 / 2048.0
f_lower float

Low-frequency cutoff in Hz

 20.0
f_final Optional[float]

High-frequency cutoff in Hz (defaults to Nyquist 1/(2*delta_t))

 None
duration Optional[float]

Duration of time series in seconds (automatically estimated if None)

 None
t_pad float

Padding duration in seconds between coalescence peak and end of window

 0.2
center_coalescence bool

If True, center merger peak at requested tc or T - t_pad

 True
device Optional[Union[str, device]]

PyTorch device

 None
dtype dtype

Tensor floating-point precision (torch.float64 recommended)

 float64
return_times bool

If True, also return 1D time coordinate tensor

 False

Returns:

Type Description
Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]]

(h_plus, h_cross) or (h_plus, h_cross, times)
Source code in src/torchwave/waveforms.py 
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
def get_td_waveform(
    approximant: str,
    mass1: Optional[Union[float, torch.Tensor]] = None,
    mass2: Optional[Union[float, torch.Tensor]] = None,
    chi1: Union[float, torch.Tensor, object] = _UNSET,
    chi2: Union[float, torch.Tensor, object] = _UNSET,
    spin1x: Union[float, torch.Tensor] = 0.0,
    spin1y: Union[float, torch.Tensor] = 0.0,
    spin1z: Optional[Union[float, torch.Tensor]] = None,
    spin2x: Union[float, torch.Tensor] = 0.0,
    spin2y: Union[float, torch.Tensor] = 0.0,
    spin2z: Optional[Union[float, torch.Tensor]] = None,
    distance: Union[float, torch.Tensor] = 100.0,
    inclination: Union[float, torch.Tensor] = 0.0,
    phic: Union[float, torch.Tensor] = 0.0,
    tc: Union[float, torch.Tensor] = 0.0,
    lambda1: Optional[Union[float, torch.Tensor]] = None,
    lambda2: Optional[Union[float, torch.Tensor]] = None,
    delta_t: float = 1.0 / 2048.0,
    f_lower: float = 20.0,
    f_final: Optional[float] = None,
    duration: Optional[float] = None,
    t_pad: float = 0.2,
    center_coalescence: bool = True,
    device: Optional[Union[str, torch.device]] = None,
    dtype: torch.dtype = torch.float64,
    return_times: bool = False,
    **kwargs,
) -> Union[
    Tuple[torch.Tensor, torch.Tensor],
    Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
]:
    """Generate conditioned time-domain gravitational-wave polarizations h_plus(t) and h_cross(t).

    Literature:
    - Allen et al. (2012), Phys. Rev. D 85, 122006 (arXiv:gr-qc/0509116)
    - Buonanno et al. (2009), Phys. Rev. D 80, 084043 (arXiv:0907.0700)

    The IMRPhenomT family is evaluated natively in time; f_ref fixes its orbital phase.
    Other models use the inverse Fourier path, which conserves signal energy:
        int |h(t)|^2 dt = 2 * int |h_tilde(f)|^2 df

    Args:
        approximant: Waveform approximant name
        mass1: Component mass of primary in solar masses
        mass2: Component mass of secondary in solar masses
        chi1: Dimensionless aligned spin of body 1
        chi2: Dimensionless aligned spin of body 2
        spin1x: Primary dimensionless spin x-component
        spin1y: Primary dimensionless spin y-component
        spin1z: Primary dimensionless spin z-component
        spin2x: Secondary dimensionless spin x-component
        spin2y: Secondary dimensionless spin y-component
        spin2z: Secondary dimensionless spin z-component
        distance: Luminosity distance in Mpc
        inclination: Orbital inclination in radians
        phic: Coalescence phase in radians
        tc: Requested coalescence peak time offset in seconds
        lambda1: Dimensionless tidal deformability of body 1
        lambda2: Dimensionless tidal deformability of body 2
        delta_t: Sampling interval in seconds (default 1/2048 s)
        f_lower: Low-frequency cutoff in Hz
        f_final: High-frequency cutoff in Hz (defaults to Nyquist 1/(2*delta_t))
        duration: Duration of time series in seconds (automatically estimated if None)
        t_pad: Padding duration in seconds between coalescence peak and end of window
        center_coalescence: If True, center merger peak at requested tc or T - t_pad
        device: PyTorch device
        dtype: Tensor floating-point precision (torch.float64 recommended)
        return_times: If True, also return 1D time coordinate tensor

    Returns:
        (h_plus, h_cross) or (h_plus, h_cross, times)
    """
    if kwargs:
        norm_kwargs = normalize_waveform_params(kwargs)
        if mass1 is None and "mass1" in norm_kwargs:
            mass1 = norm_kwargs.pop("mass1")
        if mass2 is None and "mass2" in norm_kwargs:
            mass2 = norm_kwargs.pop("mass2")
        if chi1 is _UNSET and "chi1" in norm_kwargs:
            chi1 = norm_kwargs.pop("chi1")
        if chi2 is _UNSET and "chi2" in norm_kwargs:
            chi2 = norm_kwargs.pop("chi2")
        if spin1z is None and "spin1z" in norm_kwargs:
            spin1z = norm_kwargs.pop("spin1z")
        if spin2z is None and "spin2z" in norm_kwargs:
            spin2z = norm_kwargs.pop("spin2z")
        if "spin1x" in norm_kwargs and spin1x == 0.0:
            spin1x = norm_kwargs.pop("spin1x")
        if "spin1y" in norm_kwargs and spin1y == 0.0:
            spin1y = norm_kwargs.pop("spin1y")
        if "spin2x" in norm_kwargs and spin2x == 0.0:
            spin2x = norm_kwargs.pop("spin2x")
        if "spin2y" in norm_kwargs and spin2y == 0.0:
            spin2y = norm_kwargs.pop("spin2y")
        if "distance" in norm_kwargs and distance == 100.0:
            distance = norm_kwargs.pop("distance")
        if "inclination" in norm_kwargs and inclination == 0.0:
            inclination = norm_kwargs.pop("inclination")
        if "phic" in norm_kwargs and phic == 0.0:
            phic = norm_kwargs.pop("phic")
        if "tc" in norm_kwargs and tc == 0.0:
            tc = norm_kwargs.pop("tc")
        if lambda1 is None and "lambda1" in norm_kwargs:
            lambda1 = norm_kwargs.pop("lambda1")
        if lambda2 is None and "lambda2" in norm_kwargs:
            lambda2 = norm_kwargs.pop("lambda2")
        if f_lower == 20.0 and "f_lower" in norm_kwargs:
            f_lower = norm_kwargs.pop("f_lower")
        if f_final is None and "f_final" in norm_kwargs:
            f_final = norm_kwargs.pop("f_final")
        if delta_t == 1.0 / 2048.0 and "delta_t" in norm_kwargs:
            delta_t = norm_kwargs.pop("delta_t")
        if duration is None and "duration" in norm_kwargs:
            duration = norm_kwargs.pop("duration")
        kwargs = norm_kwargs

    if mass1 is None or mass2 is None:
        raise TypeError("get_td_waveform() requires both 'mass1' and 'mass2' (or aliases 'mass_1', 'm1', etc.)")
    key = approximant.lower().replace("_", "").replace("-", "")
    if key in ("imrphenomt", "imrphenomthm", "imrphenomtp", "imrphenomtphm"):
        from ._phenomt_td import generate_td

        return generate_td(
            approximant=key, mass1=mass1, mass2=mass2,
            chi1=None if chi1 is _UNSET else chi1,
            chi2=None if chi2 is _UNSET else chi2,
            spin1z=spin1z, spin2z=spin2z, spin1x=spin1x, spin1y=spin1y,
            spin2x=spin2x, spin2y=spin2y, distance=distance,
            inclination=inclination, phic=phic, tc=tc, lambda1=lambda1, lambda2=lambda2,
            delta_t=delta_t, f_lower=f_lower, f_final=f_final, duration=duration,
            t_pad=t_pad, center_coalescence=center_coalescence,
            return_times=return_times, device=device, dtype=dtype, **kwargs,
        )

    m1_t = as_tensor(mass1, device=device, dtype=dtype)
    m2_t = as_tensor(mass2, device=device, dtype=dtype)

    # Estimate required duration if not explicitly provided
    if duration is None:
        # Newtonian 0PN chirp time formula from f_lower to merger
        m_chirp = component_masses_to_chirp_mass(m1_t, m2_t)
        m_chirp_min = torch.min(m_chirp).item()
        m_chirp_sec = m_chirp_min * MSUN_SEC
        tau0 = (5.0 / 256.0) * (m_chirp_sec ** (-5.0 / 3.0)) * ((PI * f_lower) ** (-8.0 / 3.0))
        est_duration = max(2.0, 1.2 * tau0 + 2.0 * t_pad + 1.0)
        # Round up to next power of 2 for FFT efficiency
        n_samples_raw = int(math.ceil(est_duration / delta_t))
        n_samples = 1 << (n_samples_raw - 1).bit_length()
        dur = n_samples * delta_t
    else:
        dur = float(duration)
        n_samples = int(math.ceil(dur / delta_t))

    df = 1.0 / dur
    f_nyquist = 1.0 / (2.0 * delta_t)
    f_max = f_final if f_final is not None else f_nyquist

    # Construct linear frequency evaluation grid [0, df, 2*df, ..., f_nyquist]
    freqs = torch.arange(0, n_samples // 2 + 1, dtype=dtype, device=device) * df

    # Mask frequencies above 0 and below f_max
    valid_mask = (freqs >= f_lower) & (freqs <= f_max)
    active_freqs = freqs[valid_mask]

    # Generate frequency-domain waveform
    hp_fd_active, hc_fd_active = get_fd_waveform(
        approximant=approximant,
        mass1=mass1,
        mass2=mass2,
        chi1=chi1,
        chi2=chi2,
        spin1x=spin1x,
        spin1y=spin1y,
        spin1z=spin1z,
        spin2x=spin2x,
        spin2y=spin2y,
        spin2z=spin2z,
        distance=distance,
        inclination=inclination,
        phic=phic,
        tc=0.0,  # Generate with tc=0; time centering handled cleanly in time domain
        lambda1=lambda1,
        lambda2=lambda2,
        sample_frequencies=active_freqs,
        device=device,
        dtype=dtype,
        **kwargs,
    )

    # Full Fourier grid container
    is_batched = m1_t.dim() >= 1
    complex_dtype = torch.complex128 if dtype == torch.float64 else torch.complex64

    if is_batched:
        batch_shape = m1_t.shape
        full_shape = batch_shape + (len(freqs),)
        hp_fd = torch.zeros(full_shape, dtype=complex_dtype, device=device)
        hc_fd = torch.zeros(full_shape, dtype=complex_dtype, device=device)
        hp_fd[..., valid_mask] = hp_fd_active
        hc_fd[..., valid_mask] = hc_fd_active
    else:
        hp_fd = torch.zeros(len(freqs), dtype=complex_dtype, device=device)
        hc_fd = torch.zeros(len(freqs), dtype=complex_dtype, device=device)
        hp_fd[valid_mask] = hp_fd_active
        hc_fd[valid_mask] = hc_fd_active

    # Low-frequency smooth Planck taper at f_lower over [f_lower, f_lower + df_taper]
    df_taper = max(2.0, 0.1 * f_lower)
    f_start = f_lower
    f_end = f_lower + df_taper
    z = torch.clamp((freqs - f_start) / (f_end - f_start), 1e-7, 1.0 - 1e-7)
    sigma = (1.0 / z) + (1.0 / (z - 1.0))
    w_smooth = 1.0 / (torch.exp(torch.clamp(sigma, -50.0, 50.0)) + 1.0)
    w_taper = torch.where(freqs <= f_start, torch.zeros_like(w_smooth), w_smooth)
    w_taper = torch.where(freqs >= f_end, torch.ones_like(w_taper), w_taper)

    if is_batched:
        w_taper = w_taper.unsqueeze(0)

    hp_fd = hp_fd * w_taper
    hc_fd = hc_fd * w_taper

    # Perform batched inverse FFT with exact continuous scaling:
    # h(t) = (1 / delta_t) * irfft(h_tilde) = df * N * irfft(h_tilde)
    # Satisfies exact Parseval: int |h(t)|^2 dt = 2 * int |h_tilde(f)|^2 df
    scale_factor = 1.0 / delta_t
    hp_td = scale_factor * torch.fft.irfft(hp_fd, n=n_samples, dim=-1)
    hc_td = scale_factor * torch.fft.irfft(hc_fd, n=n_samples, dim=-1)

    # Time centering: roll buffer so coalescence peak occurs at requested location
    if center_coalescence:
        tc_val = float(tc) if isinstance(tc, (int, float)) else float(tc.reshape(-1)[0].item())
        target_t = (dur - t_pad) + tc_val
        shift_samples = int(round(target_t / delta_t)) % n_samples
        hp_td = torch.roll(hp_td, shifts=shift_samples, dims=-1)
        hc_td = torch.roll(hc_td, shifts=shift_samples, dims=-1)

    times = torch.arange(n_samples, dtype=dtype, device=device) * delta_t
    if center_coalescence:
        # If centered at tc, align time array such that coalescence is at tc
        times = times - (dur - t_pad)

    if return_times:
        return hp_td, hc_td, times
    return hp_td, hc_td

torchwave.WaveformGenerator

Class-based interface for gravitational waveform generation.

Configured with static waveform parameters and grid specification, allowing repeated waveform evaluation across varying binary parameters.

Source code in src/torchwave/waveforms.py 
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
class WaveformGenerator:
    """Class-based interface for gravitational waveform generation.

    Configured with static waveform parameters and grid specification,
    allowing repeated waveform evaluation across varying binary parameters.
    """

    def __init__(
        self,
        approximant: str = "IMRPhenomD",
        f_lower: Optional[float] = None,
        f_final: Optional[float] = None,
        delta_f: Optional[float] = None,
        sample_frequencies: Optional[Any] = None,
        modes: Optional[Sequence[Tuple[int, int]]] = None,
        device: Optional[Union[str, torch.device]] = None,
        dtype: torch.dtype = torch.float64,
        **default_kwargs,
    ):
        norm_init = normalize_waveform_params(default_kwargs)
        self.approximant = norm_init.pop("approximant", approximant)
        self.f_lower = norm_init.pop("f_lower", f_lower)
        self.f_final = norm_init.pop("f_final", f_final)
        self.delta_f = norm_init.pop("delta_f", delta_f)
        sf = norm_init.pop("sample_frequencies", sample_frequencies)
        self.sample_frequencies = sf
        self.modes = norm_init.pop("modes", modes)
        self.device = norm_init.pop("device", device)
        self.dtype = norm_init.pop("dtype", dtype)
        self.default_kwargs = norm_init
        if sf is not None:
            self._cached_frequencies = torch.as_tensor(sf, dtype=self.dtype, device=self.device)
        else:
            self._cached_frequencies = None

    def generate_fd_waveform(
        self,
        parameters: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Generate frequency-domain polarizations for given parameters."""
        params = dict(self.default_kwargs)
        if parameters is not None:
            params.update(parameters)
        params.update(kwargs)
        normalized_params = normalize_waveform_params(params)

        sf = normalized_params.pop(
            "sample_frequencies",
            self._cached_frequencies if self._cached_frequencies is not None else self.sample_frequencies,
        )

        call_kwargs = {
            "approximant": normalized_params.pop("approximant", self.approximant),
            "f_lower": normalized_params.pop("f_lower", self.f_lower),
            "f_final": normalized_params.pop("f_final", self.f_final),
            "delta_f": normalized_params.pop("delta_f", self.delta_f),
            "sample_frequencies": sf,
            "modes": normalized_params.pop("modes", self.modes),
            "device": normalized_params.pop("device", self.device),
            "dtype": normalized_params.pop("dtype", self.dtype),
        }
        call_kwargs.update(normalized_params)

        return get_fd_waveform(**call_kwargs)

    get_fd_waveform = generate_fd_waveform

    def __call__(self, *args, **kwargs):
        if args and isinstance(args[0], dict):
            return self.generate_fd_waveform(parameters=args[0], **kwargs)
        return self.generate_fd_waveform(**kwargs)

generate_fd_waveform(parameters=None, **kwargs)

Generate frequency-domain polarizations for given parameters.

Source code in src/torchwave/waveforms.py 
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
def generate_fd_waveform(
    self,
    parameters: Optional[Dict[str, Any]] = None,
    **kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Generate frequency-domain polarizations for given parameters."""
    params = dict(self.default_kwargs)
    if parameters is not None:
        params.update(parameters)
    params.update(kwargs)
    normalized_params = normalize_waveform_params(params)

    sf = normalized_params.pop(
        "sample_frequencies",
        self._cached_frequencies if self._cached_frequencies is not None else self.sample_frequencies,
    )

    call_kwargs = {
        "approximant": normalized_params.pop("approximant", self.approximant),
        "f_lower": normalized_params.pop("f_lower", self.f_lower),
        "f_final": normalized_params.pop("f_final", self.f_final),
        "delta_f": normalized_params.pop("delta_f", self.delta_f),
        "sample_frequencies": sf,
        "modes": normalized_params.pop("modes", self.modes),
        "device": normalized_params.pop("device", self.device),
        "dtype": normalized_params.pop("dtype", self.dtype),
    }
    call_kwargs.update(normalized_params)

    return get_fd_waveform(**call_kwargs)

Native Time-Domain Models

See the IMRPhenomT guide for examples and conventions.

torchwave.IMRPhenomT

Bases: Module

Aligned-spin dominant (2,2) IMRPhenomT waveform in Torch.

Source code in src/torchwave/approximants/imrphenomt.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
class IMRPhenomT(nn.Module):
    """Aligned-spin dominant ``(2,2)`` IMRPhenomT waveform in Torch."""

    def __init__(self):
        super().__init__()

    @staticmethod
    def _state(m1, m2, chi1, chi2):
        M = m1 + m2
        eta = m1 * m2 / M.square()
        delta = (m1 - m2) / M
        # The fit variable is sTotR, normalized by m1^2+m2^2=1-2 eta.
        S = (chi1 * (m1 / M).square() + chi2 * (m2 / M).square()) / (1 - 2 * eta)
        dchi = chi1 - chi2
        # The helper provides the published 2017 remnant and Kerr fits.
        af = _fit_any(("IMRPhenomX_FinalSpin2017", "final_spin"), eta, chi1, chi2, delta)
        mf = _fit_any(("IMRPhenomX_FinalMass2017", "final_mass"), eta, chi1, chi2, delta)
        af, mf = _t(af, eta), _t(mf, eta)
        return M, eta, delta, S, dchi, af, mf

    @staticmethod
    def _coeff(m1, m2, chi1, chi2, final_spin=None):
        """Reconstruct the carrier, optionally using a precessing remnant spin."""
        M, eta, delta, S, dchi, af, mf = IMRPhenomT._state(m1, m2, chi1, chi2)
        t0 = _t(_fit_any(("IMRPhenomT_Inspiral_TaylorT3_t0",), eta, S, dchi, delta), eta)
        freq_cp = [_t(_fit_any(("IMRPhenomT_Inspiral_Freq_CP",), eta, S, dchi, delta, i), eta) for i in range(1, 6)]
        peak_amp = _t(_fit_any(("IMRPhenomT_PeakAmp",), eta, S, dchi, delta, 22), eta)
        amp_mid = _t(_fit_any(("IMRPhenomT_Intermediate_Amp_CP1",), eta, S, dchi, delta, 22), eta)
        tshift = torch.zeros_like(eta)
        c3a = _t(_fit_any(("IMRPhenomT_Ringdown_Amp_C3",), eta, S, dchi, 22), eta)
        omega_peak = _t(_fit_any(("IMRPhenomT_PeakFrequency",), eta, S, dchi, delta, 22), eta)
        c2f = _t(_fit_any(("IMRPhenomT_RD_Freq_C2",), eta, S, dchi, delta, 22), eta)
        c3f = _t(_fit_any(("IMRPhenomT_RD_Freq_C3",), eta, S, dchi, delta, 22), eta)
        mid_freq = _t(_fit_any(("IMRPhenomT_Intermediate_Freq_CP1",), eta, S, dchi, delta, 22), eta)
        # The QNM fits are in units of final mass; the ansatz time is in
        # units of total mass, so only division by mf is required here.
        fr = _t(_fit_any(("IMRPhenomT_fringfit",), af, 22), eta) / mf
        fd = _t(_fit_any(("IMRPhenomT_fdampfit",), af, 22), eta) / mf
        fd2 = _t(_fit_any(("IMRPhenomT_fdampn2fit",), af, 22), eta) / mf
        # The NNLO precessing remnant changes only ringdown (TPHM Sec. II C).
        # Preserve the calibrated aligned-spin inspiral and merger joins.
        rd_fr, rd_fd, rd_fd2 = fr, fd, fd2
        if final_spin is not None:
            rd_af = torch.as_tensor(final_spin, dtype=eta.dtype, device=eta.device) + torch.zeros_like(eta)
            rd_fr = _t(_fit_any(("IMRPhenomT_fringfit",), rd_af, 22), eta) / mf
            rd_fd = _t(_fit_any(("IMRPhenomT_fdampfit",), rd_af, 22), eta) / mf
            rd_fd2 = _t(_fit_any(("IMRPhenomT_fdampn2fit",), rd_af, 22), eta) / mf

        th = torch.as_tensor((0.33, 0.45, 0.55, 0.65, 0.75, 0.82), dtype=eta.dtype, device=eta.device)
        qmat = torch.stack([th ** (8 + j) for j in range(6)], -1)
        tearly = -5.0 / (eta * th[0] ** 8)
        theta0 = torch.pow(torch.clamp(eta * (t0 - tearly) / 5.0, min=1.0e-30), -1 / 8)
        vals = torch.stack([_poly_pn_frequency(theta0, eta, chi1, chi2, delta)] + freq_cp, -1)
        pvals = _poly_pn_frequency(th, eta[..., None], chi1[..., None], chi2[..., None], delta[..., None])
        rhs = 4.0 / th**3 * (vals - pvals)
        wc = torch.linalg.solve(qmat.expand(rhs.shape[:-1] + qmat.shape), rhs)

        # Merger frequency coefficients, Eqs. (28-31) of the reference.
        tcut = -26.982976386771437 / eta
        tcp = -5.0 / (eta * torch.as_tensor(0.95, dtype=eta.dtype, device=eta.device) ** 8)
        wcut = IMRPhenomT._omega_insp(tcut[..., None], eta, delta, chi1, chi2, wc)[..., 0]
        dwcut = _domega_insp(tcut[..., None], eta, chi1, chi2, delta, wc)[..., 0]
        wr = 2 * PI * fr
        alpha = 2 * PI * fd
        c1f = (1 + c3f) * (wr - omega_peak) / (c2f * c3f)
        dwpeak = -c1f * c2f.square() * c3f / (1 + c3f).square() / wr
        acut = torch.asinh(alpha * tcut)
        amer = torch.asinh(alpha * tcp)
        dencut = torch.sqrt(1 + (alpha * tcut).square())
        b0 = 1 - (wcut / wr) - (1 - omega_peak / wr) - dwpeak / alpha * acut
        b1 = 1 - mid_freq / wr - (1 - omega_peak / wr) - dwpeak / alpha * amer
        b2 = -dwcut / wr - dwpeak / dencut
        amat = torch.stack(
            (
                torch.stack((acut.square(), acut**3, acut**4), -1),
                torch.stack((amer.square(), amer**3, amer**4), -1),
                torch.stack(
                    (2 * alpha * acut / dencut, 3 * alpha * acut.square() / dencut, 4 * alpha * acut**3 / dencut), -1
                ),
            ),
            -2,
        )
        mc = torch.linalg.solve(amat, torch.stack((b0, b1, b2), -1))

        # Amplitude reconstruction is kept in the MIT-derived private helper;
        # it consumes analytic piecewise omega and domega callbacks.
        from ._phenomt_amplitude import amplitude_coefficients

        def omega_cb(tt):
            ti = torch.minimum(tt, tcut[..., None] - 1.0e-10)
            wi = IMRPhenomT._omega_insp(ti, eta, delta, chi1, chi2, wc)
            tm = torch.minimum(torch.maximum(tt, tcut[..., None]), torch.zeros_like(tt))
            x = torch.asinh(alpha[..., None] * tm)
            wm = wr[..., None] * (
                1
                - (
                    1
                    - omega_peak[..., None] / wr[..., None]
                    + x
                    * (
                        dwpeak[..., None] / alpha[..., None]
                        + x * (mc[..., 0, None] + x * (mc[..., 1, None] + x * mc[..., 2, None]))
                    )
                )
            )
            tr = torch.maximum(tt, torch.zeros_like(tt))
            er = torch.exp(-c2f[..., None] * tr)
            wd = wr[..., None] - c1f[..., None] * c2f[..., None] * c3f[..., None] * er / (1 + c3f[..., None] * er)
            return torch.where(tt < tcut[..., None], wi, torch.where(tt < 0, wm, wd))

        def domega_cb(tt):
            ti = torch.minimum(tt, tcut[..., None] - 1.0e-10)
            di = _domega_insp(ti, eta, chi1, chi2, delta, wc)
            tm = torch.minimum(torch.maximum(tt, tcut[..., None]), torch.zeros_like(tt))
            x = torch.asinh(alpha[..., None] * tm)
            dx = alpha[..., None] / torch.sqrt(1 + (alpha[..., None] * tm).square())
            dqx = (
                dwpeak[..., None] / alpha[..., None]
                + 2 * mc[..., 0, None] * x
                + 3 * mc[..., 1, None] * x.square()
                + 4 * mc[..., 2, None] * x**3
            )
            dm = -wr[..., None] * dqx * dx
            tr = torch.maximum(tt, torch.zeros_like(tt))
            er = torch.exp(-c2f[..., None] * tr)
            dd = c1f[..., None] * c2f[..., None].square() * c3f[..., None] * er / (1 + c3f[..., None] * er).square()
            return torch.where(tt < tcut[..., None], di, torch.where(tt < 0, dm, dd))

        amp_state = amplitude_coefficients(eta, delta, chi1, chi2, fd, fd2, omega_cb, domega_cb, rd_fd, rd_fd2)

        # Analytic phase offsets make all three regions C0 continuous.
        p5 = torch.as_tensor(5.0, dtype=eta.dtype, device=eta.device) ** (
            torch.arange(8, dtype=eta.dtype, device=eta.device) / 8
        )
        theta_i = torch.pow(torch.clamp(eta * (t0 - tearly), min=1.0e-30), -1 / 8)
        phoff_i = (
            _pn_phase(theta_i, eta, delta, chi1, chi2)
            - _phase_insp_raw(tearly[..., None], eta, chi1, chi2, delta, wc, p5)[..., 0]
        )
        phoff_m = (
            phoff_i
            + _phase_insp_raw(tcut[..., None], eta, chi1, chi2, delta, wc, p5)[..., 0]
            - IMRPhenomT._phase_inter_raw(tcut, alpha, wr, omega_peak, dwpeak, mc)
        )
        phoff_r = IMRPhenomT._phase_inter_raw(torch.zeros_like(eta), alpha, wr, omega_peak, dwpeak, mc) + phoff_m
        return (
            M,
            eta,
            delta,
            S,
            dchi,
            chi1,
            chi2,
            wc,
            amp_state,
            peak_amp,
            amp_mid,
            tshift,
            c3a,
            omega_peak,
            c2f,
            c3f,
            mid_freq,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwpeak,
            mc,
            tcut,
            phoff_i,
            phoff_m,
            phoff_r,
            p5,
            2 * PI * rd_fr,
            (1 + c3f) * (2 * PI * rd_fr - omega_peak) / (c2f * c3f),
        )

    @staticmethod
    def _omega_insp(t, eta, delta, s1, s2, wc):
        if t.ndim <= eta.ndim:
            t = t.reshape((1,) * eta.ndim + (-1,))
        th = torch.pow(torch.clamp(-eta[..., None] * t / 5, min=1e-30), -1 / 8)
        return _poly_pn_frequency(th, eta[..., None], s1[..., None], s2[..., None], delta[..., None]) + th**3 / 4 * sum(
            wc[..., j, None] * th ** (8 + j) for j in range(6)
        )

    @staticmethod
    def _phase_inter_raw(t, alpha, wr, op, dwp, mc):
        x = torch.asinh(alpha * t)
        x2, x3, x4 = x * x, x * x * x, x * x * x * x
        root = torch.sqrt(1 + (alpha * t).square())
        return wr * t * (
            1
            - (
                2 * mc[..., 0]
                + 24 * mc[..., 2]
                + (6 * mc[..., 1] + dwp / alpha) * x
                + (1 - op / wr)
                + (mc[..., 0] + 12 * mc[..., 2]) * x2
                + mc[..., 1] * x3
                + mc[..., 2] * x4
            )
        ) - wr / alpha * root * (
            -dwp / alpha
            - 6 * mc[..., 1]
            - x * (2 * mc[..., 0] + 24 * mc[..., 2])
            - 3 * mc[..., 1] * x2
            - 4 * mc[..., 2] * x3
        )

    @staticmethod
    def _evaluate(u, state):
        (
            M,
            eta,
            delta,
            S,
            dchi,
            s1,
            s2,
            wc,
            amp_state,
            peak,
            amid,
            tshift,
            c3a,
            op,
            c2f,
            c3f,
            midf,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwp,
            mc,
            tcut,
            phoi,
            phom,
            phor,
            p5,
            rd_wr,
            rd_c1f,
        ) = state
        ui = torch.minimum(u, tcut[..., None])
        um = torch.minimum(torch.maximum(u, tcut[..., None]), torch.zeros_like(u))
        ur = torch.maximum(u, torch.zeros_like(u))
        ex = torch.exp(-c2f[..., None] * ur)
        win = IMRPhenomT._omega_insp(ui, eta, delta, s1, s2, wc)
        # Peak-matched c4=0 branch; printed Eq. (25) conflicts with Eq. (27a).
        wring = rd_wr[..., None] - rd_c1f[..., None] * c2f[..., None] * c3f[..., None] * ex / (1 + c3f[..., None] * ex)
        x = torch.asinh(alpha[..., None] * um)
        wm = wr[..., None] * (
            1
            - (
                1
                - op[..., None] / wr[..., None]
                + x
                * (
                    dwp[..., None] / alpha[..., None]
                    + x * (mc[..., 0, None] + x * (mc[..., 1, None] + x * mc[..., 2, None]))
                )
            )
        )
        omega = torch.where(u < tcut[..., None], win, torch.where(u < 0, wm, wring))
        pi = _phase_insp_raw(ui, eta, s1, s2, delta, wc, p5) + phoi[..., None]
        pm = (
            IMRPhenomT._phase_inter_raw(
                um, alpha[..., None], wr[..., None], op[..., None], dwp[..., None], mc[..., None, :]
            )
            + phom[..., None]
        )
        er = torch.exp(-c2f[..., None] * ur)
        pr = (
            rd_c1f[..., None] * torch.log((1 + c3f[..., None] * er) / (1 + c3f[..., None]))
            + rd_wr[..., None] * ur
            + phor[..., None]
        )
        phase = torch.where(u < tcut[..., None], pi, torch.where(u < 0, pm, pr))
        return omega, phase, amp_state

    @staticmethod
    def _domega_full(u, state):
        """Analytic derivative of the piecewise IMR angular frequency."""
        (
            M,
            eta,
            delta,
            S,
            dchi,
            s1,
            s2,
            wc,
            amp_state,
            peak,
            amid,
            tshift,
            c3a,
            op,
            c2f,
            c3f,
            midf,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwp,
            mc,
            tcut,
            phoi,
            phom,
            phor,
            p5,
            rd_wr,
            rd_c1f,
        ) = state
        ui = torch.minimum(u, tcut[..., None])
        um = torch.minimum(torch.maximum(u, tcut[..., None]), torch.zeros_like(u))
        ur = torch.maximum(u, torch.zeros_like(u))
        di = _domega_insp(ui, eta, s1, s2, delta, wc)
        x = torch.asinh(alpha[..., None] * um)
        dx = alpha[..., None] / torch.sqrt(1 + (alpha[..., None] * um).square())
        dqx = (
            dwp[..., None] / alpha[..., None]
            + 2 * mc[..., 0, None] * x
            + 3 * mc[..., 1, None] * x.square()
            + 4 * mc[..., 2, None] * x**3
        )
        dm = -wr[..., None] * dqx * dx
        er = torch.exp(-c2f[..., None] * ur)
        dd = rd_c1f[..., None] * c2f[..., None].square() * c3f[..., None] * er / (1 + c3f[..., None] * er).square()
        return torch.where(u < tcut[..., None], di, torch.where(u < 0, dm, dd))

    def time_domain(
        self,
        times: TensorLike,
        mass1: TensorLike,
        mass2: TensorLike,
        chi1: TensorLike = 0.0,
        chi2: TensorLike = 0.0,
        distance: TensorLike = 100.0,
        inclination: TensorLike = 0.0,
        phic: TensorLike = 0.0,
        tc: TensorLike = 0.0,
        f_ref: TensorLike = 20.0,
        *,
        _final_spin_override=None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return h_plus, h_cross on times in seconds relative to the peak.

        Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter
        batches precede the shared time axis. This method applies no start mask.
        """
        t = torch.as_tensor(
            times,
            dtype=times.dtype if isinstance(times, torch.Tensor) and times.is_floating_point() else torch.float64,
            device=times.device if isinstance(times, torch.Tensor) else None,
        )
        m1, m2, c1, c2, d, inc, pc, tcenter, fr = [
            torch.as_tensor(v, dtype=t.dtype, device=t.device)
            for v in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref)
        ]
        m1, m2, c1, c2, d, inc, pc, tcenter, fr = torch.broadcast_tensors(m1, m2, c1, c2, d, inc, pc, tcenter, fr)
        state = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
        M = state[0]
        u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
        omega, phase, amp_state = self._evaluate(u, state)
        from ._phenomt_amplitude import amplitude

        amp = amplitude(u, amp_state, lambda tt: self._evaluate(tt, state)[0])
        # f_ref=0 follows the public convention and means the default 20 Hz
        # lower/reference frequency.  Subtract the analytic phase at that
        # frequency, as the author implementation does.
        fphase = torch.where(fr > 0, fr, torch.full_like(fr, 20.0))
        tref = self.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=_final_spin_override)
        uref = tref / (M * MSUN_SEC)
        phref = self._evaluate(uref[..., None], state)[1][..., 0]
        ph = phase - phref[..., None] + 2.0 * pc[..., None]
        scale = (M[..., None] * MSUN_SEC) / (d[..., None] * MPC_SI / C_SI)
        y22 = math.sqrt(5.0 / (4.0 * math.pi))
        ci = torch.cos(inc[..., None])
        pref = -scale * y22 * amp
        return pref * (1 + ci.square()) * 0.5 * torch.cos(ph), pref * ci * torch.sin(ph)

    forward = time_domain

    def frequency(self, times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
        """Instantaneous gravitational-wave frequency in Hz at peak-relative times."""
        t = torch.as_tensor(times, dtype=torch.float64)
        m1 = torch.as_tensor(mass1, dtype=t.dtype, device=t.device)
        m2 = torch.as_tensor(mass2, dtype=t.dtype, device=t.device)
        c1 = torch.as_tensor(chi1, dtype=t.dtype, device=t.device)
        c2 = torch.as_tensor(chi2, dtype=t.dtype, device=t.device)
        m1, m2, c1, c2 = torch.broadcast_tensors(m1, m2, c1, c2)
        M = m1 + m2
        u = t.reshape((1,) * M.ndim + (-1,)) / (M[..., None] * MSUN_SEC)
        return self._evaluate(u, self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override))[0] / (2 * PI * M[..., None] * MSUN_SEC)

    def time_at_frequency(self, frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
        """Invert the frequency evolution, preserving its implicit first derivative."""
        device = next((v.device for v in (frequency, mass1, mass2, chi1, chi2, _final_spin_override) if isinstance(v, torch.Tensor)), None)
        f = torch.as_tensor(frequency, dtype=torch.float64, device=device)
        m1 = torch.as_tensor(mass1, dtype=f.dtype, device=device)
        m2 = torch.as_tensor(mass2, dtype=f.dtype, device=device)
        c1 = torch.as_tensor(chi1, dtype=f.dtype, device=device)
        c2 = torch.as_tensor(chi2, dtype=f.dtype, device=device)
        m1, m2, c1, c2, f = torch.broadcast_tensors(m1, m2, c1, c2, f)
        M = m1 + m2
        target = 2 * PI * f * M * MSUN_SEC
        st = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
        if not bool(torch.all(torch.isfinite(f) & (f > 0))):
            raise ValueError("frequency must be finite and positive")
        # The ringdown ansatz approaches omega_R from below and has no finite
        # time solution at or above that limit.
        if bool(torch.any(target >= st[30]).item()):
            raise ValueError("frequency is above the supported IMRPhenomT range")
        # Newtonian leading-order time, buffered as in the author solver;
        # unlike a fixed -1e6 bracket this remains useful for BNS masses.
        lo = -5.0 / (256.0 * st[1] * torch.pow(target / 2.0, 8.0 / 3.0)) * 1.2
        hi = torch.zeros_like(target)
        # For frequencies above the peak, grow the upper bracket into
        # ringdown until its angular frequency exceeds the target.
        with torch.no_grad():
            for _ in range(16):
                omhi = self._evaluate(hi[..., None], st)[0][..., 0]
                need = omhi < target
                hi = torch.where(need, torch.where(hi == 0, torch.ones_like(hi), 2 * hi), hi)
            for _ in range(12):
                omlo = self._evaluate(lo[..., None], st)[0][..., 0]
                need = omlo > target
                lo = torch.where(need, lo * 2.0, lo)
            for _ in range(80):
                mid = (lo + hi) / 2
                om = self._evaluate(mid[..., None], st)[0][..., 0]
                lo = torch.where(om < target, mid, lo)
                hi = torch.where(om < target, hi, mid)
        # Bisection selects the branch robustly but is piecewise constant as a
        # function of parameters.  One implicit Newton correction restores the
        # correct d t / d parameter without differentiating through decisions.
        root = ((lo + hi) / 2).detach()
        om = self._evaluate(root[..., None], st)[0][..., 0]
        dom = self._domega_full(root[..., None], st)[..., 0]
        root = root + (target - om) / torch.clamp(dom, min=torch.finfo(root.dtype).tiny)
        return root * M * MSUN_SEC

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, *, _final_spin_override=None)

Return h_plus, h_cross on times in seconds relative to the peak.

Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter batches precede the shared time axis. This method applies no start mask.

Source code in src/torchwave/approximants/imrphenomt.py 
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
def time_domain(
    self,
    times: TensorLike,
    mass1: TensorLike,
    mass2: TensorLike,
    chi1: TensorLike = 0.0,
    chi2: TensorLike = 0.0,
    distance: TensorLike = 100.0,
    inclination: TensorLike = 0.0,
    phic: TensorLike = 0.0,
    tc: TensorLike = 0.0,
    f_ref: TensorLike = 20.0,
    *,
    _final_spin_override=None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return h_plus, h_cross on times in seconds relative to the peak.

    Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter
    batches precede the shared time axis. This method applies no start mask.
    """
    t = torch.as_tensor(
        times,
        dtype=times.dtype if isinstance(times, torch.Tensor) and times.is_floating_point() else torch.float64,
        device=times.device if isinstance(times, torch.Tensor) else None,
    )
    m1, m2, c1, c2, d, inc, pc, tcenter, fr = [
        torch.as_tensor(v, dtype=t.dtype, device=t.device)
        for v in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref)
    ]
    m1, m2, c1, c2, d, inc, pc, tcenter, fr = torch.broadcast_tensors(m1, m2, c1, c2, d, inc, pc, tcenter, fr)
    state = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
    M = state[0]
    u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
    omega, phase, amp_state = self._evaluate(u, state)
    from ._phenomt_amplitude import amplitude

    amp = amplitude(u, amp_state, lambda tt: self._evaluate(tt, state)[0])
    # f_ref=0 follows the public convention and means the default 20 Hz
    # lower/reference frequency.  Subtract the analytic phase at that
    # frequency, as the author implementation does.
    fphase = torch.where(fr > 0, fr, torch.full_like(fr, 20.0))
    tref = self.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=_final_spin_override)
    uref = tref / (M * MSUN_SEC)
    phref = self._evaluate(uref[..., None], state)[1][..., 0]
    ph = phase - phref[..., None] + 2.0 * pc[..., None]
    scale = (M[..., None] * MSUN_SEC) / (d[..., None] * MPC_SI / C_SI)
    y22 = math.sqrt(5.0 / (4.0 * math.pi))
    ci = torch.cos(inc[..., None])
    pref = -scale * y22 * amp
    return pref * (1 + ci.square()) * 0.5 * torch.cos(ph), pref * ci * torch.sin(ph)

frequency(times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None)

Instantaneous gravitational-wave frequency in Hz at peak-relative times.

Source code in src/torchwave/approximants/imrphenomt.py 
608
609
610
611
612
613
614
615
616
617
618
def frequency(self, times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
    """Instantaneous gravitational-wave frequency in Hz at peak-relative times."""
    t = torch.as_tensor(times, dtype=torch.float64)
    m1 = torch.as_tensor(mass1, dtype=t.dtype, device=t.device)
    m2 = torch.as_tensor(mass2, dtype=t.dtype, device=t.device)
    c1 = torch.as_tensor(chi1, dtype=t.dtype, device=t.device)
    c2 = torch.as_tensor(chi2, dtype=t.dtype, device=t.device)
    m1, m2, c1, c2 = torch.broadcast_tensors(m1, m2, c1, c2)
    M = m1 + m2
    u = t.reshape((1,) * M.ndim + (-1,)) / (M[..., None] * MSUN_SEC)
    return self._evaluate(u, self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override))[0] / (2 * PI * M[..., None] * MSUN_SEC)

time_at_frequency(frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None)

Invert the frequency evolution, preserving its implicit first derivative.

Source code in src/torchwave/approximants/imrphenomt.py 
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
def time_at_frequency(self, frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
    """Invert the frequency evolution, preserving its implicit first derivative."""
    device = next((v.device for v in (frequency, mass1, mass2, chi1, chi2, _final_spin_override) if isinstance(v, torch.Tensor)), None)
    f = torch.as_tensor(frequency, dtype=torch.float64, device=device)
    m1 = torch.as_tensor(mass1, dtype=f.dtype, device=device)
    m2 = torch.as_tensor(mass2, dtype=f.dtype, device=device)
    c1 = torch.as_tensor(chi1, dtype=f.dtype, device=device)
    c2 = torch.as_tensor(chi2, dtype=f.dtype, device=device)
    m1, m2, c1, c2, f = torch.broadcast_tensors(m1, m2, c1, c2, f)
    M = m1 + m2
    target = 2 * PI * f * M * MSUN_SEC
    st = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
    if not bool(torch.all(torch.isfinite(f) & (f > 0))):
        raise ValueError("frequency must be finite and positive")
    # The ringdown ansatz approaches omega_R from below and has no finite
    # time solution at or above that limit.
    if bool(torch.any(target >= st[30]).item()):
        raise ValueError("frequency is above the supported IMRPhenomT range")
    # Newtonian leading-order time, buffered as in the author solver;
    # unlike a fixed -1e6 bracket this remains useful for BNS masses.
    lo = -5.0 / (256.0 * st[1] * torch.pow(target / 2.0, 8.0 / 3.0)) * 1.2
    hi = torch.zeros_like(target)
    # For frequencies above the peak, grow the upper bracket into
    # ringdown until its angular frequency exceeds the target.
    with torch.no_grad():
        for _ in range(16):
            omhi = self._evaluate(hi[..., None], st)[0][..., 0]
            need = omhi < target
            hi = torch.where(need, torch.where(hi == 0, torch.ones_like(hi), 2 * hi), hi)
        for _ in range(12):
            omlo = self._evaluate(lo[..., None], st)[0][..., 0]
            need = omlo > target
            lo = torch.where(need, lo * 2.0, lo)
        for _ in range(80):
            mid = (lo + hi) / 2
            om = self._evaluate(mid[..., None], st)[0][..., 0]
            lo = torch.where(om < target, mid, lo)
            hi = torch.where(om < target, hi, mid)
    # Bisection selects the branch robustly but is piecewise constant as a
    # function of parameters.  One implicit Newton correction restores the
    # correct d t / d parameter without differentiating through decisions.
    root = ((lo + hi) / 2).detach()
    om = self._evaluate(root[..., None], st)[0][..., 0]
    dom = self._domega_full(root[..., None], st)[..., 0]
    root = root + (target - om) / torch.clamp(dom, min=torch.finfo(root.dtype).tiny)
    return root * M * MSUN_SEC

torchwave.IMRPhenomTHM

Bases: Module

Aligned-spin dominant and subdominant time-domain modes.

Source code in src/torchwave/approximants/imrphenomthm.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
class IMRPhenomTHM(nn.Module):
    """Aligned-spin dominant and subdominant time-domain modes."""
    supported_modes = _POSITIVE

    def __init__(self, mode_array: Iterable[tuple[int, int]] | None = None):
        super().__init__()
        self.mode_array = _validate_modes(mode_array) if mode_array is not None else _POSITIVE + tuple((l, -m) for l, m in _POSITIVE)
        self._tmodel = IMRPhenomT()

    def _positive(self, u, state22, mode, eta, delta, S, dchi, chi1, chi2, final_spin=None):
        l, m = mode
        om22, ph22, a22 = self._tmodel._evaluate(u, state22)
        if mode == (2, 2):
            from ._phenomt_amplitude import amplitude
            return amplitude(u, a22, lambda tt: self._tmodel._evaluate(tt, state22)[0]), ph22, om22
        # The published HM inspiral phase is m/2 times the cached 22 phase.
        phase = m / 2 * ph22 - _OFF[10 * l + m]
        om = m / 2 * om22
        # THM's PN amplitude uses the orbital (22) frequency.
        x = torch.clamp(om22 / 2, min=1e-30) ** (2 / 3)
        coeff = _mode_pn(10 * l + m, eta, delta, chi1, chi2)
        cp = torch.stack([_fit("_IMRPhenomT_Inspiral_Amp_CP", eta, S, dchi, delta, 10 * l + m, i) for i in (1, 2, 3)], -1)
        # Collocation times are fixed in total-mass units; evaluate from 22.
        uc = eta[..., None] * 0 + torch.as_tensor((-2000., -250., -150.), dtype=eta.dtype, device=eta.device)
        # The PN amplitude uses x=(omega_22/2)^(2/3), even for m != 2.
        oc = self._tmodel._evaluate(uc, state22)[0]
        xx = torch.clamp(oc / 2, min=1e-30) ** (2 / 3)
        b0 = _pn_amp(xx, coeff, torch.zeros(eta.shape + (3,), dtype=eta.dtype, device=eta.device)).real
        A = torch.stack((xx**4, xx**4.5, xx**5), -1)
        den = coeff[2]
        den = torch.where(den.abs() < 1e-30, torch.where(den < 0, torch.full_like(den, -1e-30), torch.full_like(den, 1e-30)), den)
        pseudo = torch.linalg.solve(A, ((cp - b0) / (den[..., None] * xx)).unsqueeze(-1))[..., 0]
        # The augmented spin supplied by precession changes the QNM rates in
        # the ringdown only.  Keep the aligned remnant for the merger
        # collocation system so its inspiral/merger portion remains identical
        # to the aligned model; amplitude and phase are then joined at their
        # respective fitted amplitude and frequency cuts.
        af_merger = _fit("IMRPhenomX_FinalSpin2017", eta, chi1, chi2, delta)
        af_rd = af_merger if final_spin is None else final_spin
        mf = _fit("IMRPhenomX_FinalMass2017", eta, chi1, chi2, delta)
        fd = _fit("IMRPhenomT_fdampfit", af_merger, 10 * l + m) / mf
        alpha = 2 * PI * fd
        fd_rd = _fit("IMRPhenomT_fdampfit", af_rd, 10 * l + m) / mf
        alpha_rd = 2 * PI * fd_rd
        peak = _fit("_IMRPhenomT_PeakAmp", eta, S, dchi, delta, 10 * l + m)
        mid = _fit("_IMRPhenomT_Intermediate_Amp_CP1", eta, S, dchi, delta, 10 * l + m)
        c3 = _fit("_IMRPhenomT_Ringdown_Amp_C3", eta, S, dchi, 10 * l + m)
        tshift = _fit("_IMRPhenomT_tshift", eta, S, dchi, 10 * l + m)
        c2 = (2 * PI * _fit("IMRPhenomT_fdampn2fit", af_merger, 10 * l + m) / mf - alpha) / 2
        c2 = torch.where(c2 > torch.abs(-.5 * alpha / torch.tanh(c3)), -.5 * alpha / torch.tanh(c3), c2)
        c2_rd = (2 * PI * _fit("IMRPhenomT_fdampn2fit", af_rd, 10 * l + m) / mf - alpha_rd) / 2
        c2_rd = torch.where(c2_rd > torch.abs(-.5 * alpha_rd / torch.tanh(c3)), -.5 * alpha_rd / torch.tanh(c3), c2_rd)
        c1, c4 = peak * alpha_rd * torch.cosh(c3).square() / c2_rd, peak - peak * alpha_rd * torch.cosh(c3).square() / c2_rd * torch.tanh(c3)
        mt = eta[..., None] * 0 + torch.as_tensor((-150., -25., 0.), dtype=eta.dtype, device=eta.device)
        amp_cut = _pn_amp(xx[..., -1:], coeff, pseudo)[..., 0]
        ic = torch.sign(amp_cut.real) * torch.abs(amp_cut)
        A4 = torch.cat((_basis(mt - tshift[..., None], alpha), _basis((mt - tshift[..., None])[..., :1], alpha, True)), -2)
        # The published matching system fixes the peak row's quadratic basis
        # entry to zero (the ansatz itself still evaluates (t-tshift)^2).
        row2 = torch.cat((torch.ones_like(A4[..., 2:3, :3]), torch.zeros_like(A4[..., 2:3, 3:])), -1)
        A4 = torch.cat((A4[..., :2, :], row2, A4[..., 3:, :]), -2)
        omega_cut = self._tmodel._evaluate(torch.as_tensor(-150., dtype=u.dtype, device=u.device), state22)[0][..., 0]
        domega_cut = self._tmodel._domega_full(torch.as_tensor(-150., dtype=u.dtype, device=u.device), state22)[..., 0]
        x_cut = torch.clamp(omega_cut / 2, min=1e-30) ** (2 / 3)
        # The source ansatz is signed by Re(A_PN) before matching its slope.
        deriv = torch.sign(amp_cut.real) * _pn_amp_abs_derivative(x_cut, coeff, pseudo, omega_cut, domega_cut)
        rhs = torch.stack((ic, mid, peak, deriv), -1)
        mc = torch.linalg.solve(A4, rhs[..., None])[..., 0]
        # The amplitude ansatz switches to ringdown at tshift (the fitted
        # ringdown cut), while the phase/frequency switch is at zero.
        tm = torch.minimum(torch.maximum(u, torch.as_tensor(-150., dtype=u.dtype, device=u.device)), tshift[..., None])
        tr = torch.maximum(u, tshift[..., None])
        am = torch.sum(_basis(tm - tshift[..., None], alpha) * mc[..., None, :], -1)
        tr = tr - tshift[..., None]
        ar = torch.exp(-alpha_rd[..., None] * tr) * (c1[..., None] * torch.tanh(c2_rd[..., None] * tr + c3[..., None]) + c4[..., None])
        amp_complex = _pn_amp(torch.clamp(x, min=1e-30), coeff, pseudo)
        # Keep every branch complex so autograd never sends a complex cotangent
        # into the real merger/RD amplitudes (a PyTorch dtype mismatch).
        am_complex = torch.complex(am, torch.zeros_like(am))
        ar_complex = torch.complex(ar, torch.zeros_like(ar))
        amp = torch.where(u < -150, amp_complex, torch.where(u < tshift[..., None], am_complex, ar_complex))

        # Independent HM merger/ringdown frequency ansatz and its analytic
        # primitive (THM Eqs. 15, 25, 28-31), joined to m/2 times the 22 phase.
        op = _fit("_IMRPhenomT_PeakFrequency", eta, S, dchi, delta, 10 * l + m)
        c2f = _fit("_IMRPhenomT_RD_Freq_C2", eta, S, dchi, delta, 10 * l + m)
        c3f = _fit("_IMRPhenomT_RD_Freq_C3", eta, S, dchi, delta, 10 * l + m)
        wr = 2 * PI * _fit("IMRPhenomT_fringfit", af_merger, 10 * l + m) / mf
        aa = 2 * PI * _fit("IMRPhenomT_fdampfit", af_merger, 10 * l + m) / mf
        c1f = (1 + c3f) * (wr - op) / (c2f * c3f)
        wr_rd = 2 * PI * _fit("IMRPhenomT_fringfit", af_rd, 10 * l + m) / mf
        c1f_rd = (1 + c3f) * (wr_rd - op) / (c2f * c3f)
        cut = torch.as_tensor(-150., dtype=u.dtype, device=u.device)
        tcp = torch.as_tensor(-25., dtype=u.dtype, device=u.device)
        w22cut = self._tmodel._evaluate(cut, state22)[0][..., 0]
        dw22cut = self._tmodel._domega_full(cut, state22)[..., 0]
        x_cut = torch.clamp(w22cut / 2, min=1e-30) ** (2 / 3)
        omega_cut_amp = -_pn_amp_phase_derivative(x_cut, coeff, pseudo, w22cut, dw22cut)
        wcut = m / 2 * w22cut + omega_cut_amp
        dwcut = m / 2 * dw22cut
        dwp = -c1f * c2f.square() * c3f / (1 + c3f).square() / wr
        ac, am = torch.asinh(aa * cut), torch.asinh(aa * tcp)
        wtcp = _fit("_IMRPhenomT_Intermediate_Freq_CP1", eta, S, dchi, delta, 10 * l + m)
        b = torch.stack((1 - wcut / wr - (1 - op / wr) - dwp / aa * ac,
                         1 - wtcp / wr - (1 - op / wr) - dwp / aa * am,
                         -dwcut / wr - dwp / torch.sqrt(1 + (aa * cut).square())), -1)
        matf = torch.stack((torch.stack((ac.square(), ac**3, ac**4), -1), torch.stack((am.square(), am**3, am**4), -1), torch.stack((2*aa*ac/torch.sqrt(1+(aa*cut).square()), 3*aa*ac.square()/torch.sqrt(1+(aa*cut).square()), 4*aa*ac**3/torch.sqrt(1+(aa*cut).square())), -1)), -2)
        cf = torch.linalg.solve(matf, b[..., None])[..., 0]
        tm = torch.clamp(u, -150., 0.)
        tr = torch.clamp(u, min=0.)
        xxm = torch.asinh(aa[..., None] * tm)
        wm = wr[..., None] * (1 - (1 - op[..., None]/wr[..., None] + xxm * (dwp[..., None]/aa[..., None] + xxm * (cf[...,0,None] + xxm*(cf[...,1,None] + xxm*cf[...,2,None])))))
        wrg = wr_rd[..., None] - c1f_rd[..., None]*c2f[..., None]*c3f[..., None]*torch.exp(-c2f[..., None]*tr)/(1+c3f[..., None]*torch.exp(-c2f[..., None]*tr))
        om = torch.where(u < -150, m/2*om22, torch.where(u < 0, wm, wrg))
        # The join phase is evaluated at the scalar cut, rather than borrowed
        # from whichever time samples the caller happened to request.
        phcut22 = self._tmodel._evaluate(cut, state22)[1][..., 0]
        # The complex PN amplitude shifts both the HM cut frequency and phase.
        # atan2 has an undefined derivative at (0, 0).  Odd modes vanish at
        # equal mass/equal spin, so evaluate the inactive branch at a safe
        # unit real amplitude before masking its phase.
        amp_nonzero = torch.abs(amp_cut) > 1e-30
        amp_re = torch.where(amp_nonzero, amp_cut.real, torch.ones_like(amp_cut.real))
        amp_im = torch.where(amp_nonzero, amp_cut.imag, torch.zeros_like(amp_cut.imag))
        phicut = torch.atan2(amp_im, amp_re) + torch.where(amp_re < 0, torch.as_tensor(PI, dtype=u.dtype, device=u.device), torch.zeros_like(amp_re))
        phicut = torch.where(amp_nonzero, phicut, torch.zeros_like(phicut))
        rawcut = m / 2 * phcut22 - phicut - _OFF[10 * l + m]
        from .imrphenomt import IMRPhenomT as _T
        inter = _T._phase_inter_raw(tm, aa[..., None], wr[..., None], op[..., None], dwp[..., None], cf[..., None, :])
        inter0 = _T._phase_inter_raw(torch.zeros_like(eta), aa, wr, op, dwp, cf)
        offm = rawcut - _T._phase_inter_raw(cut, aa, wr, op, dwp, cf)
        rphase = c1f_rd[..., None] * torch.log((1 + c3f[..., None]*torch.exp(-c2f[..., None]*tr))/(1+c3f[..., None])) + wr_rd[..., None]*tr + inter0[..., None] + offm[..., None]
        phase = torch.where(u < -150, phase, torch.where(u < 0, inter + offm[..., None], rphase))
        odd_zero = (torch.abs(delta) < 1e-14) & (torch.abs(chi1 - chi2) < 1e-14)
        odd_zero = odd_zero & (m % 2 == 1)
        amp = torch.where(odd_zero[..., None], torch.zeros_like(amp), amp)
        phase = torch.where(odd_zero[..., None], torch.zeros_like(phase), phase)
        om = torch.where(odd_zero[..., None], torch.zeros_like(om), om)
        return amp, phase, om

    def modes(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
        t = torch.as_tensor(times, dtype=times.dtype if isinstance(times, torch.Tensor) and times.is_floating_point() else torch.float64, device=times.device if isinstance(times, torch.Tensor) else None)
        vals = [torch.as_tensor(v, dtype=t.dtype, device=t.device) for v in (mass1, mass2, chi1, chi2, distance, phic, tc)]
        fref = torch.as_tensor(f_ref, dtype=t.dtype, device=t.device)
        m1, m2, c1, c2, dist, pc, tcenter, fref = torch.broadcast_tensors(*vals, fref)
        if final_spin is not None:
            final_spin = torch.broadcast_tensors(torch.as_tensor(final_spin, dtype=t.dtype, device=t.device), m1)[0]
        if modes is not None:
            if mode_array is not None:
                raise ValueError("pass only one of modes and mode_array")
            mode_array = modes
        swap = m1 < m2
        m1, m2 = torch.where(swap, m2, m1), torch.where(swap, m1, m2)
        c1, c2 = torch.where(swap, c2, c1), torch.where(swap, c1, c2)
        state = self._tmodel._coeff(m1, m2, c1, c2, final_spin)
        M, eta, delta, S, dchi = state[0], state[1], state[2], state[3], state[4]
        u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
        sel = _validate_modes(mode_array) if mode_array is not None else self.mode_array
        scale = (M[..., None] * MSUN_SEC) / (dist[..., None] * MPC_SI / C_SI)
        # Match the 22 reference phase convention used by IMRPhenomT.  A
        # zero reference frequency follows the public 20-Hz convention.
        fphase = torch.where(fref > 0, fref, torch.full_like(fref, 20.0))
        tref = self._tmodel.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=final_spin)
        uref = tref / (M * MSUN_SEC)
        phref = self._tmodel._evaluate(uref[..., None], state)[1][..., 0]
        out = {}
        positive = [(l, m) for l, m in sel if m > 0]
        for l, m in positive:
            if m < 0:
                continue
            amp, phase, _ = self._positive(u, state, (l, m), eta, delta, S, dchi, c1, c2, final_spin)
            out[(l, m)] = scale * amp * torch.exp(-1j * (phase - (m / 2) * phref[..., None] + m * pc[..., None]))
        for l, m in sel:
            if m < 0:
                if (l, -m) not in out:
                    amp, phase, _ = self._positive(u, state, (l, -m), eta, delta, S, dchi, c1, c2, final_spin)
                    out[(l, -m)] = scale * amp * torch.exp(-1j * (phase - ((-m) / 2) * phref[..., None] + (-m) * pc[..., None]))
                out[(l, m)] = (-1) ** l * out[(l, -m)].conj()
        return out

    def time_domain(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., inclination=0., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
        # Polarizations require the equatorial-symmetry partners.  A public
        # positive-mode selection therefore denotes the corresponding pair.
        if modes is not None:
            if mode_array is not None:
                raise ValueError("pass only one of modes and mode_array")
            mode_array = modes
        selected = None if mode_array is None else tuple(mode_array)
        if selected is not None:
            selected = tuple(dict.fromkeys(selected + tuple((l, -m) for l, m in selected if m > 0)))
        modes_out = self.modes(times, mass1, mass2, chi1, chi2, distance, phic, tc, f_ref=f_ref, mode_array=selected, final_spin=final_spin)
        inc = torch.as_tensor(inclination, dtype=next(iter(modes_out.values())).real.dtype, device=next(iter(modes_out.values())).device)
        # s=-2 harmonics for the five supported l,m; closed forms avoid a scipy dependency.
        c = torch.cos(inc[..., None])
        out = torch.zeros_like(next(iter(modes_out.values())).real), torch.zeros_like(next(iter(modes_out.values())).real)
        # Factorial Wigner-d form of the s=-2 spherical harmonics, evaluated
        # at the author's azimuth pi/2.  Summing signed modes preserves the
        # exact 22 projection while retaining every selected harmonic.
        for (l, m), h in modes_out.items():
            d = torch.zeros_like(c)
            for k in range(0, 20):
                a, b, q, r = l + m - k, k, 2 - m + k, l - 2 - k
                if min(a, b, q, r) < 0:
                    continue
                den = math.factorial(a)*math.factorial(b)*math.factorial(q)*math.factorial(r)
                num = math.factorial(l + 2) * math.factorial(l - 2) * math.factorial(l + m) * math.factorial(l - m)
                coef = math.sqrt(num) / den
                d = d + (-1.)**(k + m - 2) * coef * torch.cos(inc[...,None]/2)**(2*l + m - 2*k - 2) * torch.sin(inc[...,None]/2)**(2 - m + 2*k)
            y = math.sqrt((2*l+1)/(4*math.pi)) * d * ((-1) ** m) * ((1j) ** m)
            out = (out[0] + (h*y).real, out[1] - (h*y).imag)
        return out

    forward = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, mode_array=None, final_spin=None, modes=None)

Source code in src/torchwave/approximants/imrphenomthm.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def time_domain(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., inclination=0., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
    # Polarizations require the equatorial-symmetry partners.  A public
    # positive-mode selection therefore denotes the corresponding pair.
    if modes is not None:
        if mode_array is not None:
            raise ValueError("pass only one of modes and mode_array")
        mode_array = modes
    selected = None if mode_array is None else tuple(mode_array)
    if selected is not None:
        selected = tuple(dict.fromkeys(selected + tuple((l, -m) for l, m in selected if m > 0)))
    modes_out = self.modes(times, mass1, mass2, chi1, chi2, distance, phic, tc, f_ref=f_ref, mode_array=selected, final_spin=final_spin)
    inc = torch.as_tensor(inclination, dtype=next(iter(modes_out.values())).real.dtype, device=next(iter(modes_out.values())).device)
    # s=-2 harmonics for the five supported l,m; closed forms avoid a scipy dependency.
    c = torch.cos(inc[..., None])
    out = torch.zeros_like(next(iter(modes_out.values())).real), torch.zeros_like(next(iter(modes_out.values())).real)
    # Factorial Wigner-d form of the s=-2 spherical harmonics, evaluated
    # at the author's azimuth pi/2.  Summing signed modes preserves the
    # exact 22 projection while retaining every selected harmonic.
    for (l, m), h in modes_out.items():
        d = torch.zeros_like(c)
        for k in range(0, 20):
            a, b, q, r = l + m - k, k, 2 - m + k, l - 2 - k
            if min(a, b, q, r) < 0:
                continue
            den = math.factorial(a)*math.factorial(b)*math.factorial(q)*math.factorial(r)
            num = math.factorial(l + 2) * math.factorial(l - 2) * math.factorial(l + m) * math.factorial(l - m)
            coef = math.sqrt(num) / den
            d = d + (-1.)**(k + m - 2) * coef * torch.cos(inc[...,None]/2)**(2*l + m - 2*k - 2) * torch.sin(inc[...,None]/2)**(2 - m + 2*k)
        y = math.sqrt((2*l+1)/(4*math.pi)) * d * ((-1) ** m) * ((1j) ** m)
        out = (out[0] + (h*y).real, out[1] - (h*y).imag)
    return out

modes(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, phic=0.0, tc=0.0, f_ref=20.0, mode_array=None, final_spin=None, modes=None)

Source code in src/torchwave/approximants/imrphenomthm.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
def modes(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
    t = torch.as_tensor(times, dtype=times.dtype if isinstance(times, torch.Tensor) and times.is_floating_point() else torch.float64, device=times.device if isinstance(times, torch.Tensor) else None)
    vals = [torch.as_tensor(v, dtype=t.dtype, device=t.device) for v in (mass1, mass2, chi1, chi2, distance, phic, tc)]
    fref = torch.as_tensor(f_ref, dtype=t.dtype, device=t.device)
    m1, m2, c1, c2, dist, pc, tcenter, fref = torch.broadcast_tensors(*vals, fref)
    if final_spin is not None:
        final_spin = torch.broadcast_tensors(torch.as_tensor(final_spin, dtype=t.dtype, device=t.device), m1)[0]
    if modes is not None:
        if mode_array is not None:
            raise ValueError("pass only one of modes and mode_array")
        mode_array = modes
    swap = m1 < m2
    m1, m2 = torch.where(swap, m2, m1), torch.where(swap, m1, m2)
    c1, c2 = torch.where(swap, c2, c1), torch.where(swap, c1, c2)
    state = self._tmodel._coeff(m1, m2, c1, c2, final_spin)
    M, eta, delta, S, dchi = state[0], state[1], state[2], state[3], state[4]
    u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
    sel = _validate_modes(mode_array) if mode_array is not None else self.mode_array
    scale = (M[..., None] * MSUN_SEC) / (dist[..., None] * MPC_SI / C_SI)
    # Match the 22 reference phase convention used by IMRPhenomT.  A
    # zero reference frequency follows the public 20-Hz convention.
    fphase = torch.where(fref > 0, fref, torch.full_like(fref, 20.0))
    tref = self._tmodel.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=final_spin)
    uref = tref / (M * MSUN_SEC)
    phref = self._tmodel._evaluate(uref[..., None], state)[1][..., 0]
    out = {}
    positive = [(l, m) for l, m in sel if m > 0]
    for l, m in positive:
        if m < 0:
            continue
        amp, phase, _ = self._positive(u, state, (l, m), eta, delta, S, dchi, c1, c2, final_spin)
        out[(l, m)] = scale * amp * torch.exp(-1j * (phase - (m / 2) * phref[..., None] + m * pc[..., None]))
    for l, m in sel:
        if m < 0:
            if (l, -m) not in out:
                amp, phase, _ = self._positive(u, state, (l, -m), eta, delta, S, dchi, c1, c2, final_spin)
                out[(l, -m)] = scale * amp * torch.exp(-1j * (phase - ((-m) / 2) * phref[..., None] + (-m) * pc[..., None]))
            out[(l, m)] = (-1) ** l * out[(l, -m)].conj()
    return out

torchwave.IMRPhenomTP

Bases: IMRPhenomT

Precessing dominant-harmonic PhenomT.

Select the Euler-angle prescription with precession_method and its continuation with precession_version. NNLO remains the default. Rotations and mode projection use native differentiable tensors.

Source code in src/torchwave/approximants/imrphenomtp.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class IMRPhenomTP(IMRPhenomT):
    """Precessing dominant-harmonic PhenomT.

    Select the Euler-angle prescription with ``precession_method`` and its
    continuation with ``precession_version``. NNLO remains the default.
    Rotations and mode projection use native differentiable tensors.
    """

    def time_domain(
        self,
        times,
        mass1,
        mass2,
        chi1=0.0,
        chi2=0.0,
        distance=100.0,
        inclination=0.0,
        phic=0.0,
        tc=0.0,
        f_ref=20.0,
        spin1x=0.0,
        spin1y=0.0,
        spin2x=0.0,
        spin2y=0.0,
        precession_method="nnlo",
        modes=None,
        final_spin=None,
        precession_version=None,
        _precession_state=None,
    ):
        method, version = resolve_precession(precession_method, precession_version)
        if modes is not None:
            requested = tuple(_mode_key(mode) for mode in modes)
            if not requested or any((l, abs(m)) != (2, 2) for l, m in requested):
                raise ValueError("IMRPhenomTP supports only mode (2, 2)")
        return _project_t(
            self,
            times,
            mass1,
            mass2,
            chi1,
            chi2,
            distance,
            inclination,
            phic,
            tc,
            f_ref,
            spin1x,
            spin1y,
            spin2x,
            spin2y,
            final_spin=final_spin,
            precession_method=method,
            precession_version=version,
            _precession_state=_precession_state,
        )

    # IMRPhenomT assigns ``forward`` at class definition time, so bind the
    # override explicitly for nn.Module callers.
    forward = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, spin1x=0.0, spin1y=0.0, spin2x=0.0, spin2y=0.0, precession_method='nnlo', modes=None, final_spin=None, precession_version=None, _precession_state=None)

Source code in src/torchwave/approximants/imrphenomtp.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def time_domain(
    self,
    times,
    mass1,
    mass2,
    chi1=0.0,
    chi2=0.0,
    distance=100.0,
    inclination=0.0,
    phic=0.0,
    tc=0.0,
    f_ref=20.0,
    spin1x=0.0,
    spin1y=0.0,
    spin2x=0.0,
    spin2y=0.0,
    precession_method="nnlo",
    modes=None,
    final_spin=None,
    precession_version=None,
    _precession_state=None,
):
    method, version = resolve_precession(precession_method, precession_version)
    if modes is not None:
        requested = tuple(_mode_key(mode) for mode in modes)
        if not requested or any((l, abs(m)) != (2, 2) for l, m in requested):
            raise ValueError("IMRPhenomTP supports only mode (2, 2)")
    return _project_t(
        self,
        times,
        mass1,
        mass2,
        chi1,
        chi2,
        distance,
        inclination,
        phic,
        tc,
        f_ref,
        spin1x,
        spin1y,
        spin2x,
        spin2y,
        final_spin=final_spin,
        precession_method=method,
        precession_version=version,
        _precession_state=_precession_state,
    )

torchwave.IMRPhenomTPHM

Bases: IMRPhenomT

Precessing multipolar PhenomT model.

Select the Euler-angle prescription with precession_method and its continuation with precession_version. NNLO remains the default. Rotations and mode projection use native differentiable tensors.

Source code in src/torchwave/approximants/imrphenomtphm.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
class IMRPhenomTPHM(IMRPhenomT):
    """Precessing multipolar PhenomT model.

    Select the Euler-angle prescription with ``precession_method`` and its
    continuation with ``precession_version``. NNLO remains the default.
    Rotations and mode projection use native differentiable tensors.
    """

    def time_domain(
        self,
        times,
        mass1,
        mass2,
        chi1=0.0,
        chi2=0.0,
        distance=100.0,
        inclination=0.0,
        phic=0.0,
        tc=0.0,
        f_ref=20.0,
        spin1x=0.0,
        spin1y=0.0,
        spin2x=0.0,
        spin2y=0.0,
        precession_method="nnlo",
        modes=None,
        final_spin=None,
        precession_version=None,
        _precession_state=None,
    ):
        method, version = resolve_precession(precession_method, precession_version)
        requested = tuple(_mode_key(mode) for mode in modes) if modes is not None else DEFAULT_MODES
        positive = tuple(dict.fromkeys((l, abs(m)) for l, m in requested))
        if not positive:
            raise ValueError("modes must contain at least one mode")
        if any(key not in _SUPPORTED_MODES for key in positive):
            raise ValueError(f"unsupported THM mode; supported positive modes are {sorted(_SUPPORTED_MODES)}")
        t = torch.as_tensor(times, dtype=torch.float64)
        vals = [
            torch.as_tensor(x, dtype=t.dtype, device=t.device)
            for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
        ]
        if final_spin is None:
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = torch.broadcast_tensors(*vals)
            fs = None
        else:
            vals.append(torch.as_tensor(final_spin, dtype=t.dtype, device=t.device))
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = torch.broadcast_tensors(*vals)
        m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
        fs, prepared = prepare_precession(
            method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
            final_spin=fs, prepared=_precession_state,
        )
        try:
            from .imrphenomthm import IMRPhenomTHM

            thm = IMRPhenomTHM()
            selected = tuple(positive) + tuple((l, -m) for l, m in positive)
            raw = thm.modes(
                t,
                m1,
                m2,
                z1,
                z2,
                dist,
                phic=0.0,
                tc=tcenter,
                mode_array=selected,
                final_spin=fs,
                f_ref=fr,
            )
        except ImportError as exc:
            raise RuntimeError("IMRPhenomTPHM requires the native IMRPhenomTHM implementation") from exc
        cp = {}
        for key, value in raw.items():
            lm = _mode_key(key)
            if lm in selected:
                cp[lm] = _as_complex_mode(value, t.dtype, t.device)
        if not cp:
            raise ValueError("THM mode provider returned no requested positive modes")
        # Eq. 14 uses the aligned carrier for the angle series; the augmented
        # remnant enters only the post-peak QNM continuation.
        state_angle = self._coeff(m1, m2, z1, z2)
        shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
        freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
        freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
        fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
        a, b, g, thjn, zeta = precession_geometry(
            shifted,
            freq,
            m1,
            m2,
            x1,
            y1,
            z1,
            x2,
            y2,
            z2,
            inc,
            pc,
            fr,
            frequency_peak=fpeak,
            final_spin=fs,
            precession_method=method,
            precession_version=version,
            prepared=prepared,
            frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
        )
        hp_rot, hc_rot = project_modes(cp, a, b, g, thjn[..., None], zeta[..., None])
        hp_al, hc_al = thm.time_domain(
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            inc,
            phic=pc,
            tc=tcenter,
            f_ref=fr,
            mode_array=selected,
            final_spin=fs,
        )
        aligned = (x1.square() + y1.square() + x2.square() + y2.square()) == 0
        return torch.where(aligned[..., None], hp_al, hp_rot), torch.where(aligned[..., None], hc_al, hc_rot)

    forward = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, spin1x=0.0, spin1y=0.0, spin2x=0.0, spin2y=0.0, precession_method='nnlo', modes=None, final_spin=None, precession_version=None, _precession_state=None)

Source code in src/torchwave/approximants/imrphenomtphm.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def time_domain(
    self,
    times,
    mass1,
    mass2,
    chi1=0.0,
    chi2=0.0,
    distance=100.0,
    inclination=0.0,
    phic=0.0,
    tc=0.0,
    f_ref=20.0,
    spin1x=0.0,
    spin1y=0.0,
    spin2x=0.0,
    spin2y=0.0,
    precession_method="nnlo",
    modes=None,
    final_spin=None,
    precession_version=None,
    _precession_state=None,
):
    method, version = resolve_precession(precession_method, precession_version)
    requested = tuple(_mode_key(mode) for mode in modes) if modes is not None else DEFAULT_MODES
    positive = tuple(dict.fromkeys((l, abs(m)) for l, m in requested))
    if not positive:
        raise ValueError("modes must contain at least one mode")
    if any(key not in _SUPPORTED_MODES for key in positive):
        raise ValueError(f"unsupported THM mode; supported positive modes are {sorted(_SUPPORTED_MODES)}")
    t = torch.as_tensor(times, dtype=torch.float64)
    vals = [
        torch.as_tensor(x, dtype=t.dtype, device=t.device)
        for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
    ]
    if final_spin is None:
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = torch.broadcast_tensors(*vals)
        fs = None
    else:
        vals.append(torch.as_tensor(final_spin, dtype=t.dtype, device=t.device))
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = torch.broadcast_tensors(*vals)
    m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
    fs, prepared = prepare_precession(
        method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
        final_spin=fs, prepared=_precession_state,
    )
    try:
        from .imrphenomthm import IMRPhenomTHM

        thm = IMRPhenomTHM()
        selected = tuple(positive) + tuple((l, -m) for l, m in positive)
        raw = thm.modes(
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            phic=0.0,
            tc=tcenter,
            mode_array=selected,
            final_spin=fs,
            f_ref=fr,
        )
    except ImportError as exc:
        raise RuntimeError("IMRPhenomTPHM requires the native IMRPhenomTHM implementation") from exc
    cp = {}
    for key, value in raw.items():
        lm = _mode_key(key)
        if lm in selected:
            cp[lm] = _as_complex_mode(value, t.dtype, t.device)
    if not cp:
        raise ValueError("THM mode provider returned no requested positive modes")
    # Eq. 14 uses the aligned carrier for the angle series; the augmented
    # remnant enters only the post-peak QNM continuation.
    state_angle = self._coeff(m1, m2, z1, z2)
    shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
    freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
    freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
    fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
    a, b, g, thjn, zeta = precession_geometry(
        shifted,
        freq,
        m1,
        m2,
        x1,
        y1,
        z1,
        x2,
        y2,
        z2,
        inc,
        pc,
        fr,
        frequency_peak=fpeak,
        final_spin=fs,
        precession_method=method,
        precession_version=version,
        prepared=prepared,
        frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
    )
    hp_rot, hc_rot = project_modes(cp, a, b, g, thjn[..., None], zeta[..., None])
    hp_al, hc_al = thm.time_domain(
        t,
        m1,
        m2,
        z1,
        z2,
        dist,
        inc,
        phic=pc,
        tc=tcenter,
        f_ref=fr,
        mode_array=selected,
        final_spin=fs,
    )
    aligned = (x1.square() + y1.square() + x2.square() + y2.square()) == 0
    return torch.where(aligned[..., None], hp_al, hp_rot), torch.where(aligned[..., None], hc_al, hc_rot)

Bayesian Inference & Differentiable Likelihoods

torchwave.inference.fisher.fisher_information_matrix(parameters_dict, waveform_fn, psd_dict, detector_network=None, f_low=20.0, f_high=1024.0, df=None, frequencies=None, parameter_order=None)

Compute the differentiable Fisher Information Matrix Gamma_ij using PyTorch automatic differentiation.

Literature: Cutler & Flanagan (1994), Vallisneri (2008).

Formula

Gamma_ij = sum_{k in network} 4 Re int_{f_low}^{f_high} (1 / Sn^(k)(f)) * (dh_k(f)/dtheta_i) * (dh_k(f)/dtheta_j)* df

Parameters:

Name Type Description Default
parameters_dict Dict[str, Union[float, Tensor]]

Dict mapping parameter names to parameter values/tensors

 required
waveform_fn Callable[[Dict[str, Tensor]], Union[Tensor, Dict[str, Tensor], Tuple[Tensor, Tensor]]]

Function mapping parameter dictionary to frequency-domain strain(s)

 required
psd_dict Union[Tensor, Dict[str, Tensor]]

Real PSD tensor or dict of PSDs for each detector {det: psd_det}

 required
detector_network Optional[Any]

Optional DetectorNetwork instance to project (h_plus, h_cross)

 None
f_low float

Low-frequency cutoff in Hz

 20.0
f_high float

High-frequency cutoff in Hz

 1024.0
df Optional[float]

Frequency resolution in Hz

 None
frequencies Optional[Tensor]

Optional frequency vector in Hz

 None
parameter_order Optional[Sequence[str]]

Optional explicit order of parameter keys

 None

Returns:

Type Description
(fisher_matrix, param_names)

(P, P) symmetric Fisher matrix and list of parameter names.
List[str]

Degenerate parameters can make the matrix singular.
Source code in src/torchwave/inference/fisher.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def fisher_information_matrix(
    parameters_dict: Dict[str, Union[float, torch.Tensor]],
    waveform_fn: Callable[[Dict[str, torch.Tensor]], Union[torch.Tensor, Dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]],
    psd_dict: Union[torch.Tensor, Dict[str, torch.Tensor]],
    detector_network: Optional[Any] = None,
    f_low: float = 20.0,
    f_high: float = 1024.0,
    df: Optional[float] = None,
    frequencies: Optional[torch.Tensor] = None,
    parameter_order: Optional[Sequence[str]] = None,
) -> Tuple[torch.Tensor, List[str]]:
    """Compute the differentiable Fisher Information Matrix Gamma_ij using PyTorch automatic differentiation.

    Literature: Cutler & Flanagan (1994), Vallisneri (2008).

    Formula:
        Gamma_ij = sum_{k in network} 4 Re int_{f_low}^{f_high} (1 / Sn^(k)(f)) * (dh_k(f)/dtheta_i) * (dh_k(f)/dtheta_j)* df

    Args:
        parameters_dict: Dict mapping parameter names to parameter values/tensors
        waveform_fn: Function mapping parameter dictionary to frequency-domain strain(s)
        psd_dict: Real PSD tensor or dict of PSDs for each detector {det: psd_det}
        detector_network: Optional DetectorNetwork instance to project (h_plus, h_cross)
        f_low: Low-frequency cutoff in Hz
        f_high: High-frequency cutoff in Hz
        df: Frequency resolution in Hz
        frequencies: Optional frequency vector in Hz
        parameter_order: Optional explicit order of parameter keys

    Returns:
        (fisher_matrix, param_names): (P, P) symmetric Fisher matrix and list of parameter names.
        Degenerate parameters can make the matrix singular.
    """
    if parameter_order is None:
        param_names = list(parameters_dict.keys())
    else:
        param_names = list(parameter_order)

    device = None
    dtype = torch.float64
    for k in param_names:
        v = parameters_dict[k]
        if isinstance(v, torch.Tensor):
            device = v.device
            if v.is_floating_point():
                dtype = v.dtype
            break
    if device is None:
        device = torch.device("cpu")

    base_params = {}
    for k in param_names:
        val = parameters_dict[k]
        if not isinstance(val, torch.Tensor):
            base_params[k] = torch.tensor(val, dtype=dtype, device=device)
        else:
            base_params[k] = val.detach().to(device=device, dtype=dtype)

    num_params = len(param_names)

    def eval_strains(p_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        wf_out = waveform_fn(p_dict)
        if isinstance(wf_out, dict):
            return wf_out
        elif isinstance(wf_out, tuple) and len(wf_out) == 2:
            hp, hc = wf_out
            if detector_network is not None:
                ra = p_dict.get("ra", torch.tensor(0.0, device=device, dtype=dtype))
                dec = p_dict.get("dec", torch.tensor(0.0, device=device, dtype=dtype))
                psi = p_dict.get("psi", torch.tensor(0.0, device=device, dtype=dtype))
                tg = p_dict.get("t_gps", p_dict.get("tc", torch.tensor(0.0, device=device, dtype=dtype)))
                freq_vec = frequencies if frequencies is not None else torch.tensor(0.0, device=device, dtype=dtype)
                return detector_network.project_fd_strain(freq_vec, hp, hc, ra, dec, psi, tg)
            else:
                return {"det": hp}
        elif isinstance(wf_out, torch.Tensor):
            return {"det": wf_out}
        else:
            raise TypeError(f"Unsupported waveform_fn output type: {type(wf_out)}")

    det_jacobians: Dict[str, List[torch.Tensor]] = {}

    for i, p_name in enumerate(param_names):
        tangents = {k: torch.zeros_like(base_params[k]) for k in param_names}
        tangents[p_name] = torch.ones_like(base_params[p_name])

        with fwAD.dual_level():
            dual_params = {k: fwAD.make_dual(base_params[k], tangents[k]) for k in param_names}
            strains = eval_strains(dual_params)
            for det_k, strain_k in strains.items():
                unpacked = fwAD.unpack_dual(strain_k)
                dh_i = unpacked.tangent
                if dh_i is None:
                    dh_i = torch.zeros_like(unpacked.primal)
                if det_k not in det_jacobians:
                    det_jacobians[det_k] = []
                det_jacobians[det_k].append(dh_i)

    fisher = torch.zeros((num_params, num_params), dtype=dtype, device=device)

    for det_k, jac_list in det_jacobians.items():
        if isinstance(psd_dict, dict):
            if det_k not in psd_dict:
                if len(psd_dict) == 1:
                    psd_k = next(iter(psd_dict.values()))
                else:
                    continue
            else:
                psd_k = psd_dict[det_k]
        else:
            psd_k = psd_dict

        psd_safe = torch.clamp(psd_k, min=1.0e-50)
        weight = 4.0 / psd_safe

        if frequencies is not None:
            mask = torch.ones_like(frequencies, dtype=torch.bool)
            if f_low is not None:
                mask = mask & (frequencies >= f_low)
            if f_high is not None:
                mask = mask & (frequencies <= f_high)
            weight = torch.where(mask, weight, torch.zeros_like(weight))
            integration_mode = "trapezoid"
        else:
            df_val = df if df is not None else 1.0
            weight = weight * df_val
            integration_mode = "sum"

        J = torch.stack(jac_list, dim=-1)

        if integration_mode == "trapezoid":
            for i in range(num_params):
                for j in range(i, num_params):
                    integrand = weight * torch.real(jac_list[i] * torch.conj(jac_list[j]))
                    integral = torch.trapezoid(integrand, frequencies, dim=-1)
                    fisher[i, j] = fisher[i, j] + integral
                    if i != j:
                        fisher[j, i] = fisher[i, j]
        else:
            W = weight.unsqueeze(-1)
            W_J = W * J
            Gamma_det = torch.real(torch.matmul(torch.conj(J).transpose(-2, -1), W_J))
            Gamma_det = 0.5 * (Gamma_det + Gamma_det.transpose(-2, -1))
            fisher = fisher + Gamma_det

    return fisher, param_names

torchwave.inference.fisher.covariance_matrix(fisher_matrix, eps=1e-12)

Compute a regularized inverse Fisher matrix for local uncertainty forecasts.

Interpreting this as a posterior covariance requires a suitable local Gaussian approximation. Eigenvalue flooring changes the inverse in ill-conditioned directions; inspect conditioning and regularization before interpretation.

Literature: Vallisneri (2008) Phys. Rev. D 77, 042001.

Parameters:

Name Type Description Default
fisher_matrix Tensor

(P, P) symmetric Fisher information matrix Gamma

 required
eps float

Minimum relative eigenvalue floor (regularization threshold)

 1e-12

Returns:

Type Description
Tensor

(P, P) symmetric regularized inverse Fisher matrix.
Source code in src/torchwave/inference/fisher.py 
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def covariance_matrix(
    fisher_matrix: torch.Tensor,
    eps: float = 1.0e-12,
) -> torch.Tensor:
    """Compute a regularized inverse Fisher matrix for local uncertainty forecasts.

    Interpreting this as a posterior covariance requires a suitable local Gaussian
    approximation. Eigenvalue flooring changes the inverse in ill-conditioned
    directions; inspect conditioning and regularization before interpretation.

    Literature: Vallisneri (2008) Phys. Rev. D 77, 042001.

    Args:
        fisher_matrix: (P, P) symmetric Fisher information matrix Gamma
        eps: Minimum relative eigenvalue floor (regularization threshold)

    Returns:
        (P, P) symmetric regularized inverse Fisher matrix.
    """
    gamma_sym = 0.5 * (fisher_matrix + fisher_matrix.transpose(-2, -1))
    eigenvalues, eigenvectors = torch.linalg.eigh(gamma_sym)

    max_ev = torch.max(torch.abs(eigenvalues))
    threshold = torch.clamp(max_ev * eps, min=1.0e-30)
    ev_reg = torch.clamp(eigenvalues, min=threshold)

    inv_ev = 1.0 / ev_reg
    cov = torch.matmul(eigenvectors, torch.matmul(torch.diag_embed(inv_ev), eigenvectors.transpose(-2, -1)))
    return 0.5 * (cov + cov.transpose(-2, -1))

torchwave.inference.likelihood.WhittleLikelihood

Bases: Module

Differentiable Whittle Gaussian Log-Likelihood for Gravitational-Wave Detector Networks.

Source code in src/torchwave/inference/likelihood.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
class WhittleLikelihood(nn.Module):
    """Differentiable Whittle Gaussian Log-Likelihood for Gravitational-Wave Detector Networks."""

    def __init__(
        self,
        data: Dict[str, torch.Tensor],
        psds: Dict[str, torch.Tensor],
        df: Optional[float] = None,
        frequencies: Optional[torch.Tensor] = None,
        f_low: Optional[float] = None,
        f_high: Optional[float] = None,
    ):
        super().__init__()
        self.detectors = sorted(list(data.keys()))
        self.df = df
        self.f_low = f_low
        self.f_high = f_high

        for det in self.detectors:
            d_tensor = as_tensor(data[det])
            psd_tensor = as_tensor(psds[det])
            self.register_buffer(f"data_{det}", d_tensor, persistent=True)
            self.register_buffer(f"psd_{det}", psd_tensor, persistent=True)

        if frequencies is not None:
            self.register_buffer("frequencies", as_tensor(frequencies), persistent=True)
        else:
            self.frequencies = None

        # Precompute noise inner products <d | d>
        self._data_inner: Dict[str, torch.Tensor] = {}
        for det in self.detectors:
            d = getattr(self, f"data_{det}")
            psd = getattr(self, f"psd_{det}")
            self._data_inner[det] = noise_weighted_inner_product(
                d, d, psd, df=self.df, frequencies=self.frequencies, f_low=self.f_low, f_high=self.f_high
            )

    def log_likelihood(
        self,
        templates: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Compute standard Whittle log-likelihood ln L(d | theta) = -0.5 * sum_k <d_k - h_k | d_k - h_k>.

        Args:
            templates: Dictionary mapping detector prefix to frequency-domain strain tensor h_det(f)
        """
        ln_l = torch.tensor(0.0, dtype=torch.float64, device=getattr(self, f"data_{self.detectors[0]}").device)

        for det in self.detectors:
            d = getattr(self, f"data_{det}")
            psd = getattr(self, f"psd_{det}")
            h = templates[det]

            d_minus_h = d - h
            res_inner = noise_weighted_inner_product(
                d_minus_h,
                d_minus_h,
                psd,
                df=self.df,
                frequencies=self.frequencies,
                f_low=self.f_low,
                f_high=self.f_high,
            )
            ln_l = ln_l - 0.5 * res_inner

        return ln_l

    def phase_marginalized_log_likelihood(
        self,
        templates: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Analytical coalescence phase-marginalized log-likelihood (Veitch et al. 2015).

        ln L_marg = -0.5 * sum_k <h_k | h_k> + log(I0(|Z|)) - 0.5 * sum_k <d_k | d_k>
        where Z = sum_k 4 * int (d_k * h_k* / Sn_k) df.
        """
        device = getattr(self, f"data_{self.detectors[0]}").device
        z_total = torch.tensor(0.0, dtype=torch.complex128, device=device)
        hh_total = torch.tensor(0.0, dtype=torch.float64, device=device)
        dd_total = torch.tensor(0.0, dtype=torch.float64, device=device)

        for det in self.detectors:
            d = getattr(self, f"data_{det}")
            psd = getattr(self, f"psd_{det}")
            h = templates[det]

            z_k = noise_weighted_inner_product(
                d,
                h,
                psd,
                df=self.df,
                frequencies=self.frequencies,
                f_low=self.f_low,
                f_high=self.f_high,
                complex_output=True,
            )
            hh_k = noise_weighted_inner_product(
                h,
                h,
                psd,
                df=self.df,
                frequencies=self.frequencies,
                f_low=self.f_low,
                f_high=self.f_high,
            )

            z_total = z_total + z_k
            hh_total = hh_total + hh_k
            dd_total = dd_total + self._data_inner[det]

        ln_i0 = log_bessel_i0(torch.abs(z_total))
        ln_l = ln_i0 - 0.5 * hh_total - 0.5 * dd_total
        return ln_l

    def distance_marginalized_log_likelihood(
        self,
        templates: Dict[str, torch.Tensor],
        d_min: float = 10.0,
        d_max: float = 2000.0,
        d_ref: float = 100.0,
        n_quad: int = 64,
    ) -> torch.Tensor:
        """Differentiable distance-marginalized log-likelihood with prior p(dL) proportional to dL^2.

        Uses Gauss-Legendre quadrature over [d_min, d_max].
        """
        device = getattr(self, f"data_{self.detectors[0]}").device
        z_re = torch.tensor(0.0, dtype=torch.float64, device=device)
        hh_total = torch.tensor(0.0, dtype=torch.float64, device=device)
        dd_total = torch.tensor(0.0, dtype=torch.float64, device=device)

        for det in self.detectors:
            d = getattr(self, f"data_{det}")
            psd = getattr(self, f"psd_{det}")
            h = templates[det]

            dh_k = noise_weighted_inner_product(
                d,
                h,
                psd,
                df=self.df,
                frequencies=self.frequencies,
                f_low=self.f_low,
                f_high=self.f_high,
            )
            hh_k = noise_weighted_inner_product(
                h,
                h,
                psd,
                df=self.df,
                frequencies=self.frequencies,
                f_low=self.f_low,
                f_high=self.f_high,
            )

            z_re = z_re + dh_k
            hh_total = hh_total + hh_k
            dd_total = dd_total + self._data_inner[det]

        # Gauss-Legendre quadrature nodes and weights on [-1, 1]
        nodes_np, weights_np = math.erf, None
        import numpy as np
        nodes_np, weights_np = np.polynomial.legendre.leggauss(n_quad)

        # Map to [d_min, d_max]
        nodes = torch.tensor(nodes_np, dtype=torch.float64, device=device)
        weights = torch.tensor(weights_np, dtype=torch.float64, device=device)

        dl = 0.5 * (d_max - d_min) * nodes + 0.5 * (d_max + d_min)
        w_quad = weights * 0.5 * (d_max - d_min)

        # Prior p(dL) = 3 * dL^2 / (d_max^3 - d_min^3)
        vol = (d_max ** 3.0 - d_min ** 3.0) / 3.0
        p_dl = (dl * dl) / vol

        # Scale factor u = d_ref / dL
        u = d_ref / dl

        # Integrand: p(dL) * exp(u * z_re - 0.5 * u^2 * hh - 0.5 * dd)
        arg = u * z_re - 0.5 * (u * u) * hh_total - 0.5 * dd_total
        max_arg = torch.max(arg)

        integrand = p_dl * torch.exp(arg - max_arg)
        quad_sum = torch.sum(w_quad * integrand)

        ln_l = max_arg + torch.log(torch.clamp(quad_sum, min=1.0e-50))
        return ln_l

    def forward(
        self,
        templates: Dict[str, torch.Tensor],
        marginalize_phase: bool = False,
        marginalize_distance: bool = False,
        **kwargs,
    ) -> torch.Tensor:
        if marginalize_phase:
            return self.phase_marginalized_log_likelihood(templates)
        if marginalize_distance:
            return self.distance_marginalized_log_likelihood(templates, **kwargs)
        return self.log_likelihood(templates)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

distance_marginalized_log_likelihood(templates, d_min=10.0, d_max=2000.0, d_ref=100.0, n_quad=64)

Differentiable distance-marginalized log-likelihood with prior p(dL) proportional to dL^2.

Uses Gauss-Legendre quadrature over [d_min, d_max].

Source code in src/torchwave/inference/likelihood.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def distance_marginalized_log_likelihood(
    self,
    templates: Dict[str, torch.Tensor],
    d_min: float = 10.0,
    d_max: float = 2000.0,
    d_ref: float = 100.0,
    n_quad: int = 64,
) -> torch.Tensor:
    """Differentiable distance-marginalized log-likelihood with prior p(dL) proportional to dL^2.

    Uses Gauss-Legendre quadrature over [d_min, d_max].
    """
    device = getattr(self, f"data_{self.detectors[0]}").device
    z_re = torch.tensor(0.0, dtype=torch.float64, device=device)
    hh_total = torch.tensor(0.0, dtype=torch.float64, device=device)
    dd_total = torch.tensor(0.0, dtype=torch.float64, device=device)

    for det in self.detectors:
        d = getattr(self, f"data_{det}")
        psd = getattr(self, f"psd_{det}")
        h = templates[det]

        dh_k = noise_weighted_inner_product(
            d,
            h,
            psd,
            df=self.df,
            frequencies=self.frequencies,
            f_low=self.f_low,
            f_high=self.f_high,
        )
        hh_k = noise_weighted_inner_product(
            h,
            h,
            psd,
            df=self.df,
            frequencies=self.frequencies,
            f_low=self.f_low,
            f_high=self.f_high,
        )

        z_re = z_re + dh_k
        hh_total = hh_total + hh_k
        dd_total = dd_total + self._data_inner[det]

    # Gauss-Legendre quadrature nodes and weights on [-1, 1]
    nodes_np, weights_np = math.erf, None
    import numpy as np
    nodes_np, weights_np = np.polynomial.legendre.leggauss(n_quad)

    # Map to [d_min, d_max]
    nodes = torch.tensor(nodes_np, dtype=torch.float64, device=device)
    weights = torch.tensor(weights_np, dtype=torch.float64, device=device)

    dl = 0.5 * (d_max - d_min) * nodes + 0.5 * (d_max + d_min)
    w_quad = weights * 0.5 * (d_max - d_min)

    # Prior p(dL) = 3 * dL^2 / (d_max^3 - d_min^3)
    vol = (d_max ** 3.0 - d_min ** 3.0) / 3.0
    p_dl = (dl * dl) / vol

    # Scale factor u = d_ref / dL
    u = d_ref / dl

    # Integrand: p(dL) * exp(u * z_re - 0.5 * u^2 * hh - 0.5 * dd)
    arg = u * z_re - 0.5 * (u * u) * hh_total - 0.5 * dd_total
    max_arg = torch.max(arg)

    integrand = p_dl * torch.exp(arg - max_arg)
    quad_sum = torch.sum(w_quad * integrand)

    ln_l = max_arg + torch.log(torch.clamp(quad_sum, min=1.0e-50))
    return ln_l

log_likelihood(templates)

Compute standard Whittle log-likelihood ln L(d | theta) = -0.5 * sum_k .

Parameters:

Name Type Description Default
templates Dict[str, Tensor]

Dictionary mapping detector prefix to frequency-domain strain tensor h_det(f)

 required
Source code in src/torchwave/inference/likelihood.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def log_likelihood(
    self,
    templates: Dict[str, torch.Tensor],
) -> torch.Tensor:
    """Compute standard Whittle log-likelihood ln L(d | theta) = -0.5 * sum_k <d_k - h_k | d_k - h_k>.

    Args:
        templates: Dictionary mapping detector prefix to frequency-domain strain tensor h_det(f)
    """
    ln_l = torch.tensor(0.0, dtype=torch.float64, device=getattr(self, f"data_{self.detectors[0]}").device)

    for det in self.detectors:
        d = getattr(self, f"data_{det}")
        psd = getattr(self, f"psd_{det}")
        h = templates[det]

        d_minus_h = d - h
        res_inner = noise_weighted_inner_product(
            d_minus_h,
            d_minus_h,
            psd,
            df=self.df,
            frequencies=self.frequencies,
            f_low=self.f_low,
            f_high=self.f_high,
        )
        ln_l = ln_l - 0.5 * res_inner

    return ln_l

phase_marginalized_log_likelihood(templates)

Analytical coalescence phase-marginalized log-likelihood (Veitch et al. 2015).

ln L_marg = -0.5 * sum_k + log(I0(|Z|)) - 0.5 * sum_k where Z = sum_k 4 * int (d_k * h_k* / Sn_k) df.

Source code in src/torchwave/inference/likelihood.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def phase_marginalized_log_likelihood(
    self,
    templates: Dict[str, torch.Tensor],
) -> torch.Tensor:
    """Analytical coalescence phase-marginalized log-likelihood (Veitch et al. 2015).

    ln L_marg = -0.5 * sum_k <h_k | h_k> + log(I0(|Z|)) - 0.5 * sum_k <d_k | d_k>
    where Z = sum_k 4 * int (d_k * h_k* / Sn_k) df.
    """
    device = getattr(self, f"data_{self.detectors[0]}").device
    z_total = torch.tensor(0.0, dtype=torch.complex128, device=device)
    hh_total = torch.tensor(0.0, dtype=torch.float64, device=device)
    dd_total = torch.tensor(0.0, dtype=torch.float64, device=device)

    for det in self.detectors:
        d = getattr(self, f"data_{det}")
        psd = getattr(self, f"psd_{det}")
        h = templates[det]

        z_k = noise_weighted_inner_product(
            d,
            h,
            psd,
            df=self.df,
            frequencies=self.frequencies,
            f_low=self.f_low,
            f_high=self.f_high,
            complex_output=True,
        )
        hh_k = noise_weighted_inner_product(
            h,
            h,
            psd,
            df=self.df,
            frequencies=self.frequencies,
            f_low=self.f_low,
            f_high=self.f_high,
        )

        z_total = z_total + z_k
        hh_total = hh_total + hh_k
        dd_total = dd_total + self._data_inner[det]

    ln_i0 = log_bessel_i0(torch.abs(z_total))
    ln_l = ln_i0 - 0.5 * hh_total - 0.5 * dd_total
    return ln_l

torchwave.inference.relative_binning.RelativeBinningLikelihood

Bases: Module

Approximate Whittle log-likelihood using relative binning (Zackay et al. 2018).

Precomputes summary coefficients from dense data and a fiducial waveform, then evaluates candidate waveforms at M bin frequencies instead of N_f dense frequencies. The approximation uses local linear waveform ratios.

Accuracy and evaluation cost depend on binning, the reference waveform, signal properties, and parameter displacement. Validate likelihood and gradient errors against dense evaluation over the intended parameter region. Autodiff differentiates this approximation, not the dense likelihood.

Source code in src/torchwave/inference/relative_binning.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
class RelativeBinningLikelihood(nn.Module):
    """Approximate Whittle log-likelihood using relative binning (Zackay et al. 2018).

    Precomputes summary coefficients from dense data and a fiducial waveform,
    then evaluates candidate waveforms at M bin frequencies instead of N_f dense
    frequencies. The approximation uses local linear waveform ratios.

    Accuracy and evaluation cost depend on binning, the reference waveform,
    signal properties, and parameter displacement. Validate likelihood and
    gradient errors against dense evaluation over the intended parameter region.
    Autodiff differentiates this approximation, not the dense likelihood.
    """

    def __init__(
        self,
        data: Dict[str, torch.Tensor],
        psds: Dict[str, torch.Tensor],
        h0: Dict[str, torch.Tensor],
        frequencies: torch.Tensor,
        bin_edges: Optional[torch.Tensor] = None,
        n_bins: int = 50,
        f_min: Optional[float] = None,
        f_max: Optional[float] = None,
        df: Optional[float] = None,
    ):
        """Initialize RelativeBinningLikelihood with precomputed data summary.

        Args:
            data: Dictionary mapping detector name (e.g. 'H1', 'L1') to dense frequency-domain data.
            psds: Dictionary mapping detector name to PSD tensor S_n(f).
            h0: Dictionary mapping detector name to fiducial reference waveform h_0(f).
            frequencies: Dense frequency grid in Hz.
            bin_edges: Optional precomputed bin edges (M + 1,).
            n_bins: Number of bins if bin_edges is not provided.
            f_min: Minimum frequency cutoff in Hz (default: min of frequencies).
            f_max: Maximum frequency cutoff in Hz (default: max of frequencies).
            df: Dense frequency spacing in Hz.
        """
        super().__init__()
        self.detectors = sorted(list(data.keys()))
        f_dense = as_tensor(frequencies)
        self.register_buffer("frequencies", f_dense, persistent=True)

        min_f = f_min if f_min is not None else float(f_dense[0])
        max_f = f_max if f_max is not None else float(f_dense[-1])

        if bin_edges is None:
            edges = compute_relative_binning_bins(
                f_min=min_f,
                f_max=max_f,
                n_bins=n_bins,
                frequencies=f_dense,
                device=f_dense.device,
                dtype=f_dense.dtype,
            )
        else:
            edges = as_tensor(bin_edges, device=f_dense.device, dtype=f_dense.dtype)

        self.register_buffer("bin_edges", edges, persistent=True)
        self.n_bins = len(edges) - 1

        # Use the same dense-grid frequencies for fiducial and candidate evaluations
        raw_centers = 0.5 * (edges[:-1] + edges[1:])
        center_indices = torch.argmin(torch.abs(f_dense.unsqueeze(1) - raw_centers.unsqueeze(0)), dim=0)
        centers = f_dense[center_indices]
        self.register_buffer("bin_centers", centers, persistent=True)
        self.register_buffer("center_indices", center_indices, persistent=False)

        # Precompute summary coefficients for each detector
        self.d_d: Dict[str, torch.Tensor] = {}
        for det in self.detectors:
            d_tensor = as_tensor(data[det], device=f_dense.device)
            psd_tensor = as_tensor(psds[det], device=f_dense.device)
            h0_tensor = as_tensor(h0[det], device=f_dense.device)

            a0, a1, b0, b1, _, d_d_val = compute_relative_binning_summary(
                data=d_tensor,
                psd=psd_tensor,
                h0=h0_tensor,
                frequencies=f_dense,
                bin_edges=edges,
                df=df,
            )

            # Re-evaluate a1 and b1 with exact snapped centers
            safe_psd = torch.clamp(psd_tensor, min=1.0e-50)
            df_val = float(df) if df is not None else float(f_dense[1] - f_dense[0])
            integrand_a = 4.0 * ((d_tensor * torch.conj(h0_tensor)) / safe_psd) * df_val
            integrand_b = 4.0 * (torch.real(h0_tensor * torch.conj(h0_tensor)) / safe_psd) * df_val

            a1_list = []
            b1_list = []
            for k in range(self.n_bins):
                f_low = edges[k]
                f_high = edges[k + 1]
                f_bar = centers[k]
                if k == self.n_bins - 1:
                    mask = (f_dense >= f_low) & (f_dense <= f_high)
                else:
                    mask = (f_dense >= f_low) & (f_dense < f_high)
                delta_f = f_dense[mask] - f_bar
                a1_list.append(torch.sum(integrand_a[mask] * delta_f))
                b1_list.append(torch.sum(integrand_b[mask] * delta_f))

            a1 = torch.stack(a1_list)
            b1 = torch.stack(b1_list)

            self.register_buffer(f"a0_{det}", a0, persistent=True)
            self.register_buffer(f"a1_{det}", a1, persistent=True)
            self.register_buffer(f"b0_{det}", b0, persistent=True)
            self.register_buffer(f"b1_{det}", b1, persistent=True)
            self.register_buffer(f"h0_centers_{det}", h0_tensor[center_indices], persistent=True)
            self.register_buffer(f"h0_dense_{det}", h0_tensor, persistent=False)
            self.register_buffer(f"dd_{det}", d_d_val, persistent=True)

    @staticmethod
    def compute_r0_r1(
        h_centers: torch.Tensor,
        h0_centers: torch.Tensor,
        bin_centers: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute the ratio r_0 and slope r_1 at bin centers.

        Args:
            h_centers: Test waveform evaluated at bin centers, shape (..., M).
            h0_centers: Reference waveform evaluated at bin centers, shape (M,) or (..., M).
            bin_centers: Bin center frequencies, shape (M,).

        Returns:
            Tuple of (r0, r1) with same batch dimensions as h_centers.
        """
        safe_h0 = torch.where(
            torch.abs(h0_centers) < 1e-40,
            torch.ones_like(h0_centers) * 1e-40,
            h0_centers,
        )
        r0 = h_centers / safe_h0

        m = bin_centers.shape[-1]
        r1 = torch.zeros_like(r0)

        if m >= 3:
            df_interior = bin_centers[2:] - bin_centers[:-2]
            r1[..., 1:-1] = (r0[..., 2:] - r0[..., :-2]) / df_interior

            # Boundary 0
            df_0 = bin_centers[1] - bin_centers[0]
            r1[..., 0] = (r0[..., 1] - r0[..., 0]) / df_0

            # Boundary M - 1
            df_end = bin_centers[-1] - bin_centers[-2]
            r1[..., -1] = (r0[..., -1] - r0[..., -2]) / df_end
        elif m == 2:
            df_step = bin_centers[1] - bin_centers[0]
            diff = (r0[..., 1] - r0[..., 0]) / df_step
            r1[..., 0] = diff
            r1[..., 1] = diff

        return r0, r1

    def inner_product_dh(
        self,
        r0: torch.Tensor,
        r1: torch.Tensor,
        detector: str,
    ) -> torch.Tensor:
        """Compute noise-weighted inner product <d | h(theta)> from summary coefficients.

        <d | h> = Re sum_{k=0}^{M-1} [ r_0^*(f_bar_k) * A_{0, k} + r_1^*(f_bar_k) * A_{1, k} ]
        """
        a0 = getattr(self, f"a0_{detector}")
        a1 = getattr(self, f"a1_{detector}")

        r0_conj = torch.conj(r0)
        r1_conj = torch.conj(r1)

        term0 = r0_conj * a0
        term1 = r1_conj * a1

        sum_terms = torch.sum(term0 + term1, dim=-1)
        return torch.real(sum_terms)

    def inner_product_hh(
        self,
        r0: torch.Tensor,
        r1: torch.Tensor,
        detector: str,
    ) -> torch.Tensor:
        """Compute waveform self inner product <h(theta) | h(theta)> from summary coefficients.

        <h | h> = sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]
        """
        b0 = getattr(self, f"b0_{detector}")
        b1 = getattr(self, f"b1_{detector}")

        r0_sq = torch.real(r0 * torch.conj(r0))
        cross_term = 2.0 * torch.real(r0 * torch.conj(r1))

        sum_terms = torch.sum(r0_sq * b0 + cross_term * b1, dim=-1)
        return sum_terms

    def log_likelihood_from_r(
        self,
        r0_dict: Dict[str, torch.Tensor],
        r1_dict: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

        ln L = sum_{det} [ <d | h> - 0.5 * <h | h> - 0.5 * <d | d> ]
        """
        device = self.bin_centers.device
        total_ln_l = torch.tensor(0.0, dtype=torch.float64, device=device)

        for det in self.detectors:
            r0 = r0_dict[det]
            r1 = r1_dict[det]
            dh = self.inner_product_dh(r0, r1, det)
            hh = self.inner_product_hh(r0, r1, det)
            dd = getattr(self, f"dd_{det}")

            det_ln_l = dh - 0.5 * hh - 0.5 * dd
            total_ln_l = total_ln_l + det_ln_l

        return total_ln_l

    def log_likelihood(
        self,
        templates_at_bin_centers: Dict[str, torch.Tensor],
        h0_at_bin_centers: Optional[Dict[str, torch.Tensor]] = None,
        slopes: Optional[Dict[str, torch.Tensor]] = None,
    ) -> torch.Tensor:
        """Compute relative binning log-likelihood from waveforms evaluated at bin centers.

        Args:
            templates_at_bin_centers: Dictionary mapping detector name to waveform
                tensor h(theta; f_bar) evaluated ONLY at bin centers.
            h0_at_bin_centers: Optional fiducial waveform h0(f_bar). If not provided,
                sampled from stored dense h0.
            slopes: Optional explicit slopes r_1(f_bar). If None, calculated via
                centered finite differences.

        Returns:
            Scalar or batched log-likelihood tensor.
        """
        r0_dict: Dict[str, torch.Tensor] = {}
        r1_dict: Dict[str, torch.Tensor] = {}

        for det in self.detectors:
            h_det = templates_at_bin_centers[det]
            if h0_at_bin_centers is not None:
                h0_det = h0_at_bin_centers[det]
            elif hasattr(self, f"h0_centers_{det}"):
                h0_det = getattr(self, f"h0_centers_{det}")
            else:
                h0_dense = getattr(self, f"h0_dense_{det}")
                f_dense = self.frequencies
                idx = torch.searchsorted(f_dense, self.bin_centers).clamp(max=len(f_dense) - 1)
                h0_det = h0_dense[idx]

            if slopes is not None and det in slopes:
                safe_h0 = torch.where(
                    torch.abs(h0_det) < 1e-40,
                    torch.ones_like(h0_det) * 1e-40,
                    h0_det,
                )
                r0 = h_det / safe_h0
                r1 = slopes[det]
            else:
                r0, r1 = self.compute_r0_r1(h_det, h0_det, self.bin_centers)

            r0_dict[det] = r0
            r1_dict[det] = r1

        return self.log_likelihood_from_r(r0_dict, r1_dict)

    def forward(
        self,
        templates_at_bin_centers: Dict[str, torch.Tensor],
        **kwargs,
    ) -> torch.Tensor:
        return self.log_likelihood(templates_at_bin_centers, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

__init__(data, psds, h0, frequencies, bin_edges=None, n_bins=50, f_min=None, f_max=None, df=None)

Initialize RelativeBinningLikelihood with precomputed data summary.

Parameters:

Name Type Description Default
data Dict[str, Tensor]

Dictionary mapping detector name (e.g. 'H1', 'L1') to dense frequency-domain data.

 required
psds Dict[str, Tensor]

Dictionary mapping detector name to PSD tensor S_n(f).

 required
h0 Dict[str, Tensor]

Dictionary mapping detector name to fiducial reference waveform h_0(f).

 required
frequencies Tensor

Dense frequency grid in Hz.

 required
bin_edges Optional[Tensor]

Optional precomputed bin edges (M + 1,).

 None
n_bins int

Number of bins if bin_edges is not provided.

 50
f_min Optional[float]

Minimum frequency cutoff in Hz (default: min of frequencies).

 None
f_max Optional[float]

Maximum frequency cutoff in Hz (default: max of frequencies).

 None
df Optional[float]

Dense frequency spacing in Hz.

 None
Source code in src/torchwave/inference/relative_binning.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def __init__(
    self,
    data: Dict[str, torch.Tensor],
    psds: Dict[str, torch.Tensor],
    h0: Dict[str, torch.Tensor],
    frequencies: torch.Tensor,
    bin_edges: Optional[torch.Tensor] = None,
    n_bins: int = 50,
    f_min: Optional[float] = None,
    f_max: Optional[float] = None,
    df: Optional[float] = None,
):
    """Initialize RelativeBinningLikelihood with precomputed data summary.

    Args:
        data: Dictionary mapping detector name (e.g. 'H1', 'L1') to dense frequency-domain data.
        psds: Dictionary mapping detector name to PSD tensor S_n(f).
        h0: Dictionary mapping detector name to fiducial reference waveform h_0(f).
        frequencies: Dense frequency grid in Hz.
        bin_edges: Optional precomputed bin edges (M + 1,).
        n_bins: Number of bins if bin_edges is not provided.
        f_min: Minimum frequency cutoff in Hz (default: min of frequencies).
        f_max: Maximum frequency cutoff in Hz (default: max of frequencies).
        df: Dense frequency spacing in Hz.
    """
    super().__init__()
    self.detectors = sorted(list(data.keys()))
    f_dense = as_tensor(frequencies)
    self.register_buffer("frequencies", f_dense, persistent=True)

    min_f = f_min if f_min is not None else float(f_dense[0])
    max_f = f_max if f_max is not None else float(f_dense[-1])

    if bin_edges is None:
        edges = compute_relative_binning_bins(
            f_min=min_f,
            f_max=max_f,
            n_bins=n_bins,
            frequencies=f_dense,
            device=f_dense.device,
            dtype=f_dense.dtype,
        )
    else:
        edges = as_tensor(bin_edges, device=f_dense.device, dtype=f_dense.dtype)

    self.register_buffer("bin_edges", edges, persistent=True)
    self.n_bins = len(edges) - 1

    # Use the same dense-grid frequencies for fiducial and candidate evaluations
    raw_centers = 0.5 * (edges[:-1] + edges[1:])
    center_indices = torch.argmin(torch.abs(f_dense.unsqueeze(1) - raw_centers.unsqueeze(0)), dim=0)
    centers = f_dense[center_indices]
    self.register_buffer("bin_centers", centers, persistent=True)
    self.register_buffer("center_indices", center_indices, persistent=False)

    # Precompute summary coefficients for each detector
    self.d_d: Dict[str, torch.Tensor] = {}
    for det in self.detectors:
        d_tensor = as_tensor(data[det], device=f_dense.device)
        psd_tensor = as_tensor(psds[det], device=f_dense.device)
        h0_tensor = as_tensor(h0[det], device=f_dense.device)

        a0, a1, b0, b1, _, d_d_val = compute_relative_binning_summary(
            data=d_tensor,
            psd=psd_tensor,
            h0=h0_tensor,
            frequencies=f_dense,
            bin_edges=edges,
            df=df,
        )

        # Re-evaluate a1 and b1 with exact snapped centers
        safe_psd = torch.clamp(psd_tensor, min=1.0e-50)
        df_val = float(df) if df is not None else float(f_dense[1] - f_dense[0])
        integrand_a = 4.0 * ((d_tensor * torch.conj(h0_tensor)) / safe_psd) * df_val
        integrand_b = 4.0 * (torch.real(h0_tensor * torch.conj(h0_tensor)) / safe_psd) * df_val

        a1_list = []
        b1_list = []
        for k in range(self.n_bins):
            f_low = edges[k]
            f_high = edges[k + 1]
            f_bar = centers[k]
            if k == self.n_bins - 1:
                mask = (f_dense >= f_low) & (f_dense <= f_high)
            else:
                mask = (f_dense >= f_low) & (f_dense < f_high)
            delta_f = f_dense[mask] - f_bar
            a1_list.append(torch.sum(integrand_a[mask] * delta_f))
            b1_list.append(torch.sum(integrand_b[mask] * delta_f))

        a1 = torch.stack(a1_list)
        b1 = torch.stack(b1_list)

        self.register_buffer(f"a0_{det}", a0, persistent=True)
        self.register_buffer(f"a1_{det}", a1, persistent=True)
        self.register_buffer(f"b0_{det}", b0, persistent=True)
        self.register_buffer(f"b1_{det}", b1, persistent=True)
        self.register_buffer(f"h0_centers_{det}", h0_tensor[center_indices], persistent=True)
        self.register_buffer(f"h0_dense_{det}", h0_tensor, persistent=False)
        self.register_buffer(f"dd_{det}", d_d_val, persistent=True)

compute_r0_r1(h_centers, h0_centers, bin_centers) staticmethod

Compute the ratio r_0 and slope r_1 at bin centers.

Parameters:

Name Type Description Default
h_centers Tensor

Test waveform evaluated at bin centers, shape (..., M).

 required
h0_centers Tensor

Reference waveform evaluated at bin centers, shape (M,) or (..., M).

 required
bin_centers Tensor

Bin center frequencies, shape (M,).

 required

Returns:

Type Description
Tuple[Tensor, Tensor]

Tuple of (r0, r1) with same batch dimensions as h_centers.
Source code in src/torchwave/inference/relative_binning.py 
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
@staticmethod
def compute_r0_r1(
    h_centers: torch.Tensor,
    h0_centers: torch.Tensor,
    bin_centers: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the ratio r_0 and slope r_1 at bin centers.

    Args:
        h_centers: Test waveform evaluated at bin centers, shape (..., M).
        h0_centers: Reference waveform evaluated at bin centers, shape (M,) or (..., M).
        bin_centers: Bin center frequencies, shape (M,).

    Returns:
        Tuple of (r0, r1) with same batch dimensions as h_centers.
    """
    safe_h0 = torch.where(
        torch.abs(h0_centers) < 1e-40,
        torch.ones_like(h0_centers) * 1e-40,
        h0_centers,
    )
    r0 = h_centers / safe_h0

    m = bin_centers.shape[-1]
    r1 = torch.zeros_like(r0)

    if m >= 3:
        df_interior = bin_centers[2:] - bin_centers[:-2]
        r1[..., 1:-1] = (r0[..., 2:] - r0[..., :-2]) / df_interior

        # Boundary 0
        df_0 = bin_centers[1] - bin_centers[0]
        r1[..., 0] = (r0[..., 1] - r0[..., 0]) / df_0

        # Boundary M - 1
        df_end = bin_centers[-1] - bin_centers[-2]
        r1[..., -1] = (r0[..., -1] - r0[..., -2]) / df_end
    elif m == 2:
        df_step = bin_centers[1] - bin_centers[0]
        diff = (r0[..., 1] - r0[..., 0]) / df_step
        r1[..., 0] = diff
        r1[..., 1] = diff

    return r0, r1

inner_product_dh(r0, r1, detector)

Compute noise-weighted inner product from summary coefficients.

= Re sum_{k=0}^{M-1} [ r_0^(f_bar_k) * A_{0, k} + r_1^(f_bar_k) * A_{1, k} ]

Source code in src/torchwave/inference/relative_binning.py 
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
def inner_product_dh(
    self,
    r0: torch.Tensor,
    r1: torch.Tensor,
    detector: str,
) -> torch.Tensor:
    """Compute noise-weighted inner product <d | h(theta)> from summary coefficients.

    <d | h> = Re sum_{k=0}^{M-1} [ r_0^*(f_bar_k) * A_{0, k} + r_1^*(f_bar_k) * A_{1, k} ]
    """
    a0 = getattr(self, f"a0_{detector}")
    a1 = getattr(self, f"a1_{detector}")

    r0_conj = torch.conj(r0)
    r1_conj = torch.conj(r1)

    term0 = r0_conj * a0
    term1 = r1_conj * a1

    sum_terms = torch.sum(term0 + term1, dim=-1)
    return torch.real(sum_terms)

inner_product_hh(r0, r1, detector)

Compute waveform self inner product from summary coefficients.

= sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]

Source code in src/torchwave/inference/relative_binning.py 
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
def inner_product_hh(
    self,
    r0: torch.Tensor,
    r1: torch.Tensor,
    detector: str,
) -> torch.Tensor:
    """Compute waveform self inner product <h(theta) | h(theta)> from summary coefficients.

    <h | h> = sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]
    """
    b0 = getattr(self, f"b0_{detector}")
    b1 = getattr(self, f"b1_{detector}")

    r0_sq = torch.real(r0 * torch.conj(r0))
    cross_term = 2.0 * torch.real(r0 * torch.conj(r1))

    sum_terms = torch.sum(r0_sq * b0 + cross_term * b1, dim=-1)
    return sum_terms

log_likelihood(templates_at_bin_centers, h0_at_bin_centers=None, slopes=None)

Compute relative binning log-likelihood from waveforms evaluated at bin centers.

Parameters:

Name Type Description Default
templates_at_bin_centers Dict[str, Tensor]

Dictionary mapping detector name to waveform tensor h(theta; f_bar) evaluated ONLY at bin centers.

 required
h0_at_bin_centers Optional[Dict[str, Tensor]]

Optional fiducial waveform h0(f_bar). If not provided, sampled from stored dense h0.

 None
slopes Optional[Dict[str, Tensor]]

Optional explicit slopes r_1(f_bar). If None, calculated via centered finite differences.

 None

Returns:

Type Description
Tensor

Scalar or batched log-likelihood tensor.
Source code in src/torchwave/inference/relative_binning.py 
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
def log_likelihood(
    self,
    templates_at_bin_centers: Dict[str, torch.Tensor],
    h0_at_bin_centers: Optional[Dict[str, torch.Tensor]] = None,
    slopes: Optional[Dict[str, torch.Tensor]] = None,
) -> torch.Tensor:
    """Compute relative binning log-likelihood from waveforms evaluated at bin centers.

    Args:
        templates_at_bin_centers: Dictionary mapping detector name to waveform
            tensor h(theta; f_bar) evaluated ONLY at bin centers.
        h0_at_bin_centers: Optional fiducial waveform h0(f_bar). If not provided,
            sampled from stored dense h0.
        slopes: Optional explicit slopes r_1(f_bar). If None, calculated via
            centered finite differences.

    Returns:
        Scalar or batched log-likelihood tensor.
    """
    r0_dict: Dict[str, torch.Tensor] = {}
    r1_dict: Dict[str, torch.Tensor] = {}

    for det in self.detectors:
        h_det = templates_at_bin_centers[det]
        if h0_at_bin_centers is not None:
            h0_det = h0_at_bin_centers[det]
        elif hasattr(self, f"h0_centers_{det}"):
            h0_det = getattr(self, f"h0_centers_{det}")
        else:
            h0_dense = getattr(self, f"h0_dense_{det}")
            f_dense = self.frequencies
            idx = torch.searchsorted(f_dense, self.bin_centers).clamp(max=len(f_dense) - 1)
            h0_det = h0_dense[idx]

        if slopes is not None and det in slopes:
            safe_h0 = torch.where(
                torch.abs(h0_det) < 1e-40,
                torch.ones_like(h0_det) * 1e-40,
                h0_det,
            )
            r0 = h_det / safe_h0
            r1 = slopes[det]
        else:
            r0, r1 = self.compute_r0_r1(h_det, h0_det, self.bin_centers)

        r0_dict[det] = r0
        r1_dict[det] = r1

    return self.log_likelihood_from_r(r0_dict, r1_dict)

log_likelihood_from_r(r0_dict, r1_dict)

Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

ln L = sum_{det} [ - 0.5 * - 0.5 * ]

Source code in src/torchwave/inference/relative_binning.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
def log_likelihood_from_r(
    self,
    r0_dict: Dict[str, torch.Tensor],
    r1_dict: Dict[str, torch.Tensor],
) -> torch.Tensor:
    """Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

    ln L = sum_{det} [ <d | h> - 0.5 * <h | h> - 0.5 * <d | d> ]
    """
    device = self.bin_centers.device
    total_ln_l = torch.tensor(0.0, dtype=torch.float64, device=device)

    for det in self.detectors:
        r0 = r0_dict[det]
        r1 = r1_dict[det]
        dh = self.inner_product_dh(r0, r1, det)
        hh = self.inner_product_hh(r0, r1, det)
        dd = getattr(self, f"dd_{det}")

        det_ln_l = dh - 0.5 * hh - 0.5 * dd
        total_ln_l = total_ln_l + det_ln_l

    return total_ln_l

torchwave.inference.mcmc.NUTS

Pure PyTorch No-U-Turn Sampler (NUTS) with Dual-Averaging Step-Size Adaptation.

Literature: Hoffman & Gelman (2014) Algorithms 3, 5, 6; Betancourt (2017).

Source code in src/torchwave/inference/mcmc.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
class NUTS:
    """Pure PyTorch No-U-Turn Sampler (NUTS) with Dual-Averaging Step-Size Adaptation.

    Literature: Hoffman & Gelman (2014) Algorithms 3, 5, 6; Betancourt (2017).
    """

    def __init__(
        self,
        log_prob_fn: Callable[[torch.Tensor], torch.Tensor],
        target_accept: float = 0.65,
        max_depth: int = 10,
        inv_mass: Optional[torch.Tensor] = None,
    ):
        """Initialize NUTS sampler.

        Args:
            log_prob_fn: Function mapping parameter tensor theta of shape (D,) to scalar
                differentiable log-posterior ln p(theta) = ln L(theta) + ln pi(theta).
            target_accept: Desired Metropolis acceptance probability delta (default 0.65).
            max_depth: Maximum binary tree recursion depth (default 10).
            inv_mass: Diagonal inverse mass matrix elements M^(-1), shape (D,).
        """
        self.log_prob_fn = log_prob_fn
        if not 0.0 < target_accept < 1.0:
            raise ValueError("target_accept must be between zero and one")
        if not isinstance(max_depth, int) or max_depth < 1:
            raise ValueError("max_depth must be a positive integer")
        self.target_accept = target_accept
        self.max_depth = max_depth
        if inv_mass is not None:
            inv_mass = torch.as_tensor(inv_mass).detach().clone()
            if inv_mass.ndim != 1 or not torch.isfinite(inv_mass).all() or not (inv_mass > 0).all():
                raise ValueError("inv_mass must be a finite positive diagonal vector")
        self.inv_mass = inv_mass

    def _metric_like(self, theta: torch.Tensor) -> Optional[torch.Tensor]:
        if self.inv_mass is None:
            return None
        if self.inv_mass.shape != theta.shape:
            raise ValueError("inv_mass must match the parameter vector shape")
        return self.inv_mass.to(device=theta.device, dtype=theta.dtype)

    def _sample_momentum(self, theta: torch.Tensor) -> torch.Tensor:
        """Draw p ~ N(0, M), consistently with K(p) = p^T M^-1 p / 2."""
        momentum = torch.randn_like(theta)
        metric = self._metric_like(theta)
        return momentum if metric is None else momentum / torch.sqrt(metric)

    @staticmethod
    def _finite_state(*values: torch.Tensor) -> bool:
        return all(bool(torch.isfinite(value).all()) for value in values)

    def value_and_grad(self, theta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute log-probability and its gradient with respect to theta using autograd."""
        theta_var = theta.detach().clone().requires_grad_(True)
        log_p = self.log_prob_fn(theta_var)
        if not log_p.requires_grad or log_p.grad_fn is None:
            grad_v = torch.zeros_like(theta)
        else:
            grad = torch.autograd.grad(log_p, theta_var, create_graph=False)[0]
            # Hamiltonian potential energy is V(theta) = -log_p(theta), so grad_V = -grad_log_p
            grad_v = -grad.detach()
        return log_p.detach(), grad_v

    def kinetic_energy(self, p: torch.Tensor) -> torch.Tensor:
        """Compute kinetic energy K(p) = 1/2 * p^T * M^(-1) * p."""
        metric = self._metric_like(p)
        if metric is None:
            return 0.5 * torch.sum(p * p)
        return 0.5 * torch.sum(p * metric * p)

    def find_reasonable_epsilon(
        self,
        theta0: torch.Tensor,
        log_p0: torch.Tensor,
        grad_v0: torch.Tensor,
    ) -> float:
        """Heuristic initialization of step size epsilon (Hoffman & Gelman 2011 Algorithm 4)."""
        eps = 1.0
        p0 = self._sample_momentum(theta0)
        h0 = -log_p0 + self.kinetic_energy(p0)
        metric = self._metric_like(theta0)

        def log_acceptance(step_size: float) -> float:
            theta_p, p_p, log_p_p, grad_p = leapfrog_step(
                theta0, p0, grad_v0, step_size, self.value_and_grad, metric
            )
            h_p = -log_p_p + self.kinetic_energy(p_p)
            if not self._finite_state(theta_p, p_p, log_p_p, grad_p, h_p):
                return -math.inf
            return float((h0 - h_p).item())

        log_accept = log_acceptance(eps)
        # Determine whether to increase or decrease epsilon
        direction = 1 if log_accept > math.log(0.5) else -1

        count = 0
        while direction * log_accept > direction * math.log(0.5) and count < 20:
            eps = eps * (2.0 ** direction)
            log_accept = log_acceptance(eps)
            count += 1

        return float(max(1e-6, min(10.0, eps)))

    def _build_tree(
        self,
        theta: torch.Tensor,
        p: torch.Tensor,
        grad: torch.Tensor,
        log_u: float,
        v: int,
        j: int,
        epsilon: float,
        h0: float,
        diagnostics: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, bool, float, int]:
        """Recursive binary tree expansion (Hoffman & Gelman 2011 Algorithm 6).

        Returns:
            theta_minus, p_minus, grad_minus, theta_plus, p_plus, grad_plus,
            theta_prime, grad_prime, log_p_prime, n_prime, s_prime, alpha_prime, n_alpha_prime
        """
        if j == 0:
            # Base case: take a single leapfrog step in direction v
            theta_prime, p_prime, log_p_prime, grad_prime = leapfrog_step(
                theta, p, grad, v * epsilon, self.value_and_grad, self._metric_like(theta)
            )
            h_prime = -log_p_prime + self.kinetic_energy(p_prime)
            h_val = float(h_prime.item())

            # Metropolis acceptance criteria
            finite = math.isfinite(h_val) and self._finite_state(theta_prime, p_prime, log_p_prime, grad_prime)
            n_prime = int(finite and log_u <= -h_val)
            # U-turn termination and divergence check (delta H > 1000)
            divergent = not finite or (h_val - h0) > 1000.0
            s_prime = (not divergent) and (log_u < (1000.0 - h_val))
            if diagnostics is not None:
                diagnostics["divergent"] = diagnostics["divergent"] or divergent
            alpha = math.exp(min(0.0, -h_val + h0)) if finite else 0.0
            return (
                theta_prime, p_prime, grad_prime,
                theta_prime, p_prime, grad_prime,
                theta_prime, grad_prime, log_p_prime, n_prime, s_prime, alpha, 1
            )

        # Recursion: build the first half of the tree
        (
            theta_m, p_m, grad_m, theta_p, p_p, grad_p, theta_prime, grad_prime, log_p_prime,
            n_prime, s_prime, alpha_prime, n_alpha_prime
        ) = self._build_tree(theta, p, grad, log_u, v, j - 1, epsilon, h0, diagnostics)

        if s_prime:
            if v == -1:
                (
                    theta_m, p_m, grad_m, _, _, _, theta_prime2, grad_prime2, log_p_prime2,
                    n_prime2, s_prime2, alpha_prime2, n_alpha_prime2
                ) = self._build_tree(theta_m, p_m, grad_m, log_u, v, j - 1, epsilon, h0, diagnostics)
            else:
                (
                    _, _, _, theta_p, p_p, grad_p, theta_prime2, grad_prime2, log_p_prime2,
                    n_prime2, s_prime2, alpha_prime2, n_alpha_prime2
                ) = self._build_tree(theta_p, p_p, grad_p, log_u, v, j - 1, epsilon, h0, diagnostics)

            # Progressive transition probability
            total_n = n_prime + n_prime2
            if total_n > 0 and (float(torch.rand(1).item()) < float(n_prime2) / float(total_n)):
                theta_prime = theta_prime2
                grad_prime = grad_prime2
                log_p_prime = log_p_prime2

            alpha_prime = alpha_prime + alpha_prime2
            n_alpha_prime = n_alpha_prime + n_alpha_prime2

            # Check U-turn termination on full span: (theta_plus - theta_minus) . p
            diff = theta_p - theta_m
            metric = self._metric_like(theta)
            if metric is not None:
                p_p_eff = metric * p_p
                p_m_eff = metric * p_m
            else:
                p_p_eff = p_p
                p_m_eff = p_m

            u_turn_forward = torch.sum(diff * p_p_eff).item() >= 0.0
            u_turn_backward = torch.sum(diff * p_m_eff).item() >= 0.0

            s_prime = s_prime2 and u_turn_forward and u_turn_backward
            n_prime = total_n

        return (
            theta_m, p_m, grad_m, theta_p, p_p, grad_p, theta_prime, grad_prime, log_p_prime,
            n_prime, s_prime, alpha_prime, n_alpha_prime
        )

    def sample_chain(
        self,
        theta0: torch.Tensor,
        n_samples: int = 400,
        n_warmup: int = 200,
        epsilon: Optional[float] = None,
        adapt_step_size: bool = True,
        return_diagnostics: bool = False,
        progress_callback: Optional[Callable[[int, int, Dict[str, Any]], None]] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
        """Run a single NUTS chain with dual-averaging step-size adaptation.

        Args:
            theta0: Initial parameter tensor of shape (D,).
            n_samples: Number of post-warmup samples to return.
            n_warmup: Number of warmup/adaptation steps.
            epsilon: Optional initial step size. If None, chosen heuristically.
            adapt_step_size: If True, adapt step size during warmup.
            return_diagnostics: Return (samples, diagnostics) when True. Each diagnostic
                tensor has n_warmup + n_samples entries, including warmup transitions.
            progress_callback: Optional callback(iteration, total_steps, diagnostics),
                called after every completed transition with one-based iteration.

        Returns:
            Tensor of shape (n_samples, D) containing post-warmup samples by default.
            Optional diagnostics: warmup, divergent, tree_depth, hit_max_depth,
            acceptance_stat (final doubling mean), step_size (used in transition),
            energy (initial Hamiltonian after refreshing momentum), and leapfrog_count.
            hit_max_depth records termination by the cap while expansion could continue.
        """
        device = theta0.device
        dtype = theta0.dtype
        if theta0.ndim != 1 or not theta0.is_floating_point():
            raise ValueError("theta0 must be a floating point parameter vector")
        if not isinstance(n_samples, int) or n_samples < 1 or not isinstance(n_warmup, int) or n_warmup < 0:
            raise ValueError("n_samples must be positive and n_warmup nonnegative integers")
        metric = self._metric_like(theta0)
        theta = theta0.detach().clone()

        log_p, grad_v = self.value_and_grad(theta)
        if not self._finite_state(theta, log_p, grad_v):
            raise ValueError("initial parameter, log probability and gradient must be finite")

        if epsilon is None:
            eps = self.find_reasonable_epsilon(theta, log_p, grad_v)
        else:
            eps = float(epsilon)
        if not math.isfinite(eps) or eps <= 0.0:
            raise ValueError("epsilon must be finite and positive")

        # Dual-averaging parameters (Hoffman & Gelman 2011 Algorithm 5)
        gamma_da = 0.05
        t0_da = 10.0
        kappa_da = 0.75
        mu_da = math.log(10.0 * eps)
        eps_bar = 1.0
        h_bar = 0.0

        total_steps = n_warmup + n_samples
        samples: List[torch.Tensor] = []
        diagnostic_rows: List[Dict[str, Any]] = []
        collect_diagnostics = return_diagnostics or progress_callback is not None

        for m in range(1, total_steps + 1):
            # Sample momentum: p ~ N(0, M)
            p0 = self._sample_momentum(theta)
            h0 = float((-log_p + self.kinetic_energy(p0)).item())

            # Slice sampling threshold: u ~ Uniform(0, exp(-H0)) => log(u) = -H0 - Exp(1)
            exp_rnd = float(torch.empty(1, device=device).exponential_(1.0).item())
            log_u = -h0 - exp_rnd

            # Initialize tree endpoints
            theta_minus = theta.clone()
            theta_plus = theta.clone()
            p_minus = p0.clone()
            p_plus = p0.clone()
            grad_minus = grad_v.clone()
            grad_plus = grad_v.clone()
            theta_next = theta.clone()
            grad_next = grad_v.clone()
            log_p_next = log_p

            j = 0
            n = 1
            s = True

            alpha_step = 0.0
            n_alpha_step = 0
            leapfrog_count = 0
            tree_diagnostics = {"divergent": False} if collect_diagnostics else None
            transition_epsilon = eps

            while s and j < self.max_depth:
                # Direction v ~ Uniform({-1, 1})
                v = 1 if torch.rand(1, device=device).item() > 0.5 else -1

                if v == -1:
                    (
                        theta_minus, p_minus, grad_minus, _, _, _,
                        theta_prime, grad_prime, log_p_prime, n_prime, s_prime,
                        alpha_step, n_alpha_step
                    ) = self._build_tree(
                        theta_minus, p_minus, grad_minus, log_u, v, j, eps, h0,
                        diagnostics=tree_diagnostics,
                    )
                else:
                    (
                        _, _, _, theta_plus, p_plus, grad_plus,
                        theta_prime, grad_prime, log_p_prime, n_prime, s_prime,
                        alpha_step, n_alpha_step
                    ) = self._build_tree(
                        theta_plus, p_plus, grad_plus, log_u, v, j, eps, h0,
                        diagnostics=tree_diagnostics,
                    )

                if s_prime:
                    # Algorithm 6's outer transition; recursive subtree selection uses
                    # n_prime2 / (n_prime + n_prime2), a distinct probability.
                    if n_prime > 0 and (float(torch.rand(1, device=device).item()) < min(1.0, float(n_prime) / float(n))):
                        theta_next = theta_prime
                        grad_next = grad_prime
                        log_p_next = log_p_prime

                n = n + n_prime
                leapfrog_count += n_alpha_step

                # Check U-turn condition on outermost endpoints
                diff = theta_plus - theta_minus
                if metric is not None:
                    p_plus_eff = metric * p_plus
                    p_minus_eff = metric * p_minus
                else:
                    p_plus_eff = p_plus
                    p_minus_eff = p_minus

                u_turn_fwd = torch.sum(diff * p_plus_eff).item() >= 0.0
                u_turn_bwd = torch.sum(diff * p_minus_eff).item() >= 0.0

                s = s_prime and u_turn_fwd and u_turn_bwd
                j = j + 1

            theta = theta_next
            grad_v = grad_next
            log_p = log_p_next
            # Section 3.2.4 averages the FINAL doubling, not all tree expansions.
            avg_alpha = alpha_step / max(1, n_alpha_step)

            if collect_diagnostics:
                row = {
                    "warmup": m <= n_warmup,
                    "divergent": tree_diagnostics["divergent"],
                    "tree_depth": j,
                    "hit_max_depth": bool(s and j == self.max_depth),
                    "acceptance_stat": avg_alpha,
                    "step_size": transition_epsilon,
                    "energy": h0,
                    "leapfrog_count": leapfrog_count,
                }
                if return_diagnostics:
                    diagnostic_rows.append(row.copy())
                if progress_callback is not None:
                    progress_callback(m, total_steps, row.copy())

            # Dual averaging adaptation during warmup
            if m <= n_warmup and adapt_step_size:
                eta = 1.0 / (float(m) + t0_da)
                h_bar = (1.0 - eta) * h_bar + eta * (self.target_accept - avg_alpha)
                log_eps = mu_da - (math.sqrt(float(m)) / gamma_da) * h_bar
                eps = math.exp(log_eps)
                m_pow = float(m) ** (-kappa_da)
                eps_bar = math.exp(m_pow * log_eps + (1.0 - m_pow) * math.log(eps_bar))
                if m == n_warmup:
                    eps = eps_bar

            if m > n_warmup:
                samples.append(theta.detach().clone())

        result = torch.stack(samples)
        if not return_diagnostics:
            return result
        diagnostic_types = {
            "warmup": torch.bool, "divergent": torch.bool, "tree_depth": torch.int64,
            "hit_max_depth": torch.bool, "acceptance_stat": dtype, "step_size": dtype,
            "energy": dtype, "leapfrog_count": torch.int64,
        }
        diagnostics = {
            key: torch.tensor([row[key] for row in diagnostic_rows], device=device, dtype=value_dtype)
            for key, value_dtype in diagnostic_types.items()
        }
        return result, diagnostics

    def sample_parallel(
        self,
        initial_positions: torch.Tensor,
        n_samples: int = 400,
        n_warmup: int = 200,
        **chain_kwargs: Any,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
        """Run multiple NUTS chains sequentially from given initial positions.

        Args:
            initial_positions: Initial positions for each chain, shape (n_chains, D).
            n_samples: Number of post-warmup samples per chain.
            n_warmup: Number of warmup samples per chain.
            chain_kwargs: Arbitrary keyword arguments passed to sample_chain.

        Returns:
            Tensor of samples of shape (n_chains, n_samples, D) by default.
            With return_diagnostics=True, (samples, diagnostics), with each diagnostic
            tensor shaped (n_chains, n_warmup + n_samples).
        """
        n_chains = initial_positions.shape[0]
        chain_outputs = []
        chain_diagnostics = []

        for c in range(n_chains):
            theta0 = initial_positions[c]
            chain_samples = self.sample_chain(
                theta0,
                n_samples=n_samples,
                n_warmup=n_warmup,
                **chain_kwargs,
            )
            if chain_kwargs.get("return_diagnostics", False):
                samples, diagnostics = chain_samples
                chain_outputs.append(samples)
                chain_diagnostics.append(diagnostics)
            else:
                chain_outputs.append(chain_samples)

        result = torch.stack(chain_outputs, dim=0)
        if chain_diagnostics:
            return result, {
                key: torch.stack([diagnostics[key] for diagnostics in chain_diagnostics])
                for key in chain_diagnostics[0]
            }
        return result

    def __call__(self, *args, **kwargs):
        return self.sample_parallel(*args, **kwargs)

__init__(log_prob_fn, target_accept=0.65, max_depth=10, inv_mass=None)

Initialize NUTS sampler.

Parameters:

Name Type Description Default
log_prob_fn Callable[[Tensor], Tensor]

Function mapping parameter tensor theta of shape (D,) to scalar differentiable log-posterior ln p(theta) = ln L(theta) + ln pi(theta).

 required
target_accept float

Desired Metropolis acceptance probability delta (default 0.65).

 0.65
max_depth int

Maximum binary tree recursion depth (default 10).

 10
inv_mass Optional[Tensor]

Diagonal inverse mass matrix elements M^(-1), shape (D,).

 None
Source code in src/torchwave/inference/mcmc.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def __init__(
    self,
    log_prob_fn: Callable[[torch.Tensor], torch.Tensor],
    target_accept: float = 0.65,
    max_depth: int = 10,
    inv_mass: Optional[torch.Tensor] = None,
):
    """Initialize NUTS sampler.

    Args:
        log_prob_fn: Function mapping parameter tensor theta of shape (D,) to scalar
            differentiable log-posterior ln p(theta) = ln L(theta) + ln pi(theta).
        target_accept: Desired Metropolis acceptance probability delta (default 0.65).
        max_depth: Maximum binary tree recursion depth (default 10).
        inv_mass: Diagonal inverse mass matrix elements M^(-1), shape (D,).
    """
    self.log_prob_fn = log_prob_fn
    if not 0.0 < target_accept < 1.0:
        raise ValueError("target_accept must be between zero and one")
    if not isinstance(max_depth, int) or max_depth < 1:
        raise ValueError("max_depth must be a positive integer")
    self.target_accept = target_accept
    self.max_depth = max_depth
    if inv_mass is not None:
        inv_mass = torch.as_tensor(inv_mass).detach().clone()
        if inv_mass.ndim != 1 or not torch.isfinite(inv_mass).all() or not (inv_mass > 0).all():
            raise ValueError("inv_mass must be a finite positive diagonal vector")
    self.inv_mass = inv_mass

find_reasonable_epsilon(theta0, log_p0, grad_v0)

Heuristic initialization of step size epsilon (Hoffman & Gelman 2011 Algorithm 4).

Source code in src/torchwave/inference/mcmc.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def find_reasonable_epsilon(
    self,
    theta0: torch.Tensor,
    log_p0: torch.Tensor,
    grad_v0: torch.Tensor,
) -> float:
    """Heuristic initialization of step size epsilon (Hoffman & Gelman 2011 Algorithm 4)."""
    eps = 1.0
    p0 = self._sample_momentum(theta0)
    h0 = -log_p0 + self.kinetic_energy(p0)
    metric = self._metric_like(theta0)

    def log_acceptance(step_size: float) -> float:
        theta_p, p_p, log_p_p, grad_p = leapfrog_step(
            theta0, p0, grad_v0, step_size, self.value_and_grad, metric
        )
        h_p = -log_p_p + self.kinetic_energy(p_p)
        if not self._finite_state(theta_p, p_p, log_p_p, grad_p, h_p):
            return -math.inf
        return float((h0 - h_p).item())

    log_accept = log_acceptance(eps)
    # Determine whether to increase or decrease epsilon
    direction = 1 if log_accept > math.log(0.5) else -1

    count = 0
    while direction * log_accept > direction * math.log(0.5) and count < 20:
        eps = eps * (2.0 ** direction)
        log_accept = log_acceptance(eps)
        count += 1

    return float(max(1e-6, min(10.0, eps)))

kinetic_energy(p)

Compute kinetic energy K(p) = 1/2 * p^T * M^(-1) * p.

Source code in src/torchwave/inference/mcmc.py 
163
164
165
166
167
168
def kinetic_energy(self, p: torch.Tensor) -> torch.Tensor:
    """Compute kinetic energy K(p) = 1/2 * p^T * M^(-1) * p."""
    metric = self._metric_like(p)
    if metric is None:
        return 0.5 * torch.sum(p * p)
    return 0.5 * torch.sum(p * metric * p)

sample_chain(theta0, n_samples=400, n_warmup=200, epsilon=None, adapt_step_size=True, return_diagnostics=False, progress_callback=None)

Run a single NUTS chain with dual-averaging step-size adaptation.

Parameters:

Name Type Description Default
theta0 Tensor

Initial parameter tensor of shape (D,).

 required
n_samples int

Number of post-warmup samples to return.

 400
n_warmup int

Number of warmup/adaptation steps.

 200
epsilon Optional[float]

Optional initial step size. If None, chosen heuristically.

 None
adapt_step_size bool

If True, adapt step size during warmup.

 True
return_diagnostics bool

Return (samples, diagnostics) when True. Each diagnostic tensor has n_warmup + n_samples entries, including warmup transitions.

 False
progress_callback Optional[Callable[[int, int, Dict[str, Any]], None]]

Optional callback(iteration, total_steps, diagnostics), called after every completed transition with one-based iteration.

 None

Returns:

Type Description
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

Tensor of shape (n_samples, D) containing post-warmup samples by default.
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

Optional diagnostics: warmup, divergent, tree_depth, hit_max_depth,
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

acceptance_stat (final doubling mean), step_size (used in transition),
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

energy (initial Hamiltonian after refreshing momentum), and leapfrog_count.
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

hit_max_depth records termination by the cap while expansion could continue.
Source code in src/torchwave/inference/mcmc.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
def sample_chain(
    self,
    theta0: torch.Tensor,
    n_samples: int = 400,
    n_warmup: int = 200,
    epsilon: Optional[float] = None,
    adapt_step_size: bool = True,
    return_diagnostics: bool = False,
    progress_callback: Optional[Callable[[int, int, Dict[str, Any]], None]] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
    """Run a single NUTS chain with dual-averaging step-size adaptation.

    Args:
        theta0: Initial parameter tensor of shape (D,).
        n_samples: Number of post-warmup samples to return.
        n_warmup: Number of warmup/adaptation steps.
        epsilon: Optional initial step size. If None, chosen heuristically.
        adapt_step_size: If True, adapt step size during warmup.
        return_diagnostics: Return (samples, diagnostics) when True. Each diagnostic
            tensor has n_warmup + n_samples entries, including warmup transitions.
        progress_callback: Optional callback(iteration, total_steps, diagnostics),
            called after every completed transition with one-based iteration.

    Returns:
        Tensor of shape (n_samples, D) containing post-warmup samples by default.
        Optional diagnostics: warmup, divergent, tree_depth, hit_max_depth,
        acceptance_stat (final doubling mean), step_size (used in transition),
        energy (initial Hamiltonian after refreshing momentum), and leapfrog_count.
        hit_max_depth records termination by the cap while expansion could continue.
    """
    device = theta0.device
    dtype = theta0.dtype
    if theta0.ndim != 1 or not theta0.is_floating_point():
        raise ValueError("theta0 must be a floating point parameter vector")
    if not isinstance(n_samples, int) or n_samples < 1 or not isinstance(n_warmup, int) or n_warmup < 0:
        raise ValueError("n_samples must be positive and n_warmup nonnegative integers")
    metric = self._metric_like(theta0)
    theta = theta0.detach().clone()

    log_p, grad_v = self.value_and_grad(theta)
    if not self._finite_state(theta, log_p, grad_v):
        raise ValueError("initial parameter, log probability and gradient must be finite")

    if epsilon is None:
        eps = self.find_reasonable_epsilon(theta, log_p, grad_v)
    else:
        eps = float(epsilon)
    if not math.isfinite(eps) or eps <= 0.0:
        raise ValueError("epsilon must be finite and positive")

    # Dual-averaging parameters (Hoffman & Gelman 2011 Algorithm 5)
    gamma_da = 0.05
    t0_da = 10.0
    kappa_da = 0.75
    mu_da = math.log(10.0 * eps)
    eps_bar = 1.0
    h_bar = 0.0

    total_steps = n_warmup + n_samples
    samples: List[torch.Tensor] = []
    diagnostic_rows: List[Dict[str, Any]] = []
    collect_diagnostics = return_diagnostics or progress_callback is not None

    for m in range(1, total_steps + 1):
        # Sample momentum: p ~ N(0, M)
        p0 = self._sample_momentum(theta)
        h0 = float((-log_p + self.kinetic_energy(p0)).item())

        # Slice sampling threshold: u ~ Uniform(0, exp(-H0)) => log(u) = -H0 - Exp(1)
        exp_rnd = float(torch.empty(1, device=device).exponential_(1.0).item())
        log_u = -h0 - exp_rnd

        # Initialize tree endpoints
        theta_minus = theta.clone()
        theta_plus = theta.clone()
        p_minus = p0.clone()
        p_plus = p0.clone()
        grad_minus = grad_v.clone()
        grad_plus = grad_v.clone()
        theta_next = theta.clone()
        grad_next = grad_v.clone()
        log_p_next = log_p

        j = 0
        n = 1
        s = True

        alpha_step = 0.0
        n_alpha_step = 0
        leapfrog_count = 0
        tree_diagnostics = {"divergent": False} if collect_diagnostics else None
        transition_epsilon = eps

        while s and j < self.max_depth:
            # Direction v ~ Uniform({-1, 1})
            v = 1 if torch.rand(1, device=device).item() > 0.5 else -1

            if v == -1:
                (
                    theta_minus, p_minus, grad_minus, _, _, _,
                    theta_prime, grad_prime, log_p_prime, n_prime, s_prime,
                    alpha_step, n_alpha_step
                ) = self._build_tree(
                    theta_minus, p_minus, grad_minus, log_u, v, j, eps, h0,
                    diagnostics=tree_diagnostics,
                )
            else:
                (
                    _, _, _, theta_plus, p_plus, grad_plus,
                    theta_prime, grad_prime, log_p_prime, n_prime, s_prime,
                    alpha_step, n_alpha_step
                ) = self._build_tree(
                    theta_plus, p_plus, grad_plus, log_u, v, j, eps, h0,
                    diagnostics=tree_diagnostics,
                )

            if s_prime:
                # Algorithm 6's outer transition; recursive subtree selection uses
                # n_prime2 / (n_prime + n_prime2), a distinct probability.
                if n_prime > 0 and (float(torch.rand(1, device=device).item()) < min(1.0, float(n_prime) / float(n))):
                    theta_next = theta_prime
                    grad_next = grad_prime
                    log_p_next = log_p_prime

            n = n + n_prime
            leapfrog_count += n_alpha_step

            # Check U-turn condition on outermost endpoints
            diff = theta_plus - theta_minus
            if metric is not None:
                p_plus_eff = metric * p_plus
                p_minus_eff = metric * p_minus
            else:
                p_plus_eff = p_plus
                p_minus_eff = p_minus

            u_turn_fwd = torch.sum(diff * p_plus_eff).item() >= 0.0
            u_turn_bwd = torch.sum(diff * p_minus_eff).item() >= 0.0

            s = s_prime and u_turn_fwd and u_turn_bwd
            j = j + 1

        theta = theta_next
        grad_v = grad_next
        log_p = log_p_next
        # Section 3.2.4 averages the FINAL doubling, not all tree expansions.
        avg_alpha = alpha_step / max(1, n_alpha_step)

        if collect_diagnostics:
            row = {
                "warmup": m <= n_warmup,
                "divergent": tree_diagnostics["divergent"],
                "tree_depth": j,
                "hit_max_depth": bool(s and j == self.max_depth),
                "acceptance_stat": avg_alpha,
                "step_size": transition_epsilon,
                "energy": h0,
                "leapfrog_count": leapfrog_count,
            }
            if return_diagnostics:
                diagnostic_rows.append(row.copy())
            if progress_callback is not None:
                progress_callback(m, total_steps, row.copy())

        # Dual averaging adaptation during warmup
        if m <= n_warmup and adapt_step_size:
            eta = 1.0 / (float(m) + t0_da)
            h_bar = (1.0 - eta) * h_bar + eta * (self.target_accept - avg_alpha)
            log_eps = mu_da - (math.sqrt(float(m)) / gamma_da) * h_bar
            eps = math.exp(log_eps)
            m_pow = float(m) ** (-kappa_da)
            eps_bar = math.exp(m_pow * log_eps + (1.0 - m_pow) * math.log(eps_bar))
            if m == n_warmup:
                eps = eps_bar

        if m > n_warmup:
            samples.append(theta.detach().clone())

    result = torch.stack(samples)
    if not return_diagnostics:
        return result
    diagnostic_types = {
        "warmup": torch.bool, "divergent": torch.bool, "tree_depth": torch.int64,
        "hit_max_depth": torch.bool, "acceptance_stat": dtype, "step_size": dtype,
        "energy": dtype, "leapfrog_count": torch.int64,
    }
    diagnostics = {
        key: torch.tensor([row[key] for row in diagnostic_rows], device=device, dtype=value_dtype)
        for key, value_dtype in diagnostic_types.items()
    }
    return result, diagnostics

sample_parallel(initial_positions, n_samples=400, n_warmup=200, **chain_kwargs)

Run multiple NUTS chains sequentially from given initial positions.

Parameters:

Name Type Description Default
initial_positions Tensor

Initial positions for each chain, shape (n_chains, D).

 required
n_samples int

Number of post-warmup samples per chain.

 400
n_warmup int

Number of warmup samples per chain.

 200
chain_kwargs Any

Arbitrary keyword arguments passed to sample_chain.

 {}

Returns:

Type Description
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

Tensor of samples of shape (n_chains, n_samples, D) by default.
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

With return_diagnostics=True, (samples, diagnostics), with each diagnostic
Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]

tensor shaped (n_chains, n_warmup + n_samples).
Source code in src/torchwave/inference/mcmc.py 
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def sample_parallel(
    self,
    initial_positions: torch.Tensor,
    n_samples: int = 400,
    n_warmup: int = 200,
    **chain_kwargs: Any,
) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
    """Run multiple NUTS chains sequentially from given initial positions.

    Args:
        initial_positions: Initial positions for each chain, shape (n_chains, D).
        n_samples: Number of post-warmup samples per chain.
        n_warmup: Number of warmup samples per chain.
        chain_kwargs: Arbitrary keyword arguments passed to sample_chain.

    Returns:
        Tensor of samples of shape (n_chains, n_samples, D) by default.
        With return_diagnostics=True, (samples, diagnostics), with each diagnostic
        tensor shaped (n_chains, n_warmup + n_samples).
    """
    n_chains = initial_positions.shape[0]
    chain_outputs = []
    chain_diagnostics = []

    for c in range(n_chains):
        theta0 = initial_positions[c]
        chain_samples = self.sample_chain(
            theta0,
            n_samples=n_samples,
            n_warmup=n_warmup,
            **chain_kwargs,
        )
        if chain_kwargs.get("return_diagnostics", False):
            samples, diagnostics = chain_samples
            chain_outputs.append(samples)
            chain_diagnostics.append(diagnostics)
        else:
            chain_outputs.append(chain_samples)

    result = torch.stack(chain_outputs, dim=0)
    if chain_diagnostics:
        return result, {
            key: torch.stack([diagnostics[key] for diagnostics in chain_diagnostics])
            for key in chain_diagnostics[0]
        }
    return result

value_and_grad(theta)

Compute log-probability and its gradient with respect to theta using autograd.

Source code in src/torchwave/inference/mcmc.py 
151
152
153
154
155
156
157
158
159
160
161
def value_and_grad(self, theta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute log-probability and its gradient with respect to theta using autograd."""
    theta_var = theta.detach().clone().requires_grad_(True)
    log_p = self.log_prob_fn(theta_var)
    if not log_p.requires_grad or log_p.grad_fn is None:
        grad_v = torch.zeros_like(theta)
    else:
        grad = torch.autograd.grad(log_p, theta_var, create_graph=False)[0]
        # Hamiltonian potential energy is V(theta) = -log_p(theta), so grad_V = -grad_log_p
        grad_v = -grad.detach()
    return log_p.detach(), grad_v

torchwave.inference.mcmc.gelman_rubin(chains)

Compute the Gelman-Rubin potential scale reduction factor R_hat for multiple MCMC chains.

Literature: Gelman & Rubin (1992); Hoffman & Gelman (2014).

Parameters:

Name Type Description Default
chains Tensor

Tensor of MCMC samples of shape (n_chains, n_samples, n_params). Must have n_chains >= 2 and n_samples >= 4.

 required

Returns:

Type Description
Tensor

Tensor of shape (n_params,) containing R_hat for each parameter.
Tensor

Values near 1 are one diagnostic and do not establish convergence.
Tensor

This implementation uses the basic, unsplit variance-ratio statistic.
Source code in src/torchwave/inference/mcmc.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def gelman_rubin(chains: torch.Tensor) -> torch.Tensor:
    """Compute the Gelman-Rubin potential scale reduction factor R_hat for multiple MCMC chains.

    Literature: Gelman & Rubin (1992); Hoffman & Gelman (2014).

    Args:
        chains: Tensor of MCMC samples of shape (n_chains, n_samples, n_params).
            Must have n_chains >= 2 and n_samples >= 4.

    Returns:
        Tensor of shape (n_params,) containing R_hat for each parameter.
        Values near 1 are one diagnostic and do not establish convergence.
        This implementation uses the basic, unsplit variance-ratio statistic.
    """
    if chains.dim() != 3:
        raise ValueError(f"Expected chains tensor of shape (n_chains, n_samples, n_params), got {chains.shape}")

    m_chains, n_samples, n_params = chains.shape
    if m_chains < 2:
        raise ValueError(f"Need at least 2 chains for Gelman-Rubin diagnostic, got {m_chains}")
    if n_samples < 4:
        raise ValueError(f"Need at least 4 samples per chain, got {n_samples}")

    m = float(m_chains)
    n = float(n_samples)

    # Chain means: shape (m_chains, n_params)
    chain_means = torch.mean(chains, dim=1)
    # Overall mean: shape (n_params,)
    overall_mean = torch.mean(chain_means, dim=0)

    # Between-chain variance B / n: variance of the chain means
    b_over_n = torch.sum((chain_means - overall_mean) ** 2, dim=0) / (m - 1.0)
    b = b_over_n * n

    # Within-chain sample variances s_m^2: shape (m_chains, n_params)
    diffs = chains - chain_means.unsqueeze(1)
    s_m_sq = torch.sum(diffs ** 2, dim=1) / (n - 1.0)

    # Average within-chain variance W: shape (n_params,)
    w = torch.mean(s_m_sq, dim=0)
    w_safe = torch.clamp(w, min=1e-30)

    # Pooled variance estimate: var_plus = (n - 1) / n * W + 1 / n * B
    var_plus = ((n - 1.0) / n) * w + (1.0 / n) * b

    # Potential scale reduction factor R_hat = sqrt(var_plus / W)
    r_hat = torch.sqrt(torch.clamp(var_plus / w_safe, min=1.0))
    return r_hat