nd::arrfunc dynd::make_rolling_arrfunc(const nd::arrfunc &window_op,
intptr_t window_size)
{
// Validate the input arrfunc
if (window_op.is_null()) {
throw invalid_argument("make_rolling_arrfunc() 'window_op' cannot be null");
}
const arrfunc_type *window_af_tp = window_op.get_type();
if (window_af_tp->get_npos() != 1) {
stringstream ss;
ss << "To make a rolling window arrfunc, an operation with one "
"argument is required, got " << window_af_tp;
throw invalid_argument(ss.str());
}
const ndt::type &window_src_tp = window_af_tp->get_arg_type(0);
if (window_src_tp.get_ndim() < 1) {
stringstream ss;
ss << "To make a rolling window arrfunc, an operation with which "
"accepts a dimension is required, got " << window_af_tp;
throw invalid_argument(ss.str());
}
nd::string rolldimname("RollDim");
ndt::type roll_src_tp = ndt::make_typevar_dim(
rolldimname, window_src_tp.get_type_at_dimension(NULL, 1));
ndt::type roll_dst_tp =
ndt::make_typevar_dim(rolldimname, window_af_tp->get_return_type());
nd::array af = nd::empty(ndt::make_funcproto(roll_src_tp, roll_dst_tp));
arrfunc_type_data *out_af =
reinterpret_cast<arrfunc_type_data *>(af.get_readwrite_originptr());
// Create the data for the arrfunc
rolling_arrfunc_data *data = new rolling_arrfunc_data;
*out_af->get_data_as<rolling_arrfunc_data *>() = data;
out_af->free = &free_rolling_arrfunc_data;
out_af->resolve_dst_type = &resolve_rolling_dst_type;
out_af->instantiate = &instantiate_strided;
data->window_size = window_size;
data->window_op = window_op;
af.flag_as_immutable();
return af;
}
nd::arrfunc nd::functional::rolling(const nd::arrfunc &window_op,
intptr_t window_size)
{
// Validate the input arrfunc
if (window_op.is_null()) {
throw invalid_argument("make_rolling_arrfunc() 'window_op' cannot be null");
}
const ndt::arrfunc_type *window_af_tp = window_op.get_type();
if (window_af_tp->get_npos() != 1) {
stringstream ss;
ss << "To make a rolling window arrfunc, an operation with one "
"argument is required, got " << window_af_tp;
throw invalid_argument(ss.str());
}
const ndt::type &window_src_tp = window_af_tp->get_pos_type(0);
if (window_src_tp.get_ndim() < 1) {
stringstream ss;
ss << "To make a rolling window arrfunc, an operation with which "
"accepts a dimension is required, got " << window_af_tp;
throw invalid_argument(ss.str());
}
nd::string rolldimname("RollDim");
ndt::type roll_src_tp = ndt::make_typevar_dim(
rolldimname, window_src_tp.get_type_at_dimension(NULL, 1));
ndt::type roll_dst_tp =
ndt::make_typevar_dim(rolldimname, window_af_tp->get_return_type());
// Create the data for the arrfunc
std::shared_ptr<rolling_arrfunc_data> data(new rolling_arrfunc_data);
data->window_size = window_size;
data->window_op = window_op;
return arrfunc::make<rolling_ck>(
ndt::make_arrfunc(ndt::make_tuple(roll_src_tp), roll_dst_tp), data, 0);
}
void dynd::lift_reduction_arrfunc(arrfunc_type_data *out_ar,
const nd::arrfunc& elwise_reduction_arr,
const ndt::type& lifted_arr_type,
const nd::arrfunc& dst_initialization_arr,
bool keepdims,
intptr_t reduction_ndim,
const bool *reduction_dimflags,
bool associative,
bool commutative,
bool right_associative,
const nd::array& reduction_identity)
{
// Validate the input elwise_reduction arrfunc
if (elwise_reduction_arr.is_null()) {
throw runtime_error("lift_reduction_arrfunc: 'elwise_reduction' may not be empty");
}
const arrfunc_type_data *elwise_reduction = elwise_reduction_arr.get();
if (elwise_reduction->get_param_count() != 1 &&
!(elwise_reduction->get_param_count() == 2 &&
elwise_reduction->get_param_type(0) ==
elwise_reduction->get_param_type(1) &&
elwise_reduction->get_param_type(0) ==
elwise_reduction->get_return_type())) {
stringstream ss;
ss << "lift_reduction_arrfunc: 'elwise_reduction' must contain a"
" unary operation ckernel or a binary expr ckernel with all "
"equal types, its prototype is " << elwise_reduction->func_proto;
throw invalid_argument(ss.str());
}
lifted_reduction_arrfunc_data *self = new lifted_reduction_arrfunc_data;
*out_ar->get_data_as<lifted_reduction_arrfunc_data *>() = self;
out_ar->free_func = &delete_lifted_reduction_arrfunc_data;
self->child_elwise_reduction = elwise_reduction_arr;
self->child_dst_initialization = dst_initialization_arr;
if (!reduction_identity.is_null()) {
if (reduction_identity.is_immutable() &&
reduction_identity.get_type() == elwise_reduction->get_return_type()) {
self->reduction_identity = reduction_identity;
} else {
self->reduction_identity = nd::empty(elwise_reduction->get_return_type());
self->reduction_identity.vals() = reduction_identity;
self->reduction_identity.flag_as_immutable();
}
}
// Figure out the result type
ndt::type lifted_dst_type = elwise_reduction->get_return_type();
for (intptr_t i = reduction_ndim - 1; i >= 0; --i) {
if (reduction_dimflags[i]) {
if (keepdims) {
lifted_dst_type = ndt::make_strided_dim(lifted_dst_type);
}
} else {
ndt::type subtype = lifted_arr_type.get_type_at_dimension(NULL, i);
switch (subtype.get_type_id()) {
case strided_dim_type_id:
case cfixed_dim_type_id:
lifted_dst_type = ndt::make_strided_dim(lifted_dst_type);
break;
case var_dim_type_id:
lifted_dst_type = ndt::make_var_dim(lifted_dst_type);
break;
default: {
stringstream ss;
ss << "lift_reduction_arrfunc: don't know how to process ";
ss << "dimension of type " << subtype;
throw type_error(ss.str());
}
}
}
}
self->data_types[0] = lifted_dst_type;
self->data_types[1] = lifted_arr_type;
self->reduction_ndim = reduction_ndim;
self->associative = associative;
self->commutative = commutative;
self->right_associative = right_associative;
self->reduction_dimflags.init(reduction_ndim);
memcpy(self->reduction_dimflags.get(), reduction_dimflags, sizeof(bool) * reduction_ndim);
out_ar->instantiate = &instantiate_lifted_reduction_arrfunc_data;
out_ar->func_proto = ndt::make_funcproto(lifted_arr_type, lifted_dst_type);
}