<s> from __future__ import absolute_import , division , print_function <EOL> import numpy as np <EOL> from datashape import dshape , isnumeric , Record , Option <EOL> from datashape import coretypes as ct <EOL> from toolz import concat , unique <EOL> import xarray as xr <EOL> from . core import Expr <EOL> from . utils import ngjit <EOL> class Preprocess ( Expr ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , column ) : <EOL> self . column = column <EOL> @ property <EOL> def inputs ( self ) : <EOL> return ( self . column , ) <EOL> class extract ( Preprocess ) : <EOL> """<STR_LIT>""" <EOL> def apply ( self , df ) : <EOL> return df [ self . column ] . values <EOL> class category_codes ( Preprocess ) : <EOL> """<STR_LIT>""" <EOL> def apply ( self , df ) : <EOL> return df [ self . column ] . cat . codes . values <EOL> class Reduction ( Expr ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , column ) : <EOL> self . column = column <EOL> def validate ( self , in_dshape ) : <EOL> if not isnumeric ( in_dshape . measure [ self . column ] ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def out_dshape ( self , in_dshape ) : <EOL> return self . _dshape <EOL> @ property <EOL> def inputs ( self ) : <EOL> return ( extract ( self . column ) , ) <EOL> @ property <EOL> def _bases ( self ) : <EOL> return ( self , ) <EOL> @ property <EOL> def _temps ( self ) : <EOL> return ( ) <EOL> def _build_create ( self , dshape ) : <EOL> return self . _create <EOL> def _build_append ( self , dshape ) : <EOL> return self . _append <EOL> def _build_combine ( self , dshape ) : <EOL> return self . _combine <EOL> def _build_finalize ( self , dshape ) : <EOL> return self . _finalize <EOL> class OptionalFieldReduction ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , column = None ) : <EOL> self . column = column <EOL> @ property <EOL> def inputs ( self ) : <EOL> return ( extract ( self . column ) , ) if self . column else ( ) <EOL> def validate ( self , in_dshape ) : <EOL> pass <EOL> def _build_append ( self , dshape ) : <EOL> return self . _append if self . column is None else self . _append_non_na <EOL> @ staticmethod <EOL> def _finalize ( bases , ** kwargs ) : <EOL> return xr . DataArray ( bases [ <NUM_LIT:0> ] , ** kwargs ) <EOL> class count ( OptionalFieldReduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( ct . int32 ) <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg ) : <EOL> agg [ y , x ] += <NUM_LIT:1> <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append_non_na ( x , y , agg , field ) : <EOL> if not np . isnan ( field ) : <EOL> agg [ y , x ] += <NUM_LIT:1> <EOL> @ staticmethod <EOL> def _create ( shape ) : <EOL> return np . zeros ( shape , dtype = '<STR_LIT>' ) <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> return aggs . sum ( axis = <NUM_LIT:0> , dtype = '<STR_LIT>' ) <EOL> class any ( OptionalFieldReduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( ct . bool_ ) <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg ) : <EOL> agg [ y , x ] = True <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append_non_na ( x , y , agg , field ) : <EOL> if not np . isnan ( field ) : <EOL> agg [ y , x ] = True <EOL> @ staticmethod <EOL> def _create ( shape ) : <EOL> return np . zeros ( shape , dtype = '<STR_LIT:bool>' ) <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> return aggs . sum ( axis = <NUM_LIT:0> , dtype = '<STR_LIT:bool>' ) <EOL> class FloatingReduction ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( Option ( ct . float64 ) ) <EOL> @ staticmethod <EOL> def _create ( shape ) : <EOL> return np . full ( shape , np . nan , dtype = '<STR_LIT>' ) <EOL> @ staticmethod <EOL> def _finalize ( bases , ** kwargs ) : <EOL> return xr . DataArray ( bases [ <NUM_LIT:0> ] , ** kwargs ) <EOL> class sum ( FloatingReduction ) : <EOL> """<STR_LIT>""" <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg , field ) : <EOL> if not np . isnan ( field ) : <EOL> if np . isnan ( agg [ y , x ] ) : <EOL> agg [ y , x ] = field <EOL> else : <EOL> agg [ y , x ] += field <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> missing_vals = np . isnan ( aggs ) <EOL> all_empty = np . bitwise_and . reduce ( missing_vals , axis = <NUM_LIT:0> ) <EOL> set_to_zero = missing_vals & ~ all_empty <EOL> return np . where ( set_to_zero , <NUM_LIT:0> , aggs ) . sum ( axis = <NUM_LIT:0> ) <EOL> class m2 ( FloatingReduction ) : <EOL> """<STR_LIT>""" <EOL> @ property <EOL> def _temps ( self ) : <EOL> return ( sum ( self . column ) , count ( self . column ) ) <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , m2 , field , sum , count ) : <EOL> if not np . isnan ( field ) : <EOL> if count == <NUM_LIT:0> : <EOL> m2 [ y , x ] = <NUM_LIT:0> <EOL> else : <EOL> u1 = np . float64 ( sum ) / count <EOL> u = np . float64 ( sum + field ) / ( count + <NUM_LIT:1> ) <EOL> m2 [ y , x ] += ( field - u1 ) * ( field - u ) <EOL> @ staticmethod <EOL> def _combine ( Ms , sums , ns ) : <EOL> mu = np . nansum ( sums , axis = <NUM_LIT:0> ) / ns . sum ( axis = <NUM_LIT:0> ) <EOL> return np . nansum ( Ms + ns * ( sums / ns - mu ) ** <NUM_LIT:2> , axis = <NUM_LIT:0> ) <EOL> class min ( FloatingReduction ) : <EOL> """<STR_LIT>""" <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg , field ) : <EOL> if np . isnan ( agg [ y , x ] ) : <EOL> agg [ y , x ] = field <EOL> elif agg [ y , x ] > field : <EOL> agg [ y , x ] = field <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> return np . nanmin ( aggs , axis = <NUM_LIT:0> ) <EOL> class max ( FloatingReduction ) : <EOL> """<STR_LIT>""" <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg , field ) : <EOL> if np . isnan ( agg [ y , x ] ) : <EOL> agg [ y , x ] = field <EOL> elif agg [ y , x ] < field : <EOL> agg [ y , x ] = field <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> return np . nanmax ( aggs , axis = <NUM_LIT:0> ) <EOL> class count_cat ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> def validate ( self , in_dshape ) : <EOL> if not isinstance ( in_dshape . measure [ self . column ] , ct . Categorical ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def out_dshape ( self , input_dshape ) : <EOL> cats = input_dshape . measure [ self . column ] . categories <EOL> return dshape ( Record ( [ ( c , ct . int32 ) for c in cats ] ) ) <EOL> @ property <EOL> def inputs ( self ) : <EOL> return ( category_codes ( self . column ) , ) <EOL> def _build_create ( self , out_dshape ) : <EOL> n_cats = len ( out_dshape . measure . fields ) <EOL> return lambda shape : np . zeros ( shape + ( n_cats , ) , dtype = '<STR_LIT>' ) <EOL> @ staticmethod <EOL> @ ngjit <EOL> def _append ( x , y , agg , field ) : <EOL> agg [ y , x , field ] += <NUM_LIT:1> <EOL> @ staticmethod <EOL> def _combine ( aggs ) : <EOL> return aggs . sum ( axis = <NUM_LIT:0> , dtype = '<STR_LIT>' ) <EOL> def _build_finalize ( self , dshape ) : <EOL> cats = list ( dshape [ self . column ] . categories ) <EOL> def finalize ( bases , ** kwargs ) : <EOL> dims = kwargs [ '<STR_LIT>' ] + [ self . column ] <EOL> coords = kwargs [ '<STR_LIT>' ] + [ cats ] <EOL> return xr . DataArray ( bases [ <NUM_LIT:0> ] , dims = dims , coords = coords ) <EOL> return finalize <EOL> class mean ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( Option ( ct . float64 ) ) <EOL> @ property <EOL> def _bases ( self ) : <EOL> return ( sum ( self . column ) , count ( self . column ) ) <EOL> @ staticmethod <EOL> def _finalize ( bases , ** kwargs ) : <EOL> sums , counts = bases <EOL> with np . errstate ( divide = '<STR_LIT:ignore>' , invalid = '<STR_LIT:ignore>' ) : <EOL> x = sums / counts <EOL> return xr . DataArray ( x , ** kwargs ) <EOL> class var ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( Option ( ct . float64 ) ) <EOL> @ property <EOL> def _bases ( self ) : <EOL> return ( sum ( self . column ) , count ( self . column ) , m2 ( self . column ) ) <EOL> @ staticmethod <EOL> def _finalize ( bases , ** kwargs ) : <EOL> sums , counts , m2s = bases <EOL> with np . errstate ( divide = '<STR_LIT:ignore>' , invalid = '<STR_LIT:ignore>' ) : <EOL> x = m2s / counts <EOL> return xr . DataArray ( x , ** kwargs ) <EOL> class std ( Reduction ) : <EOL> """<STR_LIT>""" <EOL> _dshape = dshape ( Option ( ct . float64 ) ) <EOL> @ property <EOL> def _bases ( self ) : <EOL> return ( sum ( self . column ) , count ( self . column ) , m2 ( self . column ) ) <EOL> @ staticmethod <EOL> def _finalize ( bases , ** kwargs ) : <EOL> sums , counts , m2s = bases <EOL> with np . errstate ( divide = '<STR_LIT:ignore>' , invalid = '<STR_LIT:ignore>' ) : <EOL> x = np . sqrt ( m2s / counts ) <EOL> return xr . DataArray ( x , ** kwargs ) <EOL> class summary ( Expr ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , ** kwargs ) : <EOL> ks , vs = zip ( * sorted ( kwargs . items ( ) ) ) <EOL> self . keys = ks <EOL> self . values = vs <EOL> def __hash__ ( self ) : <EOL> return hash ( ( type ( self ) , tuple ( self . keys ) , tuple ( self . values ) ) ) <EOL> def validate ( self , input_dshape ) : <EOL> for v in self . values : <EOL> v . validate ( input_dshape ) <EOL> def out_dshape ( self , in_dshape ) : <EOL> return dshape ( Record ( [ ( k , v . out_dshape ( in_dshape ) ) for ( k , v ) <EOL> in zip ( self . keys , self . values ) ] ) ) <EOL> @ property <EOL> def inputs ( self ) : <EOL> return tuple ( unique ( concat ( v . inputs for v in self . values ) ) ) </s>
