forked from inaos/iron-array-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexpression.py
More file actions
153 lines (130 loc) · 4.82 KB
/
Copy pathexpression.py
File metadata and controls
153 lines (130 loc) · 4.82 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
###########################################################################################
# Copyright INAOS GmbH, Thalwil, 2018.
# Copyright Francesc Alted, 2018.
#
# All rights reserved.
#
# This software is the confidential and proprietary information of INAOS GmbH
# and Francesc Alted ("Confidential Information"). You shall not disclose such Confidential
# Information and shall use it only in accordance with the terms of the license agreement.
###########################################################################################
import iarray as ia
from iarray import iarray_ext as ext
from iarray import py2llvm
# The main expression class
class Expr(ext.Expression):
"""A class that is meant to hold an expression.
This is not meant to be called directly from user space.
See Also
--------
expr_from_string
expr_from_udf
"""
def __init__(self, shape, cfg=None, **kwargs):
if cfg is None:
cfg = ia.get_config_defaults()
default_shapes = check_expr_config(cfg, **kwargs)
with ia.config(cfg=cfg, shape=shape, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
self.cfg = cfg
super().__init__(self.cfg)
super().bind_out_properties(dtshape)
if default_shapes:
# Set cfg chunks and blocks to None to detect that we want the default shapes when evaluating
self.cfg.chunks = None
self.cfg.blocks = None
def eval(self) -> ia.IArray:
"""Evaluate the expression in self.
Returns
-------
:ref:`IArray`
The output array.
"""
return super().eval()
def check_inputs(inputs: list, shape):
if inputs:
first_input = inputs[0]
for input_ in inputs[1:]:
if first_input.shape != input_.shape:
raise ValueError("Inputs should have the same shape")
if first_input.dtype != input_.dtype:
raise TypeError("Inputs should have the same dtype")
return first_input.shape, first_input.dtype
else:
cfg = ia.get_config_defaults()
if shape is None:
raise AttributeError("A shape is needed")
return shape, cfg.dtype
def expr_from_string(sexpr: str, inputs: dict, cfg: ia.Config = None, **kwargs) -> Expr:
"""Create an :class:`Expr` instance from an expression in string form.
Parameters
----------
sexpr : str
An expression in string format.
inputs : dict
Map of variables in `sexpr` to actual arrays.
cfg : :class:`Config`
The configuration for running the expression.
If None (default), global defaults are used.
kwargs : dict
A dictionary for setting some or all of the fields in the :class:`Config`
dataclass that should override the current configuration.
Returns
-------
:class:`Expr`
An expression ready to be evaluated via :func:`Expr.eval`.
See Also
--------
expr_from_udf
"""
with ia.config(cfg, **kwargs):
shape, dtype = check_inputs(list(inputs.values()), None)
kwargs["dtype"] = dtype
expr = Expr(shape=shape, cfg=cfg, **kwargs)
for i in inputs:
expr.bind(i, inputs[i])
expr.compile(sexpr)
return expr
def expr_from_udf(udf: py2llvm.Function, inputs: list, shape=None, cfg=None, **kwargs) -> Expr:
"""Create an :class:`Expr` instance from an UDF function.
Parameters
----------
udf : py2llvm.Function
A User Defined Function.
inputs : list
List of arrays whose values are passed as arguments, after the output,
to the UDF function.
cfg : :class:`Config`
The configuration for running the expression.
If None (default), global defaults are used.
kwargs : dict
A dictionary for setting some or all of the fields in the :class:`Config`
dataclass that should override the current configuration.
Returns
-------
:class:`Expr`
An expression ready to be evaluated via :func:`Expr.eval`.
See Also
--------
expr_from_string
"""
with ia.config(cfg, shape=shape, **kwargs):
shape, dtype = check_inputs(inputs, shape)
kwargs["dtype"] = dtype
expr = Expr(shape=shape, cfg=cfg, **kwargs)
for i in inputs:
expr.bind("", i)
expr.compile_udf(udf)
return expr
def check_expr_config(cfg=None, **kwargs):
# Check if the chunks and blocks are explicitly set
default_shapes = False
if (cfg is not None and cfg.chunks is None and cfg.blocks is None) or cfg is None:
shape_params = {"chunks", "blocks"}
if kwargs != {}:
not_kw_shapes = all(x not in kwargs for x in shape_params)
if not_kw_shapes:
default_shapes = True
else:
default_shapes = True
return default_shapes