Add Tiled structure clients for common AFL features

AFL/automation/APIServer/Driver.py

@@ -129,7 +129,8 @@ def set_sample(self,sample_name,sample_uuid=None,**kwargs):
 
         kwargs.update({'sample_name':sample_name,'sample_uuid':sample_uuid})
         self.data.update(kwargs)
-
+        self.data.PROTECTED_SAMPLE_KEYS.update(kwargs.keys())
+
         return kwargs
 
     def get_sample(self):

AFL/automation/APIServer/data/TiledClients/CatalogOfAFLEvents.py

@@ -1,7 +1,19 @@
 import tiled
 
+from tqdm.auto import tqdm
 
 from tiled.client.container import DEFAULT_STRUCTURE_CLIENT_DISPATCH, Container
+from tiled.queries import Eq
+
+import pandas as pd
+import xarray as xr
+
+DIM_NAMES_FOR_DATA = {
+        'I': ['q'],
+        'dI': ['q'],
+        'raw': ['pix_x','pix_y']
+        }
+
 
 class CatalogOfAFLEvents(Container):
     ''' 
@@ -13,6 +25,45 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # TODO: set up properties here
 
+    def list_samples(self):
+        uuids = []
+        names = []
+        comps = []
+        al_campaigns = []
+
+        for sample in tqdm(self.task('set_sample').groupby_sample()):
+            try:
+                s = sample.items()[-1][1].metadata # just get the last entry, best metadata
+                try:
+                    uuids.append(s['sample_uuid'])
+                except KeyError:
+                    uuids.append('')
+                try:
+                    names.append(s['sample_name'])
+                except KeyError:
+                    names.append('')
+                try:
+                    comps.append(s['sample_composition'])
+                except KeyError:
+                    comps.append('')
+
+                try:
+                    al_campaigns.append(s['AL_campaign_name'])
+                except KeyError:
+                    al_campaigns.append('')
+            except Exception as e:
+                print(f'Exception {e} while loading sample')
+        # make a pd.Dataframe of these results
+        df = pd.DataFrame({'uuid':uuids,'name':names,'composition':comps,'al_campaign':al_campaigns})
+        return df
+
+    def list_drivers(self):
+        return self.distinct('driver_name')['metadata']['driver_name']
+
+    def list_tasks(self):
+        return self.distinct('task_name')['metadata']['task_name']
+
+
     def groupby_sample(self):
         sample_list = self.distinct('sample_uuid')
 
@@ -22,7 +73,7 @@ def groupby_sample(self):
         for sample_uuid in sample_list:
             sample_uuid = sample_uuid['value']
             yield self.search(Eq('sample_uuid',sample_uuid))
-        
+
     def groupby_driver(self):
         driver_list = self.distinct('driver_name')
 
@@ -34,15 +85,92 @@ def groupby_driver(self):
             yield self.search(Eq('driver_name',driver_name))
 
     def driver(self,driver_name):
-        '''shorthand to return only a particular driver
         '''
-
+        shorthand to return only a particular driver
+        '''
         return self.search(Eq('driver_name',driver_name))
 
     def sample_uuid(self,sample_uuid):
-
         return self.search(Eq('sample_uuid',sample_uuid))
+
+    def task_uuid(self,task_uuid):
+        return self.search(Eq('uuid',task_uuid))
 
-    def list_samples(self):
+    def task(self,task):
+        return self.search(Eq('task_name',task))
+
+    def al_campaign_name(self,al_campaign):
+        return self.search(Eq('AL_campaign_name',al_campaign))
+
+    def dataset_for_task(self,task_uuid):
+        '''
+        return a particular task as an xarray dataset
+        '''
+
+        # first, get a subcatalog that contains only this task
+
+        task = self.task_uuid(task_uuid)
+
+        #main_entry = task.search(Eq('array_name',''))
+
+        array_names = task.distinct('array_name')['metadata']['array_name']
+
+        data_vars = {}
+        for name in array_names:
+            name = name['value']
+            array_client = task.search(Eq('array_name',name)).values()
+            assert len(array_client)==1, f'Error: more than one array with name {name} found in task {task_uuid}.  Tyler probably wrote this Driver.'
+
+
+
+            data_vars[name] =  (
+                DIM_NAMES_FOR_DATA[name],
+                array_client[0].read(),
+                array_client[0].metadata
+                )
+
+        for var in data_vars:
+            if var[0] not in data_vars.keys():
+                try:
+
+        return xr.Dataset(data_vars, attrs = task.values()[0].metadata)
+        #return data_vars
 
