Merge pull request #3 from fbleibel-g/ffcc_python

FFCC: WIP python implementation

python/README.md

@@ -0,0 +1,3 @@
+WARNING: Do not use; this is a work in progress implementation of FFCC in python, using tensorflow.
+
+There are currently missing parts in the implementation; it's not in a usable state, but future commits will try to remedy that.

python/ellipse.py

@@ -0,0 +1,126 @@
+# Copyright 2020 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Ellipse utilities."""
+import tensorflow.compat.v1 as tf
+
+
+def _check_inputs(mat, vec):
+  """Checks the matrix and vector describing an ellipse (in either form)."""
+  mat = tf.convert_to_tensor(mat)
+  vec = tf.convert_to_tensor(vec)
+  if len(mat.shape) != len(vec.shape) + 1:
+    raise ValueError(
+        'rank(`mat`) is not one more than rank(`vec`) ({} != {}+1)'.format(
+            len(mat.shape), len(vec.shape)))
+  if mat.shape.as_list()[-1] != vec.shape.as_list()[-1]:
+    raise ValueError(
+        '`mat`s outermost dimension ({}) does not mat `vec`s ({}))'.format(
+            mat.shape.as_list()[-1],
+            vec.shape.as_list()[-1]))
+  if mat.shape.as_list()[-1] != mat.shape.as_list()[-2]:
+    raise ValueError('`mat`s last two dimensions ({}, {}) do not match)'.format(
+        mat.shape.as_list()[-1],
+        mat.shape.as_list()[-2]))
+
+
+def general_to_standard(w_mat, b_vec):
+  """Converts a quadratic from "general form" to "standard form".
+
+  That is, a quadratic given in the form of
+    (`w_mat` * `x` + `b_vec`)^2
+  is converted to a quadratic in the form of
+    (`x` - `c_vec`)^T a_mat (`x` - `c_vec`)
+
+  Args:
+    w_mat: the shear of the quadratic.
+    b_vec: the shift of the quadratic.
+
+  Returns:
+    A tuple containing:
+      (a_mat: the "inverse covariance" of the quadratic,
+       c_vec: the "center" of the quadratic)
+  """
+  _check_inputs(w_mat, b_vec)
+  a_mat = tf.linalg.matmul(w_mat, w_mat)
+  c_vec = -tf.linalg.matvec(tf.linalg.inv(w_mat), b_vec)
+  return a_mat, c_vec
+
+
+def standard_to_general(a_mat, c_vec):
+  """Converts a quadratic from "standard form" to "general form".
+
+  That is, a quadratic given in the form of
+    (`x` - `c_vec`)^T `a_mat` (`x` - `c_vec`)
+  is converted to a quadratic in the form of
+    (`w_mat` * `x` + `b_vec`)^2
+
+  Args:
+    a_mat: the "inverse covariance" of the quadratic.
+    c_vec: the "center" of the quadratic.
+
+  Returns:
+    A tuple containing:
+      (w_mat: the shear of the quadratic,
+       b_vec: the shift of the quadratic)
+  """
+  _check_inputs(a_mat, c_vec)
+  w_mat = tf.linalg.sqrtm(a_mat)
+  b_vec = -tf.linalg.matvec(w_mat, c_vec)
+  return w_mat, b_vec
+
+
+def distance(x, w_mat, b_vec):
+  """Returns distance w.r.t an ellipse.
+
+  The ellipse distance is defined as:
+    sum((`w_mat` * `x` + `b_vec`).^2)
+
+  Args:
+    x: an input real tensor of in the shape of [batch_size, dim], where each
+      x[i, :] is a vector for which distance is computed.
+    w_mat: matrix representation of a conic section, in the shape of [dim, dim].
+    b_vec: vector representation of a b_vec from the origin, in the shape of
+      [dim].
+
+  Returns:
+    An output tensor in the shape of [batch_size] that corresponds to the
+    ellipse distance of each x[i, :].
+  """
+  _check_inputs(w_mat, b_vec)
+  v = tf.linalg.matmul(x, tf.cast(w_mat, x.dtype), transpose_b=True) + b_vec
+  return tf.reduce_sum(tf.square(v), axis=-1)
+
+
+def project(x, w_mat, b_vec):
+  """Given an input point x, projects it onto the ellipse.
+
+  The ellipse is defined by: sum((`w_mat` * `x` + `b_vec`).^2) <= 1
+
+  Args:
+    x: an input tensor in the shape of [batch_size, dim], where each x[i, :] is
+      a vector to be projected.
+    w_mat: matrix representation of a conic section, in the shape of [dim, dim].
+    b_vec: vector with an offset from the origin, in the shape of [dim].
+
+  Returns:
+    An output tensor in the same shape of x that represents projected vector.
+  """
+  _check_inputs(w_mat, b_vec)
+
+  d = distance(x, w_mat, b_vec)
+  scale = tf.rsqrt(tf.maximum(d, 1))
+  _, c_vec = general_to_standard(w_mat, b_vec)
+  y = scale[:, tf.newaxis] * x + (1 - scale[:, tf.newaxis]) * tf.cast(
+      c_vec[tf.newaxis], x.dtype)
+  return y

