diff --git a/README.md b/README.md index a4f4d81..7ce79ed 100644 --- a/README.md +++ b/README.md @@ -144,6 +144,7 @@ python $SCRATCH/PINN/main.py --2d --dist --epochs 100 --bs 17672 --siren Debugging for adding the PML: ``` python $SCRATCH/PINN/main.py --2d --epochs 100 --bs 100 +python $SCRATCH/PINN/main.py --2d --epochs 100 --bs 4900 ``` ## How to run the SLURM script on NERSC diff --git a/main.py b/main.py index 8d266ab..e5d62a5 100644 --- a/main.py +++ b/main.py @@ -56,7 +56,7 @@ def get_args(): parser.add_argument('--lr', type=float, action='store', dest='learning_rate', help='learning rate', default = 1e-3) parser.add_argument('-j', type=float, action='store', dest='jitter', - help='jitter for training data', default = 0.5) + help='jitter for training data', default = 0.4) parser.add_argument('--show', action='store_true', dest='show_figures', help='show figures') @@ -66,17 +66,21 @@ def get_args(): parser.add_argument('--x_end', action='store', dest='data_x_end', help='boundary data x end', nargs='+', default = [14.0,14.0]) + # set the pml thickness + parser.add_argument('--pml_thickness', action='store', dest='pml_thickness', + help='pml thickness', nargs='+', default = [2.0,2.0]) + # set the training spacing parser.add_argument('--train_x_step', action='store', dest='training_data_x_step', - help='training data x step', nargs='+', default = [0.015,0.015]) + help='training data x step', nargs='+', default = [0.4,0.4]) # set the test spacing parser.add_argument('--test_x_step', action='store', dest='test_data_x_step', - help='test data x step', nargs='+', default = [0.3,0.3]) + help='test data x step', nargs='+', default = [0.4,0.4]) # set the evaluation region subset spacing for evaluting w parser.add_argument('--eval_x_step_subset', action='store', dest='eval_data_x_step_subset', - help='evaluation data x step', nargs='+', default = [0.15,0.15]) + help='evaluation data x step', nargs='+', default = [0.2,0.2]) # set the evaluation region spacing for final visualization parser.add_argument('--eval_x_step', action='store', dest='eval_data_x_step', @@ -158,7 +162,7 @@ def partition_dataset(args, world_size): # Test data for validation of pde loss # Test data is a list of coordinates - test_data, _ = create_data(args.test_data_x_start, args.test_data_x_end, + test_data, _ = create_data(args.data_x_start, args.data_x_end, args.test_data_x_step, args.two_d) test_partition = DataPartitioner(test_data, partition_sizes, shuffle=True) @@ -229,6 +233,10 @@ def loss_fn(data, u_scatter, data_2): device, args.use_pde_cl, args.two_d, + args.data_x_end[0]-args.data_x_start[0], + args.data_x_end[1]-args.data_x_start[1], + args.pml_thickness[0], + args.pml_thickness[1], data_2=data_2, ) @@ -291,10 +299,10 @@ def visualize(args, """ device = get_device(args) # Solve the linear system for a subset of the points, use those weights for all points - eval_data_subset, _ = create_data(args.eval_data_x_start_subset, args.eval_data_x_end_subset, + eval_data_subset, _ = create_data(args.data_x_start, args.data_x_end, args.eval_data_x_step_subset, args.two_d) - eval_data, lengths = create_data(args.eval_data_x_start, args.eval_data_x_end, + eval_data, lengths = create_data(args.data_x_start, args.data_x_end, args.eval_data_x_step, args.two_d) eval_dataloader = DataLoader(eval_data, batch_size=args.batch_size, shuffle=False) @@ -323,6 +331,10 @@ def loss_fn(data, u_scatter, data_2,w=None): device, args.use_pde_cl, args.two_d, + args.data_x_end[0]-args.data_x_start[0], + args.data_x_end[1]-args.data_x_start[1], + args.pml_thickness[0], + args.pml_thickness[1], data_2=data_2, w=w, ) diff --git a/utils.py b/utils.py index 4c0edec..163af05 100644 --- a/utils.py +++ b/utils.py @@ -188,10 +188,10 @@ def transform_linear_pde(data, model, two_d, device, - domain_size_x=14, - domain_size_z=14, - L_pml_x=2, - L_pml_z=2, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, a_0=0.25, use_vmap=True, ): @@ -276,6 +276,10 @@ def transform_affine_pde(wavelength, model, device, two_d, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, ): '''Get the right and left hand side of the PDE (del**2 + n**2*k0**2)*u_scatter = -(n**2-n_background**2)*k0**2*u_in))''' @@ -287,6 +291,10 @@ def transform_affine_pde(wavelength, model, two_d, device, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, ) if two_d: u_in = create_plane_wave_2d(data, @@ -312,6 +320,10 @@ def get_pde_loss(data, device, use_pde_cl, two_d, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, data_2=None, w=None, ): @@ -327,6 +339,10 @@ def get_pde_loss(data, model, device, two_d, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, ) @@ -356,6 +372,10 @@ def get_pde_loss(data, model, device, two_d, + domain_size_x, + domain_size_z, + L_pml_x, + L_pml_z, ) linear_pde_combine = torch.matmul(linear_pde,w)