Speed and performance improvements

.gitignore

@@ -1,3 +1,11 @@
+# Images
+*.png
+*.jpg
+*.jpeg
+
+# Ephemera
+.python-version
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

AUTHORS.rst

@@ -5,7 +5,7 @@ Credits
 Development Lead
 ----------------
 
-* Daniel Williams <d.williams.2@research.gla.ac.uk>
+* Daniel Williams <daniel.williams@glasgow.ac.uk>
 
 Contributors
 ------------

heron/inference.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+import os
 
 import click
 
@@ -13,7 +14,7 @@
 
 from heron.detector import KNOWN_IFOS
 from heron.models.lalnoise import KNOWN_PSDS
-from heron.likelihood import TimeDomainLikelihood, MultiDetector
+from heron.likelihood import TimeDomainLikelihood, MultiDetector, TimeDomainLikelihoodModelUncertainty
 import heron.priors
 
 from heron.sampling import NessaiSampler
@@ -26,13 +27,17 @@
 )
 from heron.utils import load_yaml
 
+import otter
+
 logger = logging.getLogger("heron.inference")
 
 KNOWN_LIKELIHOODS = {
     "TimeDomainLikelihood": TimeDomainLikelihood,
+    "TimeDomainUncertaintyLikelihood": TimeDomainLikelihoodModelUncertainty,
 }
 KNOWN_WAVEFORMS = {
     "IMRPhenomPv2": IMRPhenomPv2,
+    "IMRPhenomPv2_FakeUncertainty": IMRPhenomPv2_FakeUncertainty,
 }
 
 
@@ -57,6 +62,7 @@ def parse_dict(settings):
 def heron_inference(settings):
 
     settings = load_yaml(settings)
+    webdir = settings['pages directory']
     settings, other_settings = parse_dict(settings)
 
     if "logging" in other_settings:
@@ -81,12 +87,21 @@ def heron_inference(settings):
 
     data = {}
 
+    report = otter.Otter(os.path.join(
+        webdir,
+        "inference.html"),
+        author="Heron",
+        title="Heron Inference"
+    )
+
     if "data files" in settings.get("data", {}):
         # Load frame files from disk
-
-        start = settings['event time'] - settings['segment length'] + settings['after merger']
+        with report:
+            report += "# Data"
+        start = settings['event time'] - \
+            settings['segment length'] + settings['after merger']
         end = settings['event time'] + settings['after merger']
-        
+
         for ifo in settings["interferometers"]:
             print(f"Loading {ifo} data")
             logger.info(
@@ -100,11 +115,22 @@ def heron_inference(settings):
                 start=start,
                 end=end,
             )
+            with report:
+                report += f"## {ifo}"
+                f = data[ifo].plot()
+                f.savefig(os.path.join(
+                    other_settings.get('pages directory', 'pages'),
+                    f"{ifo}_data.png"
+                ))
+                report += f
+
             if data[ifo].sample_rate != settings['likelihood']['sampling rate']:
-                logger.info("Resampling the data to the likelihood sampling rate")
-                data[ifo] = data[ifo].resample(settings['likelihood']['sampling rate'])
-    #elif "injection" in other_settings:
-    #    pass
+                logger.info(
+                    "Resampling the data to the likelihood sampling rate")
+                data[ifo] = data[ifo].resample(
+                    settings['likelihood']['sampling rate'])
+    # elif "injection" in other_settings:
+    # TODO: implement handling for 'injection' in other_settings if required
 
     # Make Likelihood
     if len(settings["interferometers"]) > 1:
@@ -125,16 +151,47 @@ def heron_inference(settings):
                     ),
                 )
             )
+            print(likelihoods[-1])
         likelihood = MultiDetector(*likelihoods)
 
     priors = heron.priors.PriorDict()
     priors.from_dictionary(settings["priors"])
 
-    if settings["sampler"]["sampler"] == "nessai":
+    if settings["sampler"]["sampler"] == "naive":
+        import numpy as np
+        import matplotlib.pyplot as plt
+        # Just draw 100 points across mass ratio space
+        #prior_points = np.linspace(50, 70, 100)
+        #prior_points = np.linspace(50, 150, 100)
+        posterior = []
+        prior_points = np.linspace(settings["sampler"]["naive range"][0],
+            settings["sampler"]["naive range"][1],
+            100)
+        for mass in prior_points:
+            parameters = {
+            "total_mass": 60,#mass * u.solMass,
+            "mass_ratio": 1.0,
+            "luminosity_distance": 100,
+            "gpstime": 4000,
+            "ra": 1.0,
+            "dec": 0.3}
+
+            parameters[settings["sampler"]["naive parameter"]] = mass
+            parameters = injection_parameters_add_units(parameters)
+            posterior.append(likelihood(parameters))
+
+        posterior = np.array(posterior)
+        print(posterior)
+        f, ax = plt.subplots(1,1, dpi=300)
+        ax.plot(prior_points, posterior)
+        f.savefig(os.path.join(webdir, "posterior.png"))
+
+    elif settings["sampler"]["sampler"] == "nessai":
         nessai_model = NessaiSampler(
             likelihood,
             priors,
-            injection_parameters_add_units(other_settings["injection"]["parameters"]),
+            injection_parameters_add_units(
+                other_settings["injection"]["parameters"]),
         )
 
         fp = FlowSampler(
@@ -145,14 +202,18 @@ def heron_inference(settings):
             ),
             output=settings["name"],
             resume=settings.get("sampler", {}).get("resume", True),
-            checkpointing=settings.get("sampler", {}).get("checkpointing", True),
+            checkpointing=settings.get(
+                "sampler", {}).get("checkpointing", True),
             checkpoint_interval=settings.get("sampler", {}).get(
                 "checkpointing interval", 3600
             ),
