<s> from __future__ import print_function , division <EOL> import struct <EOL> import hashlib <EOL> from copy import deepcopy <EOL> import h5py <EOL> import numpy as np <EOL> from . . util . functions import FreezableClass , is_numpy_array , link_or_copy <EOL> from astropy import log as logger <EOL> from . grid_helpers import single_grid_dims <EOL> class OctreeGrid ( FreezableClass ) : <EOL> '''<STR_LIT>''' <EOL> _validate_cache = { } <EOL> def __init__ ( self , * args ) : <EOL> self . shape = None <EOL> self . x = None <EOL> self . y = None <EOL> self . z = None <EOL> self . dx = None <EOL> self . dy = None <EOL> self . dz = None <EOL> self . refined = None <EOL> self . quantities = { } <EOL> self . _freeze ( ) <EOL> if len ( args ) > <NUM_LIT:0> : <EOL> if isinstance ( args [ <NUM_LIT:0> ] , OctreeGrid ) : <EOL> self . set_walls ( args [ <NUM_LIT:0> ] . x , args [ <NUM_LIT:0> ] . y , args [ <NUM_LIT:0> ] . z , <EOL> args [ <NUM_LIT:0> ] . dx , args [ <NUM_LIT:0> ] . dy , args [ <NUM_LIT:0> ] . dz , <EOL> args [ <NUM_LIT:0> ] . refined ) <EOL> else : <EOL> self . set_walls ( * args ) <EOL> def set_walls ( self , x , y , z , dx , dy , dz , refined ) : <EOL> if not np . isscalar ( x ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( x ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isscalar ( y ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( y ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isscalar ( z ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( z ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isscalar ( dx ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( dx ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isscalar ( dy ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( dy ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isscalar ( dz ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not np . isreal ( dz ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if type ( refined ) in [ list , tuple ] : <EOL> refined = np . array ( refined ) <EOL> if not is_numpy_array ( refined ) or refined . ndim != <NUM_LIT:1> or refined . dtype != bool : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . x = x <EOL> self . y = y <EOL> self . z = z <EOL> self . dx = dx <EOL> self . dy = dy <EOL> self . dz = dz <EOL> self . refined = refined <EOL> self . shape = ( len ( refined ) , ) <EOL> @ property <EOL> def refined ( self ) : <EOL> return self . _refined <EOL> @ refined . setter <EOL> def refined ( self , value ) : <EOL> if value is None : <EOL> self . _refined = None <EOL> return <EOL> if not ( len ( value ) - <NUM_LIT:1> ) % <NUM_LIT:8> == <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . _refined = self . _validate ( value ) <EOL> def _validate ( self , value ) : <EOL> value_hash = hashlib . md5 ( value . tostring ( ) ) . hexdigest ( ) <EOL> if value_hash in self . _validate_cache : <EOL> return value <EOL> logger . info ( "<STR_LIT>" ) <EOL> refined_str = value . tostring ( ) <EOL> previous = '<STR_LIT>' <EOL> while True : <EOL> refined_str = refined_str . replace ( b'<STR_LIT>' , b'<STR_LIT:\x00>' ) <EOL> if refined_str == previous : <EOL> break <EOL> else : <EOL> previous = refined_str <EOL> if refined_str != b'<STR_LIT:\x00>' : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def check_recursive ( refined , current_i = <NUM_LIT:0> , max_level = <NUM_LIT:0> ) : <EOL> if refined [ current_i ] : <EOL> current_i += <NUM_LIT:1> <EOL> max_levels = [ ] <EOL> for i in range ( <NUM_LIT:8> ) : <EOL> current_i , max_level_indiv = check_recursive ( refined , current_i , max_level + <NUM_LIT:1> ) <EOL> max_levels . append ( max_level_indiv ) <EOL> max_level = max ( max_levels ) <EOL> else : <EOL> current_i += <NUM_LIT:1> <EOL> return current_i , max_level <EOL> try : <EOL> final_i , max_level = check_recursive ( value ) <EOL> except IndexError : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> logger . info ( "<STR_LIT>" . format ( max_level ) ) <EOL> if max_level > <NUM_LIT:20> : <EOL> logger . warn ( "<STR_LIT>" . format ( max_level ) ) <EOL> self . _validate_cache [ value_hash ] = True <EOL> return value <EOL> def __getattr__ ( self , attribute ) : <EOL> if attribute == '<STR_LIT>' : <EOL> n_dust = None <EOL> for quantity in self . quantities : <EOL> n_dust_q , shape_q = single_grid_dims ( self . quantities [ quantity ] , ndim = <NUM_LIT:1> ) <EOL> if n_dust is None : <EOL> n_dust = n_dust_q <EOL> elif n_dust_q is not None : <EOL> if n_dust != n_dust_q : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return n_dust <EOL> else : <EOL> return FreezableClass . __getattribute__ ( self , attribute ) <EOL> @ property <EOL> def limits ( self ) : <EOL> from hyperion . importers . _discretize_sph import _get_positions_widths <EOL> xc , yc , zc , xw , yw , zw = _get_positions_widths ( self . refined , <EOL> self . x , self . y , self . z , <EOL> self . dx , self . dy , self . dz ) <EOL> return xc - xw , xc + xw , yc - yw , yc + yw , zc - zw , zc + zw <EOL> @ property <EOL> def volumes ( self ) : <EOL> """<STR_LIT>""" <EOL> xmin , xmax , ymin , ymax , zmin , zmax = self . limits <EOL> return ( xmax - xmin ) * ( ymax - ymin ) * ( zmax - zmin ) <EOL> def _check_array_dimensions ( self , array = None ) : <EOL> '''<STR_LIT>''' <EOL> n_pop_ref = None <EOL> if isinstance ( array , OctreeGridView ) : <EOL> array = array . quantities [ array . viewed_quantity ] <EOL> for quantity in self . quantities : <EOL> if array is None : <EOL> n_pop , shape = single_grid_dims ( self . quantities [ quantity ] , ndim = <NUM_LIT:1> ) <EOL> else : <EOL> n_pop , shape = single_grid_dims ( array , ndim = <NUM_LIT:1> ) <EOL> if shape != self . shape : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> % ( shape , self . shape ) ) <EOL> if n_pop is not None : <EOL> if n_pop_ref is None : <EOL> n_pop_ref = n_pop <EOL> elif n_pop != n_pop_ref : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def read ( self , group , quantities = '<STR_LIT:all>' ) : <EOL> '''<STR_LIT>''' <EOL> self . read_geometry ( group [ '<STR_LIT>' ] ) <EOL> self . read_quantities ( group [ '<STR_LIT>' ] , quantities = quantities ) <EOL> self . _check_array_dimensions ( ) <EOL> def read_geometry ( self , group ) : <EOL> '''<STR_LIT>''' <EOL> if group . attrs [ '<STR_LIT>' ] . decode ( '<STR_LIT:utf-8>' ) != '<STR_LIT>' : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . set_walls ( group . attrs [ '<STR_LIT:x>' ] , <EOL> group . attrs [ '<STR_LIT:y>' ] , <EOL> group . attrs [ '<STR_LIT:z>' ] , <EOL> group . attrs [ '<STR_LIT>' ] , <EOL> group . attrs [ '<STR_LIT>' ] , <EOL> group . attrs [ '<STR_LIT>' ] , <EOL> group [ '<STR_LIT>' ] [ '<STR_LIT>' ] . astype ( bool ) ) <EOL> if group . attrs [ '<STR_LIT>' ] . decode ( '<STR_LIT:utf-8>' ) != self . get_geometry_id ( ) : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> def read_quantities ( self , group , quantities = '<STR_LIT:all>' ) : <EOL> '''<STR_LIT>''' <EOL> if quantities is not None : <EOL> for quantity in group : <EOL> if quantities == '<STR_LIT:all>' or quantity in quantities : <EOL> array = np . array ( group [ quantity ] ) <EOL> if array . ndim == <NUM_LIT:2> : <EOL> self . quantities [ quantity ] = [ array [ i ] for i in range ( array . shape [ <NUM_LIT:0> ] ) ] <EOL> else : <EOL> self . quantities [ quantity ] = array <EOL> self . _check_array_dimensions ( ) <EOL> def write ( self , group , quantities = '<STR_LIT:all>' , copy = True , absolute_paths = False , compression = True , wall_dtype = float , physics_dtype = float ) : <EOL> '''<STR_LIT>''' <EOL> if '<STR_LIT>' not in group : <EOL> g_geometry = group . create_group ( '<STR_LIT>' ) <EOL> else : <EOL> g_geometry = group [ '<STR_LIT>' ] <EOL> if '<STR_LIT>' not in group : <EOL> g_quantities = group . create_group ( '<STR_LIT>' ) <EOL> else : <EOL> g_quantities = group [ '<STR_LIT>' ] <EOL> g_geometry . attrs [ '<STR_LIT>' ] = np . string_ ( '<STR_LIT>' . encode ( '<STR_LIT:utf-8>' ) ) <EOL> g_geometry . attrs [ '<STR_LIT>' ] = np . string_ ( self . get_geometry_id ( ) . encode ( '<STR_LIT:utf-8>' ) ) <EOL> g_geometry . attrs [ '<STR_LIT:x>' ] = self . x <EOL> g_geometry . attrs [ '<STR_LIT:y>' ] = self . y <EOL> g_geometry . attrs [ '<STR_LIT:z>' ] = self . z <EOL> g_geometry . attrs [ '<STR_LIT>' ] = self . dx <EOL> g_geometry . attrs [ '<STR_LIT>' ] = self . dy <EOL> g_geometry . attrs [ '<STR_LIT>' ] = self . dz <EOL> dset = g_geometry . create_dataset ( "<STR_LIT>" , data = np . array ( list ( zip ( self . refined ) ) , dtype = [ ( '<STR_LIT>' , np . int32 ) ] ) , compression = compression ) <EOL> self . _check_array_dimensions ( ) <EOL> for quantity in self . quantities : <EOL> if quantities == '<STR_LIT:all>' or quantity in quantities : <EOL> if isinstance ( self . quantities [ quantity ] , h5py . ExternalLink ) : <EOL> link_or_copy ( g_quantities , quantity , self . quantities [ quantity ] , copy , absolute_paths = absolute_paths ) <EOL> else : <EOL> dset = g_quantities . create_dataset ( quantity , data = self . quantities [ quantity ] , <EOL> compression = compression , <EOL> dtype = physics_dtype ) <EOL> dset . attrs [ '<STR_LIT>' ] = np . string_ ( self . get_geometry_id ( ) . encode ( '<STR_LIT:utf-8>' ) ) <EOL> def write_single_array ( self , group , name , array , copy = True , absolute_paths = False , compression = True , physics_dtype = float ) : <EOL> '''<STR_LIT>''' <EOL> self . _check_array_dimensions ( array ) <EOL> if isinstance ( array , h5py . ExternalLink ) : <EOL> link_or_copy ( group , name , array , copy , absolute_paths = absolute_paths ) <EOL> else : <EOL> dset = group . create_dataset ( name , data = array , <EOL> compression = compression , <EOL> dtype = physics_dtype ) <EOL> dset . attrs [ '<STR_LIT>' ] = np . string_ ( self . get_geometry_id ( ) . encode ( '<STR_LIT:utf-8>' ) ) <EOL> def get_geometry_id ( self ) : <EOL> geo_hash = hashlib . md5 ( ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . x ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . y ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . z ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . dx ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . dy ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , self . dz ) ) <EOL> geo_hash . update ( self . refined . tostring ( ) ) <EOL> return geo_hash . hexdigest ( ) <EOL> def __getitem__ ( self , item ) : <EOL> return OctreeGridView ( self , item ) <EOL> def __setitem__ ( self , item , value ) : <EOL> if isinstance ( value , OctreeGridView ) : <EOL> if self . refined is None : <EOL> logger . warn ( "<STR_LIT>" ) <EOL> self . set_walls ( value . x , value . y , value . z , value . dx , value . dy , value . dz , value . refined ) <EOL> self . quantities [ item ] = deepcopy ( value . quantities [ value . viewed_quantity ] ) <EOL> elif isinstance ( value , h5py . ExternalLink ) : <EOL> self . quantities [ item ] = value <EOL> elif value == [ ] : <EOL> self . quantities [ item ] = [ ] <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def __contains__ ( self , item ) : <EOL> return self . quantities . __contains__ ( item ) <EOL> def reset_quantities ( self ) : <EOL> self . quantities = { } <EOL> def add_derived_quantity ( self , name , function ) : <EOL> if name in self . quantities : <EOL> raise KeyError ( name + '<STR_LIT>' ) <EOL> function ( self . quantities ) <EOL> def to_yt ( self , dust_id = <NUM_LIT:0> ) : <EOL> '''<STR_LIT>''' <EOL> from yt_wrappers import octree_grid_to_yt_stream <EOL> return octree_grid_to_yt_stream ( self , dust_id ) <EOL> class OctreeGridView ( OctreeGrid ) : <EOL> def __init__ ( self , grid , quantity ) : <EOL> self . viewed_quantity = quantity <EOL> OctreeGrid . __init__ ( self ) <EOL> self . set_walls ( grid . x , grid . y , grid . z , grid . dx , grid . dy , grid . dz , grid . refined ) <EOL> self . quantities = { quantity : grid . quantities [ quantity ] } <EOL> def append ( self , grid ) : <EOL> '''<STR_LIT>''' <EOL> if isinstance ( grid , OctreeGridView ) : <EOL> if self . quantities [ self . viewed_quantity ] is grid . quantities [ grid . viewed_quantity ] : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if type ( grid . quantities [ grid . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> self . _check_array_dimensions ( grid . quantities [ grid . viewed_quantity ] ) <EOL> self . quantities [ self . viewed_quantity ] . append ( deepcopy ( grid . quantities [ grid . viewed_quantity ] ) ) <EOL> elif isinstance ( grid , np . ndarray ) : <EOL> self . _check_array_dimensions ( grid ) <EOL> self . quantities [ self . viewed_quantity ] . append ( deepcopy ( grid ) ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def add ( self , grid ) : <EOL> '''<STR_LIT>''' <EOL> if type ( self . quantities [ self . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if isinstance ( grid , OctreeGridView ) : <EOL> if type ( grid . quantities [ grid . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> self . _check_array_dimensions ( grid . quantities [ grid . viewed_quantity ] ) <EOL> self . quantities [ self . viewed_quantity ] += grid . quantities [ grid . viewed_quantity ] <EOL> elif isinstance ( grid , np . ndarray ) : <EOL> self . _check_array_dimensions ( grid ) <EOL> self . quantities [ self . viewed_quantity ] += grid <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def __getitem__ ( self , item ) : <EOL> if type ( item ) is int : <EOL> grid = OctreeGridView ( self , self . viewed_quantity ) <EOL> grid . quantities = { grid . viewed_quantity : grid . quantities [ grid . viewed_quantity ] [ item ] } <EOL> return grid <EOL> else : <EOL> return OctreeGrid . __getitem__ ( self , item ) <EOL> def __getattr__ ( self , attribute ) : <EOL> if attribute == '<STR_LIT>' : <EOL> return self . quantities [ self . viewed_quantity ] <EOL> else : <EOL> return OctreeGrid . __getattr__ ( self , attribute ) </s>
