fix end-to-end multitable process

Author: antheas

src/pasteur/attribute.py

@@ -322,7 +322,7 @@ class GenerationValue(StratifiedValue):
     def __init__(self, table: str, max_len: int) -> None:
         self.table = table
         self.max_len = max_len
-        super().__init__(Grouping('ord', list(range(max_len))), 0)
+        super().__init__(Grouping('ord', list(range(max_len + 1))), 0)
 
 def _create_strat_value_cat(vals, na: bool = False, ukn_val: Any | None = None):
     arr = []

src/pasteur/extras/metrics/distr.py

@@ -48,6 +48,10 @@ def calc_marginal_1way(
         mul *= domain[col]
 
     counts = np.bincount(idx, minlength=x_dom)
+    assert (
+        len(counts) == x_dom
+    ), f"Overflow error, domain for columns `{x}` is wrong or there is a mistake in encoding."
+
     margin = counts.astype("float")
     margin /= margin.sum()
     if zero_fill is not None:
@@ -59,7 +63,9 @@ def calc_marginal_1way(
 
 
 def _visualise_cs(
-    name: str, domain: dict[str, int], data: dict[str, Summaries[dict[str, np.ndarray]]]
+    table: str,
+    domain: dict[str, int],
+    data: dict[str, Summaries[dict[str, np.ndarray]]],
 ):
     import mlflow
 
@@ -104,12 +110,12 @@ def _visualise_cs(
         split_ref="ref",
     )
 
-    fn = f"distr/cs.html" if name == "table" else f"distr/{name}_cs.html"
+    fn = f"distr/cs.html" if table == "table" else f"distr/{table}_cs.html"
     mlflow.log_text(gen_html_table(style, FONT_SIZE), fn)
 
 
 def _visualise_kl(
-    name: str, data: dict[str, Summaries[dict[tuple[str, str], np.ndarray]]]
+    table: str, data: dict[str, Summaries[dict[tuple[str, str], np.ndarray]]]
 ):
     import mlflow
 
@@ -159,7 +165,7 @@ def _visualise_kl(
         split_ref="ref",
     )
 
-    fn = f"distr/kl.html" if name == "table" else f"distr/{name}_kl.html"
+    fn = f"distr/kl.html" if table == "table" else f"distr/{table}_kl.html"
     mlflow.log_text(gen_html_table(style, FONT_SIZE), fn)
 
 

src/pasteur/extras/transformers.py

@@ -403,7 +403,7 @@ def reverse(self, data: pd.DataFrame, ref: pd.Series | None = None) -> pd.Series
             if ref is not None:
                 na_mask = pd.isna(ref) | na_mask
                 ref = ref[~na_mask]
-            vals = vals[~na_mask]
+            vals = vals[~na_mask.reindex(vals.index)]
             ofs = 1
         else:
             ofs = 0

src/pasteur/extras/views/mimic/parameters_core.yml

@@ -24,11 +24,16 @@ tables:
       admittime:
         type: seq
         ref: patients.birth_year
-        seq:
+        ctx:
           type: datetime
           span: year.hour
           max_len: 99
           bins: 32
+        seq:
+          type: datetime
+          span: year.hour
+          max_len: 10
+          bins: 10
       dischtime: datetime|day.hour:admittime
       deathtime: datetime?|day.hour:admittime
       admission_type: categorical

src/pasteur/kedro/pipelines/transform.py

@@ -129,18 +129,20 @@ def create_fit_pipeline(
             outputs=f"{view}.enc.{enc}",
             namespace=f"{view}.enc",
         )
-        for enc in encs if enc not in ('raw', 'bst')
+        for enc in encs
+        if enc not in ("raw", "bst")
     ]
 
     return PipelineMeta(
         pipeline(trn_fit_nodes + enc_fit_nodes, tags=TAGS_TRANSFORM),
         [
-            D("transformers", f"{view}.trn.{t}", ["view", view, "trn",  t], type="pkl")
+            D("transformers", f"{view}.trn.{t}", ["view", view, "trn", t], type="pkl")
             for t in view.tables
         ]
         + [
-            D("encoders", f"{view}.enc.{enc}", ["view", view, 'enc', enc], type="pkl")
-            for enc in encs if enc not in ('raw', 'bst')
+            D("encoders", f"{view}.enc.{enc}", ["view", view, "enc", enc], type="pkl")
+            for enc in encs
+            if enc not in ("raw", "bst")
         ],
     )
 
@@ -178,7 +180,7 @@ def create_transform_pipeline(
                     "split_transformed",
                     f"{view}.{split}.ctx_{t}",
                     ["view", view, split, "ctx", t],
-                    type='multi'
+                    type="multi",
                 )
             )
             outputs.append(
@@ -221,7 +223,7 @@ def create_transform_pipeline(
                 f"{view}.{split}.{enc}",
                 ["synth" if retransform else "view", view, split, enc],
                 versioned=retransform,
-                type='multi'
+                type="multi",
             )
         )
 
