Some tiny docs improvements

Signed-off-by: Johannes Mueller <[email protected]>

src/torchphysics/models/activation_fn.py

@@ -3,9 +3,9 @@
 
 
 class AdaptiveActivationFunction(nn.Module):
-    """Implementation of the adaptive activation functions used in [1].
-    Will create activations of the form: activation_fn(scaling*a * x), 
-    where activation_fn is an arbitrary function, a is the additional 
+    """Implementation of the adaptive activation functions used in [#]_.
+    Will create activations of the form: activation_fn(scaling*a * x),
+    where activation_fn is an arbitrary function, a is the additional
     hyperparameter and scaling is an additional scaling factor.
 
     Parameters
@@ -21,8 +21,8 @@ class AdaptiveActivationFunction(nn.Module):
 
     Notes
     -----
-    ..  [1] Ameya D. Jagtap, Kenji Kawaguchi and George Em Karniadakis, 
-        "Adaptive activation functions accelerate convergence in deep and 
+    ..  [#] Ameya D. Jagtap, Kenji Kawaguchi and George Em Karniadakis,
+        "Adaptive activation functions accelerate convergence in deep and
         physics-informed neural networks", 2020
     """
     def __init__(self, activation_fn, inital_a=1.0, scaling=1.0):
@@ -55,14 +55,14 @@ def backward(ctx, grad_output):
 
 
 class ReLUn(nn.Module):
-    """Implementation of a smoother version of ReLU, in the 
+    """Implementation of a smoother version of ReLU, in the
     form of relu(x)**n.
 
     Parameters
     ----------
     n : float
         The power to which the inputs should be rasied before appplying the
-        rectified linear unit function. 
+        rectified linear unit function.
     """
     def __init__(self, n):
         super().__init__()
@@ -80,4 +80,4 @@ def __init__(self):
         super().__init__()
 
     def forward(self, input):
-        return torch.sin(input)
+        return torch.sin(input)

src/torchphysics/models/deeponet/branchnets.py

@@ -11,15 +11,17 @@
 
 class BranchNet(Model):
     """A neural network that can be used inside a DeepONet-model.
+
     Parameters
     ----------
     function_space : Space
         The space of functions that can be put in this network.
     discretization_sampler : torchphysics.sampler
-        A sampler that will create the points at which the input functions should 
+        A sampler that will create the points at which the input functions should
         evaluated, to create a discrete input for the network.
         The number of input neurons will be equal to the number of sampled points.
         Therefore, the sampler should always return the same number of points!
+
     """
     def __init__(self, function_space, discretization_sampler):
         super().__init__(function_space, output_space=None)
@@ -29,39 +31,39 @@ def __init__(self, function_space, discretization_sampler):
         self.current_out = torch.empty(0)
 
     def finalize(self, output_space, output_neurons):
-        """Method to set the output space and output neurons of the network. 
+        """Method to set the output space and output neurons of the network.
         Will be called once the BranchNet is connected to the TrunkNet, so
         that both will have a fitting output shape.
 
         output_space : Space
             The space in which the final output of the DeepONet will belong to.
         output_neurons : int
-            The number of output neurons. Will be multiplied my the dimension of the 
-            output space, so each dimension will have the same number of 
+            The number of output neurons. Will be multiplied my the dimension of the
+            output space, so each dimension will have the same number of
             intermediate neurons.
         """
         self.output_neurons = output_neurons
         self.output_space = output_space
 
     def _reshape_multidimensional_output(self, output):
-        return output.reshape(-1, self.output_space.dim, 
+        return output.reshape(-1, self.output_space.dim,
                               int(self.output_neurons/self.output_space.dim))
-        
+
     @abc.abstractmethod
     def forward(self, discrete_function_batch, device='cpu'):
         """Evaluated the network at a given function batch. Should not be called
         directly, rather use the method ``.fix_input``.
-        
+
         Parameters
         ----------
         discrete_function_batch : tp.space.Points
             The points object of discrete function values to evaluate the model.
         device : str, optional
             The device where the data lays. Default is 'cpu'.
-        
+
         Notes
         -----
-        Will, in general, not return anything. The output of the network will be saved 
+        Will, in general, not return anything. The output of the network will be saved
         internally to be used multiple times.
         """
         raise NotImplementedError
@@ -75,20 +77,20 @@ def _discretize_function_set(self, function_set, device='cpu'):
         return fn_out
 
     def fix_input(self, function, device='cpu'):
-        """Fixes the branch net for a given function. The branch net will 
-        be evaluated for the given function and the output saved in ``current_out``. 
-        
+        """Fixes the branch net for a given function. The branch net will
+        be evaluated for the given function and the output saved in ``current_out``.
+
         Parameters
         ----------
-        function : callable, torchphysics.domains.FunctionSet, torch.Tensor, 
+        function : callable, torchphysics.domains.FunctionSet, torch.Tensor,
                     torchphysics.spaces.Points
             The function(s) for which the network should be evaluaded.
         device : str, optional
             The device where the data lays. Default is 'cpu'.
-        
+
         Notes
         -----
-        To overwrite the data ``current_out`` (the fixed function) just call 
+        To overwrite the data ``current_out`` (the fixed function) just call
         ``.fix_input`` again with a new function.
         """
         if isinstance(function, FunctionSet):
@@ -120,13 +122,13 @@ def fix_input(self, function, device='cpu'):
 
 class FCBranchNet(BranchNet):
     """A neural network that can be used inside a DeepONet-model.
-    
+
     Parameters
     ----------
     function_space : Space
         The space of functions that can be put in this network.
     discretization_sampler : torchphysics.sampler
-        A sampler that will create the points at which the input functions should 
+        A sampler that will create the points at which the input functions should
         evaluated, to create a discrete input for the network.
         The number of input neurons will be equal to the number of sampled points.
         Therefore, the sampler should always return the same number of points!
@@ -136,13 +138,13 @@ class FCBranchNet(BranchNet):
         of hidden layers, while the i-th entry will determine the number
         of neurons of each layer.
     activations : torch.nn or list, optional
-        The activation functions of this network. 
+        The activation functions of this network.
         Deafult is nn.Tanh().
     xavier_gains : float or list, optional
         For the weight initialization a Xavier/Glorot algorithm will be used.
-        Default is 5/3. 
+        Default is 5/3.
     """
-    def __init__(self, function_space, discretization_sampler, hidden=(20,20,20), 
+    def __init__(self, function_space, discretization_sampler, hidden=(20,20,20),
                  activations=nn.Tanh(), xavier_gains=5/3):
         super().__init__(function_space, discretization_sampler)
         self.hidden = hidden
@@ -151,16 +153,16 @@ def __init__(self, function_space, discretization_sampler, hidden=(20,20,20),
 
     def finalize(self, output_space, output_neurons):
         super().finalize(output_space, output_neurons)
-        layers = _construct_FC_layers(hidden=self.hidden, input_dim=self.input_dim, 
-                        output_dim=self.output_neurons, activations=self.activations, 
+        layers = _construct_FC_layers(hidden=self.hidden, input_dim=self.input_dim,
+                        output_dim=self.output_neurons, activations=self.activations,
                         xavier_gains=self.xavier_gains)
 
         self.sequential = nn.Sequential(*layers)
 
     def forward(self, discrete_function_batch):
         discrete_function_batch = discrete_function_batch.as_tensor.reshape(-1, self.input_dim)
         self.current_out = self._reshape_multidimensional_output(self.sequential(discrete_function_batch))
-       
+
 
 class ConvBranchNet1D(BranchNet):
     """A branch network that first applies a convolution to the input functions
@@ -171,32 +173,32 @@ class ConvBranchNet1D(BranchNet):
     function_space : Space
         The space of functions that can be put in this network.
     discretization_sampler : torchphysics.sampler
-        A sampler that will create the points at which the input functions should 
+        A sampler that will create the points at which the input functions should
         evaluated, to create a discrete input for the network.
         The number of input neurons will be equal to the number of sampled points.
         Therefore, the sampler should always return the same number of points!
     convolutional_network : torch.nn.module
-        The user defined convolutional network, that should be applied to the 
-        branch input. Inside this network, the input can be transformed arbitrary, 
-        e.g. you can also apply pooling or other layers. 
+        The user defined convolutional network, that should be applied to the
+        branch input. Inside this network, the input can be transformed arbitrary,
+        e.g. you can also apply pooling or other layers.
         We only expect that the network gets the input in the shape:
 
-        [batch_dim, function_space.output_space.dim (channels_in), 
+        [batch_dim, function_space.output_space.dim (channels_in),
          len(discretization_sampler)]
 
-        You have to make sure, that the number of output dimension is 
+        You have to make sure, that the number of output dimension is
         compatible with the following linear layers.
     hidden : list or tuple
         The number and size of the hidden layers of the neural network.
         The lenght of the list/tuple will be equal to the number
         of hidden layers, while the i-th entry will determine the number
         of neurons of each layer.
     activations : torch.nn or list, optional
-        The activation functions of this network. 
+        The activation functions of this network.
         Deafult is nn.Tanh().
     xavier_gains : float or list, optional
         For the weight initialization a Xavier/Glorot algorithm will be used.
-        Default is 5/3. 
+        Default is 5/3.
     """
     def __init__(self, function_space, discretization_sampler, convolutional_network,
                  hidden=(20,20,20), activations=nn.Tanh(), xavier_gains=5/3):
@@ -208,8 +210,8 @@ def __init__(self, function_space, discretization_sampler, convolutional_network
 
     def finalize(self, output_space, output_neurons):
         super().finalize(output_space, output_neurons)
-        layers = _construct_FC_layers(hidden=self.hidden, input_dim=self.input_dim, 
-                        output_dim=self.output_neurons, activations=self.activations, 
+        layers = _construct_FC_layers(hidden=self.hidden, input_dim=self.input_dim,
+                        output_dim=self.output_neurons, activations=self.activations,
                         xavier_gains=self.xavier_gains)
 
         self.sequential = nn.Sequential(*layers)
@@ -224,14 +226,14 @@ def _discretize_fn(self, function, device):
         return discrete_fn.unsqueeze(-1)
 
     def forward(self, discrete_function_batch):
-        # for convolution we have to change the dimension order of 
+        # for convolution we have to change the dimension order of
         # the input.
         # Pytorch conv1D needs: (batch, channels_in, length)
         # Generally we have : (batch, length, channels_in), where channels_in
-        # corresponds to the output dimension of our functions and length to the 
+        # corresponds to the output dimension of our functions and length to the
         # number of discretization points. -> switch dim. 1 and 2
         discrete_function_batch = discrete_function_batch.as_tensor
         x = self.conv_net(discrete_function_batch.permute(0, 2, 1))
         # for the linear layer transform again and remove the last dimension:
         out = self.sequential(x.permute(0, 2, 1).reshape(-1, self.input_dim))
-        self.current_out = self._reshape_multidimensional_output(out)
+        self.current_out = self._reshape_multidimensional_output(out)

src/torchphysics/models/deeponet/deeponet.py

@@ -7,39 +7,39 @@
 
 
 class DeepONet(Model):
-    """Implementation of the architecture used in the DeepONet paper [1].
+    """Implementation of the architecture used in the DeepONet paper [#]_.
     Consists of two single neural networks. One for the inputs of the function
     space (branch net) and one for the inputs of the variables (trunk net).
 
     Parameters
     ----------
     trunk_net : torchphysics.models.TrunkNet
-        The neural network that will get the space/time/... variables as an 
-        input. 
+        The neural network that will get the space/time/... variables as an
+        input.
     branch_net : torchphysics.models.BranchNet
-        The neural network that will get the function variables as an 
-        input. 
+        The neural network that will get the function variables as an
+        input.
     output_space : Space
         The space in which the final output of the DeepONet will belong to.
     output_neurons : int
-        The number of output neurons, that will be the output of the 
+        The number of output neurons, that will be the output of the
         TrunkNet and BranchNet. The corresponding outputs of both networks
-        are then connected with the inner product. 
-        For higher dimensional outputs, will be multiplied my the dimension of 
-        the output space, so each dimension will have the same number of 
+        are then connected with the inner product.
+        For higher dimensional outputs, will be multiplied my the dimension of
+        the output space, so each dimension will have the same number of
         intermediate neurons.
 
     Notes
     -----
     The number of output neurons in the branch and trunk net have to be the same!
 
-    ..  [1] Lu Lu and Pengzhan Jin and Guofei Pang and Zhongqiang Zhang
-        and George Em Karniadakis, "Learning nonlinear operators via DeepONet 
+    ..  [#] Lu Lu and Pengzhan Jin and Guofei Pang and Zhongqiang Zhang
+        and George Em Karniadakis, "Learning nonlinear operators via DeepONet
         based on the universal approximation theorem of operators", 2021
     """
     def __init__(self, trunk_net, branch_net, output_space, output_neurons):
         self._check_trunk_and_branch_correct(trunk_net, branch_net)
-        super().__init__(input_space=trunk_net.input_space, 
+        super().__init__(input_space=trunk_net.input_space,
                          output_space=output_space)
         self.trunk = trunk_net
         self.branch = branch_net
@@ -70,19 +70,19 @@ def forward(self, trunk_inputs, branch_inputs=None, device='cpu'):
         trunk_inputs : torchphysics.spaces.Points
             The inputs for the trunk net.
         branch_inputs : callable, torchphysics.domains.FunctionSet, optional
-            The function(s) for which the branch should be evaluaded. If no 
+            The function(s) for which the branch should be evaluaded. If no
             input is given, the branch net has to be fixed before hand!
         device : str, optional
             The device where the data lays. Default is 'cpu'.
-            
+
         Returns
         -------
         torchphysics.spaces.Points
             A point object containing the output.
-        
+
         """
         if not branch_inputs is None:
-            self.fix_branch_input(branch_inputs, device=device) 
+            self.fix_branch_input(branch_inputs, device=device)
         trunk_out = self.trunk(trunk_inputs)
         if len(trunk_out.shape) < 4:
             trunk_out = trunk_out.unsqueeze(0) # shape = [1, trunk_n, dim, neurons]
@@ -99,8 +99,8 @@ def _forward_branch(self, function_set, iteration_num=-1, device='cpu'):
             self.branch(discrete_fn_batch)
 
     def fix_branch_input(self, function, device='cpu'):
-        """Fixes the branch net for a given function. this function will then be used 
-        in every following forward call. To set a new function just call this method 
+        """Fixes the branch net for a given function. this function will then be used
+        in every following forward call. To set a new function just call this method
         again.
 
         Parameters
@@ -110,4 +110,4 @@ def fix_branch_input(self, function, device='cpu'):
         device : str, optional
             The device where the data lays. Default is 'cpu'.
         """
-        self.branch.fix_input(function, device=device)
+        self.branch.fix_input(function, device=device)

src/torchphysics/models/deeponet/layers.py

@@ -37,6 +37,7 @@ class TrunkLinear(torch.nn.Module):
     """Applies a linear transformation to the incoming data: :math:`y = xA^T + b`, similar
     to torch.nn.Linear, but assumes the input `x` to be identical along the first batch axis,
     since this is the case in our implementation of trunk nets.
+
     Args:
         in_features: size of each input sample
         out_features: size of each output sample
@@ -62,6 +63,7 @@ class TrunkLinear(torch.nn.Module):
         >>> output = m(input)
         >>> print(output.size())
         torch.Size([128, 30])
+
     """
     __constants__ = ['in_features', 'out_features']
 
@@ -94,4 +96,4 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
     def extra_repr(self) -> str:
         return 'in_features={}, out_features={}, bias={}'.format(
             self.in_features, self.out_features, self.bias is not None
-        )
+        )