def test_spherical_hessian_is_symmetric():
with AutoDifferentiator(persistent=True) as diff:
x = tf.reshape(tf.linspace(-4.5, 3.5, 9), (3, 3))
diff.watch(x)
x_reverse = tf.reverse(x, axis=tf.constant([1]))
y = tf.reshape(tf.einsum('ij,ij->i', x ** 2, x_reverse ** 2), (3, 1))
assert np.allclose(
diff.batch_hessian(
x, y, 0, 1, CoordinateSystem.SPHERICAL
).numpy(),
diff.batch_hessian(
x, y, 1, 0, CoordinateSystem.SPHERICAL
).numpy())
assert np.allclose(
diff.batch_hessian(
x, y, 0, 2, CoordinateSystem.SPHERICAL
).numpy(),
diff.batch_hessian(
x, y, 2, 0, CoordinateSystem.SPHERICAL
).numpy())
assert np.allclose(
diff.batch_hessian(
x, y, 1, 2, CoordinateSystem.SPHERICAL
).numpy(),
diff.batch_hessian(
x, y, 2, 1, CoordinateSystem.SPHERICAL
).numpy())
def _step(self, batch: DataBatch, optimizer: tf.optimizers.Optimizer,
diff_eq_loss_weight: float, ic_loss_weight: float,
bc_loss_weight: float) -> Loss:
"""
Performs a forward pass on the batch, computes the batch loss, and
updates the model parameters.
:param batch: the batch to compute the losses over
:param optimizer: the optimizer to use to update parameters of the
model
:param diff_eq_loss_weight: the weight of the differential equation
part of the total physics-informed loss
:param ic_loss_weight: the weight of the initial condition part of the
total physics-informed loss
:param bc_loss_weight: the weight of the boundary condition part of the
total physics-informed loss
:return: the various losses over the batch
"""
with AutoDifferentiator() as auto_diff:
loss = self._physics_informed_loss(batch, diff_eq_loss_weight,
ic_loss_weight, bc_loss_weight)
optimizer.minimize(loss.weighted_total_loss,
self.trainable_variables,
tape=auto_diff)
return loss
def _mean_squared_differential_equation_error(
self, batch: DomainDataBatch) -> tf.Tensor:
"""
Computes and returns the mean squared differential equation error.
:param batch: the domain data batch
:return: the mean squared differential equation error
"""
with AutoDifferentiator(persistent=True) as auto_diff:
auto_diff.watch(batch.t)
if batch.x is not None:
auto_diff.watch(batch.x)
y_hat = self.call((batch.u, batch.t, batch.x))
symbol_map_arg = PIDONSymbolMapArg(auto_diff, batch.t, batch.x,
y_hat)
rhs = self._symbol_mapper.map(symbol_map_arg)
diff_eq_residual = tf.concat([
self._diff_eq_lhs_functions[i](symbol_map_arg) - rhs[i]
for i in range(len(rhs))
],
axis=1)
squared_diff_eq_error = tf.square(diff_eq_residual)
return tf.reduce_mean(squared_diff_eq_error, axis=0)
def _mean_squared_boundary_condition_errors(
self, batch: BoundaryDataBatch) -> Tuple[tf.Tensor, tf.Tensor]:
"""
Computes and returns the mean squared Dirichlet boundary condition
error and the mean squared Neumann boundary condition error.
:param batch: the boundary data batch
:return: the mean squared Dirichlet and Neumann boundary condition
errors
"""
with AutoDifferentiator() as auto_diff:
auto_diff.watch(batch.x)
y_hat = self.call((batch.u, batch.t, batch.x))
d_y_over_d_n_hat = auto_diff.batch_gradient(batch.x, y_hat, batch.axes)
dirichlet_bc_error = y_hat - batch.y
dirichlet_bc_error = tf.where(tf.math.is_nan(batch.y),
tf.zeros_like(batch.y),
dirichlet_bc_error)
squared_dirichlet_bc_error = tf.square(dirichlet_bc_error)
mean_squared_dirichlet_bc_error = \
tf.reduce_mean(squared_dirichlet_bc_error, axis=0)
neumann_bc_error = d_y_over_d_n_hat - batch.d_y_over_d_n
neumann_bc_error = tf.where(tf.math.is_nan(batch.d_y_over_d_n),
tf.zeros_like(batch.d_y_over_d_n),
neumann_bc_error)
squared_neumann_bc_error = tf.square(neumann_bc_error)
mean_squared_neumann_bc_error = \
tf.reduce_mean(squared_neumann_bc_error, axis=0)
return mean_squared_dirichlet_bc_error, mean_squared_neumann_bc_error
def test_hessian_is_symmetric():
with AutoDifferentiator(persistent=True) as diff:
x = tf.zeros((3, 2), dtype=tf.float32)
diff.watch(x)
y = tf.reduce_sum(x, axis=1, keepdims=True) ** 2
assert np.allclose(
diff.batch_hessian(x, y, 0, 1).numpy(),
diff.batch_hessian(x, y, 1, 0).numpy())
def test_gradient_with_out_of_bounds_x_axis():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = c * x
with pytest.raises(ValueError):
diff.batch_gradient(x, y, 2).numpy()
def test_1d_curl():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 1), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1.], [-4.]], dtype=tf.float32)
y = c * x
with pytest.raises(ValueError):
diff.batch_curl(x, y).numpy()
def test_gradient_with_insufficient_dimensions():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2,), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[2.], [-5.]], dtype=tf.float32)
y = c * x
with pytest.raises(ValueError):
diff.batch_gradient(x, y, 0).numpy()
def test_vector_laplacian_with_non_matching_vector_field_dimension():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = tf.concat([c * x, c / x], axis=1)
with pytest.raises(ValueError):
diff.batch_vector_laplacian(x, y, 0).numpy()
def test_divergence():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = c * x
expected_divergence = [[3.], [-9.]]
actual_divergence = diff.batch_divergence(x, y).numpy()
assert np.allclose(expected_divergence, actual_divergence)
def test_laplacian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[-1.], [2.]], dtype=tf.float32)
y = c * tf.reduce_sum(x * x, axis=1, keepdims=True)
expected_laplacian = [[-4.], [8.]]
actual_laplacian = diff.batch_laplacian(x, y).numpy()
assert np.allclose(expected_laplacian, actual_laplacian)
def test_polar_laplacian_is_hessian_trace():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[-1.], [2.]], dtype=tf.float32)
y = c * tf.reduce_sum(x * x, axis=1, keepdims=True)
assert np.allclose(
diff.batch_laplacian(x, y, CoordinateSystem.POLAR).numpy(),
(diff.batch_hessian(x, y, 0, 0, CoordinateSystem.POLAR) +
diff.batch_hessian(x, y, 1, 1, CoordinateSystem.POLAR)).numpy())
def test_polar_hessian_is_symmetric():
with AutoDifferentiator(persistent=True) as diff:
x = tf.reshape(tf.linspace(-3., 2., 6), (3, 2))
diff.watch(x)
x_reverse = tf.reverse(x, axis=tf.constant([1]))
y = tf.reshape(tf.einsum('ij,ij->i', x ** 2, x_reverse ** 2), (3, 1))
assert np.allclose(
diff.batch_hessian(x, y, 0, 1, CoordinateSystem.POLAR).numpy(),
diff.batch_hessian(x, y, 1, 0, CoordinateSystem.POLAR).numpy())
def test_gradient_int_x_axis():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = c * x
x_axis = 1
expected_gradient = [[0., 2.], [0., -5.]]
actual_gradient = diff.batch_gradient(x, y, x_axis).numpy()
assert np.allclose(actual_gradient, expected_gradient)
def test_mixed_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.zeros((3, 2), dtype=tf.float32)
diff.watch(x)
y = tf.reduce_sum(x, axis=1, keepdims=True) ** 2
x_axis1 = 1
x_axis2 = 0
expected_hessian = [[2.], [2.], [2.]]
actual_hessian = diff.batch_hessian(x, y, x_axis1, x_axis2).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = c * x ** 2
x_axis1 = x_axis2 = 0
expected_hessian = [[2., 0.], [-8., 0.]]
actual_hessian = diff.batch_hessian(x, y, x_axis1, x_axis2).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_curl_with_out_of_bounds_ind():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 3), dtype=tf.float32)
diff.watch(x)
c = tf.constant(
[[1., 2., 3.], [-4., -3., -2.]],
dtype=tf.float32)
y = c * x
with pytest.raises(ValueError):
diff.batch_curl(x, y, 4).numpy()
def test_spherical_laplacian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2., 3.], [6., 5., 4.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[-1.], [2.]], dtype=tf.float32)
y = c * tf.reduce_sum(x * x, axis=1, keepdims=True)
expected_laplacian = [[-66.33603], [12.68897]]
actual_laplacian = diff.batch_laplacian(
x, y, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(expected_laplacian, actual_laplacian)
def test_gradient_tensor_x_axis():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((3, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[3., -7.], [2., -1.], [4., 5.]], dtype=tf.float32)
y = c * x
x_axis = tf.constant([1, 0, 0], dtype=tf.int32)
expected_gradient = [[0., -7.], [2., 0.], [4., 0.]]
actual_gradient = diff.batch_gradient(x, y, x_axis).numpy()
assert np.allclose(actual_gradient, expected_gradient)
def test_spherical_divergence():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((2, 3), dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., 2., 3.], [-4., -5., -6.]], dtype=tf.float32)
y = c * x
expected_divergence = [[10.303068], [-27.79453]]
actual_divergence = diff.batch_divergence(
x, y, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(expected_divergence, actual_divergence)
def test_cylindrical_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2., 3.], [4., 5., 6.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., -2., 1.5], [-4., -5., 4.5]], dtype=tf.float32)
y = c * x ** 2
x_axis1 = x_axis2 = 1
expected_hessian = [[2., -4., 0.], [-8., -.625, 0.]]
actual_hessian = diff.batch_hessian(
x, y, x_axis1, x_axis2, CoordinateSystem.CYLINDRICAL).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_polar_gradient():
with AutoDifferentiator(persistent=True) as diff:
x = tf.fill((2, 2), 2.)
diff.watch(x)
c = tf.constant([[1., 2.], [-4., -5.]], dtype=tf.float32)
y = c * x
x_axis = 1
expected_gradient = [[0., 1.], [0., -2.5]]
actual_gradient = diff.batch_gradient(
x, y, x_axis, CoordinateSystem.POLAR).numpy()
assert np.allclose(actual_gradient, expected_gradient)
def test_spherical_gradient():
with AutoDifferentiator(persistent=True) as diff:
x = tf.fill((2, 3), 2.)
diff.watch(x)
c = tf.constant([[1., 2., 3.], [-4., -5., -6.]], dtype=tf.float32)
y = c * x
x_axis = 2
expected_gradient = [[0., 0., 1.5], [0., 0., -3.]]
actual_gradient = diff.batch_gradient(
x, y, x_axis, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(actual_gradient, expected_gradient)
def test_mixed_spherical_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.reshape(tf.linspace(-4.5, 3.5, 9), (3, 3))
diff.watch(x)
x_reverse = tf.reverse(x, axis=tf.constant([1]))
y = tf.reshape(tf.einsum('ij,ij->i', x ** 2, x_reverse ** 2), (3, 1))
x_axis1 = 1
x_axis2 = 0
expected_hessian = [[-14.151262], [.4635177], [79.187874]]
actual_hessian = diff.batch_hessian(
x, y, x_axis1, x_axis2, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def value_and_gradients_function(
parameters: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
set_model_parameters(parameters)
with AutoDifferentiator() as auto_diff:
loss = self._physics_informed_loss(full_training_data_batch,
diff_eq_loss_weight,
ic_loss_weight,
bc_loss_weight)
value = tf.reduce_sum(loss.weighted_total_loss, keepdims=True)
gradients = auto_diff.gradient(value, self.trainable_variables)
flattened_gradients = tf.concat(
[tf.reshape(gradient, (1, -1)) for gradient in gradients],
axis=1)
return value, flattened_gradients
def test_mixed_polar_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.reshape(tf.linspace(-3., 2., 6), (3, 2))
diff.watch(x)
x_reverse = tf.reverse(x, axis=tf.constant([1]))
y = tf.reshape(tf.einsum('ij,ij->i', x ** 2, x_reverse ** 2), (3, 1))
x_axis1 = 1
x_axis2 = 0
expected_hessian = [[-8.], [0.], [8.]]
actual_hessian = diff.batch_hessian(
x, y, x_axis1, x_axis2, CoordinateSystem.POLAR).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_mixed_cylindrical_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.reshape(tf.linspace(-4.5, 3.5, 9), (3, 3))
diff.watch(x)
x_reverse = tf.reverse(x, axis=tf.constant([1]))
y = tf.reshape(tf.einsum('ij,ij->i', x ** 2, x_reverse ** 2), (3, 1))
x_axis1 = 1
x_axis2 = 0
expected_hessian = [[8.469135], [.22222222], [-27.777779]]
actual_hessian = diff.batch_hessian(
x, y, x_axis1, x_axis2, CoordinateSystem.CYLINDRICAL).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_spherical_hessian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.constant([[1., 2., 3.], [4., 5., 6.]], dtype=tf.float32)
diff.watch(x)
c = tf.constant([[1., -2., 1.5], [-4., -5., 4.5]], dtype=tf.float32)
y = c * x ** 2
x_axis1 = x_axis2 = 1
expected_hessian = [
[2., -200.85509, -63.13727],
[-8., -8.005327, -11.597692]
]
actual_hessian = diff.batch_hessian(
x, y, x_axis1, x_axis2, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(actual_hessian, expected_hessian)
def test_polar_vector_laplacian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((3, 2), dtype=tf.float32)
diff.watch(x)
c = tf.constant(
[[1., 2.], [-4., -5.], [0., -2.]],
dtype=tf.float32)
y = tf.concat(
[
tf.reduce_sum(c * x ** 3, axis=1, keepdims=True),
tf.reduce_sum(c * x ** 2, axis=1, keepdims=True)
], axis=1)
expected_vector_laplacian = [[7.], [-25.], [-2.]]
actual_vector_laplacian = diff.batch_vector_laplacian(
x, y, 0, CoordinateSystem.POLAR).numpy()
assert np.allclose(expected_vector_laplacian, actual_vector_laplacian)
def test_spherical_vector_laplacian():
with AutoDifferentiator(persistent=True) as diff:
x = tf.ones((3, 3), dtype=tf.float32)
diff.watch(x)
c = tf.constant(
[[1., 2., 3.], [-4., -5., -6.], [0., -2., 4.]],
dtype=tf.float32)
y = tf.concat(
[
tf.reduce_sum(c * x ** 3, axis=1, keepdims=True),
tf.reduce_sum(c * x ** 2, axis=1, keepdims=True),
tf.reduce_sum(x, axis=1, keepdims=True)
], axis=1)
expected_vector_laplacian = [[-15.359716], [17.915348], [3.6546054]]
actual_vector_laplacian = diff.batch_vector_laplacian(
x, y, 1, CoordinateSystem.SPHERICAL).numpy()
assert np.allclose(expected_vector_laplacian, actual_vector_laplacian)