@@ -262,7 +264,9 @@ def create_reverse_pipeline(view: View, alg: str, enc: str):
                     "table": f"{view}.{alg}.bst_{t}",
                     "ctx": f"{view}.{alg}.ctx_{t}",
                     "ids": f"{view}.{alg}.ids_{t}",
-                    "parents": {req: req for req in view.trn_deps.get(t, [])},
+                    "parents": {
+                        req: f"{view}.{alg}.{req}" for req in view.trn_deps.get(t, [])
+                    },
                 },
                 outputs=f"{view}.{alg}.{t}",
                 namespace=f"{view}.{alg}",
@@ -293,7 +297,7 @@ def create_reverse_pipeline(view: View, alg: str, enc: str):
                 D(
                     "synth_reversed",
                     f"{view}.{alg}.{t}",
-                    ["synth", view, alg, 'tables', t],
+                    ["synth", view, alg, "tables", t],
                     versioned=True,
                 ),
             ]

src/pasteur/metric.py

@@ -218,9 +218,17 @@ def _fit_column_metrics(
     metrics: dict[str, list[ColumnMetricFactory]],
 ):
     get_table = lazy_load_tables(tables)
+    table = get_table(name)
 
     if ref.table_has_reference():
         ids = ref.find_foreign_ids(name, get_table)
+
+        if len(table.index.symmetric_difference(ids.index)):
+            old_len = len(table)
+            table = table.reindex(ids.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table)}/{old_len} rows with missing ids."
+            )
     else:
         ids = None
 
@@ -237,22 +245,22 @@ def _fit_column_metrics(
                 m = factory.build(**col.args)
 
             if isinstance(m, ColumnMetric):
-                m.fit(name, col_name, col, get_table(name)[col_name])
+                m.fit(name, col_name, col, table[col_name])
             elif isinstance(m, RefColumnMetric):
-                ref_col = _calc_joined_refs(name, get_table, ids, col.ref)
+                ref_col = _calc_joined_refs(name, get_table, ids, col.ref, table)
                 m.fit(
                     name,
                     col_name,
                     col,
-                    RefColumnData(data=get_table(name)[col_name], ref=ref_col),
+                    RefColumnData(data=table[col_name], ref=ref_col),
                 )
             elif isinstance(m, SeqColumnMetric):
-                ref_col = _calc_unjoined_refs(name, get_table, col.ref)
+                ref_col = _calc_unjoined_refs(name, get_table, col.ref, table)
                 m.fit(
                     name,
                     col_name,
                     col,
-                    SeqColumnData(data=get_table(name)[col_name], ref=ref_col, ids=ids),
+                    SeqColumnData(data=table[col_name], ref=ref_col, ids=ids),
                 )
             else:
                 assert False, f"Unknown column metric type: {type(m)}"
@@ -272,10 +280,25 @@ def _preprocess_metrics(
 ):
     get_table_wrk = lazy_load_tables(tables_wrk)
     get_table_ref = lazy_load_tables(tables_ref)
+    table_wrk = get_table_wrk(name)
+    table_ref = get_table_ref(name)
 
     if ref.table_has_reference():
         ids_wrk = ref.find_foreign_ids(name, get_table_wrk)
         ids_ref = ref.find_foreign_ids(name, get_table_ref)
+
+        if len(table_wrk.index.symmetric_difference(ids_wrk.index)):
+            old_len = len(table_wrk)
+            table_wrk = table_wrk.reindex(ids_wrk.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table_wrk)}/{old_len} rows with missing ids."
+            )
+        if len(table_ref.index.symmetric_difference(ids_ref.index)):
+            old_len = len(table_ref)
+            table_ref = table_ref.reindex(ids_ref.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table_ref)}/{old_len} rows with missing ids."
+            )
     else:
         ids_wrk = None
         ids_ref = None
