mask.py 2.7 KB
Newer Older
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
1
2
3
4
5
6
7
8
9
10
11
12
13
#!/usr/bin/env
# encoding: utf-8
"""
Author:     Daniel Boeckenhoff
Mail:       daniel.boeckenhoff@ipp.mpg.de

part of tfields library
contains interaction methods for sympy and numpy
"""
import numpy as np
import sympy


dboe's avatar
dboe committed
14
def evalf(array, cut_expression=None, coords=None):
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
15
    """
dboe's avatar
dboe committed
16
    Linking sympy and numpy by retrieving a mask according to the cut_expression
17

Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
18
19
    Args:
        array (numpy ndarray)
dboe's avatar
dboe committed
20
        cut_expression (sympy logical expression)
21
        coord_sys (str): coord_sys to evalfuate the expression in.
22
    Returns:
dboe's avatar
dboe committed
23
        np.array: mask which is True, where cut_expression evalfuates True.
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
24
25
26
27
28
29
    Examples:
        >>> import sympy
        >>> import numpy as np
        >>> import tfields
        >>> x, y, z = sympy.symbols('x y z')

dboe's avatar
dboe committed
30
31
32
33
34
        >>> a = np.array([[1., 2., 3.], [4., 5., 6.], [1, 2, -6],
        ...               [-5, -5, -5], [1,0,-1], [0,1,-1]])
        >>> assert np.array_equal(
        ...     tfields.evalf(a, x > 0),
        ...     np.array([ True, True, True, False, True, False]))
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
35
36

        And combination
dboe's avatar
dboe committed
37
38
39
        >>> assert np.array_equal(
        ...     tfields.evalf(a, (x > 0) & (y < 3)),
        ...     np.array([True, False, True, False, True, False]))
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
40
41

        Or combination
dboe's avatar
dboe committed
42
43
44
        >>> assert np.array_equal(
        ...     tfields.evalf(a, (x > 0) | (y > 3)),
        ...     np.array([True, True, True, False, True, False]))
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
45

dboe's avatar
dboe committed
46
47
        If array of other shape than (?, 3) is given, the coords need to be
        specified
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
48
        >>> a0, a1 = sympy.symbols('a0 a1')
dboe's avatar
dboe committed
49
50
51
        >>> assert np.array_equal(
        ...     tfields.evalf([[0., 1.], [-1, 3]], a1 > 2, coords=[a0, a1]),
        ...     np.array([False,  True], dtype=bool))
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
52

53
        >= is taken care of
dboe's avatar
dboe committed
54
55
56
        >>> assert np.array_equal(
        ...     tfields.evalf(a, x >= 0),
        ...     np.array([ True, True, True, False, True, True]))
57

Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
58
59
60
    """
    if isinstance(array, list):
        array = np.array(array)
dboe's avatar
dboe committed
61
    if cut_expression is None:
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
62
63
64
65
66
        return np.full((array.shape[0]), False, dtype=bool)
    if len(array.shape) != 2:
        raise NotImplementedError("Array shape other than 2")
    if coords is None:
        if array.shape[1] == 3:
dboe's avatar
dboe committed
67
            coords = sympy.symbols("x y z")
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
68
69
70
        else:
            raise ValueError("coords are None and shape is not (?, 3)")
    elif len(coords) != array.shape[1]:
dboe's avatar
dboe committed
71
72
73
        raise ValueError(
            "Length of coords is not {0} but {1}".format(array.shape[1], len(coords))
        )
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
74

dboe's avatar
dboe committed
75
76
77
    pre_mask = sympy.utilities.lambdify(
        coords, cut_expression, modules={"&": np.logical_and, "|": np.logical_or}
    )
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
78

dboe's avatar
dboe committed
79
    mask = np.array([pre_mask(*vals) for vals in array], dtype=bool)
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
80
81
82
83

    return mask


dboe's avatar
dboe committed
84
if __name__ == "__main__":  # pragma: no cover
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
85
    import doctest
dboe's avatar
dboe committed
86

Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
87
    doctest.testmod()