postprocessing to add the NaN values is added

Author: KoomenErina

main.py

@@ -19,7 +19,10 @@ def synBar():
 
     synth_df = spop.generate(len(df))
 
+    print("synthetische data:")
     print(synth_df.head())
+
+    print("aantal NaNs:")
     print(synth_df.isna().sum())
 
 def synSD2011():
@@ -50,4 +53,4 @@ def synSD2011():
     print(synth_df.head())
 
 
-synBar()
+synSD2011()

synthpop/synthpop.py

@@ -37,7 +37,7 @@ def __init__(self,
         self.numtocat = numtocat
         self.catgroups = catgroups
         self.seed = seed
-
+        self.map_column_to_NaN_column = {}
         # check init
         self.validator.check_init()
 
@@ -54,13 +54,22 @@ def pre_preprocess(self,df,dtypes,nan_fill):
 
             nan_col_name = column+"_NaN"
             df.loc[:,nan_col_name] = maybe_nans
+            self.map_column_to_NaN_column[column] = nan_col_name
 
             dtypes[nan_col_name] = 'category'
 
 
         return df,dtypes
 
     def post_postprocessing(self,syn_df):
+        for column in syn_df:
+
+            if column in self.map_column_to_NaN_column.keys():
+                nan_col_name = self.map_column_to_NaN_column[column]
+                column_NaN_at = syn_df[nan_col_name]
+                syn_df.loc[column_NaN_at,column] = None
+                syn_df = syn_df.drop(columns=nan_col_name)
+
         return syn_df
     def fit(self, df, dtypes=None):
         # TODO check df and check/EXTRACT dtypes
@@ -117,7 +126,7 @@ def generate(self, k=None):
         # postprocess
         processed_synth_df = self.processor.postprocess(synth_df)
 
-        return processed_synth_df
+        return self.post_postprocessing(processed_synth_df)
 
     def _generate(self):
         synth_df = pd.DataFrame(data=np.zeros([self.k, len(self.visit_sequence)]), columns=self.visit_sequence.index)

tests_processing.py

@@ -17,10 +17,17 @@ def test_add_NaN_columns_for_numeric_columns(self):
         self.assertTrue(res['a_NaN'][2])
         self.assertEqual(res['a'][2], -8)
         self.assertEqual(dtype_res['a_NaN'],'category')
+        self.assertEqual(spop.map_column_to_NaN_column['a'],'a_NaN')
 
     def test_apply_and_remove_added_NaN_columns(self):
-        df = pd.DataFrame({'a':[1,2,np.nan],'a_NaN':[False,False,True], 'b':[1,1,1], 'c':['x','y',None]})
+        df = pd.DataFrame({'a':[1,2,-8],'a_NaN':[False,True,False], 'b':[1,1,1], 'c':['x','y',None]})
+
         spop = Synthpop()
+        spop.map_column_to_NaN_column = {'a':'a_NaN'}
+
+        res = spop.post_postprocessing(df)
+        self.assertTrue(np.isnan(res['a'][1]), "NaNs should be placed where indicated")
+        self.assertFalse('a_NaN' in res, "indicator columns should be removed")
 
 
 if __name__ == '__main__':