<s> from time import * <EOL> from constants import * <EOL> def get_filename ( left_bound , right_bound ) : <EOL> return str ( int ( left_bound ) ) + '<STR_LIT:_>' + str ( int ( right_bound ) ) + '<STR_LIT>' <EOL> def get_filename_dict ( ID_bounds ) : <EOL> filename_dict = { } <EOL> for ( L , R ) in ID_bounds : <EOL> for i in range ( int ( L ) , int ( R ) + <NUM_LIT:1> ) : <EOL> filename_dict [ i ] = get_filename ( L , R ) <EOL> return filename_dict <EOL> def get_ID_bounds ( halo_file , max_chunk_sz = None , chunk_dir = '<STR_LIT>' ) : <EOL> halos = np . loadtxt ( halo_file ) <EOL> if not max_chunk_sz : max_chunk_sz = halos [ <NUM_LIT:0> ] [ N200a ] <EOL> cur_sum = halos [ <NUM_LIT:0> ] [ N200a ] <EOL> left_bound = halos [ <NUM_LIT:0> ] [ ID ] <EOL> right_bound = halos [ <NUM_LIT:0> ] [ ID ] <EOL> ID_bounds = [ ] <EOL> for h in halos : <EOL> cur_id , num_p = h [ ID ] , h [ N200a ] <EOL> if cur_sum + num_p > max_chunk_sz : <EOL> ID_bounds . append ( ( left_bound , right_bound ) ) <EOL> left_bound = cur_id <EOL> right_bound = cur_id <EOL> cur_sum = num_p <EOL> else : <EOL> cur_sum += num_p <EOL> right_bound += <NUM_LIT:1> <EOL> if ( max_chunk_sz == halos [ <NUM_LIT:0> ] [ N200a ] ) : return ID_bounds [ <NUM_LIT:1> : ] <EOL> else : return ID_bounds <EOL> def make_chunks ( ID_bounds ) : <EOL> i = <NUM_LIT:0> <EOL> ( left_bound , right_bound ) = ID_bounds [ i ] <EOL> cur_chunk = [ ] <EOL> for line in open ( '<STR_LIT>' ) : <EOL> particle = [ float ( val ) for val in line . split ( ) ] <EOL> cur_ID = particle [ <NUM_LIT:1> ] <EOL> if ( cur_ID >= left_bound and cur_ID <= right_bound ) : <EOL> cur_chunk . append ( particle ) <EOL> else : <EOL> os . chdir ( chunk_dir ) <EOL> filename = get_filename ( left_bound , right_bound ) <EOL> print '<STR_LIT>' , len ( cur_chunk ) , '<STR_LIT>' , filename <EOL> print '<STR_LIT>' , <NUM_LIT:100> * ( i + <NUM_LIT:1.> ) / len ( ID_bounds ) , '<STR_LIT:%>' <EOL> assert len ( cur_chunk ) > <NUM_LIT:0> <EOL> np . savetxt ( filename , cur_chunk ) <EOL> os . chdir ( '<STR_LIT:..>' ) <EOL> i += <NUM_LIT:1> <EOL> ( left_bound , right_bound ) = ID_bounds [ i ] <EOL> cur_chunk = [ ] <EOL> if ( cur_ID >= left_bound and cur_ID <= right_bound ) : <EOL> cur_chunk . append ( particle ) <EOL> def fetch ( halo_ID , filename_dict , chunk_dir = '<STR_LIT>' ) : <EOL> P = np . loadtxt ( chunk_dir + '<STR_LIT:/>' + str ( filename_dict [ halo_ID ] ) ) <EOL> P = np . column_stack ( ( P [ : , <NUM_LIT:0> ] , P [ : , <NUM_LIT:2> : ] , P [ : , <NUM_LIT:1> ] ) ) <EOL> return P [ P [ : , H_ID ] == halo_ID ] <EOL> def test ( halo_file , filename_dict ) : <EOL> fetch_times = [ ] <EOL> length_errors = [ ] <EOL> for h in np . loadtxt ( halo_file ) : <EOL> starttime = clock ( ) <EOL> P = fetch ( h [ ID ] , filename_dict ) <EOL> fetch_times . append ( clock ( ) - starttime ) <EOL> if ( len ( P ) != h [ N200a ] ) : length_errors . append ( len ( P ) - int ( h [ N200a ] ) ) <EOL> for p in P : assert p [ H_ID ] == h [ ID ] <EOL> print '<STR_LIT>' , h [ ID ] , '<STR_LIT>' , np . average ( fetch_times ) <EOL> return fetch_times , length_errors <EOL> if ( __name__ == '<STR_LIT:__main__>' ) : <EOL> halo_file = '<STR_LIT>' <EOL> chunk_dir = '<STR_LIT>' <EOL> print '<STR_LIT>' <EOL> ID_bounds = get_ID_bounds ( halo_file ) <EOL> print '<STR_LIT>' , len ( ID_bounds ) <EOL> print '<STR_LIT>' <EOL> filename_dict = get_filename_dict ( ID_bounds ) <EOL> print '<STR_LIT>' <EOL> fetch_times , length_errors = test ( halo_file , filename_dict ) <EOL> print '<STR_LIT>' , len ( length_errors ) , '<STR_LIT>' <EOL> figure ( <NUM_LIT:0> ) <EOL> hist ( fetch_times , bins = <NUM_LIT:50> ) <EOL> xlabel ( '<STR_LIT>' , fontsize = <NUM_LIT> ) <EOL> ylabel ( '<STR_LIT>' , fontsize = <NUM_LIT> ) <EOL> title ( '<STR_LIT>' , fontsize = <NUM_LIT:30> ) <EOL> show ( ) </s>
