def test_multidim(self):
y = tf.ones((4, 5, 6), dtype=tf.float32)
integral_0 = self.evaluate(tfp_math.trapz(y, axis=0))
integral_1 = self.evaluate(tfp_math.trapz(y, axis=1))
integral_2 = self.evaluate(tfp_math.trapz(y))
self.assertTupleEqual(integral_0.shape, (5, 6))
self.assertTupleEqual(integral_1.shape, (4, 6))
self.assertTupleEqual(integral_2.shape, (4, 5))
self.assertAllClose(integral_0, np.ones((5, 6)) * 3.0)
self.assertAllClose(integral_1, np.ones((4, 6)) * 4.0)
self.assertAllClose(integral_2, np.ones((4, 5)) * 5.0)
def test_multidim_with_x(self):
y = tf.ones((4, 5), dtype=tf.float64)
v0 = tf.cast(tf.linspace(0., 1., 4), tf.float64)
v1 = tf.cast(tf.linspace(0., 4.2, 5), tf.float64)
x0, x1 = tf.meshgrid(v0, v1, indexing='ij')
integral_0 = self.evaluate(tfp_math.trapz(y, x0, axis=0))
integral_1 = self.evaluate(tfp_math.trapz(y, x1, axis=1))
self.assertTupleEqual(integral_0.shape, (5, ))
self.assertTupleEqual(integral_1.shape, (4, ))
self.assertAllClose(integral_0, np.ones((5, )) * 1.0)
self.assertAllClose(integral_1, np.ones((4, )) * 4.2)
def test_against_numpy(self, data):
dtype = data.draw(hps.sampled_from([np.float32, np.float64]))
shp = (data.draw(hps.integers(5, 10)), data.draw(hps.integers(5, 10)))
axis = data.draw(hps.integers(0, len(shp) - 1))
y = data.draw(
tfp_hps.constrained_tensors(tfp_hps.identity_fn, shp, dtype))
x_dx = data.draw(hps.sampled_from(['x', 'dx', None]))
if x_dx is None:
x = None
dx = None
elif x_dx == 'dx':
x = None
dx = data.draw(hps.floats(0.1, 10))
else:
x = data.draw(
tfp_hps.constrained_tensors(tfp_hps.identity_fn, shp, dtype))
dx = None
np_soln = np.trapz(
self.evaluate(y),
x=self.evaluate(x)
if x is not None else None, # cannot evaluate(None)
dx=dx or 1.0, # numpy default is 1.0
axis=axis)
tf_soln = tfp_math.trapz(y, x, dx, axis)
self.assertAllClose(np_soln, tf_soln)
def trapz_sin_fn(x_min, x_max, n, expected, rtol):
pi = tf.constant(np.pi, dtype=tf.float32)
x = tf.cast(tf.linspace(x_min, x_max, n), tf.float32)
y = tf.sin(pi * x)
np.testing.assert_allclose(self.evaluate(tfp_math.trapz(y, x)),
expected,
rtol=rtol)
def trapz_sin_fn(x_min, x_max, n, expected, rtol):
pi = tf.constant(np.pi, dtype=tf.float32)
s = np.linspace(0, 1, n)**1.5
x = tf.convert_to_tensor(x_min + s * (x_max - x_min), tf.float32)
y = tf.sin(pi * x)
np.testing.assert_allclose(
self.evaluate(tfp_math.trapz(y, x)), expected, rtol=rtol)
def test_multidim_with_nonhomogeneous_x(self):
# integration domain from x is different for each batch dim
y = tf.ones((3, 4), dtype=tf.float32)
x = tf.constant([[0, 1, 2, 3], [0, 2, 4, 6], [0, 3, 6, 9]],
dtype=tf.float32)
integral = self.evaluate(tfp_math.trapz(y, x, axis=-1))
self.assertTupleEqual(integral.shape, (3, ))
np.testing.assert_almost_equal(integral, [3, 6, 9])
def test_multidim_broadcast_1d_x(self):
# To use trapz() with a 1d x array, first broadcast it with the shape of y
y = tf.ones((4, 5), dtype=tf.float64)
x1 = tf.cast(tf.linspace(0., 4.2, 5), tf.float64)
x = x1 * tf.ones((4, 5), dtype=tf.float64)
integral = self.evaluate(tfp_math.trapz(y, x, axis=1))
self.assertTupleEqual(integral.shape, (4, ))
self.assertAllClose(integral, np.ones((4, )) * 4.2)
def test_incompatible_x_y_shape_raises(self):
with self.assertRaisesWithLiteralMatch(
ValueError, 'Shapes (3,) and (2, 3) are incompatible'):
tfp_math.trapz(y=tf.ones((2, 3)), x=tf.ones((3, )))
def test_provide_x_and_dx_args_raises(self):
with self.assertRaisesWithLiteralMatch(
ValueError,
'Not permitted to specify both x and dx input args.'):
tfp_math.trapz(y=[0, 0.1], x=[0, 0.1], dx=0.5)
def test_non_scalar_dx_raises(self):
with self.assertRaisesRegexp(ValueError, 'Expected dx to be a scalar'):
tfp_math.trapz(y=tf.ones((2, 3)), dx=[0.1, 0.2])
def test_provide_multiple_axes_raises(self):
with self.assertRaisesRegexp(ValueError,
'Only permitted to specify one axis'):
tfp_math.trapz(y=tf.ones((2, 3)), axis=[0, 1])
def test_simple_use_with_dx_arg(self):
y = tf.linspace(0., 10., 11)
dx = 0.1
integral = self.evaluate(tfp_math.trapz(y, dx=dx))
self.assertAllClose(integral, 5.0)
def test_simple_use_with_x_arg(self):
y = tf.linspace(0., 10., 11)
x = tf.linspace(0., 10., 11)
integral = self.evaluate(tfp_math.trapz(y, x))
self.assertAllClose(integral, 50.)
def g(y, x=None, dx=None):
return tfp_math.trapz(y, x=x, dx=dx, axis=axis, name=name)
