fab.datasources.basesources

  1from abc import ABC, abstractmethod
  2import xarray as xr
  3
  4from ..settings import cfg, get_config_object, update_context
  5
  6import logging
  7import copy
  8logger = logging.getLogger(__name__)
  9
 10
 11#TODO: more type hints
 12#TODO: move ADC slicer baseline and inversion to preprocessor
 13
 14class DataSource(ABC):
 15    ''' Represents a source of data. 
 16    
 17    Abstract class for describing a data source. It should implement
 18    the load() method for loading data. The data must be indexable 
 19    by train_id
 20    '''
 21    
 22    def __init__(self, name, *args, fillna_method = None, **kwargs):
 23        ''' Base init for a data source. 
 24        
 25        Args:
 26            name: a human readable name for the data source.
 27            fillna_method (optional): a string specifying what filling method should
 28                be used to fill missing values if the dataset needs to be reindexed.
 29                This will happen when the datasource is combined with other sources
 30                in an Instrument object that will require all source to have the same
 31                train_id index. It should be either 'ffill' for forward filling, or one
 32                of the valid method for xr.interpolate_na. Using methods other than 
 33                'ffill' on large arrays might lead to extremely high memory usage.
 34        '''
 35        self.name = name
 36        self.fillna_method = fillna_method 
 37
 38    def __init_subclass__(cls, *args, **kwargs):
 39        ''' Adds all subclasses to the context dictionary for reference in config files.
 40        '''
 41        super().__init_subclass__(*args, **kwargs)
 42        update_context(cls)
 43
 44    @staticmethod
 45    def from_config(name, source_config, **kwargs):
 46        """ Returns an datasource configured with the parameters in source_config. 
 47        
 48        This is not a classmethod, the class of the returned instance should be specified
 49        by setting the <source_type> attribute in the configuration dictionary.
 50        
 51        Extra keyword arguments are passed on verbatim to the source __init__ method
 52        
 53        Args:
 54            name: the name of the source to be created
 55            source_config: source configuration dictionary """
 56        
 57        source_config = copy.deepcopy(source_config) # Avoid modifying the original config
 58
 59        SourceClass = get_config_object(source_config.pop('__type__'))
 60        preprocess = source_config.pop('__preprocess__', {})   # dict of preprocessors
 61
 62        inst = SourceClass(name = name, **source_config, **kwargs)
 63
 64        # support for old-style preprocessors declaration
 65        if isinstance(preprocess, list):
 66            # if preprocess is a list, convert it to a dict with empty config
 67            preprocess = {preproc.pop('__name__'): {**preproc} for preproc in preprocess}
 68        
 69        #print(f"Source {name} preprocessors: {preprocess}")
 70        for preproc_name, preproc_config in preprocess.items():
 71            preprocessor = get_config_object(preproc_name)
 72            logger.debug(f"Source {name}, applying preprocessor: {preproc_name}")
 73            inst = preprocessor(inst, **preproc_config or {})
 74
 75        return inst
 76
 77    @staticmethod
 78    def from_name(name):
 79        """ Searchs for source named <name> in the loaded configuration and initializes is
 80        
 81        Args:
 82            name: the name of the source to be looked up in the config"""
 83        return DataSource.from_config(name, cfg.sources[name])
 84
 85    def _repr_rows(self):
 86        rows = []
 87        if self.fillna_method:
 88            rows.append(('Fill NA', self.fillna_method))
 89        return rows
 90
 91    def _repr_preprocessors(self):
 92        rows = []
 93        for p in getattr(self, '_preprocessors', []):
 94            parts = []
 95            if p.get('args'):
 96                parts += [str(a) for a in p['args']]
 97            if p.get('kwargs'):
 98                parts += [f'{k}={v}' for k, v in p['kwargs'].items()]
 99            rows.append((f'→ {p["name"]}', ','.join(parts)))
100        return rows
101
102    @abstractmethod
103    def load(self) -> xr.DataArray:
104        ''' Loads thata from the datasource
105
106        Returns: 
107            xr.DataArray: the loaded data, possibly represented by a lazy
108                dask.array. It must contain a dimension named 'train_id'
109        '''
110        pass
111
112    def __add__(self, other):
113        from ..instruments import Instrument
114
115        if not isinstance(other, DataSource):
116            return NotImplemented
117        return Instrument([self, other])
118
119    def _repr_html_(self):
120        rows_html = ''.join(
121            f'<tr><td style="padding:2px 8px; opacity:0.6; white-space:nowrap">{k}</td>'
122            f'<td style="padding:2px 8px; font-family:monospace">{v}</td></tr>'
123            for k, v in self._repr_rows() + self._repr_preprocessors()
124        )
125        return (
126            f'<table style="border-collapse:collapse; font-size:13px; font-family:sans-serif">'
127            f'<thead><tr><th colspan="2" style="background:rgba(100,130,150,0.35); padding:4px 10px; '
128            f'text-align:left; font-weight:normal">{type(self).__name__}</th></tr></thead>'
129            f'<tbody>{rows_html}</tbody></table>'
130        )