From edb218afb21a71dda427788f4aebe47f8a845661 Mon Sep 17 00:00:00 2001 From: Claire Jensen Date: Sat, 11 Jul 2026 15:28:46 -0600 Subject: [PATCH 1/2] Add time dependency for Fourier features. --- pinnicle/domain/domain.py | 1 + pinnicle/nn/nn.py | 60 ++++++++++++++++++++++++++++++++------- pinnicle/parameter.py | 46 +++++++++++++++++++++--------- pinnicle/pinn.py | 6 ++-- 4 files changed, 88 insertions(+), 25 deletions(-) diff --git a/pinnicle/domain/domain.py b/pinnicle/domain/domain.py index 332d3d1..ef071ca 100644 --- a/pinnicle/domain/domain.py +++ b/pinnicle/domain/domain.py @@ -86,3 +86,4 @@ def _add_midpoint(self, domain_list): newy = 0.5*(domain_list[0][1] + domain_list[1][1]) newx = 0.5*(domain_list[0][0] + domain_list[1][0]) domain_list.insert(1, [newx, newy]) + diff --git a/pinnicle/nn/nn.py b/pinnicle/nn/nn.py index f911ddd..160660b 100644 --- a/pinnicle/nn/nn.py +++ b/pinnicle/nn/nn.py @@ -15,9 +15,18 @@ def __init__(self, parameters=NNParameter()): # update necesarry parameters for fourier feature transform # NOTE: these changes will not be saved to the param file, # so that the change will not accumulate and loading the previous param file will create the same nn - if self.parameters.fft : + if self.parameters.fft and not self.parameters.time_dependent: # Then add an additional layer before the output node - self.num_neurons = parameters.num_neurons + [parameters.num_fourier_feature*parameters.sigma_size] + self.num_neurons = parameters.num_neurons + [parameters.num_space_fourier_feature*parameters.space_sigma_size] + self.num_layers = len(self.num_neurons) + # append linear transform for the output + self.activation = self.parameters.activation + [None] + + elif self.parameters.fft and self.parameters.time_dependent: + # TODO: Point-wise multiplication of Fourier features to merge + + # Add layer before output node + self.num_neurons = parameters.num_neurons + [parameters.num_space_fourier_feature*parameters.space_sigma_size + parameters.num_time_fourier_feature*parameters.time_sigma_size] self.num_layers = len(self.num_neurons) # append linear transform for the output self.activation = self.parameters.activation + [None] @@ -40,18 +49,48 @@ def __init__(self, parameters=NNParameter()): self.parameters.input_lb = bkd.as_tensor(self.parameters.input_lb, dtype=default_float_type()) self.parameters.input_ub = bkd.as_tensor(self.parameters.input_ub, dtype=default_float_type()) - if self.parameters.fft : - print(f"add Fourier feature transform to input transform") - if self.parameters.B is not None: - self.B = bkd.as_tensor(self.parameters.B, dtype=default_float_type()) + if self.parameters.fft and not self.parameters.time_dependent: + print(f"add Fourier feature transform to spatial input transform") + if self.parameters.space_B is not None: + self.space_B = bkd.as_tensor(self.parameters.space_B, dtype=default_float_type()) else: - self.B = bkd.as_tensor( - np.reshape(np.random.normal(0.0, self.parameters.sigma, [len(self.parameters.input_variables), self.parameters.num_fourier_feature, self.parameters.sigma_size]), [len(self.parameters.input_variables), self.parameters.num_fourier_feature*self.parameters.sigma_size]), + self.space_B = bkd.as_tensor( + np.reshape(np.random.normal(0.0, self.parameters.space_sigma, [len(self.parameters.input_variables), self.parameters.num_space_fourier_feature, self.parameters.space_sigma_size]), [len(self.parameters.input_variables), self.parameters.num_space_fourier_feature*self.parameters.space_sigma_size]), dtype=default_float_type()) def wrapper(x): - """a wrapper function to add fourier feature transform to the input + """a wrapper function to add fourier feature transform to the spatial input """ - return fourier_feature(minmax_scale(x, self.parameters.input_lb, self.parameters.input_ub), self.B) + return fourier_feature(minmax_scale(x, self.parameters.input_lb, self.parameters.input_ub), self.space_B) + # add to input transform + self.net.apply_feature_transform(wrapper) + elif self.parameters.fft and self.parameters.time_dependent: + print(f"add Fourier feature transform to spatial and temporal input transform") + # Spatial features + if self.parameters.space_B is not None: + self.space_B = bkd.as_tensor(self.parameters.space_B, dtype=default_float_type()) + else: + space_len = len([var for var in self.parameters.input_variables if var == 'x' or var == 'y']) + self.space_B = bkd.as_tensor( + np.reshape(np.random.normal(0.0, self.parameters.space_sigma, [space_len, self.parameters.num_space_fourier_feature, self.parameters.space_sigma_size]), [space_len, self.parameters.num_space_fourier_feature*self.parameters.space_sigma_size]), + dtype=default_float_type()) + + # Temporal features + if self.parameters.time_B is not None: + self.time_B = bkd.as_tensor(self.parameters.time_B, dtype=default_float_type()) + else: + time_len = len([var for var in self.parameters.input_variables if var == 't']) + self.time_B = bkd.as_tensor( + np.reshape(np.random.normal(0.0, self.parameters.time_sigma, [time_len, self.parameters.num_time_fourier_feature, self.parameters.time_sigma_size]), [time_len, self.parameters.num_time_fourier_feature*self.parameters.time_sigma_size]), + dtype=default_float_type()) + def wrapper(x): + """a wrapper function to add fourier feature transform to the spatial and temporal inputs separately + """ + x_scaled = minmax_scale(x, self.parameters.input_lb, self.parameters.input_ub) + x_space = x_scaled[:, :space_len] + x_time = x_scaled[:, space_len:] + space_features = fourier_feature(x_space, self.space_B) + time_features = fourier_feature(x_time, self.time_B) + return bkd.concat([space_features, time_features], 1) # add to input transform self.net.apply_feature_transform(wrapper) else: @@ -116,3 +155,4 @@ def _wrapper(dummy, x): return func(x, self.parameters.output_lb, self.parameters.output_ub) self.net.apply_output_transform(_wrapper) + diff --git a/pinnicle/parameter.py b/pinnicle/parameter.py index 1eeea60..f318558 100644 --- a/pinnicle/parameter.py +++ b/pinnicle/parameter.py @@ -221,9 +221,13 @@ def set_default(self): # fourier feature transform self.fft = False - self.num_fourier_feature = 10 - self.sigma = 1.0 - self.B = None + self.num_space_fourier_feature = 10 + self.num_time_fourier_feature = 10 + self.space_sigma = 1.0 + self.space_B = None + self.time_dependent = False + self.time_sigma = 1.0 + self.time_B = None # parallel neural network self.is_parallel = False @@ -236,13 +240,20 @@ def set_default(self): def check_consistency(self): if self.fft: - if self.input_size != self.num_fourier_feature*self.sigma_size*2: + if self.input_size != (self.num_space_fourier_feature*self.space_sigma_size*2) + (self.num_time_fourier_feature*self.time_sigma_size*2): raise ValueError("'input_size' does not match the number of fourier feature") - if self.B is not None: - if not isinstance(self.B, list): - raise TypeError("'B' matrix need to be input in a list") - if len(self.B[0]) != self.num_fourier_feature*self.sigma_size: - raise ValueError("Number of columns of 'B' matrix does not match the number of fourier feature") + # Check spatial B matrix + if not self.time_dependent and self.space_B is not None: + if not isinstance(self.space_B, list): + raise TypeError("Spatial 'B' matrix need to be input in a list") + if len(self.space_B[0]) != self.num_space_fourier_feature*self.space_sigma_size: + raise ValueError("Number of columns of spatial 'B' matrix does not match the number of fourier feature") + # Check temporal B matrix + if not self.time_dependent and self.time_B is not None: + if not isinstance(self.time_B, list): + raise TypeError("Temporal 'B' matrix need to be input in a list") + if len(self.time_B[0]) != self.num_time_fourier_feature*self.time_sigma_size: + raise ValueError("Number of columns of temporal 'B' matrix does not match the number of fourier feature") else: # input size of nn equals to dependent in physics if self.input_size != len(self.input_variables): @@ -284,12 +295,21 @@ def set_parameters(self, pdict: dict): raise ValueError("FFT currently does not support parallel nets") # always convert sigma to a list - if not isinstance(self.sigma, list): - self.sigma = [self.sigma] + if not isinstance(self.space_sigma, list): + self.space_sigma = [self.space_sigma] + + # Convert time sigma to a list + if not isinstance(self.time_sigma, list): + self.time_sigma = [self.time_sigma] # we need to know this size, to create and reshape B in nn.py - self.sigma_size = len(self.sigma) - self.input_size = self.num_fourier_feature*self.sigma_size*2 + self.space_sigma_size = len(self.space_sigma) + self.input_size = self.num_space_fourier_feature*self.space_sigma_size*2 + + # Create time-dependent input size + if self.time_dependent: + self.time_sigma_size = len(self.time_sigma) + self.input_size = (self.num_time_fourier_feature*self.time_sigma_size*2) + (self.num_space_fourier_feature*self.space_sigma_size*2) # cover num_neurons to list if not isinstance(self.num_neurons, list): diff --git a/pinnicle/pinn.py b/pinnicle/pinn.py index 10ac260..2a2e66d 100644 --- a/pinnicle/pinn.py +++ b/pinnicle/pinn.py @@ -259,8 +259,10 @@ def setup(self): # define the neural network in use self.nn = FNN(self.params.nn) # save B if it is not defined by the user and generated by FFT - if (self.params.nn.B is None) and (self.params.nn.fft): - self.params.param_dict.update({"B": dde.backend.to_numpy(self.nn.B).tolist()}) + if (self.params.nn.space_B is None) and (self.params.nn.fft): + self.params.param_dict.update({"space_B": dde.backend.to_numpy(self.nn.space_B).tolist()}) + if (self.params.nn.time_B is None) and (self.params.nn.fft) and (self.params.nn.time_dependent): + self.params.param_dict.update({"time_B": dde.backend.to_numpy(self.nn.time_B).tolist()}) # Step 7: setup the deepxde PINN model self.model = dde.Model(self.dde_data, self.nn.net) From e96d83a7962d807d55f0c6ac3628229f0ea65090 Mon Sep 17 00:00:00 2001 From: Claire Jensen Date: Sun, 12 Jul 2026 14:28:36 -0600 Subject: [PATCH 2/2] Update tests and fix bugs --- pinnicle/domain/domain.py | 1 - pinnicle/nn/nn.py | 7 ++--- pinnicle/parameter.py | 28 +++++++++++-------- pinnicle/pinn.py | 4 ++- tests/test_nn.py | 58 +++++++++++++++++++++++++++++++-------- tests/test_parameters.py | 43 ++++++++++++++++++++++------- tests/test_pinn.py | 42 +++++++++++++++++++++++++--- 7 files changed, 140 insertions(+), 43 deletions(-) diff --git a/pinnicle/domain/domain.py b/pinnicle/domain/domain.py index ef071ca..332d3d1 100644 --- a/pinnicle/domain/domain.py +++ b/pinnicle/domain/domain.py @@ -86,4 +86,3 @@ def _add_midpoint(self, domain_list): newy = 0.5*(domain_list[0][1] + domain_list[1][1]) newx = 0.5*(domain_list[0][0] + domain_list[1][0]) domain_list.insert(1, [newx, newy]) - diff --git a/pinnicle/nn/nn.py b/pinnicle/nn/nn.py index 160660b..b7d993d 100644 --- a/pinnicle/nn/nn.py +++ b/pinnicle/nn/nn.py @@ -22,8 +22,9 @@ def __init__(self, parameters=NNParameter()): # append linear transform for the output self.activation = self.parameters.activation + [None] + # Merge space and time-dependent Fourier features in second-to-last layer elif self.parameters.fft and self.parameters.time_dependent: - # TODO: Point-wise multiplication of Fourier features to merge + # TODO: Point-wise multiplication of Fourier features to merge in second-to-last layer # Add layer before output node self.num_neurons = parameters.num_neurons + [parameters.num_space_fourier_feature*parameters.space_sigma_size + parameters.num_time_fourier_feature*parameters.time_sigma_size] @@ -83,7 +84,7 @@ def wrapper(x): np.reshape(np.random.normal(0.0, self.parameters.time_sigma, [time_len, self.parameters.num_time_fourier_feature, self.parameters.time_sigma_size]), [time_len, self.parameters.num_time_fourier_feature*self.parameters.time_sigma_size]), dtype=default_float_type()) def wrapper(x): - """a wrapper function to add fourier feature transform to the spatial and temporal inputs separately + """a wrapper function to add Fourier feature transform to the spatial and temporal inputs separately """ x_scaled = minmax_scale(x, self.parameters.input_lb, self.parameters.input_ub) x_space = x_scaled[:, :space_len] @@ -154,5 +155,3 @@ def _add_output_transform(self, func): def _wrapper(dummy, x): return func(x, self.parameters.output_lb, self.parameters.output_ub) self.net.apply_output_transform(_wrapper) - - diff --git a/pinnicle/parameter.py b/pinnicle/parameter.py index f318558..aa954c8 100644 --- a/pinnicle/parameter.py +++ b/pinnicle/parameter.py @@ -218,14 +218,16 @@ def set_default(self): self.num_layers = 0 self.activation = "tanh" self.initializer = "Glorot uniform" + + # time dependency + self.time_dependent = False # fourier feature transform self.fft = False self.num_space_fourier_feature = 10 - self.num_time_fourier_feature = 10 self.space_sigma = 1.0 self.space_B = None - self.time_dependent = False + self.num_time_fourier_feature = 10 self.time_sigma = 1.0 self.time_B = None @@ -239,21 +241,24 @@ def set_default(self): self.output_ub = None def check_consistency(self): - if self.fft: - if self.input_size != (self.num_space_fourier_feature*self.space_sigma_size*2) + (self.num_time_fourier_feature*self.time_sigma_size*2): + if self.fft and not self.time_dependent: + if self.input_size != (self.num_space_fourier_feature*self.space_sigma_size*2): raise ValueError("'input_size' does not match the number of fourier feature") # Check spatial B matrix if not self.time_dependent and self.space_B is not None: if not isinstance(self.space_B, list): - raise TypeError("Spatial 'B' matrix need to be input in a list") + raise TypeError("Spatial 'B' matrix needs to be input as a list") if len(self.space_B[0]) != self.num_space_fourier_feature*self.space_sigma_size: - raise ValueError("Number of columns of spatial 'B' matrix does not match the number of fourier feature") + raise ValueError("Number of columns of spatial 'B' matrix does not match the number of Fourier features") + elif self.fft and self.time_dependent: + if self.input_size != (self.num_space_fourier_feature*self.space_sigma_size*2) + (self.num_time_fourier_feature*self.time_sigma_size*2): + raise ValueError("'input_size' does not match the number of fourier feature") # Check temporal B matrix - if not self.time_dependent and self.time_B is not None: + if self.time_dependent and self.time_B is not None: if not isinstance(self.time_B, list): - raise TypeError("Temporal 'B' matrix need to be input in a list") + raise TypeError("Temporal 'B' matrix needs to be input as a list") if len(self.time_B[0]) != self.num_time_fourier_feature*self.time_sigma_size: - raise ValueError("Number of columns of temporal 'B' matrix does not match the number of fourier feature") + raise ValueError("Number of columns of temporal 'B' matrix does not match the number of Fourier features") else: # input size of nn equals to dependent in physics if self.input_size != len(self.input_variables): @@ -294,7 +299,7 @@ def set_parameters(self, pdict: dict): if self.is_parallel: raise ValueError("FFT currently does not support parallel nets") - # always convert sigma to a list + # Convert space sigma to a list if not isinstance(self.space_sigma, list): self.space_sigma = [self.space_sigma] @@ -302,11 +307,12 @@ def set_parameters(self, pdict: dict): if not isinstance(self.time_sigma, list): self.time_sigma = [self.time_sigma] + # Create space-dependent input size # we need to know this size, to create and reshape B in nn.py self.space_sigma_size = len(self.space_sigma) self.input_size = self.num_space_fourier_feature*self.space_sigma_size*2 - # Create time-dependent input size + # Reshape for time-dependent input size if self.time_dependent: self.time_sigma_size = len(self.time_sigma) self.input_size = (self.num_time_fourier_feature*self.time_sigma_size*2) + (self.num_space_fourier_feature*self.space_sigma_size*2) diff --git a/pinnicle/pinn.py b/pinnicle/pinn.py index 2a2e66d..72714f8 100644 --- a/pinnicle/pinn.py +++ b/pinnicle/pinn.py @@ -259,9 +259,11 @@ def setup(self): # define the neural network in use self.nn = FNN(self.params.nn) # save B if it is not defined by the user and generated by FFT - if (self.params.nn.space_B is None) and (self.params.nn.fft): + if (self.params.nn.space_B is None) and (self.params.nn.fft) and not (self.params.nn.time_dependent): self.params.param_dict.update({"space_B": dde.backend.to_numpy(self.nn.space_B).tolist()}) + # Save time B if time-dependent FFT if (self.params.nn.time_B is None) and (self.params.nn.fft) and (self.params.nn.time_dependent): + self.params.param_dict.update({"space_B": dde.backend.to_numpy(self.nn.space_B).tolist()}) self.params.param_dict.update({"time_B": dde.backend.to_numpy(self.nn.time_B).tolist()}) # Step 7: setup the deepxde PINN model diff --git a/tests/test_nn.py b/tests/test_nn.py index 61f3f91..3ae0247 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -26,8 +26,8 @@ def test_upscale(): def test_fourier_feature(): x = bkd.reshape(bkd.as_tensor((np.linspace(1,100, 100)), dtype=default_float_type()), [50,2]) - B = bkd.as_tensor(np.random.normal(0.0, 10.0, [x.shape[1], 2]), dtype=default_float_type()) - y = bkd.to_numpy(fourier_feature(x, B)) + space_B = bkd.as_tensor(np.random.normal(0.0, 10.0, [x.shape[1], 2]), dtype=default_float_type()) + y = bkd.to_numpy(fourier_feature(x, space_B)) z = y**2 assert np.all((z[:,1]+z[:,3]) < 1.0+100**np.finfo(float).eps) @@ -48,8 +48,8 @@ def test_input_msfft_nn(): hp['num_neurons'] = 7 hp['num_layers'] = 3 hp['fft'] = True - hp['sigma'] = [1.0,2.0,3.0] - hp['num_fourier_feature'] = 11 + hp['space_sigma'] = [1.0,2.0,3.0] + hp['num_space_fourier_feature'] = 11 d = NNParameter(hp) d.input_lb = 1.0 d.input_ub = 10.0 @@ -59,7 +59,7 @@ def test_input_msfft_nn(): z = y**2 assert np.all(abs(z[:,1:11]+z[:,12:22]-1.0)+np.finfo(float).eps) assert y.shape[1] == 11*2*3 - assert d.sigma_size == 3 + assert d.space_sigma_size == 3 assert p.num_neurons == d.num_neurons + [11*3] assert p.num_layers == 4 assert p.activation == ['tanh']*4+[None] @@ -79,20 +79,54 @@ def test_input_fft_nn(): y = bkd.to_numpy(p.net._input_transform(x)) z = y**2 assert np.all(abs(z[:,1:10]+z[:,11:20]-1.0)+np.finfo(float).eps) - assert d.sigma_size == 1 + assert d.space_sigma_size == 1 - hp['B'] = [[1,2,3]] - hp['num_fourier_feature'] = 3 + hp['space_B'] = [[1,2,3]] + hp['num_space_fourier_feature'] = 3 d = NNParameter(hp) d.input_lb = 1.0 d.input_ub = 10.0 p = pinn.nn.FNN(d) - assert np.all(hp['B'] == bkd.to_numpy(p.B)) + assert np.all(hp['space_B'] == bkd.to_numpy(p.space_B)) - hp['sigma'] = [1.0, 10.0] - hp['B'] = [[1,2,3,4,5,6]] + hp['space_sigma'] = [1.0, 10.0] + hp['space_B'] = [[1,2,3,4,5,6]] d = NNParameter(hp) - assert d.sigma_size == 2 + assert d.space_sigma_size == 2 + +def test_input_temporal_fft_nn(): + hp={} + hp['input_variables'] = ['x'] + hp['output_variables'] = ['u'] + hp['num_neurons'] = 1 + hp['num_layers'] = 1 + hp['fft'] = True + hp["time_dependent"] = True + hp["start_time"] = 0 + hp["end_time"] = 1 + d = NNParameter(hp) + d.input_lb = 1.0 + d.input_ub = 10.0 + p = pinn.nn.FNN(d) + x = bkd.reshape(bkd.as_tensor(np.linspace(1.0, 10.0, 100), dtype=default_float_type()), [100,1]) + y = bkd.to_numpy(p.net._input_transform(x)) + z = y**2 + assert np.all(abs(z[:,1:10]+z[:,11:20]-1.0)+np.finfo(float).eps) + assert d.time_sigma_size == 1 + + # Temporal Fourier features + hp['time_B'] = [[1,2,3]] + hp['num_time_fourier_feature'] = 3 + d = NNParameter(hp) + d.input_lb = 1.0 + d.input_ub = 10.0 + p = pinn.nn.FNN(d) + assert np.all(hp['time_B'] == bkd.to_numpy(p.time_B)) + + hp['time_sigma'] = [1.0, 10.0] + hp['time_B'] = [[1,2,3,4,5,6]] + d = NNParameter(hp) + assert d.time_sigma_size == 2 def test_input_scale_nn(): hp={} diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 848ce82..c6afc9c 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -98,23 +98,46 @@ def test_nn_parameter(): assert d.input_size == 0 d = NNParameter({"fft":True}) - assert d.input_size == 2*d.num_fourier_feature - assert isinstance(d.sigma, list) + assert d.input_size == 2*d.num_space_fourier_feature + assert isinstance(d.space_sigma, list) assert d.is_input_scaling() - assert d.B is None + assert d.space_B is None - d = NNParameter({"fft":True, "num_fourier_feature":4, "B":[[1,2,3,4]]}) - assert d.B is not None + d = NNParameter({"fft":True, "num_space_fourier_feature":4, "space_B":[[1,2,3,4]]}) + assert d.space_B is not None with pytest.raises(Exception): - d = NNParameter({"fft":True, "num_fourier_feature":4, "B":1}) - d = NNParameter({"fft":True, "num_fourier_feature":4, "B":[[1,2]]}) + d = NNParameter({"fft":True, "num_space_fourier_feature":4, "space_B":1}) + d = NNParameter({"fft":True, "num_space_fourier_feature":4, "space_B":[[1,2]]}) + + d = NNParameter({"fft":True, "time_dependent":True, "num_time_fourier_feature":4, "time_B":[[1,2,3,4]]}) + assert d.time_B is not None + with pytest.raises(Exception): + d = NNParameter({"fft":True, "time_dependent":True, "num_time_fourier_feature":4, "time_B":1}) + d = NNParameter({"fft":True, "time_dependent":True, "num_time_fourier_feature":4, "time_B":[[1,2]]}) + + d = NNParameter({"fft":True, "time_dependent":True}) + assert d.input_size == (2*d.num_space_fourier_feature + 2*d.num_time_fourier_feature) + assert isinstance(d.space_sigma, list) + assert isinstance(d.time_sigma, list) + assert d.is_input_scaling() + assert d.space_B is None + assert d.time_B is None with pytest.raises(Exception): d = NNParameter({"fft":True, "is_parallel":True}) - d = NNParameter({"fft":True, "sigma":[1,10], "num_neurons":23, "num_layers":3}) - assert d.sigma_size == 2 - assert d.input_size == 2*d.num_fourier_feature*d.sigma_size + d = NNParameter({"fft":True, "space_sigma":[1,10], "num_neurons":23, "num_layers":3}) + assert d.space_sigma_size == 2 + assert d.input_size == 2*d.num_space_fourier_feature*d.space_sigma_size + assert isinstance(d.num_neurons, list) + assert d.num_layers == 3 + assert len(d.num_neurons) == d.num_layers + assert d.num_neurons[-1] == 23 + + d = NNParameter({"fft":True, "time_dependent":True, "space_sigma":[1,10], "time_sigma":[1,10], "num_neurons":23, "num_layers":3}) + assert d.space_sigma_size == 2 + assert d.time_sigma_size == 2 + assert d.input_size == (2*d.num_space_fourier_feature*d.space_sigma_size + 2*d.num_time_fourier_feature*d.time_sigma_size) assert isinstance(d.num_neurons, list) assert d.num_layers == 3 assert len(d.num_neurons) == d.num_layers diff --git a/tests/test_pinn.py b/tests/test_pinn.py index 163324b..5015712 100644 --- a/tests/test_pinn.py +++ b/tests/test_pinn.py @@ -261,7 +261,7 @@ def test_fft_training(tmp_path): hp_local["is_save"] = True hp_local["save_path"] = str(tmp_path) hp_local["num_collocation_points"] = 10 - hp_local["sigma"] = [1.0, 10.0] + hp_local["space_sigma"] = [1.0, 10.0] issm["data_size"] = {"u":10, "v":10, "s":10, "H":10, "C":None} hp_local["data"] = {"ISSM": issm} hp_local["equations"] = {"SSA":SSA} @@ -269,17 +269,51 @@ def test_fft_training(tmp_path): experiment = pinn.PINN(params=hp_local) experiment.save_setting(path=tmp_path) assert experiment.params.param_dict == experiment.load_setting(path=tmp_path) - assert experiment.params.nn.B is None + assert experiment.params.nn.space_B is None assert os.path.isdir(f"{tmp_path}/pinn/") experiment2 = pinn.PINN(loadFrom=tmp_path) assert experiment.params.param_dict == experiment2.params.param_dict - assert len(experiment2.params.nn.B) == 2 - assert len(experiment2.params.nn.B[1]) == 20 + assert len(experiment2.params.nn.space_B) == 2 + assert len(experiment2.params.nn.space_B[1]) == 20 assert experiment2.params.nn.num_layers == 4 experiment.compile() experiment.train(iterations=10) assert experiment.loss_names == ['fSSA1', 'fSSA2', 'u', 'v', 's', 'H', 'C'] experiment.load_model(path=tmp_path, epochs=hp_local['epochs']) + +@pytest.mark.skipif(backend_name in ["jax"], reason="There is a bug in deepxde (jax backend), when assigning a list of activation function, it reads in as a tuple!") +def test_fft_timedependent_training(tmp_path): + hp_local = dict(hp) + hp_local['fft'] = True + hp_local["time_dependent"] = True + hp_local["is_save"] = True + hp_local["save_path"] = str(tmp_path) + hp_local["num_collocation_points"] = 10 + hp_local["space_sigma"] = [1.0, 10.0] + hp_local["time_sigma"] = [1.0, 10.0] + hp_local["start_time"] = 0 + hp_local["end_time"] = 1 + issm["data_size"] = {"u":10, "v":10, "H":10} + hp_local["data"] = {"ISSM": issm} + hp_local["equations"] = {"Mass transport":{}} + hp_local["epochs"] = 10 + experiment = pinn.PINN(params=hp_local) + experiment.save_setting(path=tmp_path) + assert experiment.params.param_dict == experiment.load_setting(path=tmp_path) + assert experiment.params.nn.space_B is None + assert experiment.params.nn.time_B is None + assert os.path.isdir(f"{tmp_path}/pinn/") + experiment2 = pinn.PINN(loadFrom=tmp_path) + assert experiment.params.param_dict == experiment2.params.param_dict + assert len(experiment2.params.nn.space_B) == 2 + assert len(experiment2.params.nn.space_B[1]) == 20 + assert len(experiment2.params.nn.time_B) == 1 + assert len(experiment2.params.nn.time_B[0]) == 20 + assert experiment2.params.nn.num_layers == 4 + experiment.compile() + experiment.train(iterations=10) + assert experiment.loss_names == ['fMass transport', 'u', 'v', 'H'] + experiment.load_model(path=tmp_path, epochs=hp_local['epochs']) @pytest.mark.skipif(backend_name in ["jax"], reason="calving front boundary is not implemented for jax") def test_train_calvingfront():