fix more comments; update readme-page

Author: akoryagin

README.md

@@ -43,7 +43,7 @@ solver = Solver(equation=pde, ndims=2, boundary_condition=1,
 
 ```
 
-Note that we defined the architecture of the neural network by supplying `layout`, `activation` and `units` parameters. Here `layout` configures the sequence of layers: `fa fa fa f` stands for `f`ully connected architecture with four layers and three `a`ctivations. In its turn, `units` and `activation` cotrol the number of units in dense layers and activation-function. When defining neural network this way use [`ConvBlock`](https://analysiscenter.github.io/batchflow/api/batchflow.models.torch.layers.html?highlight=baseconvblock#batchflow.models.torch.layers.BaseConvBlock) from [`BatchFlow`](https://github.com/analysiscenter/batchflow). Check out its capabilities to understand more: say, making a network with attention and skip connections.
+Note that we defined the architecture of the neural network by supplying `layout`, `activation` and `units` parameters. Here `layout` configures the sequence of layers: `fa fa fa f` stands for `f`ully connected architecture with four layers and three `a`ctivations. In its turn, `units` and `activation` cotrol the number of units in dense layers and activation-function. When defining neural network this way use [`ConvBlock`](https://analysiscenter.github.io/batchflow/api/batchflow.models.torch.layers.html?highlight=baseconvblock#batchflow.models.torch.layers.BaseConvBlock) from [`BatchFlow`](https://github.com/analysiscenter/batchflow).
 
 It's time to run the optimization procedure
 
@@ -82,14 +82,14 @@ def odeparam(f, x, e):
 s = NumpySampler('uniform') & NumpySampler('uniform', low=1, high=5)
 
 solver = Solver(equation=odeparam, ndims=1, nparams=1, initial_condition=1)
-solver.fit(batch_size=1000, sampler=s, niters=5000)
+solver.fit(batch_size=1000, sampler=s, niters=5000, lr=0.01)
 # solving the whole family takes no more than a couple of seconds!
 ```
 
 Check out the result:
 
 <p align="center">
-<img src="https://raw.githubusercontent.com/analysiscenter/pydens/master/imgs/sinus_sol.gif?invert_in_darkmode" align=middle height=250.973825pt/>
+<img src="https://raw.githubusercontent.com/analysiscenter/pydens/master/imgs/sinus_parametric.gif?invert_in_darkmode" align=middle height=250.973825pt/>
 </p>
 
 ### Solving PDEs with trainable coefficients

imgs/sinus_parametric.gif

@@ -42,7 +42,6 @@ def __init__(self, initial_condition=None, boundary_condition=None, ndims=1, npa
     @abstractmethod
     def forward(self, xs):
         """ Forward of the model-network. """
-        pass
 
     def freeze_trainable(self, layers=None, variables=None):
         """ Freeze layers and trainable variables.
@@ -139,7 +138,8 @@ def reshape_and_concat(cls, tensors):
         """ Cast, reshape and concatenate sequence of incoming tensors. """
         # Determine batch size as max-len of a tensor.
         xs = list(tensors)
-        batch_size = np.max([tensor.shape[0] for tensor in xs if isinstance(tensor, (np.ndarray, torch.Tensor))])
+        sizes = [tensor.shape[0] for tensor in xs if isinstance(tensor, (np.ndarray, torch.Tensor))]
+        batch_size = np.max(sizes) if len(sizes) > 0 else 1
 
         # Perform cast and reshape of all tensors in the list.
         for i, x in enumerate(xs):
@@ -180,7 +180,7 @@ def __init__(self, equation, model=ConvBlockModel, constraints=None, **kwargs):
         _ = self.ctx.run(self.equation, u_hat, *xs)
 
 
-    def fit(self, niters, batch_size, sampler=None, losses='equation', optimizer='Adam', criterion=nn.MSELoss(),
+    def fit(self, niters, batch_size, sampler=None, loss_terms='equation', optimizer='Adam', criterion=nn.MSELoss(),
             lr=0.001, **kwargs):
         """ Fit the model inside the solver-instance. """
         # Initialize the optimizer if supplied.
@@ -206,13 +206,13 @@ def fit(self, niters, batch_size, sampler=None, losses='equation', optimizer='Ad
             u_hat = self.ctx.run(self.model, xs_concat)
 
             # Compute loss: form it summing equation-loss and constraints-loss.
-            losses = losses if isinstance(losses, (tuple, list)) else (losses, )
-            nums_constraints = [int(loss_name.replace('constraint', '').replace('_', ''))
-                                for loss_name in losses if 'constraint' in loss_name]
+            loss_terms = loss_terms if isinstance(loss_terms, (tuple, list)) else (loss_terms, )
+            nums_constraints = [int(term_name.replace('constraint', '').replace('_', ''))
+                                for term_name in loss_terms if 'constraint' in term_name]
             loss  = 0
 
             # Include equation loss.
-            if 'equation' in losses:
+            if 'equation' in loss_terms:
                 loss += criterion(self.ctx.run(self.equation, u_hat, *xs), torch.zeros_like(xs[0]))
 
             # Include additional constraints' loss.
@@ -231,7 +231,7 @@ def _forward(*xs):
             # Gather and store training stats.
             self.losses.append(loss.detach().cpu().numpy())
 
-    def solve(self, *xs):
+    def predict(self, *xs):
         """ Get approximation to a solution in a set of points.
         Points are given by a list of tensors.
         """