grids.py

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

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


def dog(grid, sigma1, sigma2):
    return gf(grid, sigma1) - gf(grid, sigma2)


def gf(grid, sigma, **kwargs):
    import scipy.ndimage
    return scipy.ndimage.gaussian_filter(grid, sigma=sigma, **kwargs)


functions = {}
functions["dog"] = dog
functions["gf"] = gf


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


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()

    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

    def normalize(self, array):
        mask = (np.isnan(array) | np.isinf(array))
        values = array * 1
        total = np.sum(values[~mask])
        return values / total

    def setter(self, key):
        def apply(value, key=key):
            self[key] = value
            return value
        return apply

    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]

    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)

    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

    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

    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


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)))