minor refactor + docs for main modules

Author: antheas

src/pasteur/cli.py

@@ -1,4 +1,4 @@
-from .kedro.utils import get_pasteur_modules
+# from .kedro.utils import get_pasteur_modules
 
 # if get_pasteur_modules():
 from .kedro.cli import cli

src/pasteur/dataset.py

@@ -1,3 +1,6 @@
+""" This module holds the definitions for the Dataset module, the initial entrypoint
+for data in Pasteur. """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Callable
@@ -99,7 +102,7 @@ def ingest(self, name, **tables: Any) -> LazyFrame:
 
         @warning: all partitioned tables should have the same partitions.
         Some tables may not be partitioned.
-        
+
         Tip: use a `match` statement to fork based on table name to per-table functions."""
         raise NotImplemented()
 

src/pasteur/encode.py

@@ -5,10 +5,24 @@
 
 
 class EncoderFactory(ModuleFactory["Encoder"]):
+    """ Factory base class for encoders. Use isinstance with this class
+    to filter the Pasteur module list into only containing Encoders. """
     ...
 
+
 class Encoder(ModuleClass):
-    """Encapsulates a special way to encode an Attribute."""
+    """Encapsulates a special way to encode an Attribute.
+    
+    One encoder is instantiated per module and its `fit` function is called to
+    fit it to the base layer data.
+    
+    After that, the module may be serialized, unserialized, and its encode
+    and decode methods may be called arbitrarily from different processes to encode
+    and decode sets of columns.
+    
+    The `data` value may contain a superset of columns than that of the encoder.
+    It is up to the encoder to filter it prior to processing. `data` should
+    not be mutated."""
 
     name: str
     attr: Attribute
@@ -21,4 +35,7 @@ def encode(self, data: pd.DataFrame) -> pd.DataFrame:
         raise NotImplementedError()
 
     def decode(self, enc: pd.DataFrame) -> pd.DataFrame:
-        raise NotImplementedError()
+        raise NotImplementedError()
+
+
+__all__ = ["EncoderFactory", "Encoder"]

src/pasteur/hierarchy.py

@@ -1,12 +1,16 @@
+""" Highly experimental and unpublished class for rebalancing Stratified Values
+with Differential Privacy.
+
+@TODO: Documentation."""
+
 import logging
 from itertools import combinations
 from math import ceil, log
 from typing import TypeVar
 
 import numpy as np
-import pandas as pd
 
-from .attribute import Attributes, IdxValue, Level, LevelValue, get_dtype
+from .attribute import Attributes, CatValue, Grouping, StratifiedValue, get_dtype
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +26,7 @@ class CatNode(list):
 
 
 def create_tree(
-    node: Level, common: int = 0, ofs: int = 0, n: int | None = None
+    node: Grouping, common: int = 0, ofs: int = 0, n: int | None = None
 ) -> list:
     """Receives the top node of the tree of a hierarchical attribute and
     converts it into the same tree structure, where the leaves have been
