332 executable smartnoise

PETsARD/synthesizer/smartnoise.py

@@ -1,8 +1,7 @@
 import time
 
 import pandas as pd
-from snsynth.transform import NoTransformer, TableTransformer
-from snsynth.transform.identity import IdentityTransformer
+from snsynth.transform import TableTransformer, MinMaxTransformer
 from snsynth import Synthesizer as SNSyn
 
 from PETsARD.error import UnfittedError, UnsupportedMethodError
@@ -17,6 +16,9 @@ class SmartNoise:
     as well as common functionality.
     """
 
+    CUBE = ['aim', 'mwem', 'mst', 'pacsynth']
+    GAN = ['dpctgan', 'patectgan']
+
     def __init__(self, data: pd.DataFrame, **kwargs) -> None:
         """
         Args:
@@ -37,14 +39,23 @@ def fit(self) -> None:
                 f"Synthesizer (SmartNoise): Fitting {self.syn_method}."
             )
 
-            # TODO - for mst and pacsynth,
-            # we need to implement other transformer rather than IdentityTransformer
-            t = TableTransformer([IdentityTransformer() 
-                                  for i in range(self.data.shape[1])])
-
             # TODO - Only support cube-style synthesizer. 
             # GAN-style synthesizer needed to be implemented.
-            self._Synthesizer.fit(self.data, transformer=t)
+
+            if self.syn_method in self.CUBE:
+                self._Synthesizer.fit(
+                    self.data,
+                    categorical_columns=self.data.columns
+                )
+            else:
+                tt = TableTransformer([
+                    MinMaxTransformer(lower=self.data[col].min(),
+                                      upper=self.data[col].max(),
+                                      negative=False) 
+                    for col in self.data.columns
+                ])
+
+                self._Synthesizer.fit(self.data, transformer=tt)
             print(
                 f"Synthesizer (SmartNoise): "
                 f"Fitting  {self.syn_method} spent "