Sawbridge modified and pylint verified

Author: SourbhBh

models/toy_sources/circle.py

@@ -1,61 +1,64 @@
+"""Circle process."""
+
 import tensorflow as tf
 import tensorflow_probability as tfp
 import numpy as np
+
 class Circle(tfp.distributions.Distribution):
-    """The "circle": [cos x, sin x] where x is uniform over [0,2pi]"""
-    def __init__(self,width=0.,dtype=tf.float32,validate_args=False,allow_nan_stats=True,name='circle'):
-        """
-        Args: 
-        width: Float in [0,1]. Allows for realizations to be uniformly distributed in a band between radius 1-width and 1+width. 
+  """The "circle": [cos x, sin x] where x is uniform over [0,2pi]"""
+  def __init__(self,width=0.,dtype=tf.float32,validate_args=False,
+                 allow_nan_stats=True,name='circle'):
+    """
+    Args:
+        width: Float in [0,1]. Allows for realizations to be uniformly
+        distributed in a band between radius 1-width and 1+width.
         dtype: Data type of returned realization. Defaults to tf.float32.
         validate_args: required by tensoflow Distribution class but unused.
         allow_nan_stats: required by tensorflow Distribution class but unused.
         name: String. Name of the created object.
-        """
-        parameters = dict(locals())
-        with tf.name_scope(name) as name:
-            self._width=width
-        super().__init__(dtype=dtype,
+    """
+    parameters = dict(locals())
+    with tf.name_scope(name) as name:
+      self._width=width
+    super().__init__(dtype=dtype,
                 reparameterization_type=tfp.distributions.NOT_REPARAMETERIZED,
                 validate_args = validate_args,
                 allow_nan_stats = allow_nan_stats,
                 parameters = parameters,
                 name=name)
-    @property
-    def width(self):
-        return self._width
+  @property
+  def width(self):
+    return self._width
 
-    def _batch_shape_tensor(self):
-        return tf.constant([],dtype=tf.int32)
-    def _batch_shape(self):
-        return tf.TensorShape([])
-    def _event_shape(self):
-        smp = tf.linspace(0.,1.,2)
-        return smp.shape
-    def _event_shape_tensor(self):
-        smp = tf.linspace(0.,1.,2)
-        return tf.shape(smp)
-    def _sample_n(self,n,seed=None):
-        angles = tf.random.uniform((1,n),minval = 0, maxval = 2*tf.constant(np.pi),seed=seed,dtype=self.dtype)
-        sample = tf.convert_to_tensor([tf.math.cos(angles),tf.math.sin(angles)])
-        if self.width>0:
-            dev = tf.random.uniform((1,n),minval = -(self.width)/2.0, maxval = (self.width)/2.0, seed=seed, dtype=self.dtype)
-            sample = sample + dev*sample
-        sample = tf.transpose(sample,[2,1,0])
-        sample = tf.reshape(sample,[n,2])
-        return sample
+  def _batch_shape_tensor(self):
+    return tf.constant([],dtype=tf.int32)
+  def _batch_shape(self):
+    return tf.TensorShape([])
+  def _event_shape(self):
+    smp = tf.linspace(0.,1.,2)
+    return smp.shape
+  def _event_shape_tensor(self):
+    smp = tf.linspace(0.,1.,2)
+    return tf.shape(smp)
+  def _sample_n(self,n,seed=None):
+    angles = tf.random.uniform((1,n),minval = 0,maxval = 2*tf.constant(np.pi)
+                               ,seed=seed,dtype=self.dtype)
+    sample = tf.convert_to_tensor([tf.math.cos(angles),tf.math.sin(angles)])
+    if self.width>0:
+      dev = tf.random.uniform((1,n),minval = -(self.width)/2.0,
+                                maxval = (self.width)/2.0, seed=seed,
+                                dtype=self.dtype)
+      sample = sample + dev*sample
+      sample = tf.transpose(sample,[2,1,0])
+      sample = tf.reshape(sample,[n,2])
+      return sample
 