-            logging_interval=settings.get("sampler", {}).get("logging interval", 10),
-            log_on_iteration=settings.get("sampler", {}).get("log on iteration", True),
+            logging_interval=settings.get(
+                "sampler", {}).get("logging interval", 10),
+            log_on_iteration=settings.get(
+                "sampler", {}).get("log on iteration", True),
             seed=settings.get("sampler", {}).get("seed", 1234),
-            flow_class=settings.get("sampler", {}).get("flow class", "GWFlowProposal"),
+            flow_class=settings.get("sampler", {}).get(
+                "flow class", "GWFlowProposal"),
             signal_handling=True,
         )
 

heron/injection.py

@@ -2,6 +2,8 @@
 Code to create injections using the various models supported by heron.
 """
 
+import os
+
 import logging
 
 import click
@@ -27,50 +29,59 @@ def make_injection(
         times=None,
         detectors=None,
         framefile=None,
-        psdfile=None
+        psdfile=None,
+        webdir="pages"
 ):
 
-    parameters = {"ra": 0, "dec": 0, "psi": 0, "theta_jn": 0, "phase": 0}
+    parameters = {"ra": 0, "dec": 0, "psi": 0, "theta_jn": 0, "phase": 0, "gpstime": 4000}
     parameters.update(injection_parameters)
-
     waveform = waveform()
 
-    if times is None:
-        times = np.linspace(parameters['gpstime']-duration+2, parameters['gpstime']+2, int(duration * sample_rate))
+    # The caller is responsible for providing an appropriate 'times' array
+    # (e.g. a time grid around 'gpstime') before invoking this function.
+    print(f"Making a waveform from {times[0]} to {times[-1]}")
     waveform = waveform.time_domain(
         parameters,
         times=times,
     )
 
+    import matplotlib
+    matplotlib.use("agg")
+    from gwpy.plot import Plot
+    f = Plot(waveform['plus'], waveform['cross'], separate=True)
+    f.savefig(os.path.join(webdir, f"injected_waveform_sep.png"))
+
     injections = {}
     for detector, psd_model in detectors.items():
+        print(f"Making injection for {detector}")
         logger.info(f"Making injection for {detector}")
         psd_model = KNOWN_PSDS[psd_model]()
         detector = KNOWN_IFOS[detector]()
         if times is None:
             times = waveform['plus'].times.value
         data = psd_model.time_series(times)
-        print(data)
-
         channel = f"{detector.abbreviation}:Injection"
+
+        print("Projecting the waveform")
         injection = data + waveform.project(detector)
         injection.channel = channel
         injections[detector.abbreviation] = injection
-        # likelihood = TimeDomainLikelihood(injection, psd=psd_model)
-        # snr = likelihood.snr(waveform.project(detector))
 
-        #logger.info(f"Optimal Filter SNR: {snr}")
+        f = Plot(data, waveform.project(detector), injection, separate=True)
+        f.savefig(os.path.join(webdir, f"{detector.abbreviation}_injected_waveform.png"))
 
         if framefile:
             filename = f"{detector.abbreviation}_{framefile}.gwf"
+            f = Plot(injection)
+            f.savefig(os.path.join(webdir, f"{detector.abbreviation}_{framefile}.png"))
             logger.info(f"Saving framefile to {filename}")
             injection.write(filename, format="gwf")
 
         if psdfile:
             # Write the PSD file to an ascii file
             filename = f"{detector.abbreviation}_{psdfile}.dat"
             psd_model.to_file(filename)
-            
+
     return injections
 
 
@@ -102,11 +113,12 @@ def make_injection_zero_noise(
         psd_model = KNOWN_PSDS[psd_model]()
         #data = psd_model.time_series(times)
 
-        # import matplotlib
-        # matplotlib.use("agg")
-        # from gwpy.plot import Plot
-        # f = Plot(data, waveform.project(detector), data+waveform.project(detector), separate=False)
-        # f.savefig(f"{detector.abbreviation}_injected_waveform.png")
+        import matplotlib
+        matplotlib.use("agg")
+        from gwpy.plot import Plot
+        webdir = settings['pages directory']
+        f = Plot(data, waveform.project(detector), data+waveform.project(detector), separate=False)
+        f.savefig(os.path.join(webdir, f"{detector.abbreviation}_injected_waveform.png"))
 
         injection = waveform.project(detector)
         injection.channel = channel
@@ -134,6 +146,9 @@ def injection_parameters_add_units(parameters):
 @click.command()
 @click.option("--settings")
 def injection(settings):
+
+    click.echo("Preparing an injection")
+
     settings = load_yaml(settings)
 
     if "logging" in settings:
@@ -147,21 +162,34 @@ def injection(settings):
         }
 
         logging.basicConfig(level=LOGGER_LEVELS[level])
-
+    webdir = settings['pages directory']
     settings = settings["injection"]
+
+    print(settings)
+
     parameters = injection_parameters_add_units(settings["parameters"])
 
     detector_dict = {
         settings["interferometers"][ifo]: settings["psds"][ifo]
         for ifo in settings["interferometers"]
     }
+
+    click.echo("Making injections")
+    t0 = settings['parameters']['gpstime']
+    duration = settings["duration"]
+    srate = settings["sample rate"]
+
+    times = np.linspace(t0-(duration-2), t0+2, duration*srate)
+
     injections = make_injection(
         waveform=IMRPhenomPv2,
         duration=settings["duration"],
+        times=times,
         sample_rate=settings["sample rate"],
         injection_parameters=parameters,
         detectors=detector_dict,
         framefile="injection",
         psdfile="psd",
+        webdir=webdir,
     )
     data = injections