@@ -286,30 +309,38 @@ def _preprocess_metrics(
             col = meta[name][col_name]
             if isinstance(m, ColumnMetric):
                 prec = m.preprocess(
-                    get_table_wrk(name)[col_name],
-                    get_table_ref(name)[col_name],
+                    table_wrk[col_name],
+                    table_ref[col_name],
                 )
             elif isinstance(m, RefColumnMetric):
                 prec = m.preprocess(
                     RefColumnData(
-                        data=get_table_wrk(name)[col_name],
-                        ref=_calc_joined_refs(name, get_table_wrk, ids_ref, col.ref),
+                        data=table_wrk[col_name],
+                        ref=_calc_joined_refs(
+                            name, get_table_wrk, ids_ref, col.ref, table_wrk
+                        ),
                     ),
                     RefColumnData(
-                        data=get_table_ref(name)[col_name],
-                        ref=_calc_joined_refs(name, get_table_ref, ids_ref, col.ref),
+                        data=table_ref[col_name],
+                        ref=_calc_joined_refs(
+                            name, get_table_ref, ids_ref, col.ref, table_ref
+                        ),
                     ),
                 )
             elif isinstance(m, SeqColumnMetric):
                 prec = m.preprocess(
                     SeqColumnData(
-                        data=get_table_wrk(name)[col_name],
-                        ref=_calc_unjoined_refs(name, get_table_wrk, col.ref),
+                        data=table_wrk[col_name],
+                        ref=_calc_unjoined_refs(
+                            name, get_table_wrk, col.ref, table_wrk
+                        ),
                         ids=ids_wrk,
                     ),
                     SeqColumnData(
-                        data=get_table_ref(name)[col_name],
-                        ref=_calc_unjoined_refs(name, get_table_ref, col.ref),
+                        data=table_ref[col_name],
+                        ref=_calc_unjoined_refs(
+                            name, get_table_ref, col.ref, table_ref
+                        ),
                         ids=ids_ref,
                     ),
                 )
@@ -334,11 +365,33 @@ def _process_metrics(
     get_table_wrk = lazy_load_tables(tables_wrk)
     get_table_ref = lazy_load_tables(tables_ref)
     get_table_syn = lazy_load_tables(tables_syn)
+    table_wrk = get_table_wrk(name)
+    table_ref = get_table_ref(name)
+    table_syn = get_table_syn(name)
 
     if ref.table_has_reference():
         ids_wrk = ref.find_foreign_ids(name, get_table_wrk)
         ids_ref = ref.find_foreign_ids(name, get_table_ref)
         ids_syn = ref.find_foreign_ids(name, get_table_syn)
+
+        if len(table_wrk.index.symmetric_difference(ids_wrk.index)):
+            old_len = len(table_wrk)
+            table_wrk = table_wrk.reindex(ids_wrk.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table_wrk)}/{old_len} rows with missing ids."
+            )
+        if len(table_ref.index.symmetric_difference(ids_ref.index)):
+            old_len = len(table_ref)
+            table_ref = table_ref.reindex(ids_ref.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table_ref)}/{old_len} rows with missing ids."
+            )
+        if len(table_syn.index.symmetric_difference(ids_syn.index)):
+            old_len = len(table_syn)
+            table_syn = table_syn.reindex(ids_syn.index)
+            logger.warn(
+                f"There are missing ids for rows in {name}, dropping {old_len-len(table_syn)}/{old_len} rows with missing ids."
+            )
     else:
         ids_wrk = None
         ids_ref = None
