ENH update stimuli hdf5 file

Author: TomDLT

tutorials/movies_3T/06_extract_motion_energy.py

@@ -64,29 +64,30 @@
 from voxelwise_tutorials.io import load_hdf5_array
 
 
-def compute_luminance(run_name, size=(96, 96)):
+def compute_luminance(run_name, size=(96, 96), batch_size=100):
 
     stimuli_file = os.path.join(directory, 'stimuli', run_name)
 
-    # get the list of batches in the stimuli file
+    # get the number of images in the stimuli file
     with h5py.File(stimuli_file, 'r') as f:
-        keys = list(f.keys())
-        keys.sort()  # sort the batches
+        n_images = f['stimuli'].shape[0]
 
     # compute the luminance on each batch
-    luminance = []
-    for key in bar(keys, title=f'compute_luminance({run_name})'):
+    luminance = np.zeros((n_images, *size))
+    for start in bar(range(0, n_images, batch_size),
+                     title=f'compute_luminance({run_name})'):
         # load the batch of images
-        images = load_hdf5_array(stimuli_file, key=key)
+        batch = slice(start, start + batch_size)
+        images = load_hdf5_array(stimuli_file, key='stimuli', slice=batch)
 
         # ``imagearray2luminance`` uses uint8 arrays
         if images.dtype != 'uint8':
             images = np.int_(np.clip(images, 0, 1) * 255).astype(np.uint8)
 
         # convert RGB images to a single luminance channel
-        luminance.append(imagearray2luminance(images, size=size))
+        luminance[batch] = imagearray2luminance(images, size=size)
 
-    return np.concatenate(luminance)
+    return luminance
 
 
 luminance_train = np.concatenate(
@@ -161,7 +162,8 @@ def compute_motion_energy(luminance,
 
 save_hdf5_dataset(
     os.path.join(features_directory, "motion_energy_recomputed.hdf"),
-    dataset=dict(X_train=motion_energy_train, X_test=motion_energy_test))
+    dataset=dict(X_train=motion_energy_train, X_test=motion_energy_test,
+                 run_onsets=np.arange(0, 3600, 300)))
 
 ###############################################################################
 # References

voxelwise_tutorials/io.py

@@ -91,15 +91,18 @@ def unpack_archive(archive_name):
     shutil.unpack_archive(archive_name, extract_dir=extract_dir)
 
 
-def load_hdf5_array(file_name, key=None):
+def load_hdf5_array(file_name, key=None, slice=slice(0, None)):
     """Function to load data from an hdf file.
 
     Parameters
     ----------
     file_name: string
-        hdf5 file name
+        hdf5 file name.
     key: string
         Key name to load. If not provided, all keys will be loaded.
+    slice: slice, or tuple of slices
+        Load only a slice of the hdf5 array. It will load `array[slice]`.
+        Use a tuple of slices to get a slice in multiple dimensions.
 
     Returns
     -------
@@ -110,10 +113,10 @@ def load_hdf5_array(file_name, key=None):
         if key is None:
             data = dict()
             for k in hf.keys():
-                data[k] = hf[k][()]
+                data[k] = hf[k][slice]
             return data
         else:
-            return hf[key][()]
+            return hf[key][slice]
 
 
 def load_hdf5_sparse_array(file_name, key):