1# distutils: include_dirs = LIB_DIR
2# distutils: libraries = STD_LIBS
3"""
4Particle Deposition onto Cells
5
6
7
8
9"""
10
11
12cimport numpy as np
13
14import numpy as np
15
16cimport cython
17from cython.view cimport memoryview as cymemview
18from libc.math cimport sqrt
19from libc.stdlib cimport free, malloc
20from oct_container cimport Oct, OctInfo, OctreeContainer
21
22from yt.utilities.lib.fp_utils cimport *
23
24from yt.utilities.lib.misc_utilities import OnceIndirect
25
26
27cdef append_axes(np.ndarray arr, int naxes):
28    if arr.ndim == naxes:
29        return arr
30    # Avoid copies
31    arr2 = arr.view()
32    arr2.shape = arr2.shape + (1,) * (naxes - arr2.ndim)
33    return arr2
34
35cdef class ParticleDepositOperation:
36    def __init__(self, nvals, kernel_name):
37        # nvals is a tuple containing the active dimensions of the
38        # grid to deposit onto and the number of grids,
39        # (nx, ny, nz, ngrids)
40        self.nvals = nvals
41        self.update_values = 0 # This is the default
42        self.sph_kernel = get_kernel_func(kernel_name)
43
44    def initialize(self, *args):
45        raise NotImplementedError
46
47    def finalize(self, *args):
48        raise NotImplementedError
49
50    @cython.boundscheck(False)
51    @cython.wraparound(False)
52    def process_octree(self, OctreeContainer octree,
53                     np.ndarray[np.int64_t, ndim=1] dom_ind,
54                     np.ndarray[np.float64_t, ndim=2] positions,
55                     fields = None, int domain_id = -1,
56                     int domain_offset = 0, lvlmax = None):
57        cdef int nf, i, j
58        if fields is None:
59            fields = []
60        nf = len(fields)
61        cdef np.float64_t[::cython.view.indirect, ::1] field_pointers
62        if nf > 0: field_pointers = OnceIndirect(fields)
63        cdef np.float64_t pos[3]
64        cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64")
65        cdef int dims[3]
66        dims[0] = dims[1] = dims[2] = (1 << octree.oref)
67        cdef int nz = dims[0] * dims[1] * dims[2]
68        cdef OctInfo oi
69        cdef np.int64_t offset, moff
70        cdef Oct *oct
71        cdef np.int64_t numpart = positions.shape[0]
72        cdef np.int8_t use_lvlmax
73        moff = octree.get_domain_offset(domain_id + domain_offset)
74        if lvlmax is None:
75            use_lvlmax = False
76            lvlmax = []
77        else:
78            use_lvlmax = True
79        cdef np.ndarray[np.int32_t, ndim=1] lvlmaxval = np.asarray(lvlmax, dtype=np.int32)
80
81        for i in range(positions.shape[0]):
82            # We should check if particle remains inside the Oct here
83            for j in range(nf):
84                field_vals[j] = field_pointers[j,i]
85            for j in range(3):
86                pos[j] = positions[i, j]
87            # This line should be modified to have it return the index into an
88            # array based on whatever cutting of the domain we have done.  This
89            # may or may not include the domain indices that we have
90            # previously generated.  This way we can support not knowing the
91            # full octree structure.  All we *really* care about is some
92            # arbitrary offset into a field value for deposition.
93            if not use_lvlmax:
94                oct = octree.get(pos, &oi)
95            else:
96                oct = octree.get(pos, &oi, max_level=lvlmaxval[i])
97            # This next line is unfortunate.  Basically it says, sometimes we
98            # might have particles that belong to octs outside our domain.
99            # For the distributed-memory octrees, this will manifest as a NULL
100            # oct.  For the non-distributed memory octrees, we'll simply see
101            # this as a domain_id that is not the current domain id.  Note that
102            # this relies on the idea that all the particles in a region are
103            # all fed to sequential domain subsets, which will not be true with
104            # RAMSES, where we *will* miss particles that live in ghost
105            # regions on other processors.  Addressing this is on the TODO
106            # list.
107            if oct == NULL or (domain_id > 0 and oct.domain != domain_id):
108                continue
109            # Note that this has to be our local index, not our in-file index.
110            offset = dom_ind[oct.domain_ind - moff]
111            if offset < 0: continue
112            # Check that we found the oct ...
113            self.process(dims, i, oi.left_edge, oi.dds,
114                         offset, pos, field_vals, oct.domain_ind)
115            if self.update_values == 1:
116                for j in range(nf):
117                    field_pointers[j][i] = field_vals[j]
118
119    @cython.boundscheck(False)
120    @cython.wraparound(False)
121    def process_grid(self, gobj,
122                     np.ndarray[np.float64_t, ndim=2] positions,
123                     fields = None):
124        cdef int nf, i, j
125        if fields is None:
126            fields = []
127        nf = len(fields)
128        cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64")
129        cdef np.float64_t[::cython.view.indirect, ::1] field_pointers
130        if nf > 0: field_pointers = OnceIndirect(fields)
131        cdef np.float64_t pos[3]
132        cdef np.int64_t gid = getattr(gobj, "id", -1)
133        cdef np.float64_t dds[3]
134        cdef np.float64_t left_edge[3]
135        cdef np.float64_t right_edge[3]
136        cdef int dims[3]
137        for i in range(3):
138            dds[i] = gobj.dds[i]
139            left_edge[i] = gobj.LeftEdge[i]
140            right_edge[i] = gobj.RightEdge[i]
141            dims[i] = gobj.ActiveDimensions[i]
142        for i in range(positions.shape[0]):
143            # Now we process
144            for j in range(nf):
145                field_vals[j] = field_pointers[j,i]
146            for j in range(3):
147                pos[j] = positions[i, j]
148            continue_loop = False
149            for j in range(3):
150                if pos[j] < left_edge[j] or pos[j] > right_edge[j]:
151                    continue_loop = True
152            if continue_loop:
153                continue
154            self.process(dims, i, left_edge, dds, 0, pos, field_vals, gid)
155            if self.update_values == 1:
156                for j in range(nf):
157                    field_pointers[j][i] = field_vals[j]
158
159    cdef int process(self, int dim[3], int ipart, np.float64_t left_edge[3],
160                     np.float64_t dds[3], np.int64_t offset,
161                     np.float64_t ppos[3], np.float64_t[:] fields,
162                     np.int64_t domain_ind) nogil except -1:
163        with gil:
164            raise NotImplementedError
165
166cdef class CountParticles(ParticleDepositOperation):
167    cdef np.int64_t[:,:,:,:] count
168    def initialize(self):
169        # Create a numpy array accessible to python
170        self.count = append_axes(
171            np.zeros(self.nvals, dtype="int64", order='F'), 4)
172
173    @cython.cdivision(True)
174    @cython.boundscheck(False)
175    cdef int process(self, int dim[3], int ipart,
176                     np.float64_t left_edge[3],
177                     np.float64_t dds[3],
178                     np.int64_t offset, # offset into IO field
179                     np.float64_t ppos[3], # this particle's position
180                     np.float64_t[:] fields,
181                     np.int64_t domain_ind
182                     ) nogil except -1:
183        # here we do our thing; this is the kernel
184        cdef int ii[3]
185        cdef int i
186        for i in range(3):
187            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
188        self.count[ii[2], ii[1], ii[0], offset] += 1
189        return 0
190
191    def finalize(self):
192        arr = np.asarray(self.count)
193        arr.shape = self.nvals
194        return arr.astype("float64")
195
196deposit_count = CountParticles
197
198cdef class SimpleSmooth(ParticleDepositOperation):
199    # Note that this does nothing at the edges.  So it will give a poor
200    # estimate there, and since Octrees are mostly edges, this will be a very
201    # poor SPH kernel.
202    cdef np.float64_t[:,:,:,:] data
203    cdef np.float64_t[:,:,:,:] temp
204
205    def initialize(self):
206        self.data = append_axes(
207            np.zeros(self.nvals, dtype="float64", order='F'), 4)
208        self.temp = append_axes(
209            np.zeros(self.nvals, dtype="float64", order='F'), 4)
210
211    @cython.cdivision(True)
212    @cython.boundscheck(False)
213    cdef int process(self, int dim[3], int ipart,
214                     np.float64_t left_edge[3],
215                     np.float64_t dds[3],
216                     np.int64_t offset,
217                     np.float64_t ppos[3],
218                     np.float64_t[:] fields,
219                     np.int64_t domain_ind
220                     ) nogil except -1:
221        cdef int ii[3]
222        cdef int ib0[3]
223        cdef int ib1[3]
224        cdef int i, j, k, half_len
225        cdef np.float64_t idist[3]
226        cdef np.float64_t kernel_sum, dist
227        # Smoothing length is fields[0]
228        kernel_sum = 0.0
229        for i in range(3):
230            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
231            half_len = <int>(fields[0]/dds[i]) + 1
232            ib0[i] = ii[i] - half_len
233            ib1[i] = ii[i] + half_len
234            if ib0[i] >= dim[i] or ib1[i] <0:
235                return 0
236            ib0[i] = iclip(ib0[i], 0, dim[i] - 1)
237            ib1[i] = iclip(ib1[i], 0, dim[i] - 1)
238        for i from ib0[0] <= i <= ib1[0]:
239            idist[0] = (ii[0] - i) * dds[0]
240            idist[0] *= idist[0]
241            for j from ib0[1] <= j <= ib1[1]:
242                idist[1] = (ii[1] - j) * dds[1]
243                idist[1] *= idist[1]
244                for k from ib0[2] <= k <= ib1[2]:
245                    idist[2] = (ii[2] - k) * dds[2]
246                    idist[2] *= idist[2]
247                    dist = idist[0] + idist[1] + idist[2]
248                    # Calculate distance in multiples of the smoothing length
249                    dist = sqrt(dist) / fields[0]
250                    with gil:
251                        self.temp[k,j,i,offset] = self.sph_kernel(dist)
252                    kernel_sum += self.temp[k,j,i,offset]
253        # Having found the kernel, deposit accordingly into gdata
254        for i from ib0[0] <= i <= ib1[0]:
255            for j from ib0[1] <= j <= ib1[1]:
256                for k from ib0[2] <= k <= ib1[2]:
257                    dist = self.temp[k,j,i,offset] / kernel_sum
258                    self.data[k,j,i,offset] += fields[1] * dist
259        return 0
260
261    def finalize(self):
262        return self.odata
263
264deposit_simple_smooth = SimpleSmooth
265
266cdef class SumParticleField(ParticleDepositOperation):
267    cdef np.float64_t[:,:,:,:] sum
268    def initialize(self):
269        self.sum = append_axes(
270            np.zeros(self.nvals, dtype="float64", order='F'), 4)
271
272    @cython.cdivision(True)
273    @cython.boundscheck(False)
274    cdef int process(self, int dim[3], int ipart,
275                     np.float64_t left_edge[3],
276                     np.float64_t dds[3],
277                     np.int64_t offset,
278                     np.float64_t ppos[3],
279                     np.float64_t[:] fields,
280                     np.int64_t domain_ind
281                     ) nogil except -1:
282        cdef int ii[3]
283        cdef int i
284        for i in range(3):
285            ii[i] = <int>((ppos[i] - left_edge[i]) / dds[i])
286        self.sum[ii[2], ii[1], ii[0], offset] += fields[0]
287        return 0
288
289    def finalize(self):
290        sum = np.asarray(self.sum)
291        sum.shape = self.nvals
292        return sum
293
294deposit_sum = SumParticleField
295
296cdef class StdParticleField(ParticleDepositOperation):
297    # Thanks to Britton and MJ Turk for the link
298    # to a single-pass STD
299    # http://www.cs.berkeley.edu/~mhoemmen/cs194/Tutorials/variance.pdf
300    cdef np.float64_t[:,:,:,:] mk
301    cdef np.float64_t[:,:,:,:] qk
302    cdef np.float64_t[:,:,:,:] i
303    def initialize(self):
304        # we do this in a single pass, but need two scalar
305        # per cell, M_k, and Q_k and also the number of particles
306        # deposited into each one
307        # the M_k term
308        self.mk = append_axes(
309            np.zeros(self.nvals, dtype="float64", order='F'), 4)
310        self.qk = append_axes(
311            np.zeros(self.nvals, dtype="float64", order='F'), 4)
312        self.i = append_axes(
313            np.zeros(self.nvals, dtype="float64", order='F'), 4)
314
315    @cython.cdivision(True)
316    @cython.boundscheck(False)
317    cdef int process(self, int dim[3], int ipart,
318                     np.float64_t left_edge[3],
319                     np.float64_t dds[3],
320                     np.int64_t offset,
321                     np.float64_t ppos[3],
322                     np.float64_t[:] fields,
323                     np.int64_t domain_ind
324                     ) nogil except -1:
325        cdef int ii[3]
326        cdef int i, cell_index
327        cdef float k, mk, qk
328        for i in range(3):
329            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
330        k = self.i[ii[2], ii[1], ii[0], offset]
331        mk = self.mk[ii[2], ii[1], ii[0], offset]
332        qk = self.qk[ii[2], ii[1], ii[0], offset]
333        if k == 0.0:
334            # Initialize cell values
335            self.mk[ii[2], ii[1], ii[0], offset] = fields[0]
336        else:
337            self.mk[ii[2], ii[1], ii[0], offset] = mk + (fields[0] - mk) / k
338            self.qk[ii[2], ii[1], ii[0], offset] = \
339                qk + (k - 1.0) * (fields[0] - mk) * (fields[0] - mk) / k
340        self.i[ii[2], ii[1], ii[0], offset] += 1
341        return 0
342
343    def finalize(self):
344        # This is the standard variance
345        # if we want sample variance divide by (self.oi - 1.0)
346        i = np.asarray(self.i)
347        std2 = np.asarray(self.qk) / i
348        std2[i == 0.0] = 0.0
349        std2.shape = self.nvals
350        return np.sqrt(std2)
351
352deposit_std = StdParticleField
353
354cdef class CICDeposit(ParticleDepositOperation):
355    cdef np.float64_t[:,:,:,:] field
356    cdef public object ofield
357    def initialize(self):
358        if not all(_ > 1 for _ in self.nvals[:-1]):
359            from yt.utilities.exceptions import YTBoundsDefinitionError
360            raise YTBoundsDefinitionError(
361                "CIC requires minimum of 2 zones in all spatial dimensions.",
362                self.nvals[:-1])
363        self.field = append_axes(
364            np.zeros(self.nvals, dtype="float64", order='F'), 4)
365
366    @cython.cdivision(True)
367    @cython.boundscheck(False)
368    cdef int process(self, int dim[3], int ipart,
369                     np.float64_t left_edge[3],
370                     np.float64_t dds[3],
371                     np.int64_t offset, # offset into IO field
372                     np.float64_t ppos[3], # this particle's position
373                     np.float64_t[:] fields,
374                     np.int64_t domain_ind
375                     ) nogil except -1:
376
377        cdef int i, j, k
378        cdef np.uint64_t ii
379        cdef int ind[3]
380        cdef np.float64_t rpos[3]
381        cdef np.float64_t rdds[3][2]
382        cdef np.float64_t fact, edge0, edge1, edge2
383        cdef np.float64_t le0, le1, le2
384        cdef np.float64_t dx, dy, dz, dx2, dy2, dz2
385
386        # Compute the position of the central cell
387        for i in range(3):
388            rpos[i] = (ppos[i]-left_edge[i])/dds[i]
389            rpos[i] = fclip(rpos[i], 0.5001, dim[i]-0.5001)
390            ind[i] = <int> (rpos[i] + 0.5)
391            # Note these are 1, then 0
392            rdds[i][1] = (<np.float64_t> ind[i]) + 0.5 - rpos[i]
393            rdds[i][0] = 1.0 - rdds[i][1]
394
395        for i in range(2):
396            for j in range(2):
397                for k in range(2):
398                    self.field[ind[2] - k, ind[1] - j, ind[0] - i, offset] += \
399                        fields[0]*rdds[0][i]*rdds[1][j]*rdds[2][k]
400
401        return 0
402
403    def finalize(self):
404        rv = np.asarray(self.field)
405        rv.shape = self.nvals
406        return rv
407
408deposit_cic = CICDeposit
409
410cdef class WeightedMeanParticleField(ParticleDepositOperation):
411    # Deposit both mass * field and mass into two scalars
412    # then in finalize divide mass * field / mass
413    cdef np.float64_t[:,:,:,:] wf
414    cdef np.float64_t[:,:,:,:] w
415    def initialize(self):
416        self.wf = append_axes(
417            np.zeros(self.nvals, dtype='float64', order='F'), 4)
418        self.w = append_axes(
419            np.zeros(self.nvals, dtype='float64', order='F'), 4)
420
421    @cython.cdivision(True)
422    @cython.boundscheck(False)
423    cdef int process(self, int dim[3], int ipart,
424                     np.float64_t left_edge[3],
425                     np.float64_t dds[3],
426                     np.int64_t offset,
427                     np.float64_t ppos[3],
428                     np.float64_t[:] fields,
429                     np.int64_t domain_ind
430                     ) nogil except -1:
431        cdef int ii[3]
432        cdef int i
433        for i in range(3):
434            ii[i] = <int>((ppos[i] - left_edge[i]) / dds[i])
435        self.w[ii[2], ii[1], ii[0], offset] += fields[1]
436        self.wf[ii[2], ii[1], ii[0], offset] += fields[0] * fields[1]
437        return 0
438
439    def finalize(self):
440        wf = np.asarray(self.wf)
441        w = np.asarray(self.w)
442        with np.errstate(divide='ignore', invalid='ignore'):
443            rv = wf / w
444        rv.shape = self.nvals
445        return rv
446
447deposit_weighted_mean = WeightedMeanParticleField
448
449cdef class MeshIdentifier(ParticleDepositOperation):
450    # This is a tricky one!  What it does is put into the particle array the
451    # value of the oct or block (grids will always be zero) identifier that a
452    # given particle resides in
453    def initialize(self):
454        self.update_values = 1
455
456    @cython.cdivision(True)
457    @cython.boundscheck(False)
458    cdef int process(self, int dim[3], int ipart,
459                      np.float64_t left_edge[3],
460                      np.float64_t dds[3],
461                      np.int64_t offset,
462                      np.float64_t ppos[3],
463                      np.float64_t[:] fields,
464                      np.int64_t domain_ind
465                      ) nogil except -1:
466        fields[0] = domain_ind
467        return 0
468
469    def finalize(self):
470        return
471
472deposit_mesh_id = MeshIdentifier
473
474cdef class CellIdentifier(ParticleDepositOperation):
475    cdef np.int64_t[:] indexes, cell_index
476    # This method stores the offset of the grid containing each particle
477    # and compute the index of the cell in the oct.
478    def initialize(self, int npart):
479        self.indexes = np.zeros(npart, dtype=np.int64, order='F') - 1
480        self.cell_index = np.zeros(npart, dtype=np.int64, order='F') - 1
481
482    @cython.cdivision(True)
483    @cython.boundscheck(False)
484    cdef int process(self, int dim[3], int ipart,
485                      np.float64_t left_edge[3],
486                      np.float64_t dds[3],
487                      np.int64_t offset,
488                      np.float64_t ppos[3],
489                      np.float64_t[:] fields,
490                      np.int64_t domain_ind
491                      ) nogil except -1:
492        cdef int i, icell
493        self.indexes[ipart] = offset
494
495        icell = 0
496        for i in range(3):
497            if ppos[i] > left_edge[i] + dds[i]:
498                icell |= 4 >> i
499
500        # Compute cell index
501        self.cell_index[ipart] = icell
502
503        return 0
504
505    def finalize(self):
506        return np.asarray(self.indexes), np.asarray(self.cell_index)
507
508deposit_cell_id = CellIdentifier
509
510cdef class NNParticleField(ParticleDepositOperation):
511    cdef np.float64_t[:,:,:,:] nnfield
512    cdef np.float64_t[:,:,:,:] distfield
513    def initialize(self):
514        self.nnfield = append_axes(
515            np.zeros(self.nvals, dtype="float64", order='F'), 4)
516        self.distfield = append_axes(
517            np.zeros(self.nvals, dtype="float64", order='F'), 4)
518        self.distfield[:] = np.inf
519
520    @cython.cdivision(True)
521    @cython.boundscheck(False)
522    cdef int process(self, int dim[3], int ipart,
523                     np.float64_t left_edge[3],
524                     np.float64_t dds[3],
525                     np.int64_t offset,
526                     np.float64_t ppos[3],
527                     np.float64_t[:] fields,
528                     np.int64_t domain_ind
529                     ) nogil except -1:
530        # This one is a bit slow.  Every grid cell is going to be iterated
531        # over, and we're going to deposit particles in it.
532        cdef int i, j, k
533        cdef int ii[3]
534        cdef np.float64_t r2
535        cdef np.float64_t gpos[3]
536        gpos[0] = left_edge[0] + 0.5 * dds[0]
537        for i in range(dim[0]):
538            gpos[1] = left_edge[1] + 0.5 * dds[1]
539            for j in range(dim[1]):
540                gpos[2] = left_edge[2] + 0.5 * dds[2]
541                for k in range(dim[2]):
542                    r2 = ((ppos[0] - gpos[0])*(ppos[0] - gpos[0]) +
543                          (ppos[1] - gpos[1])*(ppos[1] - gpos[1]) +
544                          (ppos[2] - gpos[2])*(ppos[2] - gpos[2]))
545                    if r2 < self.distfield[k,j,i,offset]:
546                        self.distfield[k,j,i,offset] = r2
547                        self.nnfield[k,j,i,offset] = fields[0]
548                    gpos[2] += dds[2]
549                gpos[1] += dds[1]
550            gpos[0] += dds[0]
551        return 0
552
553    def finalize(self):
554        nn = np.asarray(self.nnfield)
555        nn.shape = self.nvals
556        return nn
557
558deposit_nearest = NNParticleField
559