refactor; add domain

Author: akoryagin

pydens/model_torch.py

@@ -16,23 +16,34 @@
 
 class TorchModel(ABC, nn.Module):
     """ Pytorch model for solving differential equations with neural networks. """
-    def __init__(self, initial_condition=None, boundary_condition=None, ndims=1, nparams=0, **kwargs):
+    def __init__(self, ndims, initial_condition=None, boundary_condition=None, domain=(0, 1), nparams=0, **kwargs):
         _ = kwargs
         super().__init__()
 
         # Store the number of variables and parameters - dimensionality of the problem.
         self.ndims = ndims
+        self.ndims_spatial = ndims if initial_condition is None else ndims - 1
         self.nparams = nparams
         self.total = ndims + nparams
         self.variables = {}
 
-        # Parse and store initial and boundary condition.
+        # Parse and store initial condition, boundary condition and domain of the problem.
         if initial_condition is None:
             self.initial_condition = None
         else:
             self.initial_condition = (initial_condition if callable(initial_condition)
                                       else lambda *args: torch.tensor(initial_condition, dtype=torch.float32))
         self.boundary_condition = boundary_condition
+        if isinstance(domain, (tuple, list)):
+            if isinstance(domain[0], (float, int)):
+                domain = [domain] * ndims
+            elif isinstance(domain[0], (tuple, list)):
+                pass
+            else:
+                raise ValueError(f'Should be either 1d or 2d-sequence of float/ints.')
+        else:
+            raise ValueError(f'Should be either 1d or 2d-sequence of float/ints.')
+        self.domain = domain
 
         # Initialize trainable variables for anzatc-trasform to bind initial
         # and boundary conditions.
@@ -96,18 +107,23 @@ def unfreeze_trainable(self, layers=None, variables=None):
     def anzatc(self, u, xs):
         """ Anzatc-transformation of the model-output needed for binding initial and boundary conditions. """
         # Get tensor of spatial variables and time-tensor.
-        xs_spatial = xs[:, :self.ndims] if self.initial_condition is None else xs[:, :self.ndims - 1]
+        xs_spatial = xs[:, :self.ndims_spatial]
         t = xs[:, self.ndims - 1:self.ndims]
+        lower, upper = [lims[0] for lims in self.domain], [lims[1] for lims in self.domain]
+        lower_spatial, upper_spatial = [torch.Tensor(lst[:self.ndims_spatial]).reshape(1, -1).float()
+                                        for lst in (lower, upper)]
+        t0 = lower[-1]
 
         # Apply transformation to bind the boundary condition.
         if self.boundary_condition is not None:
-            u = u * (torch.prod(xs_spatial, dim=1, keepdim=True) *
-                     torch.prod((1 - xs_spatial), dim=1, keepdim=True)) + self.boundary_condition
+            u = (u * (torch.prod((xs_spatial - lower_spatial) / (upper_spatial - lower_spatial), dim=1, keepdim=True) *
+                      torch.prod((upper_spatial - xs_spatial) / (upper_spatial - lower_spatial), dim=1, keepdim=True))
+                     + self.boundary_condition)
 
         # Apply transformation to bind the initial condition.
         if self.initial_condition is not None:
             _xs_spatial = [xs_spatial[:, i] for i in range(xs_spatial.shape[1])]