@@ -37,7 +41,7 @@ def create_tree(
     for child in node:
         if ofs < common:
             out.append(None)
-        elif isinstance(child, Level):
+        elif isinstance(child, Grouping):
             out.append(create_tree(child, common, ofs, n))
         else:
             out.append(set([ofs]))
@@ -220,7 +224,7 @@ def create_node_to_group_map(tree: list, grouping: np.ndarray, ofs: int = 0):
     return ofs
 
 
-def make_grouping(counts: np.ndarray, head: Level, common: int = 0) -> np.ndarray:
+def make_grouping(counts: np.ndarray, head: Grouping, common: int = 0) -> np.ndarray:
     """Converts the hierarchical attribute level tree provided into a node-to-group
     mapping, where `group[i][j] = z`, where `i` is the height of the mapping
     `j` is node `j` and `z` is the group the node is associated at that height.
@@ -305,11 +309,11 @@ def generate_domain_list(
     return new_domains
 
 
-class RebalancedValue(IdxValue):
+class RebalancedValue(CatValue):
     def __init__(
         self,
         counts: np.ndarray,
-        col: LevelValue,
+        col: StratifiedValue,
         reshape_domain: bool = True,
         u: float = 1.3,
         fixed: list[int] = [2, 4, 5, 8, 12],
@@ -406,7 +410,7 @@ def upsample(self, column: np.ndarray, height: int, deterministic: bool = True):
 
 def rebalance_value(
     counts: np.ndarray,
-    col: LevelValue,
+    col: StratifiedValue,
     num_cols: int = 1,
     ep: float | None = None,
     gaussian: bool = False,
@@ -422,7 +426,7 @@ def rebalance_value(
         noise = np.random.laplace(scale=noise_scale, size=counts.shape)
         counts = counts + noise
 
-    assert isinstance(col, LevelValue)
+    assert isinstance(col, StratifiedValue)
     return RebalancedValue(counts, col, **kwargs)
 
 
@@ -449,7 +453,7 @@ def rebalance_attributes(
     for name, attr in attrs.items():
         cols = {}
         for col_name, col in attr.vals.items():
-            assert isinstance(col, LevelValue)
+            assert isinstance(col, StratifiedValue)
             cols[col_name] = rebalance_value(
                 counts[col_name],
                 col,

src/pasteur/metadata.py

@@ -1,10 +1,14 @@
+""" This module contains a base class `Metadata` which is used to wrap, type, 
+and check all View parameters provided to kedro.
+
+@TODO: refactor this file. """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, NamedTuple, overload
 
 if TYPE_CHECKING:
     import pandas as pd
-    from .utils import LazyFrame
 
 import logging
 
@@ -173,7 +177,7 @@ def __str__(self) -> str:
         return self.__dict__.__str__()
 
 
-class DatasetMeta:
+class ViewMeta:
     TABLE_CLS = TableMeta
 
     def __init__(
@@ -226,5 +230,5 @@ def __str__(self) -> str:
         return self.__dict__.__str__()
 
 
-class Metadata(DatasetMeta):
+class Metadata(ViewMeta):
     pass

src/pasteur/metric.py

@@ -1,6 +1,10 @@
+""" This module provides the definitions for Metric Modules.
+Metric modules can fit to a column, a table, or a whole View.
+In each case, modules are instanciated as required (for columns one is instantiated
+per column type, for tables one per table and View metrics are instantiated once)."""
+
 import logging
-from collections import defaultdict
-from typing import Generic, TypeVar, cast, TypedDict, NamedTuple, Any
+from typing import Generic, TypeVar, TypedDict, Any
 
 import pandas as pd
 
@@ -9,7 +13,7 @@
 from .module import ModuleClass, ModuleFactory
 from .table import TransformHolder
 from .utils import LazyChunk, LazyFrame
-from .utils.progress import process, process_in_parallel
+from .utils.progress import process_in_parallel
 
 logger = logging.getLogger(__name__)
 
@@ -44,9 +48,9 @@ def __init__(
         self.encodings = cls.encodings
 
 
-class DatasetMetricFactory(ModuleFactory["DatasetMetric"]):
+class ViewMetricFactory(ModuleFactory["ViewMetric"]):
     def __init__(
-        self, cls: type["DatasetMetric"], *args, name: str | None = None, **kwargs
+        self, cls: type["ViewMetric"], *args, name: str | None = None, **kwargs
     ) -> None:
         super().__init__(cls, *args, name=name, **kwargs)
         self.encodings = cls.encodings
@@ -228,8 +232,8 @@ def fit(self, table: str, meta: Metadata, data: ColumnData):
         ids = data["ids"]
         tables = data["tables"].copy()
         tables["ids"] = ids
-        part = next(iter(LazyFrame.zip_values(**tables))) # FIXME: incorrect type
-        self._fit_chunk(table, meta, part, part["ids"]) #type: ignore
+        part = next(iter(LazyFrame.zip_values(**tables)))  # FIXME: incorrect type
+        self._fit_chunk(table, meta, part, part["ids"])  # type: ignore
 
     def _process_chunk(
         self,
@@ -276,12 +280,12 @@ def preprocess(self, wrk: ColumnData, ref: ColumnData) -> Summaries:
             wrk_sum[name] = []
             ref_sum[name] = []
             for i, metric in enumerate(metrics):
-                wrk_sum[name].append(metric.combine(
-                    [chunk[name][i] for chunk in summaries_wrk]
-                ))
-                ref_sum[name].append(metric.combine(
-                    [chunk[name][i] for chunk in summaries_ref]
-                ))
+                wrk_sum[name].append(
+                    metric.combine([chunk[name][i] for chunk in summaries_wrk])
+                )
+                ref_sum[name].append(
+                    metric.combine([chunk[name][i] for chunk in summaries_ref])
+                )
 
         return Summaries(wrk_sum, ref_sum)
 
@@ -299,9 +303,9 @@ def process(
         for name, metrics in self.metrics.items():
             syn_sum[name] = []
             for i, metric in enumerate(metrics):
-                syn_sum[name].append(metric.combine(
-                    [chunk[name][i] for chunk in summaries]
-                ))
+                syn_sum[name].append(
+                    metric.combine([chunk[name][i] for chunk in summaries])
+                )
 
         return pre.replace(syn=syn_sum)
 
@@ -357,19 +361,19 @@ def unique_name(self) -> str:
         return f"{self.type}_{self.name}_{self.table}"
 
 
-class DatasetData(TypedDict):
+class ViewData(TypedDict):
     tables: dict[str, dict[str, LazyFrame]]
     ids: dict[str, LazyFrame]
 
 
-class DatasetMetric(Metric[DatasetData, _INGEST, _SUMMARY], Generic[_INGEST, _SUMMARY]):
-    _factory = DatasetMetricFactory
+class ViewMetric(Metric[ViewData, _INGEST, _SUMMARY], Generic[_INGEST, _SUMMARY]):
+    _factory = ViewMetricFactory
     type = "dst"
     table: str
     encodings: list[str] = ["raw"]
 
     def fit(
-        self, meta: Metadata, attrs: dict[str, dict[str, Attributes]], data: DatasetData
+        self, meta: Metadata, attrs: dict[str, dict[str, Attributes]], data: ViewData
     ):
         raise NotImplementedError()
 
@@ -407,10 +411,10 @@ def fit_table_metric(
 
 
 def fit_dataset_metric(
-    fs: DatasetMetricFactory,
+    fs: ViewMetricFactory,
     meta: Metadata,
     trns: dict[str, TransformHolder],
-    data: DatasetData,
+    data: ViewData,
 ):
     enc = fs.encodings
     attrs = {
@@ -432,3 +436,20 @@ def log_metric(metric: Metric[Any, Any, _SUMMARY], summary: _SUMMARY):
     mlflow_log_artifacts(
         "metrics", metric.unique_name(), metric=metric, summary=summary
     )
+
+
+DatasetMetric = ViewMetric
+DatasetMetricFactory = ViewMetricFactory
+
+__all__ = [
+    "ColumnMetricFactory",
+    "RefColumnMetricFactory",
+    "TableMetricFactory",
+    "ViewMetricFactory",
+    "Metric",
+    "Summaries",
+    "ColumnMetric",
+    "RefColumnMetric",
+    "TableMetric",
+    "ViewMetric",
+]

src/pasteur/module.py

@@ -1,3 +1,9 @@
+""" Contains the module definitions in Pasteur, the base classes all
+Pasteur modules extend from. 
+
+In Pasteur, all functionality is achieved through the use of modules.
+You should not interact with this module directly, but rather through its children."""
+
 from collections import defaultdict
 from typing import Generic, TypeVar
 
@@ -17,7 +23,11 @@ class Module:
 
 class ModuleClass:
     """Modules which need to be instantiated multiple times extend from ModuleClass and define
-    a Factory to act as their module"""
+    a Factory to act as their module.
+
+    For the module types provided by pasteur, you can call the classmethod `get_factory()`.
+    `get_factory()` also acts as a closure, allowing you to provide parameters to
+    the module's init function."""
 
     name: str
     _factory: type["ModuleFactory"]
@@ -43,7 +53,7 @@ class ModuleFactory(Module, Generic[A]):
     """Some modules (such as transformers) require multiple instances in the system. In this case,
     it's not possible to provide a module instance for them.
 
-    `ModuleFactory` is used to provide a wrapper instance to that module class."""
+    For those types, their instance is based on `ModuleFactory`."""
 
     def __init__(self, cls: type[A], *args, name: str | None = None, **kwargs) -> None:
         self._cls = cls
@@ -82,3 +92,12 @@ def get_module_dict_multiple(
         if isinstance(module, parent):
             out[module.name].append(module)
     return out
+
+
+__all__ = [
+    "Module",
+    "ModuleClass",
+    "ModuleFactory",
+    "get_module_dict",
+    "get_module_dict_multiple",
+]