1# distutils: sources = ARTIO_SOURCE 2# distutils: include_dirs = LIB_DIR_GEOM_ARTIO 3cimport cython 4 5import numpy as np 6 7cimport numpy as np 8 9import sys 10 11from libc.stdlib cimport free, malloc 12from libc.string cimport memcpy 13 14from yt.geometry.oct_container cimport OctObjectPool, SparseOctreeContainer 15from yt.geometry.oct_visitors cimport Oct 16from yt.geometry.particle_deposit cimport ParticleDepositOperation 17from yt.geometry.selection_routines cimport ( 18 AlwaysSelector, 19 OctreeSubsetSelector, 20 SelectorObject, 21) 22from yt.utilities.lib.fp_utils cimport imax 23 24import data_structures 25 26from yt.utilities.lib.misc_utilities import OnceIndirect 27 28 29cdef extern from "platform_dep.h": 30 ctypedef int int32_t 31 ctypedef long long int64_t 32 void *alloca(int) 33 34cdef extern from "cosmology.h": 35 ctypedef struct CosmologyParameters "CosmologyParameters" : 36 pass 37 CosmologyParameters *cosmology_allocate() 38 void cosmology_free(CosmologyParameters *c) 39 void cosmology_set_fixed(CosmologyParameters *c) 40 41 void cosmology_set_OmegaM(CosmologyParameters *c, double value) 42 void cosmology_set_OmegaB(CosmologyParameters *c, double value) 43 void cosmology_set_OmegaL(CosmologyParameters *c, double value) 44 void cosmology_set_h(CosmologyParameters *c, double value) 45 void cosmology_set_DeltaDC(CosmologyParameters *c, double value) 46 47 double abox_from_auni(CosmologyParameters *c, double a) nogil 48 double tcode_from_auni(CosmologyParameters *c, double a) nogil 49 double tphys_from_auni(CosmologyParameters *c, double a) nogil 50 51 double auni_from_abox(CosmologyParameters *c, double v) nogil 52 double auni_from_tcode(CosmologyParameters *c, double v) nogil 53 double auni_from_tphys(CosmologyParameters *c, double v) nogil 54 55 double abox_from_tcode(CosmologyParameters *c, double tcode) nogil 56 double tcode_from_abox(CosmologyParameters *c, double abox) nogil 57 58 double tphys_from_abox(CosmologyParameters *c, double abox) nogil 59 double tphys_from_tcode(CosmologyParameters *c, double tcode) nogil 60 61cdef extern from "artio.h": 62 ctypedef struct artio_fileset_handle "artio_fileset" : 63 pass 64 ctypedef struct artio_selection "artio_selection" : 65 pass 66 ctypedef struct artio_context : 67 pass 68 cdef extern artio_context *artio_context_global 69 70 # open modes 71 cdef int ARTIO_OPEN_HEADER "ARTIO_OPEN_HEADER" 72 cdef int ARTIO_OPEN_GRID "ARTIO_OPEN_GRID" 73 cdef int ARTIO_OPEN_PARTICLES "ARTIO_OPEN_PARTICLES" 74 75 # parameter constants 76 cdef int ARTIO_TYPE_STRING "ARTIO_TYPE_STRING" 77 cdef int ARTIO_TYPE_CHAR "ARTIO_TYPE_CHAR" 78 cdef int ARTIO_TYPE_INT "ARTIO_TYPE_INT" 79 cdef int ARTIO_TYPE_FLOAT "ARTIO_TYPE_FLOAT" 80 cdef int ARTIO_TYPE_DOUBLE "ARTIO_TYPE_DOUBLE" 81 cdef int ARTIO_TYPE_LONG "ARTIO_TYPE_LONG" 82 83 cdef int ARTIO_MAX_STRING_LENGTH "ARTIO_MAX_STRING_LENGTH" 84 85 cdef int ARTIO_PARAMETER_EXHAUSTED "ARTIO_PARAMETER_EXHAUSTED" 86 87 # grid read options 88 cdef int ARTIO_READ_LEAFS "ARTIO_READ_LEAFS" 89 cdef int ARTIO_READ_REFINED "ARTIO_READ_REFINED" 90 cdef int ARTIO_READ_ALL "ARTIO_READ_ALL" 91 cdef int ARTIO_READ_REFINED_NOT_ROOT "ARTIO_READ_REFINED_NOT_ROOT" 92 cdef int ARTIO_RETURN_CELLS "ARTIO_RETURN_CELLS" 93 cdef int ARTIO_RETURN_OCTS "ARTIO_RETURN_OCTS" 94 95 # errors 96 cdef int ARTIO_SUCCESS "ARTIO_SUCCESS" 97 cdef int ARTIO_ERR_MEMORY_ALLOCATION "ARTIO_ERR_MEMORY_ALLOCATION" 98 99 artio_fileset_handle *artio_fileset_open(char *file_prefix, int type, artio_context *context ) 100 int artio_fileset_close( artio_fileset_handle *handle ) 101 int artio_fileset_open_particle( artio_fileset_handle *handle ) 102 int artio_fileset_open_grid(artio_fileset_handle *handle) 103 int artio_fileset_close_grid(artio_fileset_handle *handle) 104 105 int artio_fileset_has_grid( artio_fileset_handle *handle ) 106 int artio_fileset_has_particles( artio_fileset_handle *handle ) 107 108 # selection functions 109 artio_selection *artio_selection_allocate( artio_fileset_handle *handle ) 110 artio_selection *artio_select_all( artio_fileset_handle *handle ) 111 artio_selection *artio_select_volume( artio_fileset_handle *handle, double lpos[3], double rpos[3] ) 112 int artio_selection_add_root_cell( artio_selection *selection, int coords[3] ) 113 int artio_selection_destroy( artio_selection *selection ) 114 int artio_selection_iterator( artio_selection *selection, 115 int64_t max_range_size, int64_t *start, int64_t *end ) 116 int64_t artio_selection_size( artio_selection *selection ) 117 void artio_selection_print( artio_selection *selection ) 118 119 # parameter functions 120 int artio_parameter_iterate( artio_fileset_handle *handle, char *key, int *type, int *length ) 121 int artio_parameter_get_int_array(artio_fileset_handle *handle, char * key, int length, int32_t *values) 122 int artio_parameter_get_float_array(artio_fileset_handle *handle, char * key, int length, float *values) 123 int artio_parameter_get_long_array(artio_fileset_handle *handle, char * key, int length, int64_t *values) 124 int artio_parameter_get_double_array(artio_fileset_handle *handle, char * key, int length, double *values) 125 int artio_parameter_get_string_array(artio_fileset_handle *handle, char * key, int length, char **values ) 126 127 # grid functions 128 int artio_grid_cache_sfc_range(artio_fileset_handle *handle, int64_t start, int64_t end) 129 int artio_grid_clear_sfc_cache( artio_fileset_handle *handle ) 130 131 int artio_grid_read_root_cell_begin(artio_fileset_handle *handle, int64_t sfc, 132 double *pos, float *variables, 133 int *num_tree_levels, int *num_octs_per_level) 134 int artio_grid_read_root_cell_end(artio_fileset_handle *handle) 135 136 int artio_grid_read_level_begin(artio_fileset_handle *handle, int level ) 137 int artio_grid_read_level_end(artio_fileset_handle *handle) 138 139 int artio_grid_read_oct(artio_fileset_handle *handle, double *pos, 140 float *variables, int *refined) 141 142 int artio_grid_count_octs_in_sfc_range(artio_fileset_handle *handle, 143 int64_t start, int64_t end, int64_t *num_octs) 144 145 #particle functions 146 int artio_fileset_open_particles(artio_fileset_handle *handle) 147 int artio_particle_read_root_cell_begin(artio_fileset_handle *handle, int64_t sfc, 148 int * num_particle_per_species) 149 int artio_particle_read_root_cell_end(artio_fileset_handle *handle) 150 int artio_particle_read_particle(artio_fileset_handle *handle, int64_t *pid, int *subspecies, 151 double *primary_variables, float *secondary_variables) 152 int artio_particle_cache_sfc_range(artio_fileset_handle *handle, int64_t sfc_start, int64_t sfc_end) 153 int artio_particle_clear_sfc_cache(artio_fileset_handle *handle) 154 int artio_particle_read_species_begin(artio_fileset_handle *handle, int species) 155 int artio_particle_read_species_end(artio_fileset_handle *handle) 156 157 158cdef extern from "artio_internal.h": 159 np.int64_t artio_sfc_index( artio_fileset_handle *handle, int coords[3] ) nogil 160 void artio_sfc_coords( artio_fileset_handle *handle, int64_t index, int coords[3] ) nogil 161 162cdef void check_artio_status(int status, char *fname="[unknown]"): 163 if status != ARTIO_SUCCESS: 164 import traceback 165 traceback.print_stack() 166 callername = sys._getframe().f_code.co_name 167 nline = sys._getframe().f_lineno 168 raise RuntimeError('failure with status', status, 'in function',fname,'from caller', callername, nline) 169 170cdef class artio_fileset : 171 cdef public object parameters 172 cdef artio_fileset_handle *handle 173 174 # cosmology parameter for time unit conversion 175 cdef CosmologyParameters *cosmology 176 cdef float tcode_to_years 177 178 # common attributes 179 cdef public int num_grid 180 cdef int64_t num_root_cells 181 cdef int64_t sfc_min, sfc_max 182 183 # grid attributes 184 cdef public int has_grid 185 cdef public int min_level, max_level 186 cdef public int num_grid_variables 187 cdef int *num_octs_per_level 188 cdef float *grid_variables 189 190 # particle attributes 191 cdef public int has_particles 192 cdef public int num_species 193 cdef int *particle_position_index 194 cdef int *num_particles_per_species 195 cdef double *primary_variables 196 cdef float *secondary_variables 197 198 def __init__(self, char *file_prefix) : 199 cdef int artio_type = ARTIO_OPEN_HEADER 200 cdef int64_t num_root 201 202 self.handle = artio_fileset_open( file_prefix, artio_type, artio_context_global ) 203 if not self.handle : 204 raise RuntimeError 205 206 self.read_parameters() 207 208 self.num_root_cells = self.parameters['num_root_cells'][0] 209 self.num_grid = 1 210 num_root = self.num_root_cells 211 while num_root > 1 : 212 self.num_grid <<= 1 213 num_root >>= 3 214 215 self.sfc_min = 0 216 self.sfc_max = self.num_root_cells-1 217 218 # initialize cosmology module 219 if "abox" in self.parameters: 220 self.cosmology = cosmology_allocate() 221 cosmology_set_OmegaM(self.cosmology, self.parameters['OmegaM'][0]) 222 cosmology_set_OmegaL(self.cosmology, self.parameters['OmegaL'][0]) 223 cosmology_set_OmegaB(self.cosmology, self.parameters['OmegaB'][0]) 224 cosmology_set_h(self.cosmology, self.parameters['hubble'][0]) 225 cosmology_set_DeltaDC(self.cosmology, self.parameters['DeltaDC'][0]) 226 cosmology_set_fixed(self.cosmology) 227 else: 228 self.cosmology = NULL 229 self.tcode_to_years = self.parameters['time_unit'][0]/(365.25*86400) 230 231 # grid detection 232 self.min_level = 0 233 self.max_level = self.parameters['grid_max_level'][0] 234 self.num_grid_variables = self.parameters['num_grid_variables'][0] 235 236 self.num_octs_per_level = <int *>malloc(self.max_level*sizeof(int)) 237 self.grid_variables = <float *>malloc(8*self.num_grid_variables*sizeof(float)) 238 if (not self.num_octs_per_level) or (not self.grid_variables) : 239 raise MemoryError 240 241 if artio_fileset_has_grid(self.handle): 242 status = artio_fileset_open_grid(self.handle) 243 check_artio_status(status) 244 self.has_grid = 1 245 else: 246 self.has_grid = 0 247 248 # particle detection 249 if ( artio_fileset_has_particles(self.handle) ): 250 status = artio_fileset_open_particles(self.handle) 251 check_artio_status(status) 252 self.has_particles = 1 253 254 for v in ["num_particle_species","num_primary_variables","num_secondary_variables"]: 255 if v not in self.parameters: 256 raise RuntimeError("Unable to locate particle header information in artio header: key=", v) 257 258 self.num_species = self.parameters['num_particle_species'][0] 259 self.particle_position_index = <int *>malloc(3*sizeof(int)*self.num_species) 260 if not self.particle_position_index : 261 raise MemoryError 262 for ispec in range(self.num_species) : 263 species_labels = "species_%02d_primary_variable_labels"% (ispec,) 264 if species_labels not in self.parameters: 265 raise RuntimeError("Unable to locate variable labels for species",ispec) 266 267 labels = self.parameters[species_labels] 268 try : 269 self.particle_position_index[3*ispec+0] = labels.index('POSITION_X') 270 self.particle_position_index[3*ispec+1] = labels.index('POSITION_Y') 271 self.particle_position_index[3*ispec+2] = labels.index('POSITION_Z') 272 except ValueError : 273 raise RuntimeError("Unable to locate position information for particle species", ispec) 274 275 self.num_particles_per_species = <int *>malloc(sizeof(int)*self.num_species) 276 self.primary_variables = <double *>malloc(sizeof(double)*max(self.parameters['num_primary_variables'])) 277 self.secondary_variables = <float *>malloc(sizeof(float)*max(self.parameters['num_secondary_variables'])) 278 if (not self.num_particles_per_species) or (not self.primary_variables) or (not self.secondary_variables) : 279 raise MemoryError 280 else: 281 self.has_particles = 0 282 283 def __dealloc__(self) : 284 if self.num_octs_per_level : free(self.num_octs_per_level) 285 if self.grid_variables : free(self.grid_variables) 286 287 if self.particle_position_index : free(self.particle_position_index) 288 if self.num_particles_per_species : free(self.num_particles_per_species) 289 if self.primary_variables : free(self.primary_variables) 290 if self.secondary_variables : free(self.secondary_variables) 291 292 if self.cosmology : cosmology_free(self.cosmology) 293 294 if self.handle : artio_fileset_close(self.handle) 295 296 def read_parameters(self) : 297 cdef char key[64] 298 cdef int type 299 cdef int length 300 cdef char ** char_values 301 cdef int32_t *int_values 302 cdef int64_t *long_values 303 cdef float *float_values 304 cdef double *double_values 305 306 self.parameters = {} 307 308 while artio_parameter_iterate( self.handle, key, &type, &length ) == ARTIO_SUCCESS : 309 310 if type == ARTIO_TYPE_STRING : 311 char_values = <char **>malloc(length*sizeof(char *)) 312 for i in range(length) : 313 char_values[i] = <char *>malloc( ARTIO_MAX_STRING_LENGTH*sizeof(char) ) 314 artio_parameter_get_string_array( self.handle, key, length, char_values ) 315 parameter = [ char_values[i] for i in range(length) ] 316 for i in range(length) : 317 free(char_values[i]) 318 free(char_values) 319 for i in range(len(parameter)): 320 parameter[i] = parameter[i].decode('utf-8') 321 elif type == ARTIO_TYPE_INT : 322 int_values = <int32_t *>malloc(length*sizeof(int32_t)) 323 artio_parameter_get_int_array( self.handle, key, length, int_values ) 324 parameter = [ int_values[i] for i in range(length) ] 325 free(int_values) 326 elif type == ARTIO_TYPE_LONG : 327 long_values = <int64_t *>malloc(length*sizeof(int64_t)) 328 artio_parameter_get_long_array( self.handle, key, length, long_values ) 329 parameter = [ long_values[i] for i in range(length) ] 330 free(long_values) 331 elif type == ARTIO_TYPE_FLOAT : 332 float_values = <float *>malloc(length*sizeof(float)) 333 artio_parameter_get_float_array( self.handle, key, length, float_values ) 334 parameter = [ float_values[i] for i in range(length) ] 335 free(float_values) 336 elif type == ARTIO_TYPE_DOUBLE : 337 double_values = <double *>malloc(length*sizeof(double)) 338 artio_parameter_get_double_array( self.handle, key, length, double_values ) 339 parameter = [ double_values[i] for i in range(length) ] 340 free(double_values) 341 else : 342 raise RuntimeError("ARTIO file corruption detected: invalid type!") 343 344 self.parameters[key.decode('utf-8')] = parameter 345 346 def abox_from_auni(self, np.float64_t a): 347 if self.cosmology: 348 return abox_from_auni(self.cosmology, a) 349 else: 350 raise RuntimeError("abox_from_auni called for non-cosmological ARTIO fileset!") 351 352 def tcode_from_auni(self, np.float64_t a): 353 if self.cosmology: 354 return tcode_from_auni(self.cosmology, a) 355 else: 356 raise RuntimeError("tcode_from_auni called for non-cosmological ARTIO fileset!") 357 358 def tphys_from_auni(self, np.float64_t a): 359 if self.cosmology: 360 return tphys_from_auni(self.cosmology, a) 361 else: 362 raise RuntimeError("tphys_from_auni called for non-cosmological ARTIO fileset!") 363 364 def auni_from_abox(self, np.float64_t v): 365 if self.cosmology: 366 return auni_from_abox(self.cosmology, v) 367 else: 368 raise RuntimeError("auni_from_abox called for non-cosmological ARTIO fileset!") 369 370 def auni_from_tcode(self, np.float64_t v): 371 if self.cosmology: 372 return auni_from_tcode(self.cosmology, v) 373 else: 374 raise RuntimeError("auni_from_tcode called for non-cosmological ARTIO fileset!") 375 376 @cython.wraparound(False) 377 @cython.boundscheck(False) 378 def auni_from_tcode_array(self, np.ndarray[np.float64_t] tcode): 379 cdef int i, nauni 380 cdef np.ndarray[np.float64_t, ndim=1] auni 381 cdef CosmologyParameters *cosmology = self.cosmology 382 if not cosmology: 383 raise RuntimeError("auni_from_tcode_array called for non-cosmological ARTIO fileset!") 384 auni = np.empty_like(tcode) 385 nauni = auni.shape[0] 386 with nogil: 387 for i in range(nauni): 388 auni[i] = auni_from_tcode(self.cosmology, tcode[i]) 389 return auni 390 391 def auni_from_tphys(self, np.float64_t v): 392 if self.cosmology: 393 return auni_from_tphys(self.cosmology, v) 394 else: 395 raise RuntimeError("auni_from_tphys called for non-cosmological ARTIO fileset!") 396 397 def abox_from_tcode(self, np.float64_t abox): 398 if self.cosmology: 399 return abox_from_tcode(self.cosmology, abox) 400 else: 401 raise RuntimeError("abox_from_tcode called for non-cosmological ARTIO fileset!") 402 403 def tcode_from_abox(self, np.float64_t abox): 404 if self.cosmology: 405 return tcode_from_abox(self.cosmology, abox) 406 else: 407 raise RuntimeError("tcode_from_abox called for non-cosmological ARTIO fileset!") 408 409 def tphys_from_abox(self, np.float64_t abox): 410 if self.cosmology: 411 return tphys_from_abox(self.cosmology, abox) 412 else: 413 raise RuntimeError("tphys_from_abox called for non-cosmological ARTIO fileset!") 414 415 def tphys_from_tcode(self, np.float64_t tcode): 416 if self.cosmology: 417 return tphys_from_tcode(self.cosmology, tcode) 418 else: 419 return self.tcode_to_years*tcode 420 421 @cython.wraparound(False) 422 @cython.boundscheck(False) 423 def tphys_from_tcode_array(self, np.ndarray[np.float64_t] tcode): 424 cdef int i, ntphys 425 cdef np.ndarray[np.float64_t, ndim=1] tphys 426 cdef CosmologyParameters *cosmology = self.cosmology 427 tphys = np.empty_like(tcode) 428 ntphys = tcode.shape[0] 429 430 if cosmology: 431 tphys = np.empty_like(tcode) 432 ntphys = tcode.shape[0] 433 with nogil: 434 for i in range(ntphys): 435 tphys[i] = tphys_from_tcode(cosmology, tcode[i]) 436 return tphys 437 else: 438 return tcode*self.tcode_to_years 439 440# @cython.boundscheck(False) 441 @cython.wraparound(False) 442 @cython.cdivision(True) 443 def read_particle_chunk(self, SelectorObject selector, int64_t sfc_start, int64_t sfc_end, fields) : 444 cdef int i 445 cdef int status 446 cdef int subspecies 447 cdef int64_t pid 448 449 cdef int num_fields = len(fields) 450 cdef np.float64_t pos[3] 451 452 # since RuntimeErrors are not fatal, ensure no artio_particles* functions 453 # called if fileset lacks particles 454 if not self.has_particles: return 455 456 data = {} 457 accessed_species = np.zeros( self.num_species, dtype="int64") 458 selected_mass = [ None for i in range(self.num_species)] 459 selected_pid = [ None for i in range(self.num_species)] 460 selected_species = [ None for i in range(self.num_species)] 461 selected_primary = [ [] for i in range(self.num_species)] 462 selected_secondary = [ [] for i in range(self.num_species)] 463 464 for species,field in fields : 465 if species < 0 or species > self.num_species : 466 raise RuntimeError("Invalid species provided to read_particle_chunk") 467 accessed_species[species] = 1 468 469 if self.parameters["num_primary_variables"][species] > 0 and \ 470 field in self.parameters["species_%02u_primary_variable_labels"%(species,)] : 471 selected_primary[species].append((self.parameters["species_%02u_primary_variable_labels"%(species,)].index(field),(species,field))) 472 data[(species,field)] = np.empty(0,dtype="float64") 473 elif self.parameters["num_secondary_variables"][species] > 0 and \ 474 field in self.parameters["species_%02u_secondary_variable_labels"%(species,)] : 475 selected_secondary[species].append((self.parameters["species_%02u_secondary_variable_labels"%(species,)].index(field),(species,field))) 476 data[(species,field)] = np.empty(0,dtype="float64") 477 elif field == "MASS" : 478 selected_mass[species] = (species,field) 479 data[(species,field)] = np.empty(0,dtype="float64") 480 elif field == "PID" : 481 selected_pid[species] = (species,field) 482 data[(species,field)] = np.empty(0,dtype="int64") 483 elif field == "SPECIES" : 484 selected_species[species] = (species,field) 485 data[(species,field)] = np.empty(0,dtype="int8") 486 else : 487 raise RuntimeError("invalid field name provided to read_particle_chunk") 488 489 # cache the range 490 status = artio_particle_cache_sfc_range( self.handle, self.sfc_min, self.sfc_max ) 491 check_artio_status(status) 492 493 for sfc in range( sfc_start, sfc_end+1 ) : 494 status = artio_particle_read_root_cell_begin( self.handle, sfc, 495 self.num_particles_per_species ) 496 check_artio_status(status) 497 498 for ispec in range(self.num_species) : 499 if accessed_species[ispec] : 500 status = artio_particle_read_species_begin(self.handle, ispec) 501 check_artio_status(status) 502 503 for particle in range( self.num_particles_per_species[ispec] ) : 504 status = artio_particle_read_particle(self.handle, 505 &pid, &subspecies, self.primary_variables, 506 self.secondary_variables) 507 check_artio_status(status) 508 509 for i in range(3) : 510 pos[i] = self.primary_variables[self.particle_position_index[3*ispec+i]] 511 512 if selector.select_point(pos) : 513 # loop over primary variables 514 for i,field in selected_primary[ispec] : 515 count = len(data[field]) 516 data[field].resize(count+1) 517 data[field][count] = self.primary_variables[i] 518 519 # loop over secondary variables 520 for i,field in selected_secondary[ispec] : 521 count = len(data[field]) 522 data[field].resize(count+1) 523 data[field][count] = self.secondary_variables[i] 524 525 # add particle id 526 if selected_pid[ispec] : 527 count = len(data[selected_pid[ispec]]) 528 data[selected_pid[ispec]].resize(count+1) 529 data[selected_pid[ispec]][count] = pid 530 531 # add mass if requested 532 if selected_mass[ispec] : 533 count = len(data[selected_mass[ispec]]) 534 data[selected_mass[ispec]].resize(count+1) 535 data[selected_mass[ispec]][count] = self.parameters["particle_species_mass"][ispec] 536 537 status = artio_particle_read_species_end( self.handle ) 538 check_artio_status(status) 539 540 status = artio_particle_read_root_cell_end( self.handle ) 541 check_artio_status(status) 542 543 return data 544 545 #@cython.boundscheck(False) 546 @cython.wraparound(False) 547 @cython.cdivision(True) 548 def read_grid_chunk(self, SelectorObject selector, int64_t sfc_start, int64_t sfc_end, fields ): 549 cdef int i 550 cdef int level 551 cdef int num_oct_levels 552 cdef int refined[8] 553 cdef int status 554 cdef int64_t count 555 cdef int64_t max_octs 556 cdef double dpos[3] 557 cdef np.float64_t left[3] 558 cdef np.float64_t right[3] 559 cdef np.float64_t dds[3] 560 561 cdef int *field_order 562 cdef int num_fields = len(fields) 563 field_order = <int*>malloc(sizeof(int)*num_fields) 564 565 # translate fields from ARTIO names to indices 566 var_labels = self.parameters['grid_variable_labels'] 567 for i, f in enumerate(fields): 568 if f not in var_labels: 569 raise RuntimeError("Field",f,"is not known to ARTIO") 570 field_order[i] = var_labels.index(f) 571 572 status = artio_grid_cache_sfc_range( self.handle, self.sfc_min, self.sfc_max ) 573 check_artio_status(status) 574 575 # determine max number of cells we could hit (optimize later) 576 #status = artio_grid_count_octs_in_sfc_range( self.handle, 577 # sfc_start, sfc_end, &max_octs ) 578 #check_artio_status(status) 579 #max_cells = sfc_end-sfc_start+1 + max_octs*8 580 581 # allocate space for _fcoords, _icoords, _fwidth, _ires 582 #fcoords = np.empty((max_cells, 3), dtype="float64") 583 #ires = np.empty(max_cells, dtype="int64") 584 fcoords = np.empty((0, 3), dtype="float64") 585 ires = np.empty(0, dtype="int64") 586 587 #data = [ np.empty(max_cells, dtype="float32") for i in range(num_fields) ] 588 data = [ np.empty(0,dtype="float64") for i in range(num_fields)] 589 590 count = 0 591 for sfc in range( sfc_start, sfc_end+1 ) : 592 status = artio_grid_read_root_cell_begin( self.handle, sfc, 593 dpos, self.grid_variables, &num_oct_levels, self.num_octs_per_level ) 594 check_artio_status(status) 595 596 if num_oct_levels == 0 : 597 for i in range(num_fields) : 598 data[i].resize(count+1) 599 data[i][count] = self.grid_variables[field_order[i]] 600 fcoords.resize((count+1,3)) 601 for i in range(3) : 602 fcoords[count][i] = dpos[i] 603 ires.resize(count+1) 604 ires[count] = 0 605 count += 1 606 607 for level in range(1,num_oct_levels+1) : 608 status = artio_grid_read_level_begin( self.handle, level ) 609 check_artio_status(status) 610 611 for i in range(3) : 612 dds[i] = 2.**-level 613 614 for oct in range(self.num_octs_per_level[level-1]) : 615 status = artio_grid_read_oct( self.handle, dpos, self.grid_variables, refined ) 616 check_artio_status(status) 617 618 for child in range(8) : 619 if not refined[child] : 620 for i in range(3) : 621 left[i] = (dpos[i]-dds[i]) if (child & (i<<1)) else dpos[i] 622 right[i] = left[i] + dds[i] 623 624 if selector.select_bbox(left,right) : 625 fcoords.resize((count+1, 3)) 626 for i in range(3) : 627 fcoords[count][i] = left[i]+0.5*dds[i] 628 ires.resize(count+1) 629 ires[count] = level 630 for i in range(num_fields) : 631 data[i].resize(count+1) 632 data[i][count] = self.grid_variables[self.num_grid_variables*child+field_order[i]] 633 count += 1 634 status = artio_grid_read_level_end( self.handle ) 635 check_artio_status(status) 636 637 status = artio_grid_read_root_cell_end( self.handle ) 638 check_artio_status(status) 639 640 free(field_order) 641 642 #fcoords.resize((count,3)) 643 #ires.resize(count) 644 # 645 #for i in range(num_fields) : 646 # data[i].resize(count) 647 648 return (fcoords, ires, data) 649 650 def root_sfc_ranges_all(self, int max_range_size = 1024) : 651 cdef int64_t sfc_start, sfc_end 652 cdef artio_selection *selection 653 654 selection = artio_select_all( self.handle ) 655 if selection == NULL : 656 raise RuntimeError 657 sfc_ranges = [] 658 while artio_selection_iterator(selection, max_range_size, 659 &sfc_start, &sfc_end) == ARTIO_SUCCESS : 660 sfc_ranges.append([sfc_start, sfc_end]) 661 artio_selection_destroy(selection) 662 return sfc_ranges 663 664 def root_sfc_ranges(self, SelectorObject selector, 665 int max_range_size = 1024): 666 cdef int coords[3] 667 cdef int64_t sfc_start, sfc_end 668 cdef np.float64_t left[3] 669 cdef np.float64_t right[3] 670 cdef np.float64_t dds[3] 671 cdef artio_selection *selection 672 cdef int i, j, k 673 674 sfc_ranges=[] 675 selection = artio_selection_allocate(self.handle) 676 for i in range(self.num_grid) : 677 # stupid cython 678 coords[0] = i 679 left[0] = coords[0] 680 right[0] = left[0] + 1.0 681 for j in range(self.num_grid) : 682 coords[1] = j 683 left[1] = coords[1] 684 right[1] = left[1] + 1.0 685 for k in range(self.num_grid) : 686 coords[2] = k 687 left[2] = coords[2] 688 right[2] = left[2] + 1.0 689 if selector.select_bbox(left,right) : 690 status = artio_selection_add_root_cell(selection, coords) 691 check_artio_status(status) 692 693 while artio_selection_iterator(selection, max_range_size, 694 &sfc_start, &sfc_end) == ARTIO_SUCCESS : 695 sfc_ranges.append([sfc_start, sfc_end]) 696 697 artio_selection_destroy(selection) 698 return sfc_ranges 699 700################################################### 701def artio_is_valid( char *file_prefix ) : 702 cdef artio_fileset_handle *handle = artio_fileset_open( file_prefix, 703 ARTIO_OPEN_HEADER, artio_context_global ) 704 if handle == NULL : 705 return False 706 else : 707 artio_fileset_close(handle) 708 return True 709 710cdef class ARTIOSFCRangeHandler: 711 cdef public np.int64_t sfc_start 712 cdef public np.int64_t sfc_end 713 cdef public artio_fileset artio_handle 714 cdef public object root_mesh_handler 715 cdef public object oct_count 716 cdef public object octree_handler 717 cdef artio_fileset_handle *handle 718 cdef np.float64_t DLE[3] 719 cdef np.float64_t DRE[3] 720 cdef np.float64_t dds[3] 721 cdef np.int64_t dims[3] 722 cdef public np.int64_t total_octs 723 cdef np.int64_t *doct_count 724 cdef np.int64_t **pcount 725 cdef float **root_mesh_data 726 cdef np.int64_t nvars[2] 727 cdef int cache_root_mesh 728 729 def __init__(self, domain_dimensions, # cells 730 domain_left_edge, 731 domain_right_edge, 732 artio_fileset artio_handle, 733 sfc_start, sfc_end, int cache_root_mesh = 0): 734 cdef int i 735 cdef np.int64_t sfc 736 self.sfc_start = sfc_start 737 self.sfc_end = sfc_end 738 self.artio_handle = artio_handle 739 self.root_mesh_handler = None 740 self.octree_handler = None 741 self.handle = artio_handle.handle 742 self.oct_count = None 743 self.root_mesh_data = NULL 744 self.pcount = NULL 745 self.cache_root_mesh = cache_root_mesh 746 747 if artio_handle.has_particles: 748 self.pcount = <np.int64_t **> malloc(sizeof(np.int64_t*) 749 * artio_handle.num_species) 750 self.nvars[0] = artio_handle.num_species 751 for i in range(artio_handle.num_species): 752 self.pcount[i] = <np.int64_t*> malloc(sizeof(np.int64_t) 753 * (self.sfc_end - self.sfc_start + 1)) 754 for sfc in range(self.sfc_end - self.sfc_start + 1): 755 self.pcount[i][sfc] = 0 756 else: 757 self.nvars[0] = 0 758 759 if artio_handle.has_grid: 760 self.nvars[1] = artio_handle.num_grid_variables 761 else: 762 self.nvars[1] = 0 763 764 for i in range(3): 765 self.dims[i] = domain_dimensions[i] 766 self.DLE[i] = domain_left_edge[i] 767 self.DRE[i] = domain_right_edge[i] 768 self.dds[i] = (self.DRE[i] - self.DLE[i])/self.dims[i] 769 770 def __dealloc__(self): 771 cdef int i 772 if self.artio_handle.has_particles: 773 for i in range(self.nvars[0]): 774 free(self.pcount[i]) 775 free(self.pcount) 776 if self.artio_handle.has_grid: 777 if self.root_mesh_data != NULL: 778 for i in range(self.nvars[1]): 779 free(self.root_mesh_data[i]) 780 free(self.root_mesh_data) 781 782 @cython.boundscheck(False) 783 @cython.wraparound(False) 784 @cython.cdivision(True) 785 def construct_mesh(self): 786 cdef int status, level, ngv 787 cdef np.int64_t sfc, oc, i 788 cdef double dpos[3] 789 cdef int num_oct_levels 790 cdef int max_level = self.artio_handle.max_level 791 cdef int *num_octs_per_level = <int *>malloc( 792 (max_level + 1)*sizeof(int)) 793 cdef int num_species = self.artio_handle.num_species 794 cdef int *num_particles_per_species 795 cdef ARTIOOctreeContainer octree 796 ngv = self.nvars[1] 797 cdef float *grid_variables = <float *>malloc( 798 ngv * sizeof(float)) 799 self.octree_handler = octree = ARTIOOctreeContainer(self) 800 if self.cache_root_mesh == 1: 801 self.root_mesh_data = <float **>malloc(sizeof(float *) * ngv) 802 for i in range(ngv): 803 self.root_mesh_data[i] = <float *>malloc(sizeof(float) * \ 804 (self.sfc_end - self.sfc_start + 1)) 805 # We want to pre-allocate an array of root pointers. In the future, 806 # this will be pre-determined by the ARTIO library. However, because 807 # realloc plays havoc with our tree searching, we can't utilize an 808 # expanding array at the present time. 809 octree.allocate_domains([], self.sfc_end - self.sfc_start + 1) 810 cdef np.ndarray[np.int64_t, ndim=1] oct_count 811 oct_count = np.zeros(self.sfc_end - self.sfc_start + 1, dtype="int64") 812 status = artio_grid_cache_sfc_range(self.handle, self.sfc_start, 813 self.sfc_end) 814 check_artio_status(status) 815 for sfc in range(self.sfc_start, self.sfc_end + 1): 816 status = artio_grid_read_root_cell_begin( self.handle, 817 sfc, dpos, grid_variables, &num_oct_levels, 818 num_octs_per_level) 819 check_artio_status(status) 820 for i in range(ngv * self.cache_root_mesh): 821 self.root_mesh_data[i][sfc - self.sfc_start] = \ 822 grid_variables[i] 823 if num_oct_levels > 0: 824 oc = 0 825 for level in range(num_oct_levels): 826 oc += num_octs_per_level[level] 827 self.total_octs += oc 828 oct_count[sfc - self.sfc_start] = oc 829 octree.initialize_local_mesh(oc, num_oct_levels, 830 num_octs_per_level, sfc) 831 status = artio_grid_read_root_cell_end(self.handle) 832 check_artio_status(status) 833 status = artio_grid_clear_sfc_cache(self.handle) 834 check_artio_status(status) 835 if self.artio_handle.has_particles: 836 num_particles_per_species = <int *>malloc( 837 sizeof(int)*num_species) 838 839 # Now particles 840 status = artio_particle_cache_sfc_range(self.handle, self.sfc_start, 841 self.sfc_end) 842 check_artio_status(status) 843 for sfc in range(self.sfc_start, self.sfc_end + 1): 844 # Now particles 845 status = artio_particle_read_root_cell_begin(self.handle, 846 sfc, num_particles_per_species) 847 check_artio_status(status) 848 849 for i in range(num_species): 850 self.pcount[i][sfc - self.sfc_start] = \ 851 num_particles_per_species[i] 852 853 status = artio_particle_read_root_cell_end(self.handle) 854 check_artio_status(status) 855 856 status = artio_particle_clear_sfc_cache(self.handle) 857 check_artio_status(status) 858 859 free(num_particles_per_species) 860 861 free(grid_variables) 862 free(num_octs_per_level) 863 self.oct_count = oct_count 864 self.doct_count = <np.int64_t *> oct_count.data 865 self.root_mesh_handler = ARTIORootMeshContainer(self) 866 867 def free_mesh(self): 868 self.octree_handler = None 869 self.root_mesh_handler = None 870 self.doct_count = NULL 871 self.oct_count = None 872 873def get_coords(artio_fileset handle, np.int64_t s): 874 cdef int coords[3] 875 artio_sfc_coords(handle.handle, s, coords) 876 return (coords[0], coords[1], coords[2]) 877 878cdef struct particle_var_pointers: 879 # The number of particles we have filled 880 np.int64_t count 881 # Number of primary variables and pointers to their indices 882 int n_p 883 int p_ind[16] # Max of 16 vars 884 # Number of secondary variables and pointers to their indices 885 int n_s 886 int s_ind[16] # Max of 16 vars 887 # Pointers to the bools and data arrays for mass, pid and species 888 int n_mass 889 np.float64_t *mass 890 int n_pid 891 np.int64_t *pid 892 int n_species 893 np.int8_t *species 894 # Pointers to the pointers to primary and secondary vars 895 np.float64_t *pvars[16] 896 np.float64_t *svars[16] 897 898cdef class ARTIOOctreeContainer(SparseOctreeContainer): 899 # This is a transitory, created-on-demand OctreeContainer. It should not 900 # be considered to be long-lasting, and during its creation it will read 901 # the index file. This means that when created it will then be able to 902 # provide coordinates, but then on subsequent IO accesses it will pass over 903 # the file again, despite knowing the indexing system already. Because of 904 # this, we will avoid creating it as long as possible. 905 906 cdef public artio_fileset artio_handle 907 cdef ARTIOSFCRangeHandler range_handler 908 cdef np.int64_t level_indices[32] 909 cdef np.int64_t sfc_start 910 cdef np.int64_t sfc_end 911 912 def __init__(self, ARTIOSFCRangeHandler range_handler): 913 self.range_handler = range_handler 914 self.sfc_start = range_handler.sfc_start 915 self.sfc_end = range_handler.sfc_end 916 self.artio_handle = range_handler.artio_handle 917 # Note the final argument is partial_coverage, which indicates whether 918 # or not an Oct can be partially refined. 919 dims, DLE, DRE = [], [], [] 920 for i in range(32): 921 self.level_indices[i] = 0 922 for i in range(3): 923 # range_handler has dims in cells, which is the same as the number 924 # of possible octs. This is because we have a forest of octrees. 925 dims.append(range_handler.dims[i]) 926 DLE.append(range_handler.DLE[i]) 927 DRE.append(range_handler.DRE[i]) 928 super(ARTIOOctreeContainer, self).__init__(dims, DLE, DRE) 929 self.artio_handle = range_handler.artio_handle 930 self.level_offset = 1 931 self.domains = OctObjectPool() 932 self.root_nodes = NULL 933 934 @cython.boundscheck(False) 935 @cython.wraparound(False) 936 @cython.cdivision(True) 937 cdef void initialize_local_mesh(self, np.int64_t oct_count, 938 int num_oct_levels, int *num_octs_per_level, 939 np.int64_t sfc): 940 # We actually will not be initializing the root mesh here, we will be 941 # initializing the entire mesh between sfc_start and sfc_end. 942 cdef np.int64_t oct_ind, tot_octs, ipos, nadded 943 cdef int i, status, level, num_root, num_octs 944 cdef int num_level_octs 945 cdef artio_fileset_handle *handle = self.artio_handle.handle 946 cdef int coords[3] 947 cdef int max_level = self.artio_handle.max_level 948 cdef double dpos[3] 949 cdef np.float64_t f64pos[3] 950 cdef np.float64_t dds[3] 951 cdef np.ndarray[np.float64_t, ndim=2] pos 952 # NOTE: We do not cache any SFC ranges here, as we should only ever be 953 # called from within a pre-cached operation in the SFC handler. 954 955 # We only allow one root oct. 956 self.append_domain(oct_count) 957 self.domains.containers[self.num_domains - 1].con_id = sfc 958 959 oct_ind = -1 960 ipos = 0 961 for level in range(num_oct_levels): 962 oct_ind = imax(oct_ind, num_octs_per_level[level]) 963 self.level_indices[level] = ipos 964 ipos += num_octs_per_level[level] 965 pos = np.empty((oct_ind, 3), dtype="float64") 966 967 # Now we initialize 968 # Note that we also assume we have already started reading the level. 969 ipos = 0 970 for level in range(num_oct_levels): 971 status = artio_grid_read_level_begin(handle, level + 1) 972 check_artio_status(status) 973 for oct_ind in range(num_octs_per_level[level]): 974 status = artio_grid_read_oct(handle, dpos, NULL, NULL) 975 for i in range(3): 976 pos[oct_ind, i] = dpos[i] 977 check_artio_status(status) 978 status = artio_grid_read_level_end(handle) 979 check_artio_status(status) 980 nadded = self.add(self.num_domains, level, pos[:num_octs_per_level[level],:]) 981 if nadded != num_octs_per_level[level]: 982 raise RuntimeError 983 984 @cython.boundscheck(False) 985 @cython.wraparound(False) 986 @cython.cdivision(True) 987 def fill_sfc(self, 988 np.ndarray[np.uint8_t, ndim=1] levels, 989 np.ndarray[np.uint8_t, ndim=1] cell_inds, 990 np.ndarray[np.int64_t, ndim=1] file_inds, 991 np.ndarray[np.int64_t, ndim=1] domain_counts, 992 field_indices, dest_fields): 993 cdef np.ndarray[np.float32_t, ndim=2] source 994 cdef np.ndarray[np.float64_t, ndim=1] dest 995 cdef int n, status, i, di, num_oct_levels, nf, ngv, max_level 996 cdef int level, j, oct_ind, si 997 cdef np.int64_t sfc, ipos 998 cdef np.float64_t val 999 cdef artio_fileset_handle *handle = self.artio_handle.handle 1000 cdef double dpos[3] 1001 # We duplicate some of the grid_variables stuff here so that we can 1002 # potentially release the GIL 1003 nf = len(field_indices) 1004 ngv = self.artio_handle.num_grid_variables 1005 max_level = self.artio_handle.max_level 1006 cdef int *num_octs_per_level = <int *>malloc( 1007 (max_level + 1)*sizeof(int)) 1008 cdef float *grid_variables = <float *>malloc( 1009 8 * ngv * sizeof(float)) 1010 cdef int* field_ind = <int*> malloc( 1011 nf * sizeof(int)) 1012 cdef np.float32_t **field_vals = <np.float32_t**> malloc( 1013 nf * sizeof(np.float32_t*)) 1014 source_arrays = [] 1015 ipos = -1 1016 for i in range(self.num_domains): 1017 ipos = imax(ipos, self.domains.containers[i].n) 1018 for i in range(nf): 1019 field_ind[i] = field_indices[i] 1020 # Note that we subtract one, because we're not using the root mesh. 1021 source = np.zeros((ipos, 8), dtype="float32") 1022 source_arrays.append(source) 1023 field_vals[i] = <np.float32_t*> source.data 1024 cdef np.int64_t level_position[32] 1025 cdef np.int64_t lp 1026 # First we need to walk the mesh in the file. Then we fill in the dest 1027 # location based on the file index. 1028 # A few ways this could be improved: 1029 # * Create a new visitor function that actually queried the data, 1030 # rather than our somewhat hokey double-loop over SFC arrays. 1031 # * Go to pointers for the dest arrays. 1032 # * Enable preloading during mesh initialization 1033 # * Calculate domain indices on the fly rather than with a 1034 # double-loop to calculate domain_counts 1035 # The cons should be in order 1036 cdef np.int64_t sfc_start, sfc_end 1037 sfc_start = self.domains.containers[0].con_id 1038 sfc_end = self.domains.containers[self.num_domains - 1].con_id 1039 status = artio_grid_cache_sfc_range(handle, sfc_start, sfc_end) 1040 check_artio_status(status) 1041 cdef np.int64_t offset = 0 1042 for si in range(self.num_domains): 1043 sfc = self.domains.containers[si].con_id 1044 status = artio_grid_read_root_cell_begin( handle, sfc, 1045 dpos, NULL, &num_oct_levels, num_octs_per_level) 1046 check_artio_status(status) 1047 lp = 0 1048 for level in range(num_oct_levels): 1049 status = artio_grid_read_level_begin(handle, level + 1) 1050 check_artio_status(status) 1051 level_position[level] = lp 1052 for oct_ind in range(num_octs_per_level[level]): 1053 status = artio_grid_read_oct(handle, dpos, grid_variables, NULL) 1054 check_artio_status(status) 1055 for j in range(8): 1056 for i in range(nf): 1057 field_vals[i][(oct_ind + lp)*8+j] = \ 1058 grid_variables[ngv*j+field_ind[i]] 1059 status = artio_grid_read_level_end(handle) 1060 check_artio_status(status) 1061 lp += num_octs_per_level[level] 1062 status = artio_grid_read_root_cell_end(handle) 1063 check_artio_status(status) 1064 # Now we have all our sources. 1065 for j in range(nf): 1066 dest = dest_fields[j] 1067 source = source_arrays[j] 1068 for i in range(domain_counts[si]): 1069 level = levels[i + offset] 1070 oct_ind = file_inds[i + offset] + level_position[level] 1071 dest[i + offset] = source[oct_ind, cell_inds[i + offset]] 1072 # Now, we offset by the actual number filled here. 1073 offset += domain_counts[si] 1074 status = artio_grid_clear_sfc_cache(handle) 1075 check_artio_status(status) 1076 free(field_ind) 1077 free(field_vals) 1078 free(grid_variables) 1079 free(num_octs_per_level) 1080 1081 def fill_sfc_particles(self, fields): 1082 # This handles not getting particles for refined sfc values. 1083 rv = read_sfc_particles(self.artio_handle, self.sfc_start, 1084 self.sfc_end, 0, fields, 1085 self.range_handler.doct_count, 1086 self.range_handler.pcount) 1087 return rv 1088 1089@cython.boundscheck(False) 1090@cython.wraparound(False) 1091@cython.cdivision(True) 1092cdef read_sfc_particles(artio_fileset artio_handle, 1093 np.int64_t sfc_start, np.int64_t sfc_end, 1094 int read_unrefined, fields, 1095 np.int64_t *doct_count, 1096 np.int64_t **pcount): 1097 cdef int status, ispec, subspecies 1098 cdef np.int64_t sfc, particle, pid, ind, vind 1099 cdef int num_species = artio_handle.num_species 1100 cdef artio_fileset_handle *handle = artio_handle.handle 1101 cdef int num_oct_levels 1102 cdef int *num_particles_per_species = <int *>malloc( 1103 sizeof(int)*num_species) 1104 cdef int *accessed_species = <int *>malloc( 1105 sizeof(int)*num_species) 1106 cdef int *total_particles = <int *>malloc( 1107 sizeof(int)*num_species) 1108 cdef particle_var_pointers *vpoints = <particle_var_pointers *> malloc( 1109 sizeof(particle_var_pointers)*num_species) 1110 cdef double *primary_variables 1111 cdef double dpos[3] 1112 cdef float *secondary_variables 1113 cdef np.int64_t tp 1114 cdef int max_level = artio_handle.max_level 1115 cdef int *num_octs_per_level = <int *>malloc( 1116 (max_level + 1)*sizeof(int)) 1117 1118 cdef np.ndarray[np.int8_t, ndim=1] npi8arr 1119 cdef np.ndarray[np.int64_t, ndim=1] npi64arr 1120 cdef np.ndarray[np.float64_t, ndim=1] npf64arr 1121 1122 if not artio_handle.has_particles: 1123 raise RuntimeError("Attempted to read non-existent particles in ARTIO") 1124 1125 # Now we set up our field pointers 1126 params = artio_handle.parameters 1127 1128 npri_vars = params["num_primary_variables"] 1129 nsec_vars = params["num_secondary_variables"] 1130 primary_variables = <double *>malloc(sizeof(double) * max(npri_vars)) 1131 secondary_variables = <float *>malloc(sizeof(float) * max(nsec_vars)) 1132 1133 cdef particle_var_pointers *vp 1134 1135 for ispec in range(num_species): 1136 total_particles[ispec] = 0 1137 accessed_species[ispec] = 0 1138 # Initialize our vpoints array 1139 vpoints[ispec].count = 0 1140 vpoints[ispec].n_mass = 0 1141 vpoints[ispec].n_pid = 0 1142 vpoints[ispec].n_species = 0 1143 vpoints[ispec].n_p = 0 1144 vpoints[ispec].n_s = 0 1145 1146 # Pass through once. We want every single particle. 1147 tp = 0 1148 cdef np.int64_t c 1149 for sfc in range(sfc_start, sfc_end + 1): 1150 c = doct_count[sfc - sfc_start] 1151 if read_unrefined == 1 and c > 0: continue 1152 if read_unrefined == 0 and c == 0: continue 1153 1154 for ispec in range(num_species): 1155 total_particles[ispec] += pcount[ispec][sfc - sfc_start] 1156 1157 # Now we allocate our final fields, which will be filled 1158 #for ispec in range(num_species): 1159 # print "In SFC %s to %s reading %s of species %s" % ( 1160 # sfc_start, sfc_end + 1, total_particles[ispec], ispec) 1161 data = {} 1162 for species, field in sorted(fields): 1163 accessed_species[species] = 1 1164 pri_vars = params.get( 1165 "species_%02u_primary_variable_labels" % (species,), []) 1166 sec_vars = params.get( 1167 "species_%02u_secondary_variable_labels" % (species,), []) 1168 tp = total_particles[species] 1169 vp = &vpoints[species] 1170 if field == "PID": 1171 vp.n_pid = 1 1172 data[(species, field)] = np.zeros(tp, dtype="int64") 1173 npi64arr = data[(species, field)] 1174 vp.pid = <np.int64_t*> npi64arr.data 1175 elif field == "SPECIES": 1176 vp.n_species = 1 1177 data[(species, field)] = np.zeros(tp, dtype="int8") 1178 npi8arr = data[(species, field)] 1179 # We fill this *now* 1180 npi8arr += species 1181 vp.species = <np.int8_t*> npi8arr.data 1182 elif npri_vars[species] > 0 and field in pri_vars : 1183 data[(species, field)] = np.zeros(tp, dtype="float64") 1184 npf64arr = data[(species, field)] 1185 vp.p_ind[vp.n_p] = pri_vars.index(field) 1186 vp.pvars[vp.n_p] = <np.float64_t *> npf64arr.data 1187 vp.n_p += 1 1188 elif nsec_vars[species] > 0 and field in sec_vars : 1189 data[(species, field)] = np.zeros(tp, dtype="float64") 1190 npf64arr = data[(species, field)] 1191 vp.s_ind[vp.n_s] = sec_vars.index(field) 1192 vp.svars[vp.n_s] = <np.float64_t *> npf64arr.data 1193 vp.n_s += 1 1194 elif field == "MASS": 1195 vp.n_mass = 1 1196 data[(species, field)] = np.zeros(tp, dtype="float64") 1197 npf64arr = data[(species, field)] 1198 # We fill this *now* 1199 npf64arr += params["particle_species_mass"][species] 1200 vp.mass = <np.float64_t*> npf64arr.data 1201 1202 status = artio_particle_cache_sfc_range( handle, 1203 sfc_start, sfc_end ) 1204 check_artio_status(status) 1205 1206 for sfc in range(sfc_start, sfc_end + 1): 1207 c = doct_count[sfc - sfc_start] 1208 check_artio_status(status) 1209 if read_unrefined == 1 and c > 0: continue 1210 if read_unrefined == 0 and c == 0: continue 1211 c = 0 1212 for ispec in range(num_species) : 1213 if accessed_species[ispec] == 0: continue 1214 c += pcount[ispec][sfc - sfc_start] 1215 if c == 0: continue 1216 status = artio_particle_read_root_cell_begin( handle, sfc, 1217 num_particles_per_species ) 1218 check_artio_status(status) 1219 for ispec in range(num_species) : 1220 if accessed_species[ispec] == 0: continue 1221 status = artio_particle_read_species_begin(handle, ispec); 1222 check_artio_status(status) 1223 vp = &vpoints[ispec] 1224 1225 for particle in range(num_particles_per_species[ispec]) : 1226 status = artio_particle_read_particle(handle, 1227 &pid, &subspecies, primary_variables, 1228 secondary_variables) 1229 check_artio_status(status) 1230 ind = vp.count 1231 1232 for i in range(vp.n_p): 1233 vind = vp.p_ind[i] 1234 vp.pvars[i][ind] = primary_variables[vind] 1235 1236 for i in range(vp.n_s): 1237 vind = vp.s_ind[i] 1238 vp.svars[i][ind] = secondary_variables[vind] 1239 1240 if vp.n_pid: 1241 vp.pid[ind] = pid 1242 1243 vp.count += 1 1244 1245 status = artio_particle_read_species_end( handle ) 1246 check_artio_status(status) 1247 1248 status = artio_particle_read_root_cell_end( handle ) 1249 check_artio_status(status) 1250 1251 status = artio_particle_clear_sfc_cache(handle) 1252 check_artio_status(status) 1253 1254 free(num_octs_per_level) 1255 free(num_particles_per_species) 1256 free(total_particles) 1257 free(accessed_species) 1258 free(vpoints) 1259 free(primary_variables) 1260 free(secondary_variables) 1261 return data 1262 1263cdef class ARTIORootMeshContainer: 1264 cdef public artio_fileset artio_handle 1265 cdef np.float64_t DLE[3] 1266 cdef np.float64_t DRE[3] 1267 cdef np.float64_t dds[3] 1268 cdef np.int64_t dims[3] 1269 cdef artio_fileset_handle *handle 1270 cdef np.uint64_t sfc_start 1271 cdef np.uint64_t sfc_end 1272 cdef public object _last_mask 1273 cdef public np.int64_t _last_selector_id 1274 cdef np.int64_t _last_mask_sum 1275 cdef ARTIOSFCRangeHandler range_handler 1276 cdef np.uint8_t *sfc_mask 1277 cdef np.int64_t nsfc 1278 1279 def __init__(self, ARTIOSFCRangeHandler range_handler): 1280 cdef int i 1281 cdef np.int64_t sfci 1282 for i in range(3): 1283 self.DLE[i] = range_handler.DLE[i] 1284 self.DRE[i] = range_handler.DRE[i] 1285 self.dims[i] = range_handler.dims[i] 1286 self.dds[i] = range_handler.dds[i] 1287 self.handle = range_handler.handle 1288 self.artio_handle = range_handler.artio_handle 1289 self._last_mask = None 1290 self._last_selector_id = -1 1291 self.sfc_start = range_handler.sfc_start 1292 self.sfc_end = range_handler.sfc_end 1293 self.range_handler = range_handler 1294 # We assume that the number of octs has been created and filled 1295 # already. We no longer care about ANY of the SFCs that have octs 1296 # inside them -- this goes for every operation that this object 1297 # performs. 1298 self.sfc_mask = <np.uint8_t *>malloc(sizeof(np.uint8_t) * 1299 self.sfc_end - self.sfc_start + 1) 1300 self.nsfc = 0 1301 for sfci in range(self.sfc_end - self.sfc_start + 1): 1302 if self.range_handler.oct_count[sfci] > 0: 1303 self.sfc_mask[sfci] = 0 1304 else: 1305 self.sfc_mask[sfci] = 1 1306 self.nsfc += 1 1307 1308 def __dealloc__(self): 1309 free(self.sfc_mask) 1310 1311 @cython.cdivision(True) 1312 cdef np.int64_t pos_to_sfc(self, np.float64_t pos[3]) nogil: 1313 # Calculate the index 1314 cdef int coords[3] 1315 cdef int i 1316 cdef np.int64_t sfc 1317 for i in range(3): 1318 coords[i] = <int>((pos[i] - self.DLE[i])/self.dds[i]) 1319 sfc = artio_sfc_index(self.handle, coords) 1320 return sfc 1321 1322 @cython.cdivision(True) 1323 cdef void sfc_to_pos(self, np.int64_t sfc, np.float64_t pos[3]) nogil: 1324 cdef int coords[3] 1325 cdef int i 1326 artio_sfc_coords(self.handle, sfc, coords) 1327 for i in range(3): 1328 pos[i] = self.DLE[i] + (coords[i] + 0.5) * self.dds[i] 1329 1330 cdef np.int64_t count_cells(self, SelectorObject selector): 1331 # We visit each cell if it is not refined and determine whether it is 1332 # included or not. 1333 cdef np.int64_t sfc 1334 cdef np.float64_t pos[3] 1335 cdef np.float64_t right_edge[3] 1336 cdef int num_cells = 0 1337 cdef int i 1338 return self.mask(selector).sum() 1339 1340 @cython.boundscheck(False) 1341 @cython.wraparound(False) 1342 @cython.cdivision(True) 1343 def icoords(self, SelectorObject selector, np.int64_t num_cells = -1, 1344 int domain_id = -1): 1345 # Note that num_octs does not have to equal sfc_end - sfc_start + 1. 1346 cdef np.int64_t sfc, sfci = -1 1347 cdef int acoords[3] 1348 cdef int i 1349 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1350 mask = self.mask(selector) 1351 num_cells = self._last_mask_sum 1352 cdef np.ndarray[np.int64_t, ndim=2] coords 1353 coords = np.empty((num_cells, 3), dtype="int64") 1354 cdef int filled = 0 1355 for sfc in range(self.sfc_start, self.sfc_end + 1): 1356 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1357 sfci += 1 1358 if mask[sfci] == 0: continue 1359 # Note that we do *no* checks on refinement here. In fact, this 1360 # entire setup should not need to touch the disk except if the 1361 # artio sfc calculators need to. 1362 artio_sfc_coords(self.handle, sfc, acoords) 1363 for i in range(3): 1364 coords[filled, i] = acoords[i] 1365 filled += 1 1366 return coords 1367 1368 @cython.boundscheck(False) 1369 @cython.wraparound(False) 1370 @cython.cdivision(True) 1371 def fcoords(self, SelectorObject selector, np.int64_t num_cells = -1, 1372 int domain_id = -1): 1373 # Note that num_cells does not have to equal sfc_end - sfc_start + 1. 1374 cdef np.int64_t sfc, sfci = -1 1375 cdef np.float64_t pos[3] 1376 cdef int acoords[3] 1377 cdef int i 1378 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1379 mask = self.mask(selector) 1380 num_cells = self._last_mask_sum 1381 cdef np.ndarray[np.float64_t, ndim=2] coords 1382 coords = np.empty((num_cells, 3), dtype="float64") 1383 cdef int filled = 0 1384 for sfc in range(self.sfc_start, self.sfc_end + 1): 1385 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1386 sfci += 1 1387 if mask[sfci] == 0: continue 1388 # Note that we do *no* checks on refinement here. In fact, this 1389 # entire setup should not need to touch the disk except if the 1390 # artio sfc calculators need to. 1391 self.sfc_to_pos(sfc, pos) 1392 for i in range(3): 1393 coords[filled, i] = pos[i] 1394 filled += 1 1395 return coords 1396 1397 @cython.boundscheck(False) 1398 @cython.wraparound(False) 1399 @cython.cdivision(True) 1400 def fwidth(self, SelectorObject selector, np.int64_t num_cells = -1, 1401 int domain_id = -1): 1402 cdef int i 1403 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1404 mask = self.mask(selector) 1405 num_cells = self._last_mask_sum 1406 cdef np.ndarray[np.float64_t, ndim=2] width 1407 width = np.zeros((num_cells, 3), dtype="float64") 1408 for i in range(3): 1409 width[:,i] = self.dds[i] 1410 return width 1411 1412 @cython.boundscheck(False) 1413 @cython.wraparound(False) 1414 @cython.cdivision(True) 1415 def ires(self, SelectorObject selector, np.int64_t num_cells = -1, 1416 int domain_id = -1): 1417 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1418 mask = self.mask(selector) 1419 num_cells = self._last_mask_sum 1420 cdef np.ndarray[np.int64_t, ndim=1] res 1421 res = np.zeros(num_cells, dtype="int64") 1422 return res 1423 1424 @cython.boundscheck(False) 1425 @cython.wraparound(False) 1426 @cython.cdivision(True) 1427 def selector_fill(self, SelectorObject selector, 1428 np.ndarray source, 1429 np.ndarray dest = None, 1430 np.int64_t offset = 0, int dims = 1, 1431 int domain_id = -1): 1432 # This is where we use the selector to transplant from one to the 1433 # other. Note that we *do* apply the selector here. 1434 cdef np.int64_t num_cells = -1 1435 cdef np.int64_t ind 1436 cdef np.int64_t sfc, sfci = -1 1437 cdef np.float64_t pos[3] 1438 cdef np.float64_t dpos[3] 1439 cdef int dim, status, filled = 0 1440 cdef int num_oct_levels, level 1441 cdef int max_level = self.artio_handle.max_level 1442 cdef int *num_octs_per_level = <int *>malloc( 1443 (max_level + 1)*sizeof(int)) 1444 cdef char *sdata = <char*> source.data 1445 cdef char *ddata 1446 cdef int ss = source.dtype.itemsize 1447 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1448 mask = self.mask(selector) 1449 if dest is None: 1450 # Note that RAMSES can have partial refinement inside an Oct. This 1451 # means we actually do want the number of Octs, not the number of 1452 # cells. 1453 num_cells = self._last_mask_sum 1454 if dims > 1: 1455 dest = np.zeros((num_cells, dims), dtype=source.dtype, 1456 order='C') 1457 else: 1458 dest = np.zeros(num_cells, dtype=source.dtype, order='C') 1459 ddata = (<char*>dest.data) + offset*ss*dims 1460 ind = 0 1461 for sfc in range(self.sfc_start, self.sfc_end + 1): 1462 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1463 sfci += 1 1464 if mask[sfci] == 0: continue 1465 memcpy(ddata, sdata + ind, dims * ss) 1466 ddata += dims * ss 1467 filled += 1 1468 ind += ss * dims 1469 if num_cells >= 0: 1470 return dest 1471 return filled 1472 1473 @cython.boundscheck(False) 1474 @cython.wraparound(False) 1475 @cython.cdivision(True) 1476 def mask(self, SelectorObject selector, np.int64_t num_cells = -1, 1477 int domain_id = -1): 1478 # We take a domain_id here to avoid subclassing 1479 cdef int i 1480 cdef np.float64_t pos[3] 1481 cdef np.int64_t sfc, sfci = -1 1482 if self._last_selector_id == hash(selector): 1483 return self._last_mask 1484 cdef np.ndarray[np.uint8_t, ndim=1] mask 1485 mask = np.zeros((self.nsfc), dtype="uint8") 1486 self._last_mask_sum = 0 1487 for sfc in range(self.sfc_start, self.sfc_end + 1): 1488 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1489 sfci += 1 1490 self.sfc_to_pos(sfc, pos) 1491 if selector.select_cell(pos, self.dds) == 0: continue 1492 mask[sfci] = 1 1493 self._last_mask_sum += 1 1494 self._last_mask = mask.astype("bool") 1495 self._last_selector_id = hash(selector) 1496 return self._last_mask 1497 1498 1499 def fill_sfc_particles(self, fields): 1500 rv = read_sfc_particles(self.artio_handle, 1501 self.sfc_start, self.sfc_end, 1502 1, fields, self.range_handler.doct_count, 1503 self.range_handler.pcount) 1504 return rv 1505 1506 @cython.boundscheck(False) 1507 @cython.wraparound(False) 1508 @cython.cdivision(True) 1509 def fill_sfc(self, SelectorObject selector, field_indices): 1510 cdef np.ndarray[np.float64_t, ndim=1] dest 1511 cdef int n, status, i, di, num_oct_levels, nf, ngv, max_level 1512 cdef np.int64_t sfc, num_cells, sfci = -1 1513 cdef np.float64_t val 1514 cdef double dpos[3] 1515 max_level = self.artio_handle.max_level 1516 cdef int *num_octs_per_level = <int *>malloc( 1517 (max_level + 1)*sizeof(int)) 1518 # We duplicate some of the grid_variables stuff here so that we can 1519 # potentially release the GIL 1520 nf = len(field_indices) 1521 ngv = self.artio_handle.num_grid_variables 1522 cdef float *grid_variables = <float *>malloc( 1523 ngv * sizeof(float)) 1524 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1525 mask = self.mask(selector, -1) 1526 num_cells = self._last_mask_sum 1527 tr = [] 1528 for i in range(nf): 1529 tr.append(np.zeros(num_cells, dtype="float64")) 1530 cdef int* field_ind = <int*> malloc( 1531 nf * sizeof(int)) 1532 cdef np.float64_t **field_vals = <np.float64_t**> malloc( 1533 nf * sizeof(np.float64_t*)) 1534 for i in range(nf): 1535 field_ind[i] = field_indices[i] 1536 # This zeros should be an empty once we handle the root grid 1537 dest = tr[i] 1538 field_vals[i] = <np.float64_t*> dest.data 1539 # First we need to walk the mesh in the file. Then we fill in the dest 1540 # location based on the file index. 1541 cdef int filled = 0 1542 cdef float **mesh_data = self.range_handler.root_mesh_data 1543 if mesh_data == NULL: 1544 status = artio_grid_cache_sfc_range(self.handle, self.sfc_start, 1545 self.sfc_end) 1546 check_artio_status(status) 1547 for sfc in range(self.sfc_start, self.sfc_end + 1): 1548 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1549 sfci += 1 1550 if mask[sfci] == 0: continue 1551 status = artio_grid_read_root_cell_begin( self.handle, 1552 sfc, dpos, grid_variables, &num_oct_levels, 1553 num_octs_per_level) 1554 check_artio_status(status) 1555 for i in range(nf): 1556 field_vals[i][filled] = grid_variables[field_ind[i]] 1557 filled += 1 1558 status = artio_grid_read_root_cell_end(self.handle) 1559 check_artio_status(status) 1560 status = artio_grid_clear_sfc_cache(self.handle) 1561 check_artio_status(status) 1562 else: 1563 for sfc in range(self.sfc_start, self.sfc_end + 1): 1564 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1565 sfci += 1 1566 if mask[sfci] == 0: continue 1567 for i in range(nf): 1568 field_vals[i][filled] = mesh_data[field_ind[i]][ 1569 sfc - self.sfc_start] 1570 filled += 1 1571 # Now we have all our sources. 1572 free(field_ind) 1573 free(field_vals) 1574 free(grid_variables) 1575 free(num_octs_per_level) 1576 return tr 1577 1578 @cython.boundscheck(False) 1579 @cython.wraparound(False) 1580 @cython.cdivision(True) 1581 def deposit(self, ParticleDepositOperation pdeposit, 1582 SelectorObject selector, 1583 np.ndarray[np.float64_t, ndim=2] positions, 1584 fields): 1585 # This implements the necessary calls to enable particle deposition to 1586 # occur as needed. 1587 cdef int nf, i, j 1588 cdef np.int64_t sfc, sfci 1589 if fields is None: 1590 fields = [] 1591 nf = len(fields) 1592 cdef np.float64_t[::cython.view.indirect, ::1] field_pointers 1593 if nf > 0: field_pointers = OnceIndirect(fields) 1594 cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64") 1595 cdef np.ndarray[np.uint8_t, ndim=1, cast=True] mask 1596 mask = self.mask(selector, -1) 1597 cdef np.ndarray[np.int64_t, ndim=1] domain_ind 1598 domain_ind = np.zeros(self.sfc_end - self.sfc_start + 1, 1599 dtype="int64") - 1 1600 j = 0 1601 sfci = -1 1602 for sfc in range(self.sfc_start, self.sfc_end + 1): 1603 if self.sfc_mask[sfc - self.sfc_start] == 0: continue 1604 sfci += 1 1605 if mask[sfci] == 0: 1606 continue 1607 domain_ind[sfc - self.sfc_start] = j 1608 j += 1 1609 cdef np.float64_t pos[3] 1610 cdef np.float64_t left_edge[3] 1611 cdef int coords[3] 1612 cdef int dims[3] 1613 dims[0] = dims[1] = dims[2] = 1 1614 cdef np.int64_t offset, moff 1615 cdef np.int64_t numpart = positions.shape[0] 1616 for i in range(positions.shape[0]): 1617 for j in range(nf): 1618 field_vals[j] = field_pointers[j][i] 1619 for j in range(3): 1620 pos[j] = positions[i, j] 1621 coords[j] = <int>((pos[j] - self.DLE[j])/self.dds[j]) 1622 sfc = artio_sfc_index(self.artio_handle.handle, coords) 1623 if sfc < self.sfc_start or sfc > self.sfc_end: continue 1624 offset = domain_ind[sfc - self.sfc_start] 1625 if offset < 0: continue 1626 # Check that we found the oct ... 1627 for j in range(3): 1628 left_edge[j] = coords[j] * self.dds[j] + self.DLE[j] 1629 pdeposit.process(dims, i, left_edge, self.dds, 1630 offset, pos, field_vals, sfc) 1631 if pdeposit.update_values == 1: 1632 for j in range(nf): 1633 field_pointers[j][i] = field_vals[j] 1634 1635cdef class SFCRangeSelector(SelectorObject): 1636 1637 cdef SelectorObject base_selector 1638 cdef ARTIOSFCRangeHandler range_handler 1639 cdef ARTIORootMeshContainer mesh_container 1640 cdef np.int64_t sfc_start, sfc_end 1641 1642 def __init__(self, dobj): 1643 self.base_selector = dobj.base_selector 1644 self.mesh_container = dobj.oct_handler 1645 self.range_handler = self.mesh_container.range_handler 1646 self.sfc_start = self.mesh_container.sfc_start 1647 self.sfc_end = self.mesh_container.sfc_end 1648 1649 @cython.boundscheck(False) 1650 @cython.wraparound(False) 1651 @cython.cdivision(True) 1652 def select_grids(self, 1653 np.ndarray[np.float64_t, ndim=2] left_edges, 1654 np.ndarray[np.float64_t, ndim=2] right_edges, 1655 np.ndarray[np.int32_t, ndim=2] levels): 1656 raise RuntimeError 1657 1658 @cython.boundscheck(False) 1659 @cython.wraparound(False) 1660 @cython.cdivision(True) 1661 cdef int select_sphere(self, np.float64_t pos[3], np.float64_t radius) nogil: 1662 return 1 1663 1664 @cython.boundscheck(False) 1665 @cython.wraparound(False) 1666 @cython.cdivision(True) 1667 cdef int select_cell(self, np.float64_t pos[3], np.float64_t dds[3]) nogil: 1668 return self.select_point(pos) 1669 1670 @cython.boundscheck(False) 1671 @cython.wraparound(False) 1672 @cython.cdivision(True) 1673 cdef int select_point(self, np.float64_t pos[3]) nogil: 1674 cdef np.int64_t sfc = self.mesh_container.pos_to_sfc(pos) 1675 if sfc > self.sfc_end: return 0 1676 cdef np.int64_t oc = self.range_handler.doct_count[ 1677 sfc - self.sfc_start] 1678 if oc > 0: return 0 1679 return 1 1680 1681 @cython.boundscheck(False) 1682 @cython.wraparound(False) 1683 @cython.cdivision(True) 1684 cdef int select_bbox(self, np.float64_t left_edge[3], 1685 np.float64_t right_edge[3]) nogil: 1686 return self.base_selector.select_bbox(left_edge, right_edge) 1687 1688 @cython.boundscheck(False) 1689 @cython.wraparound(False) 1690 @cython.cdivision(True) 1691 cdef int select_grid(self, np.float64_t left_edge[3], 1692 np.float64_t right_edge[3], np.int32_t level, 1693 Oct *o = NULL) nogil: 1694 # Because visitors now use select_grid, we should be explicitly 1695 # checking this. 1696 return self.base_selector.select_grid(left_edge, right_edge, level, o) 1697 1698 def _hash_vals(self): 1699 return (hash(self.base_selector), self.sfc_start, self.sfc_end) 1700 1701sfc_subset_selector = AlwaysSelector 1702#sfc_subset_selector = SFCRangeSelector 1703