Merge pull request #133 from SourbhBh:master

PiperOrigin-RevId: 445518637
Change-Id: I893642f3097c2fb78a7203724f91a65d15e015be

Author: Johannes Ballé

models/toy_sources/ramp.py

@@ -0,0 +1,82 @@
+# Copyright 2022 TensorFlow Compression contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Ramp process."""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+class Ramp(tfp.distributions.Distribution):
+  """The "ramp": Y(t) = (t+V) mod 1 - 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"):
+    """Initializer.
+
+    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.float32`.
+      validate_args: required by `Distribution` class but unused.
+      allow_nan_stats: required by `Distribution` class but unused.
+      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_tensor(self):
+    return tf.shape(self.index_points)
+
+  def _event_shape(self):
+    return self.index_points.shape
+
+  def _sample_n(self, n, seed=None):
+    ind = self.index_points
+    if self.phase is None:
+      phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+    else:
+      phase = tf.fill((n, 1), tf.constant(self.phase, dtype=self.dtype))
+    return (ind + phase) % 1 - .5

models/toy_sources/sawbridge.py

@@ -1,29 +1,48 @@
+# Copyright 2020 TensorFlow Compression contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
 """Sawbridge process."""
 
 import tensorflow as tf
 import tensorflow_probability as tfp
 
 
 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].
 
-  def __init__(self, index_points, stationary=True, order=1,
-               dtype=tf.float32, validate_args=False, allow_nan_stats=True,
-               name="sawbridge"):
+  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, 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] or `None` (default). Specifies realization of V.
+      drop: Float in [0,1] or `None` (default). 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.
-      dtype: Data type of the returned realization at each timestep. Defaults to
-        tf.float32.
-      validate_args: required by tensorflow Distribution class but unused.
-      allow_nan_stats: required by tensorflow Distribution class but unused.
+      dtype: Data type of the returned realization. Defaults to `tf.float32`.
+      validate_args: required by `Distribution` class but unused.
+      allow_nan_stats: required by `Distribution` class but unused.
       name: String. Name of the created object.
     """
     parameters = dict(locals())
@@ -32,6 +51,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 +74,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,16 +95,21 @@ 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)
+    if self.drop is None:
+      uniform = tf.random.uniform(
+          (self.order, n, 1), seed=seed, dtype=self.dtype)
+    else:
+      uniform = tf.fill(
+          (self.order, n, 1), tf.constant(self.drop, dtype=self.dtype))
     ind = self.index_points
-    # 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.constant(self.phase, dtype=self.dtype)
+      ind += phase
       ind %= 1.
-    less = tf.less(uniform, ind)
-    # less shape: (order, n, time)
+    less = tf.less(uniform, ind)  # shape: (order, n, time)
     # Note:
     # ind[n] == 1 -> sample[n] == 0 always
     # ind[n] == 0 -> sample[n] == 0 always

models/toy_sources/sinusoid.py

@@ -0,0 +1,86 @@
+# Copyright 2022 TensorFlow Compression contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Sinusoid process."""
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+class Sinusoid(tfp.distributions.Distribution):
+  """The "sinusoid": P(t) = sin(2pi(t+V)), 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="sinusoid"):
+    """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] or `None` (default). Specifies realization of V.
+      dtype: Data type of the returned realization at each timestep. Defaults
+        to tf.float32.
+      validate_args: required by tensorflow 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._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_tensor(self):
+    return tf.shape(self.index_points)
+
+  def _event_shape(self):
+    return self.index_points.shape
+
+  def _sample_n(self, n, seed=None):
+    ind = self.index_points
+    if self.phase is None:
+      phase = tf.random.uniform((n, 1), seed=seed, dtype=self.dtype)
+    else:
+      phase = tf.fill((n, 1), tf.constant(self.phase, dtype=self.dtype))
+    return tf.sin((2 * np.pi) * (ind + phase))

models/toy_sources/sphere.py

@@ -0,0 +1,82 @@
+# Copyright 2022 TensorFlow Compression contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Sphere process."""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+class Sphere(tfp.distributions.Distribution):
+  """Uniform distribution on the unit hypersphere."""
+
+  def __init__(self,
+               order=2,
+               width=0.,
+               dtype=tf.float32,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="sphere"):
+    """Initializer.
+
+    Arguments:
+      order: Integer >= 1. The dimensionality of the sphere.
+      width: Float in [0,1]. Allows for realizations to be approximately
+        uniformly distributed in a band between radius `1 - width` and
+        `1 + width` (for `width` << 1).
+      dtype: Data type of the returned realization. Defaults to `tf.float32`.
+      validate_args: required by `Distribution` class but unused.
+      allow_nan_stats: required by `Distribution` class but unused.
+      name: String. Name of the created object.
+    """
+    parameters = dict(locals())
+    self._order = int(order)
+    self._width = float(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 order(self):
+    return self._order
+
+  @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_tensor(self):
+    return tf.constant([self.order], dtype=tf.int32)
+
+  def _event_shape(self):
+    return tf.TensorShape([self.order])
+
+  def _sample_n(self, n, seed=None):
+    samples = tf.random.normal((n, self.order), seed=seed, dtype=self.dtype)
+    radius = tf.math.sqrt(tf.reduce_sum(tf.square(samples), -1, keepdims=True))
+    if self.width:
+      radius *= tf.random.uniform(
+          (n, 1), minval=1. - self.width / 2., maxval=1. + self.width / 2.,
+          seed=seed, dtype=self.dtype)
+    return samples / radius