<s> from __future__ import print_function , division <EOL> import hashlib <EOL> from copy import deepcopy <EOL> import h5py <EOL> import numpy as np <EOL> from . . util . meshgrid import meshgrid_nd <EOL> from . . util . functions import FreezableClass , is_numpy_array , monotonically_increasing , link_or_copy <EOL> from astropy import log as logger <EOL> from . grid_helpers import single_grid_dims <EOL> class SphericalPolarGrid ( FreezableClass ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> self . shape = None <EOL> self . r_wall = None <EOL> self . t_wall = None <EOL> self . p_wall = None <EOL> self . r = None <EOL> self . t = None <EOL> self . p = None <EOL> self . gr = None <EOL> self . gt = None <EOL> self . gp = None <EOL> self . gw = None <EOL> self . gz = None <EOL> self . volumes = None <EOL> self . areas = None <EOL> self . widths = None <EOL> self . quantities = { } <EOL> self . _freeze ( ) <EOL> if len ( args ) > <NUM_LIT:0> : <EOL> if isinstance ( args [ <NUM_LIT:0> ] , SphericalPolarGrid ) : <EOL> self . set_walls ( args [ <NUM_LIT:0> ] . r_wall , args [ <NUM_LIT:0> ] . t_wall , args [ <NUM_LIT:0> ] . p_wall ) <EOL> else : <EOL> self . set_walls ( * args ) <EOL> def set_walls ( self , r_wall , t_wall , p_wall ) : <EOL> if type ( r_wall ) in [ list , tuple ] : <EOL> r_wall = np . array ( r_wall ) <EOL> if type ( t_wall ) in [ list , tuple ] : <EOL> t_wall = np . array ( t_wall ) <EOL> if type ( p_wall ) in [ list , tuple ] : <EOL> p_wall = np . array ( p_wall ) <EOL> if not is_numpy_array ( r_wall ) or r_wall . ndim != <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not is_numpy_array ( t_wall ) or t_wall . ndim != <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not is_numpy_array ( p_wall ) or p_wall . ndim != <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not monotonically_increasing ( r_wall ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not monotonically_increasing ( t_wall ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not monotonically_increasing ( p_wall ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if np . any ( t_wall < <NUM_LIT:0.> ) or np . any ( t_wall > np . pi ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if np . any ( p_wall < <NUM_LIT:0.> ) or np . any ( p_wall > <NUM_LIT> * np . pi ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . shape = ( len ( p_wall ) - <NUM_LIT:1> , len ( t_wall ) - <NUM_LIT:1> , len ( r_wall ) - <NUM_LIT:1> ) <EOL> self . r_wall = r_wall <EOL> self . t_wall = t_wall <EOL> self . p_wall = p_wall <EOL> if r_wall [ <NUM_LIT:0> ] == <NUM_LIT:0.> : <EOL> self . r = np . zeros ( len ( r_wall ) - <NUM_LIT:1> ) <EOL> self . r [ <NUM_LIT:0> ] = r_wall [ <NUM_LIT:1> ] / <NUM_LIT> <EOL> self . r [ <NUM_LIT:1> : ] = <NUM_LIT> ** ( ( np . log10 ( r_wall [ <NUM_LIT:1> : - <NUM_LIT:1> ] ) + np . log10 ( r_wall [ <NUM_LIT:2> : ] ) ) / <NUM_LIT> ) <EOL> else : <EOL> self . r = <NUM_LIT> ** ( ( np . log10 ( r_wall [ : - <NUM_LIT:1> ] ) + np . log10 ( r_wall [ <NUM_LIT:1> : ] ) ) / <NUM_LIT> ) <EOL> self . t = ( t_wall [ : - <NUM_LIT:1> ] + t_wall [ <NUM_LIT:1> : ] ) / <NUM_LIT> <EOL> self . p = ( p_wall [ : - <NUM_LIT:1> ] + p_wall [ <NUM_LIT:1> : ] ) / <NUM_LIT> <EOL> self . gr , self . gt , self . gp = meshgrid_nd ( self . r , self . t , self . p ) <EOL> self . gz = self . gr * np . cos ( self . gt ) <EOL> self . gw = self . gr * np . sin ( self . gt ) <EOL> gr_wall_min , gt_wall_min , gp_wall_min = meshgrid_nd ( r_wall [ : - <NUM_LIT:1> ] , t_wall [ : - <NUM_LIT:1> ] , p_wall [ : - <NUM_LIT:1> ] ) <EOL> gr_wall_max , gt_wall_max , gp_wall_max = meshgrid_nd ( r_wall [ <NUM_LIT:1> : ] , t_wall [ <NUM_LIT:1> : ] , p_wall [ <NUM_LIT:1> : ] ) <EOL> dr = gr_wall_max - gr_wall_min <EOL> dr2 = gr_wall_max ** <NUM_LIT:2> - gr_wall_min ** <NUM_LIT:2> <EOL> dr3 = gr_wall_max ** <NUM_LIT:3> - gr_wall_min ** <NUM_LIT:3> <EOL> dt = gt_wall_max - gt_wall_min <EOL> dcost = np . cos ( gt_wall_min ) - np . cos ( gt_wall_max ) <EOL> dp = gp_wall_max - gp_wall_min <EOL> self . volumes = dr3 * dcost * dp / <NUM_LIT> <EOL> self . areas = np . zeros ( ( <NUM_LIT:6> , ) + self . shape ) <EOL> self . areas [ <NUM_LIT:0> , : , : , : ] = gr_wall_min ** <NUM_LIT:2> * dcost * dp <EOL> self . areas [ <NUM_LIT:1> , : , : , : ] = gr_wall_max ** <NUM_LIT:2> * dcost * dp <EOL> self . areas [ <NUM_LIT:2> , : , : , : ] = <NUM_LIT:0.5> * dr2 * np . sin ( gt_wall_min ) * dp <EOL> self . areas [ <NUM_LIT:3> , : , : , : ] = <NUM_LIT:0.5> * dr2 * np . sin ( gt_wall_max ) * dp <EOL> self . areas [ <NUM_LIT:4> , : , : , : ] = <NUM_LIT:0.5> * dr2 * dt <EOL> self . areas [ <NUM_LIT:5> , : , : , : ] = <NUM_LIT:0.5> * dr2 * dt <EOL> self . widths = np . zeros ( ( <NUM_LIT:3> , ) + self . shape ) <EOL> self . widths [ <NUM_LIT:0> , : , : , : ] = dr <EOL> self . widths [ <NUM_LIT:1> , : , : , : ] = self . gr * dt <EOL> self . widths [ <NUM_LIT:2> , : , : , : ] = self . gr * np . sin ( self . gt ) * dp <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 ] ) <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> def _check_array_dimensions ( self , array = None ) : <EOL> '''<STR_LIT>''' <EOL> n_pop_ref = None <EOL> if isinstance ( array , SphericalPolarGridView ) : <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 ] ) <EOL> else : <EOL> n_pop , shape = single_grid_dims ( array ) <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 [ '<STR_LIT>' ] [ '<STR_LIT:r>' ] , <EOL> group [ '<STR_LIT>' ] [ '<STR_LIT:t>' ] , <EOL> group [ '<STR_LIT>' ] [ '<STR_LIT:p>' ] ) <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:4> : <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> dset = g_geometry . create_dataset ( "<STR_LIT>" , data = np . array ( list ( zip ( self . r_wall ) ) , dtype = [ ( '<STR_LIT:r>' , wall_dtype ) ] ) , compression = compression ) <EOL> dset . attrs [ '<STR_LIT>' ] = np . string_ ( '<STR_LIT>' . encode ( '<STR_LIT:utf-8>' ) ) <EOL> dset = g_geometry . create_dataset ( "<STR_LIT>" , data = np . array ( list ( zip ( self . t_wall ) ) , dtype = [ ( '<STR_LIT:t>' , wall_dtype ) ] ) , compression = compression ) <EOL> dset . attrs [ '<STR_LIT>' ] = np . string_ ( '<STR_LIT>' . encode ( '<STR_LIT:utf-8>' ) ) <EOL> dset = g_geometry . create_dataset ( "<STR_LIT>" , data = np . array ( list ( zip ( self . p_wall ) ) , dtype = [ ( '<STR_LIT:p>' , wall_dtype ) ] ) , compression = compression ) <EOL> dset . attrs [ '<STR_LIT>' ] = np . string_ ( '<STR_LIT>' . encode ( '<STR_LIT:utf-8>' ) ) <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 ( self . r_wall . tostring ( ) ) <EOL> geo_hash . update ( self . t_wall . tostring ( ) ) <EOL> geo_hash . update ( self . p_wall . tostring ( ) ) <EOL> return geo_hash . hexdigest ( ) <EOL> def __getitem__ ( self , item ) : <EOL> return SphericalPolarGridView ( self , item ) <EOL> def __setitem__ ( self , item , value ) : <EOL> if isinstance ( value , SphericalPolarGridView ) : <EOL> if self . r_wall is None and self . t_wall is None and self . p_wall is None : <EOL> logger . warn ( "<STR_LIT>" ) <EOL> self . set_walls ( value . r_wall , value . t_wall , value . p_wall ) <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> class SphericalPolarGridView ( SphericalPolarGrid ) : <EOL> def __init__ ( self , grid , quantity ) : <EOL> self . viewed_quantity = quantity <EOL> SphericalPolarGrid . __init__ ( self ) <EOL> self . set_walls ( grid . r_wall , grid . t_wall , grid . p_wall ) <EOL> self . quantities = { quantity : grid . quantities [ quantity ] } <EOL> def append ( self , grid ) : <EOL> '''<STR_LIT>''' <EOL> if isinstance ( grid , SphericalPolarGridView ) : <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 , SphericalPolarGridView ) : <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 = SphericalPolarGridView ( self , self . viewed_quantity ) <EOL> grid . quantities = { grid . viewed_quantity : grid . quantities [ grid . viewed_quantity ] [ item ] } <EOL> return grid <EOL> else : <EOL> return SphericalPolarGrid . __getitem__ ( self , item ) <EOL> def __getattr__ ( self , attribute ) : <EOL> if attribute == '<STR_LIT>' : <EOL> return self . quantities [ self . viewed_quantity ] <EOL> else : <EOL> return SphericalPolarGrid . __getattr__ ( self , attribute ) </s>
