1
/*
2
 * This file implements assignment from a scalar to an ndarray.
3
 *
4
 * Written by Mark Wiebe (mwwiebe@gmail.com)
5
 * Copyright (c) 2011 by Enthought, Inc.
6
 *
7
 * See LICENSE.txt for the license.
8
 */
9

10
#define PY_SSIZE_T_CLEAN
11
#include <Python.h>
12

13
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
14
#define _MULTIARRAYMODULE
15
#include <numpy/ndarraytypes.h>
16

17
#include "npy_config.h"
18
#include "npy_pycompat.h"
19

20
#include "convert_datatype.h"
21
#include "methods.h"
22
#include "shape.h"
23
#include "lowlevel_strided_loops.h"
24

25
#include "array_assign.h"
26

27
/*
28
 * Assigns the scalar value to every element of the destination raw array.
29
 *
30
 * Returns 0 on success, -1 on failure.
31
 */
32
NPY_NO_EXPORT int
33 1
raw_array_assign_scalar(int ndim, npy_intp const *shape,
34
        PyArray_Descr *dst_dtype, char *dst_data, npy_intp const *dst_strides,
35
        PyArray_Descr *src_dtype, char *src_data)
36
{
37
    int idim;
38
    npy_intp shape_it[NPY_MAXDIMS], dst_strides_it[NPY_MAXDIMS];
39
    npy_intp coord[NPY_MAXDIMS];
40

41 1
    PyArray_StridedUnaryOp *stransfer = NULL;
42 1
    NpyAuxData *transferdata = NULL;
43 1
    int aligned, needs_api = 0;
44 1
    npy_intp src_itemsize = src_dtype->elsize;
45

46 1
    NPY_BEGIN_THREADS_DEF;
47

48
    /* Check both uint and true alignment */
49 1
    aligned = raw_array_is_aligned(ndim, shape, dst_data, dst_strides,
50 1
                                   npy_uint_alignment(dst_dtype->elsize)) &&
51 1
              raw_array_is_aligned(ndim, shape, dst_data, dst_strides,
52 1
                                   dst_dtype->alignment) &&
53 1
              npy_is_aligned(src_data, npy_uint_alignment(src_dtype->elsize) &&
54 1
              npy_is_aligned(src_data, src_dtype->alignment));
55

56
    /* Use raw iteration with no heap allocation */
57 1
    if (PyArray_PrepareOneRawArrayIter(
58
                    ndim, shape,
59
                    dst_data, dst_strides,
60
                    &ndim, shape_it,
61
                    &dst_data, dst_strides_it) < 0) {
62
        return -1;
63
    }
64

65
    /* Get the function to do the casting */
66 1
    if (PyArray_GetDTypeTransferFunction(aligned,
67
                        0, dst_strides_it[0],
68
                        src_dtype, dst_dtype,
69
                        0,
70
                        &stransfer, &transferdata,
71
                        &needs_api) != NPY_SUCCEED) {
72
        return -1;
73
    }
74

75 1
    if (!needs_api) {
76
        npy_intp nitems = 1, i;
77 1
        for (i = 0; i < ndim; i++) {
78 1
            nitems *= shape_it[i];
79
        }
80 1
        NPY_BEGIN_THREADS_THRESHOLDED(nitems);
81
    }
82

83 1
    NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
84
        /* Process the innermost dimension */
85 1
        if (stransfer(
86
                dst_data, dst_strides_it[0], src_data, 0,
87
                shape_it[0], src_itemsize, transferdata) < 0) {
88
            goto fail;
89
        }
90 1
    } NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord,
91
                            shape_it, dst_data, dst_strides_it);
92

93 1
    NPY_END_THREADS;
94 1
    NPY_AUXDATA_FREE(transferdata);
95
    return 0;
96 1
fail:
97 1
    NPY_END_THREADS;
98 1
    NPY_AUXDATA_FREE(transferdata);
99
    return -1;
100
}
101

102
/*
103
 * Assigns the scalar value to every element of the destination raw array
104
 * where the 'wheremask' value is True.
105
 *
106
 * Returns 0 on success, -1 on failure.
107
 */