-    def _prob(self, x):
-        new_x  = tf.reshape(x,[-1,tf.shape(x)[-1]])
-        flat_norm = tf.norm(new_x,axis=1)
-        if self.width>0:
-            prob = tf.where((flat_norm >= 1-self.width)&(flat_norm <= 1+self.width),
-                    1.0/(2*tf.convert_to_tensor(np.pi)*self.width),
-                    0.
-                    )
-        else:
-            prob = tf.where(flat_norm==1,
-                1.0/(2*tf.convert_to_tensor(np.pi)),
-                0.
-                )
-        return tf.reshape(prob,tf.shape(x)[:-1])
+  def _prob(self, x):
+    new_x  = tf.reshape(x,[-1,tf.shape(x)[-1]])
+    flat_norm = tf.norm(new_x,axis=1)
+    if self.width>0:
+      prob = tf.where((flat_norm >= 1-self.width)&(flat_norm <= 1+self.width),
+                    1.0/(2*tf.convert_to_tensor(np.pi)*self.width),0.)
+    else:
+      prob = tf.where(flat_norm==1,1.0/(2*tf.convert_to_tensor(np.pi)),0.)
+      return tf.reshape(prob,tf.shape(x)[:-1])

models/toy_sources/ramp.py

@@ -1,57 +1,62 @@
+"""Ramp process."""
+
 import tensorflow as tf
 import tensorflow_probability as tfp
 
 class Ramp(tfp.distributions.Distribution):
