From b81ff086c642c96d957780e474254f3c0feedd12 Mon Sep 17 00:00:00 2001
From: Gregory Ashton <gregory.ashton@ligo.org>
Date: Tue, 5 Jun 2018 17:55:46 +1000
Subject: [PATCH] Clean up example to remove the waveform generator

---
 examples/other_examples/gaussian_example.py   |  56 +++++++++
 examples/other_examples/linear_regression.py  |  55 ++++-----
 ...near_regression_with_waveform_generator.py | 115 ++++++++++++++++++
 3 files changed, 195 insertions(+), 31 deletions(-)
 create mode 100644 examples/other_examples/gaussian_example.py
 create mode 100644 examples/other_examples/linear_regression_with_waveform_generator.py

diff --git a/examples/other_examples/gaussian_example.py b/examples/other_examples/gaussian_example.py
new file mode 100644
index 000000000..b65e3bc0a
--- /dev/null
+++ b/examples/other_examples/gaussian_example.py
@@ -0,0 +1,56 @@
+#!/bin/python
+"""
+An example of how to use tupak to perform paramater estimation for
+non-gravitational wave data consisting of a Gaussian with a mean and variance
+"""
+from __future__ import division
+import tupak
+import numpy as np
+
+# A few simple setup steps
+tupak.utils.setup_logger()
+label = 'gaussian_example'
+outdir = 'outdir'
+
+# Here is minimum requirement for a Likelihood class to run with tupak. In this
+# case, we setup a GaussianLikelihood, which needs to have a log_likelihood
+# method. Note, in this case we will NOT make use of the `tupak`
+# waveform_generator to make the signal.
+
+# Making simulated data: in this case, we consider just a Gaussian
+
+data = np.random.normal(3, 4, 100)
+
+
+class GaussianLikelihood(tupak.Likelihood):
+    def __init__(self, data):
+        """
+        A very simple Gaussian likelihood
+
+        Parameters
+        ----------
+        data: array_like
+            The data to analyse
+        """
+        self.data = data
+        self.N = len(data)
+        self.parameters = {'mu': None, 'sigma': None}
+
+    def log_likelihood(self):
+        mu = self.parameters['mu']
+        sigma = self.parameters['sigma']
+        res = self.data - mu
+        return -0.5 * (np.sum((res / sigma)**2)
+                       + self.N*np.log(2*np.pi*sigma**2))
+
+
+likelihood = GaussianLikelihood(data)
+priors = dict(mu=tupak.prior.Uniform(0, 5, 'mu'),
+              sigma=tupak.prior.Uniform(0, 10, 'sigma'))
+
+# And run sampler
+result = tupak.run_sampler(
+    likelihood=likelihood, priors=priors, sampler='dynesty', npoints=500,
+    walks=10, outdir=outdir, label=label)
+result.plot_corner()
+print(result)
diff --git a/examples/other_examples/linear_regression.py b/examples/other_examples/linear_regression.py
index 208e05b97..7d659a673 100644
--- a/examples/other_examples/linear_regression.py
+++ b/examples/other_examples/linear_regression.py
@@ -1,28 +1,23 @@
 #!/bin/python
 """
 An example of how to use tupak to perform paramater estimation for
-non-gravitational wave data
+non-gravitational wave data. In this case, fitting a linear function to
+data with background Gaussian noise
+
 """
 from __future__ import division
 import tupak
 import numpy as np
 import matplotlib.pyplot as plt
+import inspect
 
 # A few simple setup steps
 tupak.utils.setup_logger()
-label = 'linear-regression'
+label = 'linear_regression'
 outdir = 'outdir'
 
-# Here is minimum requirement for a Likelihood class to run linear regression
-# with tupak. In this case, we setup a GaussianLikelihood, which needs to have
-# a log_likelihood method. Note, in this case we make use of the `tupak`
-# waveform_generator to make the signal (more on this later) But, one could
-# make this work without the waveform generator.
-
-# Making simulated data
-
 
-# First, we define our signal model, in this case a simple linear function
+# First, we define our "signal model", in this case a simple linear function
 def model(time, m, c):
     return time * m + c
 
@@ -56,8 +51,10 @@ fig.savefig('{}/{}_data.png'.format(outdir, label))
 
 
 class GaussianLikelihood(tupak.Likelihood):
-    def __init__(self, x, y, sigma, waveform_generator):
+    def __init__(self, x, y, sigma, function):
         """
+        A general Gaussian likelihood - the parameters are inferred from the
+        arguments of function
 
         Parameters
         ----------
@@ -65,19 +62,23 @@ class GaussianLikelihood(tupak.Likelihood):
             The data to analyse
         sigma: float
             The standard deviation of the noise
-        waveform_generator: `tupak.waveform_generator.WaveformGenerator`
-            An object which can generate the 'waveform', which in this case is
-            any arbitrary function
+        function:
+            The python function to fit to the data. Note, this must take the
+            dependent variable as its first argument. The other arguments are
+            will require a prior and will be sampled over (unless a fixed
+            value is given).
         """
         self.x = x
         self.y = y
         self.sigma = sigma
         self.N = len(x)
-        self.waveform_generator = waveform_generator
-        self.parameters = waveform_generator.parameters
+        self.function = function
+        parameters = inspect.getargspec(function).args
+        parameters.pop(0)
+        self.parameters = dict.fromkeys(parameters)
 
     def log_likelihood(self):
-        res = self.y - self.waveform_generator.time_domain_strain()
+        res = self.y - self.function(self.x, **self.parameters)
         return -0.5 * (np.sum((res / self.sigma)**2)
                        + self.N*np.log(2*np.pi*self.sigma**2))
 
@@ -86,18 +87,9 @@ class GaussianLikelihood(tupak.Likelihood):
                        + self.N*np.log(2*np.pi*self.sigma**2))
 
 
