Source code for pypret.mesh_data

""" This module implements an object for dealing with two-dimensional data.
"""
import numpy as np
from scipy.interpolate import RegularGridInterpolator
from . import lib
from . import io


[docs]class MeshData(io.IO): _io_store = ["data", "axes", "labels", "units", "uncertainty"]
[docs] def __init__(self, data, *axes, uncertainty=None, labels=None, units=None): """ Creates a MeshData instance. Parameters ---------- data : ndarray A at least two-dimensional array containing the data. *axes : ndarray Arrays specifying the coordinates of the data axes. Must be given in indexing order. uncertainty : ndarray An ndarray of the same size as `data` that contains some measure of the uncertainty of the meshdata. E.g., it could be the standard deviation of the data. labels : list of str, optional A list of strings labeling the axes. The last element labels the data itself, e.g. ``labels`` must have one more element than the number of axes. units : list of str, optional A list of unit strings. """ self.data = data.copy() self.axes = [np.array(a).copy() for a in axes] if uncertainty is not None: self.uncertainty = uncertainty.copy() else: self.uncertainty = None if self.ndim != len(axes): raise ValueError("Number of supplied axes is wrong!") if self.shape != tuple(ax.size for ax in self.axes): raise ValueError("Shape of supplied axes is wrong!") self.labels = labels if self.labels is None: self.labels = ["" for ax in self.axes] self.units = units if self.units is None: self.units = ["" for ax in self.axes]
@property def shape(self): """ Returns the shape of the data as a tuple. """ return self.data.shape @property def ndim(self): """ Returns the dimension of the data as integer. """ return self.data.ndim
[docs] def copy(self): """ Creates a copy of the MeshData instance. """ return MeshData(self.data, *self.axes, uncertainty=self.uncertainty, labels=self.labels, units=self.units)
[docs] def marginals(self, normalize=False, axes=None): """ Calculates the marginals of the data. axes specifies the axes of the marginals, e.g., the axes on which the sum is projected. """ return lib.marginals(self.data, normalize=normalize, axes=axes)
[docs] def normalize(self): """ Normalizes the maximum of the data to 1. """ self.scale(1.0 / self.data.max())
[docs] def scale(self, scale): if self.uncertainty is not None: self.uncertainty *= scale self.data *= scale
[docs] def autolimit(self, *axes, threshold=1e-2, padding=0.25): """ Limits the data based on the marginals. """ if len(axes) == 0: # default: operate on all axes axes = list(range(self.ndim)) marginals = lib.marginals(self.data) limits = [] for i, j in enumerate(axes): limit = lib.limit(self.axes[j], marginals[j], threshold=threshold, padding=padding) limits.append(limit) self.limit(*limits, axes=axes)
[docs] def limit(self, *limits, axes=None): """ Limits the data range of this instance. Parameters ---------- *limits : tuples The data limits in the axes as tuples. Has to match the dimension of the data or the number of axes specified in the `axes` parameter. axes : tuple or None The axes in which the limit is applied. Default is `None` in which case all axes are selected. """ if axes is None: # default: operate on all axes axes = list(range(self.ndim)) axes = lib.as_list(axes) if len(axes) != len(limits): raise ValueError("Number of limits must match the specified axes!") slices = [] for j in range(self.ndim): if j in axes: i = axes.index(j) ax = self.axes[j] x1, x2 = limits[i] # do it this way as we cannot assume them to be sorted... idx1 = np.argmin(np.abs(ax - x1)) idx2 = np.argmin(np.abs(ax - x2)) if idx1 > idx2: idx1, idx2 = idx2, idx1 elif idx1 == idx2: raise ValueError('Selected empty slice along axis %d!' % i) slices.append(slice(idx1, idx2 + 1)) else: # empty slice slices.append(slice(None)) self.axes[j] = self.axes[j][slices[-1]] self.data = self.data[(*slices,)] if self.uncertainty is not None: self.uncertainty = self.uncertainty[(*slices,)]
[docs] def interpolate(self, axis1=None, axis2=None, degree=2, sorted=False): """ Interpolates the data on a new two-dimensional, equidistantly spaced grid. """ axes = [axis1, axis2] for i in range(self.ndim): if axes[i] is None: axes[i] = self.axes[i] # FITPACK can only deal with strictly increasing axes # so sort them beforehand if necessary... orig_axes = self.axes data = self.data.copy() if self.uncertainty is not None: uncertainty = self.uncertainty.copy() if not sorted: for i in range(len(orig_axes)): idx = np.argsort(orig_axes[i]) orig_axes[i] = orig_axes[i][idx] data = np.take(data, idx, axis=i) if self.uncertainty is not None: uncertainty = np.take(uncertainty, idx, axis=i) dataf = RegularGridInterpolator(tuple(orig_axes), data, bounds_error=False, fill_value=0.0) grid = lib.build_coords(*axes) self.data = dataf(grid) self.axes = axes if self.uncertainty is not None: dataf = RegularGridInterpolator(tuple(orig_axes), uncertainty, bounds_error=False, fill_value=0.0) self.uncertainty = dataf(grid)
[docs] def flip(self, *axes): """ Flips the data on the specified axes. """ if len(axes) == 0: return axes = lib.as_list(axes) slices = [slice(None) for ax in self.axes] for ax in axes: self.axes[ax] = self.axes[ax][::-1] slices[ax] = slice(None, None, -1) self.data = self.data[slices] if self.uncertainty is not None: self.uncertainty = self.uncertainty[slices]