108
NPY_NO_EXPORT int
109 1
raw_array_wheremasked_assign_scalar(int ndim, npy_intp const *shape,
110
        PyArray_Descr *dst_dtype, char *dst_data, npy_intp const *dst_strides,
111
        PyArray_Descr *src_dtype, char *src_data,
112
        PyArray_Descr *wheremask_dtype, char *wheremask_data,
113
        npy_intp const *wheremask_strides)
114
{
115
    int idim;
116
    npy_intp shape_it[NPY_MAXDIMS], dst_strides_it[NPY_MAXDIMS];
117
    npy_intp wheremask_strides_it[NPY_MAXDIMS];
118
    npy_intp coord[NPY_MAXDIMS];
119

120 1
    PyArray_MaskedStridedUnaryOp *stransfer = NULL;
121 1
    NpyAuxData *transferdata = NULL;
122 1
    int aligned, needs_api = 0;
123 1
    npy_intp src_itemsize = src_dtype->elsize;
124

125 1
    NPY_BEGIN_THREADS_DEF;
126

127
    /* Check both uint and true alignment */
128 1
    aligned = raw_array_is_aligned(ndim, shape, dst_data, dst_strides,
129 1
                                   npy_uint_alignment(dst_dtype->elsize)) &&
130 1
              raw_array_is_aligned(ndim, shape, dst_data, dst_strides,
131 1
                                   dst_dtype->alignment) &&
132 1
              npy_is_aligned(src_data, npy_uint_alignment(src_dtype->elsize) &&
133 1
              npy_is_aligned(src_data, src_dtype->alignment));
134

135
    /* Use raw iteration with no heap allocation */
136 1
    if (PyArray_PrepareTwoRawArrayIter(
137
                    ndim, shape,
138
                    dst_data, dst_strides,
139
                    wheremask_data, wheremask_strides,
140
                    &ndim, shape_it,
141
                    &dst_data, dst_strides_it,
142
                    &wheremask_data, wheremask_strides_it) < 0) {
143
        return -1;
144
    }
145

146
    /* Get the function to do the casting */
147 1
    if (PyArray_GetMaskedDTypeTransferFunction(aligned,
148
                        0, dst_strides_it[0], wheremask_strides_it[0],
149
                        src_dtype, dst_dtype, wheremask_dtype,
150
                        0,
151
                        &stransfer, &transferdata,
152
                        &needs_api) != NPY_SUCCEED) {
153
        return -1;
154
    }
155

156 1
    if (!needs_api) {
157
        npy_intp nitems = 1, i;
158 1
        for (i = 0; i < ndim; i++) {
159 1
            nitems *= shape_it[i];
160
        }
161 1
        NPY_BEGIN_THREADS_THRESHOLDED(nitems);
162
    }
163

164 1
    NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
165
        /* Process the innermost dimension */
166 1
        stransfer(dst_data, dst_strides_it[0], src_data, 0,
167
                    (npy_bool *)wheremask_data, wheremask_strides_it[0],
168
                    shape_it[0], src_itemsize, transferdata);
169 1
    } NPY_RAW_ITER_TWO_NEXT(idim, ndim, coord, shape_it,
170
                            dst_data, dst_strides_it,
171
                            wheremask_data, wheremask_strides_it);
172

173 1
    NPY_END_THREADS;
174

175 1
    NPY_AUXDATA_FREE(transferdata);
176

177 1
    return (needs_api && PyErr_Occurred()) ? -1 : 0;
178
}
179

180
/*
181
 * Assigns a scalar value specified by 'src_dtype' and 'src_data'
182
 * to elements of 'dst'.
183
 *
184
 * dst: The destination array.
185
 * src_dtype: The data type of the source scalar.
186
 * src_data: The memory element of the source scalar.
187
 * wheremask: If non-NULL, a boolean mask specifying where to copy.
188
 * casting: An exception is raised if the assignment violates this
189
 *          casting rule.
190
 *
191
 * This function is implemented in array_assign_scalar.c.
192
 *
193
 * Returns 0 on success, -1 on failure.
194
 */
195
NPY_NO_EXPORT int
196 1
PyArray_AssignRawScalar(PyArrayObject *dst,
197
                        PyArray_Descr *src_dtype, char *src_data,
198
                        PyArrayObject *wheremask,
199
                        NPY_CASTING casting)
