Add time-dependent Fourier feature abilities

pinnicle/nn/nn.py

@@ -15,9 +15,19 @@ 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]
+
+        # 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 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]
             self.num_layers = len(self.num_neurons)
             # append linear transform for the output
             self.activation = self.parameters.activation + [None]
@@ -40,18 +50,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: 
@@ -115,4 +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)
-

pinnicle/parameter.py

@@ -218,12 +218,18 @@ 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_fourier_feature = 10
-        self.sigma = 1.0
-        self.B = None
+        self.num_space_fourier_feature = 10
+        self.space_sigma = 1.0
+        self.space_B = None
+        self.num_time_fourier_feature = 10
+        self.time_sigma = 1.0
+        self.time_B = None
 
         # parallel neural network
         self.is_parallel = False
@@ -235,14 +241,24 @@ def set_default(self):
         self.output_ub = None
 
     def check_consistency(self):
-        if self.fft:
-            if self.input_size != self.num_fourier_feature*self.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")
-            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 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 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 self.time_dependent and self.time_B is not None:
+                if not isinstance(self.time_B, 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 features")
         else:
             # input size of nn equals to dependent in physics
             if self.input_size != len(self.input_variables):
@@ -283,13 +299,23 @@ 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
-            if not isinstance(self.sigma, list):
-                self.sigma = [self.sigma]
+            # Convert space sigma to a list
+            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]
 
+            # Create space-dependent input size
             # 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
+
+            # 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)
 
             # cover num_neurons to list
             if not isinstance(self.num_neurons, list):

pinnicle/pinn.py

@@ -259,8 +259,12 @@ 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) 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
         self.model = dde.Model(self.dde_data, self.nn.net)

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={}

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