clean up code

flaring/data/iti_data_processing.py

@@ -1,4 +1,8 @@
 import collections.abc
+import shutil
+
+import pandas as pd
+
 collections.Iterable = collections.abc.Iterable
 collections.Mapping = collections.abc.Mapping
 collections.MutableSet = collections.abc.MutableSet
@@ -6,15 +10,11 @@ collections.MutableMapping = collections.abc.MutableMapping
 # Now import hyper
 import numpy as np
 from astropy.visualization import ImageNormalize, AsinhStretch
-from itipy.data.dataset import SDODataset, StackDataset, get_intersecting_files, AIADataset
-from itipy.data.editor import LoadMapEditor, NormalizeRadiusEditor, MapToDataEditor, AIAPrepEditor, \
-    BrightestPixelPatchEditor
+from itipy.data.dataset import StackDataset, get_intersecting_files, AIADataset
+from itipy.data.editor import LoadMapEditor, NormalizeRadiusEditor, MapToDataEditor, BrightestPixelPatchEditor
 import os
-import glob
-from multiprocessing import Pool, cpu_count
-import tqdm as tqdm
-import multiprocessing as mp
-from functools import partial
+from multiprocessing import Pool
+from tqdm import tqdm
 
 # Configuration for all wavelengths to process
 wavelengths = [94, 131, 171, 193, 211, 304]
@@ -31,6 +31,7 @@ sdo_norms = {
     '304': ImageNormalize(vmin=0, vmax=np.float32(401.82352), stretch=AsinhStretch(0.001), clip=True),
 }
 
+
 class SDODataset_flaring(StackDataset):
     """
     Dataset for SDO data
@@ -43,19 +44,22 @@ class SDODataset_flaring(StackDataset):
         ext (str): File extension
         **kwargs: Additional arguments
     """
+
     def __init__(self, data, patch_shape=None, wavelengths=None, resolution=2048, ext='.fits', **kwargs):
         wavelengths = [171, 193, 211, 304, 6173, ] if wavelengths is None else wavelengths
         if isinstance(data, list):
             paths = data
         else:
             paths = get_intersecting_files(data, wavelengths, ext=ext, **kwargs)
-        ds_mapping = {94:AIADataset, 131: AIADataset, 171: AIADataset, 193: AIADataset, 211: AIADataset, 304: AIADataset}
+        ds_mapping = {94: AIADataset, 131: AIADataset, 171: AIADataset, 193: AIADataset, 211: AIADataset,
+                      304: AIADataset}
         data_sets = [ds_mapping[wl_id](files, wavelength=wl_id, resolution=resolution, ext=ext)
                      for wl_id, files in zip(wavelengths, paths)]
         super().__init__(data_sets, **kwargs)
         if patch_shape is not None:
             self.addEditor(BrightestPixelPatchEditor(patch_shape))
 
+
 aia_dataset = SDODataset_flaring(data=base_input_folder, wavelengths=wavelengths, resolution=512)
 
 def save_sample(i):
@@ -66,4 +70,4 @@ def save_sample(i):
     np.save(file_path, data)
 
 with Pool(processes=90) as pool:
-    list(tqdm.tqdm(pool.imap(save_sample, range(len(aia_dataset))), total=len(aia_dataset)))
+    list(tqdm(pool.imap(save_sample, range(len(aia_dataset))), total=len(aia_dataset)))