diff --git a/pyproject.toml b/pyproject.toml index 0a0fef9d..b9f06a4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ requires-python = ">=3.11" dependencies = [ "scikit-learn==1.8", "pandas==3.0.2", - "skrub==0.8.0", + "skrub==0.11.0", "polars", "graphviz", "pygments>=2.20.0", diff --git a/stratum/_api.py b/stratum/_api.py index 79f0ec84..4264455b 100644 --- a/stratum/_api.py +++ b/stratum/_api.py @@ -1,10 +1,11 @@ import pandas as pd from dataclasses import dataclass from skrub import DataOp +from skrub._data_ops._evaluation import find_scoring_node from stratum._config import FLAGS from stratum.optimizer._optimize import SearchConfig, optimize -from stratum.optimizer.logical._scoring import resolve_scoring +from stratum.optimizer.logical._scoring import resolve_declared_scoring, resolve_scoring from stratum.runtime._scheduler import SequentialScheduler from stratum.frontend._skrub_graph import get_data from time import perf_counter @@ -20,8 +21,15 @@ class RunTimings: #TODO: Rename this file def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, env=None): - """Perform grid search with cross-validation on a DataOp DAG. ``scoring`` is required.""" - if scoring is None: + """Perform grid search with cross-validation on a DataOp DAG. + + ``scoring`` is required unless the DAG declares its metric with + ``.skb.with_scoring()``; an explicit ``scoring`` overrides that, as in skrub. + """ + declared = scoring is None and find_scoring_node(dag) is not None + if declared: + metric = resolve_declared_scoring(dag) + elif scoring is None: # A search ranks a set, so every candidate has to be measured the same way. # `.skb.make_grid_search()` still honours scoring=None, by handing the call to # skrub. See docs/adr/0004-a-search-always-names-its-metric.md. @@ -39,8 +47,8 @@ def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, en env[k] = v # The scorer is plan-time state: it decides what the plan's last operator computes, # and an unusable `scoring=` fails here rather than after a fold has been fitted. - search = SearchConfig(metric=resolve_scoring(scoring), - return_predictions=return_predictions) + search = SearchConfig(metric=metric if declared else resolve_scoring(scoring), + return_predictions=return_predictions, declared_kwargs=declared) # Resolve variables to constants at compile time, so the scheduler runs # without an environment. optimize_start = perf_counter() diff --git a/stratum/adapters/table_vectorizer.py b/stratum/adapters/table_vectorizer.py index 8533513a..d4ef768f 100644 --- a/stratum/adapters/table_vectorizer.py +++ b/stratum/adapters/table_vectorizer.py @@ -26,6 +26,7 @@ from skrub._clean_categories import CleanCategories from skrub._clean_null_strings import CleanNullStrings from skrub._drop_uninformative import DropUninformative +from skrub._duration_to_float import DurationToFloat from skrub._join_utils import pick_column_names from skrub._select_cols import Drop from skrub._table_vectorizer import ( @@ -181,6 +182,7 @@ def _preprocessors(self): drop_if_unique=self.drop_if_unique, ), ToDatetime(format=self.datetime_format), + DurationToFloat(), ToFloat(), CleanCategories(), ToStr(), diff --git a/stratum/optimizer/_optimize.py b/stratum/optimizer/_optimize.py index 3bdc7452..fe173d07 100644 --- a/stratum/optimizer/_optimize.py +++ b/stratum/optimizer/_optimize.py @@ -1,5 +1,6 @@ from skrub._data_ops._evaluation import _Graph from skrub._data_ops import DataOp +from skrub._data_ops._data_ops import Apply from skrub._data_ops._subsampling import SubsamplePreviews from collections import deque, defaultdict from dataclasses import dataclass @@ -7,8 +8,8 @@ from .logical._op_cse import apply_op_cse from .logical._dataframe_ops import extract_dataframe_op, add_splitting_op from .logical._numeric_ops import extract_numeric_op -from .logical._candidate_ops import CollectCandidatesOp, ScoreCandidatesOp -from .logical._ops import BaseEstimatorOp, ChoiceOp, Op, OperandRef, as_op, check_choices_not_shared +from .logical._candidate_ops import CollectCandidatesOp, DeclaredScoringOp, ScoreCandidatesOp +from .logical._ops import BaseEstimatorOp, ChoiceOp, Op, OperandRef, as_op, check_choices_not_shared, remap_operand_refs from .logical._split_ops import SplitOutput from ._op_utils import clone_sub_dag, find_choice_naive, replace_op_in_outputs, show_graph, topological_iterator, validate_dag from ._explain import explain_linear_plan @@ -63,6 +64,9 @@ class SearchConfig: metric: Any return_predictions: bool = False + #: The metric is the one the plan declares with `.skb.with_scoring()`, so the + #: scorer also reads the kwargs declared with it, computed on each test fold. + declared_kwargs: bool = False class OptConfig(): @@ -311,6 +315,7 @@ def convert_to_ops(dag: DataOp, env: dict = None) -> Op: check_choices_not_shared(nodes.values()) order = topological_traverse(nodes, parents, children) root_id = order[-1] + feeds_estimator = _apply_fed_nodes(nodes, parents, order) # id(DataOp) -> Op. Keyed by DataOp identity (not the graph's node keys) so # as_op's operand binder can resolve inputs found in the impl fields directly. @@ -325,7 +330,8 @@ def convert_to_ops(dag: DataOp, env: dict = None) -> Op: input_key = children.get(node_key, [])[0] ids_to_ops[id(skrub_op)] = ids_to_ops[id(nodes[input_key])] continue - ids_to_ops[id(skrub_op)] = as_op(skrub_op, ids_to_ops, env) + ids_to_ops[id(skrub_op)] = as_op(skrub_op, ids_to_ops, env, + feeds_estimator=node_key in feeds_estimator) root = ids_to_ops[id(nodes[root_id])] log_time("conversion took", start) @@ -333,6 +339,24 @@ def convert_to_ops(dag: DataOp, env: dict = None) -> Op: return root +def _apply_fed_nodes(nodes: dict, parents: dict, order: list) -> set: + """Keys of the nodes whose value some `.skb.apply()` consumes, directly or not. + + skrub decides at evaluation time whether an Apply is the last estimator: if + another Apply is still waiting on its output, it transforms (or predicts) in + every response pass instead of answering the scorer. The plan settles the same + question once, from the graph: any downstream Apply, through any of its fields, + counts. skrub's answer can depend on the order its traversal first reaches a + node shared by an Apply and a non-Apply path; this static rule ignores that. + """ + fed = set() + for key in reversed(order): + if any(isinstance(nodes[p]._skrub_impl, Apply) or p in fed + for p in parents.get(key, ())): + fed.add(key) + return fed + + def install_candidate_set(root: Op, search: SearchConfig | None) -> Op: """Terminate the plan in a candidate-set operator. @@ -350,6 +374,10 @@ def install_candidate_set(root: Op, search: SearchConfig | None) -> Op: replacing = isinstance(root, ChoiceOp) names = root.make_outcome_names() if replacing else ["default"] candidates = list(root.inputs) if replacing else [root] + # A `.skb.with_scoring()` tail is metadata for the scorer, not part of what a + # candidate produces: the candidate is the value it passes through. + declared = [c if isinstance(c, DeclaredScoringOp) else None for c in candidates] + candidates = [d.inputs[0] if d is not None else c for c, d in zip(candidates, declared)] if search is None or search.metric is None: node = CollectCandidatesOp(names) @@ -364,23 +392,59 @@ def install_candidate_set(root: Op, search: SearchConfig | None) -> Op: node = ScoreCandidatesOp(names, search.metric, response_mode=mode, emit_predictions=search.return_predictions) node.inputs = list(candidates) - for candidate in candidates: + for candidate, d in zip(candidates, declared): if replacing: # The choice is the root, so it is the only consumer being redirected. A # rebuild rather than `replace_output` because one operator may fill two # candidate slots and therefore be redirected twice. - candidate.outputs = [node if out is root else out for out in candidate.outputs] + candidate.outputs = [node if out is root or out is d else out + for out in candidate.outputs] else: + candidate.outputs = [out for out in candidate.outputs if out is not d] candidate.add_output(node) if isinstance(node, ScoreCandidatesOp): _, y_op = fold node.y_ref = OperandRef(node.add_input(y_op)) y_op.add_output(node) + if search.declared_kwargs: + if any(d is None for d in declared): + raise ValueError( + "not every candidate ends in `.skb.with_scoring()`, so the declared" + " scorer cannot score them all. Pass `scoring=` instead.") + node.kwargs_refs = [_rebind(d.scorers[0]["kwargs"], d.inputs, node) + for d in declared] + for d in declared: + if d is not None: + # The kwargs the scorer now reads have it as their consumer; the rest of + # the declaration (and a search scoring with another metric) is dead. + d.detach_and_prune() log_time("installing the candidate set took", start) return node +def _rebind(value, operands: list[Op], consumer: Op): + """``value`` with its OperandRefs into ``operands`` re-pointed at inputs of ``consumer``.""" + mapping = {} + for k in _referenced(value): + if k not in mapping: + mapping[k] = consumer.add_input(operands[k]) + operands[k].add_output(consumer) + return remap_operand_refs(value, mapping) + + +def _referenced(value): + """The operand indices ``value`` refers to, through nested tuples, lists and dicts.""" + if isinstance(value, OperandRef): + yield value.k + elif isinstance(value, (list, tuple)): + for v in value: + yield from _referenced(v) + elif isinstance(value, dict): + for v in value.values(): + yield from _referenced(v) + + def _response_mode(metric, names: list[str], candidates: list[Op]) -> str: """The response method the scoring pass runs, chosen once here. diff --git a/stratum/optimizer/logical/_aggregation_ops.py b/stratum/optimizer/logical/_aggregation_ops.py index 1967afed..39308b6e 100644 --- a/stratum/optimizer/logical/_aggregation_ops.py +++ b/stratum/optimizer/logical/_aggregation_ops.py @@ -475,10 +475,7 @@ def _detach_and_rewire(new_op: AggregateOp, df: Op, folder: _Folder, replaced. Downstream consumers are rewired by the caller. """ for node in folder.absorbed: - for inp in node.inputs: - inp.outputs = [o for o in inp.outputs if o is not node] - node.inputs = [] - node.outputs = [] + node.detach() replaced_ids = {id(node) for node in replaced if node is not None} for producer in (df, *folder.leaf_ops): producer.outputs = [o for o in producer.outputs diff --git a/stratum/optimizer/logical/_base.py b/stratum/optimizer/logical/_base.py index 16b3591e..ebb824c9 100644 --- a/stratum/optimizer/logical/_base.py +++ b/stratum/optimizer/logical/_base.py @@ -366,6 +366,32 @@ def replace_output_of_inputs(self, new_output): for in_ in self.inputs: in_.replace_output(self, new_output) + def detach(self) -> None: + """Take this node out of the graph: drop it from its producers' outputs and + clear its own edges. Its consumers are the caller's to rewire, before or after. + + Producers left without a consumer stay as they are, which is what a rewrite + wants while it is about to hand them a new consumer; see ``detach_and_prune``. + """ + for in_ in self.inputs: + in_.outputs = [out for out in in_.outputs if out is not self] + self.inputs = [] + self.outputs = [] + + def detach_and_prune(self) -> None: + """``detach`` this node, then every producer that leaves without a consumer. + + A producer nobody consumes is unreachable from the root, but it is still in + its own producers' outputs, and the topological walk rejects that edge. So the + dead cone goes with the node. Call it only once every surviving producer has + its new consumer: one whose last edge is momentarily gone would be pruned too. + """ + producers = list(dict.fromkeys(self.inputs)) + self.detach() + for in_ in producers: + if not in_.outputs: + in_.detach_and_prune() + def clone(self): if getattr(self.__class__, "fields", None) is None: raise NotImplementedError(f"Cloning of {self.__class__.__name__} objects is not implemented yet. Please implement it.") diff --git a/stratum/optimizer/logical/_candidate_ops.py b/stratum/optimizer/logical/_candidate_ops.py index 2e6e347d..60ddb44c 100644 --- a/stratum/optimizer/logical/_candidate_ops.py +++ b/stratum/optimizer/logical/_candidate_ops.py @@ -20,7 +20,7 @@ import polars as pl from stratum.optimizer.logical import _schema -from stratum.optimizer.logical._base import OutputType +from stratum.optimizer.logical._base import OutputType, _resolve_operand from stratum.optimizer.logical._ops import FITTING_MODE, Op, OperandRef from stratum.optimizer.logical._scoring import Metric, read_response @@ -30,6 +30,27 @@ _MAX_LABEL = 50 +class DeclaredScoringOp(Op): + """``pred.skb.with_scoring(...)``: ``pred`` unchanged, plus the scorers it declares. + + The scorers keep their skrub shape (``{"scoring", "kwargs", "name"}``), with every + DataOp in them bound to an operand of this op, so a ``sample_weight`` computed from + the data is part of the plan. Like skrub's ``Scoring`` node it passes ``pred`` + through; the candidate set strips it and, when the search scores with the declared + metric, reads the kwargs operands per fold (see ``install_candidate_set``). + """ + + logical_family = "DeclaredScoring" + fields = ["scorers"] + + def __init__(self, scorers: list, inputs: list | None = None): + super().__init__(inputs=list(inputs) if inputs else []) + self.scorers = scorers + + def process(self, mode: str, inputs: list): + return inputs[0] + + class CandidateSetOp(Op): """Fan-in over every candidate's output, carrying one label per candidate. @@ -100,6 +121,9 @@ def __init__(self, candidate_names: list[str], metric: Metric, self.emit_predictions = emit_predictions self.output_type = OutputType.FRAME self.y_ref: OperandRef | None = None + # One entry per candidate: the kwargs its declared scorer reads, with + # OperandRefs for the ones the plan computes on the fold; None for none. + self.kwargs_refs: list | None = None def propagate_output_schema(self): """The scored table is built column by column in ``process``: one row per @@ -133,8 +157,10 @@ def process(self, mode: str, inputs: list): ids, scores, values = [], [], [] for i, name in enumerate(self.candidate_names): produced = inputs[i] + kwargs = _resolve_operand(self.kwargs_refs[i], inputs) if self.kwargs_refs else None try: - scores.append(self.metric(y_true, read_response(produced, self.response_mode))) + scores.append(self.metric(y_true, read_response(produced, self.response_mode), + **(kwargs or {}))) except Exception as e: raise RuntimeError( f"[scoring] {self.metric!r} failed on candidate {name!r}" diff --git a/stratum/optimizer/logical/_dataframe_ops.py b/stratum/optimizer/logical/_dataframe_ops.py index 297b3861..ced7f3de 100644 --- a/stratum/optimizer/logical/_dataframe_ops.py +++ b/stratum/optimizer/logical/_dataframe_ops.py @@ -3,6 +3,7 @@ from pandas import DataFrame from polars import DataFrame as PolarsDataFrame from skrub import SelectCols +from skrub.selectors._base import All import operator from stratum.optimizer.logical import _schema import pandas as pd @@ -38,7 +39,12 @@ class ConcatOp(Op): """Logical frame concatenation. Pure config -- execution is provided by the physical impls in ``physical/_concat_execs.py`` (Pandas/PolarsConcatOp), - selected at plan time.""" + selected at plan time. + + Like skrub's ``.skb.concat`` it also stacks numpy arrays, which the physical + impls detect at run time. A graph-fed operand's kind is unknown when planning, + so the op is taken to produce a frame unless its first operand is a literal + array.""" logical_family = "Concat" fields = ["first", "others", "axis"] # Add more if needed @@ -50,7 +56,8 @@ def __init__(self, first, others: list, axis): self.first = first self.others = list(others) self.axis = axis - self.output_type = OutputType.FRAME + self.output_type = (OutputType.UNKNOWN if isinstance(first, np.ndarray) + else OutputType.FRAME) def propagate_output_schema(self): """See :func:`_schema.concat_schemas` for the per-axis rules. @@ -290,8 +297,11 @@ def extract_dataframe_op(op: Op, root: Op, selection_op = True, map_op = True, op.output_type = OutputType.FRAME # skrub implements `skb.select(cols)` as an Apply of its SelectCols # transformer; surface the selector as a dedicated ColumnSelectorOp. + # Only while the Apply hands it every column: restricted by `cols` or + # `exclude_cols`, skrub passes the other columns through untouched. if (isinstance(op, TransformerOp) and isinstance(op.estimator, SelectCols) - and not op.param_refs and not isinstance(op.y, OperandRef)): + and not op.param_refs and not isinstance(op.y, OperandRef) + and isinstance(op.cols, All) and op.exclude_cols is None): new_op = make_column_selector_op(op) elif isinstance(op, ChoiceOp): diff --git a/stratum/optimizer/logical/_map_ops.py b/stratum/optimizer/logical/_map_ops.py index 03e6a8c9..e37f3d32 100644 --- a/stratum/optimizer/logical/_map_ops.py +++ b/stratum/optimizer/logical/_map_ops.py @@ -64,10 +64,7 @@ def _detach_absorbed_and_rewire(op: Op, new_op: MapOp, folder: _Folder) -> None: Downstream consumers of ``op`` are rewired by the caller. """ for node in folder.absorbed: - for inp in node.inputs: - inp.outputs = [o for o in inp.outputs if o is not node] - node.inputs = [] - node.outputs = [] + node.detach() for producer in new_op.inputs: producer.outputs = [o for o in producer.outputs if o is not op] producer.add_output(new_op) diff --git a/stratum/optimizer/logical/_numeric_rewrites.py b/stratum/optimizer/logical/_numeric_rewrites.py index 09fb3564..dddf8e45 100644 --- a/stratum/optimizer/logical/_numeric_rewrites.py +++ b/stratum/optimizer/logical/_numeric_rewrites.py @@ -108,20 +108,20 @@ def eliminate_single_op_chain(op): def eliminate_two_op_chain(op1, op2): """Remove a redundant pair of inverse ops: y = f(op2(op1(x))) -> y = f(x). - Rewires the DAG in-place so that op1's input connects directly to op2's output. + Rewires the DAG in-place so that op1's input connects directly to op2's output, + and returns that input. """ x = op1.inputs[0] - x.outputs = [out for out in x.outputs if out is not op1] replace_op_in_outputs(op2, x) + op2.detach_and_prune() + return x def eliminate_two_op_chain_root_safe(op1: Op, op2: Op, root: Op) -> Op: """Wrapper around eliminate_two_op_chain that handles the case where op2 is the root (last node) of the DAG -- returns the updated root.""" - eliminate_two_op_chain(op1, op2) - if op2 is root: - root = op1.inputs[0] - return root + x = eliminate_two_op_chain(op1, op2) + return x if op2 is root else root def replace_two_op_chain(op1: Op, op2: Op, replacement: Op): @@ -154,9 +154,8 @@ def fold_to_zero(op: Op, root: Op) -> Op: ``x`` subgraph is never computed. """ zero_op = ValueOp(0.0) - for operand in op.inputs: - operand.outputs = [out for out in operand.outputs if out is not op] replace_op_in_outputs(op, zero_op) + op.detach_and_prune() return zero_op if op is root else root @@ -167,9 +166,8 @@ def fold_to_one(op: Op, root: Op) -> Op: and rewires downstream consumers to a :class:`ValueOp` holding ``1``. """ one_op = ValueOp(1) - for operand in op.inputs: - operand.outputs = [out for out in operand.outputs if out is not op] replace_op_in_outputs(op, one_op) + op.detach_and_prune() return one_op if op is root else root diff --git a/stratum/optimizer/logical/_op_cse.py b/stratum/optimizer/logical/_op_cse.py index ec538b90..2847692a 100644 --- a/stratum/optimizer/logical/_op_cse.py +++ b/stratum/optimizer/logical/_op_cse.py @@ -80,10 +80,7 @@ def _merge_op(op: Op, canonical: Op) -> None: for out in list(op.outputs): _rebind_consumer(out, op, canonical) canonical.add_output(out) - for inp in op.inputs: - inp.outputs = [o for o in inp.outputs if o is not op] - op.inputs = [] - op.outputs = [] + op.detach() def _rebind_consumer(out: Op, old_op: Op, new_op: Op) -> None: diff --git a/stratum/optimizer/logical/_ops.py b/stratum/optimizer/logical/_ops.py index 8c4251f9..68d4f949 100644 --- a/stratum/optimizer/logical/_ops.py +++ b/stratum/optimizer/logical/_ops.py @@ -7,7 +7,8 @@ from sklearn import clone from sklearn.base import BaseEstimator from skrub._data_ops._choosing import BaseChoice, Choice, DiscretizedNumericChoice, Match -from skrub._data_ops._data_ops import DataOp, Apply, Value, CallMethod, Call, GetAttr, GetItem, BinOp as SkrubBinOp, UnaryOp as SkrubUnaryOp, Concat, SplitX, Var, _wrap_estimator +from skrub._data_ops._data_ops import DataOp, Apply, Value, CallMethod, Call, GetAttr, GetItem, BinOp as SkrubBinOp, UnaryOp as SkrubUnaryOp, Concat, Scoring, SplitX, Var, _wrap_estimator +from skrub._data_ops._utils import NULL from skrub._utils import PassThrough from pandas import DataFrame import polars as pl @@ -248,7 +249,8 @@ def process(self, mode: str, inputs: list): class BaseEstimatorOp(Op): - fields = ["estimator", "y", "cols", "how", "allow_reject", "unsupervised", "kwargs", "param_refs"] + fields = ["estimator", "y", "cols", "exclude_cols", "no_wrap", "allow_reject", "unsupervised", "kwargs", "param_refs", + "feeds_estimator"] # skrub keys `Apply.kwargs` by the estimator method the kwargs belong to, and # evaluates only the group for the method it is about to call. Subclasses name # the two groups stratum can reach: the one for the fitting call and the one @@ -270,19 +272,21 @@ def callable_kwargs_keys(cls, estimator) -> frozenset: return frozenset({cls.fit_kwargs_key, cls.call_kwargs_key}) | { m for m in cls.response_modes if hasattr(estimator, m)} - def __init__(self, estimator: BaseEstimator, y=None, cols=None, how="no-wrap", allow_reject=False, unsupervised=False, kwargs=None, param_refs=None): + def __init__(self, estimator: BaseEstimator, y=None, cols=None, exclude_cols=None, no_wrap=False, allow_reject=False, unsupervised=False, kwargs=None, param_refs=None, + feeds_estimator=False): super().__init__() if kwargs is None: kwargs = {} self.check_kwargs(kwargs) self.estimator = estimator self.original_estimator = clone(self.estimator) - # X is the implicit primary operand (OperandRef(0)); y/cols are OperandRef - # when fed by the graph, otherwise plain values. param_refs maps the names of - # estimator hyper-parameters that are graph-fed to their OperandRef. + # X is the implicit primary operand (OperandRef(0)); y/cols/exclude_cols/no_wrap + # are OperandRefs when fed by the graph, otherwise plain values. param_refs maps + # the names of estimator hyper-parameters that are graph-fed to their OperandRef. self.y = y self.cols = cols - self.how = how + self.exclude_cols = exclude_cols + self.no_wrap = no_wrap self.allow_reject = allow_reject self.unsupervised = unsupervised # {method name: kwargs for that method}, as built by `.skb.apply()`. The @@ -291,6 +295,9 @@ def __init__(self, estimator: BaseEstimator, y=None, cols=None, how="no-wrap", a # `method_kwargs` rather than stored as ready-to-splat dicts. self.kwargs = kwargs self.param_refs = param_refs if param_refs is not None else {} + # Whether another `.skb.apply()` consumes this op's output. Such an op is not + # the last estimator, so like skrub it never answers a response pass itself. + self.feeds_estimator = feeds_estimator # Which responses this op can serve, settled here rather than by a `hasattr` in # the execution path. A parameter fed by the graph is still unresolved at this # point, so the estimator is taken at its word, as it is in skrub. @@ -302,9 +309,10 @@ def effective_mode(self, mode: str) -> str: """The method ``mode`` actually calls on this op's estimator. An op that cannot serve the requested response falls back the way skrub's - ``Apply`` does: a transformer transforms, a predictor predicts. + ``Apply`` does: a transformer transforms, a predictor predicts. So does one + that feeds another estimator, whatever responses it could serve. """ - if mode == FITTING_MODE or mode in self.supported_modes: + if mode == FITTING_MODE or (mode in self.supported_modes and not self.feeds_estimator): return mode return self.fallback_mode @@ -330,11 +338,13 @@ def clone(self): estimator=estimator_new, y=self.y, cols=self.cols, - how=self.how, + exclude_cols=self.exclude_cols, + no_wrap=self.no_wrap, allow_reject=self.allow_reject, unsupervised=self.unsupervised, kwargs=clone_value(self.kwargs), param_refs=self.param_refs, + feeds_estimator=self.feeds_estimator, ) new_op.was_cloned = True return new_op @@ -361,6 +371,8 @@ def extract_args_from_inputs(self, mode: str, inputs: list): if fitting and place_holders: estm.set_params(**place_holders) cols = inputs[self.cols.k] if isinstance(self.cols, OperandRef) else self.cols + exclude_cols = _resolve_operand(self.exclude_cols, inputs) + no_wrap = _resolve_operand(self.no_wrap, inputs) # A response pass never fits, so (like skrub) the fit group is left unevaluated. # Note the difference from skrub: skrub also never *computes* what that group # references, while our plan is eager, so a sub-DAG feeding only fit kwargs @@ -375,7 +387,8 @@ def extract_args_from_inputs(self, mode: str, inputs: list): x, y, cols, - self.how, + exclude_cols, + no_wrap, self.allow_reject, self.unsupervised, (fit_kwargs, call_kwargs), @@ -454,11 +467,12 @@ def check_estm_inputs(estimator, mode, x, y): def process_estimator_task(task_data): """ Process a predictor (EstimatorOp) task in a worker process. """ - (estimator, x, y, cols, how, allow_reject, unsupervised, kwargs, mode, parallelism) = task_data + (estimator, x, y, cols, exclude_cols, no_wrap, allow_reject, unsupervised, kwargs, mode, parallelism) = task_data fit_kwargs, call_kwargs = kwargs _, x, y = check_estm_inputs(estimator, mode, x, y) if mode == FITTING_MODE: - estimator = _wrap_estimator(estimator, cols, how=how, allow_reject=allow_reject, X=x) + estimator = _wrap_estimator(estimator, cols, exclude_cols=exclude_cols, no_wrap=no_wrap, + allow_reject=allow_reject, X=x) y_arg = () if unsupervised else (y,) estimator.fit(x, *y_arg, **fit_kwargs) result = estimator.predict(x, **call_kwargs) @@ -473,12 +487,13 @@ def process_estimator_task(task_data): def process_transformer_task(task_data): """ Process a transformer (TransformerOp) task in a worker process. """ - (estimator, x, y, cols, how, allow_reject, unsupervised, kwargs, mode, parallelism) = task_data + (estimator, x, y, cols, exclude_cols, no_wrap, allow_reject, unsupervised, kwargs, mode, parallelism) = task_data fit_transform_kwargs, transform_kwargs = kwargs converted, x, y = check_estm_inputs(estimator, mode, x, y) with estimator_parallel_config(parallelism): if mode == FITTING_MODE: - estimator = _wrap_estimator(estimator, cols, how=how, allow_reject=allow_reject, X=x) + estimator = _wrap_estimator(estimator, cols, exclude_cols=exclude_cols, no_wrap=no_wrap, + allow_reject=allow_reject, X=x) y_arg = () if unsupervised else (y,) result = estimator.fit_transform(x, *y_arg, **fit_transform_kwargs) elif mode == "transform": @@ -786,34 +801,46 @@ def _iter_nested(value): yield from _iter_nested(v) -def _apply_estimator_op(impl: Apply, estimator, ids_to_ops: dict) -> Op: +def _apply_estimator_op(impl: Apply, estimator, ids_to_ops: dict, feeds_estimator: bool = False) -> Op: """Build the TransformerOp/EstimatorOp for one concrete estimator of an Apply impl. ``estimator`` is ``impl.estimator`` itself, or one outcome of it when the Apply's estimator is a ``Choice``. Each call uses its own binder so every op gets X as OperandRef(0) and its own de-duplicated inputs list. + + ``feeds_estimator`` says another Apply consumes this one's output. skrub then + fit_transforms and transforms any estimator that can, a predictor included (a + KMeans feeds on its distances, not its labels), so such an estimator becomes a + TransformerOp here. """ if estimator is None or (isinstance(estimator, str) and estimator == "passthrough"): # Same normalization skrub's _wrap_estimator applies at fit time; needed # here already because BaseEstimatorOp clones its estimator on construction. estimator = PassThrough() - estimator_class = PredictorOp if hasattr(estimator, "predict") else TransformerOp + if hasattr(estimator, "predict") and not (feeds_estimator and hasattr(estimator, "transform")): + estimator_class = PredictorOp + else: + estimator_class = TransformerOp binder = OperandBinder(ids_to_ops) binder.ref(impl.X) # OperandRef(0) param_refs = {k: binder.ref(v) for k, v in estimator.get_params().items() if isinstance(v, DataOp) and id(v) in ids_to_ops} y = _bind_or_value(binder, impl.y) cols = _bind_or_value(binder, impl.cols) + exclude_cols = _bind_or_value(binder, impl.exclude_cols) + no_wrap = _bind_or_value(binder, impl.no_wrap) kwargs = _bind_apply_kwargs(binder, estimator_class, estimator, impl.kwargs) op = estimator_class( estimator=estimator, y=y, cols=cols, - how=impl.how, + exclude_cols=exclude_cols, + no_wrap=no_wrap, allow_reject=impl.allow_reject, unsupervised=impl.unsupervised, kwargs=kwargs, param_refs=param_refs, + feeds_estimator=feeds_estimator, ) op.inputs = binder.inputs return op @@ -991,7 +1018,7 @@ def check_choices_not_shared(data_ops) -> None: # TODO: Move this to frontend package -def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None) -> Op: +def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None, feeds_estimator: bool = False) -> Op: """Convert a single skrub DataOp into an Op, building its de-duplicated ``inputs`` list and operand references in one canonical field walk. @@ -1003,6 +1030,9 @@ def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None) -> Op: compile time. A ``Var`` whose name is bound in ``env`` is then resolved to a constant ``ValueOp`` instead of a ``VariableOp``, so the scheduler needs no environment to feed it at runtime. + + ``feeds_estimator`` says another Apply consumes ``data_op``'s value (see + ``_apply_estimator_op``). """ impl = data_op._skrub_impl is_X = is_y = False @@ -1064,7 +1094,8 @@ def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None) -> Op: # correctly if choice_unrolling later combines this choice with a # downstream one. leaves = _expand_estimator_choices(impl.estimator) - outcome_ops = [_apply_estimator_op(impl, est, ids_to_ops) for _, est in leaves] + outcome_ops = [_apply_estimator_op(impl, est, ids_to_ops, feeds_estimator) + for _, est in leaves] for est_op in outcome_ops: # The trailing edge-wiring below only covers the returned op. for in_op in est_op.inputs: @@ -1072,7 +1103,7 @@ def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None) -> Op: return_op = ChoiceOp(outcome_names=[path for path, _ in leaves], append_choice_name=False, inputs=outcome_ops) else: - return_op = _apply_estimator_op(impl, impl.estimator, ids_to_ops) + return_op = _apply_estimator_op(impl, impl.estimator, ids_to_ops, feeds_estimator) elif isinstance(impl, SplitX): for field_name in impl._fields: for child in _collect_child_data_ops(getattr(impl, field_name)): @@ -1084,8 +1115,17 @@ def as_op(data_op: DataOp, ids_to_ops: dict, env: dict | None = None) -> Op: # Resolve the variable to a compile-time constant; the runtime no # longer needs the environment to feed it. return_op = ValueOp(env[impl.name]) + elif env is not None and impl.becomes_default and impl.value is not NULL: + # `skrub.var(name, value, becomes_default=True)`: an environment that + # does not bind the variable falls back to its value, as in skrub. + return_op = ValueOp(impl.value) else: return_op = VariableOp(name=impl.name, value=impl.value) + elif isinstance(impl, Scoring): + from stratum.optimizer.logical._candidate_ops import DeclaredScoringOp + binder.ref(impl.pred) # OperandRef(0): the value passed through + return_op = DeclaredScoringOp(scorers=binder.bind(impl.scorers)) + return_op.inputs = binder.inputs elif isinstance(impl, Concat): from stratum.optimizer.logical._dataframe_ops import ConcatOp first = _bind_or_value(binder, impl.first) diff --git a/stratum/optimizer/logical/_projection_ops.py b/stratum/optimizer/logical/_projection_ops.py index eb5048a1..eaa965c3 100644 --- a/stratum/optimizer/logical/_projection_ops.py +++ b/stratum/optimizer/logical/_projection_ops.py @@ -506,6 +506,5 @@ def make_string_method_op(op: MethodCallOp) -> StringMethodOp: column.add_output(new_op) accessor.outputs = [o for o in accessor.outputs if o is not op] if not accessor.outputs: - column.outputs = [o for o in column.outputs if o is not accessor] - accessor.inputs = [] + accessor.detach_and_prune() return new_op diff --git a/stratum/optimizer/logical/_relational_rewrites.py b/stratum/optimizer/logical/_relational_rewrites.py index 13974795..3b7c5815 100644 --- a/stratum/optimizer/logical/_relational_rewrites.py +++ b/stratum/optimizer/logical/_relational_rewrites.py @@ -100,17 +100,10 @@ def _promote_to_filtering_join(op: SelectionOp, how: str, key: str, build: Op, # Every producer loses the selection, including the column operand, which the # join reads by name instead and which stays in the graph for its other # consumers. - seen: set[int] = set() - for producer in op.inputs: - if id(producer) in seen: - continue - seen.add(id(producer)) - producer.outputs = [out for out in producer.outputs if out is not op] for producer in (left, build): producer.add_output(join) op.replace_input_of_outputs(join) - op.inputs = [] - op.outputs = [] + op.detach() return join if root is op else root @@ -198,8 +191,7 @@ def _replace_counted_group_keys(op: IndexAccessOp, selection: SelectionOp, # Detach the replaced cone so later graph passes see only live consumers. for old in (op, selection, sort, agg, projection): if old is not None: - old.inputs = [] - old.outputs = [] + old.detach() return keys if root is op else root diff --git a/stratum/optimizer/logical/_scoring.py b/stratum/optimizer/logical/_scoring.py index e60a6e20..5e70e51b 100644 --- a/stratum/optimizer/logical/_scoring.py +++ b/stratum/optimizer/logical/_scoring.py @@ -22,6 +22,7 @@ """ from __future__ import annotations +import dataclasses import inspect from dataclasses import dataclass, field from typing import Any, Callable, Mapping @@ -56,8 +57,10 @@ class Metric: #: Where the implementation comes from: "stratum", "sklearn" or "callable". source: str = "sklearn" - def __call__(self, y_true, values) -> float: - return self.sign * self.fn(y_true, values, **self.kwargs) + def __call__(self, y_true, values, **kwargs) -> float: + # Per-fold kwargs (from `.skb.with_scoring(kwargs=...)`) extend the scorer's + # own, as skrub merges them into a scorer's `_kwargs`. + return self.sign * self.fn(y_true, values, **{**self.kwargs, **kwargs}) def __repr__(self) -> str: return f"Metric({self.name!r}, via={self.source})" @@ -163,6 +166,42 @@ def resolve_scoring(scoring) -> Metric: return metric +def resolve_declared_scoring(dag) -> Metric: + """The :class:`Metric` declared on ``dag`` with ``.skb.with_scoring()``. + + Only the metric is resolved here; the scorer's ``kwargs`` may be computed by the + plan (``sample_weight=X["w"]``), so they are bound as plan operands and evaluated + on each test fold (see ``DeclaredScoringOp``). A declaration Stratum cannot score + with raises ``ValueError``: several scorers (a search reports one score), a + ``scoring`` that is itself computed by the plan, or ``None``, which means the + estimator's own ``score`` (see ``resolve_scoring``). + """ + from skrub._data_ops._evaluation import find_scoring_node, needs_eval + + node = find_scoring_node(dag) + if node is None: + raise ValueError("this DataOp declares no scorer with `.skb.with_scoring()`.") + scorers = node._skrub_impl.scorers + if len(scorers) != 1: + raise ValueError( + f"this DataOp declares {len(scorers)} scorers with `.skb.with_scoring()`, but" + " a search reports a single score per candidate. Pass `scoring=` to pick one.") + [info] = scorers + scoring, kwargs = info["scoring"], info["kwargs"] + if needs_eval(scoring): + raise ValueError( + "the scoring declared with `.skb.with_scoring()` is computed by the plan;" + " Stratum resolves the metric when it builds the plan. Pass a value instead.") + if isinstance(scoring, str) and scoring in _NATIVE and kwargs not in (None, {}): + # Stratum's own metrics take no keyword arguments, scikit-learn's do. + metric = _from_sklearn_name(scoring) + else: + metric = resolve_scoring(scoring) + if info["name"] is not None: + metric = dataclasses.replace(metric, name=info["name"]) + return metric + + def _is_sklearn_scorer(scoring) -> bool: """Whether ``scoring`` is a scikit-learn scorer we can take apart.""" return hasattr(scoring, "_score_func") diff --git a/stratum/optimizer/logical/_selection_ops.py b/stratum/optimizer/logical/_selection_ops.py index e4ad7084..68775589 100644 --- a/stratum/optimizer/logical/_selection_ops.py +++ b/stratum/optimizer/logical/_selection_ops.py @@ -141,13 +141,9 @@ def make_mask_selection_op(op: GetItemOp) -> SelectionOp: # Filtering rows keeps the container's kind: a masked series is still a series. sel.output_type = src.output_type - # Detach every absorbed op from the graph (remove it from its inputs' output - # lists, then clear its edges). + # Detach every absorbed op from the graph. for node in absorbed: - for inp in node.inputs: - inp.outputs = [o for o in inp.outputs if o is not node] - node.inputs = [] - node.outputs = [] + node.detach() # The source and each kept leaf op now feed the selection in place of the mask # GetItem / the detached expression nodes they used to feed. Other (external) diff --git a/stratum/optimizer/physical/_concat_execs.py b/stratum/optimizer/physical/_concat_execs.py index 13c86533..ee631f53 100644 --- a/stratum/optimizer/physical/_concat_execs.py +++ b/stratum/optimizer/physical/_concat_execs.py @@ -8,6 +8,7 @@ """ from __future__ import annotations +import numpy as np import pandas as pd import polars as pl @@ -26,11 +27,20 @@ def _resolve(self, inputs: list): axis = inputs[self.axis.k] if isinstance(self.axis, OperandRef) else self.axis return first, others, axis + def process(self, mode: str, inputs: list): + first, others, axis = self._resolve(inputs) + if isinstance(first, np.ndarray): + # skrub concatenates arrays too (since 0.10), whatever the frame backend. + return np.concatenate([first, *others], axis=axis) + return self._concat_frames(first, others, axis) + + def _concat_frames(self, first, others: list, axis): + raise NotImplementedError + @physical_impl(of=ConcatOp, backend="pandas") class PandasConcatOp(ConcatExec): - def process(self, mode: str, inputs: list): - first, others, axis = self._resolve(inputs) + def _concat_frames(self, first, others: list, axis): return pd.concat([first, *others], axis=axis) @@ -39,6 +49,5 @@ class PolarsConcatOp(ConcatExec): # pandas concat axis (0=rows, 1=cols) -> polars `how`. axis_map = {0: "diagonal_relaxed", 1: "horizontal"} - def process(self, mode: str, inputs: list): - first, others, axis = self._resolve(inputs) + def _concat_frames(self, first, others: list, axis): return pl.concat([first, *others], how=self.axis_map[axis]) diff --git a/stratum/optimizer/physical/_predictor_execs.py b/stratum/optimizer/physical/_predictor_execs.py index ca3ed946..c5a4a1ba 100644 --- a/stratum/optimizer/physical/_predictor_execs.py +++ b/stratum/optimizer/physical/_predictor_execs.py @@ -358,7 +358,7 @@ def lower_predictor(op: PredictorOp, ctx) -> PhysicalOp | None: if family is None: return None return family( - estimator=op.estimator, y=op.y, cols=op.cols, how=op.how, - allow_reject=op.allow_reject, unsupervised=op.unsupervised, - kwargs=op.kwargs, param_refs=op.param_refs, + estimator=op.estimator, y=op.y, cols=op.cols, exclude_cols=op.exclude_cols, + no_wrap=op.no_wrap, allow_reject=op.allow_reject, unsupervised=op.unsupervised, + kwargs=op.kwargs, param_refs=op.param_refs, feeds_estimator=op.feeds_estimator, ) diff --git a/stratum/optimizer/physical/_transform_execs.py b/stratum/optimizer/physical/_transform_execs.py index ee46ec3d..650f7cbd 100644 --- a/stratum/optimizer/physical/_transform_execs.py +++ b/stratum/optimizer/physical/_transform_execs.py @@ -237,15 +237,15 @@ def lower_transformer(op: TransformerOp, ctx) -> PhysicalOp | None: """ if isinstance(op.original_estimator, _SkrubTableVectorizer): return TableVectorizerOp( - estimator=op.estimator, y=op.y, cols=op.cols, how=op.how, - allow_reject=op.allow_reject, unsupervised=op.unsupervised, - kwargs=op.kwargs, param_refs=op.param_refs, + estimator=op.estimator, y=op.y, cols=op.cols, exclude_cols=op.exclude_cols, + no_wrap=op.no_wrap, allow_reject=op.allow_reject, unsupervised=op.unsupervised, + kwargs=op.kwargs, param_refs=op.param_refs, feeds_estimator=op.feeds_estimator, ) if isinstance(op.original_estimator, _SkrubStringEncoder): return StringEncoderOp( - estimator=op.estimator, y=op.y, cols=op.cols, how=op.how, - allow_reject=op.allow_reject, unsupervised=op.unsupervised, - kwargs=op.kwargs, param_refs=op.param_refs, + estimator=op.estimator, y=op.y, cols=op.cols, exclude_cols=op.exclude_cols, + no_wrap=op.no_wrap, allow_reject=op.allow_reject, unsupervised=op.unsupervised, + kwargs=op.kwargs, param_refs=op.param_refs, feeds_estimator=op.feeds_estimator, ) op.__class__ = PassthroughTransformer return op diff --git a/stratum/patching/_gridsearch.py b/stratum/patching/_gridsearch.py index 2ea3daf1..fb30c8c4 100644 --- a/stratum/patching/_gridsearch.py +++ b/stratum/patching/_gridsearch.py @@ -1,6 +1,7 @@ from skrub._data_ops._skrub_namespace import SkrubNamespace from stratum._config import FLAGS from stratum._api import grid_search as stratum_grid_search +from stratum.optimizer.logical._scoring import resolve_declared_scoring # Store reference to original method before patching _original_make_grid_search = SkrubNamespace.make_grid_search @@ -16,8 +17,9 @@ def _stratum_make_grid_search(self, *, fitted=False, keep_subsampling=False, **k scoring = kwargs.get("scoring", None) # skrub's own default, scoring=None, means the estimator's `score`. Stratum refuses # that for a search (ADR 0004), but refusing it here would break the drop-in contract - # (ADR 0002), so the call goes to unpatched skrub unchanged instead. - if FLAGS.scheduler and scoring is not None: + # (ADR 0002), so the call goes to unpatched skrub unchanged instead. The exception is + # a metric the DataOp declares with `.skb.with_scoring()`, which skrub then uses. + if FLAGS.scheduler and (scoring is not None or _declares_usable_scoring(self._data_op)): # Use Stratum's scheduler-based grid search return stratum_grid_search( dag=self._data_op, @@ -29,5 +31,14 @@ def _stratum_make_grid_search(self, *, fitted=False, keep_subsampling=False, **k return _original_make_grid_search(self, fitted=fitted, keep_subsampling=keep_subsampling, **kwargs) +def _declares_usable_scoring(data_op) -> bool: + """Whether Stratum can score with the metric ``data_op`` declares, if any.""" + try: + resolve_declared_scoring(data_op) + except ValueError: + return False + return True + + # This will be used by the patching system to replace the method make_grid_search = _stratum_make_grid_search \ No newline at end of file diff --git a/tests/adapters/test_table_vectorizer.py b/tests/adapters/test_table_vectorizer.py index b26f3655..a7fdf3f4 100644 --- a/tests/adapters/test_table_vectorizer.py +++ b/tests/adapters/test_table_vectorizer.py @@ -347,6 +347,19 @@ def test_uninformative_columns_match_and_keep_reference_metadata(): } +def test_duration_columns_become_seconds_like_skrub(): + hours = [pd.Timedelta(hours=1), pd.Timedelta(hours=2), pd.NaT, pd.Timedelta(hours=4)] + X = pd.DataFrame({"duration": hours, "number": [1.0, 2.0, 3.0, 4.0]}) + reference, reference_output, fused, fused_output = _fit_pair(X) + _assert_pair(reference, reference_output, fused, fused_output) + assert fused_output["duration"].tolist()[:2] == [3600.0, 7200.0] + + X_pl = pl.from_pandas(X) + reference = TableVectorizer().fit(X_pl) + fused = ExactFusedTableVectorizer().fit(X_pl) + assert_polars_frame_equal(reference.transform(X_pl), fused.transform(X_pl)) + + def test_empty_and_configured_null_strings_match(): X = pd.DataFrame( { diff --git a/tests/application/test_ttt_task.py b/tests/application/test_ttt_task.py index 464982d6..a65c1e4f 100644 --- a/tests/application/test_ttt_task.py +++ b/tests/application/test_ttt_task.py @@ -1377,12 +1377,16 @@ def define_pipeline(input_dir): X = X.skb.apply( NeighbourTrackerAdoption("source_domain_id", "target_domain_id", "out"), y=y, + # skrub 0.11's ApplyToCols.transform() drops transform kwargs. + no_wrap=True, fit_transform_kwargs={"edges": out_edges}, transform_kwargs={"edges": out_edges}, ) X = X.skb.apply( NeighbourTrackerAdoption("target_domain_id", "source_domain_id", "in"), y=y, + # skrub 0.11's ApplyToCols.transform() drops transform kwargs. + no_wrap=True, fit_transform_kwargs={"edges": in_edges}, transform_kwargs={"edges": in_edges}, ) diff --git a/tests/optimizer/logical/algebraic_rewrites/test_numeric.py b/tests/optimizer/logical/algebraic_rewrites/test_numeric.py index ec76c4f2..6d8cfe8d 100644 --- a/tests/optimizer/logical/algebraic_rewrites/test_numeric.py +++ b/tests/optimizer/logical/algebraic_rewrites/test_numeric.py @@ -619,6 +619,15 @@ def test_pow_zero_root_safe(self): self.assertEqual(len(out), 1) self.assertEqual(out[0].value, 1) + def test_folded_operand_with_a_shared_producer_leaves_no_dangling_edge(self): + """The folded operand's own subgraph is pruned too: its producer feeds the plan + elsewhere, so a stale edge to it used to fail the topological walk.""" + from tests.optimizer.physical.test_source_execs import run_plan + X = st.as_data_op(pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]})) + for dag in ((X["a"] + 1) * 0 + X["b"], (X["a"] - 1) ** 0 + X["b"]): + out, *_ = optimize(dag) + pd.testing.assert_series_equal(run_plan(out), dag.skb.eval(), check_names=False) + def test_pow_zero_disabled(self): """Disabling pow_zero must leave x ** 0 untouched.""" df = st.as_data_op(5) diff --git a/tests/optimizer/logical/test_apply_kwargs.py b/tests/optimizer/logical/test_apply_kwargs.py index 5d9e855d..0e2fd435 100644 --- a/tests/optimizer/logical/test_apply_kwargs.py +++ b/tests/optimizer/logical/test_apply_kwargs.py @@ -107,7 +107,7 @@ def test_empty_groups_are_dropped(self): def test_transformer_groups_are_keyed_by_method(self): X, _ = self._source() - root = convert_to_ops(X.skb.apply(StandardScaler(), how="no_wrap", + root = convert_to_ops(X.skb.apply(StandardScaler(), no_wrap=True, transform_kwargs={"copy": True})) self.assertIsInstance(root, TransformerOp) self.assertEqual(root.call_kwargs_key, "transform") @@ -119,7 +119,7 @@ def test_group_dead_for_this_estimator_kind_is_bound_but_ignored(self): # never called. We mirror that (warn, do not raise) but still bind the # DataOps inside it, or they would be left without a consumer. X, _ = self._source() - dag = X.skb.apply(StandardScaler(), how="no_wrap", fit_kwargs={"extra": X[["a"]]}) + dag = X.skb.apply(StandardScaler(), no_wrap=True, fit_kwargs={"extra": X[["a"]]}) with self.assertLogs("stratum", level="WARNING") as logs: root = convert_to_ops(dag) diff --git a/tests/optimizer/logical/test_dataframe_ops.py b/tests/optimizer/logical/test_dataframe_ops.py index a0a9eb07..b314e1c2 100644 --- a/tests/optimizer/logical/test_dataframe_ops.py +++ b/tests/optimizer/logical/test_dataframe_ops.py @@ -11,6 +11,7 @@ import unittest from contextlib import contextmanager +import numpy as np import polars as pl from stratum._config import FLAGS from stratum.optimizer._optimize import OptConfig, optimize as optimize_ @@ -154,6 +155,33 @@ def test_polars_concat(self): result = run_op(op, pl.DataFrame({"a": [1, 2]}), pl.DataFrame({"a": [3, 4]})) self.assertEqual(4, len(result)) + def test_polars_concat_of_arrays(self): + op = ConcatOp(first=OperandRef(0), others=[OperandRef(1)], axis=1) + result = run_op(op, np.ones((2, 1)), np.zeros((2, 2))) + np.testing.assert_array_equal(result, [[1, 0, 0], [1, 0, 0]]) + + +class TestConcatOpArrays(unittest.TestCase): + """skrub's `.skb.concat` stacks numpy arrays too (since 0.10).""" + + def test_pandas_concat_of_arrays(self): + op = ConcatOp(first=OperandRef(0), others=[OperandRef(1)], axis=0) + result = run_op(op, np.ones((1, 2)), np.zeros((2, 2))) + np.testing.assert_array_equal(result, [[1, 1], [0, 0], [0, 0]]) + + def test_literal_array_is_not_a_frame(self): + self.assertIs(ConcatOp(first=np.ones(2), others=[np.ones(2)], axis=0).output_type, + OutputType.UNKNOWN) + + def test_array_pipeline_matches_skrub(self): + import pandas as pd + import stratum as st + from tests.optimizer.physical.test_source_execs import run_plan + arr = st.as_data_op(pd.DataFrame({"a": [1., 2.], "b": [3., 4.]})).skb.apply_func(np.asarray) + dag = arr.skb.concat([arr], axis=1)[:, 1:] + ops, *_ = optimize_(dag) + np.testing.assert_array_equal(run_plan(ops), dag.skb.eval()) + if __name__ == "__main__": unittest.main() diff --git a/tests/optimizer/logical/test_fit_pass_planning.py b/tests/optimizer/logical/test_fit_pass_planning.py index ac0c1466..7a9aad9f 100644 --- a/tests/optimizer/logical/test_fit_pass_planning.py +++ b/tests/optimizer/logical/test_fit_pass_planning.py @@ -111,14 +111,14 @@ def test_the_candidate_set_is_dead_while_fitting(self): def test_a_step_that_must_be_fitted_is_never_marked(self): """A transformer after the predictor is fitted on the predictor's output, so it has to run in the fitting pass even though nothing there reads its result.""" - plan = self._plan(lambda p: p.skb.apply(RecordingTail(), how="no_wrap")) + plan = self._plan(lambda p: p.skb.apply(RecordingTail(), no_wrap=True)) tails = [op for op in plan if isinstance(op, BaseEstimatorOp)] self.assertTrue(tails, "expected the estimator ops in the plan") for op in tails: self.assertFalse(op.dead_in_fit, f"{op} would go unfitted") st._api.grid_search(self._pipeline( - lambda p: p.skb.apply(RecordingTail(), how="no_wrap")), scoring=SCORING) + lambda p: p.skb.apply(RecordingTail(), no_wrap=True)), scoring=SCORING) self.assertEqual(len(RecordingTail.FITS), N_SPLITS) def test_work_feeding_the_estimator_is_never_marked(self): @@ -136,7 +136,7 @@ def test_scores_are_unchanged_by_the_skipping(self): """Equivalence: skipping work the fitting pass does not need must not move a number. `ShiftTail` carries fitted state into the score, so a step wrongly skipped would either raise or shift the result rather than pass quietly.""" - tail = lambda p: p.skb.apply(ShiftTail(), how="no_wrap") + tail = lambda p: p.skb.apply(ShiftTail(), no_wrap=True) np.testing.assert_allclose(self._scores(tail, marked=True), self._scores(tail, marked=False), rtol=1e-12) diff --git a/tests/optimizer/logical/test_ops.py b/tests/optimizer/logical/test_ops.py index ba2c2d9c..086ac7d9 100644 --- a/tests/optimizer/logical/test_ops.py +++ b/tests/optimizer/logical/test_ops.py @@ -401,6 +401,42 @@ def test_single_key_binding_is_unchanged(self): self.assertEqual(OperandRef(1), found[0].key) +def _chain(*ops): + """Wire ``ops`` producer-first into a chain and return them.""" + for producer, consumer in zip(ops, ops[1:]): + consumer.add_input(producer) + producer.add_output(consumer) + return ops + + +class TestDetach(unittest.TestCase): + def test_detach_unlinks_from_producers_only(self): + source, mid, sink = _chain(Op(), Op(), Op()) + mid.detach() + self.assertEqual((mid.inputs, mid.outputs), ([], [])) + self.assertEqual(source.outputs, []) + # The consumer is the caller's to rewire. + self.assertEqual(sink.inputs, [mid]) + + def test_detach_and_prune_removes_the_dead_cone(self): + shared, dead, other = Op(), Op(), Op() + _chain(shared, dead) + dead_tail = Op() + _chain(dead, dead_tail) + _chain(shared, other) + dead_tail.detach_and_prune() + self.assertEqual(dead.inputs, []) + # A producer with another consumer stays, minus the pruned edge. + self.assertEqual(shared.outputs, [other]) + + def test_detach_and_prune_handles_a_producer_feeding_two_slots(self): + producer, consumer = Op(), Op() + consumer.inputs = [producer, producer] + producer.add_output(consumer) + consumer.detach_and_prune() + self.assertEqual(producer.outputs, []) + + class TestEdgeDedup(unittest.TestCase): def test_add_input_dedup_returns_index(self): op, a, b = Op(), ValueOp(1), ValueOp(2) diff --git a/tests/optimizer/logical/test_optimize.py b/tests/optimizer/logical/test_optimize.py index 98b3fc8c..1b6f205f 100644 --- a/tests/optimizer/logical/test_optimize.py +++ b/tests/optimizer/logical/test_optimize.py @@ -88,6 +88,15 @@ def test_unbound_variable_stays_a_variable(self): out, *_ = optimize(data[["x"]], self._NO_REWRITES, env={"other": 1}) self.assertTrue(any(isinstance(o, VariableOp) for o in out)) + def test_unbound_default_variable_takes_its_value(self): + # skrub.var(..., becomes_default=True): an env that does not bind it falls + # back to the variable's own value instead of failing. + data = st.var("data", self.df, becomes_default=True) + out, *_ = optimize(data[["x"]], self._NO_REWRITES, env={"other": 1}) + self.assertFalse(any(isinstance(o, VariableOp) for o in out)) + [value_op] = [o for o in out if isinstance(o, ValueOp)] + pd.testing.assert_frame_equal(value_op.value, self.df) + def test_resolved_plan_runs_without_environment(self): # The whole point: once resolved, the plan executes with an empty env. data = st.var("data", self.df) diff --git a/tests/optimizer/logical/test_projection_ops.py b/tests/optimizer/logical/test_projection_ops.py index fb173b36..2ad55618 100644 --- a/tests/optimizer/logical/test_projection_ops.py +++ b/tests/optimizer/logical/test_projection_ops.py @@ -538,6 +538,19 @@ def test_select_by_names_converts(self): ops = optimize(data.skb.select(["a"]), OptConfig(dataframe_ops=True)) self.assertEqual(["a"], self._one(ops, ColumnSelectorOp).selector) + def test_select_cols_on_restricted_columns_is_not_a_projection(self): + """Applied to part of the frame, skrub's SelectCols passes the rest through.""" + from skrub import SelectCols + from tests.optimizer.physical.test_source_execs import run_plan + data = st.as_data_op(self.df) + for kwargs in ({"cols": ["a"]}, {"exclude_cols": ["a"]}): + with self.subTest(**kwargs): + dag = data.skb.apply(SelectCols([]), **kwargs) + ops = optimize(dag, OptConfig(dataframe_ops=True)) + self.assertEqual([], [o for o in ops if isinstance(o, ColumnSelectorOp)]) + self.assertEqual(sorted(dag.skb.eval().columns), + sorted(run_plan(ops).columns)) + class TestColumnSelectorProcess(unittest.TestCase): """ColumnSelectorOp resolves the selector at fit and reuses the list at predict.""" diff --git a/tests/optimizer/physical/test_generic_execs.py b/tests/optimizer/physical/test_generic_execs.py index fc40a377..0f551be2 100644 --- a/tests/optimizer/physical/test_generic_execs.py +++ b/tests/optimizer/physical/test_generic_execs.py @@ -32,7 +32,7 @@ def test_generic_ops_lower_in_place_and_execute(): def test_generic_transformer_lowers_in_place(): - op = TransformerOp(estimator=StandardScaler(), cols=selectors.all(), how="no_wrap") + op = TransformerOp(estimator=StandardScaler(), cols=selectors.all(), no_wrap=True) with pytest.raises(NotImplementedError, match="must be lowered"): op.process("fit_transform", [pd.DataFrame({"x": [1.0, 2.0]})]) original_id = id(op) diff --git a/tests/optimizer/physical/test_predictor_execs.py b/tests/optimizer/physical/test_predictor_execs.py index 6d753679..9fc53fc8 100644 --- a/tests/optimizer/physical/test_predictor_execs.py +++ b/tests/optimizer/physical/test_predictor_execs.py @@ -180,7 +180,7 @@ def predict(self, X): def test_any_estimator_passes_through(): X, y = make_regression(n_samples=20, n_features=2, random_state=0) - op = PredictorOp(estimator=_UserDefinedPredictor(), y=y, cols=All(), how="no_wrap") + op = PredictorOp(estimator=_UserDefinedPredictor(), y=y, cols=All(), no_wrap=True) select_implementations(op, _ctx()) assert type(op) is PassthroughPredictor @@ -330,7 +330,7 @@ def test_graph_fed_catboost_parameter_is_only_applied_while_fitting(): estimator=_estimator("catboost", "CatBoostRegressor", iterations=1, verbose=0), y=y, cols=All(), - how="no_wrap", + no_wrap=True, param_refs={"iterations": OperandRef(1)}, ) @@ -352,7 +352,7 @@ def set_params(self, **params): def test_graph_fed_parameter_is_applied_to_a_fresh_estimator_every_fold(): """Each fit starts from an unfitted clone, so the next fold can configure it.""" X, y = make_regression(n_samples=20, n_features=2, random_state=0) - op = PredictorOp(estimator=_RefusesRefitParams(), y=y, cols=All(), how="no_wrap", + op = PredictorOp(estimator=_RefusesRefitParams(), y=y, cols=All(), no_wrap=True, param_refs={"alpha": OperandRef(1)}) op.process("fit_transform", [X, 2.0]) diff --git a/tests/optimizer/physical/test_table_vectorizer.py b/tests/optimizer/physical/test_table_vectorizer.py index a7d05923..3ac3efbc 100644 --- a/tests/optimizer/physical/test_table_vectorizer.py +++ b/tests/optimizer/physical/test_table_vectorizer.py @@ -58,7 +58,7 @@ def _op(vectorizer): return TableVectorizerOp( estimator=vectorizer, cols=selectors.all(), - how="no_wrap", + no_wrap=True, ) diff --git a/tests/runtime/test_apply_kwargs.py b/tests/runtime/test_apply_kwargs.py index d1c87925..38e7e589 100644 --- a/tests/runtime/test_apply_kwargs.py +++ b/tests/runtime/test_apply_kwargs.py @@ -122,7 +122,7 @@ def _run_recording_pipeline(self, **apply_kwargs): data = st.as_data_op(self.df) y = data["y"].skb.mark_as_y() X = data.drop(columns=["y", "w"]).skb.mark_as_X() - transformed = X.skb.apply(RecordingTransformer(), how="no_wrap", **apply_kwargs) + transformed = X.skb.apply(RecordingTransformer(), no_wrap=True, **apply_kwargs) grid_search(transformed.skb.apply(LinearRegression(), y=y), cv=2, scoring=self.SCORING) return dict(RecordingTransformer.CALLS) @@ -146,7 +146,7 @@ def _pipeline(self): y = data["y"].skb.mark_as_y() X = data.drop(columns=["y", "w"]).skb.mark_as_X() - parts = X.skb.apply(SplitOffEvalSet(), y=y, how="no_wrap") + parts = X.skb.apply(SplitOffEvalSet(), y=y, no_wrap=True) model = LGBMRegressor(n_estimators=200, learning_rate=0.1, random_state=0, verbose=-1) return parts["X"].skb.apply( diff --git a/tests/runtime/test_scoring.py b/tests/runtime/test_scoring.py index 77a3ec94..0ba0a658 100644 --- a/tests/runtime/test_scoring.py +++ b/tests/runtime/test_scoring.py @@ -10,9 +10,10 @@ import numpy as np import pandas as pd +from sklearn.cluster import KMeans from sklearn.dummy import DummyClassifier, DummyRegressor from sklearn.ensemble import RandomForestClassifier -from sklearn.linear_model import LinearRegression +from sklearn.linear_model import LinearRegression, LogisticRegression from sklearn.metrics import accuracy_score, fbeta_score, make_scorer from sklearn.model_selection import KFold, StratifiedKFold @@ -204,5 +205,110 @@ def test_score_is_taken_after_post_processing(self): 1 - at_model.results_["scores"][0], places=9) + def _stacked_pipeline(self, inner): + """`inner` feeds a second estimator, so it is not the last step of the plan.""" + data = st.as_data_op(self.clf_df) + y = data["t"].skb.mark_as_y() + X = data[["a", "b"]].skb.mark_as_X(cv=StratifiedKFold(n_splits=3), split_kwargs={}) + features = X.skb.apply(inner, y=y).skb.apply_func( + lambda p: pd.DataFrame(np.asarray(p).reshape(len(p), -1)).add_prefix("f")) + return features.skb.apply(LogisticRegression(), y=y) + + def test_inner_predictor_predicts_in_a_probability_pass(self): + """Only the last estimator answers the scorer: one feeding another predicts in + every pass, so its test-fold features are labels as they were while fitting.""" + self.assert_matches_skrub( + lambda: self._stacked_pipeline(LogisticRegression()), "roc_auc") + + def test_inner_estimator_with_transform_transforms(self): + """An inner estimator that can transform does so, in skrub: a KMeans hands on + its distances, not its labels.""" + self.assert_matches_skrub( + lambda: self._stacked_pipeline(KMeans(n_clusters=3, n_init=1, random_state=0)), + "roc_auc") + + + # --- metrics declared with `.skb.with_scoring()` --------------------------------- + + def _declared_pipeline(self, scoring, sample_weight=None): + """A two-candidate search whose metric the DataOp declares itself. + + ``sample_weight(X, y)`` builds the scorer's kwargs from the plan's own data, so + they are computed per test fold. + """ + data = st.as_data_op(self.clf_df) + y = data["t"].skb.mark_as_y() + X = data[["a", "b"]].skb.mark_as_X(cv=StratifiedKFold(n_splits=3), split_kwargs={}) + pred = X.skb.apply(LogisticRegression(C=st.choose_from([0.01, 1.0], name="C")), y=y) + kwargs = None if sample_weight is None else {"sample_weight": sample_weight(X, y)} + return pred.skb.with_scoring(scoring, kwargs=kwargs) + + def assert_declared_matches_skrub(self, make_pipeline): + ours = st._api.grid_search(make_pipeline()) + theirs = make_pipeline().skb.make_grid_search(fitted=True) + np.testing.assert_allclose(theirs.results_["mean_test_score"], + ours.results_["scores"], rtol=1e-9) + + def test_declared_metric_is_used(self): + self.assert_declared_matches_skrub(lambda: self._declared_pipeline("accuracy")) + + def test_declared_kwargs_are_computed_on_each_test_fold(self): + weights = lambda X, y: X["a"].abs() + for scoring in ("accuracy", "roc_auc", "neg_log_loss"): + with self.subTest(scoring=scoring): + self.assert_declared_matches_skrub( + lambda: self._declared_pipeline(scoring, weights)) + + def test_declared_kwargs_follow_a_choice_per_candidate(self): + """A choice in the kwargs is a candidate dimension of its own: each candidate is + scored with the weights of its own branch.""" + def make(): + data = st.as_data_op(self.clf_df) + y = data["t"].skb.mark_as_y() + X = data[["a", "b"]].skb.mark_as_X(cv=StratifiedKFold(n_splits=3), + split_kwargs={}) + column = st.as_data_op(st.choose_from(["a", "b"], name="weights")) + return X.skb.apply(LogisticRegression(), y=y).skb.with_scoring( + "accuracy", kwargs={"sample_weight": X[column].abs()}) + self.assert_declared_matches_skrub(make) + + def test_explicit_scoring_overrides_the_declared_one(self): + """As in skrub: the declared metric and its kwargs are ignored.""" + make = lambda: self._declared_pipeline("accuracy", lambda X, y: X["a"].abs()) + ours = st._api.grid_search(make(), scoring="roc_auc") + theirs = make().skb.make_grid_search(fitted=True, scoring="roc_auc") + np.testing.assert_allclose(theirs.results_["mean_test_score"], + ours.results_["scores"], rtol=1e-9) + + def test_declarations_stratum_cannot_score_with_are_refused(self): + cases = { + "2 scorers": self._declared_pipeline("accuracy").skb.with_scoring("roc_auc"), + # The estimator's own `score`, which a search never uses (ADR 0004). + "would score each candidate with its own": self._declared_pipeline(None), + } + for message, dag in cases.items(): + with self.subTest(message), self.assertRaisesRegex(ValueError, message): + st._api.grid_search(dag) + + def test_make_grid_search_uses_the_declared_metric(self): + with st.config(scheduler=True): + search = self._declared_pipeline( + "accuracy", lambda X, y: X["a"].abs()).skb.make_grid_search(fitted=True) + # Stratum's results table, not skrub's. + self.assertIn("scores", list(search.results_.columns)) + + def test_make_grid_search_hands_an_unusable_declaration_to_skrub(self): + with st.config(scheduler=True): + search = self._declared_pipeline(None).skb.make_grid_search(fitted=True) + self.assertIn("mean_test_score", list(search.results_.columns)) + + def test_a_declared_metric_does_not_change_evaluate(self): + """Without a search the declaration is dropped, kwargs subgraph included.""" + plain = st._api.evaluate(self._declared_pipeline("accuracy")) + declared = st._api.evaluate( + self._declared_pipeline("accuracy", lambda X, y: X["a"].abs())) + np.testing.assert_array_equal(np.asarray(plain), np.asarray(declared)) + + if __name__ == "__main__": unittest.main() diff --git a/uv.lock b/uv.lock index 69986d0a..b9345fd0 100644 --- a/uv.lock +++ b/uv.lock @@ -233,6 +233,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/4c/925909008ed5a988ccbb72dcc897407e5d6d3bd72410d69e051fc0c14647/charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f", size = 53402, upload-time = "2025-10-14T04:42:31.76Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -2059,10 +2068,12 @@ wheels = [ [[package]] name = "skrub" -version = "0.8.0" +version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cloudpickle" }, { name = "jinja2" }, + { name = "joblib" }, { name = "matplotlib" }, { name = "numpy" }, { name = "pandas" }, @@ -2071,9 +2082,9 @@ dependencies = [ { name = "scikit-learn" }, { name = "scipy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c5/7f/e0c6bf2bc15a38ec1cd6c6621373353d82883d4f8b1f9ce2f32a4f4d1533/skrub-0.8.0.tar.gz", hash = "sha256:379e7eaf2dc946d3f5d025c615df46e69fdf16b405c2badb5d916953b13b8dcb", size = 3631962, upload-time = "2026-03-25T10:29:06.101Z" } +sdist = { url = "https://files.pythonhosted.org/packages/22/f6/4bcaf620df6ce360e04ae177cd1586e15e5e30a19a2edff5edca343fc70a/skrub-0.11.0.tar.gz", hash = "sha256:97e986f56a82d9e071aff566c10099827fbaee4959451b61f9bdaf71c3ad6c6d", size = 4234421, upload-time = "2026-10-01T10:03:07.232Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/07/5a/1d872525938d9da54c41fa69a3bc7fa7369047101ae702ce5013dd25ff92/skrub-0.8.0-py3-none-any.whl", hash = "sha256:6e349cee08d80e2df4dd80d8d3fdd176d54905b09ba786b81a0125127f607ebc", size = 519161, upload-time = "2026-03-25T10:29:04.399Z" }, + { url = "https://files.pythonhosted.org/packages/a0/a4/cc5181b3f29e3a9db8fd5c4bde79358d7833ae9aeb53b61f9631ef8d5a54/skrub-0.11.0-py3-none-any.whl", hash = "sha256:9cf5bda8c5d35ec3b9669a451b1ae03a5789b5f8cd328f2df224b115fa9fe931", size = 802731, upload-time = "2026-10-01T10:03:02.596Z" }, ] [[package]] @@ -2157,7 +2168,7 @@ requires-dist = [ { name = "pytest", marker = "extra == 'test'", specifier = ">=9.0.3" }, { name = "pytest-cov", marker = "extra == 'test'" }, { name = "scikit-learn", specifier = "==1.8" }, - { name = "skrub", specifier = "==0.8.0" }, + { name = "skrub", specifier = "==0.11.0" }, { name = "torch", marker = "sys_platform == 'linux' and extra == 'cpu'", specifier = ">=2.13.0", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "stratum-ai", extra = "cpu" } }, { name = "torch", marker = "sys_platform == 'linux' and extra == 'cu126'", specifier = ">=2.13.0", index = "https://download.pytorch.org/whl/cu126", conflict = { package = "stratum-ai", extra = "cu126" } }, { name = "torch", marker = "sys_platform != 'linux' and extra == 'cpu'", specifier = ">=2.13.0" },