Add docstrings and a general clean up

peyote/source.py

@@ -48,6 +48,24 @@ def LALBinaryBlackHole(
 
 
 class WaveformGenerator:
+    """ A waveform generator
+
+    Parameters
+    ----------
+    sampling_frequency: float
+        The sampling frequency to sample at
+    time_duration: float
+        Time duration of data
+    source_model: func
+        A python function taking some arguments and returning the frequency
+        domain strain. Note the first argument must be the frequencies at
+        which to compute the strain
+
+    Note: the arguments of source_model (except the first, which is the
+    frequencies at which to compute the strain) will be added to the
+    WaveformGenerator object and initialised to `None`.
+
+    """
     def __init__(self, name, sampling_frequency, time_duration, source_model):
 
         self.parameter_keys = inspect.getargspec(source_model).args
@@ -65,10 +83,12 @@ class WaveformGenerator:
         self.source_model = source_model
 
     def frequency_domain_strain(self):
+        """ Wrapper to source_model """
         kwargs = {k: self.__dict__[k] for k in self.parameter_keys}
         return self.source_model(self.frequency_array, **kwargs)
 
     def set_values(self, dictionary):
+        """ Given a dictionary of values, set the class attributes """
         for k in self.parameter_keys:
             try:
                 setattr(self, k, dictionary[k])

tutorials/BasicTutorial.py

@@ -4,8 +4,6 @@ import pylab as plt
 import dynesty.plotting as dyplot
 import corner
 import peyote
-import lal
-import lalsimulation as lalsim
 
 peyote.setup_logger()
 
@@ -28,6 +26,7 @@ simulation_parameters = dict(
     psi=2.659
     )
 
+# Create the waveformgenerator using a LAL BinaryBlackHole source function
 waveformgenerator = peyote.source.WaveformGenerator(
     'BBH', sampling_frequency, time_duration, peyote.source.LALBinaryBlackHole)
 waveformgenerator.set_values(simulation_parameters)