Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def imputer_test(self):
def _set_nulls(df):
row = df['index']
if row in [2,5]:
return np.NaN
return df.a
extract_features = ['a']
feature_extractor = FeatureExtractor(input_scalars=['a'],
output_vector='extracted_a_output',
output_vector_items=["{}_out".format(x) for x in extract_features])
imputer = Imputer(strategy='mean')
imputer.mlinit(prior_tf=feature_extractor,
output_features='a_imputed')
df2 = self.df
df2.reset_index(inplace=True)
df2['a'] = df2.apply(_set_nulls, axis=1)
imputer.fit(df2[['a']])
self.assertAlmostEqual(imputer.statistics_[0], df2.a.mean(), places = 7)
imputer.serialize_to_bundle(self.tmp_dir, imputer.name)
expected_model = {
"op": "imputer",
"attributes": {