Added a functionality to take predefined tqdm loops as input to solver.fit()

neurodiffeq/generators.py

@@ -7,20 +7,20 @@
 
 def _chebyshev_first(a, b, n):
     nodes = torch.cos(((torch.arange(n) + 0.5) / n) * np.pi)
-    nodes = ((a + b) + (b - a) * nodes) / 2
+    nodes = ((a + b) + (a - b) * nodes) / 2
     nodes.requires_grad_(True)
     return nodes
 
 
 def _chebyshev_second(a, b, n):
     nodes = torch.cos(torch.arange(n) / float(n - 1) * np.pi)
-    nodes = ((a + b) + (b - a) * nodes) / 2
+    nodes = ((a + b) + (a - b) * nodes) / 2
     nodes.requires_grad_(True)
     return nodes
 
 def _chebyshev_second_noisy(a, b, n):
     nodes = torch.cos((torch.arange(n) + (torch.rand(n) * 2 - 1)) / float(n - 1) * np.pi)
-    nodes = ((a + b) + (b - a) * nodes) / 2
+    nodes = ((a + b) + (a - b) * nodes) / 2
     nodes.requires_grad_(True)
     return nodes
 

neurodiffeq/solvers.py

@@ -440,7 +440,7 @@ def _update_best(self, key):
             self.lowest_loss = current_loss
             self.best_nets = deepcopy(self.nets)
 
-    def fit(self, max_epochs, callbacks=(), tqdm_file=sys.stderr, **kwargs):
+    def fit(self, max_epochs, callbacks=(), tqdm_file='default', **kwargs):
         r"""Run multiple epochs of training and validation, update best loss at the end of each epoch.
 
         If ``callbacks`` is passed, callbacks are run, one at a time,
@@ -471,19 +471,21 @@ def fit(self, max_epochs, callbacks=(), tqdm_file=sys.stderr, **kwargs):
             callbacks = [monitor.to_callback()] + list(callbacks)
         if kwargs:
             raise ValueError(f'Unknown keyword argument(s): {list(kwargs.keys())}')  # pragma: no cover
-
-        if tqdm_file is None:
-            loop = range(max_epochs)
-        else:
+
+        flag = True
+        if 'default' in str(tqdm_file):
             loop = tqdm(
-                range(max_epochs),
+                total = max_epochs,
                 desc='Training Progress',
                 colour='blue',
-                file=tqdm_file,
                 dynamic_ncols=True,
             )
+        elif tqdm_file is not None:
+            loop = tqdm_file
+        else:
+            flag = False
 
-        for local_epoch in loop:
+        for local_epoch in range(max_epochs):
             # stop training if self._stop_training is set to True by a callback
             if self._stop_training:
                 break
@@ -495,6 +497,8 @@ def fit(self, max_epochs, callbacks=(), tqdm_file=sys.stderr, **kwargs):
 
             for cb in callbacks:
                 cb(self)
+            if flag:
+                loop.update(1)
 
     @abstractmethod
     def get_solution(self, copy=True, best=True):