<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import , division , print_function <EOL> try : <EOL> from itertools import izip <EOL> except ImportError : <EOL> izip = zip <EOL> from functools import wraps <EOL> import numpy as np <EOL> from glue . external . six . moves import range as xrange <EOL> def check_empty ( func ) : <EOL> @ wraps ( func ) <EOL> def wrapper ( self , * args , ** kwargs ) : <EOL> if self . empty_slice : <EOL> return np . zeros ( self . shape ) * np . nan <EOL> return func ( self , * args , ** kwargs ) <EOL> return wrapper <EOL> class Aggregate ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , data , attribute , zax , slc , zlim ) : <EOL> """<STR_LIT>""" <EOL> self . data = data <EOL> self . attribute = attribute <EOL> self . zax = zax <EOL> self . slc = slc <EOL> self . zlim = min ( zlim ) , max ( zlim ) <EOL> @ property <EOL> def shape ( self ) : <EOL> """<STR_LIT>""" <EOL> s = self . data . shape <EOL> return s [ self . slc . index ( '<STR_LIT:y>' ) ] , s [ self . slc . index ( '<STR_LIT:x>' ) ] <EOL> @ property <EOL> def empty_slice ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . zlim [ <NUM_LIT:0> ] == self . zlim [ <NUM_LIT:1> ] <EOL> def _subslice ( self ) : <EOL> view = [ slice ( None , None ) for _ in self . data . shape ] <EOL> ax_collapse = self . zax <EOL> for i , s in enumerate ( self . slc ) : <EOL> if s not in [ '<STR_LIT:x>' , '<STR_LIT:y>' ] and i != self . zax : <EOL> view [ i ] = s <EOL> if i < self . zax : <EOL> ax_collapse -= <NUM_LIT:1> <EOL> view [ self . zax ] = slice ( * self . zlim ) <EOL> return view , ax_collapse <EOL> def _prepare_cube ( self , attribute = None ) : <EOL> view , ax_collapse = self . _subslice ( ) <EOL> att = attribute or self . attribute <EOL> cube = self . data [ att , view ] <EOL> return cube , ax_collapse <EOL> def _iter_slice ( self , attribute = None ) : <EOL> view , ax_collapse = self . _subslice ( ) <EOL> att = attribute or self . attribute <EOL> for z in xrange ( * self . zlim ) : <EOL> view [ self . zax ] = z <EOL> plane = self . data [ att , view ] <EOL> yield np . nan_to_num ( plane ) <EOL> def _iter_slice_index ( self ) : <EOL> """<STR_LIT>""" <EOL> att = self . data . get_world_component_id ( self . zax ) <EOL> loop = izip ( self . _iter_slice ( ) , self . _iter_slice ( att ) ) <EOL> return loop <EOL> def _finalize ( self , cube ) : <EOL> if self . slc . index ( '<STR_LIT:x>' ) < self . slc . index ( '<STR_LIT:y>' ) : <EOL> cube = cube . T <EOL> return cube <EOL> def collapse_using ( self , function ) : <EOL> """<STR_LIT>""" <EOL> cube , ax = self . _prepare_cube ( ) <EOL> result = function ( cube , axis = ax ) <EOL> return self . _finalize ( result ) <EOL> def _to_world ( self , idx ) : <EOL> args = [ None ] * self . data . ndim <EOL> y , x = np . mgrid [ : idx . shape [ <NUM_LIT:0> ] , : idx . shape [ <NUM_LIT:1> ] ] <EOL> for i , s in enumerate ( self . slc ) : <EOL> if s not in [ '<STR_LIT:x>' , '<STR_LIT:y>' ] : <EOL> args [ i ] = np . ones ( idx . size ) * s <EOL> args [ self . slc . index ( '<STR_LIT:y>' ) ] = y . ravel ( ) <EOL> args [ self . slc . index ( '<STR_LIT:x>' ) ] = x . ravel ( ) <EOL> args [ self . zax ] = idx . ravel ( ) <EOL> att = self . data . get_world_component_id ( self . zax ) <EOL> return self . data [ att , args ] . reshape ( idx . shape ) <EOL> @ staticmethod <EOL> def all_operators ( ) : <EOL> return ( Aggregate . sum , <EOL> Aggregate . mean , <EOL> Aggregate . max , <EOL> Aggregate . argmax , <EOL> Aggregate . argmin , <EOL> Aggregate . mom1 , <EOL> Aggregate . mom2 , <EOL> Aggregate . median ) <EOL> @ staticmethod <EOL> def _mean ( cube , axis ) : <EOL> s = np . nansum ( cube , axis ) <EOL> ct = np . isfinite ( cube ) . sum ( axis ) <EOL> return <NUM_LIT:1.> * s / ct <EOL> @ check_empty <EOL> def sum ( self ) : <EOL> return self . collapse_using ( np . nansum ) <EOL> @ check_empty <EOL> def mean ( self ) : <EOL> return self . collapse_using ( self . _mean ) <EOL> @ check_empty <EOL> def max ( self ) : <EOL> return self . collapse_using ( np . nanmax ) <EOL> @ check_empty <EOL> def median ( self ) : <EOL> return self . collapse_using ( np . median ) <EOL> @ check_empty <EOL> def argmax ( self ) : <EOL> """<STR_LIT>""" <EOL> idx = self . collapse_using ( np . nanargmax ) <EOL> return self . _to_world ( idx ) <EOL> @ check_empty <EOL> def argmin ( self ) : <EOL> """<STR_LIT>""" <EOL> idx = self . collapse_using ( np . nanargmin ) <EOL> return self . _to_world ( idx ) <EOL> @ check_empty <EOL> def mom1 ( self ) : <EOL> """<STR_LIT>""" <EOL> loop = self . _iter_slice_index ( ) <EOL> val , loc = next ( loop ) <EOL> val = np . maximum ( val , <NUM_LIT:0> ) <EOL> w , result = val , loc * val <EOL> for val , loc in loop : <EOL> val = np . maximum ( val , <NUM_LIT:0> ) <EOL> result += val * loc <EOL> w += val <EOL> return self . _finalize ( result / w ) <EOL> @ check_empty <EOL> def mom2 ( self ) : <EOL> """<STR_LIT>""" <EOL> loop = self . _iter_slice_index ( ) <EOL> val , loc = next ( loop ) <EOL> val = np . maximum ( val , <NUM_LIT:0> ) <EOL> w , x , x2 = val , val * loc , val * loc * loc <EOL> for val , loc in loop : <EOL> val = np . maximum ( val , <NUM_LIT:0> ) <EOL> w += val <EOL> x += loc * val <EOL> x2 += loc ** <NUM_LIT:2> * val <EOL> return self . _finalize ( np . sqrt ( x2 / w - ( x / w ) ** <NUM_LIT:2> ) ) </s>
