def __mul__(self, rhs): if isinstance(rhs, w_type): # Apply all the dimensions to get # produce a type for dim in reversed(self.dims): rhs = dim.create(rhs) return rhs elif isinstance(rhs, (str, type)): # Allow: # ndt.strided * 'int32' # ndt.strided * int rhs = w_type(rhs) for dim in reversed(self.dims): rhs = dim.create(rhs) return rhs elif isinstance(rhs, _Dim): # Combine the dimension fragments return _DimFragment(self.dims + rhs.dims) else: raise TypeError('Expected a dynd dimension or type, not %r' % rhs)
def add_computed_fields(n, fields, rm_fields=[], fnname=None): """ Adds one or more new fields to a struct array, using nd.elwise_map to create the deferred object. Each field_expr should be a string or bit of code that can be evaluated with an 'eval' call. It is called with numpy/scipy in the globals, and the input fields in the locals. Parameters ---------- n : dynd array This should have a struct data dtype. The result will be a view of this data. fields : list of (field_name, field_type, field_expr) These are the fields which are added to 'n'. rm_fields : list of string, optional For fields that are in the input, but have no expression, this removes them from the output struct instead of keeping the value. fnname : string, optional The function name, which affects how the resulting deferred expression's dtype is printed. Examples -------- >>> from dynd import nd, ndt >>> import numpy as np >>> x = np.array([(2, 0), (0, -2), (3, 5), (4, 4)], ... dtype=[('x', np.float64), ('y', np.float64)]) >>> y = nd.add_computed_fields(x, ... fields=[('r', np.float64, 'sqrt(x*x + y*y)'), ... ('theta', np.float64, 'arctan2(y, x)')], ... rm_fields=['x', 'y'], ... fnname='topolar') >>> y.dtype ndt.type('strided_dim<expr<cstruct<float64 r, float64 theta>, op0=cstruct<float64 x, float64 y>, expr=topolar(op0)>>') >>> y.eval() nd.array([[2, 0], [2, -1.5708], [5.83095, 1.03038], [5.65685, 0.785398]], strided_dim<cstruct<float64 r, float64 theta>>) >>> x[0] = (-100, 0) >>> y[0].eval() nd.array([100, 3.14159], cstruct<float64 r, float64 theta>) """ n = array(n) udt = dtype_of(n).value_type if udt.kind != 'struct': raise ValueError("parameter 'n' must have kind 'struct'") # The field names and types of the input struct field_names = as_py(udt.field_names) field_types = as_py(udt.field_types) # Put the new field names in a dict as well new_field_dict = {} for fn, ft, fe in fields: new_field_dict[fn] = w_type(ft) # Create the output struct type and corresponding expressions new_field_names = [] new_field_types = [] new_field_expr = [] for fn, ft in zip(field_names, field_types): if fn not in new_field_dict and fn not in rm_fields: new_field_names.append(fn) new_field_types.append(ft) new_field_expr.append(fn) for fn, ft, fe in fields: new_field_names.append(fn) new_field_types.append(ft) new_field_expr.append(fe) result_udt = make_cstruct(new_field_types, new_field_names) fieldexpr = FieldExpr(new_field_expr, field_names, fnname) return elwise_map([n], fieldexpr, result_udt)