Merge pull request #118 from PyLops/solver-sysout

feat: flush all prints in solver
PyLops · Nov 23, 2024 · e9bbecc · e9bbecc
2 parents 1351e7a + 307b104
commit e9bbecc
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/pylops_mpi/optimization/cls_basic.py b/pylops_mpi/optimization/cls_basic.py
@@ -1,6 +1,7 @@
 from typing import List, Optional, Tuple, Union
-import numpy as np
+import sys
 import time
+import numpy as np
 
 from pylops.optimization.basesolver import Solver
 from pylops.utils import NDArray
@@ -40,13 +41,15 @@ def _print_setup(self, xcomplex: bool = False) -> None:
         else:
             head1 = "    Itn              x[0]                  r2norm"
         print(head1)
+        sys.stdout.flush()
 
     def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> None:
         if isinstance(x, StackedDistributedArray):
             x = x.distarrays[0]
         strx = f"{x[0]:1.2e}        " if np.iscomplexobj(x.local_array) else f"{x[0]:11.4e}        "
         msg = f"{self.iiter:6g}        " + strx + f"{self.cost[self.iiter]:11.4e}"
         print(msg)
+        sys.stdout.flush()
 
     def setup(
             self,
@@ -281,6 +284,7 @@ def _print_setup(self, xcomplex: bool = False) -> None:
         else:
             head1 = "    Itn             x[0]             r1norm         r2norm"
         print(head1)
+        sys.stdout.flush()
 
     def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> None:
         if isinstance(x, StackedDistributedArray):
@@ -292,6 +296,7 @@ def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> No
             + f"{self.cost[self.iiter]:11.4e}    {self.cost1[self.iiter]:11.4e}"
         )
         print(msg)
+        sys.stdout.flush()
 
     def setup(self,
               y: Union[DistributedArray, StackedDistributedArray],