update pipeline with new api

Author: antheas

src/pasteur/encode.py

@@ -8,25 +8,25 @@
 from .module import ModuleClass, ModuleFactory
 from .utils import LazyFrame, LazyDataset
 
-ENC = TypeVar("ENC", bound="Encoder")
+ENC = TypeVar("ENC", bound=ModuleClass)
+META = TypeVar("META")
 
 
-class EncoderFactory(ModuleFactory[ENC], Generic[ENC]):
+class AttributeEncoderFactory(ModuleFactory[ENC], Generic[ENC]):
     """Factory base class for encoders. Use isinstance with this class
     to filter the Pasteur module list into only containing Encoders."""
 
     ...
 
 
-META = TypeVar("META")
-
+class EncoderFactory(ModuleFactory[ENC], Generic[ENC]):
+    """Factory base class for encoders. Use isinstance with this class
+    to filter the Pasteur module list into only containing Encoders."""
 
-class Encoder(ModuleClass, Generic[META]):
-    def get_metadata(self) -> META:
-        raise NotImplementedError()
+    ...
 
 
-class AttributeEncoder(Encoder[dict[str | tuple[str], META]], Generic[META]):
+class AttributeEncoder(ModuleClass, Generic[META]):
     """Encapsulates a special way to encode an Attribute.
 
     One encoder is instantiated per attribute and its `fit` function is called to
@@ -50,7 +50,7 @@ class AttributeEncoder(Encoder[dict[str | tuple[str], META]], Generic[META]):
     """
 
     name: str = ""
-    _factory = EncoderFactory["AttributeEncoder"]
+    _factory = AttributeEncoderFactory["AttributeEncoder"]
 
     def fit(self, attr: Attribute, data: pd.DataFrame | None):
         raise NotImplementedError()
@@ -64,10 +64,13 @@ def encode(self, data: pd.DataFrame) -> pd.DataFrame:
     def decode(self, enc: pd.DataFrame) -> pd.DataFrame:
         raise NotImplementedError()
 
+    def get_metadata(self) -> dict[str | tuple[str], META]:
+        raise NotImplementedError()
 
-class ViewEncoder(Encoder[META], Generic[META]):
+
+class Encoder(ModuleClass, Generic[META]):
     name: str = ""
