Improve omgw0 structuring and const_cs_gw

examples/const_cs_gw.py

@@ -14,7 +14,7 @@
 from pttools.bubble import lorentz
 from pttools.bubble.shock import shock_curve
 from pttools.models import ConstCSModel
-from pttools.ssmtools import Spectrum
+from pttools.omgw0 import Spectrum, omega_noise, omega_ins
 from pttools.analysis.parallel import create_spectra
 from pttools.analysis.utils import A3_PAPER_SIZE
 
@@ -25,6 +25,7 @@ def main():
     a_s = 5
     a_b = 1
     V_s = 1
+    r_star = 0.1
     v_walls: np.ndarray = np.array([0.44, 0.56, 0.92])
     alpha_ns: np.ndarray = np.array([0.07, 0.2])
     alpha_n_min = np.min(alpha_ns)
@@ -50,23 +51,30 @@ def main():
         spectra[i_model, :, :] = create_spectra(
             model=model, v_walls=v_walls, alpha_ns=alpha_ns,
             # spectrum_kwargs={"source_duration": 1},
-            # spectrum_kwargs={"z": z}
+            spectrum_kwargs={
+                "r_star": r_star
+                # "z": z
+            }
             # bubble_kwargs={"allow_invalid": False}, allow_bubble_failure=True
         )
 
     fig: plt.Figure = plt.figure(figsize=A3_PAPER_SIZE)
     fig2: plt.Figure = plt.figure(figsize=A3_PAPER_SIZE)
+    fig3: plt.Figure = plt.figure(figsize=A3_PAPER_SIZE)
     axs: np.ndarray = fig.subplots(alpha_ns.size, v_walls.size)
     axs2: np.ndarray = fig2.subplots(alpha_ns.size, v_walls.size)
+    axs3: np.ndarray = fig3.subplots(alpha_ns.size, v_walls.size)
     for i_alpha_n, alpha_n in enumerate(alpha_ns):
         for i_v_wall, v_wall in enumerate(v_walls):
             ax: plt.Axes = axs[i_alpha_n, i_v_wall]
             ax2: plt.Axes = axs2[i_alpha_n, i_v_wall]
+            ax3: plt.Axes = axs3[i_alpha_n, i_v_wall]
             for i_model, model in enumerate(models):
                 spectrum: Spectrum = spectra[i_model, i_alpha_n, i_v_wall]
                 if spectrum is not None:
                     ax.plot(spectrum.y, spectrum.pow_gw, label=model.label_latex)
                     ax2.plot(spectrum.bubble.xi, spectrum.bubble.v, label=model.label_latex)
+                    ax3.plot(spectrum.f(), spectrum.omgw0(), label=model.label_latex)
             ax.set_xscale("log")
             ax.set_yscale("log")
             ax.set_xlabel("$z = kR*$")
@@ -80,6 +88,13 @@ def main():
             ax2.grid()
             ax2.set_title(title)
 
+            ax3.set_xscale("log")
+            ax3.set_yscale("log")
+            ax3.set_xlabel(r"$f(\text{Hz})$")
+            ax3.set_ylabel(r"$\Omega$")
+            ax3.grid()
+            ax3.set_title(title + rf", r_*={r_star}")
+
     # Mu curves
     for csb2 in csb2s:
         csb = np.sqrt(csb2)
@@ -102,6 +117,16 @@ def main():
                 ax = axs2[i_alpha_n, i_v_wall]
                 ax.plot(xi_arr, vm_arr, color="k")
 
+    # Noise curves
+    f_min = np.min([spectrum.f(z=spectrum.y[0]) for spectrum in spectra.flat])
+    f_max = np.max([spectrum.f(z=spectrum.y[-1]) for spectrum in spectra.flat])
+    f: np.ndarray = np.logspace(np.log10(f_min), np.log10(f_max), num=50)
+    for i_alpha_n, alpha_n in enumerate(alpha_ns):
+        for i_v_wall, v_wall in enumerate(v_walls):
+            ax = axs3[i_alpha_n, i_v_wall]
+            ax.plot(f, omega_noise(f), label="LISA overall noise")
+            ax.plot(f, omega_ins(f), label="LISA instrument noise")
+
     # Lines
     pow_low = 9
     k_low = np.logspace(-1, -0.2, 10)