-            u = ((nn.Sigmoid()(t / torch.exp(self.log_scale)) - .5) * u
+            u = ((nn.Sigmoid()((t - t0) / torch.exp(self.log_scale)) - .5) * u
                  + self.initial_condition(*_xs_spatial).view(-1, 1))
         return u
 
@@ -119,6 +135,8 @@ class ConvBlockModel(TorchModel):
     neural networks. For purposes of solving simple differential equations we suggest using
     fully connected networks with skip connections.
 
+    See also docstring of `Solver`.
+
     Parameters
     ----------
     layout : str
@@ -137,12 +155,12 @@ class ConvBlockModel(TorchModel):
         - ``layout, units, activation = 'fa fa f', [5, 10, 1], 'Sigmoid' # Fully-conn with 2 hidden``
         - ``layout, units, activation = 'faR fa fa+ f', [5, 10, 5, 1], 'Sigmoid # Fully-conn with 3 hidden and skip'``
     """
-    def __init__(self, layout='fafaf', units=(20, 30, 1), activation='Sigmoid', **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, ndims, initial_condition=None, boundary_condition=None, domain=(0, 1), nparams=0,
+                 layout='fafaf', units=(20, 30, 1), activation='Sigmoid', **kwargs):
+        super().__init__(ndims=ndims, initial_condition=initial_condition, boundary_condition=boundary_condition,
+                         domain=domain, nparams=nparams, **kwargs)
 
         # Prepare kwargs for conv-block.
-        for key in ['initial_condition', 'ndims', 'nparams', 'boundary_condition']:
-            _ = kwargs.pop(key, None)
         kwargs.update(layout=layout, units=list(units), activation=activation)
 
         # Assemble conv-block.
@@ -207,6 +225,34 @@ def pde(f, x, e):
             So, the value of `nparams`-parameter should be set to 1, whie the value of
             `ndims` - to 2.
 
+    ndims : int
+        The dimensionality of the problem. Equals the number of variables. For instance,
+        when the problem includes `x`, `y` and `t`, the dimensionality is equal to 3.
+
+    initial_condition : callable or float
+        Defines initial condition for a problem that includes time-variable.
+
+        Examples:
+
+        - ``lambda x: x * (1 - x)``
+            Can define initial condition for a wave equation for a pertubed string with `x` and `t`
+            - variables. The initial condition is basically a perturbation of a string at `t=0`.
+        - ``lambda x, y: 10 * x * y * (1 - x) * (1 - y)``
+            Can define initial condition for a heat equation describing temperature evolution of a
+            2d plate.
+
+    boundary_condition : float
+        Defines boundary condition for the problem.
+
+    domain : tuple or list
+        Configures rectangular domain of the problem. Can be either a 2d-sequence of form
+        `[(x_low, x_high), (y_low, y_high), ..., (t_low, t_high)]` of `length=ndims`; or
+        a 1d-sequence `(low, high)`. In this case, limits for all dimensions are considered
+        to be the same.
+
+    nparams : int
+        The number of parameters with uncertainty in the model.
+
     model : class
         Class inheriting `TorchModel`. The default value is `ConvBlockModel`. The class allows
         to implement fully connected/convolutional architectures with multiple branches
@@ -229,28 +275,6 @@ def pde(f, x, e):
             This set of additional constraints can be used to (i) bind `f(0.5) = -1`
             along with minimization of a L2-norm of `f`.
 
-    ndims : int
-        The dimensionality of the problem. Equals the number of variables. For instance,
-        when the problem includes `x`, `y` and `t`, the dimensionality is equal to 3.
-
-    nparams : int
-        The number of parameters with uncertainty in the model.
-
-    initial_condition : callable or float
-        Defines initial condition for a problem that includes time-variable.
-
-        Examples:
-
-        - ``lambda x: x * (1 - x)``
-            Can define initial condition for a wave equation for a pertubed string with `x` and `t`
-            - variables. The initial condition is basically a perturbation of a string at t=0.
-        - ``lambda x, y: 10 * x * y * (1 - x) * (1 - y)``
-            Can define initial condition for a heat equation describing temperature evolution of a
-            2d plate.
-
-    boundary_condition : float
-        Defines boundary condition for the problem.
-
     kwargs : dict
         Keyword-arguments used for initialization of the model-instance. When `model`-parameter
         is set to its default value - `ConvBlockModel`, use these arguments to configure the
@@ -272,7 +296,8 @@ def pde(f, x, e):
             - ``layout, units, activation = 'fa fa f', [5, 10, 1], 'Sigmoid'``
             - ``layout, units, activation = 'faR fa fa+ f', [5, 10, 5, 1], 'Sigmoid'``
     """
-    def __init__(self, equation, model=ConvBlockModel, constraints=None, **kwargs):
+    def __init__(self, equation, ndims, initial_condition=None, boundary_condition=None, domain=(0, 1),
+                 nparams=0, model=ConvBlockModel, constraints=None, **kwargs):
         self.equation = equation
         if constraints is None:
             self.constraints = ()
@@ -284,7 +309,8 @@ def __init__(self, equation, model=ConvBlockModel, constraints=None, **kwargs):
         self.optimizer = None
 
         # Initialize neural network for solving the equation.
-        self.model = model(**kwargs)
+        self.model = model(**kwargs, ndims=ndims, initial_condition=initial_condition,
+                           boundary_condition=boundary_condition, domain=domain, nparams=nparams)
 
         # Bind created model to a global context variable.
         current_model.set(self.model)