Add tqdm to track job creation

Author: Linked-Liszt

src/simulator/diffraction_generator.py

@@ -62,33 +62,44 @@ def _generate_simulation_tasks(self, n_sims_per_file, master_seed, cleanup_worke
         job_id_counter = 0
         tasks_to_run, tasks_to_skip = [], 0
 
-        for file_path in self.file_list:
-            stem = Path(file_path).stem
-            for variation_index in range(1, n_sims_per_file + 1):
-                final_output_path = self.data_dir / f"{stem}-{variation_index}.npy"
-                if final_output_path.exists():
-                    tasks_to_skip += 1
+        total_planned = len(self.file_list) * int(n_sims_per_file)
+        disable_bar = not sys.stderr.isatty()
+        with tqdm(
+            total=total_planned,
+            desc="Preparing tasks",
+            unit="task",
+            mininterval=1.0,
+            disable=disable_bar,
+        ) as pbar:
+            for file_path in self.file_list:
+                stem = Path(file_path).stem
+                for variation_index in range(1, n_sims_per_file + 1):
+                    final_output_path = self.data_dir / f"{stem}-{variation_index}.npy"
+                    if final_output_path.exists():
+                        tasks_to_skip += 1
+                        job_id_counter += 1
+                        pbar.update(1)
+                        continue
+
+                    worker_dir = self.worker_base_dir / f'job_{job_id_counter:09d}'
+                    params = {
+                        'job_id': job_id_counter, 'worker_dir': str(worker_dir), 'input_file': str(file_path),
+                        'output_data_dir': str(self.data_dir), 'error_dir': str(self.error_dir),
+                        'instprm_file': str(self.instprm_file), 'noise_seed': rng.integers(1e9),
+                        'strain': rng.uniform(*param_ranges['strain_range']), 'size': rng.uniform(*param_ranges['size_range']),
+                        'U': rng.uniform(*param_ranges['U_range']), 'V': rng.uniform(*param_ranges['V_range']),
+                        'W': rng.uniform(*param_ranges['W_range']), 'st': rng.uniform(*param_ranges['st_range']),
+                        'en': rng.uniform(*param_ranges['en_range']), 'Npoints': int(rng.uniform(*param_ranges['Npoints_range'])),
+                        'scaler': rng.uniform(*param_ranges['scaler_range']), 'wl': rng.uniform(*param_ranges['wl_range']),
+                        'proportional_noise_range': param_ranges['proportional_noise_range'],
+                        'constant_noise_range': param_ranges['constant_noise_range'],
+                        'cleanup_worker_dir': cleanup_worker_dirs,
+                        'output_filename': f"{stem}-{variation_index}.npy",
+                        'parse_from_comment': bool(kwargs.get('parse_from_comment', False)),
+                    }
+                    tasks_to_run.append(params)
                     job_id_counter += 1
-                    continue
-                
-                worker_dir = self.worker_base_dir / f'job_{job_id_counter:09d}'
-                params = {
-                    'job_id': job_id_counter, 'worker_dir': str(worker_dir), 'input_file': str(file_path),
-                    'output_data_dir': str(self.data_dir), 'error_dir': str(self.error_dir),
-                    'instprm_file': str(self.instprm_file), 'noise_seed': rng.integers(1e9),
-                    'strain': rng.uniform(*param_ranges['strain_range']), 'size': rng.uniform(*param_ranges['size_range']),
-                    'U': rng.uniform(*param_ranges['U_range']), 'V': rng.uniform(*param_ranges['V_range']),
-                    'W': rng.uniform(*param_ranges['W_range']), 'st': rng.uniform(*param_ranges['st_range']),
-                    'en': rng.uniform(*param_ranges['en_range']), 'Npoints': int(rng.uniform(*param_ranges['Npoints_range'])),
-                    'scaler': rng.uniform(*param_ranges['scaler_range']), 'wl': rng.uniform(*param_ranges['wl_range']),
-                    'proportional_noise_range': param_ranges['proportional_noise_range'],
-                    'constant_noise_range': param_ranges['constant_noise_range'],
-                    'cleanup_worker_dir': cleanup_worker_dirs,
-                    'output_filename': f"{stem}-{variation_index}.npy",
-                    'parse_from_comment': bool(kwargs.get('parse_from_comment', False)),
-                }
-                tasks_to_run.append(params)
-                job_id_counter += 1
+                    pbar.update(1)
 
         if tasks_to_skip > 0: print(f"--> Found and skipped {tasks_to_skip} previously completed jobs.")
         for task in tasks_to_run: yield task