Spaces:
Sleeping
Sleeping
| import unittest | |
| import numpy as np | |
| from stnn.pde.common import d_dx_upwind, d2_dx2_fourth_order, d_dx_upwind_nonuniform | |
| def assert_derivative(operator, function, expected_derivative, boundary = None, rtol = None, atol = None): | |
| """ | |
| Assert the derivative of a function using a given finite-difference operator | |
| Args: | |
| operator (np.ndarray or sparse matrix): Differential operator matrix. | |
| function (np.ndarray): Values of the function to differentiate | |
| expected_derivative (np.ndarray):: Expected result of the derivative. | |
| boundary (int): Specifies if boundary elements should be excluded, and how many. | |
| rtol: Relative tolerance. | |
| atol: Absolute tolerance. | |
| """ | |
| observed_derivative = operator @ function | |
| if boundary is not None: | |
| observed_derivative = observed_derivative[boundary:-boundary] | |
| expected_derivative = expected_derivative[boundary:-boundary] | |
| np.testing.assert_allclose(observed_derivative, expected_derivative, rtol = rtol, atol = atol) | |
| class TestDifferentialOperators(unittest.TestCase): | |
| def setUp(self): | |
| # ----- Set up grids | |
| # grid for radial coordinate (non-periodic, non-uniform spacing) | |
| self.nx = 100000 | |
| zs = np.linspace(-np.pi / 2, np.pi / 2, self.nx + 2) | |
| R1, R2 = 0.5, 1.4 | |
| self.z_ = (R2 - R1) * (np.sin(zs) / 2 + 0.5) + R1 | |
| self.z = self.z_[1:-1] | |
| # grid for angular coordinates (periodic, uniform spacing) | |
| self.y = np.linspace(0, 2 * np.pi, self.nx, endpoint = False) | |
| dy = (2 * np.pi) / self.nx | |
| # ----- Set tolerances | |
| # Tolerances for "exact" tests, i.e., where the finite differences do not have truncation error | |
| self.atol = 1e-6 | |
| self.rtol = 1e-6 | |
| # Tolerances for some inexact tests | |
| self.atol_inexact = 1e-3 | |
| self.rtol_z = 3 * np.max(np.diff(self.z_)) # relative tolerance of 3*dx for first-order one-sided differences | |
| self.rtol_y = 3 * dy # relative tolerance of 3*dx for first-order one-sided differences | |
| self.rtol_y2 = 3 * dy**4 # relative tolerance of 3*dy**4 for fourth-order central differences | |
| # ----- Test inputs | |
| # "Exact" test inputs | |
| self.f1 = -2.3 * np.ones(self.nx) | |
| self.f2 = 0.8 * self.z | |
| self.f3 = 0.8 * self.y | |
| self.f4 = -0.1 * self.y * self.y | |
| # Inexact test inputs | |
| # noinspection PyRedundantParentheses | |
| self.g1 = (self.z)**2 - (self.z)**3 | |
| self.g2 = (np.cos(self.y))**2 - (np.sin(self.y))**3 | |
| self.g3 = (np.sin(2 * self.y))**2 - (np.cos(self.y))**3 | |
| def test_d_dx_upwind(self): | |
| dx_m, dx_p = d_dx_upwind(self.y, self.nx) | |
| assert_derivative(dx_m, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_p, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_m, self.f3, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_p, self.f3, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| expected_dg2 = -2 * np.cos(self.y) * np.sin(self.y) - 3 * np.sin(self.y)**2 * np.cos(self.y) | |
| assert_derivative(dx_m, self.g2, expected_dg2, rtol = self.rtol_y, atol = self.atol_inexact) | |
| assert_derivative(dx_p, self.g2, expected_dg2, rtol = self.rtol_y, atol = self.atol_inexact) | |
| expected_dg3 = 4 * np.sin(2 * self.y) * np.cos(2 * self.y) + 3 * np.cos(self.y)**2 * np.sin(self.y) | |
| assert_derivative(dx_m, self.g3, expected_dg3, rtol = self.rtol_y, atol = self.atol_inexact) | |
| assert_derivative(dx_p, self.g3, expected_dg3, rtol = self.rtol_y, atol = self.atol_inexact) | |
| def test_d2_dx2_fourth_order(self): | |
| d2x = d2_dx2_fourth_order(self.y, self.nx) | |
| assert_derivative(d2x, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol) | |
| assert_derivative(d2x, self.f3, np.zeros(self.nx), boundary = 4, rtol = self.rtol, atol = self.atol) | |
| assert_derivative(d2x, self.f4, -0.2 * np.ones(self.nx), boundary = 4, rtol = self.rtol, atol = self.atol) | |
| expected_dg2 = 2 * np.sin(self.y)**2 - 2 * np.cos(self.y)**2 - \ | |
| 6 * np.sin(self.y) * np.cos(self.y)**2 + 3 * np.sin(self.y)**3 | |
| assert_derivative(d2x, self.g2, expected_dg2, rtol = self.rtol_y2, atol = self.atol_inexact) | |
| expected_dg3 = 8 * np.cos(2 * self.y)**2 - 8 * np.sin(2 * self.y)**2 - \ | |
| 6 * np.cos(self.y) * np.sin(self.y)**2 + 3 * np.cos(self.y)**3 | |
| assert_derivative(d2x, self.g3, expected_dg3, rtol = self.rtol_y2, atol = self.atol_inexact) | |
| def test_d_dx_upwind_nonuniform(self): | |
| dx_minus, dx_plus = d_dx_upwind_nonuniform(self.z_, self.nx) | |
| assert_derivative(dx_minus, self.f1, np.zeros(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_plus, self.f1, np.zeros(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_minus, self.f2, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| assert_derivative(dx_plus, self.f2, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol) | |
| expected_dg1 = (2 * self.z - 3 * self.z**2) | |
| assert_derivative(dx_minus, self.g1, expected_dg1, boundary = 1, rtol = self.rtol_z, atol = self.atol_inexact) | |
| assert_derivative(dx_plus, self.g1, expected_dg1, boundary = 1, rtol = self.rtol_z, atol = self.atol_inexact) | |
| if __name__ == '__main__': | |
| unittest.main() | |