-        pass 
+    def __getitem__(self, key):
+        # For convenience and backward-compatiblity reasons, we support
+        # some "magic" here that is helpful in an interactive setting.
+        if isinstance(key, str):
+            # CASE 1: Interpret key as a uid or partial uid.
+            if len(key) == 36:
+                # This looks like a full uid. Try direct lookup first.
+                try:
+                    return super().__getitem__(key)
+                except KeyError:
+                    # Fall back to partial uid lookup below.
+                    pass
+            return self._lookup_by_partial_uid(key)
+        elif isinstance(key, numbers.Integral):
+            if key > 0:
+                # CASE 2: Interpret key as a scan_id.
+                return self._lookup_by_scan_id(key)
+            else:
+                # CASE 3: Interpret key as a recently lookup, as in
+                # `catalog[-1]` is the latest entry.
+                key = int(key)
+                return self.values()[key]
+        elif isinstance(key, slice):
+            if (key.start is None) or (key.start >= 0):
+                raise ValueError(
+                    "For backward-compatibility reasons, slicing here "
+                    "is limited to negative indexes. "
+                    "Use .values() to slice how you please."
+                )
+            return self.values()[key]
+        elif isinstance(key, collections.abc.Iterable):
+            # We know that isn't a str because we check that above.
+            # Recurse.
+            return [self[item] for item in key]
+        else:
+            raise ValueError(
+                "Indexing expects a string, an integer, or a collection of strings and/or integers."
+            )

AFL/automation/APIServer/data/TiledClients/DriverTask.py

@@ -0,0 +1,50 @@
+import tiled
+
+from tqdm.auto import tqdm
+
+from tiled.client.array import ArrayClient
+from tiled.queries import Eq
+
+import datetime
+
+class DriverTask(ArrayClient):
+    ''' 
+    a subclass of tiled.ArrayClient that adds accessor methods to iterate over samples, drivers,
+    and convienence methods that let you filter by sample name/driver more easily
+    '''
+
+    def __init__(self, *args, **kwargs):
+    	super().__init__(*args, **kwargs)
+    	# TODO: set up properties here
+
+    @property
+    def start_time(self):
+    	return datetime.datetime.strptime(self.metadata['meta']['started'].strip(), '%m/%d/%y %H:%M:%S-%f')
+
+    @property
+    def end_time(self):
+    	return datetime.datetime.strptime(self.metadata['meta']['ended'].strip(), '%m/%d/%y %H:%M:%S-%f')
+
+    @property
+    def queued_time(self):
+    	return datetime.datetime.strptime(self.metadata['meta']['queued'].strip(), '%m/%d/%y %H:%M:%S-%f')
+
+    @property
+    def exit_state(self):
+    	return self.metadata['meta']['exit_state']
+
+    @property
+    def return_val(self):
+    	return self.metadata['meta']['return_val']
+
+
+    @property
+    def task_name(self):
+    	return self.metadata['task_name']
+
+    @property
+    def task(self):
+    	return self.metadata['task']
+
+
+

setup.py

@@ -13,7 +13,9 @@
     scripts=['launchers/roboctl'],
     entry_points={
         'tiled.special_client': [
-                'afl.event_catalog = AFL.automation.APIServer.data.TiledClients.CatalogOfAFLEvents:CatalogOfAFLEvents'
+                'afl.event_catalog = AFL.automation.APIServer.data.TiledClients.CatalogOfAFLEvents:CatalogOfAFLEvents',
+                'afl.driver_task = AFL.automation.APIServer.data.TiledClients.DriverTask:DriverTask'
+
                 ]
     }
 )