dataset improvement

sCOCA_ML/dataset/gravpot_dataset.py

@@ -7,6 +7,7 @@ from glob import glob
 import re
 
 
+
 def read_cosmo_and_time_file(cosmo_and_time_file):
     with open(cosmo_and_time_file, 'r') as f:
         lines = f.readlines()
@@ -154,6 +155,10 @@ class GravPotDataset(Dataset):
 
     def __getitem__(self, idx):
         from pysbmy.field import read_field_chunk_3D_periodic
+        from io import BytesIO
+        from sbmy_control.low_level import stdout_redirector, stderr_redirector
+        f = BytesIO()
+        
         ID, t, ox, oy, oz = self.samples[idx]
 
         # Filepaths
@@ -165,6 +170,7 @@ class GravPotDataset(Dataset):
         style_path = os.path.join(self.root_dir, self.STYLE_DIR, f'{self.style_files}_{ID}_nforce{t}.txt')
 
         # Read 3D chunks
+        with stdout_redirector(f):
             input_arrays = [
                 read_field_chunk_3D_periodic(file, self.N,self.N,self.N, ox,oy,oz, name=varname).array
                 for file, varname in zip(input_paths, self.initial_conditions_variables)