diff --git a/examples/wave_equation/README.md b/examples/wave_equation/README.md new file mode 100644 index 0000000..df5a2b8 --- /dev/null +++ b/examples/wave_equation/README.md @@ -0,0 +1,27 @@ +## 1D Wave Equation +Given the wave equation: + +$$ u_{tt} = c^2 u_{xx}, $$ + +with wave speed $c = 1$, domain $x \in [0, 1]$ and $t \in [0, 1]$, and the initial and boundary conditions: +$$u(0, x) = \sin(\pi x), $$ +$$u_t(0, x) = 0, $$ +$$u(t, 0) = 0,$$ +$$u(t, 1) = 0.$$ + +The exact solution is $u(x, t) = \sin(\pi x) \cos(\pi t)$. + +### Problem Setup + +| 1D Wave Equation | | +|------------------------------|---| +| PDE equation | $f = u_{tt} - c^2 u_{xx}$ | +| Initial condition (displacement) | $u(0, x) = \sin(\pi x)$ | +| Initial condition (velocity) | $u_t(0, x) = 0$ | +| Dirichlet boundary conditions | $u(t, 0) = u(t, 1) = 0$ | +| The output of net | $[u(t, x)]$ | +| Layers of net | $[2] + 5 \times [50] + [1]$ | +| Sample count from collection points | $10000$ | +| Sample count from initial conditions | $100$ | +| Sample count from boundary conditions | $100$ | +| Loss function | $\text{MSE}_0 + \text{MSE}_b + \text{MSE}_c$ | diff --git a/examples/wave_equation/configs/config.yaml b/examples/wave_equation/configs/config.yaml new file mode 100644 index 0000000..9aa3332 --- /dev/null +++ b/examples/wave_equation/configs/config.yaml @@ -0,0 +1,90 @@ +defaults: + - train + - _self_ + +N0: 100 +N_b: 100 +N_f: 10_000 + +time_domain: + _target_: pinnstf2.data.TimeDomain + t_interval: [0, 1.0] + t_points: 201 + +spatial_domain: + _target_: pinnstf2.data.Interval + x_interval: [0, 1] + shape: [256, 1] + +mesh: + _target_: pinnstf2.data.Mesh + root_dir: null + read_data_fn: ??? + ub: [1.0, 1.0] + lb: [0.0, 0.0] + +train_datasets: + - mesh_sampler: + _target_: pinnstf2.data.MeshSampler + _partial_: true + num_sample: ${N_f} + collection_points: + - f + + - initial_condition: + _target_: pinnstf2.data.InitialCondition + _partial_: true + num_sample: ${N0} + solution: + - u + - v + + - dirichlet_boundary_condition: + _target_: pinnstf2.data.DirichletBoundaryCondition + _partial_: true + num_sample: ${N_b} + solution: + - u + +val_dataset: + - mesh_sampler: + _target_: pinnstf2.data.MeshSampler + _partial_: true + solution: + - u + +pred_dataset: + - mesh_sampler: + _target_: pinnstf2.data.MeshSampler + _partial_: true + solution: + - u + +net: + _target_: pinnstf2.models.FCN + layers: [2, 50, 50, 50, 50, 50, 1] + output_names: + - u + +trainer: + max_epochs: 20000 + check_val_every_n_epoch: 20001 + +model: + loss_fn: mse + +train: true +val: true +test: false +optimized_metric: + error: + - u + +plotting: null + +seed: 1234 +task_name: wave_equation + +hydra: + searchpath: + - pkg://pinnstf2/conf diff --git a/examples/wave_equation/train.py b/examples/wave_equation/train.py new file mode 100644 index 0000000..c7b6112 --- /dev/null +++ b/examples/wave_equation/train.py @@ -0,0 +1,91 @@ +from typing import Dict, Optional + +import hydra +import numpy as np +import tensorflow as tf +from omegaconf import DictConfig + +import pinnstf2 + + +C = 1.0 # wave speed + + +def read_data_fn(root_path): + """Generate analytical reference data for the 1D wave equation. + + Exact solution: u(x, t) = sin(pi*x) * cos(pi*c*t) + satisfying u_tt = c^2 * u_xx with u(0,t) = u(1,t) = 0. + + :param root_path: The root directory containing the data (unused for analytical data). + :return: Dictionary with exact solution array. + """ + + nx, nt = 256, 201 + x = np.linspace(0, 1, nx) + t = np.linspace(0, 1, nt) + X, T = np.meshgrid(x, t, indexing="ij") + exact_u = np.sin(np.pi * X) * np.cos(np.pi * C * T) + return {"u": exact_u} + + +def pde_fn( + outputs: Dict[str, tf.Tensor], + x: tf.Tensor, + t: tf.Tensor, +): + """Define the wave equation PDE residual: u_tt - c^2 * u_xx = 0. + + :param outputs: Dictionary of network outputs. + :param x: Spatial coordinate tensor. + :param t: Temporal coordinate tensor. + :return: Updated outputs with PDE residual. + """ + + u_x, u_t = pinnstf2.utils.gradient(outputs["u"], [x, t]) + u_xx = pinnstf2.utils.gradient(u_x, x) + u_tt = pinnstf2.utils.gradient(u_t, t) + outputs["f"] = u_tt - C**2 * u_xx + return outputs + + +def output_fn( + outputs: Dict[str, tf.Tensor], + x: tf.Tensor, + t: tf.Tensor, +): + """Compute velocity field for initial velocity enforcement. + + :param outputs: Dictionary of network outputs. + :param x: Spatial coordinate tensor. + :param t: Temporal coordinate tensor. + :return: Updated outputs with velocity field. + """ + + outputs["v"] = pinnstf2.utils.gradient(outputs["u"], t) + return outputs + + +@hydra.main(version_base="1.3", config_path="configs", config_name="config.yaml") +def main(cfg: DictConfig) -> Optional[float]: + """Main entry point for training. + + :param cfg: DictConfig configuration composed by Hydra. + :return: Optional[float] with optimized metric value. + """ + + pinnstf2.utils.extras(cfg) + + metric_dict, _ = pinnstf2.train( + cfg, read_data_fn=read_data_fn, pde_fn=pde_fn, output_fn=output_fn + ) + + metric_value = pinnstf2.utils.get_metric_value( + metric_dict=metric_dict, metric_names=cfg.get("optimized_metric") + ) + + return metric_value + + +if __name__ == "__main__": + main()