<s> import numpy as np <EOL> from patsy import PatsyError <EOL> from patsy . util import ( safe_isnan , safe_scalar_isnan , <EOL> no_pickling , assert_no_pickling ) <EOL> __all__ = [ "<STR_LIT>" ] <EOL> _valid_NA_types = [ "<STR_LIT:None>" , "<STR_LIT>" ] <EOL> _valid_NA_responses = [ "<STR_LIT>" , "<STR_LIT>" ] <EOL> def _desc_options ( options ) : <EOL> return "<STR_LIT:U+002CU+0020>" . join ( [ repr ( opt ) for opt in options ] ) <EOL> class NAAction ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , on_NA = "<STR_LIT>" , NA_types = [ "<STR_LIT:None>" , "<STR_LIT>" ] ) : <EOL> """<STR_LIT>""" <EOL> self . on_NA = on_NA <EOL> if self . on_NA not in _valid_NA_responses : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> % ( on_NA , _desc_options ( _valid_NA_responses ) ) ) <EOL> if isinstance ( NA_types , str ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . NA_types = tuple ( NA_types ) <EOL> for NA_type in self . NA_types : <EOL> if NA_type not in _valid_NA_types : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> % ( NA_type , _desc_options ( _valid_NA_types ) ) ) <EOL> def is_categorical_NA ( self , obj ) : <EOL> """<STR_LIT>""" <EOL> if "<STR_LIT>" in self . NA_types and safe_scalar_isnan ( obj ) : <EOL> return True <EOL> if "<STR_LIT:None>" in self . NA_types and obj is None : <EOL> return True <EOL> return False <EOL> def is_numerical_NA ( self , arr ) : <EOL> """<STR_LIT>""" <EOL> mask = np . zeros ( arr . shape , dtype = bool ) <EOL> if "<STR_LIT>" in self . NA_types : <EOL> mask |= np . isnan ( arr ) <EOL> if mask . ndim > <NUM_LIT:1> : <EOL> mask = np . any ( mask , axis = <NUM_LIT:1> ) <EOL> return mask <EOL> def handle_NA ( self , values , is_NAs , origins ) : <EOL> """<STR_LIT>""" <EOL> assert len ( values ) == len ( is_NAs ) == len ( origins ) <EOL> if len ( values ) == <NUM_LIT:0> : <EOL> return values <EOL> if self . on_NA == "<STR_LIT>" : <EOL> return self . _handle_NA_raise ( values , is_NAs , origins ) <EOL> elif self . on_NA == "<STR_LIT>" : <EOL> return self . _handle_NA_drop ( values , is_NAs , origins ) <EOL> else : <EOL> assert False <EOL> def _handle_NA_raise ( self , values , is_NAs , origins ) : <EOL> for is_NA , origin in zip ( is_NAs , origins ) : <EOL> if np . any ( is_NA ) : <EOL> raise PatsyError ( "<STR_LIT>" , origin ) <EOL> return values <EOL> def _handle_NA_drop ( self , values , is_NAs , origins ) : <EOL> total_mask = np . zeros ( is_NAs [ <NUM_LIT:0> ] . shape [ <NUM_LIT:0> ] , dtype = bool ) <EOL> for is_NA in is_NAs : <EOL> total_mask |= is_NA <EOL> good_mask = ~ total_mask <EOL> return [ v [ good_mask , ... ] for v in values ] <EOL> __getstate__ = no_pickling <EOL> def test_NAAction_basic ( ) : <EOL> from nose . tools import assert_raises <EOL> assert_raises ( ValueError , NAAction , on_NA = "<STR_LIT>" ) <EOL> assert_raises ( ValueError , NAAction , NA_types = ( "<STR_LIT>" , "<STR_LIT>" ) ) <EOL> assert_raises ( ValueError , NAAction , NA_types = "<STR_LIT>" ) <EOL> assert_no_pickling ( NAAction ( ) ) <EOL> def test_NAAction_NA_types_numerical ( ) : <EOL> for NA_types in [ [ ] , [ "<STR_LIT>" ] , [ "<STR_LIT:None>" ] , [ "<STR_LIT>" , "<STR_LIT:None>" ] ] : <EOL> action = NAAction ( NA_types = NA_types ) <EOL> for extra_shape in [ ( ) , ( <NUM_LIT:1> , ) , ( <NUM_LIT:2> , ) ] : <EOL> arr = np . ones ( ( <NUM_LIT:4> , ) + extra_shape , dtype = float ) <EOL> nan_rows = [ <NUM_LIT:0> , <NUM_LIT:2> ] <EOL> if arr . ndim > <NUM_LIT:1> and arr . shape [ <NUM_LIT:1> ] > <NUM_LIT:1> : <EOL> arr [ nan_rows , [ <NUM_LIT:0> , <NUM_LIT:1> ] ] = np . nan <EOL> else : <EOL> arr [ nan_rows ] = np . nan <EOL> exp_NA_mask = np . zeros ( <NUM_LIT:4> , dtype = bool ) <EOL> if "<STR_LIT>" in NA_types : <EOL> exp_NA_mask [ nan_rows ] = True <EOL> got_NA_mask = action . is_numerical_NA ( arr ) <EOL> assert np . array_equal ( got_NA_mask , exp_NA_mask ) <EOL> def test_NAAction_NA_types_categorical ( ) : <EOL> for NA_types in [ [ ] , [ "<STR_LIT>" ] , [ "<STR_LIT:None>" ] , [ "<STR_LIT>" , "<STR_LIT:None>" ] ] : <EOL> action = NAAction ( NA_types = NA_types ) <EOL> assert not action . is_categorical_NA ( "<STR_LIT:a>" ) <EOL> assert not action . is_categorical_NA ( <NUM_LIT:1> ) <EOL> assert action . is_categorical_NA ( None ) == ( "<STR_LIT:None>" in NA_types ) <EOL> assert action . is_categorical_NA ( np . nan ) == ( "<STR_LIT>" in NA_types ) <EOL> def test_NAAction_drop ( ) : <EOL> action = NAAction ( "<STR_LIT>" ) <EOL> in_values = [ np . asarray ( [ - <NUM_LIT:1> , <NUM_LIT:2> , - <NUM_LIT:1> , <NUM_LIT:4> , <NUM_LIT:5> ] ) , <EOL> np . asarray ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) , <EOL> np . asarray ( [ [ <NUM_LIT:1.0> , np . nan ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , np . nan ] ] ) , <EOL> ] <EOL> is_NAs = [ np . asarray ( [ True , False , True , False , False ] ) , <EOL> np . zeros ( <NUM_LIT:5> , dtype = bool ) , <EOL> np . asarray ( [ True , False , False , False , True ] ) , <EOL> ] <EOL> out_values = action . handle_NA ( in_values , is_NAs , [ None ] * <NUM_LIT:3> ) <EOL> assert len ( out_values ) == <NUM_LIT:3> <EOL> assert np . array_equal ( out_values [ <NUM_LIT:0> ] , [ <NUM_LIT:2> , <NUM_LIT:4> ] ) <EOL> assert np . array_equal ( out_values [ <NUM_LIT:1> ] , [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> assert np . array_equal ( out_values [ <NUM_LIT:2> ] , [ [ <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] ] ) <EOL> def test_NAAction_raise ( ) : <EOL> action = NAAction ( on_NA = "<STR_LIT>" ) <EOL> in_arrs = [ np . asarray ( [ <NUM_LIT> , <NUM_LIT> ] ) , <EOL> np . asarray ( [ <NUM_LIT:1> , <NUM_LIT:2> ] ) ] <EOL> is_NAs = [ np . asarray ( [ False , False ] ) ] * <NUM_LIT:2> <EOL> got_arrs = action . handle_NA ( in_arrs , is_NAs , [ None , None ] ) <EOL> assert np . array_equal ( got_arrs [ <NUM_LIT:0> ] , in_arrs [ <NUM_LIT:0> ] ) <EOL> assert np . array_equal ( got_arrs [ <NUM_LIT:1> ] , in_arrs [ <NUM_LIT:1> ] ) <EOL> from patsy . origin import Origin <EOL> o1 = Origin ( "<STR_LIT>" , <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> o2 = Origin ( "<STR_LIT>" , <NUM_LIT:2> , <NUM_LIT:3> ) <EOL> in_idx = np . arange ( <NUM_LIT:2> ) <EOL> in_arrs = [ np . asarray ( [ <NUM_LIT> , <NUM_LIT> ] ) , <EOL> np . asarray ( [ <NUM_LIT:1.0> , np . nan ] ) ] <EOL> is_NAs = [ np . asarray ( [ False , False ] ) , <EOL> np . asarray ( [ False , True ] ) ] <EOL> try : <EOL> action . handle_NA ( in_arrs , is_NAs , [ o1 , o2 ] ) <EOL> assert False <EOL> except PatsyError as e : <EOL> assert e . origin is o2 </s>
