Merge pull request #33 from radionets-project/tqdm_support

Changes in vis_loop

docs/changes/33.feature.rst

@@ -0,0 +1,3 @@
+Changes to `vis_loop` function in `visibility.py`:
+- add a an optional tqdm progress bar to get a visual confirmation the calculation is still running
+- add optional `batch_size` parameter to control memory consumption

pyvisgen/simulation/visibility.py

@@ -4,6 +4,8 @@
 
 import pyvisgen.simulation.scan as scan
 
+from tqdm import tqdm
+
 
 @dataclass
 class Visibilities:
@@ -37,7 +39,7 @@ def add(self, visibilities):
         ]
 
 
-def vis_loop(obs, SI, num_threads=10, noisy=True, mode="full"):
+def vis_loop(obs, SI, num_threads=10, noisy=True, mode="full", batch_size=100, show_progress=False):
     torch.set_num_threads(num_threads)
     torch._dynamo.config.suppress_errors = True
 
@@ -93,7 +95,12 @@ def vis_loop(obs, SI, num_threads=10, noisy=True, mode="full"):
     else:
         raise ValueError("Unsupported mode!")
 
-    for p in torch.arange(bas[:].shape[1]).split(1000):
+    batches = torch.arange(bas[:].shape[1]).split(batch_size)
+
+    if show_progress:
+        batches = tqdm(batches)
+
+    for p in batches:
         bas_p = bas[:][:, p]
 
         int_values = torch.cat(