Source code for vaex.grids

__author__ = 'maartenbreddels'
import numpy as np
import vaex.vaexfast
from vaex.utils import filesize_format
import logging
import vaex.utils
import scipy.ndimage
total_bytes = 0

logger = logging.getLogger("vaex.grids")

[docs]def dog(grid, sigma1, sigma2): return gf(grid, sigma1) - gf(grid, sigma2)
[docs]def gf(grid, sigma, **kwargs): return scipy.ndimage.gaussian_filter(grid, sigma=sigma, **kwargs)
functions = {} functions["dog"] = dog functions["gf"] = gf
[docs]def grid_average(scope, counts_name="counts", weighted_name="weighted"): counts = scope.evaluate(counts_name) weighted = scope.evaluate(weighted_name) logger.debug("evaluating average: counts=%r weighted=%r" % (counts, weighted)) mask = counts==0 average = np.zeros(counts.shape, dtype=np.float64) average[~mask] = weighted[~mask]/counts[~mask] average[mask] = np.nan return average
[docs]class GridScope(object): def __init__(self, locals=None, globals=None): #self.locals = locals or {} #if locals: # self.__dict__.update(locals) self.globals = globals or {} self.lazy = {} self.lazy["average"] = grid_average self.globals["cumulative"] = self.cumulative self.globals["normalize"] = self.normalize self.globals.update(functions) self.user_added = set()
[docs] def cumulative(self, array, normalize=True): mask = (np.isnan(array) | np.isinf(array)) values = array * 1 values[mask] = 0 c = np.cumsum(values) if normalize: return c/c[-1] else: return c
[docs] def normalize(self, array): mask = (np.isnan(array) | np.isinf(array)) values = array * 1 total = np.sum(values[~mask]) return values/total
[docs] def setter(self, key): def apply(value, key=key): self[key] = value return value return apply
[docs] def add_lazy(self, key, f): self.lazy[key] = f
def __contains__(self, item): return (item in self.__dict__) or (item in self.lazy) def __setitem__(self, key, value): #logger.debug("%r.__setitem__(%r, %r)" % (self, key, value)) self.__dict__[key] = value self.user_added.add(key) def __getitem__(self, key): logger.debug("%r.__getitem__(%r)" % (self, key)) if key in self.lazy and key not in self.__dict__: logger.debug("%r lazy item %r" % (self, key)) f = self.lazy[key] self.__dict__[key] = f(self) return self.__dict__[key]
[docs] def evaluate(self, expression): if 0: locals = dict(self.__dict__) del locals["globals"] logger.debug("evaluating: %r locals=%r", expression, locals) return eval(expression, self.globals, self)
[docs] def slice(self, slice): gridscope = GridScope(globals=self.globals) for key in self.user_added: value = self[key] if isinstance(value, np.ndarray): grid = value sliced = np.sum(grid[slice,...], axis=0) logger.debug("sliced %s from %r to %r", key, grid.shape, sliced.shape) gridscope[key] = sliced else: gridscope[key] = value return gridscope
[docs] def disjoined(self): gridscope = GridScope(globals=self.globals) for key in self.user_added: value = self[key] if isinstance(value, np.ndarray): grid = vaex.utils.disjoined(value) gridscope[key] = grid else: gridscope[key] = value return gridscope
[docs] def marginal2d(self, i, j): gridscope = GridScope(globals=self.globals) for key in self.user_added: value = self[key] if isinstance(value, np.ndarray): dimension = len(value.shape) axes = list(range(dimension)) axes.remove(i) axes.remove(j) grid = vaex.utils.multisum(value, axes) gridscope[key] = grid else: gridscope[key] = value return gridscope
[docs]def add_mem(bytes, *info): global total_bytes total_bytes += bytes added = filesize_format(bytes) total = filesize_format(total_bytes) print(("MEMORY USAGE: added %s, total %s (%r)" % (added, total, info)))