<s> from __future__ import print_function , division <EOL> import os <EOL> import struct <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 , link_or_copy <EOL> from astropy import log as logger <EOL> from . grid_helpers import single_grid_dims <EOL> def zero_density ( grid , xmin = - np . inf , xmax = np . inf , ymin = - np . inf , ymax = np . inf , zmin = np . inf , zmax = np . inf ) : <EOL> for ilevel , level in enumerate ( grid . levels ) : <EOL> for igrid , grid in enumerate ( level . grids ) : <EOL> wx = np . linspace ( grid . xmin , grid . xmax , grid . nx + <NUM_LIT:1> ) <EOL> wy = np . linspace ( grid . ymin , grid . ymax , grid . ny + <NUM_LIT:1> ) <EOL> wz = np . linspace ( grid . zmin , grid . zmax , grid . nz + <NUM_LIT:1> ) <EOL> x = <NUM_LIT:0.5> * ( wx [ : - <NUM_LIT:1> ] + wx [ <NUM_LIT:1> : ] ) <EOL> y = <NUM_LIT:0.5> * ( wy [ : - <NUM_LIT:1> ] + wy [ <NUM_LIT:1> : ] ) <EOL> z = <NUM_LIT:0.5> * ( wz [ : - <NUM_LIT:1> ] + wz [ <NUM_LIT:1> : ] ) <EOL> gx , gy , gz = meshgrid_nd ( x , y , z ) <EOL> reset = ( gx < xmin ) | ( gx > xmax ) | ( gy < ymin ) | ( gy > ymax ) | ( gz < zmin ) | ( gz > zmax ) <EOL> grid . data [ reset ] = <NUM_LIT:0.> <EOL> return grid <EOL> class Grid ( FreezableClass ) : <EOL> def __init__ ( self ) : <EOL> self . quantities = { } <EOL> self . xmin , self . xmax = None , None <EOL> self . ymin , self . ymax = None , None <EOL> self . zmin , self . zmax = None , None <EOL> self . nx , self . ny , self . nz = None , None , None <EOL> self . _freeze ( ) <EOL> def __getattr__ ( self , attribute ) : <EOL> if attribute == '<STR_LIT>' : <EOL> return ( self . nz , self . ny , self . nx ) <EOL> else : <EOL> return FreezableClass . __getattribute__ ( self , attribute ) <EOL> class Level ( FreezableClass ) : <EOL> def __init__ ( self ) : <EOL> self . grids = [ ] <EOL> self . _freeze ( ) <EOL> def add_grid ( self ) : <EOL> grid = Grid ( ) <EOL> self . grids . append ( grid ) <EOL> return grid <EOL> class AMRGrid ( FreezableClass ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , amr_grid = None ) : <EOL> self . levels = [ ] <EOL> self . _freeze ( ) <EOL> if amr_grid is not None : <EOL> for level in amr_grid . levels : <EOL> level_ref = self . add_level ( ) <EOL> for grid in level . grids : <EOL> grid_ref = level_ref . add_grid ( ) <EOL> grid_ref . nx = grid . nx <EOL> grid_ref . ny = grid . ny <EOL> grid_ref . nz = grid . nz <EOL> grid_ref . xmin , grid_ref . xmax = grid . xmin , grid . xmax <EOL> grid_ref . ymin , grid_ref . ymax = grid . ymin , grid . ymax <EOL> grid_ref . zmin , grid_ref . zmax = grid . zmin , grid . zmax <EOL> grid_ref . quantities = { } <EOL> def remove_level ( self , level_id ) : <EOL> self . levels . pop ( level_id ) <EOL> def add_level ( self ) : <EOL> level = Level ( ) <EOL> self . levels . append ( level ) <EOL> return level <EOL> def __getattr__ ( self , attribute ) : <EOL> if attribute == '<STR_LIT>' : <EOL> n_dust = None <EOL> for level in self . levels : <EOL> for grid in level . grids : <EOL> for quantity in grid . quantities : <EOL> n_dust_q , shape_q = single_grid_dims ( grid . 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 , amr_grid = None ) : <EOL> '''<STR_LIT>''' <EOL> if amr_grid is None : <EOL> amr_grid = self <EOL> n_pop_ref = None <EOL> for ilevel , level_ref in enumerate ( self . levels ) : <EOL> level = amr_grid . levels [ ilevel ] <EOL> for igrid , grid_ref in enumerate ( level_ref . grids ) : <EOL> grid = level . grids [ igrid ] <EOL> for quantity in grid . quantities : <EOL> n_pop , shape = single_grid_dims ( grid . quantities [ quantity ] ) <EOL> if shape != grid_ref . shape : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> % ( shape , grid_ref . 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 Exception ( "<STR_LIT>" ) <EOL> self . levels = [ ] <EOL> for ilevel in range ( group . attrs [ '<STR_LIT>' ] ) : <EOL> level_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> g_level = group [ level_path ] <EOL> level = self . add_level ( ) <EOL> for igrid in range ( g_level . attrs [ '<STR_LIT>' ] ) : <EOL> grid_path = '<STR_LIT>' % ( igrid + <NUM_LIT:1> ) <EOL> g_grid = g_level [ grid_path ] <EOL> grid = level . add_grid ( ) <EOL> grid . xmin = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . xmax = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . ymin = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . ymax = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . zmin = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . zmax = g_grid . attrs [ '<STR_LIT>' ] <EOL> grid . nx = int ( g_grid . attrs [ '<STR_LIT>' ] ) <EOL> grid . ny = int ( g_grid . attrs [ '<STR_LIT>' ] ) <EOL> grid . nz = int ( g_grid . attrs [ '<STR_LIT>' ] ) <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> for ilevel , level in enumerate ( self . levels ) : <EOL> level_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> for igrid , grid in enumerate ( level . grids ) : <EOL> grid_path = '<STR_LIT>' % ( igrid + <NUM_LIT:1> ) <EOL> g_grid_quantities = group [ level_path ] [ grid_path ] <EOL> for quantity in g_grid_quantities : <EOL> if quantities == '<STR_LIT:all>' or quantity in quantities : <EOL> array = np . array ( g_grid_quantities [ quantity ] ) <EOL> if array . ndim == <NUM_LIT:4> : <EOL> grid . quantities [ quantity ] = [ array [ i ] for i in range ( array . shape [ <NUM_LIT:0> ] ) ] <EOL> else : <EOL> grid . 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>' ] = len ( self . levels ) <EOL> self . _check_array_dimensions ( ) <EOL> for ilevel , level in enumerate ( self . levels ) : <EOL> level_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> g_level = g_geometry . create_group ( level_path ) <EOL> q_level = g_quantities . create_group ( level_path ) <EOL> g_level . attrs [ '<STR_LIT>' ] = len ( level . grids ) <EOL> for igrid , grid in enumerate ( level . grids ) : <EOL> grid_path = '<STR_LIT>' % ( igrid + <NUM_LIT:1> ) <EOL> g_grid = g_level . create_group ( grid_path ) <EOL> q_grid = q_level . create_group ( grid_path ) <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . xmin <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . xmax <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . ymin <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . ymax <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . zmin <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . zmax <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . nx <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . ny <EOL> g_grid . attrs [ '<STR_LIT>' ] = grid . nz <EOL> for quantity in grid . quantities : <EOL> if quantities == '<STR_LIT:all>' or quantity in quantities : <EOL> if isinstance ( grid . quantities [ quantity ] , h5py . ExternalLink ) : <EOL> link_or_copy ( q_grid , quantity , grid . quantities [ quantity ] , copy , absolute_paths = absolute_paths ) <EOL> else : <EOL> q_grid . create_dataset ( quantity , data = grid . quantities [ quantity ] , <EOL> compression = compression , <EOL> dtype = physics_dtype ) <EOL> g_geometry . attrs [ '<STR_LIT>' ] = np . string_ ( self . get_geometry_id ( ) . encode ( '<STR_LIT:utf-8>' ) ) <EOL> def write_single_array ( self , group , name , amr_grid , copy = True , absolute_paths = False , compression = True , physics_dtype = float ) : <EOL> '''<STR_LIT>''' <EOL> if not isinstance ( amr_grid , AMRGridView ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> for ilevel , level in enumerate ( self . levels ) : <EOL> level_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> if level_path in group : <EOL> q_level = group [ level_path ] <EOL> else : <EOL> q_level = group . create_group ( level_path ) <EOL> for igrid , grid in enumerate ( level . grids ) : <EOL> grid_path = '<STR_LIT>' % ( igrid + <NUM_LIT:1> ) <EOL> if grid_path in q_level : <EOL> q_grid = q_level [ grid_path ] <EOL> else : <EOL> q_grid = q_level . create_group ( grid_path ) <EOL> array = amr_grid . levels [ ilevel ] . grids [ igrid ] . quantities [ amr_grid . viewed_quantity ] <EOL> if isinstance ( array , h5py . ExternalLink ) : <EOL> link_or_copy ( q_grid , name , array , copy , absolute_paths = absolute_paths ) <EOL> else : <EOL> q_grid . create_dataset ( name , data = array , <EOL> compression = compression , <EOL> dtype = physics_dtype ) <EOL> def get_geometry_id ( self ) : <EOL> geo_hash = hashlib . md5 ( ) <EOL> for level in self . levels : <EOL> for grid in level . grids : <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . xmin ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . xmax ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . ymin ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . ymax ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . zmin ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . zmax ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . nx ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . ny ) ) <EOL> geo_hash . update ( struct . pack ( '<STR_LIT>' , grid . nz ) ) <EOL> return geo_hash . hexdigest ( ) <EOL> def __getitem__ ( self , item ) : <EOL> return AMRGridView ( self , item ) <EOL> def __setitem__ ( self , item , value ) : <EOL> if isinstance ( value , AMRGridView ) : <EOL> if self . levels == [ ] and value . levels != [ ] : <EOL> logger . warn ( "<STR_LIT>" ) <EOL> for level in value . levels : <EOL> level_ref = self . add_level ( ) <EOL> for grid in level . grids : <EOL> grid_ref = level_ref . add_grid ( ) <EOL> grid_ref . nx = grid . nx <EOL> grid_ref . ny = grid . ny <EOL> grid_ref . nz = grid . nz <EOL> grid_ref . xmin , grid_ref . xmax = grid . xmin , grid . xmax <EOL> grid_ref . ymin , grid_ref . ymax = grid . ymin , grid . ymax <EOL> grid_ref . zmin , grid_ref . zmax = grid . zmin , grid . zmax <EOL> grid_ref . quantities = { } <EOL> for ilevel , level_ref in enumerate ( self . levels ) : <EOL> level = value . levels [ ilevel ] <EOL> for igrid , grid_ref in enumerate ( level_ref . grids ) : <EOL> grid = level . grids [ igrid ] <EOL> grid_ref . quantities [ item ] = deepcopy ( grid . quantities [ value . viewed_quantity ] ) <EOL> elif isinstance ( value , h5py . ExternalLink ) : <EOL> filename = value . filename <EOL> base_path = os . path . dirname ( value . path ) <EOL> array_name = os . path . basename ( value . path ) <EOL> for ilevel , level_ref in enumerate ( self . levels ) : <EOL> level_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> for igrid , grid_ref in enumerate ( level_ref . grids ) : <EOL> grid_path = '<STR_LIT>' % ( ilevel + <NUM_LIT:1> ) <EOL> grid_ref . quantities [ item ] = h5py . ExternalLink ( filename , os . path . join ( base_path , level_path , grid_path , array_name ) ) <EOL> elif value == [ ] : <EOL> for level in self . levels : <EOL> for grid in level . grids : <EOL> grid . quantities [ item ] = [ ] <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def __contains__ ( self , item ) : <EOL> if len ( self . levels ) > <NUM_LIT:0> : <EOL> if len ( self . levels [ <NUM_LIT:0> ] . grids ) > <NUM_LIT:0> : <EOL> return item in self . levels [ <NUM_LIT:0> ] . grids [ <NUM_LIT:0> ] . quantities <EOL> else : <EOL> return False <EOL> else : <EOL> return False <EOL> def reset_quantities ( self ) : <EOL> self . quantities = { } <EOL> for level in self . levels : <EOL> for grid in level . grids : <EOL> grid . quantities = { } <EOL> def add_derived_quantity ( self , name , function ) : <EOL> for level in self . levels : <EOL> for grid in level . grids : <EOL> if name in grid . quantities : <EOL> raise KeyError ( name + '<STR_LIT>' ) <EOL> function ( grid . quantities ) <EOL> def to_yt ( self , dust_id = <NUM_LIT:0> ) : <EOL> '''<STR_LIT>''' <EOL> from yt_wrappers import amr_grid_to_yt_stream <EOL> return amr_grid_to_yt_stream ( self . levels , dust_id ) <EOL> @ classmethod <EOL> def from_yt ( cls , ds , quantity_mapping = { } ) : <EOL> """<STR_LIT>""" <EOL> import yt <EOL> from distutils . version import LooseVersion <EOL> if not LooseVersion ( yt . __version__ ) >= LooseVersion ( '<STR_LIT:3>' ) : <EOL> raise ImportError ( "<STR_LIT>" ) <EOL> from yt_wrappers import yt_dataset_to_amr_grid <EOL> return yt_dataset_to_amr_grid ( ds , quantity_mapping = quantity_mapping ) <EOL> class AMRGridView ( AMRGrid ) : <EOL> def __init__ ( self , amr_grid , quantity ) : <EOL> self . viewed_quantity = quantity <EOL> AMRGrid . __init__ ( self ) <EOL> for level_ref in amr_grid . levels : <EOL> level = self . add_level ( ) <EOL> for grid_ref in level_ref . grids : <EOL> grid = level . add_grid ( ) <EOL> grid . nx = grid_ref . nx <EOL> grid . ny = grid_ref . ny <EOL> grid . nz = grid_ref . nz <EOL> grid . xmin , grid . xmax = grid_ref . xmin , grid_ref . xmax <EOL> grid . ymin , grid . ymax = grid_ref . ymin , grid_ref . ymax <EOL> grid . zmin , grid . zmax = grid_ref . zmin , grid_ref . zmax <EOL> grid . quantities = { } <EOL> grid . quantities [ quantity ] = grid_ref . quantities [ quantity ] <EOL> def append ( self , amr_grid_view ) : <EOL> '''<STR_LIT>''' <EOL> if not isinstance ( amr_grid_view , AMRGridView ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . _check_array_dimensions ( amr_grid_view [ amr_grid_view . viewed_quantity ] ) <EOL> for ilevel , level_ref in enumerate ( self . levels ) : <EOL> level = amr_grid_view . levels [ ilevel ] <EOL> for igrid , grid_ref in enumerate ( level_ref . grids ) : <EOL> grid = level . grids [ igrid ] <EOL> if grid_ref . quantities [ self . viewed_quantity ] is grid . quantities [ amr_grid_view . viewed_quantity ] : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if type ( grid . quantities [ amr_grid_view . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> grid_ref . quantities [ self . viewed_quantity ] . append ( deepcopy ( grid . quantities [ amr_grid_view . viewed_quantity ] ) ) <EOL> def add ( self , amr_grid_view ) : <EOL> '''<STR_LIT>''' <EOL> if not isinstance ( amr_grid_view , AMRGridView ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . _check_array_dimensions ( amr_grid_view [ amr_grid_view . viewed_quantity ] ) <EOL> for ilevel , level_ref in enumerate ( self . levels ) : <EOL> level = amr_grid_view . levels [ ilevel ] <EOL> for igrid , grid_ref in enumerate ( level_ref . grids ) : <EOL> grid = level . grids [ igrid ] <EOL> if type ( grid_ref . quantities [ self . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if type ( grid . quantities [ amr_grid_view . viewed_quantity ] ) is list : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> grid_ref . quantities [ self . viewed_quantity ] += grid . quantities [ amr_grid_view . viewed_quantity ] <EOL> def __getitem__ ( self , item ) : <EOL> if type ( item ) is int : <EOL> amr_grid = AMRGridView ( self , self . viewed_quantity ) <EOL> for level in amr_grid . levels : <EOL> for grid in level . grids : <EOL> grid . quantities = { amr_grid . viewed_quantity : grid . quantities [ amr_grid . viewed_quantity ] [ item ] } <EOL> return amr_grid <EOL> else : <EOL> return AMRGrid . __getitem__ ( self , item ) </s>