-    """The "ramp": Y(t) = (t+V)mod1 - 0.5, where V is uniform over [0,1]. """
-    def __init__(self,index_points,phase=None,dtype=tf.float32,validate_args=False,allow_nan_stats=True,name='ramp'):
-        """
-        Args:
-        index_points: 1-D Tensor representing the locations at which to evaluate the process.
-        phase: Float in [0,1]. Specifies a realization of V.
-        dtype: Data type of the returned realization. Defaults to tf.flaot32.
-        validate_args: required by tensorflow Distribution class but not used.
-        allow_nan_stats: required by tensorflow Distribution class but not used.
-        name: String. Name of the created object. 
-        """
-        parameters = dict(locals())
-        with tf.name_scope(name) as name:
-            self._index_points = tf.convert_to_tensor(index_points,dtype_hint= dtype,name='index_points')
-            self._phase = phase
-        super().__init__(
+  """The "ramp": Y(t) = (t+V)mod1 - 0.5, where V is uniform over [0,1]. """
+  def __init__(self,index_points,phase=None,dtype=tf.float32,
+                 validate_args=False,allow_nan_stats=True,name='ramp'):
+    """
+    Args:
+    index_points: 1-D Tensor representing the locations at which to evaluate
+    the process.
+    phase: Float in [0,1]. Specifies a realization of V.
+    dtype: Data type of the returned realization. Defaults to tf.flaot32.
+    validate_args: required by tensorflow Distribution class but not used.
+    allow_nan_stats: required by tensorflow Distribution class but not used.
+    name: String. Name of the created object.
+    """
+    parameters = dict(locals())
+    with tf.name_scope(name) as name:
+      self._index_points = tf.convert_to_tensor(index_points,dtype_hint= dtype,
+                                                name='index_points')
+      self._phase = phase
+    super().__init__(
                 dtype=dtype,
                 reparameterization_type = tfp.distributions.NOT_REPARAMETERIZED,
                 validate_args = validate_args,
                 allow_nan_stats = allow_nan_stats,
                 parameters = parameters,
                 name=name,)
-    @property
-    def index_points(self):
-        return self._index_points
-
-    @property 
-    def phase(self):
-        return self._phase
-
-    def _batch_shape_tensor(self):
-        return tf.constant([],dtype=tf.int32)
-
-    def _batch_shape(self):
-        return tf.TensorShape([])
-
-    def _event_shape(self):
-        return self.index_points.shape
-
-    def _event_shape_tensor(self):
-        return tf.shape(self.index_points)
-
-    def _sample_n(self,n,seed=None):
-        ind = self.index_points
-        if self._phase is None:
-            ind += tf.random.uniform ((n,1),seed=seed,dtype=self.dtype)
-            ind %= 1
-        else:
-            phase = tf.repeat(self._phase,repeats=n,axis=0)
-            phase = tf.reshape(phase,[n,1])
-            ind += phase
-            ind %= 1
-        return ind - 0.5
+  @property
+  def index_points(self):
+    return self._index_points
+
+  @property
+  def phase(self):
+    return self._phase
+
+  def _batch_shape_tensor(self):
+    return tf.constant([],dtype=tf.int32)
+
+  def _batch_shape(self):
+    return tf.TensorShape([])
+
+  def _event_shape(self):
+    return self.index_points.shape
+
+  def _event_shape_tensor(self):
+    return tf.shape(self.index_points)
+
+  def _sample_n(self,n,seed=None):
+    ind = self.index_points
+    if self._phase is None:
+      ind += tf.random.uniform ((n,1),seed=seed,dtype=self.dtype)
+      ind %= 1
+    else:
+      phase = tf.repeat(self._phase,repeats=n,axis=0)
+      phase = tf.reshape(phase,[n,1])
+      ind += phase
+      ind %= 1
+      return ind - 0.5

models/toy_sources/sawbridge.py

@@ -5,18 +5,22 @@
 
 
 class Sawbridge(tfp.distributions.Distribution):
-  """The "sawbridge": B(t) = t - 1(t > Z), where Z is uniform over [0,1]."""
+  """The "sawbridge": B(t) = t - 1(t > Z), where Z is uniform over [0,1].
+  The stationary sawbridge is B(t+V mod 1) where V is uniform over [0,1]
+  and Z,V are independent."""
 
-  def __init__(self, index_points, stationary=True, order=1,
-               dtype=tf.float32, validate_args=False, allow_nan_stats=True,
-               name="sawbridge"):
+  def __init__(self, index_points, phase=None, drop = None, stationary=True,
+               order=1,dtype=tf.float32, validate_args=False,
+               allow_nan_stats=True,name="sawbridge"):
     """Initializer.
 
     Args:
       index_points: 1-D `Tensor` representing the locations at which to evaluate
         the process. The intent is that all locations are in [0,1], but the
         process has a natural extrapolation outside this range so no error is
         thrown.
+      phase: Float in [0,1]. Specifies realization of V.
+      drop: Float in [0,1]. Specifies realization of Z.
       stationary: Boolean. Whether or not to "scramble" phase.
       order: Integer >= 1. The resulting process is a linear combination of
         `order` sawbridges.
@@ -32,6 +36,8 @@ def __init__(self, index_points, stationary=True, order=1,
           index_points, dtype_hint=dtype, name="index_points")
       self._stationary = bool(stationary)
       self._order = int(order)
+      self._phase = phase
+      self._drop = drop
     super().__init__(
         dtype=dtype,
         reparameterization_type=tfp.distributions.NOT_REPARAMETERIZED,
@@ -53,6 +59,14 @@ def stationary(self):
   def order(self):
     return self._order
 
+  @property
+  def phase(self):
+    return self._phase
+
+  @property
+  def drop(self):
+    return self._drop
+
   def _batch_shape_tensor(self):
     return tf.constant([], dtype=tf.int32)
 
@@ -66,13 +80,22 @@ def _event_shape(self):
     return self.index_points.shape
 
   def _sample_n(self, n, seed=None):
-    uniform = tf.random.uniform((self.order, n), seed=seed, dtype=self.dtype)
     ind = self.index_points
+    if self._drop is None:
+      uniform = tf.random.uniform((self.order,n),seed=seed,dtype = self.dtype)
+    else:
+      uniform = tf.repeat(self._drop,repeats = n*self.order)
+      uniform = tf.reshape(uniform,[self.order,n])
     # ind shape: (time)
     uniform = tf.expand_dims(uniform, -1)
     # uniform shape: (order, n, 1)
     if self.stationary:
-      ind += tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+      if self._phase is None:
+        phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+      else:
+        phase = tf.repeat(self._phase,repeats=n,axis=0)
+        phase = tf.reshape(phase, [n,1])
+      ind += phase
       ind %= 1.
     less = tf.less(uniform, ind)
     # less shape: (order, n, time)