200
{
201 1
    int allocated_src_data = 0;
202
    npy_longlong scalarbuffer[4];
203

204 1
    if (PyArray_FailUnlessWriteable(dst, "assignment destination") < 0) {
205
        return -1;
206
    }
207

208
    /* Check the casting rule */
209 1
    if (!can_cast_scalar_to(src_dtype, src_data,
210
                            PyArray_DESCR(dst), casting)) {
211 1
        npy_set_invalid_cast_error(
212
                src_dtype, PyArray_DESCR(dst), casting, NPY_TRUE);
213 1
        return -1;
214
    }
215

216
    /*
217
     * Make a copy of the src data if it's a different dtype than 'dst'
218
     * or isn't aligned, and the destination we're copying to has
219
     * more than one element. To avoid having to manage object lifetimes,
220
     * we also skip this if 'dst' has an object dtype.
221
     */
222 1
    if ((!PyArray_EquivTypes(PyArray_DESCR(dst), src_dtype) ||
223 1
            !(npy_is_aligned(src_data, npy_uint_alignment(src_dtype->elsize)) &&
224 1
              npy_is_aligned(src_data, src_dtype->alignment))) &&
225 1
                    PyArray_SIZE(dst) > 1 &&
226 1
                    !PyDataType_REFCHK(PyArray_DESCR(dst))) {
227
        char *tmp_src_data;
228

229
        /*
230
         * Use a static buffer to store the aligned/cast version,
231
         * or allocate some memory if more space is needed.
232
         */
233 1
        if ((int)sizeof(scalarbuffer) >= PyArray_DESCR(dst)->elsize) {
234
            tmp_src_data = (char *)&scalarbuffer[0];
235
        }
236
        else {
237 1
            tmp_src_data = PyArray_malloc(PyArray_DESCR(dst)->elsize);
238 1
            if (tmp_src_data == NULL) {
239 0
                PyErr_NoMemory();
240
                goto fail;
241
            }
242
            allocated_src_data = 1;
243
        }
244

245 1
        if (PyDataType_FLAGCHK(PyArray_DESCR(dst), NPY_NEEDS_INIT)) {
246 1
            memset(tmp_src_data, 0, PyArray_DESCR(dst)->elsize);
247
        }
248

249 1
        if (PyArray_CastRawArrays(1, src_data, tmp_src_data, 0, 0,
250
                            src_dtype, PyArray_DESCR(dst), 0) != NPY_SUCCEED) {
251
            src_data = tmp_src_data;
252
            goto fail;
253
        }
254

255
        /* Replace src_data/src_dtype */
256 1
        src_data = tmp_src_data;
257 1
        src_dtype = PyArray_DESCR(dst);
258
    }
259

260 1
    if (wheremask == NULL) {
261
        /* A straightforward value assignment */
262
        /* Do the assignment with raw array iteration */
263 1
        if (raw_array_assign_scalar(PyArray_NDIM(dst), PyArray_DIMS(dst),
264 1
                PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
265
                src_dtype, src_data) < 0) {
266
            goto fail;
267
        }
268
    }
269
    else {
270
        npy_intp wheremask_strides[NPY_MAXDIMS];
271

272
        /* Broadcast the wheremask to 'dst' for raw iteration */
273 1
        if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
274 1
                    PyArray_NDIM(wheremask), PyArray_DIMS(wheremask),
275 1
                    PyArray_STRIDES(wheremask), "where mask",
276
                    wheremask_strides) < 0) {
277
            goto fail;
278
        }
279

280
        /* Do the masked assignment with raw array iteration */
281 1
        if (raw_array_wheremasked_assign_scalar(
282 1
                PyArray_NDIM(dst), PyArray_DIMS(dst),
283 1
                PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
284
                src_dtype, src_data,
285 1
                PyArray_DESCR(wheremask), PyArray_DATA(wheremask),
286
                wheremask_strides) < 0) {
287
            goto fail;
288
        }
289
    }
290

291 1
    if (allocated_src_data) {
292 1
        PyArray_free(src_data);
293
    }
294

295
    return 0;
296

297 1
fail:
298 1
    if (allocated_src_data) {
299 0
        PyArray_free(src_data);
300
    }
301

302
    return -1;
303
}

Read our documentation on viewing source code .

Loading