Merge pull request #112 from TomWagg/h_c

More flexible `plot_sources_on_sc`

changelog.md

@@ -166,6 +166,10 @@ having two when we could just set `f_dom=2 f_orb`
 - Ensure `theta` is a co-latitude in `amplitude_modulation`
 
 ## 0.4.6
-*TW 13/04/22*
+*TW 13/04/23*
 - [Issue [#109](https://github.com/TeamLEGWORK/LEGWORK/issues/109)] Upgrade `numpy` version
-    - Upgrade required to avoid error in `np.nan_to_num` usage - thanks to Jakob Stegmann for raising this
+    - Upgrade required to avoid error in `np.nan_to_num` usage - thanks to Jakob Stegmann for raising this
+
+## 0.4.7
+*TW 12/10/23*
+- [Issue [#111](https://github.com/TeamLEGWORK/LEGWORK/issues/111)] Allow `plot_sources_on_sc` to change the underlying sensitivity curve directly with `sc_vis_settings`

legwork/_version.py

@@ -1 +1 @@
-__version__ = "0.4.6"
+__version__ = "0.4.7"

legwork/evol.py

@@ -15,7 +15,7 @@
            'determine_stationarity']
 
 
-@jit
+@jit(nopython=True)
 def de_dt(e, times, beta, c_0):                             # pragma: no cover
     """Compute eccentricity time derivative
 

legwork/source.py

@@ -1192,7 +1192,7 @@ def plot_source_variables(self, xstr, ystr=None, which_sources=None,
             return vis.plot_1D_dist(x=x[which_sources], weights=weights, **kwargs)
 
     def plot_sources_on_sc(self, snr_cutoff=0, fig=None, ax=None, show=True,
-                           label="Stationary sources", **kwargs):  # pragma: no cover
+                           label="Stationary sources", sc_vis_settings={}, **kwargs):  # pragma: no cover
         """Plot all sources in the class on the sensitivity curve
 
         Parameters
@@ -1214,6 +1214,10 @@ def plot_sources_on_sc(self, snr_cutoff=0, fig=None, ax=None, show=True,
         label : `str`
             Label to use for the plotted points
 
+        sc_vis_settings : `dict`
+            Dictionary of parameters to pass to :meth:`~legwork.visualisation.plot_sensitivity_curve`,
+            e.g. {"y_quantity": "h_c"} will plot characteristic strain instead of ASD
+
         **kwargs : `various`
             Keyword arguments to be passed to plotting functions
 
@@ -1250,7 +1254,7 @@ def plot_sources_on_sc(self, snr_cutoff=0, fig=None, ax=None, show=True,
             weights = self.weights[stat] if self.weights is not None else None
             fig, ax = vis.plot_sources_on_sc(f_dom=f_dom, snr=self.snr[stat], weights=weights,
                                              snr_cutoff=snr_cutoff, show=show, fig=fig, ax=ax,
-                                             label=label, **self._sc_params, **kwargs)
+                                             label=label, sc_vis_settings=sc_vis_settings, **self._sc_params, **kwargs)
 
         # show warnings for evolving sources
         circ_evol = self.get_source_mask(circular=True, stationary=False)

legwork/tests/test_evol.py

@@ -137,7 +137,7 @@ def test_t_merge_special_cases(self):
             single_time = evol.get_t_merge_ecc(beta=beta[0], a_i=a_i[0],
                                                ecc_i=ecc_i[0],
                                                large_e_tol=0.9).to(u.yr)
-            self.assertTrue(array_time[0] == single_time)
+            self.assertTrue(np.isclose(array_time[0], single_time))
 
     def test_mandel_fit(self):
         """checks that the Mandel fit to the Peters timescale is the same

legwork/visualisation.py

@@ -374,7 +374,7 @@ def plot_sensitivity_curve(frequency_range=None, y_quantity="ASD", fig=None, ax=
 
 def plot_sources_on_sc(f_dom, snr, weights=None, snr_cutoff=0, t_obs="auto",
                        instrument="LISA", custom_psd=None, L="auto", approximate_R=False,
-                       confusion_noise="auto", fig=None, ax=None, show=True, **kwargs):
+                       confusion_noise="auto", fig=None, ax=None, show=True, sc_vis_settings={}, **kwargs):
     """Overlay *stationary* sources on the LISA sensitivity curve.
 
     Each source is plotted at its max snr harmonic frequency such that that its height above the curve is
@@ -429,6 +429,10 @@ def plot_sources_on_sc(f_dom, snr, weights=None, snr_cutoff=0, t_obs="auto",
     show : `boolean`
         Whether to immediately show the plot or only return the Figure and Axis
 
+    sc_vis_settings : `dict`
+        Dictionary of parameters to pass to :meth:`~legwork.visualisation.plot_sensitivity_curve`,
+        e.g. {"y_quantity": "h_c"} will plot characteristic strain instead of ASD
+
     **kwargs : `various`
         This function is a wrapper on :func:`legwork.visualisation.plot_2D_dist` and each kwarg is passed
         directly to this function. For example, you can write `disttype="kde"` for a kde density plot
@@ -446,7 +450,7 @@ def plot_sources_on_sc(f_dom, snr, weights=None, snr_cutoff=0, t_obs="auto",
     if fig is None or ax is None:
         fig, ax = plot_sensitivity_curve(show=False, t_obs=t_obs, instrument=instrument, L=L,
                                          custom_psd=custom_psd, approximate_R=approximate_R,
-                                         confusion_noise=confusion_noise)
+                                         confusion_noise=confusion_noise, **sc_vis_settings)
 
     # work out which binaries are above the cutoff
     detectable = snr > snr_cutoff
@@ -455,12 +459,14 @@ def plot_sources_on_sc(f_dom, snr, weights=None, snr_cutoff=0, t_obs="auto",
         return fig, ax
 
     # calculate asd that makes it so height above curve is snr
-    asd = snr[detectable] * np.sqrt(psd.power_spectral_density(f_dom[detectable]))
+    asd = (snr[detectable] * np.sqrt(psd.power_spectral_density(f_dom[detectable]))).to(u.Hz**(-1/2))
+    h_c = asd * np.sqrt(f_dom[detectable])
+    use_h_c = ("y_quantity" in sc_vis_settings and sc_vis_settings["y_quantity"] == "h_c")
 
     # plot either a scatter or density plot of the detectable binaries
     ylims = ax.get_ylim()
     weights = weights[detectable] if weights is not None else None
-    fig, ax = plot_2D_dist(x=f_dom[detectable], y=asd.to(u.Hz**(-1/2)), weights=weights,
+    fig, ax = plot_2D_dist(x=f_dom[detectable], y=h_c if use_h_c else asd, weights=weights,
                            fig=fig, ax=ax, show=False, **kwargs)
     ax.set_ylim(ylims)