@@ -358,34 +411,46 @@ def _process_metrics(
             elif isinstance(m, RefColumnMetric):
                 proc = m.process(
                     RefColumnData(
-                        data=get_table_wrk(name)[col_name],
-                        ref=_calc_joined_refs(name, get_table_wrk, ids_wrk, col.ref),
+                        data=table_wrk[col_name],
+                        ref=_calc_joined_refs(
+                            name, get_table_wrk, ids_wrk, col.ref, table_wrk
+                        ),
                     ),
                     RefColumnData(
-                        data=get_table_ref(name)[col_name],
-                        ref=_calc_joined_refs(name, get_table_ref, ids_ref, col.ref),
+                        data=table_ref[col_name],
+                        ref=_calc_joined_refs(
+                            name, get_table_ref, ids_ref, col.ref, table_ref
+                        ),
                     ),
                     RefColumnData(
-                        data=get_table_syn(name)[col_name],
-                        ref=_calc_joined_refs(name, get_table_syn, ids_syn, col.ref),
+                        data=table_syn[col_name],
+                        ref=_calc_joined_refs(
+                            name, get_table_syn, ids_syn, col.ref, table_syn
+                        ),
                     ),
                     prec,
                 )
             elif isinstance(m, SeqColumnMetric):
                 proc = m.process(
                     SeqColumnData(
-                        data=get_table_wrk(name)[col_name],
-                        ref=_calc_unjoined_refs(name, get_table_wrk, col.ref),
+                        data=table_wrk[col_name],
+                        ref=_calc_unjoined_refs(
+                            name, get_table_wrk, col.ref, table_wrk
+                        ),
                         ids=ids_wrk,
                     ),
                     SeqColumnData(
-                        data=get_table_ref(name)[col_name],
-                        ref=_calc_unjoined_refs(name, get_table_ref, col.ref),
+                        data=table_ref[col_name],
+                        ref=_calc_unjoined_refs(
+                            name, get_table_ref, col.ref, table_ref
+                        ),
                         ids=ids_ref,
                     ),
                     SeqColumnData(
-                        data=get_table_syn(name)[col_name],
-                        ref=_calc_unjoined_refs(name, get_table_syn, col.ref),
+                        data=table_syn[col_name],
+                        ref=_calc_unjoined_refs(
+                            name, get_table_syn, col.ref, table_syn
+                        ),
                         ids=ids_syn,
                     ),
                     prec,

src/pasteur/synth.py

@@ -8,6 +8,8 @@
 from functools import partial, wraps
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
+from pasteur.utils import LazyDataset
+
 from .encode import ViewEncoder
 from .metadata import Metadata
 from .module import ModuleClass, ModuleFactory
@@ -147,7 +149,7 @@ def synth_fit(
 
     tracker = PerformanceTracker.get("synth")
 
-    tracker.ensemble("total", "preprocess", "bake", "fit", "sample")
+    tracker.ensemble("total", "preprocess", "bake", "fit")
 
     meta = encoder.get_metadata()
     args = {**metadata.algs.get(factory.name, {}), **metadata.alg_override}
@@ -185,10 +187,10 @@ class IdentSynth(Synth):
     def preprocess(self, meta: Any, data: dict[str, LazyDataset]):
         pass
 
-    def bake(self, meta: Any, data: dict[str, LazyDataset]):
+    def bake(self, data: dict[str, LazyDataset]):
         pass
 
-    def fit(self, meta: Any, data: dict[str, LazyDataset]):
+    def fit(self, data: dict[str, LazyDataset]):
         self.data = data
 
     def sample(self, n: int | None = None):