@@ -124,8 +149,10 @@ def main():
 
     fig.tight_layout()
     fig2.tight_layout()
+    fig3.tight_layout()
     utils.save_and_show(fig, "const_cs_gw.png")
     utils.save_and_show(fig2, "const_cs_gw_v.png")
+    utils.save_and_show(fig3, "const_cs_gw_omgw0.png")
 
 
 if __name__ == "__main__":

examples/plot_const_cs.py

@@ -12,7 +12,7 @@
 from pttools.bubble import Bubble
 from pttools.models import ConstCSModel
 from pttools.analysis.plot_model import ModelPlot
-from pttools.ssmtools import Spectrum, NucType
+from pttools.ssmtools import SSMSpectrum, NucType
 
 
 def main():
@@ -26,7 +26,7 @@ def main():
     bubble_fig = bubble.plot()
     save(bubble_fig, "const_cs_bubble.png")
 
-    spectrum = Spectrum(bubble, nuc_type=NucType.EXPONENTIAL)
+    spectrum = SSMSpectrum(bubble, nuc_type=NucType.EXPONENTIAL)
     spectrum_fig = spectrum.plot_multi()
     save(spectrum_fig, "const_cs_spectrum.png")
 

examples/plot_old_new.py

@@ -17,7 +17,7 @@
 from pttools.bubble import relativity
 from pttools.models.model import Model
 from pttools.models.bag import BagModel
-from pttools.ssmtools.spectrum import Spectrum, power_gw_scaled_bag, spec_den_v_bag, power_v_bag
+from pttools.ssmtools.spectrum import SSMSpectrum, power_gw_scaled_bag, spec_den_v_bag, power_v_bag
 from tests.paper.plane import xiv_plane
 from tests.paper.plot_plane import plot_plane
 