-    _factory = EncoderFactory["ViewEncoder"]
+    _factory = EncoderFactory["Encoder"]
 
     def fit(
         self,
@@ -97,5 +100,10 @@ def decode(
     ]:
         raise NotImplementedError()
 
+    def get_metadata(self) -> META:
+        raise NotImplementedError()
+
+
+ViewEncoder = Encoder
 
 __all__ = ["EncoderFactory", "Encoder", "ViewEncoder", "AttributeEncoder"]

src/pasteur/extras/__init__.py

@@ -35,7 +35,7 @@ def get_recommended_datasets() -> list[Dataset | View]:
 
 def get_recommended_system_modules() -> list[Module]:
     from .encoders import IdxEncoder, NumEncoder
-    from .metrics.distr import ChiSquareMetric, KullbackLeiblerMetric
+    from .metrics.distr import DistributionMetric
     from .metrics.visual import (
         CategoricalHist,
         DateHist,
@@ -83,8 +83,7 @@ def get_recommended_system_modules() -> list[Module]:
         # AimSynth.get_factory(),
         # PrivMrfSynth.get_factory(),
         # Metrics
-        ChiSquareMetric.get_factory(),
-        KullbackLeiblerMetric.get_factory(),
+        DistributionMetric.get_factory(),
         NumericalHist.get_factory(),
         OrdinalHist.get_factory(),
         CategoricalHist.get_factory(),

src/pasteur/extras/metrics/distr.py

@@ -192,7 +192,7 @@ class DistributionMetric(
         Summaries[dict[str, tuple[dict[str, ndarray], dict[tuple[str, str], ndarray]]]],
     ]
 ):
-    name = "cs"
+    name = "dstr"
     encodings = "idx"
 
     def fit(

src/pasteur/extras/synth/pgm/aim.py

@@ -5,11 +5,9 @@
 
 import pandas as pd
 
-from pasteur.attribute import Attribute
-from pasteur.utils import LazyDataset
-
-from ....synth import Synth, data_to_tables, make_deterministic, tables_to_data
-from ....utils import LazyFrame
+from ....attribute import Attributes
+from ....synth import Synth, make_deterministic
+from ....utils import LazyFrame, data_to_tables, tables_to_data
 
 if TYPE_CHECKING:
     from ....attribute import Attributes

src/pasteur/extras/synth/pgm/mst.py

@@ -5,8 +5,8 @@
 
 import pandas as pd
 
-from ....synth import Synth, data_to_tables, make_deterministic, tables_to_data
-from ....utils import LazyFrame
+from ....synth import Synth, make_deterministic
+from ....utils import LazyFrame, data_to_tables, tables_to_data
 
 if TYPE_CHECKING:
     from ....attribute import Attributes

src/pasteur/extras/synth/privbayes/__init__.py

@@ -5,8 +5,8 @@
 from typing import TYPE_CHECKING, Any
 
 from ....marginal import MarginalOracle
-from ....synth import Synth, data_to_tables, make_deterministic, tables_to_data
-from ....utils import LazyFrame
+from ....synth import Synth, make_deterministic
+from ....utils import LazyFrame, data_to_tables, tables_to_data
 
 if TYPE_CHECKING:
     import pandas as pd

src/pasteur/kedro/pipelines/main.py

@@ -12,16 +12,17 @@
 from ...view import View
 from .dataset import create_dataset_pipeline
 from .meta import DatasetMeta, PipelineMeta
-# from .metrics import (
-#     create_metrics_ingest_pipeline,
-#     create_metrics_model_pipeline,
-#     get_metrics_types,
-# )
+
+from .metrics import (
+    create_metrics_ingest_pipeline,
+    create_metrics_model_pipeline,
+    get_metrics_types,
+)
 from .synth import create_synth_pipeline
 from .transform import (
+    create_fit_pipeline,
     create_reverse_pipeline,
     create_transform_pipeline,
-    create_transformer_pipeline,
 )
 from .utils import list_unique
 from .views import (
@@ -97,15 +98,8 @@ def generate_pipelines(
     # Wrk, ref splits are transformed to all types
     # Synthetic data is transformed only to syn_types (as required by metrics currently)
     alg_types = _get_alg_types(algs)
-    # msr_types = get_metrics_types(modules)
-
-    all_types = alg_types# list_unique(alg_types, msr_types)
-    encoders = {
-        k: v
-        for k, v in get_module_dict(EncoderFactory, modules).items()
-        if k in all_types
-    }
-    transformers = get_module_dict(TransformerFactory, modules)
+    msr_types = get_metrics_types(modules)
+    all_types = list_unique(alg_types, msr_types)
 
     wrk_split = WRK_SPLIT
     ref_split = REF_SPLIT
@@ -124,20 +118,20 @@ def generate_pipelines(
         # To make debugging metrics easier, it's bundled with `.measure` pipelines
         # as well. That way, only `.measure` needs to run when changes are made
         # to fit functions
-        # pipe_metrics_fit = create_metrics_ingest_pipeline(
-        #     view, modules, wrk_split, ref_split
-        # )
+        pipe_metrics_fit = create_metrics_ingest_pipeline(
+            view, modules, wrk_split, ref_split
+        )
 
         # Create view transform pipeline that can run as part of ingest
         pipe_transform = (
-            create_transformer_pipeline(view, transformers, encoders, wrk_split)
+            create_fit_pipeline(view, all_types, modules, wrk_split)
             + create_transform_pipeline(
                 view,
                 wrk_split,
                 all_types,
             )
-            # + create_transform_pipeline(view, ref_split, msr_types)
-            # + pipe_metrics_fit
+            + create_transform_pipeline(view, ref_split, msr_types)
+            + pipe_metrics_fit
         )
 
         # Metadata needs to be created every time to allow for overrides
@@ -166,11 +160,11 @@ def generate_pipelines(
                 view, wrk_split, cls
             ) + create_reverse_pipeline(view, alg, cls.type)
 
-            # pipe_measure = create_transform_pipeline(
-            #     view, alg, msr_types, retransform=True
-            # ) + create_metrics_model_pipeline(view, alg, wrk_split, ref_split, modules)
+            pipe_measure = create_transform_pipeline(
+                view, alg, msr_types, retransform=True
+            ) + create_metrics_model_pipeline(view, alg, wrk_split, ref_split, modules)
 
-            complete_pipe = pipe_ds_ingest + pipe_ingest + pipe_synth #+ pipe_measure
+            complete_pipe = pipe_ds_ingest + pipe_ingest + pipe_synth + pipe_measure
 
             if "ident" in alg:
                 # Hide ident pipelines