New plotting types

gwinteract/waveforms/forms.py

@@ -1,5 +1,7 @@
 from django import forms
 
+from collections import OrderedDict
+
 import lalsimulation
 
 def _is_useable_psd(func):
@@ -18,7 +20,13 @@ def _get_psds(module):
 
 class WaveformForm(forms.Form):
 
-    _PSDS = dict(_get_psds(lalsimulation))
+    _PSDS = OrderedDict(_get_psds(lalsimulation))
+    _REPR = OrderedDict((("time", "time domain"),
+             ("freq", "frequency domain"),
+             ("tf", "time-frequency")))
+
+    representation = forms.ChoiceField(label="representation", \
+                            choices=_REPR.items())
 
     m1 = forms.FloatField(label='m1')#, default=1.4)
     m2 = forms.FloatField(label='m2')#, default=1.4)

gwinteract/waveforms/views.py

@@ -8,6 +8,7 @@
 
 import numpy
 from gwpy.frequencyseries import FrequencySeries
+from gwpy.timeseries import TimeSeries
 from matplotlib import pyplot
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 
@@ -23,8 +24,40 @@
     "distance": 100., "inclination": 0.,
 }
 
-def gen_waveform(event_params, flow=10.0, deltaf=0.125, fhigh=2048., fref=10.,
-                    approximant="IMRPhenomPv2"):
+def gen_waveform_td(event_params, flow=10.0, fhigh=2048.,
+                    fref=10., approximant="IMRPhenomPv2"):
+    """
+    Generate the h_+ and h_x polarizations for an event, as well as an associated frequency array.
+    """
+    eprm = _defaults.copy()
+    eprm.update(event_params)
+
+    delta_t = 0.5 / fhigh
+
+    params = None
+    hp, hx = lalsimulation.SimInspiralTD(
+                # Masses
+                eprm["m1"] * lal.MSUN_SI, eprm["m2"] * lal.MSUN_SI, \
+                # Spins
+                eprm["spin1x"], eprm["spin1y"], eprm["spin1z"], \
+                eprm["spin1x"], eprm["spin1y"], eprm["spin1z"], \
+                # distance and inclination
+                eprm["distance"] * 1e6 * lal.PC_SI, eprm["inclination"],
+                # These are eccentricity and other orbital parameters
+                0.0, 0.0, 0.0, 0.0,
+                # frequency binning params
+                delta_t, flow, fref, \
+                # Other extraneous options
+                params,
+                lalsimulation.SimInspiralGetApproximantFromString(approximant))
+
+    time_ar = numpy.arange(0, len(hp.data.data)) * delta_t
+    time_ar += float(hp.epoch)
+
+    return time_ar, hp.data.data, hx.data.data
+
+def gen_waveform_fd(event_params, deltaf=0.125, flow=10.0, fhigh=2048.,
+                    fref=10., approximant="IMRPhenomPv2"):
     """
     Generate the h_+ and h_x polarizations for an event, as well as an associated frequency array.
     """
@@ -58,6 +91,73 @@ def gen_psd(form, freq_ar, asd=True):
     psd[0] = psd[1]
     return numpy.sqrt(psd) if asd else numpy.asarray(psd)
 
+def plot_fd(form, plot_flow=10., plot_fhigh=2048., ax=None):
+
+    f, hp, hx = gen_waveform_fd(form.cleaned_data)
+    assert len(f) == len(hp), "{0}, {1}".format(len(f), len(hp))
+
+    hp = FrequencySeries(hp, frequencies=f)
+    hx = FrequencySeries(hx, frequencies=f)
+
+    asd = FrequencySeries(gen_psd(form, f), frequencies=f) if \
+                        form.cleaned_data["psd"] != "None" else None
+
+    #if ax is not None:
+        #plot.sca(ax)
+    plot = hp.abs().plot(color='red', label=r"$|h_+|(f)$", yscale='log')
+    ax = plot.gca()
+    ax.plot(hx.abs(), color='black', linestyle='-.', label=r"$|h_{\times}|(f)$")
+    if asd is not None:
+        ax.plot(asd)
+
+    ax.set_xlim(plot_flow, plot_fhigh)
+
+    # For now just plot |h|
+    # FIXME: These will most likely overlap
+    ax.set_ylabel(r'GW strain [strain$/\sqrt{\mathrm{Hz}}$]')
+    plot.legend()
+    return plot
+
+def plot_td(form, ax=None):
+
+    t, hp, hx = gen_waveform_td(form.cleaned_data, flow=10)
+    assert len(t) == len(hp), "{0}, {1}".format(len(t), len(hp))
+
+    hp = TimeSeries(hp, t0=t[0], dt=t[1] - t[0])
+    hx = TimeSeries(hx, t0=t[0], dt=t[1] - t[0])
+
+    #if ax is not None:
+        #plot.sca(ax)
+    plot = hp.plot(color='red', label=r"$h_+(t)$")
+    ax = plot.gca()
+    ax.plot(hx, color='black', label=r"$h_{\times}(t)$")
+
+    # For now just plot |h|
+    # FIXME: These will most likely overlap
+    ax.set_ylabel(r'GW strain [strain]')
+    ax.set_xlabel(r'time [s]')
+    plot.legend()
+    return plot
+
+def plot_tf(form, ax=None):
+
+    t, hp, hx = gen_waveform_td(form.cleaned_data, flow=10)
+    assert len(t) == len(hp), "{0}, {1}".format(len(t), len(hp))
+
+    hp = TimeSeries(hp, t0=t[0], dt=t[1] - t[0])
+    hx = TimeSeries(hx, t0=t[0], dt=t[1] - t[0])
+
+    q = hp.q_transform()
+    med, sig = numpy.median(q), numpy.std(q)
+
+    #if ax is not None:
+        #plot.sca(ax)
+    plot = q.plot(vmin=(med - 5 * sig))
+    ax = plot.gca()
+    ax.set_yscale('log')
+
+    return plot
+
 #######
 
 def index(request):
@@ -83,26 +183,15 @@ def plot(request):
         form = WaveformForm(request.GET)
         # check whether it's valid:
         if form.is_valid():
-           # This is where waveform generation code will go Chris P
-           # You parse the request for as such
-            f, hp, hx = gen_waveform(form.cleaned_data)
-            assert len(f) == len(hp), "{0}, {1}".format(len(f), len(hp))
-
-            hp = FrequencySeries(hp, frequencies=f)
-            hx = FrequencySeries(hx, frequencies=f)
-
-            asd = FrequencySeries(gen_psd(form, f), frequencies=f) if form.cleaned_data["psd"] != "None" else None
-
-            plot = hp.abs().plot(color='red', label=r"$|h_+|(f)$", yscale='log')
-            ax = plot.gca()
-            ax.plot(hx.abs(), color='black', linestyle='-.', label=r"$|h_x|(f)$")
-            if asd is not None:
-                ax.plot(asd)
-
-            # For now just plot |h|
-            # FIXME: These will most likely overlap
-            ax.set_ylabel(r'GW strain [strain$/\sqrt{\mathrm{Hz}}$]')
-            plot.legend()
+
+            if form.cleaned_data["representation"] == "freq":
+                plot = plot_fd(form)
+            elif form.cleaned_data["representation"] == "time":
+                plot = plot_td(form)
+            else:
+                plot = plot_tf(form)
+
+            # Package result and export through HTTP response
             buf = io.BytesIO()
             canvas = FigureCanvas(plot)
             canvas.print_png(buf)