python/ellipse_test.py

@@ -0,0 +1,107 @@
+# Copyright 2020 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Tests for ellipse.py."""
+
+from . import ellipse
+import numpy as np
+from scipy.linalg import expm
+import tensorflow.compat.v1 as tf
+
+
+class EllipseTest(tf.test.TestCase):
+
+  def _get_random_ellipse(self, dim):
+    x = np.random.uniform(size=(dim, dim))
+    w_mat = expm((np.transpose(x) + x) / 2.0)
+    b_vec = np.random.uniform(size=(dim))
+    return w_mat, b_vec
+
+  def _eval(self, tensor, feed_dict=None):
+    with tf.Session() as sess:
+      return sess.run(tensor, feed_dict=feed_dict)
+
+  def setUp(self):
+    super(EllipseTest, self).setUp()
+    np.random.seed(0)
+
+  def testStandardToGeneralSpecialCases(self):
+    """Sanity-check a few easy to verify cases of standard_to_general()."""
+    a_mat = np.float32(np.array([[1, 0], [0, 1]]))
+    c_vec = np.float32(np.array([0, 0]))
+    w_mat, b_vec = ellipse.standard_to_general(a_mat, c_vec)
+    self.assertAllClose(w_mat, a_mat)
+    self.assertAllClose(b_vec, c_vec)
+
+    a_mat = np.float32(np.array([[1, 0], [0, 1]]))
+    c_vec = np.float32(np.random.normal(size=2))
+    w_mat, b_vec = ellipse.standard_to_general(a_mat, c_vec)
+    self.assertAllClose(w_mat, a_mat)
+    self.assertAllClose(b_vec, -c_vec)
+
+    a_mat = np.float32(np.array([[2., 0.], [0., 0.5]]))
+    c_vec = np.float32(np.array([1., 1.]))
+    w_mat, b_vec = ellipse.standard_to_general(a_mat, c_vec)
+    self.assertAllClose(w_mat, np.array([[np.sqrt(2), 0], [0, np.sqrt(0.5)]]))
+    self.assertAllClose(b_vec, np.array([-np.sqrt(2), -np.sqrt(0.5)]))
+
+  def testGeneraltoStandardRoundTrip(self):
+    for _ in range(10):
+      w_mat, b_vec = self._get_random_ellipse(2)
+      a_mat, c_vec = ellipse.general_to_standard(w_mat, b_vec)
+      w_mat_recon, b_vec_recon = ellipse.standard_to_general(a_mat, c_vec)
+      self.assertAllClose(w_mat, w_mat_recon)
+      self.assertAllClose(b_vec, b_vec_recon)
+
+  def testStandardToGeneralDistancesMatch(self):
+    """Check distance() against the standard parametrization's distance."""
+    for _ in range(10):
+      num_dims = 2
+      w_mat, b_vec = self._get_random_ellipse(num_dims)
+      a_mat, c_vec = ellipse.general_to_standard(w_mat, b_vec)
+      data = np.random.normal(size=(100, num_dims))
+      dist_general = ellipse.distance(data, w_mat, b_vec)
+      dist_standard = tf.reduce_sum(
+          (data - c_vec) * tf.linalg.matmul((data - c_vec), a_mat), axis=-1)
+      self.assertAllClose(dist_general, dist_standard)
+
+  def testProject(self):
+    for _ in range(10):
+      dim = np.random.randint(low=1, high=10)
+      w_mat, b_vec = self._get_random_ellipse(dim)
+
+      batch_size = 100
+      x = np.random.normal(size=(batch_size, dim))
+      proj = self._eval(ellipse.project(x, w_mat, b_vec))
+      x_distance = self._eval(ellipse.distance(x, w_mat, b_vec))
+      proj_distance = self._eval(ellipse.distance(proj, w_mat, b_vec))
+      self.assertTrue(np.all(np.less_equal(proj_distance, 1.0 + 1e-8)))
+
+      # Check points with distance < 1 have not changed.
+      mask = x_distance < 1
+      self.assertTrue(np.all(np.equal(x[mask, :], proj[mask, :])))
+
+      # Check points with distance >= 1 have a distance of 1 after projection.
+      np.testing.assert_allclose(proj_distance[x_distance >= 1], 1.)
+
+      # Check that the projected points are scaled versions of the input points.
+      center = -np.matmul(np.linalg.inv(w_mat), b_vec)
+      delta_ratio = (x - center) / (proj - center)
+      avg_delta_ratio = np.tile(
+          np.mean(delta_ratio, axis=-1, keepdims=True),
+          (1, delta_ratio.shape[1]))
+      np.testing.assert_allclose(delta_ratio, avg_delta_ratio)
+
+
+if __name__ == '__main__':
+  tf.test.main()