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