@@ -80,7 +80,7 @@ def main():
     sol_types = [SolutionType.SUB_DEF, SolutionType.HYBRID, SolutionType.DETON]
 
     spectra = [
-        Spectrum(
+        SSMSpectrum(
             Bubble(bag, v_wall=v_walls[i], alpha_n=alpha_ns[i], sol_type=sol_types[i])
         )
         for i in range(len(v_walls))

pttools/analysis/parallel.py

@@ -6,10 +6,9 @@
 import numpy as np
 
 from pttools.bubble.bubble import Bubble
-from pttools.bubble import fluid_reference, shock_curve
+from pttools.bubble import fluid_reference
 from pttools.bubble.integrate import precompile
-from pttools.ssmtools.spectrum import DEFAULT_NUC_TYPE, NucType, Spectrum
-from pttools.ssmtools.const import Z_ST_THRESH
+from pttools.omgw0.omgw0_ssm import Spectrum
 from pttools.speedup import options
 from pttools.speedup import parallel
 if tp.TYPE_CHECKING:

pttools/analysis/plot_spectra.py

@@ -2,10 +2,10 @@
 
 import matplotlib.pyplot as plt
 
-from pttools.ssmtools.spectrum import Spectrum
+from pttools.ssmtools.spectrum import SSMSpectrum
 
 
-def plot_spectra(spectra: tp.List[Spectrum], fig: plt.Figure = None):
+def plot_spectra(spectra: tp.List[SSMSpectrum], fig: plt.Figure = None):
     if fig is None:
         fig = plt.figure(figsize=(11.7, 8.3))
 

pttools/analysis/plot_spectrum.py

@@ -2,24 +2,24 @@
 import numpy as np
 
 from pttools.analysis.utils import A4_PAPER_SIZE, FigAndAxes, create_fig_ax
-from pttools.ssmtools.spectrum import Spectrum
+from pttools.ssmtools.spectrum import SSMSpectrum
 
 
 # TODO: set proper labels for these
 
 def plot_spectrum(
-        spectrum: Spectrum,
+        spectrum: SSMSpectrum,
         fig: plt.Figure = None,
         ax: plt.Axes = None,
         path: str = None,
         **kwargs) -> FigAndAxes:
     fig, ax = create_fig_ax(fig, ax)
     ax.plot(spectrum.y, spectrum.pow_gw, **kwargs)
     ax.set_ylabel("pow_gw")
-    return plot_spectrum_common(spectrum, fig, ax, path)
+    return plot_spectrum_common(SSMSpectrum, fig, ax, path)
 
 
-def plot_spectrum_common(spectrum: Spectrum, fig: plt.Figure, ax: plt.Axes, path: str = None) -> FigAndAxes:
+def plot_spectrum_common(spectrum: SSMSpectrum, fig: plt.Figure, ax: plt.Axes, path: str = None) -> FigAndAxes:
     ax.set_xlabel("$z$")
     ax.set_xlim(np.min(spectrum.y), np.max(spectrum.y))
     ax.set_xscale("log")
@@ -30,51 +30,51 @@ def plot_spectrum_common(spectrum: Spectrum, fig: plt.Figure, ax: plt.Axes, path
     return fig, ax
 
 
-def plot_spectrum_multi(spectrum: Spectrum, fig: plt.Figure = None, path: str = None, **kwargs) -> plt.Figure:
+def plot_spectrum_multi(spectrum: SSMSpectrum, fig: plt.Figure = None, path: str = None, **kwargs) -> plt.Figure:
     if fig is None:
         fig = plt.figure(figsize=A4_PAPER_SIZE)
     axs = fig.subplots(2, 2)
-    plot_spectrum_spec_den_v(spectrum, fig, axs[0, 0], **kwargs)
-    plot_spectrum_spec_den_gw(spectrum, fig, axs[1, 0], **kwargs)
-    plot_spectrum_v(spectrum, fig, axs[0, 1], **kwargs)
-    plot_spectrum(spectrum, fig, axs[1, 1], **kwargs)
+    plot_spectrum_spec_den_v(SSMSpectrum, fig, axs[0, 0], **kwargs)
+    plot_spectrum_spec_den_gw(SSMSpectrum, fig, axs[1, 0], **kwargs)
+    plot_spectrum_v(SSMSpectrum, fig, axs[0, 1], **kwargs)
+    plot_spectrum(SSMSpectrum, fig, axs[1, 1], **kwargs)
     fig.tight_layout()
     if path is not None:
         fig.savefig(path)
     return fig
 
 
 def plot_spectrum_v(
-        spectrum: Spectrum,
+        spectrum: SSMSpectrum,
         fig: plt.Figure = None,
         ax: plt.Axes = None,
         path: str = None,
         **kwargs) -> FigAndAxes:
     fig, ax = create_fig_ax(fig, ax)
     ax.plot(spectrum.y, spectrum.pow_v, **kwargs)
     ax.set_ylabel("pow_v")
-    return plot_spectrum_common(spectrum, fig, ax, path)
+    return plot_spectrum_common(SSMSpectrum, fig, ax, path)
 
 
 def plot_spectrum_spec_den_gw(
-        spectrum: Spectrum,
+        spectrum: SSMSpectrum,
         fig: plt.Figure = None,
         ax: plt.Axes = None,
         path: str = None,
         **kwargs) -> FigAndAxes:
     fig, ax = create_fig_ax(fig, ax)
     ax.plot(spectrum.y, spectrum.spec_den_gw, **kwargs)
     ax.set_ylabel("spec_den_gw")
-    return plot_spectrum_common(spectrum, fig, ax, path)
+    return plot_spectrum_common(SSMSpectrum, fig, ax, path)
 
 
 def plot_spectrum_spec_den_v(
-        spectrum: Spectrum,
+        spectrum: SSMSpectrum,
         fig: plt.Figure = None,
         ax: plt.Axes = None,
         path: str = None,
         **kwargs) -> FigAndAxes:
     fig, ax = create_fig_ax(fig, ax)
     ax.plot(spectrum.y, spectrum.spec_den_v, **kwargs)
     ax.set_ylabel("spec_den_v")
-    return plot_spectrum_common(spectrum, fig, ax, path)
+    return plot_spectrum_common(SSMSpectrum, fig, ax, path)

pttools/omgw0/const.py

@@ -22,11 +22,16 @@
 #: :caprini_2020:`\ ` p. 12
 GS0 = 3.91
 
-#: Hubble constant, TODO
-H0: float = None
+#: Parsec to meters
+PC_TO_M: float = 3.0857e16
+
+#: Hubble constant, Planck value
+H0_KM_S_MPC: float = 67.4
+#: Hubble constant, Planck value in Hz (about 2.27e-18 Hz)
+H0_HZ: float = H0_KM_S_MPC * 1e3 / (PC_TO_M * 1e6)
 
 #: LISA arm length (m)
-LISA_ARM_LENGTH: float = 2.5e8
+LISA_ARM_LENGTH: float = 2.5e9
 
 #: LISA observation time (s)
 LISA_OBS_TIME: float = 126227808.

pttools/omgw0/noise.py

@@ -39,22 +39,26 @@ def ft(L: th.FloatOrArr = const.LISA_ARM_LENGTH) -> th.FloatOrArr:
 #     f_frac = f/ft
 #     cos_f_frac = np.cos(f_frac)
 #     W = 1 - np.exp(-2j*f_frac)
-#     return ((4 + 2*cos_f_frac)*P_oms(f, L) + 8*(1 + cos_f_frac + cos_f_frac**2 * P_acc(f, L))) * np.abs(W)**2
+#     return ((4 + 2*cos_f_frac)*P_oms(L) + 8*(1 + cos_f_frac + cos_f_frac**2 * P_acc(f, L))) * np.abs(W)**2
+
+
+def omega(f: th.FloatOrArr, S: th.FloatOrArr) -> th.FloatOrArr:
+    return 4*np.pi**2 / (3*const.H0_HZ**2) * f**3 * S
 
 
 def omega_eb(f: th.FloatOrArr, f_ref_eb: float = 25, omega_ref_eb: float = 8.9e-10) -> th.FloatOrArr:
     return omega_ref_eb * (f/f_ref_eb)**(2/3)
 
 
 def omega_gb(f: th.FloatOrArr) -> th.FloatOrArr:
-    return 4*np.pi**2 / (3*const.H0**2) * f**3 * S_gb(f)
+    return omega(f, S_gb(f))
 
 
 def omega_ins(f: th.FloatOrArr) -> th.FloatOrArr:
     r"""LISA instrument noise
     $$\Omega_\text{ins} = \left( \frac{4 \pi^2}{3 H_0^2} f^3 S_A(f)$$
     """
-    return (4*np.pi**2 / (3*const.H0**2)) * f**3 * S_AE(f)
+    return omega(f, S_AE(f))
 
 
 def omega_noise(f: th.FloatOrArr) -> th.FloatOrArr:
@@ -69,27 +73,29 @@ def P_acc(f: th.FloatOrArr, L: th.FloatOrArr = const.LISA_ARM_LENGTH) -> th.Floa
     return (3e-15 / ((2*np.pi*f)**2 * L))**2 * (1 + (0.4e-3/f)**2)
 
 
-def P_oms(f: th.FloatOrArr, L: th.FloatOrArr = const.LISA_ARM_LENGTH) -> th.FloatOrArr:
+def P_oms(L: th.FloatOrArr = const.LISA_ARM_LENGTH) -> th.FloatOrArr:
     r"""
     LISA optical metrology noise
     $$P_\text{oms}(f) = \left( \frac{1.5 \cdot 10^{-11} \text{m}}{L} \right)^2 \text{Hz}^{-1}$$
     :gowling_2021:`\ ` eq. 3.2
+    This is white noise and therefore independent of the frequency.
     """
-    # Todo: Why is the frequency not used?
     return (1.5e-11 / L)**2
 
 
-def S_AE(f: th.FloatOrArr, L: th.FloatOrArr = const.LISA_ARM_LENGTH) -> th.FloatOrArr:
+def S_AE(f: th.FloatOrArr, L: th.FloatOrArr = const.LISA_ARM_LENGTH, both_channels: bool = True) -> th.FloatOrArr:
     r"""Approximate noise power spectral density for the LISA A and E channels
     $$S_A = S_E$
     """
-    return 40/3 * (P_oms(f, L) + 4*P_acc(f, L)) * (1 + (f/(4*ft(L)/3))**2)
+    ret = 40/3 * (P_oms(L) + 4*P_acc(f, L)) * (1 + (3*f/(4*ft(L)))**2)
+    if both_channels:
+        return 1/np.sqrt(2) * ret
+    return ret
 
 
-# TODO: fix the units for A
 def S_gb(
         f: th.FloatOrArr,
-        A: float = 9e-38,
+        A: float = 9e-35,  # 1/mHz -> 10³
         f_ref_gb: float = 1,
         fk: float = 1.13e-3,
         a: float = 0.138,

pttools/omgw0/omgw0_ssm.py

@@ -7,20 +7,64 @@
 @author: chloeg, markh
 """
 
+import functools
+
 import numpy as np
 
+from pttools.bubble.boundary import Phase
 import pttools.bubble.ke_frac_approx as K
 import pttools.omgw0.suppression as sup
 from pttools.ssmtools.const import NPTDEFAULT, NptType
 import pttools.ssmtools.spectrum as ssm
 import pttools.type_hints as th
 from . import const
+from . import noise
+
+
+class Spectrum(ssm.SSMSpectrum):
+    def f(self, z: np.ndarray = None) -> th.FloatOrArr:
+        if z is None:
+            z = self.y
+        return f(z=z, r_star=self.r_star, f_star0=self.f_star0)
+
+    @functools.cached_property
+    def f_star0(self) -> float:
+        return f_star0(
+            T_n=self.bubble.Tn,
+            g_star=self.g_star
+        )
+
+    def F_gw0(self, g0: float = const.G0, gs0: float = const.GS0) -> float:
+        return F_gw0(
+            g_star=self.g_star,
+            g0=g0,
+            gs0=gs0,
+            gs_star=self.bubble.model.gs(w=self.bubble.va_enthalpy_density, phase=Phase.BROKEN)
+        )
+
+    @functools.cached_property
+    def g_star(self) -> float:
+        return self.bubble.model.gp(w=self.bubble.va_enthalpy_density, phase=Phase.BROKEN)
+
+    def noise(self) -> np.ndarray:
+        return noise.omega_noise(self.f())
+
+    def noise_ins(self) -> np.ndarray:
+        return noise.omega_ins(self.f())
+
+    def omgw0(
+            self,
+            g0: float = const.G0,
+            gs0: float = const.GS0,
+            suppression: sup.SuppressionMethod = sup.SuppressionMethod.DEFAULT) -> np.ndarray:
+        supp = sup.get_suppression_factor(vw=self.bubble.v_wall, alpha=self.bubble.alpha_n, method=suppression)
+        # The r_star compensates the fact that the pow_gw includes a correction factor that is J without r_star
+        return self.r_star * self.F_gw0(g0=g0, gs0=gs0) * self.pow_gw * supp
+
+    def signal_to_noise_ratio(self) -> float:
+        return noise.signal_to_noise_ratio(f=self.f(), signal=self.omgw0(), noise=self.noise())
 
 
-################################
-# SGWB as calculated by PTtools (SSM)
-################################
-
 def f(z: th.FloatOrArr, r_star: th.FloatOrArr, f_star0: th.FloatOrArr) -> th.FloatOrArr:
     r"""Convert the dimensionless wavenumber $z$ to frequency today by taking into account the redshift.
     $$f = \frac{z}{r_*} f_{*,0}$$,