def test_polar_stream_plot(): file_path = 'polar_stream_plot' mesh = Mesh([(.1, 1.1), (0., 2 * np.pi)], [.2, .4 * np.pi], CoordinateSystem.POLAR) plot = StreamPlot(np.random.rand(2, 5, 5, 2), mesh, False) plot.save(file_path).close() os.remove(f'{file_path}.gif')
def test_fdm_operator_on_2d_pde(): diff_eq = NavierStokesEquation(5000.) mesh = Mesh([(0., 10.), (0., 10.)], [1., 1.]) bcs = [(DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (1., .1, None, None)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (0., 0., None, None)), is_static=True)), (DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (0., 0., None, None)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (0., 0., None, None)), is_static=True))] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, lambda x: np.zeros((len(x), 4))) ivp = InitialValueProblem(cp, (0., 10.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .25) solution = op.solve(ivp) assert solution.vertex_oriented assert solution.d_t == .25 assert solution.discrete_y().shape == (40, 11, 11, 4) assert solution.discrete_y(False).shape == (40, 10, 10, 4)
def test_spherical_quiver_plot(): file_path = 'spherical_quiver_plot' mesh = Mesh([(.1, 1.1), (0., 2 * np.pi), (.05 * np.pi, .95 * np.pi)], [.2, .4 * np.pi, .1 * np.pi], CoordinateSystem.SPHERICAL) plot = QuiverPlot(np.random.rand(2, 6, 6, 10, 3), mesh, True) plot.save(file_path).close() os.remove(f'{file_path}.gif')
def test_discrete_initial_condition_2d_pde(): diff_eq = WaveEquation(2) mesh = Mesh([(0., 2.), (0., 2.)], [1., 1.]) bcs = [(DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (0., 2.)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (1., 2.)), is_static=True)), (DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (3., 2.)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (4., 2.)), is_static=True))] cp = ConstrainedProblem(diff_eq, mesh, bcs) initial_condition = DiscreteInitialCondition(cp, np.zeros((3, 3, 2)), True) y = initial_condition.y_0(np.array([1.5, .5]).reshape((1, 2))) assert np.allclose(y, [1.75, 1.5]) y_0_vertices = initial_condition.discrete_y_0(True) assert y_0_vertices.shape == (3, 3, 2) assert np.all(y_0_vertices[0, 1:-1, 0] == 0.) assert np.all(y_0_vertices[0, 1:-1, 1] == 2.) assert np.all(y_0_vertices[-1, 1:-1, 0] == 1.) assert np.all(y_0_vertices[-1, 1:-1, 1] == 2.) assert np.all(y_0_vertices[:, 0, 0] == 3.) assert np.all(y_0_vertices[:, 0, 1] == 2.) assert np.all(y_0_vertices[:, -1, 0] == 4.) assert np.all(y_0_vertices[:, -1, 1] == 2.) assert np.all(y_0_vertices[1:-1, 1:-1, :] == 0.) y_0_cell_centers = initial_condition.discrete_y_0(False) assert y_0_cell_centers.shape == (2, 2, 2)
def test_continuous_initial_condition_pde_with_wrong_shape(): diff_eq = WaveEquation(1) mesh = Mesh([(0., 10.)], [1.]) bcs = [(DirichletBoundaryCondition(lambda x: np.zeros((len(x), 2))), ) * 2] cp = ConstrainedProblem(diff_eq, mesh, bcs) with pytest.raises(ValueError): ContinuousInitialCondition(cp, lambda x: np.zeros((3, 2)))
def test_beta_initial_condition_with_wrong_number_of_alpha_and_betas(): diff_eq = DiffusionEquation(1) mesh = Mesh([(0., 1.)], [.1]) bcs = [(NeumannBoundaryCondition(lambda x: np.zeros((len(x), 1))), ) * 2] cp = ConstrainedProblem(diff_eq, mesh, bcs) with pytest.raises(ValueError): BetaInitialCondition(cp, [(1., 1.), (1., 1)])
def test_discrete_initial_condition_pde_with_wrong_shape(): diff_eq = WaveEquation(1) mesh = Mesh([(0., 10.)], [1.]) bcs = [(DirichletBoundaryCondition(lambda x: np.zeros((len(x), 2))), ) * 2] cp = ConstrainedProblem(diff_eq, mesh, bcs) with pytest.raises(ValueError): DiscreteInitialCondition(cp, np.zeros((10, 2)), vertex_oriented=True)
def test_gaussian_initial_condition_2d_pde(): diff_eq = WaveEquation(2) mesh = Mesh([(0., 2.), (0., 2.)], [1., 1.]) bcs = [(DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (0., 2.)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (1., 2.)), is_static=True)), (DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (3., 2.)), is_static=True), DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: (4., 2.)), is_static=True))] cp = ConstrainedProblem(diff_eq, mesh, bcs) initial_condition = GaussianInitialCondition(cp, [ (np.array([1., 1.]), np.array([[1., 0.], [0., 1.]])), (np.array([1., 1.]), np.array([[.75, .25], [.25, .75]])), ], [1., 2.]) x_coordinates = np.array([[1., 1.], [.5, 1.5]]) expected_y_0 = [[.15915494, .45015816], [.12394999, .27303472]] actual_y_0 = initial_condition.y_0(x_coordinates) assert np.allclose(actual_y_0, expected_y_0) expected_vertex_discrete_y_0 = [[[3., 2.], [0., 2.], [4., 2.]], [[3., 2.], [.15915494, .45015816], [4., 2.]], [[3., 2.], [1., 2.], [4., 2.]]] actual_vertex_discrete_y_0 = initial_condition.discrete_y_0(True) assert np.allclose(actual_vertex_discrete_y_0, expected_vertex_discrete_y_0) expected_cell_discrete_y_0 = [[[.12394999, .35058353], [.12394999, .27303472]], [[.12394999, .27303472], [.12394999, .35058353]]] actual_cell_discrete_y_0 = initial_condition.discrete_y_0(False) assert np.allclose(actual_cell_discrete_y_0, expected_cell_discrete_y_0)
def test_gaussian_initial_condition_pde_with_wrong_means_and_cov_length(): diff_eq = WaveEquation(1) mesh = Mesh([(0., 10.)], [1.]) bcs = [(DirichletBoundaryCondition(lambda x: np.zeros((len(x), 2))), ) * 2] cp = ConstrainedProblem(diff_eq, mesh, bcs) with pytest.raises(ValueError): GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 1)
def test_continuous_initial_condition_1d_pde(): diff_eq = DiffusionEquation(1) mesh = Mesh([(0., 20.)], [.1]) bcs = [(DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True), DirichletBoundaryCondition(lambda x, t: np.full((len(x), 1), 1.5), is_static=True))] cp = ConstrainedProblem(diff_eq, mesh, bcs) initial_condition = ContinuousInitialCondition( cp, lambda x: np.exp(-np.square(np.array(x) - 10.) / (2 * 5**2))) assert np.isclose(initial_condition.y_0(np.full((1, 1), 10.)), 1.) assert np.isclose( initial_condition.y_0(np.full((1, 1), np.sqrt(50) + 10.)), np.e**-1) assert np.allclose(initial_condition.y_0(np.full((5, 1), 10.)), np.ones((5, 1))) y_0_vertices = initial_condition.discrete_y_0(True) assert y_0_vertices.shape == (201, 1) assert y_0_vertices[0, 0] == 0. assert y_0_vertices[-1, 0] == 1.5 assert y_0_vertices[100, 0] == 1. assert np.all(0. < y_0_vertices[1:100, 0]) \ and np.all(y_0_vertices[1:100, 0] < 1.) assert np.all(0. < y_0_vertices[101:-1, 0]) \ and np.all(y_0_vertices[101:-1, 0] < 1.) y_0_cell_centers = initial_condition.discrete_y_0(False) assert y_0_cell_centers.shape == (200, 1) assert np.all(0. < y_0_cell_centers) and np.all(y_0_cell_centers < 1.)
def test_fdm_operator_on_pde_with_t_and_x_dependent_rhs(): class TestDiffEq(DifferentialEquation): def __init__(self): super(TestDiffEq, self).__init__(2, 1) @property def symbolic_equation_system(self) -> SymbolicEquationSystem: return SymbolicEquationSystem([ self.symbols.t / 100. * (self.symbols.x[0] + self.symbols.x[1])**2 ]) diff_eq = TestDiffEq() mesh = Mesh([(-5., 5.), (0., 3.)], [2., 1.]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))))] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, lambda x: np.zeros((len(x), 1))) ivp = InitialValueProblem(cp, (0., 5.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .25) solution = op.solve(ivp) y = solution.discrete_y() assert solution.vertex_oriented assert solution.d_t == .25 assert y.shape == (20, 6, 4, 1)
def test_pidon_operator_on_spherical_pde(): set_random_seed(0) diff_eq = DiffusionEquation(3) mesh = Mesh( [(1., 11.), (0., 2 * np.pi), (.25 * np.pi, .75 * np.pi)], [2., np.pi / 5., np.pi / 4], CoordinateSystem.SPHERICAL) bcs = [ (DirichletBoundaryCondition( lambda x, t: np.ones((len(x), 1)), is_static=True), DirichletBoundaryCondition( lambda x, t: np.full((len(x), 1), 1. / 11.), is_static=True)), (NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 1)), is_static=True), NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 1)), is_static=True)), (NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 1)), is_static=True), NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 1)), is_static=True)) ] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, lambda x: 1. / x[:, :1]) t_interval = (0., .5) ivp = InitialValueProblem(cp, t_interval, ic) sampler = UniformRandomCollocationPointSampler() pidon = PIDONOperator(sampler, .001, True) training_loss_history, test_loss_history = pidon.train( cp, t_interval, training_data_args=DataArgs( y_0_functions=[ic.y_0], n_domain_points=20, n_boundary_points=10, n_batches=1 ), model_args=ModelArgs( latent_output_size=20, branch_hidden_layer_sizes=[30, 30], trunk_hidden_layer_sizes=[30, 30], ), optimization_args=OptimizationArgs( optimizer=optimizers.Adam(learning_rate=2e-5), epochs=3, verbose=False ) ) assert len(training_loss_history) == 3 for i in range(2): assert np.all( training_loss_history[i + 1].weighted_total_loss.numpy() < training_loss_history[i].weighted_total_loss.numpy()) solution = pidon.solve(ivp) assert solution.d_t == .001 assert solution.discrete_y().shape == (500, 6, 11, 3, 1)
def test_solution_generate_plots_for_2d_pde_with_scalar_and_vector_fields(): diff_eq = ShallowWaterEquation(.5) mesh = Mesh([(0., 5.), (0., 5.)], [1., 1.]) bcs = [(NeumannBoundaryCondition( vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True), NeumannBoundaryCondition( vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True))] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition( cp, [(np.array([2.5, 1.25]), np.array([[1., 0.], [0., 1.]]))] * 3, [1., .0, .0]) ivp = InitialValueProblem(cp, (0., 20.), ic) t_coordinates = np.array([10., 20.]) discrete_y = np.arange(216).reshape((2, 6, 6, 3)) solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True) plots = list(solution.generate_plots()) try: assert len(plots) == 4 assert isinstance(plots[0], QuiverPlot) assert isinstance(plots[1], StreamPlot) assert isinstance(plots[2], ContourPlot) assert isinstance(plots[3], SurfacePlot) finally: for plot in plots: plot.close()
def test_beta_initial_condition_with_more_than_1d_pde(): diff_eq = DiffusionEquation(2) mesh = Mesh([(0., 1.), (0., 1.)], [.1, .1]) bcs = [(NeumannBoundaryCondition(lambda x: np.zeros((len(x), 1))), ) * 2 ] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) with pytest.raises(ValueError): BetaInitialCondition(cp, [(1., 1.), (1., 1)])
def test_pidon_operator_on_pde_system(): set_random_seed(0) diff_eq = NavierStokesEquation() mesh = Mesh([(-2.5, 2.5), (0., 4.)], [1., 1.]) ic_function = vectorize_ic_function(lambda x: [ 2. * x[0] - 4., 2. * x[0] ** 2 + 3. * x[1] - x[0] * x[1] ** 2, 4. * x[0] - x[1] ** 2, 2. * x[0] * x[1] - 3. ]) bcs = [ (DirichletBoundaryCondition( lambda x, t: ic_function(x), is_static=True), DirichletBoundaryCondition( lambda x, t: ic_function(x), is_static=True)) ] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, ic_function) t_interval = (0., .5) ivp = InitialValueProblem(cp, t_interval, ic) sampler = UniformRandomCollocationPointSampler() pidon = PIDONOperator(sampler, .001, True) training_loss_history, test_loss_history = pidon.train( cp, t_interval, training_data_args=DataArgs( y_0_functions=[ic.y_0], n_domain_points=20, n_boundary_points=10, n_batches=1 ), model_args=ModelArgs( latent_output_size=20, branch_hidden_layer_sizes=[20, 20], trunk_hidden_layer_sizes=[20, 20], ), optimization_args=OptimizationArgs( optimizer=optimizers.Adam(learning_rate=1e-5), epochs=3, verbose=False ) ) assert len(training_loss_history) == 3 for i in range(2): assert np.all( training_loss_history[i + 1].weighted_total_loss.numpy() < training_loss_history[i].weighted_total_loss.numpy()) solution = pidon.solve(ivp) assert solution.d_t == .001 assert solution.discrete_y().shape == (500, 6, 5, 4)
def test_fdm_operator_on_pde_with_t_and_x_dependent_rhs(): class TestDiffEq(DifferentialEquation): def __init__(self): super(TestDiffEq, self).__init__(2, 1) @property def symbolic_equation_system(self) -> SymbolicEquationSystem: return SymbolicEquationSystem([ self.symbols.t / 100. * (self.symbols.x[0] + self.symbols.x[1]) ** 2 ]) diff_eq = TestDiffEq() mesh = Mesh([(-1., 1.), (0., 2.)], [2., 1.]) bcs = [ (NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)))) ] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, lambda x: np.zeros((len(x), 1))) t_interval = (0., 1.) ivp = InitialValueProblem(cp, t_interval, ic) sampler = UniformRandomCollocationPointSampler() pidon = PIDONOperator(sampler, .05, True) training_loss_history, test_loss_history = pidon.train( cp, t_interval, training_data_args=DataArgs( y_0_functions=[ic.y_0], n_domain_points=20, n_boundary_points=10, n_batches=1 ), model_args=ModelArgs( latent_output_size=20, branch_hidden_layer_sizes=[30, 30], trunk_hidden_layer_sizes=[30, 30], ), optimization_args=OptimizationArgs( optimizer=optimizers.Adam(learning_rate=2e-5), epochs=3, verbose=False ) ) assert len(training_loss_history) == 3 for i in range(2): assert np.all( training_loss_history[i + 1].weighted_total_loss.numpy() < training_loss_history[i].weighted_total_loss.numpy()) solution = pidon.solve(ivp) assert solution.d_t == .05 assert solution.discrete_y().shape == (20, 2, 3, 1)
def test_cp_pde_with_wrong_boundary_constraint_width(): diff_eq = WaveEquation(2) mesh = Mesh([(0., 5.), (-5., 5.)], [.1, .2]) bcs = [(DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True), ) * 2] * 2 with pytest.raises(ValueError): ConstrainedProblem(diff_eq, mesh, bcs) bcs = [(DirichletBoundaryCondition( vectorize_bc_function(lambda x, t: [0.]), is_static=True), ) * 2] * 2 with pytest.raises(ValueError): ConstrainedProblem(diff_eq, mesh, bcs)
def test_mesh(): x_intervals = [(-10., 10.), (0., 50.)] d_x = [.1, .2] mesh = Mesh(x_intervals, d_x) assert mesh.vertices_shape == (201, 251) assert mesh.cells_shape == (200, 250) assert mesh.shape(True) == mesh.vertices_shape assert mesh.shape(False) == mesh.cells_shape for axis in range(2): assert np.array_equal( mesh.vertex_axis_coordinates[axis], np.linspace(*x_intervals[axis], mesh.vertices_shape[axis])) assert np.array_equal( mesh.cell_center_axis_coordinates[axis], np.linspace(x_intervals[axis][0] + d_x[axis] / 2., x_intervals[axis][1] - d_x[axis] / 2., mesh.cells_shape[axis])) assert np.array_equal( mesh.axis_coordinates(True)[axis], mesh.vertex_axis_coordinates[axis]) assert np.array_equal( mesh.axis_coordinates(False)[axis], mesh.cell_center_axis_coordinates[axis]) all_vertex_x = mesh.all_index_coordinates(True) assert np.allclose(all_vertex_x[2, 3], (-9.8, .6)) all_vertex_x_flattened = mesh.all_index_coordinates(True, True) assert np.array_equal(all_vertex_x[2, 3], all_vertex_x_flattened[2 * 251 + 3]) all_cell_x = mesh.all_index_coordinates(False) assert np.allclose(all_cell_x[2, 3], (-9.75, .7)) all_cell_x_flattened = mesh.all_index_coordinates(False, True) assert np.array_equal(all_cell_x[2, 3], all_cell_x_flattened[2 * 250 + 3])
def test_ode_operator_on_pde(): diff_eq = DiffusionEquation(1, 1.5) mesh = Mesh([(0., 10.)], [.1]) bcs = [ (NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))), DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 1)))), ] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition(cp, [(np.array([5.]), np.array([[2.5]]))], [20.]) ivp = InitialValueProblem(cp, (0., 10.), ic) op = ODEOperator('RK23', 2.5e-3) with pytest.raises(ValueError): op.solve(ivp)
def test_cp_pde_with_wrong_boundary_constraint_length(): diff_eq = DiffusionEquation(2) mesh = Mesh([(0., 5.), (-5., 5.)], [.1, .2]) static_bcs = [(DirichletBoundaryCondition(lambda x, t: np.zeros((13, 1)), is_static=True), ) * 2] * 2 with pytest.raises(ValueError): ConstrainedProblem(diff_eq, mesh, static_bcs) dynamic_bcs = [ (DirichletBoundaryCondition(lambda x, t: np.zeros((13, 1))), ) * 2 ] * 2 cp = ConstrainedProblem(diff_eq, mesh, dynamic_bcs) with pytest.raises(ValueError): cp.create_boundary_constraints(True, 0.)
def test_pidon_operator_in_ar_mode_on_pde(): set_random_seed(0) diff_eq = WaveEquation(1) mesh = Mesh([(0., 1.)], (.2,)) bcs = [ (NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 2)), is_static=True), NeumannBoundaryCondition( lambda x, t: np.zeros((len(x), 2)), is_static=True)), ] cp = ConstrainedProblem(diff_eq, mesh, bcs) t_interval = (0., 1.) ic = BetaInitialCondition(cp, [(3.5, 3.5), (3.5, 3.5)]) ivp = InitialValueProblem(cp, t_interval, ic) training_y_0_functions = [ BetaInitialCondition(cp, [(p, p), (p, p)]).y_0 for p in [2., 3., 4., 5.] ] sampler = UniformRandomCollocationPointSampler() pidon = PIDONOperator(sampler, .25, False, auto_regression_mode=True) assert pidon.auto_regression_mode pidon.train( cp, (0., .25), training_data_args=DataArgs( y_0_functions=training_y_0_functions, n_domain_points=50, n_boundary_points=20, n_batches=2 ), model_args=ModelArgs( latent_output_size=50, branch_hidden_layer_sizes=[50, 50], trunk_hidden_layer_sizes=[50, 50], ), optimization_args=OptimizationArgs( optimizer=optimizers.Adam(learning_rate=1e-4), epochs=2, ic_loss_weight=10., verbose=False ) ) sol = pidon.solve(ivp) assert np.allclose(sol.t_coordinates, [.25, .5, .75, 1.]) assert sol.discrete_y().shape == (4, 5, 2)
def test_solution_pde_with_no_vertex_orientation_defined(): diff_eq = WaveEquation(1) mesh = Mesh([(0., 2.)], [1.]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 2) ivp = InitialValueProblem(cp, (0., 2.), ic) t_coordinates = np.array([1., 2.]) discrete_y = np.arange(12).reshape((2, 3, 2)) with pytest.raises(ValueError): Solution(ivp, t_coordinates, discrete_y)
def test_fdm_operator_on_1d_pde(): diff_eq = BurgerEquation(1, 1000.) mesh = Mesh([(0., 10.)], [.1]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True))] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition(cp, [(np.array([2.5]), np.array([[1.]]))]) ivp = InitialValueProblem(cp, (0., 50.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .25) solution = op.solve(ivp) assert solution.vertex_oriented assert solution.d_t == .25 assert solution.discrete_y().shape == (200, 101, 1) assert solution.discrete_y(False).shape == (200, 100, 1)
def _verify_pde_solution_shape_matches_problem( y: np.ndarray, mesh: Mesh, vertex_oriented: bool, expected_x_dims: Union[int, Tuple[int, int]], is_vector_field: bool): """ Verifies that the shape of the input array representing the solution of a partial differential equation over the provided mesh and with the specified vertex orientation matches expectations. :param y: an array representing the solution of the partial differential equation :param mesh: the spatial mesh over which the solution is evaluated :param vertex_oriented: whether the solution is evaluated over the vertices or the cell centers of the mesh :param expected_x_dims: the expected number of spatial dimensions :param is_vector_field: whether the solution is supposed to be a vector field or a scalar field """ if isinstance(expected_x_dims, int): if mesh.dimensions != expected_x_dims: raise ValueError(f'mesh must be {expected_x_dims} dimensional') elif not (expected_x_dims[0] <= mesh.dimensions <= expected_x_dims[1]): raise ValueError( f'mesh must be between {expected_x_dims[0]} and ' f'{expected_x_dims[1]} dimensional') if y.ndim != mesh.dimensions + 2: raise ValueError( f'number of y axes ({y.ndim}) must be two larger than mesh ' f'dimensions ({mesh.dimensions})') if y.shape[1:-1] != mesh.shape(vertex_oriented): raise ValueError( f'y shape {y.shape} must be compatible with mesh shape ' f'{mesh.shape(vertex_oriented)}') if is_vector_field: if y.shape[-1] != mesh.dimensions: raise ValueError( f'number of y components ({y.shape[-1]}) must match ' f'x dimensions {mesh.dimensions}') elif y.shape[-1] != 1: raise ValueError( f'number of y components ({y.shape[-1]}) must be one')
def test_fdm_operator_on_spherical_pde(): diff_eq = DiffusionEquation(3) mesh = Mesh([(1., 11.), (0., 2. * np.pi), (.1 * np.pi, .9 * np.pi)], [2., np.pi / 5., np.pi / 5], CoordinateSystem.SPHERICAL) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)), is_static=True))] * 3 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = ContinuousInitialCondition(cp, lambda x: 1. / x[:, :1]) ivp = InitialValueProblem(cp, (0., 5.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .1) solution = op.solve(ivp) assert solution.vertex_oriented assert solution.d_t == .1 assert solution.discrete_y().shape == (50, 6, 11, 5, 1) assert solution.discrete_y(False).shape == (50, 5, 10, 4, 1)
def test_fdm_operator_on_3d_pde(): diff_eq = CahnHilliardEquation(3) mesh = Mesh([(0., 5.), (0., 5.), (0., 10.)], [.5, 1., 2.]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)), is_static=True), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)), is_static=True))] * 3 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = DiscreteInitialCondition( cp, .05 * np.random.uniform(-1., 1., cp.y_shape(True)), True) ivp = InitialValueProblem(cp, (0., 5.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .05) solution = op.solve(ivp) assert solution.vertex_oriented assert solution.d_t == .05 assert solution.discrete_y().shape == (100, 11, 6, 6, 2) assert solution.discrete_y(False).shape == (100, 10, 5, 5, 2)
def test_cp_1d_pde(): diff_eq = DiffusionEquation(1) mesh = Mesh([(0., 1.)], [.1]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((1, 1)), is_static=True), ) * 2] cp = ConstrainedProblem(diff_eq, mesh, bcs) assert cp.are_all_boundary_conditions_static assert not cp.are_there_boundary_conditions_on_y assert cp.y_shape(True) == (11, 1) assert cp.y_shape(False) == (10, 1) assert cp.differential_equation == diff_eq assert cp.mesh == mesh assert np.array_equal(cp.boundary_conditions, bcs) y_vertex_constraints = cp.static_y_vertex_constraints assert y_vertex_constraints.shape == (1, ) assert np.all(y_vertex_constraints[0].mask == [False]) assert np.all(y_vertex_constraints[0].values == []) vertex_boundary_constraints = cp.static_boundary_constraints(True) y_vertex_boundary_constraints = vertex_boundary_constraints[0] assert y_vertex_boundary_constraints.shape == (1, 1) assert y_vertex_boundary_constraints[0, 0][0] is None assert y_vertex_boundary_constraints[0, 0][1] is None d_y_vertex_boundary_constraints = vertex_boundary_constraints[1] assert d_y_vertex_boundary_constraints.shape == (1, 1) assert np.all(d_y_vertex_boundary_constraints[0, 0][0].mask == [True]) assert np.all(d_y_vertex_boundary_constraints[0, 0][0].values == [0.]) assert np.all(d_y_vertex_boundary_constraints[0, 0][1].mask == [True]) assert np.all(d_y_vertex_boundary_constraints[0, 0][1].values == [0.]) cell_boundary_constraints = cp.static_boundary_constraints(False) y_cell_boundary_constraints = cell_boundary_constraints[0] assert y_cell_boundary_constraints.shape == (1, 1) assert y_cell_boundary_constraints[0, 0][0] is None assert y_cell_boundary_constraints[0, 0][1] is None d_y_cell_boundary_constraints = cell_boundary_constraints[1] assert d_y_cell_boundary_constraints.shape == (1, 1) assert np.all(d_y_cell_boundary_constraints[0, 0][0].mask == [True]) assert np.all(d_y_cell_boundary_constraints[0, 0][0].values == [0.]) assert np.all(d_y_cell_boundary_constraints[0, 0][1].mask == [True]) assert np.all(d_y_cell_boundary_constraints[0, 0][1].values == [0.])
def test_auto_regression_operator_on_pde(): set_random_seed(0) diff_eq = WaveEquation(2) mesh = Mesh([(-5., 5.), (-5., 5.)], [1., 1.]) bcs = [(DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 2)), is_static=True), DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 2)), is_static=True))] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition( cp, [(np.array([0., 2.5]), np.array([[.1, 0.], [0., .1]]))] * 2, [3., .0]) ivp = InitialValueProblem(cp, (0., 10.), ic) oracle = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .1) ref_solution = oracle.solve(ivp) ml_op = AutoRegressionOperator(2.5, True) ml_op.train( ivp, oracle, SKLearnKerasRegressor( DeepONet([ np.prod(cp.y_shape(True)).item(), 100, 50, diff_eq.y_dimension * 10 ], [1 + diff_eq.x_dimension, 50, 50, diff_eq.y_dimension * 10], diff_eq.y_dimension), optimizer=optimizers.Adam( learning_rate=optimizers.schedules.ExponentialDecay( 1e-2, decay_steps=500, decay_rate=.95)), batch_size=968, epochs=500, ), 20, lambda t, y: y + np.random.normal(0., t / 75., size=y.shape)) ml_solution = ml_op.solve(ivp) assert ml_solution.vertex_oriented assert ml_solution.d_t == 2.5 assert ml_solution.discrete_y().shape == (4, 11, 11, 2) diff = ref_solution.diff([ml_solution]) assert np.all(diff.matching_time_points == np.linspace(2.5, 10., 4)) assert np.max(np.abs(diff.differences[0])) < .5
def test_fdm_operator_on_polar_pde(): diff_eq = ShallowWaterEquation(.5) mesh = Mesh([(1., 11.), (0., 2 * np.pi)], [2., np.pi / 5.], CoordinateSystem.POLAR) bcs = [(NeumannBoundaryCondition( vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True), NeumannBoundaryCondition( vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True))] * 2 cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition( cp, [(np.array([-6., 0.]), np.array([[.25, 0.], [0., .25]]))] * 3, [1., .0, .0]) ivp = InitialValueProblem(cp, (0., 5.), ic) op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .1) solution = op.solve(ivp) assert solution.vertex_oriented assert solution.d_t == .1 assert solution.discrete_y().shape == (50, 6, 11, 3) assert solution.discrete_y(False).shape == (50, 5, 10, 3)
def test_solution_pde(): diff_eq = WaveEquation(1) mesh = Mesh([(0., 2.)], [1.]) bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))), NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))] cp = ConstrainedProblem(diff_eq, mesh, bcs) ic = GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 2) ivp = InitialValueProblem(cp, (0., 2.), ic) t_coordinates = np.array([1., 2.]) discrete_y = np.arange(12).reshape((2, 3, 2)) solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True) assert solution.initial_value_problem == ivp assert np.array_equal(solution.t_coordinates, t_coordinates) assert np.isclose(solution.d_t, 1.) assert solution.vertex_oriented x_coordinates = np.array([[.5], [1.]]) expected_y = [[[1., 2.], [2., 3.]], [[7., 8.], [8., 9.]]] assert np.allclose(solution.y(x_coordinates), expected_y) expected_cell_y = [[[1., 2.], [3., 4.]], [[7., 8.], [9., 10.]]] assert np.allclose(solution.discrete_y(False), expected_cell_y) assert np.allclose(solution.discrete_y(), discrete_y) other_solutions = [ Solution(ivp, np.linspace(.5, 2., 4), np.arange(16).reshape((4, 2, 2)), vertex_oriented=False) ] expected_differences = [ [[3., 3.], [3., 3.], [3., 3.]], [[5., 5.], [5., 5.], [5., 5.]], ] diff = solution.diff(other_solutions) assert np.allclose(diff.matching_time_points, [1., 2.]) assert np.allclose(diff.differences, expected_differences)