Esempio n. 1
0
 def _build_header_database(self, database=':memory:'):
     """
     Builds a look-up table for trace header attributes.
     """
     # Make sure abcissa key is in the table (trace in file = rowid)
     header_keys = [self._get_header_alias(self.ABSCISSA_KEY)]
     # Add pick primary keys if we need them
     if self.pickdb is not None:
         pick_keys = self.pickdb._get_primary_fields(
             self.pickdb.TRACE_TABLE)
         for k in [self._get_header_alias(k) for k in pick_keys]:
             if k not in header_keys:
                 header_keys.append(k)
     # Build a database table for looking up header values
     self.sdb = SEGYHeaderDatabase(self.segy, *header_keys, 
                                   database=database)
Esempio n. 2
0
class SEGYPlotManager(object):
    """
    Plot manager for displaying SEG-Y data and picks.

    Handles plot layers, parameters, and unit conversions.
    """
    def __init__(self, ax, segy, pickdb=None,
                 trace_header_database=':memory:', **kwargs):
        """
        Set default plotter configuration parameters

        :param ax: :class:`matplotlib.axes.Axes` instance to manage.
        :param segy: :class:`SEGYFile` with data to plot.
        :param pickdb: Optional. :class:`PickDatabaseConnection` with picks 
            to plot. Default is ``None``.
        :param trace_header_database: Optional. Filename for the
            trace header attribute look-up table database.  Default                        is to create the database in memory.
        """
        # Layer managers
        self.ACTIVE_LINES = {}
        self.INACTIVE_LINES = {}
        self.ACTIVE_PATCHES = {}
        self.INACTIVE_PATCHES = {}
        # General parameters
        self.ABSCISSA_KEY = kwargs.pop('abscissa_key',
                'source_receiver_offset_in_m')
        # Data plotting parameters
        self.GAIN = kwargs.pop('gain', 1.0)
        self.CLIP = kwargs.pop('clip', 1.0) 
        self.NORMALIZATION_METHOD = kwargs.pop('normalization_method',
                'trace')
        self.OFFSET_GAIN_POWER = kwargs.pop('offset_gain_power', 0.0)
        self.WIGGLE_PEN_COLOR = kwargs.pop('wiggle_pen_color', 'black')
        self.WIGGLE_PEN_WIDTH = kwargs.pop('wiggle_pen_width', 0.5)
        self.NEG_FILL_COLOR = kwargs.pop('neg_fill_color', 'red') 
        self.POS_FILL_COLOR = kwargs.pop('pos_fill_color', 'black')
        # Unit parameters
        self.DISTANCE_UNIT = kwargs.pop('distance_unit', 'km')
        self.TIME_UNIT = kwargs.pop('time_unit', 's')
        self.SEGY_TIME_UNITS = {'time':'ms',
                                'delay':'ms',
                                'sample_interval':'microseconds'}
        self.SEGY_DISTANCE_UNITS = {'offset':'m'}
        # Define aliases for SEGY header words
        self.SEGY_HEADER_ALIASES = {
            'shot':'original_field_record_number',
            'channel':'trace_number_within_the_original_field_record',
            'ensemble':'ensemble_number',
            'trace':'trace_number_within_the_ensemble',
            'sample_interval':'sample_interval_in_ms_for_this_trace',
            'delay':'delay_recording_time_in_ms',
            'trace_in_file':'trace_sequence_number_within_segy_file',
            'npts':'number_of_samples_in_this_trace',
            'offset':'source_receiver_offset_in_m'}
        # Link to a segy file and pick database
        self.segy = segy
        self.pickdb = pickdb
        # Build header look up table, if we have picks to plot
        if self.pickdb is not None:
            self._build_header_database(database=trace_header_database)
        # Attach axes
        self.ax = ax
        # Set global range parameters and build interpolators
        self._set_global_ranges()
        self._init_interpolators()

    def _activate_line(self, plot_item):
        """
        Moves a line from the inactive to active lines dictionary and adds
        it to the plot axes.

        :param plot_item: Name of plot item.
        """
        self.ACTIVE_LINES[plot_item] = self.INACTIVE_LINES.pop(plot_item)
        for line in self.ACTIVE_LINES[plot_item]:
            self.ax.lines.append(line)

    def _deactivate_line(self, plot_item):
        """
        Moves a line from the active to inactive lines dictionary and removes
        it from the plot axes.

        :param plot_item: Name of plot item.
        """
        self.INACTIVE_LINES[plot_item] = self.ACTIVE_LINES.pop(plot_item)
        for line in self.INACTIVE_LINES[plot_item]:
            self.ax.lines.remove(line)

    def _activate_patch(self, plot_item):
        """
        Moves a patch from the inactive to active lines dictionary and adds
        it to the plot axes.

        :param plot_item: Name of plot item.
        """
        self.ACTIVE_PATCHES[plot_item] = self.INACTIVE_PATCHES.pop(plot_item)
        for patch in self.ACTIVE_PATCHES[plot_item]:
            self.ax.patches.append(patch)

    def _deactivate_patch(self, plot_item):
        """
        Moves a patch from the active to inactive lines dictionary and removes
        it from the plot axes.

        :param plot_item: Name of plot item.
        """
        self.INACTIVE_PATCHES[plot_item] = self.ACTIVE_PATCHES.pop(plot_item)
        for patch in self.INACTIVE_PATCHES[plot_item]:
            self.ax.patches.remove(patch)

    def _delete_line(self, plot_item):
        """
        Removes line from the layer manager dictionaries and the plot.
        
        :param plot_item: Name of plot item.
        """
        if plot_item in self.ACTIVE_LINES:
            for line in self.ACTIVE_LINES.pop(plot_item):
                self.ax.lines.remove(line)
        if plot_item in self.INACTIVE_LINES:
            del self.INACTIVE_LINES[plot_item]

    def _delete_patch(self, plot_item):
        """
        Removes patch from the layer manager dictionaries and the plot.
        
        :param plot_item: Name of plot item.
        """
        if plot_item in self.ACTIVE_PATCHES:
            for patch in self.ACTIVE_PATCHES.pop(plot_item):
                self.ax.patches.remove(patch)
        if plot_item in self.INACTIVE_PATCHES:
            del self.INACTIVE_PATCHES[plot_item]

    def _manage_layers(self, force_new=False, **kwargs):
        """
        Activate or deactivate plot items and return a dictionary of
        ``True``/``False`` values for items that need to be replotted.

        :param force_new: Optional, bool. If ``True``, delete plot items from
            the dictionary and return ``True`` for all items.
        :param **kwargs: Keyword=Bool arguments for plot items to
            activate/deactivate.  If plot item is not currently in one of the
            plot item dictionaries, it is set ``True``.
        :returns: Dictionary of ``True`` or ``False`` values for plot items. 
            If the plot item needs to be plotted, returns ``True``, else
            ``False``.
        """
        for plot_item in kwargs:
            if plot_item in self.ACTIVE_LINES or self.INACTIVE_LINES:
                if kwargs[plot_item] and force_new:
                    self._delete_line(plot_item)
                elif kwargs[plot_item] and plot_item in self.ACTIVE_LINES:
                    # item is already plotted
                    kwargs[plot_item] = False
                elif not kwargs[plot_item] and plot_item in self.ACTIVE_LINES:
                    # remove from current plot
                    kwargs[plot_item] = False
                    self._deactivate_line(plot_item)
                elif kwargs[plot_item] and plot_item in self.INACTIVE_LINES:
                    # just reactivate it
                    kwargs[plot_item] = False
                    self._activate_line(plot_item)
            elif plot_item in self.ACTIVE_PATCHES or self.INACTIVE_PATCHES:
                if kwargs[plot_item] and force_new:
                    self._delete_patch(plot_item)
                elif kwargs[plot_item] and plot_item in self.ACTIVE_PATCHES:
                    # item is already plotted
                    kwargs[plot_item] = False
                elif not kwargs[plot_item] and plot_item in self.ACTIVE_PATCHES:
                    # remove from current plot
                    kwargs[plot_item] = False
                    self._deactivate_patch(plot_item)
                elif kwargs[plot_item] and plot_item in self.INACTIVE_PATCHES:
                    # just reactivate it
                    kwargs[plot_item] = False
                    self._activate_patch(plot_item)
        logging.debug('active lines: %s', self.ACTIVE_LINES)
        logging.debug('inactive lines: %s', self.INACTIVE_LINES)
        logging.debug('active patches: %s', self.ACTIVE_PATCHES)
        logging.debug('inactive patches: %s', self.INACTIVE_PATCHES)
        logging.debug('items to plot: %s', kwargs)
        return kwargs


    def _build_header_database(self, database=':memory:'):
        """
        Builds a look-up table for trace header attributes.
        """
        # Make sure abcissa key is in the table (trace in file = rowid)
        header_keys = [self._get_header_alias(self.ABSCISSA_KEY)]
        # Add pick primary keys if we need them
        if self.pickdb is not None:
            pick_keys = self.pickdb._get_primary_fields(
                self.pickdb.TRACE_TABLE)
            for k in [self._get_header_alias(k) for k in pick_keys]:
                if k not in header_keys:
                    header_keys.append(k)
        # Build a database table for looking up header values
        self.sdb = SEGYHeaderDatabase(self.segy, *header_keys, 
                                      database=database)

    def get_time_array(self, header, convert_units=True):
        """
        Get an array of time for a single trace header.

        :param header: single ``SEGYTraceHeader`` instance
        :param convert_units: Bool, optional.  If true, convert units to plot
            units.
        :returns: ``list`` of time values
        """
        delay = self.get_header_value(header,'delay', 
                                      convert_units=convert_units)
        npts = self.get_header_value(header,'npts', 
                                     convert_units=convert_units)
        dt = self.get_header_value(header, 'sample_interval',
                                   convert_units=convert_units)
        return  list(delay + np.asarray(range(0,npts)) * dt)

    def get_header_value(self, header, key, convert_units=True):
        """
        Gets a value from trace header.  

        :param header: single ``SEGYTraceHeader`` object
        :param key: ``bool``, name of trace header attribute
        :param convert_units: ``bool``, convert units to plot units if true
        :returns: value of trace header attribute
        """
        _key = self._get_header_alias(key)
        try:   
            # use the unpacked value
            value = header.__getattribute__(_key)
        except AttributeError:
            value = header.__getattr__(_key)

        if convert_units is True:
            return self._convert_units(key, [value])[0]
        else:
            return value

    def _get_header_alias(self, key):
        """
        Finds the header attribute name for an alias.  If none exists, returns
        input ``key``.

        :param key: keyword ``str`` in ``HEADER_ALIASES``
        """
        try:
            return self.SEGY_HEADER_ALIASES[key.lower()]
        except KeyError:
            return key

    def _convert_units(self, key, values):
        """
        Convert SEGY units to plot units.
        """
        _units = self._get_units(key)
        if _units is not None:
            scl =  float(units.__getattribute__(_units[0])\
                         /units.__getattribute__(_units[1]))
            return [v * scl for v in values]
        else:
            return values

    def _get_units(self, key):
        """
        :param key: Trace header attribute name.
        :returns: (<segy units>, <plot units>) if ``key`` is in SEGY_TIME_UNITS
            or SEGY_DISTANCE_UNITS, or ``None`` if units are not specified.
        """
        _key = self._get_header_alias(key)
        if key in self.SEGY_TIME_UNITS:
            return (self.SEGY_TIME_UNITS[key], self.TIME_UNIT)
        elif _key in self.SEGY_TIME_UNITS:
            return (self.SEGY_TIME_UNITS[_key], self.TIME_UNIT)
        elif key in self.SEGY_DISTANCE_UNITS:
            return (self.SEGY_DISTANCE_UNITS[key], self.DISTANCE_UNIT)
        elif _key in self.SEGY_DISTANCE_UNITS:
            return (self.SEGY_DISTANCE_UNITS[_key], self.DISTANCE_UNIT)
        else:
            return None

    def _get_abscissa(self, keys, values, convert_units=True):
        """
        Get plot abscissa values from the trace header database.
        """
        _keys = [self._get_header_alias(k) for k in keys]
        sql = "SELECT " + self._get_header_alias(self.ABSCISSA_KEY)
        sql += " FROM " + self.sdb.TRACE_TABLE
        sql += " WHERE " + " and ".join(['%s=?' %k for k in _keys])
        x = []
        for row in values:
            msg = 'calling sdb.execute(' + str(sql) + ', '
            msg += str(row) + ')'
            logging.debug(msg)
            _x = self.sdb.execute(sql, row).fetchone()
            if _x:
                x.append(_x[0])
            else:
                x.append(None)
        return x

    def _set_global_ranges(self):
        """
        Set parameters based on the range of all traces.
        """
        self.X = [self.get_header_value(tr.header, self.ABSCISSA_KEY) \
                  for tr in self.segy.traces]
        self.DX = abs(np.mean(np.diff(self.X)))
        self.T = self.get_time_array(self.segy.traces[0].header)
        self.XLIMITS = [self.X[0], self.X[-1]]
        self.TLIMITS = [self.T[0], self.T[-1]]
        self.AMP_MAX = np.max([tr.data for tr in self.segy.traces])

    def _init_interpolators(self):
        """
        Creates the methods ``abscissa2trace(xplt)`` and 
        ``trace2abscissa(idx)`` for converting an x-coordinate to
        a trace index and visa versa.
        """
        ntrc = len(self.segy.traces)
        idx = range(0, ntrc)
        self.trace2abscissa = interp1d(idx, self.X, kind='nearest')
        self.abscissa2trace = interp1d(self.X, idx, kind='nearest')