import h5py
import numpy as np

class FileWriter:
    """Write data in European XFEL HDF5 format

    This is intended to allow copying a subset of data into a smaller,
    more portable file.
    """
    def __init__(self, path, data):
        self.file = h5py.File(path, 'w')
        self.data = data
        self.indexes = {} # {path: (first, count)}
        self.data_sources = set()

    def add_control_dataset(self, source, key):
        a = self.data.get_array(source, key)
        path = 'CONTROL/{}/{}'.format(source, key.replace('.', '/'))
        self.file[path] = a.values

        if source not in self.indexes:
            n = a.shape[0]
            self.indexes[source] = \
                (np.arange(n, dtype='u8'), np.ones(n, dtype='u8'))
            self.data_sources.add('CONTROL/' + source)

    def add_instrument_dataset(self, source, key):
        a = self.data.get_array(source, key)
        path = 'INSTRUMENT/{}/{}'.format(source, key.replace('.', '/'))
        self.file[path] = a.values

        index_path = source + '/' + key.partition('.')[0]
        if index_path not in self.indexes:
            # Convert an array of train IDs to first/count for each train
            data_tids = a.coords['trainId']
            first, count = [], []
            for tid in self.data.train_ids:
                matches = (data_tids == tid)
                first.append(matches.nonzero()[0][0])
                count.append(matches.sum())

            self.indexes[index_path] = \
                (np.array(first, dtype='u8'), np.array(count, dtype='u8'))
            self.data_sources.add('INSTRUMENT/' + index_path)

    def write_indexes(self):
        for groupname, (first, count) in self.indexes.items():
            self.file['INDEX/{}/first'.format(groupname)] = first
            self.file['INDEX/{}/count'.format(groupname)] = count

    def write_metadata(self):
        vlen_bytes = h5py.special_dtype(vlen=bytes)
        data_sources = sorted(self.data_sources)
        N = len(data_sources)

        sources_ds = self.file.create_dataset('METADATA/dataSourceId', (N,),
                                              dtype=vlen_bytes, maxshape=(None,))
        sources_ds[:] = data_sources

        root_ds = self.file.create_dataset('METADATA/root', (N,),
                                           dtype=vlen_bytes, maxshape=(None,))
        root_ds[:] = [ds.split('/', 1)[0] for ds in data_sources]

        devices_ds = self.file.create_dataset('METADATA/deviceId', (N,),
                                              dtype=vlen_bytes, maxshape=(None,))
        devices_ds[:] = [ds.split('/', 1)[1] for ds in data_sources]

    def write(self):
        d = self.data
        self.file.create_dataset('INDEX/trainId', data=d.train_ids, dtype='u8')

        for source in d.control_sources:
            for key in d._keys_for_source(source):
                self.add_control_dataset(source, key)

        for source in d.instrument_sources:
            for key in d._keys_for_source(source):
                self.add_instrument_dataset(source, key)

        self.write_indexes()
        self.write_metadata()

@@ -1123,18 +1123,242 @@
         
-# 
+data_tids = a.coords%5B'trainId'%5D.values%0A            self.indexes%5Bindex_path%5D = self._generate_index(data_tids)%0A            self.data_sources.add('INSTRUMENT/' + index_path)%0A%0A    def _generate_index(self, data_tids):%0A        %22%22%22
 Convert 
@@ -1408,16 +1408,19 @@
 ch train
+%22%22%22
 %0A       
@@ -1424,82 +1424,127 @@
     
-    data_tids = a.coords%5B'trainId'%5D%0A            first, count = %5B%5D, %5B%5D%0A    
+first = np.zeros_like(self.data.train_ids, dtype='u8')%0A        count = np.zeros_like(self.data.train_ids, dtype='u8')%0A%0A
     
@@ -1554,16 +1554,20 @@
  for
+ ix,
  tid in 
 self
@@ -1558,24 +1558,34 @@
  ix, tid in 
+enumerate(
 self.data.tr
@@ -1591,16 +1591,17 @@
 rain_ids
+)
 :%0A      
@@ -1606,20 +1606,16 @@
         
-   
- 
 matches 
@@ -1635,16 +1635,46 @@
 == tid)%0A
+            if matches.any():%0A
         
@@ -1686,24 +1686,23 @@
    first
-.append(
+%5Bix%5D = 
 matches.
@@ -1716,17 +1716,16 @@
 ()%5B0%5D%5B0%5D
-)
 %0A       
@@ -1742,16 +1742,15 @@
 ount
-.append(
+%5Bix%5D = 
 matc
@@ -1758,17 +1758,16 @@
 es.sum()
-)
 %0A%0A      
@@ -1772,177 +1772,27 @@
     
-  
-  self.indexes%5Bindex_path%5D = %5C%0A                (np.array(first, dtype='u8'), np.array(count, dtype='u8'))%0A            self.data_sources.add('INSTRUMENT/' + index_path)
+return first, count
 %0A%0A  

Handle data with missing trains