-# Here, we make a `tupak` waveform_generator. In this case, of course, the
-# name doesn't make so much sense. But essentially this is an objects that
-# can generate a signal. We give it information on how to make the time series
-# and the model() we wrote earlier.
-
-waveform_generator = tupak.WaveformGenerator(
-    time_duration=time_duration, sampling_frequency=sampling_frequency,
-    time_domain_source_model=model)
-
-# Now lets instantiate a version of out GravitationalWaveTransient, giving it
-# the time, data and waveform_generator
-likelihood = GaussianLikelihood(time, data, sigma, waveform_generator)
+# Now lets instantiate a version of our GaussianLikelihood, giving it
+# the time, data and signal model
+likelihood = GaussianLikelihood(time, data, sigma, model)
 
 # From hereon, the syntax is exactly equivalent to other tupak examples
 # We make a prior
@@ -109,5 +101,6 @@ priors['c'] = tupak.prior.Uniform(-2, 2, 'c')
 result = tupak.run_sampler(
     likelihood=likelihood, priors=priors, sampler='dynesty', npoints=500,
     walks=10, injection_parameters=injection_parameters, outdir=outdir,
-    label=label, plot=True)
+    label=label)
+result.plot_corner()
 print(result)
diff --git a/examples/other_examples/linear_regression_with_waveform_generator.py b/examples/other_examples/linear_regression_with_waveform_generator.py
new file mode 100644
index 000000000..f09658c6c
--- /dev/null
+++ b/examples/other_examples/linear_regression_with_waveform_generator.py
@@ -0,0 +1,115 @@
+#!/bin/python
+"""
+An example of how to use tupak to perform paramater estimation for
+non-gravitational wave data. In this case, fitting a linear function to
+data with background Gaussian noise. This example illustrates how the
+waveform_generator could be used.
+"""
+from __future__ import division
+import tupak
+import numpy as np
+import matplotlib.pyplot as plt
+
+# A few simple setup steps
+tupak.utils.setup_logger()
+label = 'linear-regression'
+outdir = 'outdir'
+
+# Here is minimum requirement for a Likelihood class to run linear regression
+# with tupak. In this case, we setup a GaussianLikelihood, which needs to have
+# a log_likelihood method. Note, in this case we make use of the `tupak`
+# waveform_generator to make the signal (more on this later) But, one could
+# make this work without the waveform generator.
+
+# Making simulated data
+
+
+# First, we define our signal model, in this case a simple linear function
+def model(time, m, c):
+    return time * m + c
+
+
+# New we define the injection parameters which we make simulated data with
+injection_parameters = dict(m=0.5, c=0.2)
+
+# For this example, we'll use standard Gaussian noise
+sigma = 1
+
+# These lines of code generate the fake data. Note the ** just unpacks the
+# contents of the injection_paramsters when calling the model function.
+sampling_frequency = 10
+time_duration = 10
+time = np.arange(0, time_duration, 1/sampling_frequency)
+N = len(time)
+data = model(time, **injection_parameters) + np.random.normal(0, sigma, N)
+
+# We quickly plot the data to check it looks sensible
+fig, ax = plt.subplots()
+ax.plot(time, data, 'o', label='data')
+ax.plot(time, model(time, **injection_parameters), '--r', label='signal')
+ax.set_xlabel('time')
+ax.set_ylabel('y')
+ax.legend()
+fig.savefig('{}/{}_data.png'.format(outdir, label))
+
+
+# Parameter estimation: we now define a Gaussian Likelihood class relevant for
+# our model.
+
+
+class GaussianLikelihood(tupak.Likelihood):
+    def __init__(self, x, y, sigma, waveform_generator):
+        """
+
+        Parameters
+        ----------
+        x, y: array_like
+            The data to analyse
+        sigma: float
+            The standard deviation of the noise
+        waveform_generator: `tupak.waveform_generator.WaveformGenerator`
+            An object which can generate the 'waveform', which in this case is
+            any arbitrary function
+        """
+        self.x = x
+        self.y = y
+        self.sigma = sigma
+        self.N = len(x)
+        self.waveform_generator = waveform_generator
+        self.parameters = waveform_generator.parameters
+
+    def log_likelihood(self):
+        res = self.y - self.waveform_generator.time_domain_strain()
+        return -0.5 * (np.sum((res / self.sigma)**2)
+                       + self.N*np.log(2*np.pi*self.sigma**2))
+
+    def noise_log_likelihood(self):
+        return -0.5 * (np.sum((self.y / self.sigma)**2)
+                       + self.N*np.log(2*np.pi*self.sigma**2))
+
+
+# Here, we make a `tupak` waveform_generator. In this case, of course, the
+# name doesn't make so much sense. But essentially this is an objects that
+# can generate a signal. We give it information on how to make the time series
+# and the model() we wrote earlier.
+
+waveform_generator = tupak.WaveformGenerator(
+    time_duration=time_duration, sampling_frequency=sampling_frequency,
+    time_domain_source_model=model)
+
+# Now lets instantiate a version of out GravitationalWaveTransient, giving it
+# the time, data and waveform_generator
+likelihood = GaussianLikelihood(time, data, sigma, waveform_generator)
+
+# From hereon, the syntax is exactly equivalent to other tupak examples
+# We make a prior
+priors = {}
+priors['m'] = tupak.prior.Uniform(0, 5, 'm')
+priors['c'] = tupak.prior.Uniform(-2, 2, 'c')
+
+# And run sampler
+result = tupak.run_sampler(
+    likelihood=likelihood, priors=priors, sampler='dynesty', npoints=500,
+    walks=10, injection_parameters=injection_parameters, outdir=outdir,
+    label=label, plot=True)
+print(result)
-- 
GitLab